# Imports
import numpy as np
import pandas as pd
import xarray as xr

# Special functions
# from scipy.special import sph_harm
import spherical_functions as sf
import quaternion

# Performance & benchmarking libraries
# from joblib import Memory
# import xyzpy as xyz
import numba as nb

# Timings with ttictoc or time
# https://github.com/hector-sab/ttictoc
# from ttictoc import TicToc
import time

# Package fns.
# For module testing, include path to module here
import sys
import os
# modPath = r'D:\code\github\ePSproc'  # Win test machine
modPath = r'/home/femtolab/github/ePSproc/'  # Linux test machine
sys.path.append(modPath)
import epsproc as ep
# TODO: tidy this up!
from epsproc.util import matEleSelector
from epsproc.geomFunc import geomCalc, geomUtils
from epsproc.geomFunc.lfblmGeom import lfblmXprod

"""
Explanation: ePSproc LF/AF function verification & tests
19/06/20
For LF and AF calculations, trying to get to the bottom of issues with magnitudes and/or phases and/or formalism differences with raw ePS matrix elements.
Formalism
Test cases: 

ePS matrix elements with formalism from [1], for LF cross-sections and $\beta_{2}$
ePSproc AF calculations, for LF cross-sections and $\beta_{2}$. 

The AF calculations should reduce to the LF case for an isotropic ensemble, and both cases should match the "direct" ePS GetCro outputs (LF). Hopefully this should clear up any outstanding issues with normalisation, units, scale-factors, phase conventions etc. For details of the AF code, see the method dev notes.
(For MF verification, see the MFPADs and associated $\beta_{LM}$ notebooks, where the numerics are verified for the NO2 test case, although the total cross-sections may still have issues (for more discussion, see the Matlab code release software paper). The geometric tensor version of the MF calculations is also verified against the same test case.)
[1] Cross section and asymmetry parameter calculation for sulfur 1s photoionization of SF6, A. P. P. Natalense and R. R. Lucchese, J. Chem. Phys. 111, 5344 (1999), http://dx.doi.org/10.1063/1.479794
[2] Reid, Katharine L., and Jonathan G. Underwood. “Extracting Molecular Axis Alignment from Photoelectron Angular Distributions.” The Journal of Chemical Physics 112, no. 8 (2000): 3643. https://doi.org/10.1063/1.480517.
[3] Underwood, Jonathan G., and Katharine L. Reid. “Time-Resolved Photoelectron Angular Distributions as a Probe of Intramolecular Dynamics: Connecting the Molecular Frame and the Laboratory Frame.” The Journal of Chemical Physics 113, no. 3 (2000): 1067. https://doi.org/10.1063/1.481918.
[4] Stolow, Albert, and Jonathan G. Underwood. “Time-Resolved Photoelectron Spectroscopy of Non-Adiabatic Dynamics in Polyatomic Molecules.” In Advances in Chemical Physics, edited by Stuart A. Rice, 139:497–584. Advances in Chemical Physics. Hoboken, NJ, USA: John Wiley & Sons, Inc., 2008. https://doi.org/10.1002/9780470259498.ch6.
Formalism: LF case with CG terms
As given in ref. [1].

The matrix elements $I_{\mathbf{k},\hat{n}}^{(L,V)}$ of Eqs. (8)
and (9) can be expanded in terms of the $X_{lh}^{p\mu}$ functions
of Eq. (7) as$^{14}$
\begin{equation}
I_{\mathbf{k},\hat{n}}^{(L,V)}=\left[\frac{4\pi}{3}\right]^{1/2}\sum_{p\mu lhv}I_{lhv}^{p\mu(L,V)}X_{lh}^{p\mu}(\hat{k})X_{1v}^{p_{v}\mu_{v}}(\hat{n}).
\end{equation}
{[}Note here the final term gives polarization (dipole) terms, with
$l=1$, $h=v$, corresponding to a photon with one unit of angular
momentum and projections $v=-1,0,1$, correlated with irreducible
representations $p_{v}\mu_{v}$.{]}
The differential cross section is given by
\begin{equation}
\frac{d\sigma^{L,V}}{d\Omega_{\mathbf{k}}}=\frac{\sigma^{L,V}}{4\pi}[1+\beta_{\mathbf{k}}^{L,V}P_{2}(\cos\theta)],
\end{equation}
where the asymmetry parameter can be written as$^{14}$
\begin{eqnarray}
\beta_{\mathbf{k}}^{L,V} & = & \frac{3}{5}\frac{1}{\sum_{p\mu lhv}|I_{\mathbf{k},\hat{n}}^{(L,V)}|^{2}}\sum_{\stackrel{p\mu lhvmm_{v}}{p'\mu'l'h'v'm'm'{v}}}(-1)^{m'-m{v}}I_{\mathbf{k},\hat{n}}^{(L,V)}\nonumber \
 & \times & \left(I_{\mathbf{k},\hat{n}}^{(L,V)}\right)^{}b_{lhm}^{p\mu}b_{l'h'm'}^{p'\mu'}b_{1vm_{v}}^{p_{v}\mu_{v}}b_{1v'm'{v}}^{p'{v}\mu'{v}*}\nonumber \
 & \times & [(2l+1)(2l'+1)]^{1/2}(1100|20)(l'l00|20)\nonumber \
 & \times & (11-m'{v}m_{v}|2M')(l'l-m'm|2-M'),
\end{eqnarray}
and the $(l'lm'm|L'M')$ are the usual Clebsch--Gordan coefficients.
The total cross section is
\begin{equation}
\sigma^{L,V}=\frac{4\pi^{2}}{3c}E\sum_{p\mu lhv}|I_{\mathbf{k},\hat{n}}^{(L,V)}|^{2},
\end{equation}
where c is the speed of light.

AF formalism
The original (full) form for the AF equations, as implemented in ePSproc.afblm:
\begin{eqnarray}
\beta_{L,-M}^{\mu_{i},\mu_{f}} & = & \sum_{l,m,\mu}\sum_{l',m',\mu'}(-1)^{M}(-1)^{m}(-1)^{(\mu'-\mu_{0})}\left(\frac{(2l+1)(2l'+1)(2L+1)}{4\pi}\right)^{1/2}\left(\begin{array}{ccc}
l & l' & L\
0 & 0 & 0
\end{array}\right)\left(\begin{array}{ccc}
l & l' & L\
-m & m' & -M
\end{array}\right)\nonumber \
 & \times & I_{l,m,\mu}^{p_{i}\mu_{i},p_{f}\mu_{f}}(E)I_{l',m',\mu'}^{p_{i}\mu_{i},p_{f}\mu_{f}*}(E)\
 & \times & \sum_{P,R,R'}(2P+1)(-1)^{(R'-R)}\left(\begin{array}{ccc}
1 & 1 & P\
\mu_{0} & -\mu_{0} & R
\end{array}\right)\left(\begin{array}{ccc}
1 & 1 & P\
\mu & -\mu' & R'
\end{array}\right)\
 & \times & \sum_{K,Q,S}(2K+1)^{1/2}(-1)^{K+Q}\left(\begin{array}{ccc}
P & K & L\
R & -Q & -M
\end{array}\right)\left(\begin{array}{ccc}
P & K & L\
R' & -S & S-R'
\end{array}\right)A_{Q,S}^{K}(t)
\end{eqnarray}
Where $I_{l,m,\mu}^{p_{i}\mu_{i},p_{f}\mu_{f}}(E)$ are the energy-dependent dipole matrix elements, and $A_{Q,S}^{K}(t)$ define the alignment parameters.
In terms of the geometric parameters, this can be rewritten as:
\begin{eqnarray}
\beta_{L,-M}^{\mu_{i},\mu_{f}} & =(-1)^{M} & \sum_{P,R',R}{[P]^{\frac{1}{2}}}{E_{P-R}(\hat{e};\mu_{0})}\sum_{l,m,\mu}\sum_{l',m',\mu'}(-1)^{(\mu'-\mu_{0})}{\Lambda_{R'}(\mu,P,R')B_{L,-M}(l,l',m,m')}I_{l,m,\mu}^{p_{i}\mu_{i},p_{f}\mu_{f}}(E)I_{l',m',\mu'}^{p_{i}\mu_{i},p_{f}\mu_{f}*}(E)\sum_{K,Q,S}\Delta_{L,M}(K,Q,S)A_{Q,S}^{K}(t)\label{eq:BLM-tidy-prod-2}
\end{eqnarray}
See the method dev notebook for more details. Both methods gave the same results for N2 test cases, so are at least consistent, but do not currently match ePS GetCro outputs for the LF case.
Numerics
In both LF and AF cases, the numerics tested herein are based on the geometric tensor expansion code, which has been verified for the MF case as noted above (for PADs at a single energy).
A few additional notes on the implementations...

The matrix elements used are taken from the DumpIdy output segments of the ePS output file, which provide "phase corrected and properly normalized dynamical coefs".
The matrix elements output by ePS are assumed to correspond to $I_{lhv}^{p\mu(L,V)}$ as defined above.
The Scale Factor (SF) "to sqrt Mbarn" output with the matrix elements is assumed to correspond to the $\frac{4\pi^{2}}{3c}E$ term defined above, plus any other required numerical factors ($4\pi$ terms and similar).
The SF is energy dependent, but not continuum (or partial wave) dependent.
If correct, then using matrix elements * scale factor, should give correct results (as a function of $E$), while omitting the scale factor should still give correct PADs at any given $E$, but incorrect total cross-section and energy scaling.
This may be incorrect, and some other assumptions are tested herein.


The AF and LF case should match for an isotropic distribution, defined as $A^{0}_{0,0}=1$. Additional normalisation required here...?
A factor of $\sqrt{(2K+1)}/8\pi^2$ might be required for correct normalisation, although shouldn't matter in this case. (See eqn. 47 in [4].)
For the LF case, as defined above, conversion from Legendre-normalised $\beta$ to spherical harmonic normalised $\beta$ is required for comparison with the AF formalism, where $\beta^{Sph}_{L,0} = \sqrt{(2L+1)/4\pi}\beta^{Lg}$



Set up
End of explanation
"""


