# -*- coding: utf-8 -*-
"""
------------------------------------------------------------------
 Program: Mikro Funkstrecken - Swimming Pixel
 Karlsruhe University of Applied Science
 Authors: Jonas Felhauer, Matthias Essig
 Created on Wed Nov 11 16:43:34 2020
 Description: This code should run on the PC (Host) zu analyse the
 the recorded audio from the pixel. The class ChirpSignal is used to
 create the chirp signal. The recorded audio data gets send to the PC
 over the MQTT protocol. Make sure to have a network connection to the
 same WiFi as your pixel.
------------------------------------------------------------------
"""

"""
NOTES: 
    Due to no real synchronization between the esp32 and the pc the
    calculation of the x and y position can diversify. With measuring 
    the latency of the protocol a correction gets calculated.
    Unfortunately this is not totaly sufficient. For hihger
    accuracy a better synchronization is needed. A possibility could be
    to use GPS timestamps to snchronize the time on the PC and the 
    ESP32.
"""

#--------------------------- imports -----------------------------

import paho.mqtt.client as mqtt #MQTT
import datetime #calculate latency
import time #sleep time
import scipy.io.wavfile as wav #run .wav files
import numpy as np 
import sounddevice as sd #audio output
import scipy #cross correlation
from pydub import AudioSegment
import logging #logging

#import class ChirpSinal
from ChirpSignal import ChirpSignal


#-------------------- Change logging level ------------------------

logging.basicConfig(level=logging.INFO)

#------------------ Variables for Calculation ---------------------

speed_of_sound = 343.421 #speed in m/s
distance_speaker = 1 #distance in meter

#--------------------- Chirp-Stereo-Signal ------------------------

samplerate = 44100 # sampling frequency in Hz
duration = 0.2 # duration time of chirp signal time in sec
start_freq = 3000 # start frequency in Hz
end_freq = 400 # end frequency in Hz
start_amp = 0.8 # start amplitude of signal
end_amp = 0.1 # end amplitude of signal

dtype = np.float32 #create signal with float32 datatype

#-------------------- Calculation Lateration ------------------------
######################## NOT CORRECT #################################
def lateration(distance_left, distance_right):
    #cordinate origin in the left speaker
    distance_left = float(distance_left)
    distance_right = float(distance_right)

    x = (distance_left**2-distance_right**2+distance_speaker**2)/(2*distance_speaker)
    y = np.sqrt(abs(distance_left**2-x**2)) #only positiv result has to get considered
    return x, y

#------------------- callback functions (MQTT) ---------------------

def on_log(client, userdata, level, buf):
    logging.info(buf)

def on_connect(client, userdata, flags, rc): #called when broker responds to connection request
    if rc==0:
        logging.info("Connected with result code "+str(rc))       
    else:
        logging.warning("Bad connection result code ",str(rc))
        
def on_disconnect(client, userdata, flags, rc=0): #called when broker disconnects from broker
    logging.info("DisConnected result code "+str(rc)+'\n')
    
def on_subscribe(client, userdata, mid, granted_qos): #called when broker responds to subscribe request
    logging.info("Subscribed: "+str(mid)+" "+str(granted_qos)+'\n')
    
def on_unsubscribe(client, userdata, mid): #called when broker responds to unsubscribe request
    logging.info("Unsubscribed: "+str(mid)+'\n')

def on_publish(client,userdata,rc=0): #for QoS 0 --> called when message left client
    global start_time
    start_time = datetime.datetime.now()
    logging.info("Start_Flag PUBLISHED!"+'\n')
        
def on_message_data(client, userdata, msg): #calles when message has been received on subscribed topic  
    message = msg.payload
    f.write(message)
    
def on_message_echo(client, userdata, msg): #calles when message has been received on subscribed topic  
    global end_time
    end_time = datetime.datetime.now()

    
 
#------------------------- main function --------------------------

