"""
@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-22 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
"""

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

# noinspection PyUnresolvedReferences
from pmsco.calculators.calculator import InternalAtomicCalculator
# noinspection PyUnresolvedReferences
from pmsco.calculators.edac import EdacCalculator
# noinspection PyUnresolvedReferences
from pmsco.calculators.phagen.runner import PhagenCalculator
# noinspection PyUnresolvedReferences
from pmsco.cluster import Cluster, ClusterGenerator
# noinspection PyUnresolvedReferences
import pmsco.elements.bindingenergy
from pmsco.graphics.scan import render_scan
from pmsco.project import Project, ModelSpace, CalculatorParams
# noinspection PyUnresolvedReferences
from pmsco.scan import ScanKey, ScanLoader, ScanCreator
# noinspection PyUnresolvedReferences
from pmsco.dispatch import CalcID
from pmsco.helpers import BraceMessage as BMsg

logger = logging.getLogger(__name__)


class FCC111Project(Project):
    def __init__(self):
        """
        initialize a project instance

        the element attribute must be set directly after creation, e.g. via run-file.

        unlike previous versions, the current version of this class does not define a scan_dict.
        the scans should be declared in the run-file using the ScanLoader and ScanCreator classes.
        the demo_holo_scan.etpi and demo_alpha_scan.etpai files can be used as templates.
        """
        super(FCC111Project, self).__init__()
        self.element = "Ni"
        self.scan_dict = {}
        self.phase_files = {}
        self.rme_files = {}

    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 = 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 = CalculatorParams()

        params.title = "fcc(111)"
        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 = pt.elements.symbol(emitter).binding_energy[initial_state]
        params.polarization = "H"
        params.z_surface = model['Zsurf']
        params.inner_potential = model['V0']
        params.work_function = 4.5
        params.polar_incidence_angle = 60.0
        params.azimuthal_incidence_angle = 0.0
        params.angular_resolution = 5.0
        params.experiment_temperature = 300.0
        params.debye_temperature = 400.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

        # edac_interface only
        params.emitters = []
        params.lmax = 15
        params.dmax = 5.0
        params.orders = [25]
        # params.phase_output_classes = self.cluster_generator.create_cluster(model, index).get_atom_count()

        return params

    def create_model_space(self):
        """
        define the model space of the optimization parameters.
        """
        dom = 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
