# To control display of warnings...
# WORKING IN SOME CASES... but not all, maybe ipython/env version issues?
# https://stackoverflow.com/a/9031848
import warnings
warnings.filterwarnings('ignore')  # set 'once' or 'ignore'
warnings.simplefilter('ignore')

# 16/08/21 ADDED TO util.env.isnotebook(), not yet tested.

# For ipython and other deprecation warnings, https://stackoverflow.com/questions/63413807/deprecation-warning-from-jupyter-should-run-async-will-not-call-transform-c
# warnings.filterwarnings("ignore", category=DeprecationWarning)
# warnings.filterwarnings("ignore", category=FutureWarning)



# Standard libs
import sys
import os
from pathlib import Path
import numpy as np
import xarray as xr

from datetime import datetime as dt
timeString = dt.now()



# For local module testing, include path to module here, otherwise use global installation
# For editable 
local = False

if local:
    import sys
    if sys.platform == "win32":
        modPath = r'D:\code\github\ePSproc'  # Win test machine
        winFlag = True
    else:
        modPath = r'~/github/ePSproc/'  # Linux test machine
        winFlag = False

    sys.path.append(modPath)

# Import base package
import epsproc as ep  

# Set data path
# Note this is set here from ep.__path__, but may not be correct in all cases - depends on where the Github repo is.
epDemoDataPath = Path(ep.__path__[0]).parent/'data'



"""
Explanation: Method dev: metrics
16/08/21
Using geometric tensor formalism to compose & calculate some system/problem metrics.
To consider:

Aligment metrics
$cos^n(\theta)$

$cos^{3D}$


Information content

Channel functions (from alignment?)
Information/sensitivity (alignment x matE functions)
Continuum density matrix

Setup
End of explanation
"""


# Class dev code
from epsproc.classes.multiJob import ePSmultiJob
from epsproc.classes.base import ePSbase

# Instantiate class object.
# Minimally this needs just the dataPath, if verbose = 1 is set then some useful output will also be printed.
dataPath = os.path.join(epDemoDataPath, 'photoionization')
data = ePSbase(dataPath, verbose = 1)

# ScanFiles() - this will look for data files on the path provided, and read from them.
data.scanFiles()

"""
Explanation: Load test data
Class version
End of explanation
"""


# Set ADMs for increasing alignment...
tPoints = 10
# inputADMs = np.asarray([[0,0,0, *np.ones(10)], [2,0,0, *np.linspace(0,1,tPoints)], [4,0,0, *np.linspace(0,0.5,tPoints)]])
inputADMs = [[0,0,0, *np.ones(10)], [2,0,0, *np.linspace(0,1,tPoints)], [4,0,0, *np.linspace(0,0.8,tPoints)], [6,0,0, *np.linspace(0,0.5,tPoints)]]

ADMs = ep.setADMs(ADMs = inputADMs)
ep.lmPlot(ADMs, xDim='t');

# data.setADMs(ADMs = ADMs['ADM'], t=ADMs['time'].squeeze(), KQSLabels = ADMs['ADMlist'], addS = True)  
# data.setADMs(ADMs = ADMs)   # setADMs not yet in class, but see PEMtk fitClass.py for method example

# Set to separate dataset for testing
ADMs.attrs['jobLabel'] = 'Alignment data testing'
data.data['ADM'] = {'ADM': ADMs}

data.ADMplot(keys = 'ADM')   #padPlot(keys = 'orb11', dataType='AFBLM', selDims = {'Labels':'z'}, pStyle='grid', reducePhi='sum', facetDims=['t', 'Eke'])

# Axis distributions P(theta,phi) vs. t
data.padPlot(keys = 'ADM', dataType='ADM', Etype = 't', pStyle='grid', facetDims=['t'])

# Axis distributions P(theta) vs. t (summed over phi)
Pt = data.padPlot(keys = 'ADM', dataType='ADM', Etype = 't', pStyle='grid', reducePhi='sum', returnFlag = True)  #, facetDims=['t', 'Eke'])

# The returned data array contains the plotted values
Pt

# Check norm
Pt.sum('Theta').plot()

# Renorm by pop
Pt = Pt/Pt.sum('Theta')

