"""
@package projects.twoatom
Two-atom demo scattering calculation project
"""

import logging
import math
import numpy as np
import periodictable as pt

# noinspection PyUnresolvedReferences
from pmsco.pmsco import main
from pmsco.calculators.calculator import InternalAtomicCalculator
from pmsco.calculators.edac import EdacCalculator
from pmsco.calculators.phagen.runner import PhagenCalculator
from pmsco.cluster import Cluster, ClusterGenerator
from pmsco.project import CalculatorParams, ModelSpace, Project
from pmsco.helpers import BraceMessage as BMsg

# the following imports are used in run files
# noinspection PyUnresolvedReferences
import pmsco.data
# noinspection PyUnresolvedReferences
from pmsco.scan import ScanKey, ScanLoader, ScanCreator, HoloScanCreator

logger = logging.getLogger(__name__)


class TwoatomCluster(ClusterGenerator):
    """
    cluster of two atoms.

    atom A (top) is set at position (0, 0, 0), atom B (bottom) at (-dx, -dy, -dz)
    where dx, dy and dz are calculated from model parameters.
    the type of the atoms is set upon construction.

    the model parameters are:
    @arg @c model['dAB']   : distance between the two atoms in Angstrom.
    @arg @c model['th']    : polar angle of the connection line, 0 = on top geometry.
    @arg @c model['ph']    : azimuthal angle of the connection line, 0 = polar angle affects X coordinate.

    the class is designed to be reusable in various projects.
    object attributes refine the atom types and the mapping of project-specific model parameters.
    """

    ## @var atom_types (dict)
    # chemical element numbers of the cluster atoms.
    #
    # atom 'A' is the top atom, 'B' the bottom one.
    # upon construction both atoms are set to oxygen.
    # to customize, call @ref set_atom_type.

    ## @var model_dict (dict)
    # mapping of model parameters to cluster parameters
    #
    # the default model parameters used by the cluster are 'dAB', 'th' and 'ph'.
    # if the project uses other parameter names, e.g. 'dCO' instead of 'dAB',
    # the project-specific names can be declared here.
    # in the example, set model_dict['dAB'] = 'dCO'.

    def __init__(self, project):
        """
        initialize the cluster generator.

        the atoms and model dictionary are given default values.
        see @ref set_atom_type and @ref model_dict for customization.

        @param project: project instance.
        """
        super(TwoatomCluster, self).__init__(project)

        self.atom_types = {'A': pt.O.number, 'B': pt.O.number}
        self.model_dict = {'dAB': 'dAB', 'th': 'th', 'ph': 'ph'}

    def set_atom_type(self, atom, element):
        """
        set the type (chemical element) of an atom.

        @param atom: atom key, 'A' (top) or 'B' (bottom).
        @param element: chemical element number or symbol.
        """
        try:
            self.atom_types[atom] = int(element)
        except ValueError:
            self.atom_types[atom] = pt.elements.symbol(element.strip()).number

    def count_emitters(self, model, index):
        """
        return the number of emitter configurations.

        this cluster supports only one configuration.

        @param model:
        @param index:
        @return 1
        """
        return 1

    def create_cluster(self, model, index):
        """
        create a cluster given the model parameters and index.

        @param model:
        @param index:
        @return a pmsco.cluster.Cluster object containing the atomic coordinates.
        """
        r = model[self.model_dict['dAB']]
        try:
            th = math.radians(model[self.model_dict['th']])
        except KeyError:
            th = 0.
        try:
            ph = math.radians(model[self.model_dict['ph']])
        except KeyError:
            ph = 0.

        dx = r * math.sin(th) * math.cos(ph)
        dy = r * math.sin(th) * math.sin(ph)
        dz = r * math.cos(th)

        clu = Cluster()
        clu.comment = "{0} {1}".format(self.__class__, index)
        clu.set_rmax(r * 2.0)

        a_top = np.array((0.0, 0.0, 0.0))
        a_bot = np.array((-dx, -dy, -dz))

        clu.add_atom(self.atom_types['A'], a_top, 1)
        clu.add_atom(self.atom_types['B'], a_bot, 0)

        return clu