# Load data from modPath\data
dataPath = os.path.join(modPath, 'data', 'photoionization')
dataFile = os.path.join(dataPath, 'n2_3sg_0.1-50.1eV_A2.inp.out')  # Set for sample N2 data for testing

# Scan data file
dataSet = ep.readMatEle(fileIn = dataFile)
dataXS = ep.readMatEle(fileIn = dataFile, recordType = 'CrossSection')  # XS info currently not set in NO2 sample file.

"""
Explanation: Test N2 case
Load data
End of explanation
"""


# Plot cross sections using Xarray functionality
dataXS[0].sel({'Type':'L', 'XC':'SIGMA'}).plot.line(x='Eke');

# Plot B2
dataXS[0].sel({'Type':'L', 'XC':'BETA'}).plot.line(x='Eke');

"""
Explanation: Reference results from GetCro
End of explanation
"""


# Set parameters
SFflag = False  # Multiply matrix elements by SF?
symSum = False  # Sum over symmetries?
phaseConvention = 'S'
thres = 1e-2
selDims = {'it':1, 'Type':'L'}
thresDims = 'Eke'

# Set terms for testing - NOTE ORDERING HERE may affect CG term!!!
dlistMatE = ['lp', 'l', 'L', 'mp', 'm', 'M']  # Match published terms
dlistP = ['p1', 'p2', 'L', 'mup', 'mu', 'M']
# dlistMatE = ['l', 'lp', 'L', 'm', 'mp', 'M']  # Standard terms
# dlistP = ['p1', 'p2', 'L', 'mu', 'mup', 'M']