if __name__ == "__main__":

    f = open("Data_P1.raw", "wb") #file to save the audio-data to
    
    #create sinus chirp with class ChirpSignal
    Signal_Stereo = ChirpSignal(samplerate, start_freq, end_freq, start_amp, end_amp, duration, dtype)
    Signal_Mono = ChirpSignal(samplerate, start_freq, end_freq, start_amp, end_amp, duration, dtype)
    #------------------- MQTT SETUP -------------------     
    #Pub- & Sub-Topics 
    PubTopic = "pixel/Start_Flag" #used to tell subsriber when to start recording
    SubTopic_0 = "pixel/Data_P1" #audio data gets transmitted over this topic
    SubTopic_1 = "pixel/Echo" #used to measure publish latency. Subscriber echos back
        
    Start_Flag = True #set Start Flag = True
  
    MQTT_BROKER_ADDRESS = "YOUR IP ADDRESS" #broker adress (local network)
    PORT = 1883 #port number
    
    client = mqtt.Client() #create new instanz
    client.connect(MQTT_BROKER_ADDRESS,PORT) #connect to broker
    
    #bind callback functions
    client.on_connect = on_connect 
    client.on_disconnect = on_disconnect
    client.on_subscribe = on_subscribe
    client.on_unsubscribe = on_unsubscribe
    client.on_publish = on_publish
    #bind different topic revered callback functions when message arrives
    client.message_callback_add("pixel/Data_P1", on_message_data)
    client.message_callback_add("pixel/Echo", on_message_echo)
    client.on_log = on_log
    
    #subscribe to topics
    client.subscribe(SubTopic_0) #subscribed topic
    client.subscribe(SubTopic_1) #subscribed topic
    
    time.sleep(1)
    client.loop_start() #start client loop (new thread)
    
    #------------- AUDIO OUTPUT & RECORDING -------------
    Start_Button = input("Hit Enter to start: ") #hit button to start ultrasonic positioning
    time.sleep(1.5)
    client.publish(PubTopic, Start_Flag) #send start flag for recording to pixel    

    
    start_play = datetime.datetime.now()
    sd.play(Signal_Stereo.mSignal,samplerate) #play chirp sound    
    stop_play = datetime.datetime.now()
    
    time.sleep(1) #record audio for 2 sec
    
    client.disconnect() #disconnect  
    client.loop_stop() #stop client loop  
    Start_Flag = False
    
    play_diff = stop_play.microsecond - stop_play.microsecond
    #calculation of latency    
    latency = (end_time.microsecond - start_time.microsecond)/2
   
    if(latency < 0):
        #correction when time is shifts to next second
        latency = ((1*10**(6) - start_time.microsecond) + end_time.microsecond)/2
        
    logging.info('Measured latency in micro-seconds: {}'.format(str(latency))+'\n')
    
    if(latency <= 0):
        logging.error('Error in calculating latency!')
    
    else:
        #-------------- CROSS CORRELATION ---------------
        #convert reorded .raw audio to .wav
        raw_audio = AudioSegment.from_file("Data_P1.raw", format="raw",
                                       frame_rate=44100, channels=1, sample_width=2)
    
        raw_audio.export("Data_P1.wav", format="wav") #export .wav to working directory
        wave_rec = 'Data_P1.wav'
        sr_rec, samps_rec = wav.read(wave_rec)
        
        #console output to check format
        samples_per_second_rec = sr_rec
        total_samples_rec = len(samps_rec)
        time_seconds_rec = total_samples_rec / samples_per_second_rec
        logging.info('Sampling rate: {}'.format(samples_per_second_rec))
        logging.info('Total number of samples: {}'.format(total_samples_rec))
        logging.info('Total time in seconds: {}'.format(time_seconds_rec)+'\n')
        
        if total_samples_rec > 0:         
            #-------------- prepare signals --------------
            # Cross-correlation of two arrays; arrays have to be the same size!
            # Solve this problem by adding zeros to the end of the frame until both files are the same size
            # The zero-padding is also needed because FFT method is actually circular cross-correlation,
            # meaning the signal wraps around at the ends. Thus enough zeros have to be added to get rid of
            # the overlap. At least have of the wave should be 'blank'
            
            #Downsweep
            #sd.play(Signal_Mono.mChirp,samplerate)
            samps_source_ch_1 = Signal_Mono.mChirp
            #sd.wait()
            #Upsweep
            #sd.play(Signal_Mono.mReverseChirp,samplerate)
            samps_source_ch_2 = Signal_Mono.mReverseChirp
            
            #calculate the difference between array length of recorded audio and source audio
            len_samps_rec = len(samps_rec)
            len_samps_source_ch_1 = len(samps_source_ch_1)
            len_samps_source_ch_2 = len(samps_source_ch_2)
            len_diff_ch_1 = len_samps_rec - len_samps_source_ch_1
            len_diff_ch_2 = len_samps_rec - len_samps_source_ch_2
            
            #error handling log
            logging.debug('Total number of samples: {}'.format(len_samps_rec))
            logging.debug('Total number of samples Ch1: {}'.format(len_samps_source_ch_1))
            logging.debug('Total number of samples CH2: {}'.format(len_samps_source_ch_2))
        
            #add zeros to both channels --> length of recorded & source audio is equal
            samps_source_zeros_ch_1 = np.pad(samps_source_ch_1, (0,len_diff_ch_1), 'constant')
            samps_source_zeros_ch_2 = np.pad(samps_source_ch_2, (0,len_diff_ch_2), 'constant')
            
            #error handling log
            logging.debug('Total number of samples: {}'.format(len_samps_rec))
            logging.debug('Total number of samples Ch1 with zeros: {}'.format(len(samps_source_zeros_ch_1)))
            logging.debug('Total number of samples Ch1 with zeros: {}'.format(len(samps_source_zeros_ch_2)))
            
            #------------ cross correlation Ch1 -------------
            # To get the cross-correlation, one of the signals needs to get time reversed
            samps_source_zeros_ch_1_rev = np.flip(samps_source_zeros_ch_1)
            
            corr_ch_1 = scipy.fft.ifft(scipy.fft.fft(samps_rec)*scipy.fft.fft(samps_source_zeros_ch_1_rev))
            
            corr_ch_1_abs = np.absolute(corr_ch_1)
            corr_max_ch_1 = np.argmax(corr_ch_1_abs) # position in array with greatest value
            logging.debug('Max position in aray: {}'.format(corr_max_ch_1))
            
            corr_max_time_ch_1 = corr_max_ch_1 / samples_per_second_rec  # time in array with greatest value
            logging.debug('Time of Max position: {} sec'.format(corr_max_time_ch_1)+'\n')
        
            #------------ cross correlation Ch2 -------------
            # To get the cross-correlation, one of the signals needs to get time reversed
            samps_source_zeros_ch_2_rev = np.flip(samps_source_zeros_ch_2)
            
            corr_ch_2 = scipy.fft.ifft(scipy.fft.fft(samps_rec)*scipy.fft.fft(samps_source_zeros_ch_2_rev))
            
            corr_ch_2_abs = np.absolute(corr_ch_2)
            corr_max_ch_2 = np.argmax(corr_ch_2_abs) # position in array with greatest value
            logging.debug('Max position in aray: {}'.format(corr_max_ch_2))
            
            corr_max_time_ch_2 = corr_max_ch_2 / samples_per_second_rec  # time in array with greatest value
            logging.debug('Time of Max position: {} sec'.format(corr_max_time_ch_2)+'\n')
            
            #--------------- time difference ---------------
            #calculate time difference between up- and downsweep in recorded audio
            time_diff = corr_max_time_ch_2 - corr_max_time_ch_1
            logging.info('Time difference: {}'.format(abs(time_diff))+'\n')
            
            #--------------- latency correction -------------
            #trying to correct time due to no available synchronisation between pixel and PC
            #using latency, time difference when playing sound signal, and a const delay
            corr_max_time_ch_1 = corr_max_time_ch_1 - 0.58 + latency*(10**(-6)) + play_diff*(10**(-6))
            corr_max_time_ch_2 = corr_max_time_ch_2 - 0.58 + latency*(10**(-6)) + play_diff*(10**(-6))
            logging.debug('Time of Max position Ch1: {} sec'.format(corr_max_time_ch_1))
            logging.debug('Time of Max position Ch2: {} sec'.format(corr_max_time_ch_2)+'\n')
            
            #--------------- calculate distance -------------
            ch_1_dis = corr_max_time_ch_1 * speed_of_sound
            ch_2_dis = corr_max_time_ch_2 * speed_of_sound
            
            logging.info('Distance Ch1: {} m'.format(ch_1_dis))
            logging.info('Distnace Ch2: {} m'.format(ch_2_dis)+'\n')
            
            if ch_1_dis < 0 or ch_2_dis < 0:
                logging.warning("There might be a problem with the audio ouput! Please check!")
            
            x,y = lateration(ch_1_dis, ch_2_dis)
            logging.info('X-Position: {} m'.format(x))
            logging.info('Y-Positon: {} m'.format(y)+'\n')
            
        else:
            logging.error("No audio recorded or received!")