import matplotlib,numpy
matplotlib.use('Agg')
from astropy.units        import Quantity
from matplotlib           import pyplot
from gwpy.plotter         import SegmentPlot,TimeSeriesPlot,SpectrumPlot,SpectrogramPlot
from gwpy.segments        import SegmentList
from gwpy.spectrum        import Spectrum
from gwpy.spectrogram     import Spectrogram
from gwpy.table.lsctables import SnglBurstTable
from gwpy.timeseries      import TimeSeries

def plot_activity(full_seglist):
    """
    Plot full activity period for station.

    Parameters
    ----------
    full_seglist : dictionary
      Continuous list of available data in the selected time period
    """
    # Import gwpy tools
    plot = SegmentPlot()
    # Initialize plotting figure
    ax = plot.gca()
    # Plot all segment in figure
    ax.plot(full_seglist)
    # Save figure
    pyplot.savefig("activity.png",dpi=500)

def plot_time_series(station,ts_list,seglist=None,hp=None):
    """
    Generate a plot of the whole data time series
    """
    plot = TimeSeriesPlot()
    ax = plot.gca()
    # Loop over all the time series 
    for ts in ts_list:
        # Plot time series for each segment
        ax.plot(ts, color='black')
    # Display title
    ax.set_title('$\mathrm{'+station+'}$')
    # Plot fake signal
    if hp!=None:
        ax.plot(hp, color='red')
    # Plot activity segments
    if seglist!=None:
        plot.add_state_segments(SegmentList(seglist[station].active),plotargs={'label':'data present','facecolor': 'g','edgecolor': 'k'})
    # Fix exceeded cell block limit error
    matplotlib.pyplot.rcParams['agg.path.chunksize'] = 20000
    # Save figure
    plot.savefig('time_series.png',dpi=500)

def plot_asd(station,ts_list):
    """
    Plot Amplitude Spectral Density. AGG complexity starts to complain
    with large numbers of points. And we somehow invoke precision issues
    that need to be ameliorated.
    """
    if station!='fake':
        for d in ts_list:
            d.x0 = Quantity(int(d.x0.value * 500), d.xunit)
            d.dx = Quantity(1, d.xunit)
        ts_list.coalesce()
        for d in ts_list:
            d.x0 = Quantity(d.x0.value / 500, d.xunit)
            d.dx = Quantity(0.002, d.xunit)
    # Initialize plotting functionality
    plot = SpectrumPlot()
    # Loop over all the time series
    for d in ts_list:
        # Generate 8 seconds per FFT with 4 second (50%) overlap
        spectrum = d.asd(8, 4)
        # Create plotting axis
        ax = plot.gca()
        # Plot square root of the spectrum
        ax.plot(numpy.sqrt(spectrum))
    # Set x axis to log scale
    ax.set_xscale('log')
    # Set y axis to log scale
    ax.set_yscale('log')
    # Set x axis limits
    ax.set_xlim(1e-1, 500)
    # Save figure
    plot.savefig("asd.png",dpi=500)
    
def plot_whitening(station,ts_list,seglist=None):
    """
    Generate a spectrogram plot and normalized spectrogram
    norm: \sqrt{S(f,t)} / \sqrt{\overbar{S(f)}}
    """
    stride,fftlength,overlap = 20,6,3
    plot = SpectrogramPlot()
    ax = plot.gca()
    white_plot = SpectrogramPlot()
    wax = white_plot.gca()
    for ts in ts_list:
        if (len(ts) * ts.dt).value < stride:
            continue
        spec = ts.spectrogram(stride, fftlength=fftlength, overlap=overlap)
        ax.plot(spec)
        wspec = spec.ratio('median')
        wax.plot(wspec, vmin=0.1, vmax=100)
    ax.set_title('$\mathrm{'+station+'}$')
    ax.set_ylim(0.1, ts.sample_rate.value/2.)
    ax.set_yscale('log')
    wax.set_title('$\mathrm{'+station+'}$')
    wax.set_ylim(0.1, ts.sample_rate.value/2.)
    wax.set_yscale('log')
    plot.add_colorbar(label='Amplitude')
    white_plot.add_colorbar(label='Amplitude')
    if seglist!=None:
        plot.add_state_segments(SegmentList(seglist[station].active),plotargs={'label':'data present','facecolor':'g','edgecolor':'k'})
        white_plot.add_state_segments(SegmentList(seglist[station].active),plotargs={'label':'data present','facecolor':'g','edgecolor':'k'})
    plot.savefig("spectrogram.png",dpi=500)
    white_plot.savefig("whitened.png",dpi=500)

