from ase import Atoms
from ase.build import add_adsorbate, bulk

from msspec.calculator import MSSPEC, RFACTOR
from msspec.utils import hemispherical_cluster

import numpy as np


def create_cluster(height=0.8, theta=45, phi=0, bond_length=1.15):
    iron = bulk('Fe', cubic=True)
    cluster = hemispherical_cluster(iron, diameter=5, planes=2, emitter_plane=1)

    t = np.radians(theta)
    p = np.radians(phi)

    z = bond_length * np.cos(t)
    x = bond_length * np.sin(t) * np.cos(p)
    y = bond_length * np.sin(t) * np.sin(p)
    CO=Atoms('CO',positions=[(0,0,0),(x,y,z)])

    add_adsorbate(cluster,CO, height=height)

    # Store some information in the cluster object
    cluster.info.update(adsorbate={'theta': theta, 'phi': phi, 'height': height, 'bond_length': bond_length})
    return cluster


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

    calc.calculation_parameters.scattering_order = 1

    polar_angles = np.arange(-5, 85, 0.5)
    # set the Carbon as absorber and compute the polar scan
    cluster.absorber = cluster.get_chemical_symbols().index('C')
    data = calc.get_theta_scan(level='1s', theta=polar_angles, kinetic_energy=1202)

    return data




results    = []
parameters = {'theta': [], 'phi': [], 'height': []}
for theta in np.arange(53, 57, 1):
    for phi in [0, 30, 45]:
        for height in np.arange(0.1, 0.8, 0.1):
            # Create the cluster
            cluster = create_cluster(theta=theta, phi=phi, height=height,
                                    bond_length=1.157)
            # Compute
            data = compute_polar_scan(cluster)

            # Update lists of results and parameters
            results.append(data[-1].theta.copy())
            results.append(data[-1].cross_section.copy())
            parameters['theta'].append(theta)
            parameters['phi'].append(phi)
            parameters['height'].append(height)


# Load the experimental data
exp_data = np.loadtxt('experimental_data.txt')

# Create an R-Factor calculator
rfc = RFACTOR()
rfc.set_references(exp_data[:,0], exp_data[:,1])

# Perform the R-Factor analysis
data = rfc.run(*results, **parameters)
data.view()