"""
@package pmsco.projects.demo.molecule
scattering calculation project for single molecules or coordinate files from other programs

the atomic positions are read from a molecule (.xyz) file.
emitters are selected by chemical element symbol or by an additional molecule file (emitter file)
that contains only those atoms of the cluster file which are inequivalent emitters.

cluster, emitters, initial state and kinetic energy are specified on the command line.
there are no structural parameters.

example 1: molecule from XYZ file
---------------------------------

the cluster file contains all atomic positions necessary for calculating the diffraction pattern.
emitters are selected by chemical element symbol.
the cluster is not trimmed.
normal emission is along the z-axis.


example 2: periodic structure from external program (e.g. Vesta)
----------------------------------------------------------------

the cluster file contains the unit cells spanned by at least 3 unit vectors in the surface plane.
the emitter file contains only one unit cell as a sub-set of the cluster file.
emitters can be narrowed down further by chemical element symbol.
an rmax parameter can be specified to trim the cluster.
normal emission is along the z-axis.


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

@copyright (c) 2015-20 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 math
import numpy as np
from pathlib import Path
import periodictable as pt
import logging

# 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
import pmsco.cluster as cluster
import pmsco.data
# noinspection PyUnresolvedReferences
import pmsco.elements.bindingenergy
from pmsco.helpers import BraceMessage as BMsg
import pmsco.project

logger = logging.getLogger(__name__)


class MoleculeFileCluster(cluster.ClusterGenerator):
    """
    cluster generator based on external file.

    work in progress.
    """
    def __init__(self, project):
        super(MoleculeFileCluster, self).__init__(project)
        self.base_cluster = None
        self.emitter_cluster = None

    def get_base_cluster(self):
        """
        return the project-defined atom coordinates, load from file if necessary.

        the file path is set in self.project.cluster_file.
        the file must be in XYZ (.xyz) or PMSCO cluster (.clu) format (cf. pmsco.cluster module).

        @return: Cluster object (also referenced by self.base_cluster)
        """
        if self.base_cluster is None:
            clu = cluster.Cluster()
            clu.set_rmax(np.inf)
            p = Path(self.project.cluster_file)
            ext = p.suffix
            if ext == ".xyz":
                fmt = cluster.FMT_XYZ
            elif ext == ".clu":
                fmt = cluster.FMT_PMSCO
            else:
                raise ValueError(f"unknown cluster file extension {ext}")
            clu.load_from_file(self.project.cluster_file, fmt=fmt)
            self.base_cluster = clu

        return self.base_cluster

    def get_emitter_cluster(self):
        """
        return the project-defined emitter coordinates, load from file if necessary.

        the file path is set in self.project.emitter_file.
        if the file path is None, the method loads the base cluster self.project.cluster_file.
        the file must be in XYZ (.xyz) or PMSCO cluster (.clu) format (cf. pmsco.cluster module).

        @return: Cluster object (also referenced by self.emitter_cluster).
            None if no emitter file was specified.
        """
        if self.emitter_cluster is None:
            clu = cluster.Cluster()
            clu.set_rmax(np.inf)
            try:
                p = Path(self.project.emitter_file)
            except TypeError:
                p = Path(self.project.cluster_file)
            ext = p.suffix
            if ext == ".xyz":
                fmt = cluster.FMT_XYZ
            elif ext == ".clu":
                fmt = cluster.FMT_PMSCO
            else:
                raise ValueError(f"unknown cluster file extension {ext}")
            clu.load_from_file(self.project.emitter_file, fmt=fmt)
            self.emitter_cluster = clu

        return self.emitter_cluster

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

        the method creates the full cluster and counts the emitters.

        @param model: model parameters.
        @param index: scan and domain are used by the create_cluster() method,
            emit decides whether the method returns the number of configurations (-1),
            or the number of emitters in the specified configuration (>= 0).
        @return: number of emitter configurations.
        """
        if index.emit == -1:
            clu = self.get_emitter_cluster()
            sel_emit = clu.data['s'] == self.project.scans[index.scan].emitter
            return np.sum(sel_emit)
        elif index.emit >= 0:
            return 1
        else:
            raise ValueError(f"invalid emitter index {index.emit}")

    def create_cluster(self, model, index):
        """
        import a cluster from a coordinate file (XYZ format).

        the method does the following:
        - load the cluster file specified by self.cluster_file.
        - trim the cluster according to model['rmax'].
        - mark the 6 nitrogen atoms at the center of the trimer as emitters.

        @param model: rmax is the trim radius of the cluster in units of the surface lattice constant.

        @param index (named tuple CalcID) calculation index.
            this method uses the domain index to look up domain parameters in
            `pmsco.project.Project.domains`.
            `index.emit` selects whether a single-emitter (>= 0) or all-emitter cluster (== -1) is returned.

        @return pmsco.cluster.Cluster object
        """
        clu = cluster.Cluster()
        clu.copy_from(self.get_base_cluster())
        clu.comment = f"{self.__class__}, {index}"
        dom = self.project.domains[index.domain]

        # emitter selection
        ems = cluster.Cluster()
        ems.copy_from(self.get_emitter_cluster())
        ems.set_rmax(model['rmax'] + 0.1)
        ems.trim_cylinder(clu.rmax, clu.rmax)

        idx_emit = np.where(ems.data['s'] == self.project.scans[index.scan].emitter)
        assert isinstance(idx_emit, tuple)
        idx_emit = idx_emit[0]
        if index.emit >= 0:
            idx_emit = idx_emit[index.emit]
            origin = ems.get_position(idx_emit)
            clu.translate(-origin)
            clu.data['e'] = 0
            clu.set_emitter(pos=np.array((0.0, 0.0, 0.0)))
        else:
            for idx in idx_emit:
                clu.set_emitter(pos=ems.get_position(idx))

        # rotation
        if 'xrot' in model:
            clu.rotate_x(model['xrot'])
        elif 'xrot' in dom:
            clu.rotate_x(dom['xrot'])
        if 'yrot' in model:
            clu.rotate_y(model['yrot'])
        elif 'yrot' in dom:
            clu.rotate_y(dom['yrot'])
        if 'zrot' in model:
            clu.rotate_z(model['zrot'])
        elif 'zrot' in dom:
            clu.rotate_z(dom['zrot'])

        # trim
        clu.set_rmax(model['rmax'] + 0.1)
        clu.trim_paraboloid(clu.rmax, -clu.rmax)

        logger.info(f"cluster for calculation {index}: "
                    f"{clu.get_atom_count()} atoms, {clu.get_emitter_count()} emitters")

        return clu