# Set matrix elements
matE = dataSet[0].copy()

# Calculate betas
BetaNormXS, BetaNorm, BetaRaw, XSmatE = lfblmXprod(matE, symSum = symSum, SFflag = SFflag, 
                              thres = thres, thresDims = thresDims, selDims = selDims, 
                              phaseConvention = phaseConvention,
                              dlistMatE = dlistMatE, dlistP = dlistP)

plotThres = None

BetaNorm.SF.real.plot()
ep.util.matEleSelector(BetaNorm.XS, thres = plotThres, dims='Eke', sq=False, drop=False).real.plot.line(x='Eke', col='Sym');
ep.util.matEleSelector(BetaNorm.XS2, thres = plotThres, dims='Eke', sq=False, drop=False).real.plot.line(x='Eke', col='Sym');
ep.util.matEleSelector(BetaNorm, thres = plotThres, dims='Eke', sq=True, drop=True).real.plot.line(x='Eke', col='Sym');
ep.util.matEleSelector(BetaRaw, thres = plotThres, dims='Eke', sq=True, drop=True).real.plot.line(x='Eke', col='Sym');
ep.util.matEleSelector(BetaRaw/BetaRaw.sel({'L':0,'M':0}).drop('LM'), thres = plotThres, dims='Eke', sq=True, drop=True).real.plot.line(x='Eke', col='Sym');
ep.util.matEleSelector(XSmatE, thres = plotThres, dims='Eke', sq=True, drop=True).real.plot.line(x='Eke', col='Sym');
ep.util.matEleSelector(BetaNormXS, thres = plotThres, dims='Eke', sq=True, drop=True).real.plot.line(x='Eke', col='Sym');