# Check norm
Pt.sum('Theta').plot()

# Check renorm version
Pt.plot()

"""
Explanation: Alignment metrics
Here we'll start with the molecular alignment defined in terms of axis distribution moments (ADMs), given by parameters $A^K_{Q,S}(t)$:
$ P(\Omega,t) = \sum_{K,Q,S} A^K_{Q,S}(t)D^K_{Q,S}(\Omega)$
Where $P(\Omega,t)$ is the full axis distribution probability, expanded for Euler angles $\Omega$.
This reduces to the 2D case if $S=0$:
$ P(\theta,\phi,t) = \sum_{K,Q} A^K_{Q,0}(t)D^K_{Q,0}(\Omega) = \sum_{K,Q} A^K_{Q}(t)Y_{K,Q}(\Omega)$
In the current code (v1.3.0, 16/08/21), the ADMs can be set directly, and expanded trivially for the 2D case (3D case to do). (Note that the 3D case is supported for AF-BLM calculations, but not yet for direct expansion & visualisation.)
End of explanation
"""


# Generate basis fn.

import xarray as xr

# cn = xr.zeros_like(Pt)  # Set to same form as Pt
# # cn = np.cos(cn.Theta)**2  # Set values - only keeps Theta dim in this case
# cn = (np.cos(cn.Theta)**2).expand_dims('t')
# # cn = np.cos(cn)**2  # Set values
# # cn.apply_ufunc(np.cos)

N = 2
cn = xr.ones_like(Pt) * (np.cos(Pt.Theta)**N)  # Use multiplication to keep t coord

cn.plot()
# cn

Jt = np.sin(Pt.Theta) # Jacobian
cnMetric = (Pt * cn * Jt).sum('Theta')  # Include Jacobian?
# cnMetric = 2*(Pt * cn).sum('Theta')  # x2 for 0:2pi range correction?

# norm = Pt.Theta.size/(cn.sum('Theta')+Pt.sum('Theta')) # Renorm by # elements/sum
norm = (Pt * Jt).sum('Theta') # From old matlab code, should be correct renorm: sum(P.*cos(q).^2.*sin(q))/sum(P.*sin(q)) 
                            # Should also have sum(cos(q)) on the bottom? OK if already normalise?
                            # Actually, should just renorm Pt correctly above, i.e. by sum(Pt.*sin(q))! (Currently by only sum(Pt))
(cnMetric/norm).plot()

# Equivalent result with dot product function

# Pt.dot(cn, dims='Theta').plot()  #.sum('Theta')  # No normalization - gives non-linear result, but looks OK, although min value = 0.53 (should be 1/3)

(Pt.dot(cn * Jt, dims='Theta')/Pt.dot(Jt, dims='Theta')).plot()

cn.sum('Theta')

"""
Explanation: $\langle cos^n(\theta, t)\rangle$
Expectation values of $\langle cos^n(\theta, t)\rangle$ can be calculated directly from the $P(\theta,t)$ distributions:
$\langle cos^n(\theta, t)\rangle = \int_{\theta} cos^n(\theta)P(\theta,t) d\theta$
Refs:
<div class="csl-bib-body" style="line-height: 1.35; ">
  <div class="csl-entry" style="clear: left; ">
    <div class="csl-left-margin" style="float: left; padding-right: 0.5em;text-align: right; width: 1em;">[1]</div><div class="csl-right-inline" style="margin: 0 .4em 0 1.5em;">P. Hockett, “General phenomenology of ionization from aligned molecular ensembles,” <i>New Journal of Physics</i>, vol. 17, no. 2, p. 023069, Feb. 2015, doi: <a href="https://doi.org/10.1088/1367-2630/17/2/023069">10.1088/1367-2630/17/2/023069</a>.</div>
  </div>
  <span class="Z3988" title="url_ver=Z39.88-2004&amp;ctx_ver=Z39.88-2004&amp;rfr_id=info%3Asid%2Fzotero.org%3A2&amp;rft_id=info%3Adoi%2F10.1088%2F1367-2630%2F17%2F2%2F023069&amp;rft_val_fmt=info%3Aofi%2Ffmt%3Akev%3Amtx%3Ajournal&amp;rft.genre=article&amp;rft.atitle=General%20phenomenology%20of%20ionization%20from%20aligned%20molecular%20ensembles&amp;rft.jtitle=New%20Journal%20of%20Physics&amp;rft.volume=17&amp;rft.issue=2&amp;rft.aufirst=Paul&amp;rft.aulast=Hockett&amp;rft.au=Paul%20Hockett&amp;rft.date=2015-02&amp;rft.pages=023069&amp;rft.issn=1367-2630"></span>
</div>
End of explanation
"""


