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

this file is specific to the project and the state of the data analysis,
as it contains particular parameter values.
"""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import argparse
import logging
import math
import numpy as np
import os.path
import periodictable as pt

import pmsco.cluster as mc
import pmsco.project as mp
from pmsco.helpers import BraceMessage as BMsg

logger = logging.getLogger(__name__)


class TwoatomCluster(mc.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 = mc.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(mp.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.scan_dict = {}
        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'

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

        @param model: (dict) optimizable parameters
        """
        params = mp.Params()

        params.title = "two-atom demo"
        params.comment = "{0} {1}".format(self.__class__, index)
        params.cluster_file = ""
        params.output_file = ""
        params.initial_state = self.scans[index.scan].initial_state
        params.spherical_order = 2
        params.polarization = "H"
        params.scattering_level = 5
        params.fcut = 15.0
        params.cut = 15.0
        params.angular_resolution = 0.0
        params.lattice_constant = 1.0
        params.z_surface = model['Zsurf']
        params.phase_files = {self.cluster_generator.atom_types['A']: "",
                              self.cluster_generator.atom_types['B']: ""}
        params.msq_displacement = {self.cluster_generator.atom_types['A']: 0.01,
                                   self.cluster_generator.atom_types['B']: 0.0}
        params.planewave_attenuation = 1.0
        params.inner_potential = model['V0']
        params.work_function = 3.6
        params.symmetry_range = 360.0
        params.polar_incidence_angle = 60.0
        params.azimuthal_incidence_angle = 0.0
        params.vibration_model = "P"
        params.substrate_atomic_mass = 58.69
        params.experiment_temperature = 300.0
        params.debye_temperature = 356.0
        params.debye_wavevector = 1.7558
        params.rme_minus_value = 0.0
        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_domain(self):
        """
        define the domain of the optimization parameters.
        """
        dom = mp.Domain()

        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


def create_project():
    """
    create a new TwoatomProject calculation project.

    the default experimental data file is @c twoatom_hemi_scan_250e.etpi
    in the same directory as this Python module.
    it defines a classic hemispherical angle scan grid
    but does not include measured data for optimization.

    @return project instance.
    """

    project = TwoatomProject()

    project_dir = os.path.dirname(os.path.abspath(__file__))
    project.data_dir = project_dir

    # scan dictionary
    # to select any number of scans, add their dictionary keys as scans option on the command line
    project.scan_dict['ea'] = {'filename': os.path.join(project_dir, "twoatom_energy_alpha.etpai"),
                               'emitter': "N", 'initial_state': "1s"}
    project.scan_dict['et0p'] = {'filename': os.path.join(project_dir, "twoatom_energy_theta_0p.etpi"),
                                 'emitter': "N", 'initial_state': "1s"}
    project.scan_dict['et180p'] = {'filename': os.path.join(project_dir, "twoatom_energy_theta_180p.etpi"),
                                   'emitter': "N", 'initial_state': "1s"}
    project.scan_dict['tp215e'] = {'filename': os.path.join(project_dir, "twoatom_hemi_215e.etpi"),
                                   'emitter': "N", 'initial_state': "1s"}
    project.scan_dict['tp250e'] = {'filename': os.path.join(project_dir, "twoatom_hemi_250e.etpi"),
                                   'emitter': "N", 'initial_state': "1s"}

    return project


def set_project_args(project, project_args):
    """
    set the project-specific arguments.

    @param project: project instance

    @param project_args: (Namespace object) project arguments.
    """

    scans = ['tp250e']
    try:
        if project_args.scans:
            scans = project_args.scans
        else:
            logger.warning(BMsg("missing scan argument, using {0}", scans[0]))
    except AttributeError:
        logger.warning(BMsg("missing scan argument, using {0}", scans[0]))

    for scan_key in scans:
        scan_spec = project.scan_dict[scan_key]
        project.add_scan(**scan_spec)
        logger.info(BMsg("add scan {filename} ({emitter} {initial_state})", **scan_spec))

    project.add_symmetry({'default': 0.0})


def parse_project_args(_args):
    """
    parse project-specific command line arguments.

    @param _args: list of project-specific arguments from the command line.
        this is typically the unknown_args return value from argparse.ArgumentParser.parse_known_args().

    @return: namespace object containing the specified arguments as attributes.
    """
    parser = argparse.ArgumentParser()

    # main arguments
    parser.add_argument('-s', '--scans', nargs="*", default=['tp250e'],
                        help="nick names of scans to use in calculation (see create_project function)")

    parsed_args = parser.parse_args(_args)

    return parsed_args