#!/usr/bin/env python
if True:
    import matplotlib
    matplotlib.use('Agg')

    import numpy as np
    import matplotlib.pyplot as plt
    import argparse
    from matplotlib.gridspec import GridSpec

    parser = argparse.ArgumentParser()
    parser.add_argument('freq')
    parser.add_argument('samp_rate')
    parser.add_argument('nchan')
    parser.add_argument('nbin')
    args = parser.parse_args()

    def decibel(x):
        #return 10.0*np.log10(x)
        return x

    def shift_col_up(chan_num,n_rows):
        z[:,chan_num] = np.roll(z[:,chan_num], -n_rows)
    print(1)
    if __name__ == "__main__":
        print(2)
        #Observation parameters
        exec(args.freq)
        exec(args.samp_rate)
        exec(args.nchan)
        exec(args.nbin)
        fname = "/home/pictor/Desktop/pictortelescope/sim_pulsar.dat"

        #Load data
        z = np.fromfile(fname, dtype="float32").reshape(-1, nchan)/nbin
        #z = z*10000
        #z[z > 35000] = 30000

        #Dedisperse
        nchan = z.shape[1]
        DM = 100 #2.96927 #pc/cm^3
        bandwidth_mhz = 16
        bandwidth = 16000000
        f_center=freq
        f_center_mhz=0.000001*f_center
        nbins=nbin
        deltaF = float(bandwidth_mhz)/nchan #bandwidth in MHz
        print(deltaF)
        f_max = f_center_mhz+(float(bandwidth_mhz)/2)
        f_max_mhz = f_max #*0.000001
        t_int = float(nbins*nchan)/bandwidth
        for chan_num in range(nchan):
            f_start = f_center_mhz-(float(bandwidth_mhz)/2)
            f_chan = f_start+(chan_num)*deltaF
            print(f_chan)
            deltaT = 4149*DM*((1/(f_chan**2))-(1/(f_max_mhz**2)))
            print('deltaT:',deltaT)
            print('f_chan:',f_chan)
            print('f_max_mhz:',f_max_mhz)
            n = int(round(float(100*deltaT)/t_int))
            print('n:',n)
            print('t_int:',t_int)
            shift_col_up(chan_num,n)

        #Define numpy array for Power vs Time plot
        w = np.mean(a=z, axis=1)

        #Number of sub-integrations
        nsub = z.shape[0]

        #Compute average spectrum
        zmean = np.mean(z, axis=0)

        #Compute time axis
        tint = float(nbin*nchan)/samp_rate
        t = tint*np.arange(nsub)

        v = np.arange(0, np.max(t)+tint, tint)

        #Compute frequency axis (convert to MHz)
        freq = np.linspace(freq-0.5*samp_rate, freq+0.5*samp_rate, nchan, endpoint=False)*1e-6

        #Initialize plot
        fig = plt.figure(figsize=(25,20))
        gs = GridSpec(2,2)

        #Plot average spectrum
        ax1 = fig.add_subplot(gs[0,0])
        ax1.plot(freq, decibel(zmean))
        ax1.set_xlim(np.min(freq), np.max(freq))
        ax1.ticklabel_format(useOffset=False)
        ax1.set_xlabel("Frequency (MHz)")
        ax1.set_ylabel("Relative Power")
        ax1.set_title("Averaged Spectrum")

        #Plot dynamic spectrum
        ax2 = fig.add_subplot(gs[0,1])
        ax2.imshow(decibel(z), origin="lower", interpolation="None", aspect="auto",
                   extent=[np.min(freq), np.max(freq), np.min(t), np.max(t)])
        ax2.ticklabel_format(useOffset=False)
        ax2.set_xlabel("Frequency (MHz)")
        ax2.set_ylabel("Time (s)")
        ax2.set_title("Dynamic Spectrum (Waterfall)")

        #Plot Power vs Time
        ax3 = fig.add_subplot(gs[1,:])
        ax3.plot(v,w)
        ax3.set_xlim(0,np.max(t)+tint)
        ax3.set_xlabel("Time (s)")
        ax3.set_ylabel("Relative Power")
        ax3.set_title("Power vs Time")

        plt.tight_layout()
        plt.savefig("/home/pictor/Desktop/pictortelescope/plot.png")
else:
    print(0)