import pandas as pd
import matplotlib.pylab as plt
from larch import *
from larch_plugins.xafs import pre_edge
from larch_plugins.xafs import mback

scandata_f = []

#Working for SigScans

def rdin(filename):
    scandata_f = pd.read_csv(filename, sep='\t', skiprows=15)
    if not ("Counter 0" in scandata_f.columns):
	scandata_f = pd.read_csv(filename, sep='\t', skiprows=14)
    print scandata_f.columns

    if not ("Counter 0" in scandata_f.columns):
	print ("Problem with header. skipping 12 or 10 lines did not make it. Check input file.")
	return None
    return scandata_f


scandata_f = rdin("/Users/SeanFackler/data/ALS092016/160923/SigScan26827.txt")
# print scandata_f.Energy


def prepare_scan(scandata_f, datacounter="TEY_up", reference_counter='Izero'):
    # Preparing Scan (normalization)
    if 'Counter 4' in scandata_f.columns:
	clockname = 'Counter 4'
    elif 'Counter 6' in scandata_f.columns:
	clockname = 'Counter 6'
    else:
	print("No counter for clock found (looked for 'Counter 4' and 'Counter 6'). Defaulting to 'Counter 0'.")
	clockname = 'Counter 0'

    scandata_f["I_Norm0"] = scandata_f[datacounter].astype(float) / scandata_f[reference_counter].astype(float)
    scandata_f["I_Normt"] = scandata_f[datacounter].astype(float) / scandata_f[clockname].astype(float)
    scandata_f["Energy"] = scandata_f["Energy"].round(1)
    # scandata_f["Z"] = scandata_f["Z"].round(2)
    return (scandata_f)


prepare_scan(scandata_f)

mylarch = Interpreter()

scandata_f.e = scandata_f.Energy#
scandata_f.mu = scandata_f.I_Norm0#


mback(scandata_f.e, scandata_f.mu, group=scandata_f, z=8, edge='K', order=1, whiteline=6, fit_erfc=True,
      _larch=mylarch)
print scandata_f.mback_params.e0

plt.plot(scandata_f.e, scandata_f.f2)
plt.plot(scandata_f.e, scandata_f.fpp)
plt.plot(scandata_f.e, 5*scandata_f.mu)
plt.show()
