import os
import sys
import struct
import csv
from math import exp, sqrt
import numpy as np
from math import log, exp
from numpy import arange
from scipy.interpolate import interp1d
from matplotlib import pyplot as plt
from mpl_toolkits.mplot3d.axes3d import Axes3D
from matplotlib import cm


"""
Show ARPES data: raw data, 1st differential, 2nd differential
"""


#================================
# global parameters
#================================
infile = "PbSe_ARPES_25K.pxt"
rawcsvfile    = None
rawgpltfile   = None
diff1csvfile  = None
diff1gpltfile = None
diff2csvfile  = None
diff2gpltfile = None

# mode: 'r', '1', '2'
mode = 'r'

#================================
# for raw/smoothing data plots
#================================
# offset of Intensity for raw/smooth data plots: ratio of (max(Ilist) - min(Ilist))
rIoffset = 0.3

# # of E data ignored for plot
nEoffset = 30

# # of angle data to be plotted
nAplot = 20

# # of order for adaptive smoothing
nSmoothRaw = 15
nSmooth1   = 31
nSmooth2   = 51

#================================
# for contour plots
#================================
# # of contours
nlevels = 20

# colormap for contour plot: line, hsv, cool, iwnter, gray, gist_gray, cividis etc
#colormap = 'line'
colormap = 'cool'
#colormap = 'hsv'
#colormap = 'Spectral'


#===================================
# figure configuration
#===================================
figsize = [15, 10]
fontsize        = 12
legend_fontsize = 8

show_legend = 0


#=============================
# Treat argments
#=============================
def pfloat(str):
    try:
        return float(str)
    except:
        return None

def pint(str):
    try:
        return int(str)
    except:
        return None

def getarg(position, defval = None):
    try:
        return sys.argv[position]
    except:
        return defval

def getfloatarg(position, defval = None):
    return pfloat(getarg(position, defval))

def getintarg(position, defval = None):
    return pint(getarg(position, defval))

def usage():
    print("")
    print("Usage:")
    print(" python {} infile mode (nEoffset rIoffset nAplot nSmoothRaw nSmooth1 nSmooth2 nlevels colormap show_legend)"
                .format(sys.argv[0]))
    print("   mode: 'a'")
    print("   nEoffset: Number of angle data removed from the top")
    print("   rIoffset: Offset in the I axis for raw/smooth data plot expressed in the ratio to (max(I) - min(I))")
    print("   nAplot  : Maximum number of angle data for raw data plots")
    print("   nSmooth*: Number of smooth points for raw, 1st differential and 2nd differential data by adaptive smoothing")
    print("   colormap: 'line', 'cool', 'hsv', 'Spectral', 'PuOr', 'semismic', etc")
    print(" (1) ex: python {} {} {} {} {} {} {} {} {} {} {} {}"
                .format(sys.argv[0], infile, 'a', nEoffset, rIoffset, nAplot, nSmoothRaw, nSmooth1, nSmooth2, 
                        nlevels, colormap, show_legend))


def terminate(message = None):
    if message is not None:
        print("")
        print(message)

    print("")
    usage()
    print("")
    exit()

def updatevars():
    global infile, mode
    global nEoffset, rIoffset, nAplot
    global nSmoothRaw, nSmooth1, nSmooth2
    global nlevels, colormap, show_legend
    global rawcsvfile, diff1csvfile, diff2csvfile
    global rawgpltfile, diff1gpltfile, diff2gpltfile

    argv = sys.argv
