"""
@package pmsco.projects.fcc
scattering calculation project for the (111) surface of an arbitrary face-centered cubic crystal

@author Matthias Muntwiler, matthias.muntwiler@psi.ch

@copyright (c) 2015-19 by Paul Scherrer Institut @n
Licensed under the Apache License, Version 2.0 (the "License"); @n
  you may not use this file except in compliance with the License.
  You may obtain a copy of the License at
  http://www.apache.org/licenses/LICENSE-2.0
"""

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

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

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

logger = logging.getLogger(__name__)


class FCC111Project(mp.Project):
    def __init__(self):
        """
        initialize a project instance
        """
        super(FCC111Project, self).__init__()
        self.scan_dict = {}
        self.element = "Ni"

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

        @param model  (dict)          optimizable parameters
            @arg    model['dlat']     bulk lattice constant in Angstrom
            @arg    model['dl1l2']    distance between top and second layer (may deviate from bulk)
            @arg    model['rmax']     cluster radius
            @arg    model['phi']      azimuthal rotation angle in degrees
        """
        clu = mc.Cluster()
        clu.comment = "{0} {1}".format(self.__class__, index)
        clu.set_rmax(model['rmax'])
        # fcc lattice constant
        a_lat = model['dlat']
        # surface lattice constant of the (111) surface
        a_surf = a_lat / math.sqrt(2.0)

        # lattice vectors
        # a1 and a2 span the (111) surface
        a1 = np.array((a_surf, 0.0, 0.0))
        a2 = np.array((a_surf / 2.0, a_surf * math.sqrt(3.0) / 2.0, 0.0))
        a3 = np.array((0.0, a_surf * math.sqrt(3.0) / 3.0, a_lat * math.sqrt(3.0) / 3))

        a_l1 = np.array((0.0, 0.0, 0.0))
        a_l2 = np.array(((a1[0] + a2[0]) * 2.0 / 3.0,
            (a1[1] + a2[1]) * 2.0 / 3.0,
            -(model['dl1l2'])))
        a_l3 = np.array(((a1[0] + a2[0]) / 3.0,
            (a1[1] + a2[1]) / 3.0,
            -(a3[2] + model['dl1l2'])))
        a_bulk = np.array((0.0, 0.0,
            -(2.0 * a3[2] + model['dl1l2'])))

        clu.add_layer(self.element, a_l1, a1, a2)
        clu.add_layer(self.element, a_l2, a1, a2)
        clu.add_layer(self.element, a_l3, a1, a2)
        clu.add_bulk(self.element, a_bulk, a1, a2, a3, a_bulk[2] + 0.01)

        clu.set_emitter(a_l1)

        clu.rotate_z(model['phi'])

        return clu

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

        par = optimizable parameters
        par['V0']  = inner potential
        par['Zsurf'] = position of surface
        """
        params = mp.CalculatorParams()

        params.title = "fcc(111)"
        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.atom_types = 3
        params.atomic_number = [pt.elements.symbol(self.element).number]
        params.phase_file = []
        params.msq_displacement = [0.00]
        params.planewave_attenuation = 1.0
        params.inner_potential = model['V0']
        params.work_function = 4.5
        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 = pt.elements.symbol(self.element).mass
        params.experiment_temperature = 300.0
        params.debye_temperature = 400.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_model_space(self):
        """
        define the model space of the optimization parameters.
        """
        dom = mp.ModelSpace()

        if self.mode == "single":
            dom.add_param('rmax',     5.00,    5.00, 15.00, 2.50)
            dom.add_param('phi',      0.00,    0.00,  0.00, 0.00)
            dom.add_param('dlat',     3.52,    2.00,  5.00, 0.10)
            dom.add_param('dl1l2',    2.03,    1.80,  2.20, 0.05)
            dom.add_param('V0',      10.00,    0.00, 20.00, 1.00)
            dom.add_param('Zsurf',    1.00,    0.00,  2.00, 0.50)
        elif self.mode == "swarm":
            dom.add_param('rmax',     7.50,    5.00, 15.00, 2.50)
            dom.add_param('phi',      0.00,    0.00,  0.00, 0.00)
            dom.add_param('dlat',     3.52,    2.00,  5.00, 0.10)
            dom.add_param('dl1l2',    2.03,    1.80,  2.20, 0.05)
            dom.add_param('V0',      10.00,    0.00, 20.00, 1.00)
            dom.add_param('Zsurf',    1.00,    0.00,  2.00, 0.50)
        elif self.mode == "grid":
            dom.add_param('rmax',     7.50,    5.00, 15.00, 2.50)
            dom.add_param('phi',      0.00,    0.00,  0.00, 0.00)
            dom.add_param('dlat',     3.52,    2.00,  5.00, 0.10)
            dom.add_param('dl1l2',    2.03,    1.80,  2.20, 0.05)
            dom.add_param('V0',      10.00,    0.00, 20.00, 1.00)
            dom.add_param('Zsurf',    1.00,    0.00,  2.00, 0.50)
        else:
            dom.add_param('rmax',     7.50,    5.00, 15.00, 2.50)
            dom.add_param('phi',      0.00,    0.00,  0.00, 0.00)
            dom.add_param('dlat',     3.52,    2.00,  5.00, 0.10)
            dom.add_param('dl1l2',    2.03,    1.80,  2.20, 0.05)
            dom.add_param('V0',      10.00,    0.00, 20.00, 1.00)
            dom.add_param('Zsurf',    1.00,    0.00,  2.00, 0.50)

        return dom


def create_project():
    """
    create an FCC111Project calculation project.
    """

    project = FCC111Project()

    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['default'] = {'filename': os.path.join(project_dir, "demo_holo_scan.etp"),
                                    'emitter': "Ni", 'initial_state': "3s"}
    project.scan_dict['holo'] = {'filename': os.path.join(project_dir, "demo_holo_scan.etp"),
                                 'emitter': "Ni", 'initial_state': "3s"}
    project.scan_dict['alpha'] = {'filename': os.path.join(project_dir, "demo_alpha_scan.etp"),
                                  'emitter': "Ni", 'initial_state': "3s"}

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

    return project


def set_project_args(project, project_args):
    """
    set the project arguments of a MnGeTeProject calculation project.

    @param project: project instance

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

    scans = ['default']
    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))

    try:
        if project_args.element:
            project.element = project_args.element
            for scan in project.scans:
                scan.emitter = project_args.element
            logger.warning(BMsg("override emitters to {0}", project.emitter))
    except AttributeError:
        pass

    try:
        if project_args.initial_state:
            for scan in project.scans:
                scan.initial_state = project_args.initial_state
            logger.warning(f"override initial states of all scans to {project_args.initial_state}")
    except AttributeError:
        pass

    try:
        if project_args.energy:
            for scan in project.scans:
                scan.energies = np.asarray((project_args.energy, ))
            logger.warning(BMsg("override scan energy, set to {0}", project_args.energy))
    except AttributeError:
        pass


def parse_project_args(_args):
    parser = argparse.ArgumentParser()

    # main arguments
    parser.add_argument('-e', '--element', help="chemical element symbol")
    parser.add_argument('-s', '--scans', nargs="*", default=['default'],
                        help="nick names of scans to use in calculation (see create_project function)")
    parser.add_argument('-i', '--initial-state',
                        help="inital state of photoelectron")
    parser.add_argument('--energy', type=float,
                        help="kinetic energy of photoelectron (override scan file)")

    parsed_args = parser.parse_args(_args)
    return parsed_args