import matplotlib.pyplot as plt
import pandas as pd
import numpy as np
import datetime
import math
import re

import larch
from larch import Interpreter, Group
from larch_plugins.math.mathutils import index_of
from larch.fitting import minimize, param, fit_report
from larch_plugins.math.lineshapes import gaussian, erf
from larch_plugins.xafs import pre_edge

mylarch = Interpreter()


def rdin(filename):
    scandata_f = pd.read_csv(filename, sep='\t', skiprows=10)
    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


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


def is_stacked(scandata_f, num_of_scans=0, stacked=0):
    """Determine stacked or single scan file(s) and informs user"""
    your_list = scandata_f.Energy
    # Number of scans in file using set to find repeating numbers in energy column
    num_of_scans = len(your_list) / len(set(your_list))
    # Inform user whether file is stacked (multiple scans in one file, or not)
    if len(your_list) != len(set(your_list)):
	stacked = 1
	print 'Stacked scan file'
	# Inform user about number of scans in stack
	print 'Number of scans in stack is', num_of_scans
    else:
	print 'Individual scan file'
	stacked = 0

    return num_of_scans, your_list, stacked


def unstacking(scandata_f, num_of_scans, scan_rng=0, unstacked = pd.DataFrame(np.array([]))):
    """Unstack a stacked scan file"""
    if stacked:
	bins = len(set(your_list))
	print bins, 'bins in scan'
	scan_rng = range(1, num_of_scans + 1, 1)
	unstacked = pd.DataFrame(np.array(scandata_f.Energy[0 * bins:(1 * bins)]), columns=['Energy'])
	for idx in range(1, num_of_scans + 1, 1):
	    unstacked[str(idx)] = np.array(scandata_f.I_Norm0.iloc[(idx - 1) * bins:(idx * bins)])
	return scan_rng, unstacked, bins


def make_model(pars, data, components=False):
    """make model of spectra: 3 peak functions, 1 erf function, offset"""
    p1 = pars.amp1 * gaussian(data.e, pars.cen1, pars.wid1)
    p2 = pars.amp2 * gaussian(data.e, pars.cen2, pars.wid2)
    p3 = pars.amp3 * gaussian(data.e, pars.cen3, pars.wid3)

    e1 = pars.off + pars.erf_amp * erf(
	pars.erf_wid * (data.e - pars.erf_cen))  # pars.off + pars.erf_amp * erf(pars.erf_wid*(mdat.e - pars.erf_cen))
    sum = p1 + p2 + p3 + e1
    if components:
	return sum, p1, p2, p3, e1
    return sum


def resid(pars, data):
    """fit residual"""
    return make_model(pars, data) - data.y

file_name = "TrajScan21930.txt"

scandata_f = rdin(file_name)

prepare_scan(scandata_f)

num_of_scans, your_list, stacked = is_stacked(scandata_f, num_of_scans=0, stacked=0)

scan_rng, unstacked, bins = unstacking(scandata_f, num_of_scans, scan_rng=0, unstacked = pd.DataFrame(np.array([])))

_try = "try2_TFY_"  # define the naming for results!

AMP1 = []
AMP2 = []
AMP3 = []
CEN1 = []
CEN2 = []
CEN3 = []
ERF_AMP = []
ERF_WID = []
OFF = []
WID1 = []
WID2 = []
WID3 = []

# create group for parameters
print 'defining parameters'
params = larch.Group(
    cen1=param(532.2, vary=True, min=531.8, max=532.7),
    cen2=param(536.5, vary=True, min=536.3, max=538),
    cen3=param(543, vary=True, min=541.2, max=543.4),

    amp1=param(0.25, vary=True, min=0, max=3),
    amp2=param(0.25, vary=True, min=0, max=3),
    amp3=param(0.25, vary=True, min=0, max=3),

    wid1=param(0.2, vary=True, min=0.05, max=1),
    wid2=param(0.2, vary=True, min=0.05, max=5),
    wid3=param(0.2, vary=True, min=0.05, max=5),

    off=param(0.05, vary=True, min=0, max=0.1),

    erf_amp=param(0.05, vary=True, min=0, max=.12),
    erf_wid=param(0.2, vary=True, min=0.1, max=5),
    erf_cen=531.1  # param(531, vary=True, min=530, max=532)
)