#    if len(argv) == 1:
#        terminate()

    infile      = getarg     ( 1, infile)
    mode        = getarg     ( 2, mode)
    nEoffset    = getintarg  ( 3, nEoffset)
    rIoffset    = getfloatarg( 4, rIoffset)
    nAplot      = getintarg  ( 5, nAplot)
    nSmoothRaw  = getintarg  ( 6, nSmoothRaw)
    nSmooth1    = getintarg  ( 7, nSmooth1)
    nSmooth2    = getintarg  ( 8, nSmooth2)
    nlevels     = getintarg  ( 9, nlevels)
    colormap    = getarg     (10, colormap)
    show_legend = getintarg  (11, show_legend)

    header, ext     = os.path.splitext(infile)
    filebody        = os.path.basename(header)
    rawcsvfile      = filebody + '-raw.csv'
    diff1csvfile    = filebody + '-diff1.csv'
    diff2csvfile    = filebody + '-diff2.csv'
    rawgpltfile      = filebody + '-raw-gplt.txt'
    diff1gpltfile    = filebody + '-diff1-gplt.txt'
    diff2gpltfile    = filebody + '-diff2-gplt.txt'


def make_matrix2(nrow, ncolumn):
    matrix = []
    for ir in range(nrow):
        matrix.append([])
        for ic in range(ncolumn):
             matrix[ir].append(0.0)

    return matrix

def transpose(matrix):
    nrow = len(matrix)
    ncolumn =len(matrix[0])
    t = make_matrix2(ncolumn, nrow)
#    print("n1=", nrow, ncolumn)
#    print("n2=", len(t), len(t[0]))
    for i in range(nrow):
        for j in range(ncolumn):
            t[j][i] = matrix[i][j]

    return t

def savecsv(outfile, header, datalist):
    try: 
        print("Write to [{}]".format(outfile))
        f = open(outfile, 'w')
    except:
#    except IOError:
        print("Error: Can not write to [{}]".format(outfile))
    else:
        fout = csv.writer(f, lineterminator='\n')
        fout.writerow(header)
#        fout.writerows(data)
        for i in range(0, len(datalist[0])):
            a = []
            for j in range(len(datalist)):
                a.append(datalist[j][i])
            fout.writerow(a)
        f.close()

def save_ARPESdata(outfile, inf):
    try:
        f = open(outfile, 'w')
    except:
#    except IOError:
        print("Error: Can not write to [{}]".format(outfile))
    else:
        ylist = inf['ylist']
        f.write("{}\n".format(inf["Sample"]))
        for iQ in range(inf["nQ"]):
            Q = inf["Q0"] + iQ * inf["dQ"]
            f.write(",{}".format(Q))
        f.write("\n")
        
        for iE in range(inf["nE"]):
            E = inf["E0"] + iE * inf["dE"]
            f.write("{}".format(E))
            for iQ in range(inf["nQ"]):
                Q = inf["Q0"] + iQ * inf["dQ"]
                f.write(",{}".format(ylist[iQ][iE]))
            f.write("\n")

        f.close()

def read_ARPESdata(infile, is_print = 0):
    inf = {}
    inf["filename"] = infile

    f = open(infile, 'rb')
    if f is None:
        return None

    data = f.read()
    f.close()
#    print("data:", data)

    s = ''
    for i in range(40):
        c, = struct.unpack_from('B', data, 0x64 + i)
        if c == 0:
            break
        s += chr(c)
    inf["Sample"] = s

    inf["nE"], = struct.unpack_from('<i', data, 0x8C)
    inf["nQ"], = struct.unpack_from('<i', data, 0x90)
    inf["dE"], = struct.unpack_from('<d', data, 0x9C)
    inf["dQ"], = struct.unpack_from('<d', data, 0xA4)
    inf["E0"], = struct.unpack_from('<d', data, 0xBC)
    inf["Q0"], = struct.unpack_from('<d', data, 0xC4)

    fsize        = os.path.getsize(infile)
    inf["fsize"] = fsize
    data0        = 0x188
    lendata      = inf["nE"] * inf["nQ"] * 8
    meta0        = 0x188 + lendata
    lenmeta      = len(data) - meta0
    
    if is_print:
        print("file size:", fsize)
        print("Sample   :", inf["Sample"])
        print("nE       :", inf["nE"])
        print("nQ       :", inf["nQ"])
        print("E0       :", inf["E0"])
        print("Q0       :", inf["Q0"])
        print("position of data     :", data0, hex(data0))
        print("length of data       :", lendata)
        print("position of meta data:", meta0, hex(meta0))
        print("length of meta data  :", lenmeta)

    s = ''
    for i in range(meta0, fsize):
        c, = struct.unpack_from('B', data, i)
        if c == 0 or c == 0x0D or c == 0x0A:
            c = '\n'
        else:
            c = chr(c)
        s += c
    inf["meta"] = s
    
    if is_print:
        print("meta:\n", inf["meta"])

    nQ = inf["nQ"]
    nE = inf["nE"]
    ylist = make_matrix2(nQ, nE)
