import matplotlib,numpy
matplotlib.use('Agg')
from .retrieve             import time_convert
from astropy.units         import Quantity
from matplotlib            import pyplot
from gwpy.table            import EventTable
from gwpy.plotter          import SegmentPlot,TimeSeriesPlot
from gwpy.plotter          import FrequencySeriesPlot,SpectrogramPlot
from gwpy.segments         import SegmentList
from gwpy.frequencyseries  import FrequencySeries
from gwpy.spectrogram      import Spectrogram
from gwpy.timeseries       import TimeSeries
from pylab                 import *
from scipy                 import signal

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=300)

def plot_asd(station,data):
    """
    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 data:
            d.x0 = Quantity(int(d.x0.value * 500), d.xunit)
            d.dx = Quantity(1, d.xunit)
        data.coalesce()
        for d in data:
            d.x0 = Quantity(d.x0.value / 500, d.xunit)
            d.dx = Quantity(0.002, d.xunit)
    # Initialize plotting functionality
    plot = FrequencySeriesPlot()
    # Loop over all the time series
    for d in data:
        # 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')
    ax.set_xlabel('Frequency [Hz]')
    # Set y axis to log scale
    ax.set_yscale('log')
    ax.set_ylabel('Amplitude [pT]')
    # Set x axis limits
    ax.set_xlim(1e-1, 500)
    import matplotlib.ticker as ticker
    x = ax.get_xticklabels()
    def myticks(x,pos):
        if x == 0: return "$0$"
        exponent = int(np.log10(x))
        coeff = x/10**exponent
        if coeff==1:
            return r"$10^{{ {:2d} }}$".format(exponent)
        else:
            return r"${:2.0f} \times 10^{{ {:2d} }}$".format(coeff,
                                                             exponent)
    ax.xaxis.set_major_formatter(ticker.FuncFormatter(myticks))
    # Save figure
    plot.savefig("asd.png",dpi=300)
    
def plot_time_series(data,station='station-name',t0=None,t1=None,seglist=None,burst=None,fname='time_series'):
    """
    Generate a plot of the whole data time series
    """
    if type(data[0])==float64:
        data = [TimeSeries(data,sample_rate=data.sample_rate,epoch=data.start_time)]
    plot = TimeSeriesPlot()
    ax = plot.gca()
    # Loop over all the time series
    for ts in data:
        # Plot time series for each segment
        ax.plot(ts, color='black')
    # Display title
    ax.set_title('$\mathrm{'+station+'}$')
    ax.set_ylabel('Magnetic Field')
    # Plot fake signal
    if burst is not None:
        ax.plot(burst, color='red')
    # Plot activity segments
    if seglist!=None:
        activity = SegmentList(seglist[station].active)
        plotargs = {'label':'data present','facecolor':'g','edgecolor':'k'}
        plot.add_state_segments(activity,plotargs=plotargs)
    # Set limits
    if t0 is not None and t1 is not None:
        t0,t1 = time_convert(t0,t1)
        plot.axes[0].set_epoch(t0)
        plot.axes[1].set_epoch(t0)
        ax.set_xlim(t0,t1)
    # Fix exceeded cell block limit error
    matplotlib.pyplot.rcParams['agg.path.chunksize'] = 20000
    # Save figure
    plot.savefig('%s.png'%fname)
    
def plot_whitening(data,station='station-name',t0=None,t1=None,stride=20,fftlength=6,overlap=3,seglist=None):
    """
    Generate a spectrogram plot and normalized spectrogram
    norm: \sqrt{S(f,t)} / \sqrt{\overbar{S(f)}}
    """
    if type(data[0])==float64:
        data = [TimeSeries(data,sample_rate=data.sample_rate,epoch=data.start_time)]
    # Setup plots
    plot = SpectrogramPlot()
    ax = plot.gca()
    white_plot = SpectrogramPlot()
    wax = white_plot.gca()
    # Loop through available time series
    for ts in data:
        if (len(ts) * ts.dt).value < stride:
            continue
        spec = ts.spectrogram(stride,fftlength=fftlength,overlap=overlap)
        wspec = spec.ratio('median')
        ax.plot(spec, cmap='jet')
        wax.plot(wspec, vmin=0.1, vmax=100,cmap='jet')
    # Define y axis and title
    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')