def computeCNs(Pt, N = None, Nrange = [2,10], Nstep = 2):
    """
    Function to compute <cos^N> expectation values from P(theta) distributions.
    
    NOTE: currently assumes 1D case. Should generalise!
    
    """
    
    if N is None:
        N = np.arange(Nrange[0], Nrange[1], Nstep)
    
    # Generate cos^N basis
    #  cn = xr.ones_like(Pt) * (np.cos(Pt.Theta)**N)  # Use multiplication to keep t coord - this fails for array-like N with broadcast errors
    cn = xr.ones_like(Pt) * np.cos(Pt.Theta).expand_dims({'N':N}).T.pipe(np.power,N)  # OK with transpose
    
    Jt = np.sin(Pt.Theta) # Jacobian
    
    norm = (Pt * Jt).sum('Theta')
    cnMetric = (Pt * cn * Jt).sum('Theta')/norm
    
#     cnMetric = Pt.dot(cn, dims='Theta') # Multiply and sum over Theta - dot version (no Jacobian or renorm)
#     cnMetric = (Pt.dot(cn * Jt, dims='Theta')/Pt.dot(Jt, dims='Theta'))  # Dot version with renorm
    
    return cnMetric

# Compute and plot for a range of N
cnMetric = computeCNs(Pt)

# cnMetric.plot()  # Basic plot
ep.lmPlot(cnMetric, xDim='t');  # With lmPlot()

cnMetric.min()

cnMetric.sel(N=2)  # = 0.33 for isotropic case OK

"""
Explanation: To check: missing additional 1/2 or other norm factor here? Should have 0.33 for isotropic case. OR need to renorm Pt and cn by area.
Functionalise and extend to higher orders...
End of explanation
"""


N = np.arange(0,10,2)
xr.ones_like(Pt) * (np.cos(Pt.Theta)**N)

# np.power(np.cos(Pt.Theta),N)
# temp = np.cos(Pt.Theta).expand_dims({'N':N})  #.pipe(np.power,N)
# temp = np.cos(Pt.Theta).pipe(np.power,N)
temp = np.cos(Pt.Theta).expand_dims({'N':N}).T.pipe(np.power,N)  # OK with transpose
temp

 N = np.arange(0, 10, 2)  # , ndmin=2)  # Can't set ndmin here?
np.power(temp.T,N)  # AH, works with transpose.

# Above should work...? Not sure why broadcasting is broken!
# Example from https://numpy.org/doc/stable/reference/generated/numpy.power.html

x1 = np.arange(6)
x2 = np.array([[1, 2, 3, 3, 2, 1], [1, 2, 3, 3, 2, 1]])

x1**x2

# Try flipping - also OK, UNLESS x1.T is used!
x2 = np.arange(6)
x1 = np.array([[1, 2, 3, 3, 2, 1], [1, 2, 3, 3, 2, 1]])

x1**x2  # OK
# (x1.T)**x2  # Dim broadcast error in this case.

x1.shape

x2.shape

xr.ones_like(Pt) * np.cos(Pt.Theta).expand_dims({'N':N}).T.pipe(np.power,N)

"""
Explanation: CN function dev scratch
End of explanation
"""


import scooby
scooby.Report(additional=['epsproc', 'holoviews', 'hvplot', 'xarray', 'matplotlib', 'bokeh'])

# Check current Git commit for local ePSproc version
from pathlib import Path
!git -C {Path(ep.__file__).parent} branch
!git -C {Path(ep.__file__).parent} log --format="%H" -n 1

# Check current remote commits
!git ls-remote --heads git://github.com/phockett/ePSproc

"""
Explanation: Versions
End of explanation
"""


