#! /usr/bin/env python

import sys
import csv
import matplotlib.pyplot as plt
from matplotlib.mlab import griddata
import numpy as np

print "***********************************************************************************"
print "** This script generates a quick plot with the csv file from OCLIMAX.            **"
print "** Usage: pclimax.py [-unit_tag] *.csv                                           **"
print "**        -m : meV,  -w : wavenumber (cm-1),  -t : THz   (default: unit in csv)  **"
print "***********************************************************************************"

if sys.argv[1].startswith('-'):
    if sys.argv[1]=='-m':
        plt_unit = 'meV'
    elif sys.argv[1]=='-w':
        plt_unit = 'cm-1'
    elif sys.argv[1]=='-t':
        plt_unit = 'THz'
    else:
        print "-unit_tag not recognized"
        sys.exit()
    csvfile = sys.argv[2]
else:
    csvfile = sys.argv[1]

tab = list(csv.reader(open(csvfile,'r')))

INSTR = tab[2][1]
dE_UNIT = tab[2][4]
if int(dE_UNIT) == 0:
    csv_unit = 'cm-1'
elif int(dE_UNIT) == 1:
    csv_unit = 'meV'
elif int(dE_UNIT) == 2:
    csv_unit = 'THz'

try:
    plt_unit
except NameError:
    scale = 1.0
    unit = csv_unit
else:
    unit = plt_unit
    if plt_unit==csv_unit:
        scale = 1.0
    elif plt_unit=='cm-1' and csv_unit=='meV':
        scale = 8.06554
    elif plt_unit=='cm-1' and csv_unit=='THz':
        scale = 33.35641
    elif plt_unit=='meV' and csv_unit=='cm-1':
        scale = 0.12399
    elif plt_unit=='meV' and csv_unit=='THz':
        scale = 4.13567
    elif plt_unit=='THz' and csv_unit=='cm-1':
        scale = 0.02998
    elif plt_unit=='THz' and csv_unit=='meV':
        scale = 0.24180

i=0
while i<len(tab):
    if not tab[i]:
        del tab[i]
        continue
    try:
        float(tab[i][0])
    except ValueError:
        del tab[i]
        continue
    i=i+1

if int(INSTR) in [0,1,2]:
    ttab=[[tab[row][col] for row in range(len(tab))] for col in range(len(tab[0]))]
    ns = len(tab[0])-1
    ttab[0]=[float(ttab[0][i])*scale for i in range(len(ttab[0]))]
    for i in range(ns):
        plt.plot(ttab[0],ttab[i+1],label='Spec'+str(i+1))
    plt.xlabel('Energy transfer ('+unit+')')
    plt.ylabel('Intensity (a.u.)')
    plt.title(csvfile)
    plt.legend()
    plt.show()
elif int(INSTR)==3:
    ttab=[[tab[row][col] for row in range(len(tab))] for col in range(len(tab[0]))]
    x=[float(ttab[0][i]) for i in range(len(ttab[0]))]
    y=[float(ttab[1][i])*scale for i in range(len(ttab[1]))]
    z=[float(ttab[2][i]) for i in range(len(ttab[2]))]
    xi=np.sort(np.asarray(list(set(x))))
    yi=np.sort(np.asarray(list(set(y))))
    Z = np.zeros((len(yi),len(xi)))
    for i in range(len(yi)):
        for j in range(len(xi)):
            k = i+j*len(yi)
            Z[i][j]=z[k]
    #X, Y =np.meshgrid(xi,yi)
    #Z = griddata(x,y,z,xi,yi,interp='linear')
    asp = float(max(xi))/float(max(yi))
    plt.imshow(Z,origin='lower',extent=[min(x),max(x),min(y),max(y)],aspect=asp)
    plt.colorbar()
    plt.xlabel(r'Momentum transfer (1/$\AA$)')
    plt.ylabel('Energy transfer ('+unit+')')
    plt.title(csvfile)
    plt.show()