#!/usr/bin/python

import os,sys,numpy
from optparse import OptionParser

usage =  "usage: %prog -i <ionexfile>"
parser = OptionParser(usage)
parser.add_option("-i", "--ionexfile", dest="ionexfile", default="",
                  help="ionexfile to plot")
parser.add_option("-d", "--difffile", dest="difffile", default="",
                  help="file to diff ionexfile with")
parser.add_option("-m", "--makemovie", dest="domovie", default=False,
                  action="store_true", help="Make an avi")
parser.add_option("--latitude", dest="latitude", default="-90,90",
                  help="comma separated latitude range (deg)")
parser.add_option("--longitude", dest="longitude", default="-180,180",
                  help="comma separated longitude range (deg)")
(options, junk) = parser.parse_args()
ionexfile = options.ionexfile
difffile = options.difffile
latitudes = options.latitude.split(',')
longitudes = options.longitude.split(',')
domovie = options.domovie

if ionexfile == "":
    print "You must supply an ionex file - aborting!!!"
    sys.exit(1)

ionexin = open(ionexfile)
ionexlines = ionexin.readlines()
ionexin.close()
if not difffile == "":
    diffin = open(difffile)
    difflines = diffin.readlines()
    diffin.close()
at = 0
at2 = 0

while at < len(ionexlines) and (len(ionexlines[at]) < 21 or
                                not ionexlines[at][-21:].rstrip() == \
                                'START OF TEC MAP'):
    at = at + 1
if not difffile == "":
    while at2 < len(difflines) and (len(difflines[at2]) < 21 or
                                    not difflines[at2][-21:].rstrip() == \
                                    'START OF TEC MAP'):
        at2 = at2 + 1

if at == len(ionexlines) or (not difffile == "" and at2 == len(difflines)):
    print "Couldn't find start of map - aborting!"
    sys.exit(1)

ionexlines = ionexlines[at:]
if not difffile == "":
    difflines = difflines[at2:]
nummaps = len(ionexlines)/(429*2)
print "Number of maps is " + str(nummaps)

numlats = 71
numlongs = 73
tecvals = numpy.zeros((nummaps,numlats,numlongs), numpy.float64)
times = []
at = 0
for i in range(nummaps):
    at = at+1
    timesplit = ionexlines[at].split()
    times.append(timesplit[:6])
    at = at+2
    for j in range(numlats):
        line = ""
        dline = ""
        for addline in ionexlines[at:at+5]:
            line = line + addline
        if not difffile == "":
            for addline in difflines[at:at+5]:
                dline = dline + addline
            dtecsplit = dline.split()
        tecsplit = line.split()
        k = 0
        for tec in tecsplit:
            tecvals[i][j][k] = float(tec)
            if not difffile == "":
                tecvals[i][j][k] = tecvals[i][j][k] - float(dtecsplit[k])
            k = k+1
        at = at+6

maxtecval = 0
for i in range(nummaps):
    tempout = open("movie.%03i.frame" % (i), 'w')
    for j in range(numlats):
        for k in range(numlongs):
            if tecvals[i][j][k] > maxtecval:
                maxtecval = tecvals[i][j][k]
            tempout.write(str(87.5-float(j)*2.5) + " " + \
                          str(-180+k*5) + " " + \
                          str(tecvals[i][j][k]) + "\n")
        tempout.write("\n")
    tempout.close()

if difffile == "":
    tempfile = "plot." + ionexfile.split('/')[-1] + ".gnuplot"
else:
    tempfile = "plot." + ionexfile.split('/')[-1] + \
               "-" + difffile.split('/')[-1] + ".gnuplot"
tempout = open(tempfile, 'w')
tempout.write("set pm3d\n")
tempout.write("set term x11\n")
for i in range(nummaps):
    tempout.write("set border 895\n")
    tempout.write("set ztics mirror\n")
    tempout.write("set title \"" + times[i][0] + "/" + times[i][1] + "/" + \
                  times[i][2] + " " + times[i][3] + "h" + times[i][4] +"m\"\n")
    tempout.write("set xlabel \"Latitude (deg)\"\n")
    tempout.write("set xrange [" + latitudes[0] + ":" + latitudes[1] + "]\n")
    tempout.write("set ylabel \"Longitude (deg)\"\n")
    tempout.write("set yrange [" + longitudes[0] + ":" + longitudes[1] + "]\n")
    tempout.write("set zlabel \"TEC\"\n")
    tempout.write("set zrange [0:" + str(maxtecval) + "]\n")
    if domovie:
        tempout.write("set term png gray\n")
        tempout.write("set output \"movie.%03i.frame.png\"\n" % (i))
    tempout.write("splot \"movie.%03i.frame\" u 1:2:3 w lines\n" % (i))
    if not domovie:
        tempout.write("pause 1\n")
tempout.close()
os.system("gnuplot < " + tempfile)
if domovie:
    os.system("mencoder \"mf://movie*.frame.png\" -mf w=800:h=600:fps=1:type=png -ovc copy -oac copy -o " + ionexfile.split('/')[0] + "." + ionexfile.split('/')[-1] + ".avi")
    os.system("rm -f movie*.frame.png\n")
os.system("rm -f movie*frame\n")