from ase.build import mx2
from ase.visualize import view
from msspec.calculator import MSSPEC
from msspec.utils import hemispherical_cluster, get_atom_index

# some constants
a=3.535;c=6.004;d=3.450;D=2.554

TiSe2 = mx2(formula='TiSe2', kind='1T', a=a, thickness=d, size=(1, 1, 1), vacuum=None)
TiSe2.cell[2] = [0, 0, 6.004]
TiSe2.rotate(60, 'z', rotate_cell=True)

for i in range(3): 
    TiSe2[i].tag=i

data = None
cluster = hemispherical_cluster(TiSe2, emitter_tag=0, emitter_plane=1, planes=5)
cluster.emitter = get_atom_index(cluster, 0, 0, 0)


calc = MSSPEC(spectroscopy='PED', algorithm='expansion')
calc.set_atoms(cluster)

for ndif in range(1,4):    
    calc.calculation_parameters.scattering_order = ndif
    data = calc.get_theta_phi_scan(level='2p', kinetic_energy=1030, data=data)

data.view()