import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import pylab
from pylab import *

import larch
from larch import Interpreter
from larch_plugins.xafs import autobk, xftf, xftr, feffit, pre_edge, find_e0
from larch_plugins.io import read_ascii
mylarch = Interpreter(with_plugins=False)

label_k     = r'$k \rm\,/\rm{\AA}^{-1}$'
label_r     = r'$R \rm\,/\rm{\AA}$'
label_chikw = r'$k^3\chi(k) /\rm{\AA}^{-3}$'
label_chir  = r'|FT| of $k^3\chi(k) /\rm{\AA}^{-4}$'

hc = 12398.52  # ? might be 12398.4193

si311 = 1.63747

def deg2ev311(de):
    ev = hc/(2*si311*np.sin(np.pi*de/180))
    return ev


### reading raw data ###
dat = read_ascii('./PdO.q',labels='angc ango ts i0 i1', _larch=mylarch)
dat.e = deg2ev311(dat.ango)
dat.mu=np.log(dat.i0/dat.i1)
date=24355.9037323484

pre_edge(dat.e, dat.mu, group=dat, e0=date,
         pre1=-150, pre2=-60, norm1=150, norm2=1100, nnorm=2, _larch=mylarch)
autobk(dat.e, dat.mu, group=dat, e0=dat.e0, rbkg=1,
       nclamp=5, clamp_lo=0, clamp_hi=24,
       kstep=0.05, kmin=0, kmax=17, kweight=2, win='hanning', _larch=mylarch)
xftf(dat.k, dat.chi, group=dat, kmin=3, kmax=15, dk=0.5, window='hanning', kweight=3, _larch=mylarch)

### reading Athena EXAFS data ###
dat2 = read_ascii('./AthenaLarch_PdO.q.chik',labels='k chi chik chik2 chik3 win energy', _larch=mylarch)
xftf(dat2.k, dat2.chi, group=dat2, kmin=3, kmax=15, dk=0.5, window='hanning', kweight=3, _larch=mylarch)

dat3 = read_ascii('./AthenaIfeffit_PdO.q.chik',labels='k chi chik chik2 chik3 win energy', _larch=mylarch)
xftf(dat3.k, dat3.chi, group=dat3, kmin=3, kmax=15, dk=0.5, window='hanning', kweight=3, _larch=mylarch)

## find energies in mu(e) data between 15 an 15.05 Ang^-1
from larch.utils import index_of
from larch_plugins.xafs import KTOE

e1 = KTOE*(15.00)**2
e2 = KTOE*(15.10)**2

ie1 = index_of(dat.e, date+e1)
ie2 = index_of(dat.e, date+e2)

print '%i energy points between k=15.00 and 15.10' % (ie2-ie1)

### plotting ###

plt.figure(figsize=(20,6))
plt.subplot(121)
plt.plot(dat.k,       dat.chi*dat.k**3, label='larch')
plt.plot(dat2.k, -5 + dat2.chi*dat2.k**3, label='athena/Larch'),
plt.plot(dat3.k,  5 + dat3.chi*dat3.k**3, label='athena/Ifeffit'),
plt.legend(), xlabel(label_k), ylabel(label_chikw), xlim(0,18), ylim(-30,30)
plt.subplot(122)
plt.plot(dat.r, dat.chir_mag, label='larch')
plt.plot(dat2.r, dat2.chir_mag, label='athena/Larch')
plt.plot(dat3.r, dat3.chir_mag, label='athena/Ifeffit')
plt.legend(), xlabel(label_r), ylabel(label_chir), xlim(0,6)
plt.show()