import os, datetime, numpy as np
from utils import *
from EmbedNet import *
import torchvision.transforms as transforms
from detectors import S3FD
import argparse

def createParser():
    parser = argparse.ArgumentParser(description = "FaceNet");

    parser.add_argument('--config',         type=str,   default=None,   help='Config YAML file');

    ## Data loader
    parser.add_argument('--batch_size',         type=int, default=200,	help='Batch size, number of classes per batch');
    parser.add_argument('--max_img_per_cls',    type=int, default=500,	help='Maximum number of images per class per epoch');
    parser.add_argument('--nDataLoaderThread',  type=int, default=5, 	help='Number of loader threads');

    ## Training details
    parser.add_argument('--test_interval',  type=int,   default=5,     help='Test and save every [test_interval] epochs');
    parser.add_argument('--max_epoch',      type=int,   default=100,    help='Maximum number of epochs');
    parser.add_argument('--trainfunc',      type=str,   default="softmax",  help='Loss function');

    ## Optimizer
    parser.add_argument('--optimizer',      type=str,   default="adam", help='sgd or adam');
    parser.add_argument('--scheduler',      type=str,   default="steplr", help='Learning rate scheduler');
    parser.add_argument('--lr',             type=float, default=0.001,  help='Learning rate');
    parser.add_argument("--lr_decay",       type=float, default=0.90,   help='Learning rate decay every [test_interval] epochs');
    parser.add_argument('--weight_decay',   type=float, default=0,      help='Weight decay in the optimizer');

    ## Loss functions
    parser.add_argument("--hard_prob",      type=float, default=0.5,    help='Hard negative mining probability, otherwise random, only for some loss functions');
    parser.add_argument("--hard_rank",      type=int,   default=10,     help='Hard negative mining rank in the batch, only for some loss functions');
    parser.add_argument('--margin',         type=float, default=0.1,    help='Loss margin, only for some loss functions');
    parser.add_argument('--scale',          type=float, default=30,     help='Loss scale, only for some loss functions');
    parser.add_argument('--nPerClass',      type=int,   default=1,      help='Number of images per class per batch, only for metric learning based losses');
    parser.add_argument('--nClasses',       type=int,   default=8700,   help='Number of classes in the softmax layer, only for softmax-based losses');

    ## Load and save
    parser.add_argument('--initial_model',  type=str,   default="./models/amsoft_model.model",     help='Initial model weights');
    parser.add_argument('--save_path',      type=str,   default="exps/exp1", help='Path for model and logs');

    ## Training and test data
    parser.add_argument('--train_path',     type=str,   default="data/vggface2", help='Absolute path to the train set');
    parser.add_argument('--train_ext',      type=str,   default="jpg",          help='Training files extension');
    parser.add_argument('--test_path',      type=str,   default="data/test",    help='Absolute path to the test set');
    parser.add_argument('--test_list',      type=str,   default="data/test_list.csv",   help='Evaluation list');

    ## Model definition
    parser.add_argument('--model',          type=str,   default="ResNet18", help='Name of model definition');
    parser.add_argument('--nOut',           type=int,   default=512,        help='Embedding size in the last FC layer');

    ## For test only
    parser.add_argument('--eval',           dest='eval', action='store_true', help='Eval only')

    ## For server
    parser.add_argument('--server',             dest='server',  action='store_true',    help='Server mode')
    parser.add_argument('--feat_save_path',     type=str,       default='saved_feats',  help='Absolute path to the feature')
    parser.add_argument('--img_save_path',     type=str,       default='saved_img',  help='Absolute path to the image')
    parser.add_argument('--port',               type=int,       default=10000,          help='Port for the server')

    ## Distributed and mixed precision training
    parser.add_argument('--mixedprec',      dest='mixedprec',   action='store_true', help='Enable mixed precision training')

    args = parser.parse_args()

    return args

def loadParameters(model, path):
    state = model.state_dict()
    loaded_state = torch.load(path)

    for name, param in loaded_state.items():
        origname = name;
        if name not in state:
            if name not in state:
                print("%s is not in the model."%origname);
                continue;

        if state[name].size() != loaded_state[origname].size():
            print("Wrong parameter length: %s, model: %s, loaded: %s"%(origname, state[name].size(), loaded_state[origname].size()));
            continue;

        state[name].copy_(param);

class your_dataset(torch.utils.data.Dataset):
    def __init__(self, files):

        self.data   = files

        print('{:d} files in the dataset'.format(len(self.data)))

    def __getitem__(self, index):

      fname = self.data[index]
    
      try:
        # return image if read is successful
        image = cv2.imread(fname)
        image_np = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
        return image, image_np, fname
      except:
        # return empty if not successful
        return np.array([]), np.array([]), fname

    def __len__(self):
      return len(self.data)


DET = S3FD(device='cuda')
app = Flask(__name__)
args = createParser()
UNKNOWN_THRESHOLD = 0.5


s = EmbedNet(**vars(args)).cuda()
transform = transforms.Compose(
        [transforms.ToTensor(),
         transforms.Resize(256),
         transforms.CenterCrop([224,224]),
         transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])])

# trainer  = ModelTrainer(s, **vars(args))

