# -*- encoding: utf-8 -*-
# vim: set fdm=indent ts=4 sw=4 sts=4 et ai tw=80 cc=+0 mouse=a nu : #

from msspec.calculator import MSSPEC, XRaySource
from msspec.utils import *

from ase.build import fcc111, add_adsorbate
from ase.visualize import view
from msspec.iodata import cols2matrix

from matplotlib import pyplot as plt
import numpy as np
import sys

data = None
all_z = np.arange(1.10, 1.50, 0.02)
all_z=(1.1,)
for zi, z0 in enumerate(all_z):
    # construct the cluster
    cluster = fcc111('Rh', size = (2,2,1))
    add_adsorbate(cluster, 'O', z0, position = 'fcc')
    cluster.pop(3)
    #cluster.rotate('z',np.pi/3.)
    #view(cluster)

    cluster.absorber = len(cluster) - 1

    calc = MSSPEC(spectroscopy = 'PED', folder = './RhO_z')
    calc.set_atoms(cluster)
    calc.calculation_parameters.scattering_order = 3
    calc.calculation_parameters.RA_cutoff = 1
    calc.source_parameters.energy = XRaySource.AL_KALPHA

    # compute
    level = '1s'
    ke = 723.


    data = calc.get_theta_phi_scan(level=level, kinetic_energy=ke, data=data)

    # OPTIONAL, this will create an image of the data that you can combine
    # in an animated gif
    dset = data[-1]
    theta, phi, Xsec = cols2matrix(dset.theta, dset.phi, dset.cross_section)
    X, Y = np.meshgrid(np.radians(phi), 2*np.tan(np.radians(theta/2.)))
    fig = plt.figure()
    ax = fig.add_subplot(111, projection='polar')
    im = ax.pcolormesh(X, Y, Xsec)
    theta_ticks = np.arange(0, 91, 15)
    plt.yticks(2 * np.tan(np.radians(theta_ticks/2.)), theta_ticks)
    plt.title(r"$z_0 = {:.2f} \AA$".format(z0))
    plt.savefig('image{:03d}.png'.format(zi))

data.view()