# coding: utf8

import numpy as np
from ase.build import bulk
from ase.lattice.tetragonal import SimpleTetragonalFactory
from msspec.calculator import MSSPEC, XRaySource
from msspec.utils import hemispherical_cluster, get_atom_index
import logging

logging.basicConfig(level=logging.INFO)


do_ped = False

# Define a Rocksalt Factory class (to tetragonalize the unit cell)
class RocksaltFactory(SimpleTetragonalFactory):
    bravais_basis = [[0, 0, 0], [0.5, 0.5, 0], [0.5, 0, 0.5], [0, 0.5, 0.5], 
                     [0, 0, 0.5], [0.5, 0, 0], [0, 0.5, 0], [0.5, 0.5, 0.5]]
    element_basis = (0, 0, 0, 0, 1, 1, 1, 1)
Rocksalt = RocksaltFactory()


a0 = 4.09
a_perp = 4.25
MgO = Rocksalt(('Mg', 'O'), 
               latticeconstant={'a': a0, 'c/a': a_perp/a0}, 
               size=(1,1,1))


for atom in MgO:
   atom.set('mean_square_vibration', 0.01)
   atom.set('forward_angle', 20.)
   if atom.symbol == 'Mg':
      atom.tag = 1
      atom.set('mt_radius', 0.63)
   else:
      atom.tag = 2
      atom.set('mt_radius', 1.415)


#cluster = hemispherical_cluster(MgO, emitter_tag=1, emitter_plane=1, planes=4, diameter=4.5*a0)
cluster = hemispherical_cluster(MgO, emitter_tag=1, emitter_plane=1, planes=2)
cluster.absorber = get_atom_index(cluster, 0, 0, 0)
#cluster.edit()
#exit()

if do_ped:
    calc = MSSPEC(spectroscopy='PED', algorithm='inversion')
else:
    calc = MSSPEC(spectroscopy='AED', algorithm='inversion')
calc.set_atoms(cluster)

calc.muffintin_parameters.ionicity = {'Mg': 0.1, 'O': -0.1}

calc.tmatrix_parameters.exchange_correlation = 'hedin_lundqvist_complex'
calc.tmatrix_parameters.lmax_mode            = 'true_ke'
#calc.tmatrix_parameters.tl_threshold         = 1e-6

calc.source_parameters.energy = XRaySource.AL_KALPHA
calc.source_parameters.theta  = -55.
calc.source_parameters.phi    = -55.

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

calc.calculation_parameters.scattering_order    = 4
calc.calculation_parameters.RA_cutoff           = 2
calc.calculation_parameters.path_filtering      = 'forward_scattering'
calc.calculation_parameters.off_cone_events     = 1
calc.calculation_parameters.vibrational_damping = 'averaged_tl'
calc.calculation_parameters.vibration_scaling   = 3.

if do_ped:
    calc.muffintin_parameters.interstitial_potential = 14
    data = calc.get_theta_scan(phi=0, theta=np.arange(-5, 60.5, 0.5), 
                               level='2p',
                               kinetic_energy=1200)
else:
    data = calc.get_theta_scan(phi=0, theta=np.arange(-5, 60.5, 0.5), 
                               edge='KL2L2', multiplet='1D2',
                               kinetic_energy=1200)

data.save('results.hdf5')