# coding: utf8

import logging
import numpy as np
import os
import sys

from ase.build import bulk, diamond111

from msspec.calculator import MSSPEC, XRaySource
from msspec.iodata import Data
from msspec.utils import hemispherical_cluster, get_atom_index, cut_sphere, cut_plane

logging.basicConfig(level=logging.INFO)

def create_clusters(nplanes=6, direction='100', a=5.43, c=5.43, radius=30):
    clusters = []
    Si = bulk('Si', a=a, cubic=True)

    # Get scaled positions and cell
    p = Si.get_scaled_positions()
    cell = Si.get_cell()

    # Resize
    covera = c / a
    cell[2,:] *= covera
    p[:,2] *= covera
    Si.set_scaled_positions(p)
    Si.set_cell(cell)

    if direction in ('001', '010', '100'):
        pass
    elif direction in ('111'):
        Si.rotate([1,1,1], [1,1,0], rotate_cell=True)
        Si.rotate([1,1,0], [1,0,0], rotate_cell=True)
        Si.rotate([1,0,0], [0,0,1], rotate_cell=True)

    for iplane in range(nplanes):
    #for iplane in (nplanes-1,):
        logging.info(f'Building cluster #{iplane:d}/{nplanes-1:d}')
        cluster = hemispherical_cluster(Si, #planes=max(iplane+1, 4),
                                        diameter=70,
                                        emitter_plane=iplane,
                                        shape = 'cylindrical',
                                       )

        cluster = cut_sphere(cluster, radius = radius)
        cluster = cut_plane(cluster, z = -a/4)

        for atom in cluster:
            atom.set('mean_square_vibration', 0.006)
            atom.set('mt_radius', 1.1)
        cluster.absorber = get_atom_index(cluster, 0, 0, 0)
        cluster.info.update(emitter_plane=iplane)
        clusters.append(cluster)
    return clusters

def compute_polar_scan(cluster, data=None, phi=0, folder='calc', RA_cutoff=2,
                       max_ndif=6, level='2p', kinetic_energy=1382.28):
    emitter_plane = cluster.info['emitter_plane']

    calc = MSSPEC(spectroscopy='PED', algorithm='expansion',
                  folder=folder)

    calc.muffintin_parameters.interstitial_potential = 0.

    #calc.tmatrix_parameters.exchange_correlation = 'x_alpha_real'
    calc.tmatrix_parameters.exchange_correlation = 'hedin_lundqvist_complex'

    calc.source_parameters.energy = XRaySource.AL_KALPHA
    calc.source_parameters.theta  = -54.7
    calc.source_parameters.phi    = 90
    calc.spectroscopy_parameters.final_state = 1
    calc.detector_parameters.angular_acceptance = 1.0
    calc.detector_parameters.average_sampling   = 'high'

    calc.calculation_parameters.scattering_order = min(max_ndif,
                                                       max(1, emitter_plane))
    #calc.tmatrix_parameters.lmax_mode ='imposed'
    #calc.tmatrix_parameters.lmaxt = 10
    calc.tmatrix_parameters.tl_threshold = 1e-4
    #calc.calculation_parameters.scattering_order = 6
    #calc.calculation_parameters.mean_free_path = 10.
    calc.calculation_parameters.vibrational_damping = 'averaged_tl'
    calc.calculation_parameters.RA_cutoff = RA_cutoff
    my_filters = ('forward_scattering', 'distance_cutoff')
    calc.calculation_parameters.path_filtering = my_filters
    #calc.calculation_parameters.path_filtering   = 'forward_scattering'
    calc.calculation_parameters.distance = 30
    calc.calculation_parameters.off_cone_events   = 0
    calc.calculation_parameters.RA_cutoff_damping = 1
    [a.set('forward_angle', 40) for a in cluster]

    calc.phagen_parameters.noproto = '.false.'

    calc.set_atoms(cluster)

    data = calc.get_theta_scan(level=level,
                               theta=np.arange(-30., 80., 0.5),
                               kinetic_energy = kinetic_energy,
                               phi=phi, data=data)

    return data

def sum_planes(data):
    cs = None
    ds = None
    for dset in data:
        theta = dset.theta
        phi = dset.phi
        try:
            cs += dset.cross_section
            ds += dset.direct_signal
        except:
            cs = dset.cross_section.copy()
            ds = dset.direct_signal.copy()

    dset = data.add_dset('Substrate signal')
    dset.add_columns(theta=theta, phi=phi, cross_section=cs,
                     direct_signal=ds, chi=(cs-ds)/ds)

    phi_values = np.unique(phi)

    view = dset.add_view('Substrate signal',
                         title=r'Si(2p)',
                         xlabel=r'$\Theta (\degree$)',
                         ylabel='Signal (a.u.)')

    for phi_value in phi_values:
        condition = f'phi == {phi_value:.1f}'
        legend = '$\Phi = $' + f'{phi_value:.1f}' + '$^\circ$'
        view.select('theta', 'cross_section', where=condition, legend=legend)
    view.set_plot_options(autoscale=True)


# Create the list of clusters
RA   = 2   #int(sys.argv[1])
NDIF = 5   #int(sys.argv[2])
RAD  = 20. #float(sys.argv[3])
PHI  = float(sys.argv[1])
KE   = float(sys.argv[2])
#basename = f'RA{RA:d}_NDIF{NDIF:d}_RAD{RAD:.1f}_PHI{PHI:.1f}'
basename = f'PHI{PHI:.1f}_KE{KE:.2f}'
data = Data(basename)
clusters = create_clusters(nplanes=15, radius=RAD)

if 'view' in sys.argv:
    for cluster in clusters:
        cluster.edit()
    exit(0)

for cluster in clusters:
#for cluster in (clusters[-1],):
    data = compute_polar_scan(cluster, data, phi=PHI,
                              folder='calc_' + basename,
                              RA_cutoff=RA,
                              max_ndif=NDIF,
                              kinetic_energy=KE
                              )

sum_planes(data)
data.save('simulation_' + basename + '.hdf5', append=False)
data.export('simulation_' + basename)
#data.view()