#!/usr/bin/python
import sys, os, datetime
import matplotlib
matplotlib.use('Agg')
from matplotlib.pylab import *
from scipy import *

def init_guess(p, x):
    func = (p[3] - p[0] * tanh(p[1] * (x - p[2])))
    return func

def residual(p, y, x):
    err = y - (p[3] - p[0] * tanh(p[1] * (x - p[2])))
    return err

def TECanalysis(sat_name, datename, timename):

    ### Site name ###
    sitename = 'Fukui'

    ### Basic constants define ###
    A_coef = 80.6
    c_speed = 2.998e8   # speed of light
    fr = 50.0e6         # basic frequency
    
    WL400 = c_speed / 400.e6     #wavelength of 400MHz
    WL150 = c_speed / 150.e6    
    sat_vel = 0.0
    
    skip_sec = 0.0
    
    read_sec = 999.0
    show_sec = read_sec
    
    tune_error = 0.0 # No error for Shigaraki after April 18, 2008

###    f_off400 = 2650.0   # frequency offset at 400MHz (Shigaraki)
###    f_off400 = 4200.0   # frequency offset at 400MHz (Shionomisaki)
    f_off400 = -600.0   # frequency offset at 400MHz (Fukui)

    sat_dop = sat_vel * -1.e3 / WL400    #initial Doppler of 400MHz signal

    #########################

    root_dir = '/home/beacon/beaconRX/'

    # Check parameters sent to this function
    print 'sat_name = ', sat_name, ', datename = ', datename, ', timename = ', timename

    iy=int(datename[0:4])
    im=int(datename[4:6])
    id=int(datename[6:8])

    ih=int(timename[0:2])
    imin=int(timename[2:4])
    isec=int(timename[4:6])
    
    #Date string
    dt = datetime.datetime(iy, im, id, ih, imin, isec)
    dt_string = dt.strftime('%Y-%b-%d %H:%M:%S UT')

    #filename set
    base_dir =root_dir + datename + '/'
    
