import os
import torch
import numpy as np
import asyncio
import json
import base64
import websockets
from io import BytesIO

import pymysql
import datetime

from PIL import Image, ImageDraw
from IPython import display

from models.mtcnn import MTCNN
from models.inception_resnet_v1 import InceptionResnetV1

device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
print('Running on device: {}'.format(device))

model = InceptionResnetV1().eval().to(device)
attendance_db = pymysql.connect(
    user='root', 
    passwd='5978', 
    host='localhost', 
    db='attendance', 
    charset='utf8'
)

lock = asyncio.Lock()
clients = set()
#processes = []

async def get_embeddings(face_list):
  global model
  x = torch.Tensor(face_list).to(device)
  yhat = model(x)
  return yhat

async def get_distance(arr1, arr2):
    distance = (arr1 - arr2).norm().item()
    return distance

def get_argmin(someone, database):
    distance = get_distance(someone, database)
    for i in range(len(distance)):
        return np.argmin(distance)
    return -1

async def register(websocket):
    global lock
    global clients
    async with lock:
        clients.add(websocket)
        remote_ip = websocket.remote_address[0]
        msg='[{ip}] connected'.format(ip=remote_ip)
        print(msg)

async def unregister(websocket):
    global lock
    global clients
    async with lock:
        clients.remove(websocket)
        remote_ip = websocket.remote_address[0]
        msg='[{ip}] disconnected'.format(ip=remote_ip)
        print(msg)

async def thread(websocket, path):
    # register(websocket) sends user_event() to websocket
    await register(websocket)
    try:
        # await websocket.send(state_event())
        async for message in websocket:
            data = json.loads(message)
            remote_ip = websocket.remote_address[0]
            if data['action'] == 'register':
                # log
                msg='[{ip}] register face'.format(ip=remote_ip)
                print(msg)

                # load json
                student_id = data['student_id']
                student_name = data['student_name']
                face = np.asarray(data['MTCNN'], dtype = np.float32)
                face = face.reshape((1,3,160,160))

                # DB에 연결
                cursor = attendance_db.cursor(pymysql.cursors.DictCursor)

                # 학생을 찾음
                sql = "SELECT student_id FROM student  WHERE student_id = %s;"
                cursor.execute(sql, (student_id))

                # DB에 학생이 없으면 등록
                if not cursor.fetchone():
                    sql = "insert into student(student_id, student_name) values (%s, %s)"
                    cursor.execute(sql, (student_id, student_name))
                    attendance_db.commit()

                # student_embedding Table에 등록
                embedding = await get_embeddings(face)
                embedding = embedding.detach().numpy().tobytes()
                embedding_date = datetime.now().strftime('%Y-%m-%d %H:%M:%S')
                sql = "insert into student_embedding(student_id, embedding_date, embedding) values (%s, %s, %s)"
                cursor.execute(sql, (student_id, embedding_date, embedding))
                attendance_db.commit()
                await websocket.send('{id} registered'.format(id=student_id))
            elif data['action'] == "verify":
                # log
                msg='[{ip}] verify face'.format(ip=remote_ip)
                print(msg)
                ###############

                # load json
                face = np.asarray(data['MTCNN'], dtype = np.float32)
                face = face.reshape((1,3,160,160))

                # embedding 구하기
                embedding = await get_embeddings(face)
                embedding = embedding.detach().numpy()
                # embedding.numpy()
                # [1, 512] numpy()임
                # np.frombuffer()로 불러오는 것이 좋을 듯.
                # DB에 연결
                cursor = attendance_db.cursor(pymysql.cursors.DictCursor)

                # 학생을 찾음
                sql = "SELECT student_id, embedding FROM student_embedding;"
                cursor.execute(sql)
                result = cursor.fetchall()
                verified_id = '0000000000'
                distance_min = 1
                for row_data in result:
                    db_embedding = np.frombuffer(row_data['embedding'], dtype=np.float32)
                    db_embedding = db_embedding.reshape((1,512))
                    distance = get_distance(embedding, db_embedding)
                    if (distance < distance_min):
                        verified_id = row_data['student_id']
                        distance_min = distance

                # 출석 데이터 전송
                data = ''
                if distance_min >= 0.6:
                    # 해당하는 사람 DB에 없음
                    print('verification failed: not in DB')
                    data = json.dumps({'state': 'fail'})
                else:
                    # 해당하는 사람 DB에 있음
                    print('verification success:', verified_id)
                    data = json.dumps({'state': 'success', 'id': verified_id})
                await websocket.send(data)
            elif data['action'] == "save_image":
                # 출석이 제대로 이뤄지지 않으면 이미지를 저장하여
                # 나중에 교강사가 출석을 확인할 수 있도록 한다
                arr = np.asarray(data['image'], dtype = np.uint8)
                # 이 데이터는 데이터베이스(과목명/일자/undefined)에 저장하는 것이 좋을듯
                # image = Image.fromarray(arr)
                # image.save('face.jpg')# storage에 데이터 저장
                remote_ip = websocket.remote_address[0]
                msg='[{ip}] save image'.format(ip=remote_ip)
                print(msg)
                ###
                await websocket.send('정해갑')
            else:
                print("unsupported event: {}", data)
    finally:
        await unregister(websocket)

print('run verification server')
start_server = websockets.serve(thread, '0.0.0.0', 8765)
asyncio.get_event_loop().run_until_complete(start_server)
asyncio.get_event_loop().run_forever()