import os
import fnmatch
import mcvqoe.base
import argparse
import pandas as pd
import math
import re
import numpy as np
import statistics

class Diagnose():
    """
   Diagnose Class to perform diagnostic evaluation of received
    audio files and confirm data integrity.
     
    Use a-weight, FSF, and clipping checks to inform
    user of potential problems in collected data. Flags
    trials that may require further investigation.
    
    Parameters
    ----------
    wav_dir : string
        directory of WAV files
    
    Attributes
    ----------
    tx_filename : list
        names of tx audio files
    tx_wavs   :  list
        audio for tx audio
    Trials : int
        number of trials 
    rx_rec : list 
        audio for rx recordings   
    fs : int 
        sampling rate of rx recordings       
    rx_dat : list
        names of rx recordings
    
    Methods
   ----------

    See Also
    --------

    Examples
    --------
    
    Returns
    -------
   
    """
    def __init__(self, 
                 wav_dir = ''):
        self.wav_dir = wav_dir
        # Read in a directory of test trial wav files.
        # Get all the Rx wav files 
        dir_files = os.listdir(self.wav_dir)
        naming = 'Rx*'
        all_wavs = fnmatch.filter(dir_files, naming)
        # Get total number of trials
        self.trials = len(all_wavs)
        # Create empty list for rx recordings 
        self.rx_rec = []
        self.rx_dat = []
        # Cycle through, in order
        # TODO: be clear about testing file formats. This gets confused if
        # the RX numbers reset for each file like in access time
        for n in range(1,self.trials+1):
            start = 'Rx'+str(n)+'_'
            rx_name = [s for s in all_wavs if start in s]
            rx_path = self.wav_dir + '/' + rx_name[0]
            # Check how many channels we have 
            self.fs,y_rec = mcvqoe.base.audio_read(rx_path)
            self.rx_rec.append(y_rec[:]) 
            self.rx_dat.append(rx_name[0])
            # Find all the Tx files in the wav_dir, strip 
            # the Tx and .wav off 
            tx_names = 'Tx*'
            tx_obj = fnmatch.filter(dir_files, tx_names)
            self.tx_filename =fnmatch.filter(tx_obj, '*.wav')
            # Create empty list for tx wavs 
            self.tx_wavs = []
            if tx_obj:
                # Cycle through and get all the TX files
                for k in range(0,len(self.tx_filename)):
                    tx_path = self.wav_dir + '/' + self.tx_filename[k]
                    self.fs,tx_wavfile = mcvqoe.base.audio_read(tx_path)
                    self.tx_wavs.append(tx_wavfile)
            # TODO be robust to scenarios where TX audio is not
            # saved to the data folder        
            # if there are no TX files in the dat dir, find them
    
    def aw_calc(self):
        """
        Calculates the a-weight (dBA) of each trial.
    
        Returns
        -------
        A_Weight : array
            A_Weight of every trial 
        """
        # Create empty list for a-weight values     
        a_weight = []
          
        # Get A-weight
        for k in range(0,self.trials): 
            # Calculate the A-weighted power of each recording 
            aw = mcvqoe.base.a_weighted_power(self.rx_rec[k], self.fs) 
            a_weight.append(aw)
    
        return a_weight
    
    def fsf_calc(self):   
        """
        Calculate FSF scores of each trial. 
    
        Returns
        -------
        FSF_all : list
           FSF scores of every trial 
        """
        # Create empty list for FSF scores
        fsf_all = []
        
        # Cycle through, match RX names to tx names
        # Get just the names of tx files, remove the Tx and .wav 
        tx_base = []
        for n in range(0,len(self.tx_filename)):
            tx_justname = re.sub('\.wav$', '', self.tx_filename[n])
            tx_justname = tx_justname[3:]
            tx_base.append(tx_justname)   
    
        for j in range(0,self.trials): 
            # Find RX files with the matching tx name, create groups 
            match_wavs = re.match(r'(Rx\d+_(?P<tx_base_name>[^.]+))',self.rx_dat[j])
            # find the index of the TX and RX clips to match with the lists of 
            # wav data
            tx_idx = tx_base.index(match_wavs.group('tx_base_name'))
            tx_wav = self.tx_wavs[tx_idx]
            rx_wav = self.rx_rec[j]
            # Get FSF scores for each tx-rx pair
            get_fsf = mcvqoe.base.fsf(rx_wav,tx_wav,self.fs)
            # Get just the FSF score
            fsf_all.append(get_fsf[0])   
    
        return fsf_all

    def peak_amp_calc(self):
        """
        Cycle through all rx recordings and get peak amplitude.

        Returns
        ------- 
        peak_dbfs : list
            peak amplitude of each trial, dB relative to full 
            scale 
        """
        # Create empty list for peak volume
        peak_dbfs = []
        for n in range(0,self.trials):
            # check for positive and negative clipping
            peak = max(abs(self.rx_rec[n]))
            peak_db = round(20 * math.log10(peak), 2)
            peak_dbfs.append(peak_db)
        
        return peak_dbfs
    
    def clip_flag(self,peak_dbfs):
        """
        Cycle through the peak amplitude and check
        for clipping.  
    
        Returns
        -------
        clip_flag : list
            Trials that clipped 
    
        """
        # Set up warning threshold
        vol_high = -1
        # Create empty list for setting flags
        clip_flag = []
        # Check for positive and negative clipping
        peak = np.array(peak_dbfs)
        for n in range(0,self.trials):
            # If approaching clipping, flag as a bad trial
            if peak[n] > vol_high:
                # Flag as a clipped trial
                clip_wav = 1
                # Add it to the bad list
                clip_flag.append(clip_wav) 
            elif peak[n] < vol_high:   
                # Flag as a non-clipped trial
                clip_wav = 0
                # Add it to the bad list
                clip_flag.append(clip_wav)         
        
        return clip_flag   

    def fsf_flag(self,fsf_all):
        """
        Check for trials with an FSF a certain distance
        from the mean. Use this info to find trials that 
        may have lost audio.
            
        Returns
        -------
        fsf_flag : list
            Trials that have low FSF scores or otherwise deviate 
            from the patterns of the dataset
        """
        # Create empty list for FSF flag    
        fsf_flag = []
        # Gather metrics for FSF scores
        fsf_array = np.array(fsf_all)
        fsf_mean = round(statistics.mean(fsf_array),3)
        fsf_std = round(statistics.stdev(fsf_array),3)
        for m in range(0,self.trials):
            # Cycle through FSF scores, look for trials where the score
            # stands out by being a certain distance from the mean
            if abs(fsf_array[m]-fsf_mean) > 2*fsf_std:
                # Flag the trial due to FSF score
                fsf_flag_wav = 1
                # Add it to the bad list
                fsf_flag.append(fsf_flag_wav)
            elif abs(fsf_array[m]-fsf_mean) < 2*fsf_std:
                # Do not flag
                fsf_flag_wav = 0
                # Add it to the bad list
                fsf_flag.append(fsf_flag_wav)   
            # Add a flag for low FSF scores that may be a sign of 
            # dropped audio
            if fsf_array[m] < 0.25:
                # Flag the trial due to low FSF
                fsf_flag_wav = 1
                # Add it to the bad list
                fsf_flag.add(fsf_flag_wav)      
        
        return fsf_flag

    def aw_flag(self,a_weight):
        """
        Check for trials with a dBA less than -60 dBA. Check 
        for trials with a dBA a certain distance from the 
        mean.   
    
        Returns
        -------     
        aw_flag : list
            Trials that have low dBA values (and likely lost audio)
            or otherwise deviate from the patterns of the dataset 
    
        """
        # Create empty list for a-weight flag     
        aw_flag = []
        # Calculate mean a-weight, standard deviation.
        # Use this info to find trials that may have lost
        # audio. 
        aw_array = np.array(a_weight)
        aw_lin = 10**(aw_array/20)
        aw_low = 10**(-60/20)
        aw_mean = round(statistics.mean(aw_lin),3)
        aw_std = round(statistics.stdev(aw_lin),3)
        for m in range(0,self.trials):
            # Cycle through AW values, look for trials where the values
            # stand out by being a certain distance from the mean
            if abs(aw_lin[m]-aw_mean) > 2*aw_std:
                aw_flagged = 1
                # Add it to the bad list
                aw_flag.append(aw_flagged)
            elif abs(aw_lin[m]-aw_mean) < 2*aw_std:
                aw_flagged = 0
                aw_flag.append(aw_flagged)     
            # Add a flag if the a-weight is below -60 dBA
            if aw_lin[m] < aw_low:
                aw_low = 1
                # Add it to the bad list
                aw_flag.add(aw_low)   
        
        return aw_flag
        
    def gather_diagnostics(self,a_weight,fsf_all,peak_dbfs,clip_flag,fsf_flag,aw_flag):
        """
        Create a dataframe of all diagnostic data. A-weight,
        FSF scores, max clip amplitude. Convert to json, csv
        
        Parameters
        ----------
        A_Weight : list    
            A-Weight across all RX trials 
        FSF_all : list
            FSF scores across all RX trials
        peak_dbfs : list     
            Peak amplitude across all RX trials   
        clip_flag : list
            Trials that clipped    
        fsf_flag : list
            Trials that have low FSF scores or otherwise deviate 
            from the patterns of the dataset
        aw_flag : list
            Trials that have low dBA values (and likely lost audio)
            or otherwise deviate from the patterns of the dataset         
        
        Returns
        -------
        diagnostics_csv : csv 
            CSV containing all the dat for diagnostics measurements
       """
       # Create dataframe of info
        df_diagnostics = pd.DataFrame({"RX_Name":self.rx_dat, 
                        "A_Weight":a_weight,
                        "FSF_Scores":fsf_all,
                        "Peak_Amplitude":peak_dbfs,
                        "AW_flag":aw_flag,
                        "Clip_flag":clip_flag,
                        "FSF_flag":fsf_flag})
        # Get session name and use that to name files 
        test_path, test_name =re.split("wav/+", self.wav_dir)

        # Create csv    
        diagnostics_csv = df_diagnostics.to_csv(index=False, line_terminator='\n')
        with open(os.path.join(self.wav_dir,'diagnostics.csv'),'w') as f:
            f.write(diagnostics_csv)
            
        return diagnostics_csv 

def main():   
    """
    Read in a directory of WAV files. 
    Use a-weight, FSF, and clipping checks to inform
    user of potential problems in collected data. Flags
    trials that may require further investigation.

    Parameters
    ----------
    Wav_Dir : string
        directory of WAV files

    Returns
    -------

    """
    # Input parsing
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "wav_dir",
        default=None,
        type=str,
        help="Directory to test data wav files",
    )
    # Parse input arguments
    args = parser.parse_args()
    obj = Diagnose(args.wav_dir)
    print("Measuring a-weight") 
    a_weight = obj.aw_calc()   
    aw_flag = obj.aw_flag(a_weight)
    print("Measuring FSF") 
    fsf_all = obj.fsf_calc()
    fsf_flag = obj.fsf_flag(fsf_all)
    print("Measuring peak amplitude") 
    peak_dbfs = obj.peak_amp_calc()
    clip_flag = obj.clip_flag(peak_dbfs)
    print("Creating json and csv") 
    obj.gather_diagnostics(a_weight,fsf_all,peak_dbfs,clip_flag,fsf_flag,aw_flag)

if __name__ == "__main__":
    main()