# -*- coding: utf-8 -*-
"""
Created on Fri Mar 18 18:01:56 2016
v 24-Mar-2016

@author: Jacob_2
"""
#!/bin/python
# Filename em27_signoise.py


# this function takes ascii spectra and calculates the signal to noise
# ratios (spectra are seperated into directories with signal spectra and
# noise spectra)

# test values
sig_dir = 'sig/'
noise_dir = 'noise/'

def em27_signoise(sig_dir, noise_dir):

    # load in python packages
    import re
    import os
    import scipy.stats as sps
    import numpy as np
    from os import path
    curr_dir = os.getcwd() #get current working directory
    os.chdir('../python_subs')
    import egi_subs as egs
    import read_bin2asc_spectrum as rb2as
    os.chdir(curr_dir)
    
    plotting = 0 #indicator if plots should be made or not showing cutoffs
    pnum = 55 #which increment to plot (I haven't figured out an interative way to plot all in python yet without making a mess...I used "pause" in matlab)
    
    sig_spec_filenames = [f for f in os.listdir(sig_dir) if path.isfile(path.join(sig_dir, f))] #lists files in specified directory
    noise_spec_filenames = [f for f in os.listdir(noise_dir) if path.isfile(path.join(noise_dir, f))] 
    
    start_date = noise_spec_filenames[0][6:14] #stard date in string form
    end_date = noise_spec_filenames[-1][6:14]
    
    sig_dir_up1_ind = [m.start() for m in re.finditer('/', sig_dir)] #list of all occurences of '/'
    if len(sig_dir_up1_ind) > 1:
        sig_dir_up1 = sig_dir[0:sig_dir_up1_ind[-2]+1] #returns path name of where "sig" and "noise" directories are at
    else:
        sig_dir_up1 = '' 
    
    if len(sig_spec_filenames) != len(noise_spec_filenames):
        raise ValueError('number of noise spectra is not equal to number of signal spectra')
    
    ex_stds = 3 #exclude data outside of this std range
    tail_prop = 1-sps.norm.cdf(ex_stds) #this is proportion of data that should be in tails if data is gaussian
    
    sig_nois = egs.nan_ary((len(sig_spec_filenames),2)) #pre-allocate numpy numerical array
    
    
    i1 = 0
    for sigfname_i1 in sig_spec_filenames:
        noiseffname_i1 = path.join(noise_dir, noise_spec_filenames[i1]) #noise filename with path
        noise_spec_i1 = rb2as.read_bin2asc_spectrum((noiseffname_i1)) #read in noise file
        
        nsignal = noise_spec_i1['Spectrum']['Signal'] #just the signal component of the noise part spectrum file
        std_mm = egs.nan_ary((20, 3)) #should converge within 20 rows, pre-allocating for speed - standard deviation, mean, median
        
        std_mm[0,0] = np.std(nsignal)
        std_mm[0,1] = np.mean(nsignal)
        std_mm[0,2] = np.median(nsignal)
        
        out_points = np.floor(len(nsignal)*tail_prop) #this is the number of points allowed outside of range, can be set manually
        
        if plotting==1: #if making plots
           import matplotlib.pyplot as mplt
           nfreq = noise_spec_i1['Spectrum']['Frequency'] #get frequency values of noise spectrum
           nfreq_limits = [np.min(nfreq), np.max(nfreq)] #frequency limits of spectrum
           if i1 == pnum: #only want 1 plot... 
               mplt.figure()       
               mplt.plot(nfreq, nsignal)
               mplt.plot(nfreq_limits, [std_mm[0,2]-std_mm[0,0], std_mm[0,2]-std_mm[0,0]], 'k-') #plot -1 std
               mplt.plot(nfreq_limits, [std_mm[0,2]+std_mm[0,0]*ex_stds, std_mm[0,2]+std_mm[0,0]*ex_stds], 'k-')
               
        i2=1
        # calculate noise and remove residual point (not completely blacked out lines)
        while np.sum(nsignal>(std_mm[i2-1,2]+std_mm[i2-1,0]*ex_stds)) >out_points or np.sum(nsignal<(std_mm[i2-1,2]-std_mm[i2-1,0]*ex_stds)) >out_points:
            wthin_lim = np.all([nsignal<std_mm[i2-1,2]+std_mm[i2-1,0]*ex_stds, nsignal>(std_mm[i2-1,2]-std_mm[i2-1,0]*ex_stds)], axis=0)
            nsignal = nsignal[wthin_lim] # trim some data outside of 3 sigma range
            
            # recalculate means and stds
            std_mm[i2,0] = np.std(nsignal)
            std_mm[i2,1] = np.mean(nsignal)
            std_mm[i2,2] = np.median(nsignal)
            
            if plotting ==1:
                if i1 == pnum: #only want 1 plot... 
                    mplt.plot(nfreq_limits, [std_mm[i2,2]-std_mm[i2,0]*ex_stds, std_mm[i2,2]-std_mm[i2,0]*ex_stds], 'g-') #plot upper-cutoff bar
                    mplt.plot(nfreq_limits, [std_mm[i2,2]+std_mm[i2,0]*ex_stds, std_mm[i2,2]+std_mm[i2,0]*ex_stds], 'g-') #plot lower-cutoff bar
            
            i2=i2+1
        
        i2=i2-1 #last iteration number didn't count
        # currently removing outliers based on an assumed gaussian
        # distribution. Could also remove based on difference between mean and
        # meadian relative to std.
        
        # calculate signal
        signal_spec_i1 = rb2as.read_bin2asc_spectrum(path.join(sig_dir, sig_spec_filenames[i1])) #read in noise spectra
        signal_ordered = np.sort(signal_spec_i1['Spectrum']['Signal'],axis=0)
        near_max_sig = signal_ordered[np.int(len(signal_ordered)*0.99)] #take the 99% highest point, saying top 2% is noise
        
        
        if plotting == 1:
            if i1 == pnum: #only want 1 plot... 
                mplt.plot(nfreq_limits, [std_mm[i2,2]-std_mm[i2,0], std_mm[i2,2]-std_mm[i2,0]], 'r-') #plot +1-sigma from noise level
                mplt.plot(nfreq_limits, [std_mm[i2,2]+std_mm[i2,0], std_mm[i2,2]+std_mm[i2,0]], 'r-') 
                mplt.figure() #open a new figure
                mplt.hist(nsignal, bins=20) #plot a histogram of the noise signal (within limits)
                mplt.title(['number iterations = ', i2])
                
                mplt.figure() #open a new figure
                mplt.plot(signal_spec_i1['Spectrum']['Frequency'], signal_spec_i1['Spectrum']['Signal'], '.-')
                freq_limits = [np.min(signal_spec_i1['Spectrum']['Frequency']), np.max(signal_spec_i1['Spectrum']['Frequency'])]
                mplt.plot(freq_limits, [1, 1]*near_max_sig, 'k-')
        
        
        sig_nois[i1,0] = near_max_sig-std_mm[i2,2] #subtracting noise median from signal
        sig_nois[i1,1] = std_mm[i2,0] #noise
        
        i1 += 1 #increments by 1
        
    sig_nois_rto = sig_nois[:,0]/sig_nois[:,1] #signal to noise ratio
    
    
    # make/write snr file to insert into daily runlogs
    snr_rl_filename = sig_dir_up1 + start_date + '_' + end_date + '.snrrl' #for direct use in runlog
    snr_filename = sig_dir_up1 + start_date + '_' + end_date + '.snr' #with more decimal places
    
    ## Actual file writing
    fileID_rlin = open(snr_rl_filename, 'w') #open file for writing
    fileID_snrin = open(snr_filename, 'w') #open file for writing
    
    fileID_rlin.write(' SNR' + '\n') #adds column header
    
    
    i1 = 1
    for snr_i1 in sig_nois_rto:
        
        fileID_rlin.write('%4.f' % snr_i1 + '\n') #right size for snr spot in runlogs  
        
        if i1==len(sig_nois_rto): #don't print new line at end of file
            fileID_snrin.write('%7.2f' % snr_i1)
        else:
            fileID_snrin.write('%7.2f' % snr_i1 + '\r\n')
        
        i1 +=1
    
    fileID_snrin.close()
    fileID_rlin.close()
    
    
