"""
@package pmsco.calculators.edac
Garcia de Abajo EDAC program interface.

@author Matthias Muntwiler

@copyright (c) 2015-18 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 numpy as np
import os

import pmsco.calculators.calculator as calculator
from pmsco.compat import open
import pmsco.data as md
import pmsco.cluster as mc
import pmsco.edac.edac as edac
from pmsco.helpers import BraceMessage as BMsg

logger = logging.getLogger(__name__)


class EdacCalculator(calculator.Calculator):
    def write_input_file(self, params, scan, filepath):
        """
        write parameters to an EDAC input file

        EDAC will calculate results on a rectangular grid.
        the grid is constructed from the limits of the scan coordinates
        and the number of positions in each respective dimension.
        to avoid any confusion, the input scan should be rectangular with equidistant steps.

        the following scans are supported:
        (energy), (energy, theta), (energy, phi), (energy, alpha), (theta, phi) holo.
        except for the holo scan, each scan dimension must be linear.
        the holo scan is translated to a rectangular (theta, phi) scan
        where theta is copied
        and phi is replaced by a linear scan from the minimum to the maximum phi at 1 degree steps.
        the scan type is detected from the scan file.

        if alpha is defined, theta is implicitly set to normal emission! (to be generalized)

        @param params: a pmsco.project.Params object with all necessary values except cluster and output files set.

        @param scan: a pmsco.project.Scan() object describing the experimental scanning scheme.

        @param filepath: (str) name and path of the file to be created.

        @return dictionary of created files {filename: category}
        """
        files = {}

        with open(filepath, "w") as f:
            f.write("verbose off\n")
            f.write("cluster input {0}\n".format(params.cluster_file))
            f.write("emitters {0:d} l(A)\n".format(len(params.emitters)))
            for em in params.emitters:
                f.write("{0:f} {1:f} {2:f} {3:d}\n".format(*em))

            en = scan.energies + params.work_function
            en_min = en.min()
            en_max = en.max()
            if en.shape[0] <= 1:
                en_num = 1
            else:
                de = np.diff(en)
                de = de[de >= 0.01]
                de = de.min()
                en_num = int(round((en_max - en_min) / de)) + 1
                if en_num != en.shape[0]:
                    logger.warning("energy scan length mismatch: EDAC {0}, scan {1}".format(en_num, en.shape[0]))
                assert en_num < en.shape[0] * 10, \
                    "linearization of energy scan causes excessive oversampling {0}/{1}".format(en_num, en.shape[0])
            f.write("emission energy E(eV) {en0:f} {en1:f} {nen:d}\n".format(en0=en_min, en1=en_max, nen=en_num))

            if params.fixed_cluster:
                th = scan.alphas
                ph = np.remainder(scan.phis + 90.0, 360.0)
                f.write("fixed cluster\n")
                if np.abs(scan.thetas).max() > 0.0:
                    logger.warning("theta angle implicitly set to zero due to alpha scan.")
            else:
                th = np.unique(scan.thetas)
                ph = scan.phis
                f.write("movable cluster\n")

            th_min = th.min()
            th_max = th.max()
            if th.shape[0] <= 1:
                th_num = 1
            else:
                dt = np.diff(th)
                dt = dt[dt >= 0.1]
                dt = dt.min()
                if ph.shape[0] > 1:
                    # hemispherical scan
                    if th_min < 0:
                        th_min = max(th_min - dt, -90.0)
                    else:
                        th_min = max(th_min - dt, 0.0)
                    if th_max > 0:
                        th_max = min(th_max + dt, 90.0)
                    else:
                        th_max = min(th_max + dt, 0.0)
                th_num = int(round((th_max - th_min) / dt)) + 1
                assert th_num < th.shape[0] * 10, \
                    "linearization of theta scan causes excessive oversampling {0}/{1}".format(th_num, th.shape[0])

            f.write("beta {0}\n".format(params.polar_incidence_angle))
            f.write("incidence {0} {1}\n".format(params.polar_incidence_angle, params.azimuthal_incidence_angle))
            f.write("emission angle theta {th0:f} {th1:f} {nth:d}\n".format(th0=th_min, th1=th_max, nth=th_num))

            ph_min = ph.min()
            ph_max = ph.max()
            if th.shape[0] <= 1:
                # azimuthal scan
                ph_num = ph.shape[0]
            elif ph.shape[0] <= 1:
                # polar scan
                ph_num = 1
            else:
                # hemispherical scan
                dp = np.diff(ph)
                dp = dp[dp >= 0.1]
                dp = dp.min()
                ph_min = max(ph_min - dp, 0.0)
                ph_max = min(ph_max + dp, 360.0)
                dt = (th_max - th_min) / (th_num - 1)
                dp = min(dp, dt)
                ph_num = int(round((ph_max - ph_min) / dp)) + 1
                assert ph_num < ph.shape[0] * 10, \
                    "linearization of phi scan causes excessive oversampling {0}/{1}".format(ph_num, ph.shape[0])

            f.write("emission angle phi {ph0:f} {ph1:f} {nph:d}\n".format(ph0=ph_min, ph1=ph_max, nph=ph_num))

            f.write("initial state {0}\n".format(params.initial_state))
            polarizations = {'H': 'LPx', 'V': 'LPy', 'L': 'LCP', 'R': 'RCP'}
            f.write("polarization {0}\n".format(polarizations[params.polarization]))

            scatterers = ["scatterer {at} {fi}\n".format(at=at, fi=fi)
                          for (at, fi) in params.phase_files.items()
                          if os.path.isfile(fi)]
            rme = ["rmat {fi}\n".format(fi=fi)
                   for (at, fi) in params.rme_files.items()
                   if at == params.emitters[0][3] and os.path.isfile(fi)] or \
                  ["rmat inline 1 regular1 {l0} {pv} {pd} {mv} {md}\n".format(l0=params.l_init,
                   pv=params.rme_plus_value, pd=params.rme_plus_shift,
                   mv=params.rme_minus_value, md=params.rme_minus_shift)]
            if scatterers and rme:
                for scat in scatterers:
                    f.write(scat)
                f.write(rme[0])
            else:
                f.write("muffin-tin\n")

            f.write("V0 E(eV) {0:f}\n".format(params.inner_potential))
            f.write("cluster surface l(A) {0:f}\n".format(params.z_surface))
            f.write("imfp SD-UC\n")
            f.write("temperature {0:f} {1:f}\n".format(params.experiment_temperature, params.debye_temperature))
            f.write("iteration recursion\n")
            f.write("dmax l(A) {0:f}\n".format(params.dmax))
            f.write("lmax {0:d}\n".format(params.lmax))
            f.write("orders {0:d} ".format(len(params.orders)))
            f.write(" ".join(format(order, "d") for order in params.orders) + "\n")
            f.write("emission angle window {0:F}\n".format(params.angular_resolution / 2.0))

            # scattering factor output (see project.Params.phase_output_classes)
            if params.phase_output_classes is not None:
                fn = "{0}.clu".format(params.output_file)
                f.write("cluster output l(A) {fn}\n".format(fn=fn))
                files[fn] = "output"
                try:
                    cls = (cl for cl in params.phase_output_classes)
                except TypeError:
                    cls = range(params.phase_output_classes)
                for cl in cls:
                    fn = "{of}.{cl}.scat".format(cl=cl, of=params.output_file)
                    f.write("scan scatterer {cl} phase-shifts {fn}\n".format(cl=cl, fn=fn))
                    files[fn] = "output"

            f.write("scan pd {0}\n".format(params.output_file))
            files[params.output_file] = "output"
            f.write("end\n")

        return files

    def run(self, params, cluster, scan, output_file):
        """
        run EDAC with the given parameters and cluster.

        @param params: a pmsco.project.Params object with all necessary values except cluster and output files set.

        @param cluster: a pmsco.cluster.Cluster(format=FMT_EDAC) object with all atom positions set.

        @param scan: a pmsco.project.Scan() object describing the experimental scanning scheme.

        @param output_file: base name for all intermediate and output files

        @return: (str, dict) result_file, and dictionary of created files {filename: category}
        """

        # set up scan
        params.fixed_cluster = 'a' in scan.mode

        # generate file names
        base_filename = output_file
        clu_filename = base_filename + ".clu"
        out_filename = base_filename + ".out"
        par_filename = base_filename + ".par"
        dat_filename = out_filename
        if params.fixed_cluster:
            etpi_filename = base_filename + ".etpai"
        else:
            etpi_filename = base_filename + ".etpi"

        # fix EDAC particularities
        params.cluster_file = clu_filename
        params.output_file = out_filename
        params.data_file = dat_filename
        params.emitters = cluster.get_emitters(['x', 'y', 'z', 'c'])

        # save parameter files
        logger.debug("writing cluster file %s", clu_filename)
        cluster.save_to_file(clu_filename, fmt=mc.FMT_EDAC)
        logger.debug("writing input file %s", par_filename)
        files = self.write_input_file(params, scan, par_filename)

        # run EDAC
        logger.info("calling EDAC with input file %s", par_filename)
        edac.run_script(par_filename)

        # load results and save in ETPI or ETPAI format
        logger.debug("importing data from output file %s", dat_filename)
        result_etpi = md.load_edac_pd(dat_filename, energy=scan.energies[0] + params.work_function,
            theta=scan.thetas[0], phi=scan.phis[0],
            fixed_cluster=params.fixed_cluster)
        result_etpi['e'] -= params.work_function

        if 't' in scan.mode and 'p' in scan.mode:
            hemi_tpi = scan.raw_data.copy()
            hemi_tpi['i'] = 0.0
            try:
                hemi_tpi['s'] = 0.0
            except ValueError:
                pass
            result_etpi = md.interpolate_hemi_scan(result_etpi, hemi_tpi)

        if params.fixed_cluster:
            expected_shape = max(scan.energies.shape[0], 1) * max(scan.alphas.shape[0], 1)
        else:
            expected_shape = max(scan.energies.shape[0], 1) * max(scan.phis.shape[0], scan.thetas.shape[0], 1)
        if result_etpi.shape[0] != expected_shape:
            logger.warning(BMsg("possible scan length mismatch: EDAC result: {result}, expected: {expected}",
                              result=result_etpi.shape[0], expected=expected_shape))

        logger.debug("save result to file %s", etpi_filename)
        md.save_data(etpi_filename, result_etpi)

        files[clu_filename] = 'input'
        files[par_filename] = 'input'
        files[dat_filename] = 'output'
        files[etpi_filename] = 'region'

        return etpi_filename, files