# reads/writes data from/to a socket and answer with a water peak obtained with HLSVDPro
# Code based and modified from Amir shamaei repository: "https://github.com/amirshamaei/HLSVDPro4J"
# Authors: Federico Turco.
# Email: federico.turco@extern.insel.ch
# Date:    04.2021
import socket
import sys
import struct
import hlsvdpro
import numpy as np
import math

# reads n bytes from socket sc and stores them in the buffer recv_buf
def recv_bytes( sc, toread):
    try:
        # The bytearray() method returns a bytearray object which is an array 
        # of the given bytes.
        recv_buf = bytearray(toread)
        # returns a memory view object of the given argument
        view = memoryview(recv_buf)
       
        while toread>0:
            nbytes = sc.recv_into(view,toread)
            if nbytes == 0:
                break # the connection is closed   
            #_endif
            view = view[nbytes:]
            toread-=nbytes
        #_endwhile
    except Exception as err:   ##    except ConnectionResetError: in python 3...
        print('Exception caught: %s\nClosing...' % err)
         # ends connection with the client, waits for a need connection (client = sock.accept()[0] )    
    return recv_buf
    #_enddef   


# returns array  of fmt format 
def recv_array( sc, ftm, nItems):
       # get all bytes for nItems of fmt format from socket sc
        recv_buf=recv_bytes(sc, nItems*struct.calcsize(">"+ftm))
        
        # unpack bytes to tuple of nItmes of the type ftm

        tuple = struct.unpack_from(">%d"%nItems+ftm, recv_buf)  
        # convert tuple to array      
        arr = np.asarray(tuple)
        return arr
     

# returns string 
def read_string ( sc):
        strlen = read_int(sc)
        recv_buf=recv_bytes(sc,strlen)
        stringdata = recv_buf.decode('utf-8')
        return stringdata
      
# returns integer    
def read_int( sc):        
        return recv_array(sc,'i',1)[0]

# returns integer
def read_double( sc):
        return recv_array(sc,'d',1)[0]
    
# returns array of doubles
def read_1DArrayDouble( sc):
        nItems = read_int(sc) 
        return  recv_array(sc,'d',nItems)
    
def read_2DArrayDouble( sc, ftm):                     
        size= recv_array(sc,'i',2)
        return recv_array(sc,'d',size[0]*size[1]).reshape(size[0],size[1]) 
    
# writes double array to socket
# sc.. socket
# arr ... numpy array
def write_1DArrayDouble( sc, arr):
        #  sends a bytes object containing doubles packed (big-endian).
        sc.sendall(struct.pack('>%d'%arr.shape[0]+'d',*(arr[:])))


def write_int(sc, intg):
    #  sends a bytes object containing doubles packed (big-endian).
    sc.send(struct.pack(">i", intg))

def check_result(original, results, ndp, dwell_time):
    final_spectra = hlsvdpro.create_hlsvd_fids(results, ndp, dwell_time, sum_results=True)
    cumul = 0
    cumul_orig = 0
    for i in range(0, len(original)):
        cumul += np.abs(final_spectra[i].real)
        cumul_orig += np.abs(original[i].real)
    if cumul < 0.5*cumul_orig:
        FD = np.fft.fft(hlsvdpro._chop(final_spectra)).real
        FD_orig = np.fft.fft(hlsvdpro._chop(original)).real
        cumul = 0
        cumul_orig = 0
        for i in range(0, len(FD)):
            cumul += np.abs(FD[i].real)
            cumul_orig += np.abs(FD_orig[i].real)
        if cumul < 0.05*cumul_orig:
            return False
    return True