def plot_bank(fdb):
    pyplot.figure()
    for i, fdt in enumerate(fdb[:5]):
        pyplot.plot(fdt.frequencies, fdt, 'k-')
    pyplot.grid()
    pyplot.xlabel("frequency [Hz]")
    pyplot.savefig('bank.png')
    pyplot.close()

def plot_filters(tdb,flow,band):
    pyplot.figure()
    pyplot.subplots_adjust(left=0.2,right=0.95,bottom=0.15,top=0.95,hspace=0,wspace=0)
    for i, tdt in enumerate(tdb[:8:3]):
        ax = pyplot.subplot(3, 1, i+1)
        ax.plot(tdt.times.value - 2., numpy.real_if_close(tdt.value), 'k-')
        c_f = flow+band/2 + 3 * (band*i) + 2.
        ax.set_xlabel("Time (s)")
        ax.set_ylabel("%d Hz" % c_f)
        ax.set_xlim(25.0, 31.0)
        ax.set_ylim([-max(tdt.value), max(tdt.value)])
        if i!=2: pyplot.setp(ax.get_xticklabels(), visible=False)
    pyplot.savefig('filters.png')
    pyplot.close()
    
def plot_ts(ts, fname="ts.png"):
    plot = TimeSeriesPlot()
    ax = plot.gca()
    ax.plot(TimeSeries(ts, sample_rate=1.0/ts.delta_t, epoch=ts.start_time))
    ax.set_xlim(ts.start_time,ts.end_time)
    pyplot.savefig(fname)
    pyplot.close()

def plot_spectrum(fd_psd):
    plot = SpectrumPlot()
    ax = plot.gca()
    ax.plot(Spectrum(fd_psd, df=fd_psd.delta_f))
    #pyplot.ylim(1e-10, 1e-3)
    pyplot.xlim(0.1, 500)
    pyplot.loglog()
    pyplot.savefig("psd.png")
    pyplot.close()

def plot_spectrogram(spec,dt,df,sample_rate,start_time,end_time,fname="specgram.png"):
    plot = SpectrogramPlot()
    ax = plot.gca()
    ax.plot(Spectrogram(spec,dt=dt,df=df,epoch=start_time), cmap='viridis')
    plot.add_colorbar(label='Amplitude')
    pyplot.xlim(start_time,end_time)
    pyplot.ylim(0,sample_rate/2.)
    pyplot.savefig(fname)
    pyplot.close()

def plot_spectrogram_from_ts(ts):
    plot = SpectrogramPlot()
    ax = plot.gca()
    ax.plot(Spectrogram(spec))
    #pyplot.ylim(1e-9, 1e-2)
    #pyplot.xlim(0.1, 500)
    #pyplot.loglog()
    pyplot.savefig("specgram.png")
    pyplot.close()

def plot_triggers():
    events = SnglBurstTable.read('excesspower.xml.gz')
    #plot = events.plot('time', 'central_freq', "duration", "bandwidth", color='snr')
    plot = events.plot('time','central_freq',color='snr',edgecolor='none')
    #plot.set_xlim(time_start,time_end)
    #plot.set_ylim(band, sample_rate/2.)
    plot.set_ylabel('Frequency [Hz]')
    plot.set_yscale('log')
    #plot.set_title('GNOME '+station+' station event triggers')
    plot.add_colorbar(cmap='copper_r',label='Tile Energy')
    pyplot.savefig("triggers.png",dpi=400)

def plot_tiles():    
    bins = numpy.linspace(0, 40, 100)
    cnt = numpy.zeros(bins.shape[0]-1)
    for i, tdf in enumerate(tdb[:nchans]):
        us_rate = int(1.0 / (2 * band*nc_sum * ts_data.dt.value))
        pyplot.figure(0, figsize=(10, 10))
        pyplot.subplot(nchans, 1, i+1)
        white = tmp_ts_data.whiten(64, 32, asd=numpy.sqrt(cdata_psd_tmp), window='boxcar') * sample_rate/4
        snr_1dof = numpy.convolve(tdf, white, "valid")
        # Undersample the data
        snr_1dof = snr_1dof[::us_rate]**2
        # Sum semi-adjacent samples to get 2 DOF tiles
        snr_2dof = numpy.convolve(snr_1dof, numpy.array([1, 0, 1, 0]))
        t = TimeSeries(snr_2dof, epoch=white.epoch, sample_rate=int(1.0/(us_rate * tmp_ts_data.dt.value)))
        pyplot.plot(t.times + len(tdf)/2 * tdf.dt, snr_2dof, 'k-')
        pyplot.axvline(random_time)
        tmp, _ = numpy.histogram(snr_2dof, bins=bins)
        cnt += tmp
    plot_spectrogram(dof_tiles.T,fname='%s/tf_%ichans_%02idof.png'%(segfolder,nc_sum+1,2*j))
    plot.savefig("%s/bands.png"%(segfolder))