class MoleculeProject(pmsco.project.Project):
    """
    general molecule project.

    the following model parameters are used:

    @arg `model['zsurf']`   : position of surface above molecule (angstrom)
    @arg `model['Texp']`    : experimental temperature (K)
    @arg `model['Tdeb']`    : debye temperature (K)
    @arg `model['V0']`      : inner potential (eV)
    @arg `model['rmax']`    : cluster radius (angstrom)
    @arg `model['ares']`    : angular resolution (degrees, FWHM)
    @arg `model['distm']`   : dmax for EDAC (angstrom)

    the following domain parameters are used.
    they can also be specified as model parameters.

    @arg `'xrot'`           : rotation about x-axis (applied first) (deg)
    @arg `'yrot'`           : rotation about y-axis (applied after x) (deg)
    @arg `'zrot'`           : rotation about z-axis (applied after x and y) (deg)

    the project parameters are:

    @arg `cluster_file`    : name of cluster file of template molecule.
                              default: "dpdi-trimer.xyz"
    """
    def __init__(self):
        """
        initialize a project instance
        """
        super(MoleculeProject, self).__init__()
        self.model_space = pmsco.project.ModelSpace()
        self.cluster_file = "demo-cluster.xyz"
        self.emitter_file = None
        self.cluster_generator = MoleculeFileCluster(self)
        self.atomic_scattering_factory = PhagenCalculator
        self.multiple_scattering_factory = EdacCalculator
        self.phase_files = {}
        self.rme_files = {}

    def validate(self):
        """
        Validate project parameters

        Resolve paths of cluster and emitter files after calling the inherited method.

        @return: None
        """

        super().validate()
        self.cluster_file = self.directories.resolve_path(self.cluster_file)
        self.emitter_file = self.directories.resolve_path(self.emitter_file)
        logger.warning(f"cluster_file: {self.cluster_file}")
        logger.warning(f"emitter_file: {self.emitter_file}")

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

        @param model: (dict) optimization parameters
            this method requires zsurf, V0, Texp, Tdeb, ares and distm.

        @param index (named tuple CalcID) calculation index.
            this method formats the index into the comment line.
        """
        params = pmsco.project.CalculatorParams()

        params.title = "molecule demo"
        params.comment = f"{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 = model['ares']
        params.experiment_temperature = model['Texp']
        params.debye_temperature = model['Tdeb']

        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 = model['distm']
        params.orders = [20]

        return params

    def create_model_space(self):
        """
        define the range of model parameters.

        see the class description for a list of parameters.
        """

        return self.model_space