class ClientThread():

    def __init__( self, client, sock):
        self.client  = client
        self.sock    = sock


    def run(self):
        count = 0
        try:
            order = read_string(self.client)
            if order=='PING':
                write_int(self.client, 1);
                return
            if order=='TERMINATE':
                exit()
            while True:
                count += 1
                if order == 'FINISHED':
                    return
                if order == 'DATA':
                    hlsvd_min_freq = read_double(self.client)
                    hlsvd_max_freq = read_double(self.client)
                    realdata = read_1DArrayDouble(self.client)
                    imagdata = read_1DArrayDouble(self.client)
                    ndp = read_int(self.client)
                    nsv_sought = read_int(self.client)
                    dwell_time = read_double(self.client)
                    m = read_int(self.client)
                order = read_string(self.client)
                if order == 'PROCESS':
                    filterType = read_string(self.client)
                    extreme_width = 0
                    if filterType in ["WATER+BROAD", "BROAD"]:
                        extreme_width = read_double(self.client)
                    correct = False
                    incorrect_1 = False
                    count = 0
                    m_orig = m
                    nsv_sought_orig = nsv_sought
                    data = realdata + 1j * imagdata
                    npts = len(realdata)
                    if filterType in ["WATER+BROAD", "WATER"]:
                        while correct==False:
                            count += 1
                            removed = False
                            result = hlsvdpro.hlsvdpro(data, nsv_sought, m=ndp-m-1)
                            nsv_found, sing_vals, freq, damp, ampl, phas, _, _ = result
                            results_copy = result

                            for i in range(0, len(freq)):
                                if hlsvd_min_freq/ndp < result[2][i] < hlsvd_max_freq/ndp:
                                    removed = True
                                    pass
                                else:
                                    result[4][i] = 0

                            water = hlsvdpro.create_hlsvd_fids(result, ndp, dwell_time, sum_results=True)
                            water = np.array(water.flatten())
                            if removed == False:
                                correct = False
                            else:
                                if check_result(data, results_copy, ndp, dwell_time)==True:
                                    correct = True

                            if correct == False:
                                if incorrect_1 == True:
                                    m = m_orig
                                    if nsv_sought > 5:
                                        nsv_sought = nsv_sought - 5
                                    else:
                                        correct = True
                                else:
                                    incorrect_1 = True
                                    m = ndp - nsv_sought
                        if filterType == "WATER+BROAD":
                            result = results_copy
                            for i in range(0, len(result[3])):
                                if (np.abs(result[3][i])*1000.0/math.pi>=extreme_width) or (hlsvd_min_freq/ndp < result[2][i] < hlsvd_max_freq/ndp):
                                    pass
                                else:
                                    result[4][i] = 0

                            broad_line = hlsvdpro.create_hlsvd_fids(result, ndp, dwell_time, sum_results=True)
                            broad_line = np.array(broad_line.flatten())
                            water = broad_line

                    elif filterType == "BROAD":
                        result = hlsvdpro.hlsvdpro(data, nsv_sought, m=ndp-m-1)
                        nsv_found, sing_vals, freq, damp, ampl, phas, _, _ = result

                        for i in range(0, len(damp)):
                            if np.abs(damp[i])*1000.0/math.pi>=extreme_width:
                                pass
                            else:
                                result[4][i] = 0

                        water = hlsvdpro.create_hlsvd_fids(result, ndp, dwell_time, sum_results=True)
                        water = np.array(water.flatten())

                    if nsv_sought != nsv_sought_orig:
                        zeros = np.zeros(nsv_sought_orig - nsv_sought)
                        freq = np.concatenate([freq, zeros])
                        damp = np.concatenate([damp, zeros])
                        ampl = np.concatenate([ampl, zeros])
                        phas = np.concatenate([phas, zeros])
                    write_1DArrayDouble(self.client,freq)
                    write_1DArrayDouble(self.client,damp)
                    write_1DArrayDouble(self.client,ampl)
                    write_1DArrayDouble(self.client,phas)
                    write_1DArrayDouble(self.client,water.real)
                    write_1DArrayDouble(self.client,water.imag)
                    #data = data - water
                    #plt.plot(np.fft.fft(hlsvdpro._chop(data)).real)
                    #plt.show()
                    #input()
                #order = read_string(self.client)
        except Exception as err:
            print('Exception caught: %s\nClosing...' % err)


def accept(sock):
    client = sock.accept()[0]
    ClientThread(client,sock).run()

def main (port, host):
    #Open TCP sock
    #this next 6 lines are exactly the same idea than a socket in C
    sock = socket.socket (socket.AF_INET, socket.SOCK_STREAM)
    sock.setsockopt (socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
    sock.bind ((host, port))
    sock.listen(1)
    while True:
        try:
            accept(sock)
        except Exception as err:
            print('Exception caught: %s Closing...' % err)



if "__main__" == __name__:
    port = 33356 #some random big port, add it to the config file.
    host = 'localhost'
    main(port, host)
    print("Finished")


#26, 19, 3