#    ylist = np.empty([nQ, nE], dtype = 'float')
    pos = 0x188
    for iQ in range(nQ):
        for iE in range(nE):
            ylist[iQ][iE], = struct.unpack_from('<d', data, pos)
            pos += 8
#            print("y:", ylist[iQ][iE])

    inf["ylist"] = ylist

    return inf


def save_3Ddata(outfile, header, Elist, Alist, matrixdata):
    try: 
        print("Write to [{}]".format(outfile))
        f = open(outfile, 'w')
    except:
#    except IOError:
        print("Error: Can not write to [{}]".format(outfile))
    else:
# make transposed list of matrixdata
        matrixdata_t = [list(x) for x in zip(*matrixdata)]
        fout = csv.writer(f, lineterminator='\n')
        fout.writerow(header)
        fout.writerow([''] + Alist)
        for i in range(len(Elist)):
            fout.writerow([Elist[0]] + matrixdata_t[i])
        f.close()

def save_3D_gplt_data(outfile, header, Elist, Alist, matrixdata):
    try: 
        print("Write to [{}]".format(outfile))
        f = open(outfile, 'w')
    except:
#    except IOError:
        print("Error: Can not write to [{}]".format(outfile))
    else:
# make transposed list of matrixdata
        matrixdata_t = [list(x) for x in zip(*matrixdata)]
        f.write(header + "\n")
        for ia in range(len(Alist)):
            for ie in range(len(Elist)):
                f.write("{:10.6f} {:8.4f} {:16.8g}\n".format(Elist[ie], Alist[ia], matrixdata[ia][ie]))
            f.write("\n")
        f.close()


def SmoothingByPolynomialFit(y, n):
    m = int(n / 2);
    W23 = (4.0 * m * m - 1.0) * (2.0 * m + 3.0) / 3.0;
    w23j = [0.0]*n
    for j in range(-m, m+1):
        w23j[j + m] = (3.0 * m * (m+1.0) - 1.0 - 5.0 * j * j) / W23

    ndata = len(y)
    ys = []
    for i in range(0, ndata):
        c = 0.0;
        ys.append(0.0);
        for j in range(-m, m+1):
            k = i + j
            if k < 0 or k >= ndata:
                continue
            ys[i] += w23j[j+m] * y[k]
            c += w23j[j+m]
        if c > 0:
            ys[i] /= c
        else:
            ys[i] = y[i]
    return ys;

def Diff1(x, y):
    ndata = len(y)
    diff = []
    for i in range(0, ndata):
        if i == 0:
            d = (y[1] - y[0]) / (x[1] - x[0])
        elif i == ndata - 1:
            d = (y[ndata-1] - y[ndata-2]) / (x[ndata-1] - x[ndata-2])
        else:
            d = (y[i+1] - y[i-1]) / (x[i+1] - x[i-1])

        diff.append(d)

    return diff

def Diff2(x, y):
    ndata = len(y)
    diff = []
    for i in range(0, ndata):
        if i <= 1:
            d = (y[2] - 2*y[1] + y[0]) / (x[1] - x[0])**2
        elif i >= ndata - 2:
            d = (y[ndata-1] - 2*y[ndata-2] + y[ndata-3]) / (x[ndata-1] - x[ndata-2])**2
        else:
            d = (y[i+1] - 2*y[i] + y[i-1]) / (x[i] - x[i-1])**2

        diff.append(d)

    return diff


