# 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 struct
import hlsvdpro
import numpy as np
import math
import multiprocessing
import threading
import itertools
import argparse

_debug_ = False

# 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)
    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


def freq_in_range(freq, ranges_string):
    if ranges_string.find(";")==-1:
        ranges = ranges_string.split("#")
        if float(ranges[0]) < freq < float(ranges[1]):
            return True
    else:
        ranges = ranges_string.split(";")
        for i in range(0, len(ranges)):
            interval = ranges[i].split("#")
            if float(interval[0]) < freq < float(interval[1]):
                return True
    return False


def compute_hlsvd(const, data):
    nsv_sought, m, ndp, dwell_time, filterType, extreme_width, hlsvd_freq_ranges = const[0:]
    correct = False
    incorrect_1 = False
    count = 0
    m_orig = m
    nsv_sought_orig = nsv_sought
    npts = len(data)
    if filterType == "WATER":
        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(freq)):
            ###if hlsvd_min_freq < result[2][i]/dwell_time < hlsvd_max_freq:
            if freq_in_range(result[2][i]/dwell_time, hlsvd_freq_ranges)==True:
                pass
            else:
                result[4][i] = 0

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

    elif filterType == "WATER+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(result[3])):
            if ((np.abs(result[3][i])*1000.0/math.pi>=extreme_width) or (freq_in_range(result[2][i]/dwell_time, hlsvd_freq_ranges)==True)):
                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())

    nro_eing = len(freq)
    if nro_eing != nsv_sought_orig:
        zeros = np.zeros(nsv_sought_orig - nro_eing)
        freq = np.concatenate([freq, zeros])
        damp = np.concatenate([damp, zeros])
        ampl = np.concatenate([ampl, zeros])
        phas = np.concatenate([phas, zeros])

    return [water, freq, damp, ampl, phas]


class ClientThread():

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

    def run(self):
        count = 0
        try:
            order = read_string(self.client)
            if order=='PING':
                write_int(self.client, 1);
                return
            if order == 'TERMINATE':
                self.terminate_event.set()  # Signal termination
                self.sock.close()
                return
            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)
                    hlsvd_freq_ranges = read_string(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)
                    if _debug_:
                        print("freq range:", hlsvd_freq_ranges)
                        print("dwell_time:",dwell_time)
                    m = read_int(self.client)   
                    batch_size = read_int(self.client)
                    batch_size = min(batch_size, int(len(realdata)/ndp))
                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)
                    final_freq = []
                    final_ampl = []
                    final_damp = []
                    final_phas = []
                    final_water = []
                    data = []
                    for i in range(0, batch_size):
                        data.append(realdata[i*ndp:(i+1)*ndp] + 1j * imagdata[i*ndp:(i+1)*ndp])

                    ###const = [nsv_sought, m, ndp, dwell_time, filterType, extreme_width, hlsvd_max_freq, hlsvd_min_freq]
                    const = [nsv_sought, m, ndp, dwell_time, filterType, extreme_width, hlsvd_freq_ranges]
                    #water, freq, damp, ampl, phas= compute_hlsvd(data, nsv_sought, m, ndp, dwell_time, filterType, extreme_width, hlsvd_max_freq, hlsvd_min_freq)
                    num_cores = min(multiprocessing.cpu_count(), batch_size)
                    print("Number of cores:", num_cores)
                    pool = multiprocessing.Pool(num_cores)
                    result = pool.starmap(compute_hlsvd, zip(itertools.repeat(const), data))
                    pool.close()
                    pool.join()
                    
                    for i in range(0, batch_size):
                        water, freq, damp, ampl, phas = result[i][0:]
                        final_freq = np.concatenate([final_freq, freq])
                        final_ampl = np.concatenate([final_ampl, ampl])
                        final_damp = np.concatenate([final_damp, damp])
                        final_phas = np.concatenate([final_phas, phas])
                        final_water = np.concatenate([final_water, water])

                        if _debug_:
                            f = open("0_freqs.txt", "a")
                            f2 = open("0_ampl.txt", "a")
                            f3 = open("0_damp.txt", "a")
                            f4 = open("0_phas.txt", "a")
                            f5 = open("0_number_of_findings.txt", "a")
                            number_of_findings = len(freq)
                            np.savetxt(f, freq/dwell_time, newline="\n")
                            np.savetxt(f2, ampl, newline="\n")
                            np.savetxt(f3, damp, newline="\n")
                            np.savetxt(f4, phas, newline="\n")
                            f5.write(str(number_of_findings)+'\n')

                            f.close()
                            f2.close()
                            f3.close()
                            f4.close()
                            f5.close()
                    write_1DArrayDouble(self.client,final_freq)
                    write_1DArrayDouble(self.client,final_damp)
                    write_1DArrayDouble(self.client,final_ampl)
                    write_1DArrayDouble(self.client,final_phas)
                    write_1DArrayDouble(self.client,final_water.real)
                    write_1DArrayDouble(self.client,final_water.imag)
                    
                #order = read_string(self.client)
        except Exception as err:
            print('Exception caught: %s\nClosing...' % err)
            if not self.terminate_event.is_set():
                self.terminate_event.set()
                

def accept(sock, terminate_event):
    while not terminate_event.is_set():
        client, _ = sock.accept()
        ClientThread(client, sock, terminate_event).run()


def main (port, host, debug):
    #Open TCP sock
    #this next 6 lines are exactly the same idea than a socket in C
    global _debug_
    _debug_ = debug
    terminate_event = threading.Event()  # Create a threading event for termination
    sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    sock.bind ((host, port))
    sock.listen(1)
    accept(sock, terminate_event)  # Pass the termination event to the accept function
    sock.close()



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


#26, 19, 3