"""
Explanation: Test LF calculations (CG version)
End of explanation
"""


BetaNormXS.sel({'Eke':0.1}).unstack('LM').sel({'L':2}).sum('M')

ep.util.matEleSelector(BetaNormXS.unstack('LM').sum('M'), thres = plotThres, dims='Eke', sq=True, drop=True).real.plot.line(x='Eke', col='Sym');

"""
Explanation: 22/06/20: 

BetaNormXS version B2(SU) now correct (with or without SF), but PU case incorrect! 
Phase switch tests (mup, mp) doesn't fix this. Could be something in photon or degeneracy terms? Or in EPR function?
Total XS with or without SF still incorrect as B00, just missing a factor...?
End of explanation
"""


dataXS[0].sel({'Type':'L', 'XC':'BETA'}).plot.line(x='Eke');
ep.util.matEleSelector(BetaNormXS.unstack('LM').sum('M'), thres = plotThres, dims='Eke', sq=True, drop=True).real.plot.line(x='Eke', col='Sym');
ep.util.matEleSelector(BetaNormXS.unstack('LM').sum('M').sum('Sym'), thres = plotThres, dims='Eke', sq=True, drop=True).real.plot.line(x='Eke');

"""
Explanation: SUM OVER M IS KEY HERE - M' in formalism is NOT same as BLM!!!
Next question is why... this is just MF degeneracy term?
In 1986 paper, this is given label $\mu''$, which seems more appropriate (eqn. 4.13), and there are some different phase choices.
End of explanation
"""


# Test for XS...
# ep.util.matEleSelector((BetaNorm*BetaNorm.SF).unstack('LM').sum('M'), thres = plotThres, dims='Eke', sq=True, drop=True).real.plot.line(x='Eke', col='Sym');
ep.util.matEleSelector((BetaNorm.unstack('LM').sel({'L':0})), thres = plotThres, dims='Eke', sq=True, drop=True).real.plot.line(x='Eke', col='Sym');

dataXS[0].sel({'Type':'L', 'XC':'BETA'})

ep.util.matEleSelector(BetaNormXS.unstack('LM').sum('M'), thres = plotThres, dims='Eke', sq=True, drop=True)

"""
Explanation: TO DO

23/06/20 Betas per continuum now OK above. Test also sym sym case.
Get to bottom of XS scaling.
Check AF formalism - may be missing a term linking MF to LF degen?  3j in mu/lambda...? Check derivations!!!
End of explanation
"""






# Set matrix elements
matE = dataSet[0].copy()

# Setting this here gives usual problem (non-matching XS values to getCro)
# Now testing multiplication at end of calc...
# if SFflag:
#     matE.values = matE * matE.SF

matEthres = ep.matEleSelector(matE, thres = thres, inds = selDims, dims = thresDims, sq = True, drop = True)


