"""
@package pmsco.projects.demo.molecule
scattering calculation project for single molecules

the atomic positions are read from a molecule file.
cluster file, emitter (by chemical symbol), initial state and kinetic energy are specified on the command line.
there are no structural parameters.

@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
import os.path
from pathlib import Path
import periodictable as pt
import argparse
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
from pmsco.data import calc_modfunc_loess
# noinspection PyUnresolvedReferences
import pmsco.elements.bindingenergy
from pmsco.helpers import BraceMessage as BMsg
import pmsco.project as 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

    def load_base_cluster(self):
        """
        load and cache the project-defined coordinate file.

        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(120.0)
            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 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.
        """
        clu = self.create_cluster(model, index)
        return clu.get_emitter_count()

    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
        """
        self.load_base_cluster()
        clu = cluster.Cluster()
        clu.copy_from(self.base_cluster)
        clu.comment = f"{self.__class__}, {index}"
        dom = self.project.domains[index.domain]

        # trim
        clu.set_rmax(model['rmax'])
        clu.trim_sphere(clu.rmax)

        # emitter selection
        idx_emit = np.where(clu.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]
        clu.data['e'][idx_emit] = 1

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

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

        return clu


class MoleculeProject(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 = project.ModelSpace()
        self.scan_dict = {}
        self.cluster_file = "demo-cluster.xyz"
        self.cluster_generator = MoleculeFileCluster(self)
        self.atomic_scattering_factory = PhagenCalculator
        self.multiple_scattering_factory = EdacCalculator
        self.phase_files = {}
        self.rme_files = {}
        self.modf_smth_ei = 0.5

    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 = 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']
        params.phase_files = self.phase_files
        params.rme_files = self.rme_files
        # 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

    # noinspection PyUnusedLocal
    def calc_modulation(self, data, model):
        """
        calculate the modulation function with project-specific smoothing factor

        see @ref pmsco.pmsco.project.calc_modulation.

        @param data: (numpy.ndarray) experimental data in ETPI, or ETPAI format.

        @param model: (dict) model parameters of the calculation task. not used.

        @return copy of the data array with the modulation function in the 'i' column.
        """
        return calc_modfunc_loess(data, smth=self.modf_smth_ei)


def create_model_space(mode):
    """
    define the model space.
    """
    dom = project.ModelSpace()

    if mode == "single":
        dom.add_param('zsurf',   1.20)
        dom.add_param('Texp',  300.00)
        dom.add_param('Tdeb',  100.00)
        dom.add_param('V0',     10.00)
        dom.add_param('rmax',   50.00)
        dom.add_param('ares',    5.00)
        dom.add_param('distm',   5.00)
        dom.add_param('wdom1', 1.0)
        dom.add_param('wdom2', 1.0)
        dom.add_param('wdom3', 1.0)
        dom.add_param('wdom4', 1.0)
        dom.add_param('wdom5', 1.0)
    else:
        raise ValueError(f"undefined model space for {mode} optimization")

    return dom


def create_project():
    """
    create the project instance.
    """

    proj = MoleculeProject()
    proj_dir = os.path.dirname(os.path.abspath(__file__))
    proj.project_dir = proj_dir

    # scan dictionary
    # to select any number of scans, add their dictionary keys as scans option on the command line
    proj.scan_dict['empty'] = {'filename': os.path.join(proj_dir, "../common/empty-hemiscan.etpi"),
                               'emitter': "N", 'initial_state': "1s"}

    proj.mode = 'single'
    proj.model_space = create_model_space(proj.mode)
    proj.job_name = 'molecule0000'
    proj.description = 'molecule demo'

    return proj


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

    @param project: project instance

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

    scans = []
    try:
        if project_args.scans:
            scans = project_args.scans
        else:
            logger.error("missing scan argument")
            exit(1)
    except AttributeError:
        logger.error("missing scan argument")
        exit(1)

    for scan_key in scans:
        scan_spec = project.scan_dict[scan_key]
        project.add_scan(**scan_spec)

    try:
        project.cluster_file = os.path.abspath(project_args.cluster_file)
        project.cluster_generator = MoleculeFileCluster(project)
    except (AttributeError, TypeError):
        logger.error("missing cluster-file argument")
        exit(1)

    try:
        if project_args.emitter:
            for scan in project.scans:
                scan.emitter = project_args.emitter
            logger.warning(f"override emitters of all scans to {project_args.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(f"override scan energy of all scans to {project_args.energy}")
    except AttributeError:
        pass

    try:
        if project_args.symmetry:
            for angle in np.linspace(0, 360, num=project_args.symmetry, endpoint=False):
                project.add_domain({'xrot': 0., 'yrot': 0., 'zrot': angle})
                logger.warning(f"override rotation symmetry to {project_args.symmetry}")
    except AttributeError:
        pass


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

    # main arguments
    parser.add_argument('--scans', nargs="*",
                        help="nick names of scans to use in calculation (see create_project function)")
    parser.add_argument('--cluster-file',
                        help="path name of molecule file (xyz format).")

    # conditional arguments
    parser.add_argument('--emitter',
                        help="emitter: chemical symbol")
    parser.add_argument('--initial-state',
                        help="initial state term: e.g. 2p1/2")
    parser.add_argument('--energy', type=float,
                        help="kinetic energy (eV)")
    parser.add_argument('--symmetry', type=int, default=1,
                        help="n-fold rotational symmetry")

    parsed_args = parser.parse_args(_args)
    return parsed_args