class TwoatomProject(Project):
    """
    two-atom calculation project class.

    the cluster contains a nitrogen in the top layer,
    and a nickel atom in the second layer.
    The layer distance and the angle can be adjusted by parameters.

    the model parameters are:
    @arg @c model['dNNi']    : vertical distance N - Ni in Angstrom.
    @arg @c model['pNNi']    : polar angle of axis N - Ni in degrees. 0 = on top geometry.
    @arg @c model['V0']      : inner potential
    @arg @c model['Zsurf']   : position of surface
    """

    def __init__(self):
        super(TwoatomProject, self).__init__()
        self.cluster_generator = TwoatomCluster(self)
        self.cluster_generator.set_atom_type('A', 'N')
        self.cluster_generator.set_atom_type('B', 'Ni')
        self.cluster_generator.model_dict['dAB'] = 'dNNi'
        self.cluster_generator.model_dict['th'] = 'pNNi'
        self.cluster_generator.model_dict['ph'] = 'aNNi'
        self.atomic_scattering_factory = PhagenCalculator
        self.multiple_scattering_factory = EdacCalculator
        self.phase_files = {}
        self.rme_files = {}
        self.bindings = {}
        self.bindings['N'] = {'1s': 409.9}
        self.bindings['B'] = {'1s': 188.0}
        self.bindings['Ni'] = {'2s': 1008.6,
                               '2p': (870.0 + 852.7) / 2, '2p1/2': 870.0, '2p3/2': 852.7,
                               '3s': 110.8,
                               '3p': (68.0 + 66.2) / 2, '3p1/2': 68.0, '3p3/2': 66.2}

    def create_params(self, model, index):
        """
        set a specific set of parameters given the optimizable parameters.

        @param model: (dict) optimizable parameters
        """
        params = CalculatorParams()

        params.title = "two-atom demo"
        params.comment = "{0} {1}".format(self.__class__, index)
        params.cluster_file = ""
        params.output_file = ""
        initial_state = self.scans[index.scan].initial_state
        params.initial_state = initial_state
        emitter = self.scans[index.scan].emitter
        params.binding_energy = self.bindings[emitter][initial_state]
        params.polarization = "H"
        params.z_surface = model['Zsurf']
        params.inner_potential = model['V0']
        params.work_function = 3.6
        params.polar_incidence_angle = 60.0
        params.azimuthal_incidence_angle = 0.0
        params.experiment_temperature = 300.0
        params.debye_temperature = 356.0

        if self.phase_files:
            state = emitter + initial_state
            try:
                params.phase_files = self.phase_files[state]
            except KeyError:
                params.phase_files = {}
                logger.warning("no phase files found for {} - using default calculator".format(state))

        params.rme_files = {}
        params.rme_minus_value = 0.1
        params.rme_minus_shift = 0.0
        params.rme_plus_value = 1.0
        params.rme_plus_shift = 0.0

        # used by EDAC only
        params.emitters = []
        params.lmax = 15
        params.dmax = 5.0
        params.orders = [25]

        return params

    def create_model_space(self):
        """
        define the domain of the optimization parameters.
        """
        dom = ModelSpace()

        if self.mode == "single":
            dom.add_param('dNNi',     2.109,  2.000,  2.250, 0.050)
            dom.add_param('pNNi',    15.000,  0.000, 30.000, 1.000)
            dom.add_param('V0',      21.966, 15.000, 25.000, 1.000)
            dom.add_param('Zsurf',    1.449,  0.500,  2.000, 0.250)
        elif self.mode == "swarm":
            dom.add_param('dNNi',     2.109,  2.000,  2.250, 0.050)
            dom.add_param('pNNi',    15.000,  0.000, 30.000, 1.000)
            dom.add_param('V0',      21.966, 15.000, 25.000, 1.000)
            dom.add_param('Zsurf',    1.449,  0.500,  2.000, 0.250)
        elif self.mode == "grid":
            dom.add_param('dNNi',     2.109,  2.000,  2.250, 0.050)
            dom.add_param('pNNi',    15.000,  0.000, 30.000, 1.000)
            dom.add_param('V0',      21.966, 15.000, 25.000, 1.000)
            dom.add_param('Zsurf',    1.449,  0.500,  2.000, 0.250)
        else:
            dom.add_param('dNNi',     2.109,  2.000,  2.250, 0.050)
            dom.add_param('pNNi',    15.000,  0.000, 30.000, 1.000)
            dom.add_param('V0',      21.966, 15.000, 25.000, 1.000)
            dom.add_param('Zsurf',    1.449,  0.500,  2.000, 0.250)

        return dom


def example_intensity(e, t, p, a):
    """
    arbitrary intensity pattern for example data

    this function can be used to calculate the intensity in example scan files.
    the function implements an arbitrary modulation function

    @param e: energy
    @param t: theta
    @param p: phi
    @param a: alpha
    @return intensity
    """
    i = np.random.random() * 1e6 * \
        np.cos(np.radians(t)) ** 2 * \
        np.cos(np.radians(a)) ** 2 * \
        np.cos(np.radians(p)) ** 2 * \
        np.sin(e / 1000. * np.pi * 0.1 / np.sqrt(e)) ** 2
    return i