# Calculate 3j sets

#*** For (l,lp) terms, use code from betaTerm(), but CG version
# Set QNs
QNs = geomUtils.genllpMatE(matEthres, phaseConvention = phaseConvention)
# Set phase convention for term (lp,l,-mp,m|2,-M) - NOTE also ordering of lp,l etc.!
dlistMatE = ['lp', 'l', 'L', 'mp', 'm', 'M']
# QNs[:,3] *= -1  # mp > -mp  # Including this restricts things to mp=0 only? Phase con choice? Doesn't seem correct!
                # Full calculation including this term sends PU continuum to zero/NaN.
QNs[:,5] *= -1  # M > -M

# Calculate main term + term with m=mp=M=0
CGmatE = CG(QNs=QNs, dlist = dlistMatE)
CGmatEM0 = CG(QNs = geomUtils.genllLList(QNs, uniqueFlag = True, mFlag = False), 
              dlist = dlistMatE, form='xdaLM')

# Multiply terms (as per betaTerm() code)
CGmatEmult = CGmatE.unstack()*CGmatEM0.drop('mSet').squeeze().unstack()


#*** For photon terms
# See EPR() and MFproj() for existing 3j prototypes.
dlistP = ['p1', 'p2', 'L', 'mup', 'mu', 'M']  # Note p1=p2=1 and will be dropped
QNsP=geomUtils.genllL(Lmin=1,Lmax=1)  # Generates all possible l1=l2=1 terms, no phase convention applied.

QNsP[:,3] *= -1  # mup > -mup
CGP = CG(QNs=QNsP, dlist = dlistP, form='xdaLM')
CGPM0 = CG(QNs = geomUtils.genllLList(QNsP, uniqueFlag = True, mFlag = False), 
              dlist = dlistP, form='xdaLM')

# Multiply terms (as per betaTerm() code)
CGPmult = CGP.unstack()*CGPM0.drop('mSet').squeeze().unstack()

# Define matrix element multiplication term (code from mfblmGeom())
matEconj = matEthres.copy().conj()
# matEconj = matEconj.unstack().rename({'l':'lp','m':'mp','mu':'mup'})  # Full unstack
# matEmult = matEthres.unstack() * matEconj
matEconj = matEconj.unstack('LM').rename({'l':'lp','m':'mp','mu':'mup'})  # Unstack LM only.
matEmult = matEthres.unstack('LM') * matEconj
matEmult.attrs['dataType'] = 'multTest'

# Threshold product and drop dims.
# matEmult = ep.util.matEleSelector(matEmult, thres = thres, dims = thresDims)
matEmult = matEleSelector(matEmult, thres = thres, dims = thresDims)

matEmult.dims

# Multiply matE * photon terms
CGPM = CGmatEmult * CGPmult.squeeze(['p1','p2']).drop(['p1','p2'])

# Set phase = (-1)^(m'-mu)
CGPM *= np.power(-1, np.abs(CGPM.mp - CGPM.mu))

# Big mult
sumDims = ['mu', 'mup', 'l','lp','m','mp']
xDim = {'LM':['L','M']}

# Set any other terms...
degenllp = np.sqrt((2*matEmult.l+1)*(2*matEmult.lp+1))

Betas = (CGPM * matEmult * degenllp).sum(sumDims).stack(xDim)

# Renorm
if SFflag:
#     normXS = np.abs(Betas.SF)
    normXS = np.abs(Betas.SF)**2
else:
    normXS = 1

# normBLM = np.sqrt(4*np.pi/3) # 3/5  # Terms from 1999 SF6 paper
# normBLM = (4*np.pi**2)/5  # Term from 1985 long paper



BetasNorm = Betas.copy() * normXS
BetasNorm['XS'] = BetasNorm.sel({'L':0,'M':0}).drop('LM').copy()  # This basically works, and keeps all non-summed dims... but may give issues later...? Make sure to .copy(), otherwise it's just a pointer.
BetasNorm /= BetasNorm.sel({'L':0,'M':0}).drop('LM')
BetasNorm *= normBLM