def main():
    global infile, tvalue

    updatevars()

    print("")
    print("=========================================================")
    print(" View ARPES data and save them to csv files")
    print("=========================================================")
    print("Input data:")
    print("  infile     :", infile)
    print("  mode       :", mode)
    print("  Nomber of energy data removed to plot:")
    print("    nEoffset   :", nEoffset)
    print("  Offset to view different angle data in the ratio to the maximum intensity")
    print("    rIoffset   :", rIoffset)
    print("  Number of angle data to plot raw/smoothing data")
    print("    nAplot     :", nAplot)
    print("  Numbers of smoothing points for adaptive smoothing")
    print("    nSmoothRaw :", nSmoothRaw)
    print("    nSmooth1   :", nSmooth1)
    print("    nSmooth2   :", nSmooth2)
    print("  Contour plot configuration:")
    print("    nlevels    :", nlevels)
    print("    colormap   :", colormap)
    print("  Flag to plot legend for raw/smoothing plots: 0 or 1")
    print("    show_legend:", show_legend)
    print("")
    print("Output files:")
    print("  raw csv file  :", rawcsvfile)
    print("  diff1 csv file:", diff1csvfile)
    print("  diff2 csv file:", diff2csvfile)
    print("  raw gnuplot file  :", rawgpltfile)
    print("  diff1 gnuplot file:", diff1gpltfile)
    print("  diff2 gnuplot file:", diff2gpltfile)
    
    print("")
    print("Read [{}]".format(infile))
    inf = read_ARPESdata(infile)

    print("file size:", inf["fsize"])
    print("Sample   :", inf["Sample"])
    print("nE       :", inf["nE"])
    print("nQ       :", inf["nQ"])
    print("E0       :", inf["E0"])
    print("Q0       :", inf["Q0"])
    print("meta:\n", inf["meta"])
    print("")

    print("")
    print("Save data to [{}]".format(rawcsvfile))
    save_ARPESdata(rawcsvfile, inf)

    nE = inf["nE"]
    nQ = inf["nQ"]
    E0 = inf["E0"]
    Q0 = inf["Q0"]
    dE = inf["dE"]
    dQ = inf["dQ"]
    Elist = [E0 + iE * dE for iE in range(nE)]
    Alist = [Q0 + iQ * dQ for iQ in range(nQ)]
    Ilist = inf["ylist"]

    print("")
    print("{}-order adaptive smoothing".format(nSmoothRaw))
    Islist = []
    for i in range(len(Ilist)):
        Islist.append(np.array(SmoothingByPolynomialFit(Ilist[i], nSmoothRaw)))

    print("")
    print("1st differential smoothed by {}-th order".format(nSmooth1))
    Idiff1list = []
    for i in range(len(Ilist)):
        array = Diff1(Elist, Islist[i])
        array = SmoothingByPolynomialFit(array, nSmooth1)
        Idiff1list.append(np.array(array))

    print("")
    print("2nd differential smoothed by {}-th order".format(nSmooth2))
    Idiff2list = []
    for i in range(len(Ilist)):
        array = Diff2(Elist, Islist[i])
        array = SmoothingByPolynomialFit(array, nSmooth2)
        Idiff2list.append(np.array(array))

