import os
import sys
import struct
import numpy as np
from matplotlib import pyplot as plt


"""
Convert ARPES raw data file to text file
"""


#================================
# global parameters
#================================
infile     = "PbSe_ARPES_25K.pxt"
outcsvfile = None

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

rIoffset = 0.5

#=============================
# 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"
                .format(sys.argv[0]))
    print(" ex: python {} {}".format(sys.argv[0], infile))



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

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

def updatevars():
    global infile, outcsvfile

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

    infile    = getarg   ( 1, infile)

    header, ext     = os.path.splitext(infile)
    filebody        = os.path.basename(header)
    infile     = filebody + '.pxt'
    outcsvfile = filebody + '.csv'

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

    try:
        f = open(infile, 'rb')
    except:
        terminate("Error: Can not read [{}]".format(infile))

    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 = 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_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 main():
    global infile, inifile, outfile

    updatevars()

    print("")
    print("=========================================================")
    print(" Convert ARPES raw data (.pxt) file to text file")
    print("=========================================================")
    print("Input data:")
    print("  infile     :", infile)
    print("Output data:")
    print("  output csv file :", outcsvfile)

    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(outcsvfile))
    save_ARPESdata(outcsvfile, inf)


#=============================
# Plot graphs
#=============================
    print("")
    print("plot")

    fig = plt.figure(figsize = figsize)

    ax1 = fig.add_subplot(1, 1, 1)

    ylist = inf["ylist"]
    Elist = [inf["E0"] + iE * inf["dE"] for iE in range(inf["nE"])]
    offset = 0.0
    for iQ in range(inf["nQ"]):
        ax1.plot(Elist, ylist[iQ] + offset, linestyle = '-', linewidth = 0.3, color = 'black')
        offset += rIoffset * (max(ylist[iQ]) - min(ylist[iQ]))

    ax1.set_xlabel('$E$ (eV)', fontsize = fontsize)
    ax1.set_ylabel('Inensity', fontsize = fontsize)
    ax1.tick_params(labelsize = fontsize)

    plt.tight_layout()

    plt.pause(0.1)

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

    terminate()


if __name__ == '__main__':
    main()