"""
Script to analyse XLD and XMCD spectra
"""
import matplotlib.lines as mlines
import matplotlib.pyplot as plt
import numpy as np
import math
import copy
import random
import h5py
import copy

FILE_LC = "../../Data/i09-75643.nxs"
FILE_RC = "../../Data/i09-75644.nxs"
delta_E = 0.05 #eV, absolute
delta_int = 0.01 #relative

import csv
def save(energies, spectrum, ext):
    name = 'LMO4MNin'
    with open('csv/'+name+'-'+ext+'.csv', 'wb') as f:
        writer = csv.writer(f)
        writer.writerow(('energy','intensity'))
        for n, s in enumerate(spectrum):
            writer.writerow((energies[n], s))
        f.close()

def getBackground(energies, spectrum):
    
    points = [3, 360, 530, 620, 765, 950, 1090]
    B = 1
    answ = "n"

    # Substracting pre-edge
    minimum = min(spectrum[points[5]:points[6]])
    spectrum = [s-minimum for s in spectrum]

    # Simulating step-function
    step = []
    start = spectrum[points[5]]
    mid = spectrum[points[3]]
    end = spectrum[points[1]]
    try:
        for e in energies:
            # fermi-dirac distribution-like normalisation function
            n = diff1/(math.exp((energies[points[2]]-e)/B)+1)
            m = diff2/(math.exp((energies[points[4]]-e)/B)+1)
            step.append(n+m)
    except:
        diff1 = mid - start
        diff2 = end - mid
        global diff1
        global diff2
        for e in energies:
            # fermi-dirac distribution-like normalisation function
            n = diff1/(math.exp((energies[points[2]]-e)/B)+1)
            m = diff2/(math.exp((energies[points[4]]-e)/B)+1)
            step.append(n+m)
    fitspect = [s - step[n] for n, s in enumerate(spectrum)]

##    while answ == "n":
##        plt.plot(energies, spectrum)
##        plt.plot(energies, [s+start for s in step])
##        for p in points:
##            plt.axvline(energies[p], color='k', linestyle='solid')
##        plt.show()
##            
##        answ = raw_input("Are the lines placed correctly? y/n: ")
##        while answ != "n" and answ != "y":
##            answ = raw_input("Use 'y' or 'n'. Are the lines correct?: ")
##
##        if answ == "y": break
##        print "Previous lines were places at ",points
##        for n, p in enumerate(points):
##            try: points[n] = int(raw_input("Line "+str(n+1)+" to: "))
##            except:
##                print "Invalid syntax"
##                points[n] = int(raw_input("Line "+str(n+1)+" to: "))
##        print points

    # Fitting 4th order polynomial to bg without step
    bg = range(points[0],points[1]) + [points[3]] + range(points[5],points[6])
    x, y = [energies[n] for n in bg], [fitspect[n] for n in bg]
    a, b, c, d, e =  np.polyfit(x, y, 4)
    norm = [a*i**4 + b*i**3 + c*i**2 + d*i + e for i in energies]
    spectNorm = [s-norm[n]+step[n] for n, s in enumerate(fitspect)]

    reg = [points[1],points[3],points[5]]
    return norm, spectNorm, reg

def getError(intensity, energies, delta_E, delta_int):
    
    error = []
    for n, e in enumerate(energies):
        try:
            a = (intensity[n+1] - intensity[n-1])/(energies[n+1]-energies[n-1])
            i_err = intensity[n] * delta_int
            e_err = a * delta_E
            error.append(math.sqrt(i_err**2 + e_err**2))
        except: error.append(0)

    return error