Betas.coords

# Betas.squeeze().plot.line(x='Eke',col='Sym')
# BetasNorm.XS.real.plot.line(x='Eke', col='Sym');
BetasNorm.SF.real.plot()
ep.util.matEleSelector(BetasNorm.XS, thres = 0.001, dims='Eke', sq=False, drop=False).real.plot.line(x='Eke', col='Sym');
ep.util.matEleSelector(BetasNorm, thres = 0.1, dims='Eke', sq=False, drop=False).real.plot.line(x='Eke', col='Sym');


np.abs((matE*matE.SF).sel({'Eke':0.1}))

# Try alternative XS definition... looking at 1999 SF6 paper formalism
# XS = (matE.SF * (matE)**2).sum(['LM','mu','it'])  # This has real and imag parts
# XS = (matE.SF * (matE*matE.conj())).sum(['LM','mu','it'])  # This has real and imag parts, peaks at ~50 in SU channel.
# XS = ((matE.SF * matE) * (matE.SF * matE).conj()).sum(['LM','mu','it'])  # Re only, peaks at ~300 in SU channel.
# XS = (np.abs(matE.SF * matE)).sum(['LM','mu','it'])  # Re only, POOR SHAPE. peaks at ~24 in SU channel.
# XS = ((matE*matE.conj())).sum(['LM','mu','it'])  # Real only, peak ~18. Shape looks good.
# XS = (matE.SF * (np.abs(matE))).sum(['LM','mu','it']) # Poor shape
# XS = (matE.SF * (np.abs(matE)**2)).sum(['LM','mu','it']) # Pretty similar to previous attempts, peaks ~50

# Testing additional terms/normalisation
# XS = ((matE*matE.conj())).sum(['LM','mu','it'])/(4*np.pi)  # Real only, peak ~1.5. Shape looks good.
# XS = ((matE*matE.conj())).sum(['LM','mu','it'])/(4*np.pi)*((4*np.pi**2)/3)  # Real only, peak ~18. Shape looks good.
# XS = ((matE*matE.conj())).sum(['LM','mu','it'])*((4*np.pi**2)/3) # Re only, peaks ~ 240
# XS = ((matE.SF * matE) * (matE.SF * matE).conj()).sum(['LM','mu','it'])*((4*np.pi**2)/3)  # Re only, peaks at ~3500 in SU channel.
XS = ((matE.SF * matE) * (matE.SF * matE).conj()).sum(['LM','mu','it'])*((4*np.pi)/3)

#** XS = (matE.SF * np.abs(matE)).sum(['LM','mu','it'])  # UGH
# XS = (matE.SF * np.abs(matE**2)).sum(['LM','mu','it'])  # Re+Im, Peaks at ~50
# XS = (np.abs(matE.SF) * np.abs(matE**2)).sum(['LM','mu','it'])  # Re only, Peaks at ~70

# Investigating E/c term - so far assume this is included in SF... but maybe not? Not sure on units.
# XS = (matE.SF * (matE*matE.conj())).sum(['LM','mu','it']) * matE.Eke  # shape looks poor, peaks at ~500
# XS = ((matE*matE.conj())).sum(['LM','mu','it']) * matE.Eke  # shape looks poor, peaks at ~250

XS.real.plot.line(x='Eke',col='Type')
XS.imag.plot.line(x='Eke',col='Type')
XS.pipe(np.abs).plot.line(x='Eke',col='Type')
# XS

# Check terms
# Set parameters to restack the Xarray into (L,M) pairs
xDim = {'LM':['L','M']}

# Plot with ep.lmPlot(), real values
daPlot, daPlotpd, legendList, gFig = ep.lmPlot(CGmatEmult, xDim=xDim, pType = 'r')
daPlot, daPlotpd, legendList, gFig = ep.lmPlot(CGPmult, xDim=xDim, pType = 'r')
# daPlot, daPlotpd, legendList, gFig = ep.lmPlot(BetasNorm, xDim=xDim, pType = 'r')
daPlotpd

"""
Explanation: SORTING BELOW
End of explanation
"""


