import os
import sys
import csv
import re
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


"""
Search possible elements and levels for photoemission spectroscopy
"""

argv = sys.argv
prog_path = argv[0]

Source1 = 'XPS-BE-Source.csv'
#Source2 = 'XPS-BE-Small.csv'
if not os.path.isfile(Source1):
    dirname     = os.path.dirname(prog_path)
    Source1 = os.path.join(dirname, Source1)

SearchVal = '' #510; # eV
SearchErr = '' #  5; # eV


#=============================
# General functions
#=============================
def pfloat(s, strict = True, defval = 0.0):
    try:
        return float(s)
    except:
        pass
    if s is None or strict:
        return defval

    match = re.search(r'([\+\-]?[\d\.eE]+)', s)
    if not match:
        return defval

    val = match.group()
    try:
        return float(val)
    except:
        if defval is None:
            return s
        return defval

def Search(p, t, e):
    print("Hit for {} - {} - {} eV ...".format(t - e, t, t + e))
    for key in p.keys():
        v = p[key]
#        print("comp: $v ($t +- $e ?)")
        if t - e <= v <= t + e:
            an, name, st = key.split()
            print("  {:>2} {:>2} {:>8}   {:>8.2f} eV".format(an, name, st, v))

def ReadCSV(infile):
    try: 
        print("Read [{}]".format(infile))
        f = open(infile, 'r')
    except:
#    except IOError:
        print("Error: Can not read [{}]".format(infile))
    else:
        hash = {}
        header = None
        st = []
        val = []
        while 1:
            line = f.readline()
            if not line:
                break

            if 's' in line and 'p' in line and 'd' in line:
                st = line.split(',')
#                print("st=", st)
#                exit()
            else:
                val = line.split(',')
#                print("  v0=", val)
#                exit()
                for i in range(len(val)):
#                    print("val[{}]={}".format(i,val[i]))
                    if val[i] == '':
                        val[i] = None
                    elif re.search(r'([A-Z][a-z]?)', val[i]):
                        pass
                    else:
                        val[i] = pfloat(val[i], False, 0.0)
#                print("  v=", val)
                for i in range(2, len(val)):
                    if val[i]:
                        o = st[i]
                        key = "{} {} {}".format(val[0], val[1], o)
                        hash[key] = pfloat(val[i])
                        print("  {} {}".format(key, hash[key]))
        f.close()

    return hash


def main():
    global SearchVal, SearchErr

    print("")
    print("====================================")
    print("  Search possible elements and peaks from binding energy and possible range")
    print("====================================")
    print("")

    if SearchVal == '':
        print("Input search value (Def: 510 eV): ", end = '')
        line = input().strip()
        
        if line == '':
            SearchVal = 510.0
        else:
            SearchVal = pfloat(line)

    if SearchErr == '':
        print("Input error value (Def: 10 eV): ", end = '')
        line = input().strip()
        if line == '':
            SearchErr = 10.0
        else:
            SearchErr = pfloat(line)

    print("Read [{}] =====================================".format(Source1))
    pBE1 = ReadCSV(Source1);
#    print "Read [{}] =====================================".format(Source2));
#    pBE2 = ReadCSV(Source2);
    print("File read end=====")

    print("")
    print("=====================================")
    Search(pBE1, SearchVal, SearchErr)

    print("")


if __name__ == '__main__':
    main()