import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
from matplotlib import dates
import datetime
import numpy as np
import math
import os
import re
import time
import glob
import calendar

ZERO = datetime.timedelta(0)

class UTC(datetime.tzinfo):
    """UTC"""

    def utcoffset(self, dt):
        return ZERO

    def tzname(self, dt):
        return "UTC"

    def dst(self, dt):
        return ZERO

utc = UTC()

def dfunc(arg):
    return datetime.datetime.fromtimestamp(arg, tz=utc)

def nameToTime(name):
    m = re.match( r'.*/spectra/rfispectrum_(.*)\+0000.txt', name, re.M|re.I)
    tm = time.strptime(m.group(1), "%Y-%m-%dT%H:%M:%S")
    return (calendar.timegm(tm))
    

def readData(rDict):
    files = glob.glob("/data/MENTOK_2/ste616/EB500/Software/Utilities/TCPExample/spectra/rfispectrum_*.txt")
    rfiles = files[::-1]
    ftime = nameToTime(rfiles[0])
    for f in rfiles:
        tt = nameToTime(f)
        if ((ftime - tt) < 3600): # Keep within an hour.
            rDict['rawdata'].append(np.loadtxt(f))
            rDict['times'].append(tt)
        else:
            break
    ts = np.array(rDict['rawdata'])
    rDict['rawdata'] = ts + abs(ts.min()) + 1

def reformatData(rDict):
    rDict['freqs'] = []
    for i in xrange(802, 3000, 2):
        rDict['freqs'].append(float(i) / 1000.0)

    rDict['vals'] = []
    for az in xrange(0, len(rDict['freqs'])):
        for za in xrange(0, len(rDict['times'])):
            rDict['vals'].append(math.log10(rDict['rawdata'][za][az][1]))

def plot_spectrum(values, freqs, times):
    values = np.array(values)
    values = values.reshape(len(freqs), len(times))

    mspect = []
    for f in xrange(0, len(freqs)):
        mspect.append(values[f, :].max())

    fig, ax = plt.subplots()
    ax.plot(freqs, values[:,1], 'b', label='Last (%s)' % datetime.datetime.fromtimestamp(times[1], tz=utc).strftime("%d/%m %H:%M:%S"))
    ax.plot(freqs, values[:,0], 'r', label='Latest (%s)' % datetime.datetime.fromtimestamp(times[0], tz=utc).strftime("%d/%m %H:%M:%S"))
    ax.plot(freqs, mspect, 'g', label='Last hour max')
    ax.set_xlabel("Frequency [GHz]")
    ax.set_ylabel("Log10 Amplitude");
    ax.set_autoscaley_on(False)
    ax.set_ylim([1.0, 3.0])
    ax.set_autoscalex_on(False)
    ax.set_xlim([0.8, 3.0])
    plt.legend(prop={'size':10})
    if __name__ == "__main__":
        fout = 'spectrumtest.png'
        plt.savefig(fout)

def plot_contour(values, freqs, times):
    """Plot a polar contour plot, with 0 degrees at the North.
 
    Arguments:
 
     * `values` -- A list (or other iterable - eg. a NumPy array) of the values to plot on the
     contour plot (the `z` values)
     * `azimuths` -- A list of azimuths (in degrees)
     * `zeniths` -- A list of zeniths (that is, radii)
 
    The shapes of these lists are important, and are designed for a particular
    use case (but should be more generally useful). The values list should be `len(azimuths) * len(zeniths)`
    long with data for the first azimuth for all the zeniths, then the second azimuth for all the zeniths etc.
 
    This is designed to work nicely with data that is produced using a loop as follows:
 
    values = []
    for azimuth in azimuths:
      for zenith in zeniths:
        # Do something and get a result
        values.append(result)
 
    After that code the azimuths, zeniths and values lists will be ready to be passed into this function.
 
    """
    values = np.array(values)
    values = values.reshape(len(freqs), len(times))

    dts = map(dfunc, times)
    fds = dates.date2num(dts)
    hfmt = dates.DateFormatter("%d/%m\n%H:%M", tz=utc)
 
    r, theta = np.meshgrid(fds, freqs)
    fig, ax = plt.subplots()
    v = np.arange(1, 3, 0.01)

    cax = ax.contourf(theta, r, values, v, cmap=plt.cm.jet) # Reasonably detailed

    ax.yaxis.set_major_locator(dates.MinuteLocator(byminute=range(0, 60, 10), tz=utc))
    ax.yaxis.set_major_formatter(hfmt)

    cb = fig.colorbar(cax)
    cb.set_label("Log10 Amplitude")
    ax.set_xlabel("Frequency [GHz]")
    ax.set_ylabel("UTC")
    
    if __name__ == "__main__":
        #fout = '/n/ste616/Documents/rfi_skyplots/rfiwaterfall_latest.png'
        fout = 'rfiwaterfall_test.png'
        #print fout
        plt.savefig(fout)
    else:
        plt.ion()
        plt.show()
 
    return fig, ax, cax

def main():
    r = {}
    r['rawdata'] = []
    r['times'] = []
    readData(r)
    reformatData(r)
    plot_contour(r['vals'], r['freqs'], r['times'])
    plot_spectrum(r['vals'], r['freqs'], r['times'])
    return r

if __name__ == "__main__":
    t = main()