def getMagnetism(intensity1, intensity1_norm, intensity2, intensity2_norm, energies, reg):
    # Getting error expectation for each data point    
    intensity1_error = getError(intensity1, energies, delta_E, delta_int)
    intensity2_error = getError(intensity2, energies, delta_E, delta_int)

    # Putting together the XMCD spectrum    
    s = []
    error = []
    for n, e in enumerate(energies):
        new = intensity1_norm[n] - intensity2_norm[n]
        err = math.sqrt(intensity1_error[n]**2 + intensity2_error[n]**2)
        s.append(new)
        error.append(err)

        
    # Read data
    A, B, A_err, B_err = 0, 0, 0, 0
    for n in range(len(energies)):
        if n > reg[0] and n < reg[1]:
            A += s[n]
            A_err += error[n]**2
        elif n > reg[1] and n < reg[2]:
            B += s[n]
            B_err += error[n]**2
    A_err = math.sqrt(A_err)
    B_err = math.sqrt(B_err)

    C = 10.0         # beV
    # Calculate m_s
    m_s = 1/C*(2*B - A)
    m_s_err = 1/C* math.sqrt((2*B_err)**2 + A_err**2)

    # Calculate m_o
    m_o = -2/3/C*(A + B)
    m_o_err = -2/3/C * math.sqrt(A_err**2 + B_err**2)
    
    print "m_s has a value of "+str(m_s)+" +- "+str(abs(m_s_err/m_s))+"%"
    print "m_o has a value of "+str(m_o)+" +- "+str(abs(m_o_err/m_o))+"%"
    print "Ratio: "+str(m_o/m_s)+" +- "+str(math.sqrt((m_s_err/m_s)**2 + (m_o_err/m_o)**2))+" %"

    return m_o, m_o_err, m_s, m_s_err

if __name__ == '__main__':
    print "Gathering measurements data"
    
    with h5py.File(FILE_LC, 'r') as f:
        energies = f["/entry1/SX_Mn2p/energies"][0]
        
        spectrum = f["/entry1/SX_Mn2p/spectrum_data"][0]

        LCEnergies = [e for e in energies]
        LCSpectrum = [s for n, s in enumerate(spectrum)]

    with h5py.File(FILE_RC, 'r') as f:
        energies = f["/entry1/SX_Mn2p/energies"][0]
        
        spectrum = f["/entry1/SX_Mn2p/spectrum_data"][0]
        RCEnergies = [e for e in energies]
        RCSpectrum = [s for n, s in enumerate(spectrum)]

##    LCSpectrum = scaleIntensity(LCSpectrum, LCEnergies, RCEnergies)

    print "============"
    print "First background"
    print "============"
    LCNorm, LCSpectrum_norm, reg = getBackground(RCEnergies, LCSpectrum)
    print "============"
    print "Second background"
    print "============"
    RCNorm, RCSpectrum_norm, reg = getBackground(RCEnergies, RCSpectrum)

    XMCD = [(LCSpectrum_norm[n]-s)/max(LCSpectrum_norm) for n, s in enumerate(RCSpectrum_norm)]

    # Get magnetism
    m_o, m_o_err, m_s, m_s_err = getMagnetism(RCSpectrum, RCSpectrum_norm, LCSpectrum, LCSpectrum_norm, RCEnergies, reg)



    # Save data to CSV
    spectra = [LCSpectrum_norm, RCSpectrum_norm, XMCD]
    for n, s in enumerate(spectra):
        save(RCEnergies, s, str(n+1))

    plt.plot(RCEnergies, [0 for e in RCEnergies])
##    plt.plot(RCEnergies, LCSpectrum)
##    plt.plot(RCEnergies, LCNorm)
    plt.plot(RCEnergies, LCSpectrum_norm, label='$\sigma^+H^+$')
    plt.plot(RCEnergies, RCSpectrum_norm, label='$\sigma^-H^+$')
    plt.plot(RCEnergies, XMCD, label='$\sigma^+-\sigma^-$ x4 in $H^+$')
    
    legend = plt.legend(loc='upper right', shadow=True)
    title = plt.title("Mn L XMCD spectrum of in-situ STO|LMO(4) on Inconel plate")
    plt.xlabel("Energy (eV)")
    plt.ylabel("Intensity (a.u.)")
##    plt.text(RCEnergies[reg[1]]-5,min(XMCD), "m_s = "+str(m_s)+" +- "+str(abs(m_s_err/m_s))+"%\n"+
##                                    "m_o = "+str(m_o)+" +- "+str(abs(m_o_err/m_o))+"%\n"+
##                                    "Ratio = "+str(m_o/m_s)+" +- "+str(math.sqrt((m_s_err/m_s)**2 + (m_o_err/m_o)**2))+" %",
##             fontsize=10)
    
    plt.show()   


