"""
@package pmsco.edac_calculator
Garcia de Abajo EDAC program interface.

@author Matthias Muntwiler

@copyright (c) 2015 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 division
import os
import logging
import math
import numpy as np
import calculator
import data as md
import cluster as mc
import edac.edac as edac

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)

        TODO: some parameters are still hard-coded.
        """
        with open(filepath, "w") as f:
            f.write("verbose off\n")
            f.write("cluster input %s\n" % (params.cluster_file))
            f.write("emitters %u l(A)\n" % (len(params.emitters)))
            for em in params.emitters:
                f.write("%g %g %g %u\n" % em)
                #for iat in range(params.atom_types):
                #pf = params.phase_file[iat]
                #pf = pf.replace(".pha", ".edac.pha")
                #f.write("scatterer %u %s\n" % (params.atomic_number[iat], pf))

            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, params.azimuthal_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]))
            f.write("muffin-tin\n")
            f.write("V0 E(eV) {0}\n".format(params.inner_potential))
            f.write("cluster surface l(A) {0}\n".format(params.z_surface))
            f.write("imfp SD-UC\n")
            f.write("temperature %g %g\n" % (params.experiment_temperature, params.debye_temperature))
            f.write("iteration recursion\n")
            f.write("dmax l(A) %g\n" % (params.dmax))
            f.write("lmax %u\n" % (params.lmax))
            f.write("orders %u " % (len(params.orders)))
            for order in params.orders:
                f.write("%u " % (order))
            f.write("\n")
            f.write("emission angle window 1\n")
            f.write("scan pd %s\n" % (params.output_file))
            f.write("end\n")

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

        @param params: a msc_param.Params() object with all necessary values except cluster and output files set.

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

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

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

        @return: result_file, files_cats
        """

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

        # 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)
        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 result_etpi.shape[0] != scan.raw_data.shape[0]:
            logger.error("scan length mismatch: EDAC result: %u, scan data: %u", result_etpi.shape[0], scan.raw_data.shape[0])
        logger.debug("save result to file %s", etpi_filename)
        md.save_data(etpi_filename, result_etpi)

        files = {clu_filename: 'input', par_filename: 'input', dat_filename: 'output',
                 etpi_filename: 'region'}
        return etpi_filename, files