for idx in range(1, num_of_scans + 1, 1):
    print 'for loop number', idx

    # Now we fit, save, and plot results

    mdat = unstacked
    mdat.x = mdat.Energy
    mdat.mu = unstacked[str(idx)]

    # do pre-processing steps, here XAFS pre-edge removal
    print 'leveling'

    # Manually set white line position (e0) for algorithm result stability (false edge detection)
    pre_edge(mdat.x, mdat.mu, pre1=-10, pre2=-2, e0=531.1, group=mdat, _larch=mylarch)
    # time.sleep(4)
    # print 'pause'
    mdat.y = mdat.mu - mdat.pre_edge
    mdat[str(idx)] = mdat.y
    # select data range to be considered in the fit here.
    # note that this could be done inside the objective function,
    # but doing these steps here means it done only once.
    i1, i2 = index_of(mdat.x, unstacked.Energy[50]), index_of(mdat.x, unstacked.Energy.iloc[-1])
    mdat.e = mdat.x[i1 + 1:i2 + 1]
    mdat.y = mdat.y[i1 + 1:i2 + 1]

    print datetime.datetime.now().time().isoformat()
    print 'minimizing least squares...'
    m = minimize(resid, params, args=(mdat,), _larch=mylarch, maxfev=120000)
    print 'minimizing done!'
    # Export results
    full_report = fit_report(params, show_correl=False, _larch=mylarch)
    report = "\r \n" + 'Sample ' + str(idx) + "\n" + file_name + "\n" + full_report
    print report

    # Extract numbers from report and write to files
    results = re.findall("[-+]?\d+[\.]?\d*[eE]?[-+]?\d*", full_report)
    print results

    amp1 = float(results[7])
    amp2 = float(results[11])
    amp3 = float(results[15])
    cen1 = float(results[19])
    cen2 = float(results[23])
    cen3 = float(results[27])
    erf_amp = float(results[30])
    erf_wid = float(results[33])
    off = float(results[36])
    wid1 = float(results[40])
    wid2 = float(results[44])
    wid3 = float(results[48])

    # wid variables are like sigma of guassian, and amp variables meaning amplitude
    # need to be divided by 1/sigma*sqrt(2pi)

    AMP1.append(amp1)
    AMP2.append(amp2)
    AMP3.append(amp3)
    CEN1.append(cen1)
    CEN2.append(cen2)
    CEN3.append(cen3)
    ERF_AMP.append(erf_amp)
    ERF_WID.append(erf_wid)
    OFF.append(off)
    WID1.append(wid1)
    WID2.append(wid2)
    WID3.append(wid3)

    # now plot results
    final, f1, f2, f3, e1 = make_model(params, mdat, components=True)
    # raw_input("Press Enter to continue...")  # Use this to pause to briefly inspect fit
    #

    plt.figure(idx)
    plt.clf()
    # plt.plot(mdat.x, unstacked[str(idx)])
    plt.plot(mdat.e, mdat.y, color='red', linewidth=2.0)  # ,
    plt.plot(mdat.e, final, color='purple')
    plt.plot(mdat.e, mdat.y - final, color='grey')  # ,
    # plt.plot(mdat.x, mdat.pre_edge, color='blue')
    # For plotting gaussian peak positions I needed to divide out square root term
    g_scale1 = 1 / (wid1 * (2 * math.pi) ** 0.5)
    g_scale2 = 1 / (wid2 * (2 * math.pi) ** 0.5)
    g_scale3 = 1 / (wid3 * (2 * math.pi) ** 0.5)

    plt.plot([cen1, cen2, cen3], [(amp1 * g_scale1) + 2 * erf_amp,
				  (amp2 * g_scale2) + 2 * erf_amp,
				  (amp3 * g_scale3) + 2 * erf_amp], 'ro')
    plt.plot(mdat.e, f1)  # label='peak1'
    plt.plot(mdat.e, f2)  # label='peak2'
    plt.plot(mdat.e, f3)  # label='peak2'
    plt.plot(mdat.e, e1)  # label='erf +offset'
    plt.ylabel('XAS, fits')
    plt.xlabel('Energy (eV)')
    plt.savefig('results/' + _try + str(idx) + '.png')  # define the naming for results!
    print 'saved figure'
    # plt.show()
    plt.close()
    # report.to_csv('my_csv.csv', mode='a', header=False)
    with open(_try + '_gaussian.txt', "a") as text_file:  # define the naming for results!
	text_file.write(report)
	#
	# end of examples/fitting/doc_example2a.lar

	# raw_input("Press Enter to continue...")  # Use this to pause to briefly inspect fit