# save output files
    print("")
    print("Save smoothed raw data to {}".format(rawcsvfile))
    save_3Ddata(rawcsvfile,   ['v:Energy (eV)', '>:Angle (deg)'], Elist, Alist, Islist)
    print("Save 1st differential data to {}".format(diff1csvfile))
    save_3Ddata(diff1csvfile, ['v:Energy (eV)', '>:Angle (deg)'], Elist, Alist, Idiff1list)
    print("Save 2nd differential data to {}".format(diff2csvfile))
    save_3Ddata(diff2csvfile, ['v:Energy (eV)', '>:Angle (deg)'], Elist, Alist, Idiff2list)
    print("Save smoothed raw data to {}".format(rawgpltfile))
    save_3D_gplt_data(rawgpltfile,   "# Energy (eV), Angle (deg), Intensity", Elist, Alist, Islist)
    print("Save 1st differential data to {}".format(diff1gpltfile))
    save_3D_gplt_data(diff1gpltfile, "# Energy (eV), Angle (deg), dI/dE", Elist, Alist, Idiff1list)
    print("Save 2nd differential data to {}".format(diff2gpltfile))
    save_3D_gplt_data(diff2gpltfile, "# Energy (eV), Angle (deg), d2I/dE2", Elist, Alist, Idiff2list)
    
# remove the first nEoffset data
    print("")
    print("Remove the first {} data".format(nEoffset))
    Elist         = Elist[nEoffset:]
    for i in range(nQ):
        Ilist[i]      = Ilist[i][nEoffset:]
        Islist[i]     = Islist[i][nEoffset:]
        Idiff1list[i] = Idiff1list[i][nEoffset:]
        Idiff2list[i] = Idiff2list[i][nEoffset:]


#=============================
# Plot graphs
#=============================
    print("")
    print("plot")
    nASkip = int(len(Ilist) / nAplot)
    print("nASkip:", nASkip)

    fig = plt.figure(figsize = figsize)

    ax1 = fig.add_subplot(3, 3, 1)
    ax2 = fig.add_subplot(3, 3, 2)
    ax3 = fig.add_subplot(3, 3, 3, projection='3d')
    ax4 = fig.add_subplot(3, 3, 4)
    ax5 = fig.add_subplot(3, 3, 5)
    ax6 = fig.add_subplot(3, 3, 6, projection='3d')
    ax7 = fig.add_subplot(3, 3, 7)
    ax8 = fig.add_subplot(3, 3, 8)
    ax9 = fig.add_subplot(3, 3, 9, projection='3d')

# Raw/smoothing data plots
    Earray = np.array(Elist)
    offset = 0.0
    for i in range(0, len(Ilist), nASkip):
        Iarray = np.array(Ilist[i])
        ax1.plot(Earray, Iarray + offset, label = '{} $^\circ$'.format(Alist[i]), 
                linestyle = '', marker = 'o', markersize = 1.5, markeredgewidth = 0, markerfacecolor = 'black')
        ax1.plot(Earray, Islist[i] + offset, color = 'red', linewidth = 0.5)
        offset += rIoffset * (max(Iarray) - min(Iarray))
    ax1.set_xlim([min(Elist), max(Elist)])
    ax1.set_ylim([0.0, ax1.get_ylim()[1]])
    ax1.set_xlabel('$E$ (eV)', fontsize = fontsize)
    ax1.set_ylabel('Inensity', fontsize = fontsize)
    if show_legend:
#        ax1.legend(fontsize = legend_fontsize)
        ax1.legend(bbox_to_anchor=(1.05, 1.0), loc='upper left', borderaxespad = 0, fontsize = legend_fontsize)
    ax1.tick_params(labelsize = fontsize)

    aax, eax = np.meshgrid(Alist, Elist)

    Islist_t = np.array(Islist).T
    if colormap == 'line':
        ax2.contour(aax, eax, Islist_t, levels = nlevels, colors = ['black'], linewidths = 0.2)
    else:
        ax2.contourf(aax, eax, Islist_t, levels = nlevels, cmap = colormap)
    ax2.set_xlabel('Angle ($^\circ$)', fontsize = fontsize)
    ax2.set_ylabel('$E$ (eV)',         fontsize = fontsize)
    ax2.tick_params(labelsize = fontsize)

    ax3.plot_wireframe(aax, eax, Islist_t, color = 'blue', linewidth = 0.3)#, rstride=2, cstride=2)
    ax3.set_xlabel('Angle ($^\circ$)', fontsize = fontsize)
    ax3.set_ylabel('$E$ (eV)'        , fontsize = fontsize)
    ax3.set_zlabel('Intensity'       , fontsize = fontsize)
    ax3.tick_params(labelsize = fontsize)

