# coding: utf-8

from msspec.calculator import MSSPEC, XRaySource
from msspec.utils import *
from msspec.iodata import Data
from ase.build  import bulk

import numpy as np
import sys

a0        = 3.6                          # The lattice parameter in angstroms
phi       = np.arange(0, 100)            # An array of phi angles
all_T     = np.arange(300, 1000, 100)    # All temperatures used for the calculation
all_theta = np.array([45, 83])           # 2 polar angles, 83° is grazing detection
eps       = 0.01                         # a useful small value

DATA_FNAME = 'all_data.hdf5'             # Where to store all the data
ANALYSIS_FNAME = 'analysis.hdf5'


def compute(filename):
    """Will compute a phi scan for an emitter in the first, second and third plane
    of a copper substrate for various polar angles and temperatures.
    In a second step (outside this function), intensities from all these emitters
    are added to get  the signal of a substrate."""
    calc = MSSPEC(spectroscopy='PED')
    calc.source_parameters.energy = XRaySource.AL_KALPHA
    calc.muffintin_parameters.interstitial_potential = 14.1

    calc.calculation_parameters.vibrational_damping = 'averaged_tl'
    calc.calculation_parameters.use_debye_model = True
    calc.calculation_parameters.debye_temperature = 343
    calc.calculation_parameters.vibration_scaling = 1.2

    calc.detector_parameters.average_sampling = 'high'
    calc.detector_parameters.angular_acceptance = 5.7

    filters = ['forward_scattering', 'backward_scattering']
    calc.calculation_parameters.path_filtering = filters

    calc.calculation_parameters.RA_cutoff = 3
    
    # You can also choose a real potential and manually define the
    # electron mean free path
    #
    # calc.tmatrix_parameters.exchange_correlation = 'hedin_lundqvist_real'
    # calc.calculation_parameters.mean_free_path = 10.

    data = None # init an empty data object
    for temperature in all_T:
        for plane in range(3):
	    # We create a cylindrical cluster here with one plane below the emitter
            # and 0, 1 or to planes above the emitter
            cluster = bulk('Cu', a=a0, cubic=True)
            cluster = cluster.repeat((20, 20, 8))
            center_cluster(cluster)
            cluster = cut_cylinder(cluster, radius=1.5 * a0 + eps)
            cluster = cut_plane(cluster, z=-( a0/2 + eps))
            cluster = cut_plane(cluster, z=(plane) * a0 / 2 + eps)
            cluster.absorber = get_atom_index(cluster, 0, 0, 0)

            calc.calculation_parameters.temperature = temperature
	    # the scattering order depends on the number of planes above
            # the emitter to speed up calculations
            calc.calculation_parameters.scattering_order = 1 + plane

            calc.set_atoms(cluster)
            # Here starts the calculation
            data = calc.get_phi_scan(level='2p3/2', theta=all_theta, phi=phi,
                                     kinetic_energy=560, data=data)

    data.save(filename)

def process_data(datafile, outputfile):
    """Will create another Data file with some phi-scans corresponding to 3
       planes at different temperatures and the anisotropy curve for 45° and
       grazing detection.
    """
    def get_signal(datafile, T=300, theta=45):
        data = Data.load(datafile)
        total = None
        for dset in data:
            p = {_['group'] + '.' + _['name']: _['value'] for _ in dset.parameters()}
            temperature = np.float(p['CalculationParameters.temperature'])
            if temperature != T: continue
            i = np.where(dset.plane == -1)
            j = np.where(dset.theta[i] == theta)
            signal = dset.cross_section[i][j]
            try:
                total += signal
            except:
                total = signal
            phi = dset.phi[i][j]
        return phi, total

    analysis = Data('Temperature tutorial')
    scans_dset = analysis.add_dset('Phi scans')
    scans_view = scans_dset.add_view('Figure 1',
                                     title=r'Cu(2p) Azimuthal scan for $\theta = 83\degree$',
                                     xlabel=r'$\Phi (\degree$)',
                                     ylabel='Signal (a.u.)')
    anisotropy_dset = analysis.add_dset('Anisotropies')
    anisotropy_view = anisotropy_dset.add_view('Figure 2',
                title='Relative anisotropies for Cu(2p)',
                marker='o',
                xlabel='T (K)',
                ylabel=r'$\frac{\Delta I / I_{max}(T)}{\Delta I_{300} / I_{max}(300)} (\%)$')
    for theta in all_theta:
        for T in all_T:
            PHI, SIGNAL = get_signal(datafile, T=T, theta=theta)
            for phi, signal in zip(PHI, SIGNAL):
                scans_dset.add_row(phi=phi, signal=signal, theta=theta, temperature=T)

            middleT = all_T[np.abs(all_T - np.mean(all_T)).argmin()]
            if theta == 83 and T in [np.min(all_T), middleT, np.max(all_T)]:
                scans_view.select('phi', 'signal',
                              where='temperature == {:f} and theta == {:f}'.format(T, theta),
                              legend='{:.0f} K'.format(T))

            PHI, SIGNAL0 = get_signal(datafile, T=np.min(all_T), theta=theta)
            Imax = np.max(SIGNAL)
            Imax0 = np.max(SIGNAL0)
            dI = Imax - np.min(SIGNAL)
            dI0 = Imax0 - np.min(SIGNAL0)
            ani = (dI / Imax) / (dI0 / Imax0)
            anisotropy_dset.add_row(temperature=T, anisotropy=ani, theta=theta)

        anisotropy_view.select('temperature', 'anisotropy',
                               where='theta == {:f}'.format(theta),
                               legend=r'$\theta = {:.0f} \degree$'.format(theta))

    analysis.save(outputfile)


# A convenient way to run the script, just specify one or more of these calc, analyse or
# view keywords as arguments
# ... to calculate all the data, analyse the data and view the results, run
# python  Cu_temperature.py calc analyse view
# ... to just view the results, simply run
# python  Cu_temperature.py view

if 'calc' in sys.argv:
    compute(DATA_FNAME)
if 'analyse' in sys.argv:
    process_data(DATA_FNAME, ANALYSIS_FNAME)
if 'view' in sys.argv:
    data = Data.load(ANALYSIS_FNAME)
    data.view()