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

this module runs the PHAGEN program to calculate scattering factors and radial matrix element.

@author Matthias Muntwiler

@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 logging
import os
import shutil
import tempfile

from pmsco.calculators.calculator import AtomicCalculator
from pmsco.calculators.phagen.phagen import libmain
from pmsco.calculators.phagen.translator import Translator
import pmsco.cluster

logger = logging.getLogger(__name__)


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, cluster, scan, output_file):
        """
        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 = os.getcwd()
        try:
            with tempfile.TemporaryDirectory() as temp_dir:
                os.chdir(temp_dir)
                os.mkdir("div")
                os.mkdir("div/wf")
                os.mkdir("plot")
                os.mkdir("data")

                # prepare input for phagen
                infile = "phagen.in"
                outfile = "phagen.out"

                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))

                # call phagen
                libmain(infile, outfile)

                # collect results
                try:
                    phafile = outfile + ".pha"
                    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:
                    radfile = outfile + ".rad"
                    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:
                    clufile = outfile + ".clu"
                    phagen_cluster.load_from_file(clufile, pmsco.cluster.FMT_PHAGEN_OUT)
                except IOError:
                    logger.error("error loading phagen cluster file {fi}".format(fi=clufile))

                try:
                    listfile = outfile + ".list"
                    report_listfile = os.path.join(prev_wd, output_file + ".phagen.list")
                    shutil.copy(listfile, report_listfile)
                    files[report_listfile] = "log"
                except IOError:
                    logger.error("error loading phagen list file {fi}".format(fi=listfile))

        finally:
            os.chdir(prev_wd)

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

        rmefile = output_file + ".rme"
        transl.write_edac_emission(rmefile)
        files[rmefile] = "atomic"

        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