# 1st differential plots
    offset = 0.0
    for i in range(0, len(Idiff1list), nASkip):
        Idiff1array = np.array(Idiff1list[i])
        ax4.plot(Earray, Idiff1array + offset, color = 'red', linewidth = 0.5)
        offset += rIoffset * (max(Idiff1array) - min(Idiff1array))
    ax4.set_xlim([min(Elist), max(Elist)])
    ax4.set_xlabel('$E$ (eV)', fontsize = fontsize)
    ax4.set_ylabel('$dI/dE$', fontsize = fontsize)
    ax4.tick_params(labelsize = fontsize)

    Idiff1list_t = np.array(Idiff1list).T
    if colormap == 'line':
        ax5.contour(aax, eax, Idiff1list_t, levels = nlevels, colors = ['black'], linewidths = 0.2)
    else:
        ax5.contourf(aax, eax, Idiff1list_t, levels = nlevels, cmap = colormap)
    ax5.set_xlabel('Angle ($^\circ$)', fontsize = fontsize)
    ax5.set_ylabel('$E$ (eV)',         fontsize = fontsize)
    ax5.tick_params(labelsize = fontsize)

    ax6.plot_wireframe(aax, eax, Idiff1list_t, color = 'blue', linewidth = 0.3)#, rstride=2, cstride=2)
    ax6.set_xlabel('Angle ($^\circ$)', fontsize = fontsize)
    ax6.set_ylabel('$E$ (eV)'        , fontsize = fontsize)
    ax6.set_zlabel('$dI/dE$'         , fontsize = fontsize)
    ax6.tick_params(labelsize = fontsize)

# 2nd differential plots
    offset = 0.0
    for i in range(0, len(Idiff2list), nASkip):
        Idiff2array = np.array(Idiff2list[i])
        ax7.plot(Earray, Idiff2array + offset, color = 'red', linewidth = 0.5)
        offset += rIoffset * (max(Idiff2array) - min(Idiff2array))
    ax7.set_xlim([min(Elist), max(Elist)])
    ax7.set_xlabel('$E$ (eV)', fontsize = fontsize)
    ax7.set_ylabel('$d^2I/dE^2$', fontsize = fontsize)
    ax7.tick_params(labelsize = fontsize)

    Idiff2list_t = np.array(Idiff2list).T
    if colormap == 'line':
        cs8 = ax8.contour(aax, eax, Idiff2list_t, levels = nlevels, colors = ['black'], linewidths = 0.2)
    else:
        cs8 = ax8.contourf(aax, eax, Idiff2list_t, levels = nlevels, cmap = colormap)
#    cbar8 = ax8.colorbar(cs8)
#    cbar8.set_label('color bar')
    ax8.set_xlabel('Angle ($^\circ$)', fontsize = fontsize)
    ax8.set_ylabel('$E$ (eV)',         fontsize = fontsize)
    ax8.tick_params(labelsize = fontsize)

    ax9.plot_wireframe(aax, eax, Idiff2list_t, color = 'blue', linewidth = 0.3)#, rstride=2, cstride=2)
    ax9.set_xlabel('Angle ($^\circ$)', fontsize = fontsize)
    ax9.set_ylabel('$E$ (eV)'        , fontsize = fontsize)
    ax9.set_zlabel('$d^2I/dE^2$'     , fontsize = fontsize)
    ax9.tick_params(labelsize = fontsize)

    plt.tight_layout()

    plt.show()
    print("")
    print("Close graph window to proceed")

    """
    plt.pause(0.1)

    print("")
    print("Press ENTER to exit>>", end = '')
    input()
    """

    terminate()


if __name__ == '__main__':
    main()