# Make plots of fitted results like gaussian position and height

fig = plt.figure()
ax1 = fig.add_subplot(111)
ax1.plot(CEN1, scan_rng, 'b-')  #
ax1.plot(CEN2, scan_rng, 'b-')  #
ax1.plot(CEN3, scan_rng, 'b-')  #
ax1.set_xlabel('Energy (eV)')
ax1.set_ylabel('Samples position (mm)')
ax2 = ax1.twinx()
s2_1 = np.array(mdat['6'])
s2_2 = np.array(mdat['39'])
s2_3 = np.array(mdat['45'])
ax2.plot(unstacked.Energy, 2 * s2_1 - 0.3, 'k-')
ax2.plot(unstacked.Energy, 6 * s2_2 + 0.9, 'r-')
ax2.plot(unstacked.Energy, s2_3 + 1.45, 'g-')
ax2.set_ylabel('XAS')
fig.savefig('results/' + _try + '_CEN_spectra.png')  # define the naming for results!
print 'saved figure'
plt.show()
plt.close()

plt.clf()
plt.plot(CEN1, scan_rng)  #
plt.plot(CEN2, scan_rng)  #
plt.plot(CEN3, scan_rng)  #
plt.xlabel('Energy (eV)')
plt.ylabel('XAS, fits')
plt.savefig('results/' + _try + '_CEN.png')  # define the naming for results!
print 'saved figure'
plt.show()
plt.close()

# plot ratio of peaks from ratio calculation below

G_SCALE1 = 1 / (np.array(WID1) * (2 * math.pi) ** 0.5)
G_SCALE2 = 1 / (np.array(WID2) * (2 * math.pi) ** 0.5)
G_SCALE3 = 1 / (np.array(WID3) * (2 * math.pi) ** 0.5)

HGHT1 = (np.array(AMP1) * G_SCALE1) + 2 * np.array(ERF_AMP)
HGHT2 = (np.array(AMP2) * G_SCALE2) + 2 * np.array(ERF_AMP)
HGHT3 = (np.array(AMP3) * G_SCALE3) + 2 * np.array(ERF_AMP)

HGHT12_ratio = HGHT1/HGHT2

plt.clf()
plt.plot(HGHT12_ratio, scan_rng)  #
plt.xlabel('Energy (eV)')
plt.ylabel('XAS peak ratio second/first')
plt.savefig('results/' + _try + '_AMP_ratio.png')  # define the naming for results!
print 'saved figure'
plt.show()
plt.close()

delta_t2g_eg = np.array(CEN2) - np.array(CEN1)
# print delta_t2g_eg
# np.savetxt('results/delta_t2g_egtest.txt', delta_t2g_eg, fmt='%.6f'))

delta_t2g_d = np.array(CEN3) - np.array(CEN1)
# print delta_t2g_d
# np.savetxt('results/delta_t2g_dtest.txt', delta_t2g_d, fmt='%.6f')

matrix = np.matrix([CEN1, CEN2, CEN3, delta_t2g_eg, delta_t2g_d, AMP1, AMP2, AMP3,
		    WID1, WID2, WID3, ERF_AMP, ERF_WID, OFF])
np.savetxt('results/' + _try + 'O_K-edge_results.txt', matrix.T, fmt='%.6f',
	   header='cen1(eV)    cen2(eV)    cen3(eV)    t2g_eg(eV)  t2g_d(eV)  amp1    amp2    amp3  wid1(eV)'
		  '    wid2(eV)    wid3(eV)    erf_amp erf_wid offset')