# coding: utf8

import numpy as np
from ase.build import bulk

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


# Create the cluster
a = 5.43
Si = bulk('Si', a=a, cubic=True)
cluster = hemispherical_cluster(Si,
                                diameter=30, planes=4,
                                emitter_plane=3,
                                shape = 'cylindrical',
                                )
for atom in cluster:
    atom.set('mean_square_vibration', 0.006)
    atom.set('mt_radius', 1.1)
cluster.emitter = get_atom_index(cluster, 0, 0, 0)

# Create a calculator and set parameters
calc = MSSPEC(spectroscopy='PED', algorithm='expansion')

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.tmatrix_parameters.tl_threshold = 1e-4
calc.calculation_parameters.vibrational_damping = 'averaged_tl'
calc.calculation_parameters.RA_cutoff = 2

my_filters = ('forward_scattering', 'distance_cutoff')
calc.calculation_parameters.path_filtering = my_filters
calc.calculation_parameters.distance = 30
calc.calculation_parameters.off_cone_events = 0
[a.set('forward_angle', 40) for a in cluster]

calc.set_atoms(cluster)

# Compute and add previous data for comparison
data = None
for i in range(1,4):
    calc.calculation_parameters.scattering_order = i
    data = calc.get_theta_scan(level='2p',
                                theta=np.arange(-30., 80., 0.5),
                                phi=0,
                                kinetic_energy=1382.28, data=data)

# Put all orders on the same view
dset = data[0]
dset.add_columns(cs2=data[1].cross_section)
dset.add_columns(cs3=data[2].cross_section)
view1 = dset.views[0]
view1.select('theta', 'cross_section', index=0, legend="Scattering order 1")
view1.select('theta', 'cs2', legend="Scattering order 2")
view1.select('theta', 'cs3', legend="Scattering order 3")


no_filters = Data.load('path_filtering.hdf5')
dset.add_columns(no_filters=no_filters[0].cross_section)
view2 = dset.add_view('comparison', **view1._plotopts)
view2._plotopts['legend'] = []
view2.select('theta', 'cs3', legend="With path filtering")
view2.select('theta', 'no_filters', legend="Without path filtering")

data.view()