import os
print os.getcwd()
import sys
print sys.version_info
print sys.executable
import numpy as np
import matplotlib.pyplot as plt
os.chdir('/Users/SeanFackler/anaconda/pkgs/xraylarch-0.9.29-py27_0/share/larch/plugins/math')
print os.getcwd()
from larch_plugins.math.lineshapes import gaussian
from larch.builtins import _group
from larch_plugins.wx.plotter import *
from larch.fitting import guess
from larch.fitting import minimize
from larch.interpreter import *


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=12)
    if not ("Counter 0" in scandata_f.columns):
	scandata_f = pd.read_csv(filename, sep='\t', skiprows=8)  # TrajScan files need 8 header lines somehow?
    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/151020/TrajScan21930_swf.txt")
# print scandata_f.Energy


def prepare_scan(scandata_f, datacounter="Counter 3", reference_counter='Counter 2'):
    # 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)


# had to import group and gaussian functions by finding them in the python library and making sure they were referred to
# correctly, e.g. changing group to _group and importing gaussian from _plugins.math.lineshapes.

mylarch = Interpreter()

# examples/fitting/doc_example1.lar
# create mock data

# mdat = _group(_larch=mylarch)
# mdat.x = np.linspace(-10, 10, 201)
# mdat.y = 1.0 + 12.0 * gaussian(mdat.x, 1.5, 2.0) + \
#          np.random.normal(size=len(mdat.x), scale=0.050)
#
# plt.plot(x,y)
# plt.show()

mdat = scandata_f
mdat.x = mdat.Energy
mdat.y = mdat.I_Norm0

print mdat.x
print mdat.y
# create a group of fit parameters
params = _group(off=guess(0),
		amp=guess(5, min=0),
		cen=guess(535),
		wid=guess(1, min=0), _larch=mylarch)  # Need an edge value for the first guess for the center of the
# guassian: 'cen='. Can use the edge-finder for that

init = params.off + params.amp * \
		    gaussian(mdat.x, params.cen, params.wid)


# define objective function for fit residual
def resid(p, data):
    return data.y - (p.off + p.amp * gaussian(data.x, p.cen, p.wid))


# perform fit
minimize(resid, params, args=(mdat,), _larch=mylarch)

# make final array
final = params.off + params.amp * \
		     gaussian(mdat.x, params.cen, params.wid)

# plot results
plt.plot(mdat.x, mdat.y)
plt.plot(mdat.x, init)
plt.plot(mdat.x, final)

if __name__ == "__main__":
    import sys
    import_fit(int(sys.argv[1]))