loadParameters(s, args.initial_model)
s.eval()


@app.route('/cal', methods=['POST'])
def calculate():
    # unpack the received data
    data = pickle.loads(request.get_data())

    iname = data['name']

    image_save_path = os.path.join(args.img_save_path, iname)

    files = glob.glob(os.path.join(image_save_path, '{}*.jpg'.format(iname)))

    dataset = your_dataset(files)
    loader = torch.utils.data.DataLoader(dataset, batch_size=1, shuffle=False, num_workers=10)

    embedding_list = ''
    for data in loader:
        image     = data[0][0].numpy()
        image_np  = data[1][0].numpy()
        fname     = data[2][0].split('/')[1]
        bboxes = DET.detect_faces(image_np, conf_th=0.9, scales=[0.5])

        bsi = 100

        sx = int((bboxes[0][0]+bboxes[0][2])/2) + bsi
        sy = int((bboxes[0][1]+bboxes[0][3])/2) + bsi
        ss = int(max((bboxes[0][3]-bboxes[0][1]),(bboxes[0][2]-bboxes[0][0]))/2)

        image = numpy.pad(image,((bsi,bsi),(bsi,bsi),(0,0)), 'constant', constant_values=(110,110))

        face = image[int(sy-ss):int(sy+ss),int(sx-ss):int(sx+ss)]
        face = cv2.resize(face,(240,240))
        
        im1 = Image.fromarray(cv2.cvtColor(face, cv2.COLOR_BGR2RGB))

        inp1 = transform(im1).cuda()

        ref_feat = s(inp1).detach().cpu()

        embedding_list = torch.cat([ref_feat, embedding_list]) if embedding_list != '' else ref_feat
    
    embedding_mean    = torch.mean(embedding_list,dim=0,keepdim=True)
    torch.save(embedding_mean, os.path.join(args.feat_save_path,'{}.pt'.format(fname)))
    
    return 'success'

@app.route('/query', methods=['POST'])
def query():
    # unpack the received data
    data = pickle.loads(request.get_data())

    image = data['img']
    image_np = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)

    bboxes = DET.detect_faces(image_np, conf_th=0.9, scales=[0.5])

    if len(bboxes) != 1:
        return "fail"

    bsi = 100

    sx = int((bboxes[0][0]+bboxes[0][2])/2) + bsi
    sy = int((bboxes[0][1]+bboxes[0][3])/2) + bsi
    ss = int(max((bboxes[0][3]-bboxes[0][1]),(bboxes[0][2]-bboxes[0][0]))/2)

    image = numpy.pad(image,((bsi,bsi),(bsi,bsi),(0,0)), 'constant', constant_values=(110,110))

    face = image[int(sy-ss):int(sy+ss),int(sx-ss):int(sx+ss)]
    face = cv2.resize(face,(240,240))

    im1 = Image.fromarray(cv2.cvtColor(face, cv2.COLOR_BGR2RGB))

    inp1 = transform(im1).cuda()

    com_feat = s(inp1).detach().cpu()

    files = glob.glob(os.path.join(args.feat_save_path, '*.pt'))
    
    max_score = 0
    pname = 'none'
    for file in files:

        ref_feat = torch.load(file)
        score = F.cosine_similarity(ref_feat, com_feat)
        if(score>max_score) :
            max_score = score.item()
            pname = file.split('/')[1].split('.')[0]

        print('{} {:.2f}'.format(file,score.item()))

    if max_score < 0.1:
        max_score = 0
        pname = "Unknown"
            
    return {
        "file":pname,
        "score":max_score,
        "x1":bboxes[0][0],
        "y1":bboxes[0][1],
        "x2":bboxes[0][2],
        "y2":bboxes[0][3]
    }

@app.route('/enroll', methods=['POST'])
def enroll():
    # unpack the received data
    data = pickle.loads(request.get_data())

    iname = data['name']

    image = data['img']
    image_np = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)

    bboxes = DET.detect_faces(image_np, conf_th=0.9, scales=[0.5])

    if len(bboxes) != 1:
        return "fail"

    bsi = 100

    sx = int((bboxes[0][0]+bboxes[0][2])/2) + bsi
    sy = int((bboxes[0][1]+bboxes[0][3])/2) + bsi
    ss = int(max((bboxes[0][3]-bboxes[0][1]),(bboxes[0][2]-bboxes[0][0]))/2)

    image = numpy.pad(image,((bsi,bsi),(bsi,bsi),(0,0)), 'constant', constant_values=(110,110))

    face = image[int(sy-ss):int(sy+ss),int(sx-ss):int(sx+ss)]
    face = cv2.resize(face,(240,240))

    now = datetime.datetime.now().strftime('%y-%m-%d-%H-%M-%f')
    image_save_path = os.path.join(args.img_save_path, iname)

    if not(os.path.exists(args.feat_save_path)):
        os.makedirs(args.feat_save_path)

    if not(os.path.exists(image_save_path)):
        os.makedirs(image_save_path)

    cv2.imwrite(os.path.join(image_save_path, '{}_{}.jpg'.format(iname, now)), face)

    return "success"


if __name__ == "__main__":
    app.run(host='0.0.0.0', debug=True, port=args.port, threaded=False)