#!/usr/bin/env python

"""
@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 division
import sys
import os
import math
import numpy as np
import periodictable as pt
import argparse
import logging

# adjust the system path so that the main PMSCO code is found
base_dir = os.path.dirname(__file__) or '.'
package_dir = os.path.join(base_dir, '../..')
package_dir = os.path.abspath(package_dir)
sys.path.append(package_dir)

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

logger = logging.getLogger(__name__)


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 = {}

    def create_cluster(self, model, index):
        """
        calculate a specific set of atom positions given the optimizable parameters.

        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.

        @param model: (dict) optimizable parameters
        """
        clu = mc.Cluster()
        clu.comment = "{0} {1}".format(self.__class__, index)
        clu.set_rmax(10.0)

        a_N = np.array((0.0, 0.0, 0.0))
        rad_pNNi = math.radians(model['pNNi'])
        a_Ni1 = np.array((0.0,
            -model['dNNi'] * math.sin(rad_pNNi),
            -model['dNNi'] * math.cos(rad_pNNi)))

        clu.add_atom(pt.N.number, a_N, 1)
        clu.add_atom(pt.Ni.number, a_Ni1, 0)

        return clu

    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_broadening = 0.0
        params.lattice_constant = 1.0
        params.z_surface = model['Zsurf']
        params.atom_types = 3
        params.atomic_number = [7, 28]
        params.phase_file = ["hbn_n.pha", "ni.pha"]
        params.msq_displacement = [0.01, 0.01, 0.00]
        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 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


def main():
    args, unknown_args = pmsco.pmsco.parse_cli()
    if unknown_args:
        project_args = parse_project_args(unknown_args)
    else:
        project_args = None

    project = create_project()
    pmsco.pmsco.set_common_args(project, args)
    set_project_args(project, project_args)
    pmsco.pmsco.run_project(project)

if __name__ == '__main__':
    main()
    sys.exit(0)