#!/usr/bin/env python3

import numpy as np
import os
from ivector_PLDA_OSI import iv_OSI
from ivector_PLDA_CSI import iv_CSI
from ivector_PLDA_SV import iv_SV
from gmm_ubm_OSI import gmm_OSI
from gmm_ubm_CSI import gmm_CSI
from gmm_ubm_SV import gmm_SV
from scipy.io.wavfile import read
import pickle
import socket
import time
from multiprocessing import Process, Queue, Pipe, SimpleQueue
import select

N_JOBS = 1
HOST = ''
PORT = 12943
END_STRING_MARKING = b'=========end'

gmm_model_list = []


test_dir = "./data/test-set"
adversarial_dir = "./adversarial-audio"
illegal_dir = "./data/illegal-set"

model_dir = "model"
spk_id_list = ["1580", "2830", "4446", "5142", "61"]  # Change to your own spk ids !!!!
gmm_model_paths = [os.path.join(model_dir, spk_id + ".gmm") for spk_id in spk_id_list]

for path in gmm_model_paths:
    with open(path, "rb") as reader:
        model = pickle.load(reader)
        gmm_model_list.append(model)


pre_model_dir = "pre-models"
ubm = os.path.join(pre_model_dir, "final.dubm")

spk_id = model[0]
print("spk_id:", spk_id)
spk_id_extra = "test-gmm-SV-" + spk_id
gmm_sv_model = gmm_SV(spk_id_extra, model, ubm)



#audio_list = []
#spk_dir = os.path.join(test_dir, spk_id)
#audio_iter = os.listdir(spk_dir)
#for i, audio_name in enumerate(audio_iter):
#    path = os.path.join(spk_dir, audio_name)
#    _, audio = read(path)
#    audio_list.append(audio)


def send_score_old(score):
        host = socket.gethostname()
        print("host:", host, flush=True)
        sh = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        sh.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
        sh.connect((host, PORT + 1))
        send_data = pickle.dumps(score)
        sh.sendall(send_data)
        sh.close()



def send_score_this_file_server(s, score):
    try: 
        conn, addr = accept_connection(s, timeout=0.1)
        send_data = pickle.dumps(score)
        print('sending', flush=True)
        conn.sendall(send_data)
    except socket.timeout as e:
        print('Exception when trying to send score:')
        print(e)

    conn.close()



def socket_bind(host, port):
    s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
    s.bind((host, port))
    return s 


def start_server():
    s = socket_bind(host, port)
    return s


def accept_connection(s, timeout=1):
    #print('Waiting for connection')
    #s.settimeout(timeout)
    #print('listen', flush=True)
    #s.listen(1)
    print('accept', flush=True)
    conn, addr = s.accept()
    print('return', flush=True)
    return conn, addr

#s = socket_bind('', PORT)
#timeout = 1
##s.settimeout(timeout)
#s.listen(1)
#s2 = socket_bind('', PORT + 1)

def pipe_full(conn):
   r, w, x = select.select([], [conn], [], 0.0)
   print("len(w):", len(w))
   return 0 == len(w)


def start_prediction_loop(conn, mp_conn):
        with conn:
            while True:
                try:
                    received_data_all = b''
                    while True:
                        #print('recv')
                        received_data = conn.recv(1024)
                        #print("received_data:", received_data)
                        received_data_all += received_data
                        if received_data_all[-len(END_STRING_MARKING):] == END_STRING_MARKING:
                            received_data_all = received_data_all[:-len(END_STRING_MARKING)]
                            break
                        elif not received_data:
                            break

                    #print("received_data_all:", received_data_all)
                    if not received_data_all: break 
                    print('outside receive loop', flush=True)
                    audio = pickle.loads(received_data_all)
                    #print("audio:", audio)
                    #audios = [audio]
                    time1 = time.time()
                    #audio_score_queue.put(audio)
                    mp_conn.send(audio)
                    score = mp_conn.recv()
                    time_to_evaluate = time.time() - time1
                    print("time_to_evaluate:", time_to_evaluate)
                    print("score:", score)


                    #send_score_this_file_server(s, score)
                    score_encoded = pickle.dumps(score)
                    conn.sendall(score_encoded)


                except (EOFError, socket.timeout) as exception:
                    print(exception)


def score_worker(mp_queue):
    print('started score worker')
    mp_connections = []
    while True:
        #for e in iter(mp_queue.get, 'STOP'):
        #    print("e:", e)
        #    mp_connections.append(e)

        while True:
            if mp_queue.empty(): break
            print('Adding mp connection')
            mp_connections.append(mp_queue.get())
        #print("mp_connections:", mp_connections)
        audios = []
        conn_indices = []
        for i, mp_conn in enumerate(mp_connections):
            recv_value = []
            #recv_value = mp_connection.recv()
            #if not audio_score_queue.empty():
            #    recv_value = audio_score_queue.get()
            if mp_conn.poll(0.001):
            #if pipe_full(mp_conn):
                #print('recv_audio_mp')
                try:
                    recv_value = mp_conn.recv()
                except EOFError as e:
                    print(e)
                    mp_connections.remove(mp_conn)
            #print("recv_value:", recv_value)
            if len(recv_value) > 0:
                #if len(recv_value) > 0:
                audios.append(recv_value)
                conn_indices.append(i)
        print("audios:", audios)

        if conn_indices:
            print("conn_indices:", conn_indices)

            _, scores = gmm_sv_model.make_decisions(audios, n_jobs=N_JOBS, debug=False)
            print("type(scores):", type(scores))
            if type(scores) != list and type(scores) != np.ndarray:
                scores = [scores]
            print("scores:", scores)
            for i, score in zip(conn_indices, scores):
                mp_connections[i].send(score)

        time.sleep(0.01)



#socket.setdefaulttimeout(1)
with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
    s.bind((HOST, PORT))
    s.listen(100)
    mp_queue = Queue()
    Process(target=score_worker, args=(mp_queue,)).start()
    while True:
        conn, addr = s.accept()
        parent_conn, child_conn = Pipe()
        #audio_score_queue = SimpleQueue()
        print('client connected')
        p = Process(target=start_prediction_loop, args=(conn, child_conn))
        p.start()
        print('adding to queue')

        #mp_queue.put(audio_score_queue)
        mp_queue.put(parent_conn)


        #mp_queue.put('STOP')

    #time.sleep(1)



