from ase.build import bulk, add_adsorbate, fcc100
from msspec.calculator import MSSPEC
from msspec.utils import cut_cylinder
import numpy as np

cluster = fcc100('Ni', size=(2, 2, 2))
add_adsorbate(cluster, 'Se', 1.55, 'hollow')
#add_adsorbate(cluster, 'Se', 2.34, 'ontop')
cluster = cluster.repeat((10,10,1))


def get_index_close_to_origin():
    iSe = [atom.index for atom in cluster if atom.symbol == 'Se']
    R = np.linalg.norm(cluster.positions[iSe], axis=1)
    i = iSe[np.argmin(R)]
    return i

cluster.center(about=(0,0,0))
i = get_index_close_to_origin()
cluster.positions -= cluster[i].position

cluster = cut_cylinder(cluster, radius=10)
emitter = get_index_close_to_origin()
cluster.emitter = emitter

if False:
    cluster.edit()
    exit()

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

calc.source_parameters.theta = -60
calc.calculation_parameters.scattering_order = 3
calc.calculation_parameters.path_filtering = 'forward_scattering'
[a.set('forward_angle', 30) for a in cluster] 
calc.muffintin_parameters.interstitial_potential = 13.2

data = calc.get_energy_scan(level='3d', kinetic_energy=[50,110,20], theta=0, phi=0)
data = calc.get_energy_scan(level='3d', kinetic_energy=[50,110,20], theta=0, phi=45, data=data)

data.view()


#cluster.edit()