"""
@package pmsco.calculators.phagen.runner
Natoli/Sebilleau PHAGEN interface

This module runs the PHAGEN program to calculate scattering factors and radial matrix elements.

Requires PHAGEN version 2.2 from https://git.ipr.univ-rennes.fr/epsi/msspec_python3.git (contained in subprojects).

@author Matthias Muntwiler

@copyright (c) 2015-23 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 os
import shutil
import tempfile
from pathlib import Path
import sys

from pmsco.calculators.calculator import AtomicCalculator

from pmsco.calculators.phagen.translator import Translator
import pmsco.cluster
from pmsco.helpers import stdout_redirected
import pmsco.project

logger = logging.getLogger(__name__)

try:
    import phagen
except (ImportError, ModuleNotFoundError) as e:
    phagen = None
    logger.critical("Error importing the phagen package.", exc_info=e)


class PhagenCalculator(AtomicCalculator):
    """
    use the PHAGEN program to calculate scattering factors and radial matrix element.

    this produces scatterer, radial matrix element and cluster files for EDAC.
    """

    def run(self,
            params: pmsco.project.CalculatorParams,
            cluster: pmsco.cluster.Cluster,
            scan: pmsco.project.Scan,
            output_file: str):
        """
        create the input file, run PHAGEN, and translate the output to EDAC format.

        the following files are created in the job work directory:
        - scattering factor files in EDAC format.
          their names are `output_file + "_{atomclass}.scat"`.
        - radial matrix element file in EDAC format.
          its name is `output_file + ".rme"`.
        - cluster file in PMSCO format.
          its name is `output_file + ".clu"`.

        the cluster and params objects are updated and linked to the scattering files
        so that they can be passed to EDAC without further modification.
        the radial matrix element is currently not used.

        note that the scattering files are numbered according to the atomic environment and not chemical element.
        this means that the updated cluster (cluster object or ".clu" file)
        must be used in the scattering calculation.
        atomic index is not preserved - atoms in the input and output clusters can only be related by coordinate!

        because PHAGEN generates a lot of files with hard-coded names,
        the function creates a temporary directory for PHAGEN and deletes it before returning.

        @param params: pmsco.project.CalculatorParams object.
            the phase_files attribute is updated with the paths of the scattering files.

        @param cluster: pmsco.cluster.Cluster object.
            the cluster is updated with the one returned from PHAGEN.
            the atom classes are linked to the scattering files.

        @param scan: pmsco.project.Scan object.
            the scan object is used to determine the kinetic energy range.

        @param output_file: base path and name of the output files.

        @return (None, dict) where dict is a list of output files with their category.
            the category is "atomic" for all output files.
        """
        assert cluster.get_emitter_count() == 1, "PHAGEN cannot handle more than one emitter at a time"

        transl = Translator()
        transl.params.set_params(params)
        transl.params.set_cluster(cluster)
        transl.params.set_scan(scan)
        phagen_cluster = pmsco.cluster.Cluster()

        files = {}
        prev_wd = Path.cwd()
        try:
            with tempfile.TemporaryDirectory() as temp_dir:
                temp_path = Path(temp_dir)
                in_path = temp_path / "input"
                in_path.mkdir(exist_ok=True)
                out_path = temp_path / "output"
                out_path.mkdir(exist_ok=True)

                infile = in_path / "input.ms"
                try:
                    transl.write_input(infile)
                    report_infile = os.path.join(prev_wd, output_file + ".phagen.in")
                    shutil.copy(infile, report_infile)
                    files[report_infile] = "input"
                except IOError:
                    logger.warning("error writing phagen input file {fi}.".format(fi=infile))

                report_listfile = os.path.join(prev_wd, output_file + ".phagen.list")
                files[report_listfile] = "log"

                # call phagen, redirect stdout (unit 6)
                os.chdir(out_path)
                with open(report_listfile, "wb") as f:
                    with stdout_redirected(f):
                        phagen.phagen()

                phafile = out_path / "div" / "phases.dat"
                radfile = out_path / "fort.55"
                # tlfile = out_path / "fort.35"
                clufile = out_path / "clus" / "clus.out"

                # collect results
                try:
                    transl.parse_phagen_phase(phafile)
                    report_phafile = os.path.join(prev_wd, output_file + ".phagen.pha")
                    shutil.copy(phafile, report_phafile)
                    files[report_phafile] = "output"
                except IOError:
                    logger.error("error loading phagen phase file {fi}".format(fi=phafile))

                try:
                    transl.parse_radial_file(radfile)
                    report_radfile = os.path.join(prev_wd, output_file + ".phagen.rad")
                    shutil.copy(radfile, report_radfile)
                    files[report_radfile] = "output"
                except IOError:
                    logger.error("error loading phagen radial file {fi}".format(fi=radfile))

                try:
                    phagen_cluster.load_from_file(clufile, pmsco.cluster.FMT_PHAGEN_OUT)
                except IOError:
                    logger.error("error loading phagen cluster file {fi}".format(fi=clufile))

        finally:
            os.chdir(prev_wd)

        # write edac files
        scatfile = output_file + "_{}.scat"
        scatfiles = transl.write_edac_scattering(scatfile)
        params.phase_files = scatfiles.copy()
        files.update({f: "atomic" for f in params.phase_files.values()})

        rmefile = output_file + "_{}.rme"
        rmefiles = transl.write_edac_emission(rmefile)
        params.rme_files = rmefiles.copy()
        files.update({f: "atomic" for f in params.rme_files.values()})

        cluster.update_atoms(phagen_cluster, {'c'})
        clufile = output_file + ".pmsco.clu"
        cluster.save_to_file(clufile, pmsco.cluster.FMT_PMSCO)
        files[clufile] = "cluster"

        return None, files