##    datename = ('%4d' % iy) + ('%02d' % im) + ('%02d' % id)
##    timename = ('%02d' % ih) + ('%02d' % imin) + ('%02d' % isec)
    
    file150 = base_dir + sat_name + '_' + datename + '_' + timename + '_data150.dat'
    file400 = base_dir + sat_name + '_' + datename + '_' + timename + '_data400.dat'

    print '150MHz datafile = ', file150
    print '400MHz datafile = ', file400
    
    #check data file and size
    if os.path.exists(file400)& os.path.exists(file150):
        print '400MHz data file:', file400
        data_size = os.stat(file400)[6] # Get filesisze
        print '150MHz data file:', file150
        data_size = min(data_size, os.stat(file150)[6]) # Get filesize
        print 'Data size in bytes', data_size
        print '### Start analysis ###'
    else:
        print 'Datafiles were not found. Stop analysis.'
        return
    
    # analysis parameters set
    sample_f = 32000.0 
    nfft = 8192
    delta_f = sample_f / float(nfft)  #frequency resolution
    
    # Calculate length of observations in seconds
    if read_sec == 999.0:
        read_sec = float((data_size // (8 * sample_f)) - skip_sec)
        show_sec = read_sec
    
    # Calculate reading parameters
    
    nskip_bytes = int((sample_f * skip_sec) // nfft) * nfft * 8
    
    ntdata = int((sample_f * read_sec) // nfft) * nfft

    nparam = ntdata / nfft * 2 -1 # Number of spectral data (half interleave)
    
    npfact = 128
    npft = nfft/npfact
    npst = nfft/npfact/4
    n_phase = ntdata / npfact      #Number of phase data for TEC evaluation
    
    print 'total data number = ', ntdata, ', read_seconds = ', ntdata/sample_f
    print 'number of spectra = ', nparam
    print 'skip second = ', skip_sec, 'second of data presentation = ', show_sec
    
    # Array def
    
    x_sec = arange(nparam, dtype = 'f') * (float(nfft)/sample_f/2.0) + skip_sec
    freq = arange(nfft, dtype = 'f') * delta_f - sample_f/2.0
    
    buf = arange(ntdata, dtype = 'F')
    
    ser = arange(nfft, dtype = 'D')
    fft_result = arange(nfft, dtype = 'D')
    pow = arange(nfft, dtype = 'd')
    
    if150 = arange(nparam, dtype = 'l')
    if400 = arange(nparam, dtype = 'l')
    freq150 = arange(nparam, dtype = 'f')
    freq400 = arange(nparam, dtype = 'f')
    sigpow150 = arange(nparam, dtype = 'f')
    sigpow400 = arange(nparam, dtype = 'f')
    dif_f = arange(nparam, dtype = 'f')
    fTEC = arange(nparam, dtype = 'f')
    
    fil_ser150 = arange(n_phase, dtype = 'F')
    fil_ser400 = arange(n_phase, dtype = 'F')
    dif_p = arange(n_phase, dtype = 'f')
    pTEC = arange(n_phase, dtype = 'f')
    phase_sec = arange(n_phase, dtype = 'f') * (float(npfact)/sample_f) + skip_sec
    
    # record power spectrum
    nf_fact = 32
    nt_fact = 4
    nf_num = nfft/nf_fact
    nt_num = nparam/nt_fact
    f_spc = arange(nf_num, dtype='f')
    x_spc = arange(nt_num, dtype='f')
    spc_db400 = zeros(shape=(nt_num, nf_num), dtype='f')
    spc_db150 = zeros(shape=(nt_num, nf_num), dtype='f')
    spc_buf = zeros(shape=(nf_num), dtype='f')
    
    print 'nf_num=',nf_num,', nt_num=',nt_num
    
    for ii in range(0, nf_num):
        i1 = ii*nf_fact
        i2 = i1 + nf_fact
        f_spc[ii] = freq[i1:i2].mean()
    
    for ii in range(0, nt_num):
        i1 = ii*nt_fact
        i2 = i1 + nt_fact
        x_spc[ii] = x_sec[i1:i2].mean()

    # Multplied spectrum
    mp_fact = 8
    mp_nfft = nfft/mp_fact
    mp_tnum = nparam/2
    time_mp = arange(mp_tnum, dtype='f')
    spc_mp400 = arange(mp_nfft, dtype='f')
    spc_mp150 = arange(mp_nfft, dtype='f')
    freq_mp = arange(mp_nfft, dtype='f')
    spc_mp = zeros(shape=(mp_tnum, mp_nfft), dtype='f')
    tpeak_mp = arange(mp_tnum, dtype='f')
    fpeak_mp = arange(mp_tnum, dtype='f')
##    tpeak2 = arange(mp_tnum//10, dtype='f')
##    fpeak2 = arange(mp_tnum//10, dtype='f')
    
    # Open data files and skip
    f150 = open(file150, 'rb')
    f150.seek(nskip_bytes)
    
    f400 = open(file400, 'rb')
    f400.seek(nskip_bytes)

    # 400MHz data read (repositioning and read again)
    buf = io.fread(f400, ntdata, 'F')

    # Calculate mulplied spectrum and find initial values
    for ip in range(0, mp_tnum):

    # 400MHz spectrum
        noffset = ip * nfft
        ser = buf[noffset: noffset+nfft]
##        ser = io.fread(f400, nfft, 'F')
        fft_result = fft(ser)/nfft
        fft_result = fftpack.fftshift(fft_result)   #Rotate freq.
        pow = fft_result.real **2 + fft_result.imag **2     #Signal power
        if (ip % 1000) == 0:
            print 'mp_tnum = ',mp_tnum,', ip = ',ip
        for jf in range(0, mp_nfft):
            jf400 = jf * 8
            spc_mp400[jf] = pow[jf400: jf400+8].max()
            freq_mp[jf] = freq[jf400: jf400+8].mean()

    # 150MHz spectrum
        ser = io.fread(f150, nfft, 'F')
        fft_result = fft(ser)/nfft
        fft_result = fftpack.fftshift(fft_result)   #Rotate freq.
        pow = fft_result.real **2 + fft_result.imag **2     #Signal power
        for jf in range(0, mp_nfft):
            jf150 = jf * 3 + 2560
            spc_mp150[jf] = pow[jf150: jf150+3].max()
    # multpy spectra
        spc_mp[ip] = spc_mp400 * spc_mp150
        time_mp[ip] = x_sec[ip*2]

    spc_mp_min = spc_mp.min()
    spc_mp = 10.0*log10(spc_mp / spc_mp_min)

    #Find first guess of initial guess
    n_mp_peak = 0
    for ip in range(0, mp_tnum):
        temp_peak = spc_mp[ip].max()
        if temp_peak > 60.0:
            i_argmax = spc_mp[ip].argmax()
            tpeak_mp[n_mp_peak] = time_mp[ip]
            fpeak_mp[n_mp_peak] = freq_mp[i_argmax]
            n_mp_peak = n_mp_peak + 1

    print 'Number of peaks = ', n_mp_peak

    #Median filterring of initial guess
    n_mp2 = n_mp_peak//10
    tpeak2 = range(0)
    fpeak2 = range(0)
    for iip in range(0, n_mp2):
        ip = iip * 10
        tpeak2.append(tpeak_mp[ip:ip+10].mean())
        temp_f = fpeak_mp[ip:ip+10]
        temp_f.sort()
        fpeak2.append(temp_f[5])
##        print 'iip = ',iip,', t = ',tpeak2[iip],', f = ',fpeak2[iip]

    print 'Number of sample (0) = ', len(tpeak2)

    #Fitting to init_guess function
    #Init of parameters: P0 = (amp, slope, phase, offset)

    #1st attempt
    P0 = [8500.0, 0.008, x_sec[nparam/2], f_off400]
    plsq = optimize.leastsq(residual, P0, args=(fpeak2, tpeak2))
    
    #Remove outliers (1)
    iip = 0
    for j in range(len(tpeak2)):
        if abs(fpeak2[iip] - init_guess(plsq[0], tpeak2[iip])) > 5000.0:
            temp = tpeak2.pop(iip)
            temp = fpeak2.pop(iip)
        else:
            iip = iip + 1
        if (len(tpeak2) - iip) <= 0:
            break

    print 'Number of sample (1) = ', len(tpeak2)

    #2nd attempt
    plsq = optimize.leastsq(residual, P0, args=(fpeak2, tpeak2))

    #Remove outliers (2)
    iip = 0
    for j in range(len(tpeak2)):
        if abs(fpeak2[iip] - init_guess(plsq[0], tpeak2[iip])) > 1000.0:
            temp = tpeak2.pop(iip)
            temp = fpeak2.pop(iip)
        else:
            iip = iip + 1
        if (len(tpeak2) - iip) <= 0:
            break

    print 'Number of sample (2) = ', len(tpeak2)
    
    #3rd attempt
    plsq = optimize.leastsq(residual, P0, args=(fpeak2, tpeak2))
    
    print 'P0 = ', P0
    print 'Result = ', plsq[0]

    fig0 = figure()

    plot(tpeak_mp[0: n_mp_peak], fpeak_mp[0: n_mp_peak])
    plot(tpeak2[0:n_mp2], fpeak2[0:n_mp2])

    tpeak3 = range(0)
    fpeak3 = range(0)
    for iip in range(mp_tnum//10):
        tpeak3.append(time_mp[iip*10])
        fpeak3.append(init_guess(plsq[0], tpeak3[iip]))

    plot(tpeak3, fpeak3)
    xlabel('time after ' + dt_string + ' (s)')
    ylabel('initial guess result (Hz)')
    title(sat_name + '   ' + dt_string +' UT')

    fig_filename0 = base_dir + sat_name + '_' + datename + '_' + timename + '_initguess.png'
    savefig(fig_filename0)

##    fig1 = figure()
##    V = arange(9, dtype='f')*5.0 - 40.0
##    contourf(freq_mp, time_mp, spc_mp)
##    colorbar()
##    show()
    
    # 400MHz data read (repositioning and read again)
##    f400.seek(nskip_bytes, 0)
##    buf = io.fread(f400, ntdata, 'F')
    
    fil_ser400[0:npst] = 1.0 + 0.0j
    fil_ser400[-npst:] = 1.0 + 0.0j
    ifst = nfft/4
    ifen = ifst * 3

    # 400MHz data processing

    sat_found = 0

    for ip in range(0, nparam):

        if (ip % 1000) == 0:
            print '400MHz: nparam = ',nparam,', ip = ',ip

        noffset = ip * nfft / 2
        ser = buf[noffset: noffset+nfft]
    
        fft_result = fft(ser)/nfft
    
        fft_result = fftpack.fftshift(fft_result)   #Rotate freq.
        pow = fft_result.real **2 + fft_result.imag **2     #Signal power

        if sat_found == 0:
            if_init = int((init_guess(plsq[0], x_sec[ip]) + sample_f/2.0)/delta_f) - 75
            if400[ip] = pow[if_init: if_init+150].argmax() + if_init
        else:
            if_init = if400[ip-1] - 25
            if400[ip] = pow[if_init: if_init + 50].argmax() + if_init
    
        sigpow400[ip] = 0.0
        freq400[ip] = 0.0
        for j in range(if400[ip]-2, if400[ip]+3):
            sigpow400[ip] = sigpow400[ip] + pow[j]
            freq400[ip] = freq400[ip] + freq[j] * pow[j]
    
        freq400[ip] = freq400[ip] / sigpow400[ip]
    
        if (10.*log10(sigpow400[ip])) >= -10.:
            if abs(freq400[ip] - init_guess(plsq[0], x_sec[ip])) <= 200.0:
                sat_found = 1
            else:
                sat_found = 0
        else:
            sat_found = 0
        
        ist = ip*npst*2 + npst
        ien = ist + npst*2
    
        if sigpow400[ip] > 0.0:
            fft_result[:if400[ip]-3] = 0.0 + 0.0j
            fft_result[if400[ip]+4:] = 0.0 + 0.0j
         
            fft_result = fftpack.ifftshift(fft_result)  #Rotate freq. back
            
            ser = ifft(fft_result)
            
            fil_ser400[ist:ien] = ser[ifst:ifen:npfact]
        else:
            fil_ser400[ist:ien] = 1.0 + 0.0j
    
    # record power spectrum
        for ii in range(0, nf_num):
            i1 = ii*nf_fact
            i2 = i1 + nf_fact
            spc_buf[ii] = spc_buf[ii] + pow[i1:i2].mean()
    
        if (ip % nt_fact) == (nt_fact - 1):
            spc_db400[ip/nt_fact] = spc_buf/float(nt_fact)
            spc_buf.fill(0.0)
    
    spc_db400 = 10.0 * log10(spc_db400)
    
    # 150MHz data read
    f150.seek(nskip_bytes, 0)
    buf = io.fread(f150, ntdata, 'F')
    
    fil_ser150[0:npst] = 1.0 + 0.0j
    fil_ser150[-npst:] = 1.0 + 0.0j
    ifst = nfft/4
    ifen = ifst * 3
    
    spc_buf.fill(0.0)
    
    for ip in range(0, nparam):

        if (ip % 1000) == 0:
            print '150MHz: nparam = ',nparam,', ip = ',ip

        noffset = ip * nfft / 2
        ser = buf[noffset: noffset+nfft]
    
        fft_result = fft(ser)/nfft
    
        fft_result = fftpack.fftshift(fft_result)   #Rotate freq.
        pow = fft_result.real **2 + fft_result.imag **2     #Signal power
    
    #plot power spectrum
   
        if_init = int(((freq400[ip]*3.0/8.0) + sample_f/2.0)/delta_f) - 5
        if150[ip] = pow[if_init: if_init+11].argmax() + if_init
    
        sigpow150[ip] = 0.0
        freq150[ip] = 0.0
        for j in range(if150[ip]-2, if150[ip]+3):
            sigpow150[ip] = sigpow150[ip] + pow[j]
            freq150[ip] = freq150[ip] + freq[j] * pow[j]
    
        freq150[ip] = freq150[ip] / sigpow150[ip]
    
        ist = ip*npst*2 + npst
        ien = ist + npst*2
    
        if sigpow150[ip] > 0.0:
            fft_result[:if150[ip]-1] = 0.0 + 0.0j
            fft_result[if150[ip]+2:] = 0.0 + 0.0j
            
            fft_result = fftpack.ifftshift(fft_result)  #Rotate freq. back
    
            ser = ifft(fft_result)
            
            fil_ser150[ist:ien] = ser[ifst:ifen:npfact]
        else:
            fil_ser150[ist:ien] = 1.0 + 0.0j
    
    # record power spectrum
        for ii in range(0, nf_num):
            i1 = ii*nf_fact
            i2 = i1 + nf_fact
            spc_buf[ii] = spc_buf[ii] + pow[i1:i2].mean()
    
        if (ip % nt_fact) == (nt_fact - 1):
            spc_db150[ip/nt_fact] = spc_buf/float(nt_fact)
            spc_buf.fill(0.0)
    
    spc_db150 = 10.0 * log10(spc_db150)
    
    # Calculate differential phase
    dif_p = angle((fil_ser150 ** 8) / (fil_ser400 ** 3))
    
    # Unwrap phse
    pTEC = unwrap(dif_p)
    
    # Tune error removal
    for ip in range(0, len(phase_sec)):
        pTEC[ip] = pTEC[ip] - tune_error * (phase_sec[ip] - phase_sec[0])
    
    # calculate TEC value
    coef = pi * A_coef / fr / c_speed * 55.0 / 24.0
    pTEC = pTEC / coef
    
    # 400MHz spectrum contour
    fig1 = figure()
    
    V = arange(9, dtype='f')*5.0 - 40.0
    
    contourf(f_spc, x_spc, spc_db400, V)
    xlabel('frequency (Hz)')
    ylabel('time after ' + dt_string + ' (s)')
    title(sat_name + ' 400 MHz signal PSD (' + sitename + ')')
    colorbar()
    fig_filename1 = base_dir + sat_name + '_' + datename + '_' + timename + '_400.png'
    savefig(fig_filename1)
    
    # 150MHz spectrum contour
    fig2 = figure()
    
    V = arange(9, dtype='f')*5.0 - 40.0
    
    contourf(f_spc, x_spc, spc_db150, V)
    xlabel('frequency (Hz)')
    ylabel('time after ' + dt_string + ' (s)')
    title(sat_name + ' 150 MHz signal PSD (' + sitename + ')')
    colorbar()
    fig_filename2 = base_dir + sat_name + '_' + datename + '_' + timename + '_150.png'
    savefig(fig_filename2)
    
    # File close
    f150.close
    f400.close
    del buf
    
    ### file output ###
    
    Out_filename = base_dir + sat_name + '_' + datename + '_' + timename + '_summary.txt'
    
    f_out1 = open(Out_filename, 'w')
    
    f_out1.write('time(s) 400MHz(pow freq) 150MHz(pow freq) TEC(TECu) STD_TEC (Total %d lines)\n'% nparam)
    
    for ip in range(0, nparam):
        ist = ip*npst*2 + npst
        ien = ist + npst*2
        ave_TEC = pTEC[ist:ien].mean()/1.0e16    # averaged TEC in TECu
        std_TEC = pTEC[ist:ien].std()/1.0e16    # its standard deviation in TECu
        write_buf = '%g, %g, %g, %g, %g, %g, %g\n' % (x_sec[ip], sigpow400[ip], freq400[ip], sigpow150[ip], freq150[ip], ave_TEC, std_TEC)
        f_out1.write(write_buf)
    
    f_out1.close
    
    Out_filename = base_dir + sat_name + '_' + datename + '_' + timename + '_TEConly.txt'
    
    f_out2 = open(Out_filename, 'w')
    
    f_out2.write('time TEC (Total %d lines)\n' % n_phase)
    
    for ip in range(0, n_phase):
        write_buf = '%g, %g\n' % (phase_sec[ip], pTEC[ip]/1.0e16)
        f_out2.write(write_buf)
    
    f_out2.close
    
def _main():
    
    if len(sys.argv) > 1:
        if len(sys.argv) >= 4:
            sat_name = sys.argv[1]
            datename = sys.argv[2]
            timename = sys.argv[3]
        else:
            print 'TEC_analysis: not enough parameters! exit!'
            sys.exit()
    else:
##    sat_name = raw_input('Please enter sat_name: ')
##    datename = raw_input('Please enter date (YYYYMMDD): ')
##    timename = raw_input('Please enter time (hhmmss): ')

##    sat_name = 'DMSPF15'
##    datename = '20080728'
##    timename = '092835'

##    sat_name = 'FM6'
##    datename = '20080728'
##    timename = '024629'

        sat_name = 'OSCAR32'
        datename = '20080728'
        timename = '052303'

##    sat_name = 'COSMOS2407'
##    datename = '20080727'
##    timename = '213235'

##    sat_name = 'RADCAL'
##    datename = '20080726'
##    timename = '092612'

# Execute analysis
    print 'TEC_analysis: sat_name=',sat_name,', datename=',datename,', timename=',timename
    TECanalysis(sat_name, datename, timename)

if __name__ == '__main__':
    _main()