from ase.build import bulk
from msspec.calculator import MSSPEC
from msspec.utils import hemispherical_cluster, get_atom_index, cut_plane
import numpy as np
from matplotlib import pyplot as plt


Ag = bulk('Ag', cubic=True)
Ag.rotate((1,1,1), (0,0,1), rotate_cell=True)
Ag.rotate(15, 'z', rotate_cell=True)

cluster = hemispherical_cluster(Ag, diameter=20, emitter_plane=0)

cluster = cut_plane(cluster, z=-4.8)
cluster.emitter = get_atom_index(cluster, 0,0,0)
cluster[cluster.emitter].symbol = 'Sb'

cluster.edit()

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

calc.source_parameters.theta = 0
calc.source_parameters.phi = 0

calc.detector_parameters.angular_acceptance = 1
calc.detector_parameters.average_sampling = 'low'

calc.muffintin_parameters.interstitial_potential = 0
data = calc.get_phi_scan(level='4d', theta=40, phi=np.linspace(0,240,121), kinetic_energy=45)

# normalize data between [0,1]
dset = data[0]
dset.cross_section -= dset.cross_section.min()
dset.cross_section /= dset.cross_section.max()

# Add experimental data points in the dataset
x, y = np.loadtxt('data.txt').T
dset.add_columns(experiment=y)

# Add points to view
view = dset.views[0]
view.select('phi', 'experiment', legend='Exp. data')

data.view()

calc.shutdown()