#!/usr/bin/env python

"""Handle Monte Carlo simulation run settings"""

from __future__ import absolute_import, division

import pisa.utils.fileio as fileio
import pisa.utils.flavInt as flavInt
from pisa.utils import resources as resources
from pisa.utils.cross_sections import CrossSections

# Following "import *" is intentionally done so that `eval` called in
# translate_source_dict will execute with direct access to numpy namespace
from numpy import * # pylint: disable=wildcard-import, unused-wildcard-import, redefined-builtin

__all__ = ['MCSimRunSettings', 'DetMCSimRunsSettings']

__author__ = 'J.L. Lanfranchi'

__license__ = '''Copyright (c) 2014-2017, The IceCube Collaboration

 Licensed under the Apache License, Version 2.0 (the "License");
 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

 Unless required by applicable law or agreed to in writing, software
 distributed under the License is distributed on an "AS IS" BASIS,
 WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 See the License for the specific language governing permissions and
 limitations under the License.'''

# TODO: make ability to serialize instantiated MCSimRunSettings and
# DetMCSimRunsSettings objects back to JSON format from which they were
# generated

# TODO: put more thought into the MCSimRunSettings class

# TODO: make sure user is forced to choose run & detector here

class MCSimRunSettings(dict):
    """Handle Monte Carlo run settings

    Parameters
    ----------
    run_settings : string or dict
    run
    detector : string or None

    Notes
    -----
    run_settings dictionary format (and same for corresponding JSON file, e.g.
    resources/events/mc_sim_run_settings.json); example is for PINGU but should
    generalize to DeepCore, etc. Note that a JSON file will use null and not
    None.

    (Also note that expressions for numerical values utilizing the Python/numpy
    namespace are valid, e.g. "2*pi" will be evaluated within the code to its
    decimal value.)::

        {

          # Specify the detector name, lower case
          "pingu": {

            # Monte Carlo run number for this detector
            "388": {

              # Version of geometry, lower-case. E.g., ic86 for IceCube/DeepCore
              "geom": "v36",

              # A straightforward way of computing aeff/veff/meff is to keep all
              # simulated events and compare the #analysis/#simulated. If all
              # simulated events are kept, then the filename containing these is
              # recorded here.
              "all_gen_events_file": None,

              # Max and min azimuth angle simulated (rad)
              "azimuth_max": "2*pi",
              "azimuth_min": 0,

              # Max and min energy simulated (GeV)
              "energy_max": 80,
              "energy_min": 1,

              # GENIE simulates some un-physica events (interactions that will not
              # occur in nature). The number below was arrived at by Ken Clark, so
              # ask him for more info.
              "physical_events_fract": 0.8095,

              # GENIE has a prescale factor (TODO: generalize or eliminate for
              # other xsec?)
              "genie_prescale_factor": 1.2,

              # Neutrino flavors simulated
              "flavints": "nutau,nutaubar",

              # #nu / (#nu + #nubar) simulated
              "nu_to_total_fract": 0.5,

              # Number of events simulated per I3 file
              "num_events_per_file": 250000,

              # Number of I3 files used
              "num_i3_files": 195,

              # Simulated spectral inde gamma; value of 1 => E*{-1}
              "sim_spectral_index": 1,

              # Version of neutrino/ice cross sections used for the simulation
              "xsec_version": "genie_2.6.4",

              # Max and min zenith angle simulated (rad)
              "zenith_max": "pi",
              "zenith_min": 0
            }
          }
        }

    """
    def __init__(self, run_settings, run=None, detector=None):
        super().__init__()
        # TODO: clean up this constructor!
        #if isinstance(run_settings, str):
        #    rsd = jsons.from_json(resources.find_resource(run_settings))
        if isinstance(run_settings, dict):
            rsd = run_settings
        else:
            raise TypeError('Unhandled run_settings type passed in arg: %s'
                            %type(run_settings))
        #if detector is not None:
        #    try:
        #        rsd = rsd[detector]
        #    except:
        #        pass
        rsd = self.translate_source_dict(rsd)
        if not detector is None:
            detector = str(detector).strip()
        self.detector = detector
        self.run = run
        self.update(rsd)

    @staticmethod
    def translate_source_dict(d):
        d['tot_gen'] = d['num_events_per_file'] * d['num_i3_files']

        # TODO: does the following logic actually work with both old and new
        # conventions?

        # NOTE: the ',' --> '+' mapping is necessary since some data files
        # were saved prior to the convention that commas exclusively separate
        # groups while plus signs indicate flav/ints grouped together

        d['flavints'] = flavInt.NuFlavIntGroup(d['flavints'].replace(',', '+'))

        # Numeric fields are allowed to be expressions that get evaluated
        numeric_fields = [
            'azimuth_max',
            'azimuth_min',
            'energy_max',
            'energy_min',
            'physical_events_fract',
            'genie_prescale_factor',
            'nu_to_total_fract',
            'num_events_per_file',
            'num_i3_files',
            'sim_spectral_index',
            'zenith_max',
            'zenith_min',
        ]
        for f in numeric_fields:
            if isinstance(d[f], str):
                d[f] = eval(d[f])

        return d

    def consistency_checks(self, data, flav=None):
        # TODO: implement!
        pass

    def barnobarfract(self, barnobar=None, is_particle=None,
                      flav_or_flavint=None):
        """Fraction of events generated (either particles or antiparticles).

        Specifying whether you want the fraction for particle or
        antiparticle is done by specifying one (and only one) of the three
        parameters:
            `barnobar`, `is_particle` or `flav_or_flavint`

        Parameters
        ----------
        barnobar : None or int
            -1 for antiparticle, +1 for particle
        is_particle : None or bool
            True for particle, false for antiparticle
        flav_or_flavint : None or convertible to NuFlav or NuFlavInt
            Particle or antiparticles is determined from the flavor or flavint
            passed

        """
        nargs = sum([(not barnobar is None),
                     (not is_particle is None),
                     (not flav_or_flavint is None)])
        if nargs != 1:
            raise ValueError('One and only one of `barnobar`, `is_particle`,'
                             ' and `flav_or_flavint` must be specified. Got'
                             ' %d non-None args instead.' % nargs)

        if flav_or_flavint is not None:
            is_particle = flavInt.NuFlavInt(flav_or_flavint).particle
        elif barnobar is not None:
            is_particle = barnobar > 0

        if is_particle:
            return self['nu_to_total_fract']
        return 1 - self['nu_to_total_fract']

    def get_num_gen(self, barnobar=None, is_particle=None,
                    flav_or_flavint=None, include_physical_fract=True):
        """Return the number of events generated.

        Parameters
        ----------
        barnobar : None or int
            -1 for antiparticle or +1 for particle

        is_particle : None or bool

        flav_or_flavint : None or convertible to NuFlav or NuFlavInt
            If one of `barnobar`, `is_particle`, or `flav_or_flavint` is
            specified, returns only the number of particles or antiparticles
            generated. Otherwise (if none of those is specified), return the
            total number of generated events.

        include_physical_fract : bool
            Whether to include the "GENIE physical fraction", which accounts
            for events that are generated but are un-physical and therefore
            will never be detectable. These are removed to not penalize
            detection efficiency.

        """
        nargs = sum([(not barnobar is None),
                     (not is_particle is None),
                     (not flav_or_flavint is None)])
        if flav_or_flavint is not None:
            if (flav_or_flavint not in self.get_flavs()
                    and flav_or_flavint not in self.get_flavints()):
                return 0
        barnobarfract = 1
        if nargs > 0:
            barnobarfract = self.barnobarfract(
                barnobar=barnobar, is_particle=is_particle,
                flav_or_flavint=flav_or_flavint
            )
        physical_fract = 1
        if include_physical_fract:
            physical_fract = self['physical_events_fract']
        return self['tot_gen'] * barnobarfract * physical_fract

    def get_flavints(self):
        return self['flavints'].get_flavints()

    def get_flavs(self):
        return self['flavints'].get_flavs()

    def get_energy_range(self):
        """(min, max) energy in GeV"""
        return self['energy_min'], self['energy_max']

    def get_spectral_index(self):
        """Spectral index (positive number for negative powers of energy)"""
        return self['sim_spectral_index']

    def get_xsec_version(self):
        """Cross sectons version name used in generating the MC"""
        return self['xsec_version']

    def get_xsec(self, xsec=None):
        """Instantiated CrossSections object"""
        if xsec is None:
            return CrossSections(ver=self['xsec_version'])
        return CrossSections(ver=self['xsec_version'], xsec=xsec)

class DetMCSimRunsSettings(dict):
    """Handle Monte Carlo run settings for a detector (i.e., without specifying
    which run as is required for the MCSimRunSettings object)

    Since run is not specified at instantiation, method calls require the user
    to specify a run ID.

    Parameters
    ----------
    run_settings : string or dict
    detector : string or None

    See Also
    --------
    MCSimRunSettings : Same, but specifies a specific run at instantiation; see
                       class docstring for specification of run_settings dict /
                       JSON file

    """
    def __init__(self, run_settings, detector=None):
        super().__init__()
        if isinstance(run_settings, str):
            rsd = fileio.from_file(resources.find_resource(run_settings))
        elif isinstance(run_settings, dict):
            rsd = run_settings
        else:
            raise TypeError('Unhandled run_settings type passed in arg: ' +
                            type(run_settings))

        if detector:
            detector = str(detector).strip()
        self.detector = detector

        # Determine how deeply nested runs are in the dict to allow for
        # user to specify a dict that has multiple detectors in it OR
        # a dict with just a single detector in it
        if 'flavints' in rsd.values()[0]:
            runs_d = rsd
        elif 'flavints' in rsd.values()[0].values()[0]:
            if self.detector is None:
                if len(rsd) == 1:
                    runs_d = rsd.values()[0]
                else:
                    raise ValueError('Must specify which detector; detectors '
                                     'found: ' + str(rsd.keys()))
            else:
                runs_d = rsd[self.detector.strip()]
        else:
            raise Exception('dict must either be 3 levels: '
                            '{DET:{RUN:{...}}}; or 2 levels: {RUN:{...}}')

        # Force run numbers to be strings (JSON files cannot have an int as
        # a key, so it is a string upon import, and it's safest to keep it as
        # a string considering how non-standardized naming is in IceCube) and
        # convert actual run settings dict to MCSimRunSettings instances
        runs_d = {str(k): MCSimRunSettings(v) for k, v in runs_d.items()}

        # Save the runs_d to this object instance, which behaves like a dict
        self.update(runs_d)

    def consistency_checks(self, data, run, flav=None):
        pass

    def barnobarfract(self, run, barnobar=None, is_particle=None,
                      flav_or_flavint=None):
        return self[str(run)].barnobarfract(barnobar=barnobar,
                                            is_particle=is_particle,
                                            flav_or_flavint=flav_or_flavint)

    def get_num_gen(self, run, barnobar=None, is_particle=None,
                    flav_or_flavint=None, include_physical_fract=True):
        """Return the total number of events generated"""
        return self[str(run)].get_num_gen(
            barnobar=barnobar, is_particle=is_particle,
            flav_or_flavint=flav_or_flavint,
            include_physical_fract=include_physical_fract
        )

    def get_flavints(self, run):
        return self[str(run)].get_flavints()

    def get_flavs(self, run):
        return self[str(run)].get_flavs()

    def get_energy_range(self, run):
        """(min, max) energy in GeV"""
        return self[str(run)].get_energy_range()

    def get_spectral_index(self, run):
        """Spectral index (positive number for negative powers of energy)"""
        return self[str(run)].get_spectral_index()

    def get_xsec_version(self, run):
        """Cross sectons version name used in generating the MC"""
        return self[str(run)].get_xsec_version()

    def get_xsec(self, run, xsec=None):
        """Instantiated CrossSections object"""
        return self[str(run)].get_xsec(xsec)


"""`BVT24`"""
import numpy as np
from skimage.transform import SimilarityTransform

from .base import ProsthesisSystem
from .electrodes import DiskElectrode
from .electrode_arrays import ElectrodeArray

class BVT24(ProsthesisSystem):
    """24-channel suprachoroidal retinal prosthesis

    This class creates a 24-channel suprachoroidal retinal prosthesis
    [Layton2014]_, which was developed by the Bionic Vision Australia
    Consortium and commercialized by Bionic Vision Technologies (BVT).
    The center of the array is located at (x,y,z), given in microns, and the
    array is rotated by rotation angle ``rot``, given in degrees.

    The array consists of:

    -   33 platinum stimulating electrodes:

        -   30 electrodes with 600um diameter (Electrodes C1-20 (except
            C9, C17, C19) and Electrodes C21a-m),

        -   3 electrodes with 400um diameter (Electrodes C9, C17, C19)

    -   2 return electrodes with 2000um diameter (Electrodes R1, R2)

    Electrodes C21a-m are typically being ganged to provide an external
    ring for common ground. The center of the array is assumed to lie
    between Electrodes C7, C8, C9, and C13.

    .. note::

        Column order for electrode numbering is reversed in a left-eye
        implant.

    .. versionadded:: 0.6

    Parameters
    ----------
    x/y/z : double
        3D location of the center of the electrode array.
        The coordinate system is centered over the fovea.
        Positive ``x`` values move the electrode into the nasal retina.
        Positive ``y`` values move the electrode into the superior retina.
        Positive ``z`` values move the electrode away from the retina into the
        vitreous humor (sometimes called electrode-retina distance).
        ``z`` can either be a list with 35 entries or a scalar that is applied
        to all electrodes.
    rot : float
        Rotation angle of the array (deg). Positive values denote
        counter-clock-wise (CCW) rotations in the retinal coordinate
        system.
    eye : {'RE', 'LE'}, optional
        Eye in which array is implanted.
    preprocess : bool or callable, optional
        Either True/False to indicate whether to execute the implant's default
        preprocessing method whenever a new stimulus is assigned, or a custom
        function (callable).
    safe_mode : bool, optional
        If safe mode is enabled, only charge-balanced stimuli are allowed.

    """
    # Frozen class: User cannot add more class attributes
    __slots__ = ()

    def __init__(self, x=0, y=0, z=0, rot=0, eye='RE', stim=None,
                 preprocess=False, safe_mode=False):
        self.eye = eye
        self.preprocess = preprocess
        self.safe_mode = safe_mode
        self.earray = ElectrodeArray([])
        n_elecs = 35

        # the positions of the electrodes 1-20, 21a-21m, R1-R2
        x_arr = np.array([1275.0, 850.0, 1275.0, 850.0, 1275.0,
                          425.0, 0, 425.0, 0, 425.0,
                          -425.0, -850.0, -425.0, -850.0, -425.0,
                          -1275.0, -1700.0, -1275.0, -1700.0, -1275.0,
                          850.0, 0, -850.0, -1700.0, -2125.0,
                          -2550.0, -2125.0, -2550.0, -2125.0, -1700.0,
                          -850.0, 0, 850.0, -7000.0, -9370.0])
        y_arr = np.array([1520.0, 760.0, 0, -760.0, -1520.0,
                          1520.0, 760.0, 0, -760.0, -1520.0,
                          1520.0, 760.0, 0, -760.0, -1520.0,
                          1520.0, 760.0, 0, -760.0, -1520.0,
                          2280.0, 2280.0, 2280.0, 2280.0, 1520.0,
                          760.0, 0.0, -760.0, -1520.0, -2280.0,
                          -2280.0, -2280.0, -2280.0, 0, 0])
        if isinstance(z, (list, np.ndarray)):
            # Specify different height for every electrode in a list:
            z_arr = np.asarray(self.z).flatten()
            if z_arr.size != n_elecs:
                raise ValueError("If `z` is a list, it must have %d entries, "
                                 "not %d." % (n_elecs, len(z)))
        else:
            # If `z` is a scalar, choose same height for all electrodes:
            z_arr = np.ones(n_elecs, dtype=float) * z

        # the position of the electrodes 1-20, 21a-21m, R1-R2 for left eye
        if eye == 'LE':
            x_arr = np.negative(x_arr)

        # the radius of all the electrodes in the implants
        r_arr = [300.0] * n_elecs
        # the radius of electrodes 9, 17, 19 is 200.0 um
        r_arr[8] = r_arr[16] = r_arr[18] = 200.0
        # the radius of the return electrodes is 1000.0 um
        r_arr[33] = r_arr[34] = 1000.0
        # the names of the electrodes C1-20, C21a-21m, R1 and R2
        names = ["C%s" % name for name in range(1, 21)]
        names.extend(['C21a', 'C21b', 'C21c', 'C21d', 'C21e',
                      'C21f', 'C21g', 'C21h', 'C21i', 'C21j',
                      'C21k', 'C21l', 'C21m'])
        names.extend(['R1', 'R2'])

        # Rotate the grid and center at (x,y):
        tf = SimilarityTransform(rotation=np.deg2rad(rot), translation=[x, y])
        x_arr, y_arr = tf(np.vstack([x_arr.ravel(), y_arr.ravel()]).T).T

        for x, y, z, r, name in zip(x_arr, y_arr, z_arr, r_arr, names):
            self.earray.add_electrode(name, DiskElectrode(x, y, z, r))

        # Beware of race condition: Stim must be set last, because it requires
        # indexing into self.electrodes:
        self.stim = stim


# Fit to multiple relative binding free energy edges
import sys
from argparse import ArgumentParser
import jax
from jax import numpy as jnp
import numpy as np
import datetime
import timemachine
from training.dataset import Dataset

# forcefield handlers
from ff import Forcefield
from ff.handlers.serialize import serialize_handlers
from ff.handlers.deserialize import deserialize_handlers
from ff.handlers.nonbonded import AM1CCCHandler, LennardJonesHandler

# free energy classes
from fe.free_energy import (
    RelativeFreeEnergy,
    construct_lambda_schedule,
    RBFETransformIndex,
)
from fe.estimator import SimulationResult
from fe.model import RBFEModel
from fe.loss import pseudo_huber_loss  # , l1_loss, flat_bottom_loss

# MD initialization
from md import builders

# parallelization across multiple GPUs
from parallel.client import CUDAPoolClient, GRPCClient
from parallel.utils import get_gpu_count

from collections import namedtuple, defaultdict

from pickle import load, dump

from optimize.step import truncated_step
from optimize.utils import flatten_and_unflatten

from typing import Tuple, Dict, List, Union, Any

from pathlib import Path
from time import time

array = Union[np.array, jnp.array]

Handler = Union[AM1CCCHandler, LennardJonesHandler]  # TODO: relax this assumption

NUM_GPUS = get_gpu_count()

# how much MD to run
# TODO: rename this to something more descriptive than "Configuration"...
#   want to distinguish later between an "RBFE configuration"
#       (which describes the computation of a single edge)
#   and a "training configuration"
#       (which describes the overall training loop)
Configuration = namedtuple(
    "Configuration", ["num_complex_windows", "num_solvent_windows", "num_equil_steps", "num_prod_steps"]
)

# define a couple configurations: one for quick tests, and one for production
production_configuration = Configuration(
    num_complex_windows=60,
    num_solvent_windows=60,
    num_equil_steps=10000,
    num_prod_steps=100000,
)

intermediate_configuration = Configuration(
    num_complex_windows=60,
    num_solvent_windows=60,
    num_equil_steps=10000,
    num_prod_steps=10000,
)

testing_configuration = Configuration(
    num_complex_windows=10,
    num_solvent_windows=10,
    num_equil_steps=1000,
    num_prod_steps=1000,
)

# locations relative to project root
root = Path(timemachine.__file__).parent.parent

def _save_forcefield(fname, ff_params):
    with open(fname, "w") as fh:
        fh.write(ff_params)

def _compute_step_lower_bound(loss: float, blown_up: bool, relative_improvement_bound: float = 0.8) -> float:
    """problem this addresses: on a small fraction of steps, the free energy estimate may be grossly unreliable
    away from target, typically indicating an instability was encountered.
    detect if this occurs, and don't allow a step.

    """
    if not blown_up:
        return loss * relative_improvement_bound
    else:
        return loss  # don't move!

def _results_to_arrays(results: List[SimulationResult]):
    """each result object was constructed by SimulationResult(xs=xs, du_dls=full_du_dls, du_dps=grads)

    for each field, concatenate into an array
    """

    xs = np.array([r.xs for r in results])
    du_dls = np.array([r.du_dls for r in results])
    # without dtype=object get a warning about ragged arrays (inconsistent size/type)
    du_dps = np.array([r.du_dps for r in results], dtype=object)

    return xs, du_dls, du_dps

def _blew_up(results: List[SimulationResult]) -> bool:
    """if stddev(du_dls) for any window exceeded 1000 kJ/mol, don't trust result enough to take a step
    if du_dls contains any nans, don't trust result enough to take a step"""
    du_dls = _results_to_arrays(results)[1]

    # TODO: adjust this threshold a bit, move reliability calculations into fe/estimator.py or fe/model.py
    return np.isnan(du_dls).any() or (du_dls.std(1).max() > 1000)

def loss_fxn(ff_params, batch: List[Tuple[RelativeFreeEnergy, RBFEModel]]):
    index = RBFETransformIndex()
    index.build([edge[0] for edge in batch])
    indices = []
    all_results = []
    preds = []
    for rfe, model in batch:
        indices.append(list(index.get_transform_indices(rfe)))
        pred_ddG, stage_results = model.predict(ff_params, rfe.mol_a, rfe.mol_b, rfe.core)
        all_results.extend(list(stage_results))
        preds.append(pred_ddG)
    labels = jnp.asarray([rfe.label for rfe, _ in batch])
    loss = pseudo_huber_loss(jnp.asarray(preds) - labels)
    # Aggregate the pseudo huber loss using mean
    loss = jnp.mean(loss)
    return loss, (preds, all_results)

def run_validation_edges(validation: Dataset, params, systems, epoch, inference: bool = False):
    if len(validation) <= 0:
        return
    message_prefix = "Validation"
    if inference:
        message_prefix = "Inference"
    val_loss = np.zeros(len(validation))
    for i, rfe in enumerate(validation.data):
        if getattr(rfe, "complex_path", None) is not None:
            model = systems[rfe.complex_path]
        else:
            model = systems[protein_path]
        start = time()

        loss, (preds, stage_results) = loss_fxn(params, [(rfe, model)])
        elapsed = time() - start
        print(f"{message_prefix} edge {i}: time={elapsed:.2f}s, loss={loss:.2f}")
        du_dls_dict = {stage: _results_to_arrays(results)[1] for stage, results in stage_results}
        np.savez(output_path.joinpath(f"validation_du_dls_snapshot_{epoch}_{i}.npz"), **du_dls_dict)
        val_loss[i] = loss
    np.savez(output_path.joinpath(f"validation_edge_losses_{epoch}.npz"), loss=val_loss)

def equilibrate_edges(datasets: List[Dataset], systems: List[Dict[str, Any]], num_steps: int, cache_path: str):
    model_set = defaultdict(list)
    for dataset in datasets:
        for rfe in dataset.data:
            if getattr(rfe, "complex_path", None) is not None:
                model_set[rfe.complex_path].append(rfe)
            else:
                model_set[protein_path].append(rfe)
    for path, edges in model_set.items():
        model = systems[path]
        model.equilibrate_edges(
            [(edge.mol_a, edge.mol_b, edge.core) for edge in edges],
            equilibration_steps=num_steps,
            cache_path=cache_path,
        )

if __name__ == "__main__":
    default_output_path = f"results_{datetime.datetime.now().strftime('%Y_%m_%d_%H_%M_%S')}"
    parser = ArgumentParser(description="Fit Forcefield parameters to hif2a")
    parser.add_argument(
        "--num_gpus",
        default=None,
        type=int,
        help=f"Number of GPUs to run against, defaults to {NUM_GPUS} if no hosts provided",
    )
    parser.add_argument("--hosts", nargs="*", default=None, help="Hosts running GRPC worker to use for compute")
    parser.add_argument("--param_updates", default=1000, type=int, help="Number of updates for parameters")
    parser.add_argument("--seed", default=2021, type=int, help="Seed for shuffling ordering of transformations")
    parser.add_argument("--config", default="intermediate", choices=["intermediate", "production", "test"])
    parser.add_argument("--batch_size", default=1, type=int, help="Number of items to batch together for training")

    parser.add_argument("--path_to_ff", default=str(root.joinpath("ff/params/smirnoff_1_1_0_ccc.py")))
    parser.add_argument(
        "--path_to_edges",
        default=["relative_transformations.pkl"],
        nargs="+",
        help="Path to pickle file containing list of RelativeFreeEnergy objects",
    )
    parser.add_argument("--split", action="store_true", help="Split edges into train and validation set")
    parser.add_argument(
        "--pre_equil",
        default=None,
        help="Number of pre equilibration steps or path to cached equilibrated edges, if not provided no pre equilibration performed",
    )
    parser.add_argument("--hmr", action="store_true", help="Enable HMR")
    parser.add_argument("--output_path", default=default_output_path, help="Path to output directory")
    parser.add_argument("--protein_path", default=None, help="Path to protein if edges don't provide protein")
    parser.add_argument(
        "--inference_only", action="store_true", help="Disable training, run all edges as validation edges"
    )
    # TODO: also make configurable: forces_to_refit, optimizer params, path_to_protein, path_to_protein_ff, ...
    args = parser.parse_args()
    protein_path = None
    if args.protein_path:
        prot_path = Path(args.protein_path).expanduser()
        protein_path = prot_path.as_posix()
        if not prot_path.is_file():
            print(f"Unable to find path: {protein_path}")
            sys.exit(1)

    # xor num_gpus and hosts args
    if args.num_gpus is not None and args.hosts is not None:
        print("Unable to provide --num-gpus and --hosts together")
        sys.exit(1)

    # which force field components we'll refit
    forces_to_refit = [AM1CCCHandler, LennardJonesHandler]

    # how much computation to spend per refitting step
    configuration = None
    if args.config == "intermediate":  # goldilocks
        configuration = intermediate_configuration
    elif args.config == "test":
        configuration = testing_configuration  # a little
    elif args.config == "production":
        configuration = production_configuration  # a lot
    assert configuration is not None, "No configuration provided"

    if not args.hosts:
        num_gpus = args.num_gpus
        if num_gpus is None:
            num_gpus = NUM_GPUS
        # set up multi-GPU client
        client = CUDAPoolClient(max_workers=num_gpus)
    else:
        # Setup GRPC client
        client = GRPCClient(hosts=args.hosts)
    client.verify()

    # load and construct forcefield
    with open(args.path_to_ff) as f:
        ff_handlers = deserialize_handlers(f.read())

    forcefield = Forcefield(ff_handlers)

    relative_transformations: List[RelativeFreeEnergy] = []
    # load pre-defined collection of relative transformations
    for edge_path in args.path_to_edges:
        with open(edge_path, "rb") as f:
            relative_transformations.extend(load(f))

    # if older transformation lack a complex_path, rely on --protein_path to set
    protein_paths = set(x.complex_path for x in relative_transformations if hasattr(x, "complex_path"))
    if protein_path is not None:
        protein_paths.add(protein_path)
    if len(protein_paths) == 0:
        print("No proteins provided by edges or with --protein_path")
        sys.exit(1)

    # create path if it doesn't exist
    output_path = Path(args.output_path)
    output_path.mkdir(parents=True, exist_ok=True)
    print(f"Storing results in {output_path}")

    dataset = Dataset(relative_transformations)
    if not args.inference_only:
        if args.split:
            # TODO: More physically meaningful split
            # 80, 20 split on transformations
            training, validation = dataset.random_split(0.8)
        else:
            validation = Dataset([])
            training = dataset
    else:
        validation = dataset
        training = Dataset([])

    with open(output_path.joinpath("training_edges.pk"), "wb") as ofs:
        dump(training.data, ofs)
    if len(validation):
        with open(output_path.joinpath("validation_edges.pk"), "wb") as ofs:
            dump(validation.data, ofs)

    # Build all of the different protein systems
    systems = {}
    for prot_path in protein_paths:
        # build the complex system
        # note: "complex" means "protein + solvent"
        complex_system, complex_coords, _, _, complex_box, _ = builders.build_protein_system(prot_path)

        # build the water system
        solvent_system, solvent_coords, solvent_box, _ = builders.build_water_system(4.0)

        systems[prot_path] = RBFEModel(
            client=client,
            ff=forcefield,
            complex_system=complex_system,
            complex_coords=complex_coords,
            complex_box=complex_box,
            complex_schedule=construct_lambda_schedule(configuration.num_complex_windows),
            solvent_system=solvent_system,
            solvent_coords=solvent_coords,
            solvent_box=solvent_box,
            solvent_schedule=construct_lambda_schedule(configuration.num_solvent_windows),
            equil_steps=configuration.num_equil_steps,
            prod_steps=configuration.num_prod_steps,
            pre_equilibrate=args.pre_equil is not None,
            hmr=args.hmr,
        )

    # TODO: how to get intermediate results from the computational pipeline encapsulated in binding_model.loss ?
    #   e.g. stage_results, and further diagnostic information
    #   * x trajectories,
    #   * d U / d parameters trajectories,
    #   * matrix of U(x; lambda) for all x, lambda
    #   * the deltaG pred
    #   (proper way is probably something like has_aux=True https://jax.readthedocs.io/en/latest/jax.html#jax.value_and_grad)

    ordered_params = forcefield.get_ordered_params()
    ordered_handles = forcefield.get_ordered_handles()

    # compute and save the sequence of relative_transformation indices
    num_epochs = int(np.ceil(args.param_updates / len(relative_transformations)))
    np.random.seed(args.seed)

    batch_size = args.batch_size
    step_inds = []
    for epoch in range(num_epochs):
        inds = np.arange(len(training.data))
        np.random.shuffle(inds)
        batched_inds = []
        num_steps = (len(inds) + batch_size - 1) // batch_size
        for i in range(num_steps):
            offset = i * batch_size
            batched_inds.append(inds[offset : offset + batch_size])
        step_inds.append(np.asarray(batched_inds, dtype=object))

    np.save(output_path.joinpath("step_indices.npy"), np.hstack(step_inds)[: args.param_updates])

    pre_equil = args.pre_equil
    if pre_equil is not None:
        steps = 0
        cache_path = output_path.joinpath("equilibration_cache.pkl")
        if pre_equil.isdigit():
            steps = int(pre_equil)
        elif Path(pre_equil).is_file():
            cache_path = pre_equil
        else:
            print(f"Must provide either an integer or a valid path for --pre_equil, got {pre_equil}")
            sys.exit(1)

        equilibrate_edges([training, validation], systems, steps, cache_path)

    flatten, unflatten = flatten_and_unflatten(ordered_params)

    step = 0
    # in each optimizer step, look at one transformation from relative_transformations
    for epoch in range(num_epochs):
        # Run Validation edges at start of epoch. Unlike NNs we have a reasonable starting
        # point that is worth knowing
        run_validation_edges(validation, ordered_params, systems, epoch + 1, inference=args.inference_only)
        print(f"Epoch: {epoch+1}/{num_epochs}")
        for batch in step_inds[epoch]:
            batch_data = []
            for i in batch:
                rfe = training.data[i]
                if getattr(rfe, "complex_path", None):
                    model = systems[rfe.complex_path]
                else:
                    model = systems[protein_path]
                batch_data.append((rfe, model))
            # compute a batch, measuring total wall-time
            t0 = time()

            (loss, (predictions, stage_results)), loss_grads = jax.value_and_grad(loss_fxn, argnums=0, has_aux=True)(
                ordered_params, batch_data
            )

            results_this_step = {stage: result for stage, result in stage_results}

            print(f"at optimizer step {step}, loss={loss:.3f}")

            # check if it's probably okay to take an optimizer step on the basis of this result
            # TODO: move responsibility for returning error flags / simulation uncertainty estimates further upstream
            blown_up = False
            for stage, results in results_this_step.items():
                if _blew_up(results):
                    blown_up = True
                    print(f"step {step} blew up in {stage} stage")

            flat_loss_grad = flatten(loss_grads)
            flat_theta = flatten(ordered_params)

            # based on current estimate of (loss, grad, and simulation stability), return a conservative step to take in parameter space
            theta_increment = truncated_step(
                flat_theta, loss, flat_loss_grad, step_lower_bound=_compute_step_lower_bound(loss, blown_up)
            )
            param_increments = unflatten(theta_increment)

            # for any parameter handler types being updated, update in place
            for handle, increment in zip(ordered_handles, param_increments):
                handle_type = type(handle)
                if handle_type in forces_to_refit:

                    # TODO: careful -- this must be a "+=" or "-=" not an "="!
                    handle.params += increment

                    nonzero_increments = increment[increment != 0]
                    min_update = 0.0
                    max_update = 0.0
                    if len(nonzero_increments):
                        min_update = np.min(nonzero_increments)
                        max_update = np.max(nonzero_increments)
                    # TODO: replace with a function that knows what to report about each handle type
                    print(
                        f"updated {len(nonzero_increments)} {handle_type.__name__} params by between {min_update:.4f} and {max_update:.4f}"
                    )

            t1 = time()
            elapsed = t1 - t0

            print(f"completed forcefield-updating step {step} in {elapsed:.3f} s !")

            # save du_dls snapshot
            path_to_du_dls = output_path.joinpath(f"du_dls_snapshot_{step}.npz")
            print(f"saving du_dl trajs to {path_to_du_dls}")
            du_dls_dict = dict()  # keywords here must be strings
            for stage, results in results_this_step.items():
                du_dls_dict[stage] = _results_to_arrays(results)[1]
            np.savez(path_to_du_dls, **du_dls_dict)

            # also save information about this step's parameter gradient and parameter update
            # results to npz
            path_to_npz = output_path.joinpath(f"theta_grad_loss_snapshot_{step}.npz")
            print(f"saving theta, grad, loss snapshot to {path_to_npz}")
            np.savez(path_to_npz, theta=np.array(flat_theta), grad=np.array(flat_loss_grad), loss=loss)

            # TODO: same for xs and du_dps snapshots

            # save updated forcefield .py files after every gradient step
            step_params = serialize_handlers(ff_handlers)
            # TODO: consider if there's a more modular way to keep track of ff updates
            _save_forcefield(output_path.joinpath(f"forcefield_checkpoint_{step}.py"), ff_params=step_params)
            step += 1
            if step >= args.param_updates:
                break
    if not args.inference_only:
        run_validation_edges(validation, ordered_params, systems, epoch + 1)


import numpy as np
from .weighted_yield import SSP, lifetime_Argast, lifetime_Raiteri
from .imf import IMF
from .yields import SN2_feedback, AGB_feedback, SN1a_feedback, Hypernova_feedback
import copy

class SSP_wrap():
	'''
	This is the wrapper around the SSP function. It preloads the needed classes and calls all nucleosynthetic enrichment processes when the enrichment is calculated.
	'''
	def __init__(self, a):
		'''
		Upon initialization the default IMF, CC-SN yields, SN Ia yields and AGB yields is loaded.

		INPUT:

		   a = Modelparameter class. So the default IMF etc are loaded. If we want other yield sets etc. loaded we need to specify that in paramter.py
		'''

		## loading the IMF and the yieldsets prescribed in a (containing all the model parameters)
		basic_imf = IMF(a.mmin,a.mmax,a.mass_steps)
		getattr(basic_imf, a.imf_type_name)(a.imf_parameter)
		basic_sn2 = SN2_feedback()
		getattr(basic_sn2, a.yield_table_name_sn2)()
		basic_1a = SN1a_feedback()
		getattr(basic_1a, a.yield_table_name_1a)()
		basic_agb = AGB_feedback()
		getattr(basic_agb, a.yield_table_name_agb)()
		### mixing of Nomoto CC-SN and HN yields
		if a.yield_table_name_sn2 == 'Nomoto2013':
			basic_hn = Hypernova_feedback()
			getattr(basic_hn, a.yield_table_name_hn)()
			for item in basic_sn2.metallicities:
				x = np.copy(basic_sn2.table[item])
				y = np.copy(basic_hn.table[item])
				for jtem in basic_hn.masses:
					basic_sn2.table[item]['mass_in_remnants'][np.where(basic_sn2.table[item]['Mass']==jtem)] = a.sn2_to_hn * (x['mass_in_remnants'][np.where(x['Mass']==jtem)]) + (1-a.sn2_to_hn) * (y['mass_in_remnants'][np.where(y['Mass']==jtem)])
					basic_sn2.table[item]['unprocessed_mass_in_winds'][np.where(basic_sn2.table[item]['Mass']==jtem)] = a.sn2_to_hn * (x['unprocessed_mass_in_winds'][np.where(x['Mass']==jtem)]) + (1-a.sn2_to_hn) * (y['unprocessed_mass_in_winds'][np.where(y['Mass']==jtem)])
					hn_mass = []
					sn_mass = []
					for stem in basic_sn2.elements:
						sn_mass.append(x[stem][np.where(x['Mass']==jtem)])
						hn_mass.append(y[stem][np.where(y['Mass']==jtem)])
						basic_sn2.table[item][stem][np.where(basic_sn2.table[item]['Mass']==jtem)]= a.sn2_to_hn * (x[stem][np.where(x['Mass']==jtem)]) + (1-a.sn2_to_hn) * (y[stem][np.where(y['Mass']==jtem)])
		## to pass the information on to the feedback calculation
		self.a = a
		self.imf = basic_imf
		self.sn2 = basic_sn2
		self.sn1a = basic_1a
		self.agb = basic_agb

	def calculate_feedback(self, z, elements, element_fractions, time_steps, ssp_mass):
		'''
		The feedback is calculated for the initializes SSP.

		INPUT:

		   z = metallicity of the SSP in mass fraction (not normed to solar!)

		   elements = which elements to follow

		   element_fractions = the birth material of the SSP in the same order as 'elements'

		   time_steps = the time-steps for which the enrichment of the SSP should be calculated (usually the time-steps until the end of the chempy simulation)
		'''
		if self.a.stochastic_IMF:
			imf_copy = copy.copy(self.imf)
			imf_copy.stochastic_sampling(ssp_mass)
		else:
			imf_copy = copy.copy(self.imf)
		basic_ssp = SSP(False, float(z), imf_copy.x, imf_copy.dm, imf_copy.dn, np.copy(time_steps), list(elements), str(self.a.stellar_lifetimes), str(self.a.interpolation_scheme), bool(self.a.only_net_yields_in_process_tables))
		basic_ssp.sn2_feedback(list(self.sn2.elements), dict(self.sn2.table), np.copy(self.sn2.metallicities), float(self.a.sn2mmin), float(self.a.sn2mmax),list(element_fractions))
		basic_ssp.agb_feedback(list(self.agb.elements), dict(self.agb.table), list(self.agb.metallicities), float(self.a.agbmmin), float(self.a.agbmmax),np.hstack(element_fractions))
		basic_ssp.sn1a_feedback(list(self.sn1a.elements), list(self.sn1a.metallicities), dict(self.sn1a.table), str(self.a.time_delay_functional_form), float(self.a.sn1ammin), float(self.a.sn1ammax), self.a.sn1a_parameter, ssp_mass, bool(self.a.stochastic_IMF))
		basic_ssp.bh_feedback(float(self.a.bhmmin),float(self.a.bhmmax),list(elements), np.hstack(element_fractions) , float(self.a.percentage_of_bh_mass))
		# exposing these tables to the outside wrapper
		self.table = basic_ssp.table
		self.sn2_table = basic_ssp.sn2_table
		self.agb_table = basic_ssp.agb_table
		self.sn1a_table = basic_ssp.sn1a_table
		self.bh_table = basic_ssp.bh_table
		self.inverse_imf = basic_ssp.inverse_imf

def initialise_stuff(a):
	'''
	Convenience function initialising the solar abundance, SFR and infall with the default values provided in parameter.py as a
	'''
	from .solar_abundance import solar_abundances
	from .sfr import SFR
	from .infall import INFALL

	basic_solar = solar_abundances()
	getattr(basic_solar, a.solar_abundance_name)()

	basic_sfr = SFR(a.start,a.end,a.time_steps)
	if a.basic_sfr_name == 'gamma_function':
		getattr(basic_sfr, a.basic_sfr_name)(S0 = a.S_0 * a.mass_factor,a_parameter = a.a_parameter, loc = a.sfr_beginning, scale = a.sfr_scale)
	elif a.basic_sfr_name == 'model_A':
		basic_sfr.model_A(a.mass_factor*a.S_0,a.t_0,a.t_1)
	elif a.basic_sfr_name == 'prescribed':
		basic_sfr.prescribed(a.mass_factor, a.name_of_file)
	elif a.basic_sfr_name == 'doubly_peaked':
		basic_sfr.doubly_peaked(S0 = a.mass_factor*a.S_0, peak_ratio = a.peak_ratio, decay = a.sfr_decay, t0 = a.sfr_t0, peak1t0 = a.peak1t0, peak1sigma = a.peak1sigma)
	elif a.basic_sfr_name == 'normal':
		basic_sfr.normal(S0=a.mass_factor*a.S_0, loc=a.sfr_peak, scale=a.sfr_scale)
	elif a.basic_sfr_name == 'step':
		basic_sfr.step(S0=a.mass_factor*a.S_0, loc=a.sfr_cutoff)
	elif a.basic_sfr_name == 'non_parametric':
		basic_sfr.non_parametric(S0=a.mass_factor*a.S_0, breaks=a.sfr_breaks, weights=a.sfr_weights)
	basic_sfr.sfr = a.total_mass * np.divide(basic_sfr.sfr, sum(basic_sfr.sfr))
	basic_infall = INFALL(np.copy(basic_sfr.t), np.copy(basic_sfr.sfr))
	if a.basic_infall_name == 'exponential':
		getattr(basic_infall, a.basic_infall_name)((a.infall_amplitude,a.tau_infall,a.infall_time_offset,a.c_infall,a.norm_infall))
	elif a.basic_infall_name == 'gamma_function':
		getattr(basic_infall, a.basic_infall_name)(mass_factor = a.norm_infall, a_parameter = a.infall_a_parameter, loc = a.infall_beginning, scale = a.infall_scale)
	elif a.basic_infall_name == 'sfr_related':
		getattr(basic_infall, a.basic_infall_name)()

	return basic_solar, basic_sfr, basic_infall

def Chempy(a):
	'''
	Chemical evolution run with the default parameters using the net yields.

	INPUT:

	   a = ModelParameters() from parameter.py

	OUTPUT:

	   cube = The ISM evolution class

	   abundances = The abundances of the ISM
	'''
	from .infall import PRIMORDIAL_INFALL
	from .time_integration import ABUNDANCE_MATRIX
	from .making_abundances import mass_fraction_to_abundances
	from numpy.lib.recfunctions import append_fields
	basic_solar, basic_sfr, basic_infall = initialise_stuff(a)
	elements_to_trace = a.elements_to_trace
	basic_primordial = PRIMORDIAL_INFALL(list(elements_to_trace),np.copy(basic_solar.table))
	basic_primordial.primordial()
	# Needed a rescaling for the shortened sfr
	gas_reservoir_mass_factor = a.gas_reservoir_mass_factor / a.shortened_sfr_rescaling
	#sfr_factor_for_cosmic_accretion	= a.sfr_factor_for_cosmic_accretion / a.shortened_sfr_rescaling
	#gas_at_start = a.gas_at_start / a.shortened_sfr_rescaling
	cube = ABUNDANCE_MATRIX(np.copy(basic_sfr.t),np.copy(basic_sfr.sfr),np.copy(basic_infall.infall),list(elements_to_trace),
list(basic_primordial.symbols),list(basic_primordial.fractions),float(a.gas_at_start),list(basic_primordial.symbols),list(basic_primordial.fractions),
float(gas_reservoir_mass_factor),float(a.outflow_feedback_fraction),bool(a.check_processes),float(a.starformation_efficiency),float(a.gas_power),
 float(a.sfr_factor_for_cosmic_accretion), list(basic_primordial.symbols), list(basic_primordial.fractions))
	basic_ssp = SSP_wrap(a)
	for i in range(len(basic_sfr.t)-1):
		ssp_mass = float(basic_sfr.sfr[i])
		#print(ssp_mass)
		j = len(basic_sfr.t)-i
		element_fractions = []
		for item in elements_to_trace:
			element_fractions.append(float(np.copy(cube.cube[item][max(i-1,0)]/cube.cube['gas'][max(i-1,0)])))## gas element fractions from one time step before
		metallicity = float(cube.cube['Z'][i])
		time_steps = np.copy(basic_sfr.t[:j])
		basic_ssp.calculate_feedback(float(metallicity), list(elements_to_trace), list(element_fractions), np.copy(time_steps), ssp_mass)
		cube.advance_one_step(i+1,np.copy(basic_ssp.table),np.copy(basic_ssp.sn2_table),np.copy(basic_ssp.agb_table),np.copy(basic_ssp.sn1a_table),np.copy(basic_ssp.bh_table))
		if cube.cube['gas'][i] < 0:
			print(i, basic_sfr.t[i])
			print('gas became negative. returning -inf')
			return -np.inf, [0]
		if cube.gas_reservoir['gas'][i] < 0:
			print('gas_reservoir became negative. returning -inf')
			return -np.inf, [0]

	abundances,elements,numbers = mass_fraction_to_abundances(np.copy(cube.cube),np.copy(basic_solar.table))
	weights = cube.cube['sfr']
	abundances = append_fields(abundances,'weights',weights)
	abundances = append_fields(abundances,'time', cube.cube['time'])
	abundances = np.array(abundances)

	return cube, abundances

def Chempy_gross(a):
	'''
	Chemical evolution run with the default parameters but now using solar scaled material (testing the worse case when total yields provided).

	INPUT:

	   a = ModelParameters() from parameter.py

	OUTPUT:

	   cube = The ISM evolution class

	   abundances = The abundances of the ISM
	'''
	from infall import PRIMORDIAL_INFALL
	from time_integration import ABUNDANCE_MATRIX
	from making_abundances import mass_fraction_to_abundances
	from numpy.lib.recfunctions import append_fields
	basic_solar, basic_sfr, basic_infall = initialise_stuff(a)
	elements_to_trace = a.elements_to_trace
	basic_primordial = PRIMORDIAL_INFALL(list(elements_to_trace),np.copy(basic_solar.table))
	basic_primordial.primordial(0)
	gas_reservoir_mass_factor = a.gas_reservoir_mass_factor / a.shortened_sfr_rescaling
	cube = ABUNDANCE_MATRIX(np.copy(basic_sfr.t),np.copy(basic_sfr.sfr),np.copy(basic_infall.infall),list(elements_to_trace),list(basic_primordial.symbols),
list(basic_primordial.fractions),float(a.gas_at_start),list(basic_primordial.symbols),list(basic_primordial.fractions),float(gas_reservoir_mass_factor),
float(a.outflow_feedback_fraction),bool(a.check_processes),float(a.starformation_efficiency),float(a.gas_power), float(a.sfr_factor_for_cosmic_accretion),
list(basic_primordial.symbols), list(basic_primordial.fractions))
	basic_ssp = SSP_wrap(a)
	for i in range(len(basic_sfr.t)-1):
		j = len(basic_sfr.t)-i
		metallicity = float(cube.cube['Z'][i])
		solar_scaled_material = PRIMORDIAL_INFALL(list(elements_to_trace),np.copy(basic_solar.table))
		solar_scaled_material.solar(np.log10(metallicity/basic_solar.z))
		element_fractions = list(solar_scaled_material.fractions)
		for item in elements_to_trace:
			element_fractions.append(float(np.copy(cube.cube[item][max(i-1,0)]/cube.cube['gas'][max(i-1,0)])))## gas element fractions from one time step before
		time_steps = np.copy(basic_sfr.t[:j])
		basic_ssp.calculate_feedback(float(metallicity), list(elements_to_trace), list(element_fractions), np.copy(time_steps))
		cube.advance_one_step(i+1,np.copy(basic_ssp.table),np.copy(basic_ssp.sn2_table),np.copy(basic_ssp.agb_table),np.copy(basic_ssp.sn1a_table),np.copy(basic_ssp.bh_table))
	abundances,elements,numbers = mass_fraction_to_abundances(np.copy(cube.cube),np.copy(basic_solar.table))
	weights = cube.cube['sfr']
	abundances = append_fields(abundances,'weights',weights)
	abundances = np.array(abundances)

	return cube, abundances

def multi_star_optimization():
	'''
	This function will optimize the parameters of all stars in a hierachical manner (similar to gibbs sampling)

	INPUT:

	   a = will be loaded from parameter.py (prepare all variables there)

	OUTPUT:

	   log_list = a list of intermediate results (so far only for debugging)
	'''
	import time
	import multiprocessing as mp
	from .optimization import minimizer_initial, minimizer_global, minimizer_local
	from .cem_function import global_optimization_error_returned
	from .parameter import ModelParameters

	a = ModelParameters()
	print(a.stellar_identifier_list)
	start_time = time.time()

	log_list = []
	# I: Minimization for each star seperately
	# 1: for each star make initial conditions (each star needs other model parameters)
	parameter_list = []
	for item in a.stellar_identifier_list:
		parameter_list.append(item)
	# 2: call posterior_function_for_minimization with scipy.optimize.minimize in multiprocess for each star and recover the found parameters
	p = mp.Pool(len(parameter_list))
	t = p.map(minimizer_initial, parameter_list)
	p.close()
	p.join()
	result = np.vstack(t)

	log_list.append(np.copy(result))
	log_list.append('initial minimization')
	initial = time.time()
	print('first minimization for each star separately took: %2.f seconds' %(initial - start_time))

	# IV: repeat II and III until posterior does not change much
	result[:,:len(a.SSP_parameters)] = np.mean(result[:,:len(a.SSP_parameters)], axis = 0)
	posteriors = []
	counter = 0
	while True:
		counter += 1
		if len(posteriors) > 1:
			if np.abs(posteriors[-1] - posteriors[-2]) < a.gibbs_sampler_tolerance:
				break
			if len(posteriors) > a.gibbs_sampler_maxiter:
				break

		initial = time.time()
		# II: Global parameter minimization:
		# 1: only SSP parameters free. Use mean SSP parameter values and individual (but fixed ISM parameter values)
		changing_parameter = result[0,:len(a.SSP_parameters)]
		# 2: Call each star in multiprocess but only return the predictions
		# 3: Calculate the likelihood for each star and optimize the common model error (is all done within minimizer global, which is calling 'global optimization')
		x = minimizer_global(changing_parameter,  a.tol_minimization, a.maxiter_minimization, a.verbose, result)

		# 4: return global SSP parameters and common model error
		posterior, error_list, elements = global_optimization_error_returned(x, result)
		posteriors.append(posterior)
		print(posteriors)

		global_iteration1 = time.time()
		print('step %d global minimization took: %2.f seconds' %(counter, global_iteration1 - initial))

		# III: Local parameter minimization:
		# 1: Use fixed global parameters and fixed common errors make initial conditions
		result[:,:len(a.SSP_parameters)] = x

		log_list.append((np.copy(x),posterior))
		log_list.append('step %d global minimization' %(counter))

		p0_list = []
		parameter_list = []
		x_list = []
		error_list_mp = []
		element_list_mp = []

		for i,item in enumerate(a.stellar_identifier_list):
			parameter_list.append(item)
			p0_list.append(result[i,len(a.SSP_parameters):])
			x_list.append(x)
			error_list_mp.append(error_list)
			element_list_mp.append(elements)

		args = zip(p0_list,parameter_list,x_list,error_list_mp,element_list_mp)

		# 2: Minimize each star ISM parameters in multiprocess
		p = mp.Pool(len(parameter_list))
		t = p.map(minimizer_local, args)
		p.close()
		p.join()
		local_parameters = np.vstack(t)
		result[:,len(a.SSP_parameters):] = local_parameters

		log_list.append(np.copy(result))
		log_list.append('step %d local minimization' %(counter))
		local_iteration1 = time.time()
		print('step %d local minimization took: %2.f seconds' %(counter, local_iteration1 - global_iteration1))

	log_list.append(posteriors)
	print(log_list)

	# V: MCMC run
	## reshape the result to have global parameters in the front and the local parameters following
	changing_parameter = list(result[0,:len(a.SSP_parameters)])
	for i in range(result.shape[0]):
		changing_parameter.append(list(result[i,len(a.SSP_parameters):]))
	changing_parameter = np.hstack(changing_parameter)
	## jitter the parameters to initialise the chain (add a validation later, i.e. testing that the particular parameters yield a result)
	mcmc_multi(changing_parameter, error_list, elements)
	# 1: Free all parameters and optimize common error (SSP should be the same for all stars)
	# 2: Plug everything into emcee and sample the posterior
	return log_list

def mcmc(a):
	'''
	Convenience function to use the MCMC. A subdirectory mcmc/ will be created in the current directory and intermediate chains will be stored there.
	The chains are not actually flattened as the file name suggests.

	The MCMC will sample the volume of best posterior for the likelihood functions that are declared in parameter.py. Default is ['sol_norm','gas_reservoir','sn_ratio'] which corresponds to 'Sun+' from the paper.
	'''
	import time
	import os
	import multiprocessing as mp
	from .optimization import creating_chain, posterior_probability
	import emcee

	start1 = time.time()
	directory = 'mcmc/'
	if os.path.exists(directory):
		if a.verbose:
			print('%s already existed. Content might be overwritten' %(directory))
	else:
		os.makedirs(directory)

	a.check_processes = False
	a.number_of_models_overplotted = 1
	a.only_net_yields_in_process_tables = False
	a.testing_output = False
	a.summary_pdf = False
	a.nthreads = mp.cpu_count()
	if a.nthreads == 4:
		a.nthreads = 2

	chain = creating_chain(a,np.copy(a.p0))
	sampler = emcee.EnsembleSampler(a.nwalkers,a.ndim,posterior_probability,threads=a.nthreads, args = [a])
	pos,prob,state,blobs = sampler.run_mcmc(chain,a.mburn)

	mean_prob = mean_prob_beginning = np.zeros((a.m))
	posterior_list = []
	posterior_std_list = []
	for i in range(a.m):
		print('step ', i+1 , 'of ',a.m)
		pos, prob, state, blobs = sampler.run_mcmc(pos, a.save_state_every, rstate0=state, lnprob0=prob, blobs0 = blobs, storechain = True)
		np.save('%s/flatchain' %(directory),sampler.chain)
		np.save('%s/flatlnprobability' %(directory),sampler.lnprobability)
		np.save('%s/flatblobs' %(directory),sampler.blobs)
		posterior = np.load('%s/flatlnprobability.npy' %(directory))
		posterior_list.append(np.mean(posterior, axis = 0)[-1])
		posterior_std_list.append(np.std(posterior, axis = 0)[-1])
		np.save('%s/flatmeanposterior' %(directory), posterior_list)
		np.save('%s/flatstdposterior' %(directory), posterior_std_list)
		print(np.mean(posterior, axis = 0)[0], np.mean(posterior, axis = 0)[-1])

		if i>202:
			print('posterior -1, -100, -200',np.mean(posterior, axis = 0)[-1], np.mean(posterior, axis = 0)[-100], np.mean(posterior, axis = 0)[-200])
			print('posterior 0, 100, 200',np.mean(posterior, axis = 0)[0], np.mean(posterior, axis = 0)[100], np.mean(posterior, axis = 0)[200])
		#print("Mean acceptance fraction:", sampler.acceptance_fraction)
		elapsed1 = (time.time() - start1)
		print('calculation so far took', elapsed1, ' seconds')
		if i>a.min_mcmc_iterations and np.abs(np.mean(posterior, axis = 0)[-1] - np.mean(posterior, axis = 0)[-100]) < a.mcmc_tolerance and np.abs(np.mean(posterior, axis = 0)[-1] - np.mean(posterior, axis = 0)[-200]) < a.mcmc_tolerance:
			break
	if a.send_email:
		send_email(a.nthreads, i, np.mean(posterior, axis = 0)[0], np.mean(posterior, axis = 0)[-1], a, elapsed1)

def mcmc_multi(changing_parameter, error_list, elements):
	'''
	Convenience function to use the MCMC for multiple zones (and therefore multiple observations). A subdirectory mcmc/ will be created in the current directory and intermediate chains will be stored there.
	The MCMC will sample the volume of best posterior for the likelihood functions that are declared in parameter.py.
	Default is a list of Proto-sun, Arcturus and B-stars. The MCMC uses many walkers and can use multiple threads. Each walker will evaluate a series of Chempy zones and add their posterior together which then will be returned.

	INPUT:

	   changing_parameter = the parameter vector for initialization (will usually be found from minimization before). The initial chain will be created by jittering slightly the initial parameter guess

	   error_list = the vector of element errors

	   elements = the corresponding element symbols

	OUTPUT:

	   The function will create a folder and store the chain as well as the predicted element values

	The MCMC stops when the convergence criteria is met, which is when the median posterior of all walkers does not change much inbetween 200 steps anymore.
	'''
	import time
	import os
	import multiprocessing as mp
	from .cem_function import  posterior_function_many_stars
	from .parameter import ModelParameters
	import emcee

	a = ModelParameters()
	start1 = time.time()
	directory = 'mcmc/'
	if os.path.exists(directory):
		if a.verbose:
			print('%s already existed. Content might be overwritten' %(directory))
	else:
		os.makedirs(directory)

	nthreads = mp.cpu_count()
	if nthreads == 4:
		nthreads = 2
	ndim = len(changing_parameter)
	a.nwalkers = max(a.nwalkers, int(ndim*2))
	chain = np.empty(shape = (a.nwalkers,ndim))
	for i in range(a.nwalkers):
		result = -np.inf
		while result == -np.inf:
			jitter = np.random.normal(loc = 0, scale = 0.001, size = ndim)
			result, dummy = posterior_function_many_stars(changing_parameter + jitter,error_list,elements)
		chain[i] = changing_parameter + jitter

	sampler = emcee.EnsembleSampler(a.nwalkers,ndim,posterior_function_many_stars,threads=nthreads, args = [error_list,elements])
	pos,prob,state,blobs = sampler.run_mcmc(chain,a.mburn)

	mean_prob = mean_prob_beginning = np.zeros((a.m))
	posterior_list = []
	posterior_std_list = []
	for i in range(a.m):
		print('step ', i+1 , 'of ',a.m)
		pos, prob, state, blobs = sampler.run_mcmc(pos, a.save_state_every, rstate0=state, lnprob0=prob, blobs0 = blobs, storechain = True)
		np.save('%s/flatchain' %(directory),sampler.chain)
		np.save('%s/flatlnprobability' %(directory),sampler.lnprobability)
		np.save('%s/flatblobs' %(directory),sampler.blobs)
		posterior = np.load('%s/flatlnprobability.npy' %(directory))
		posterior_list.append(np.mean(posterior, axis = 0)[-1])
		posterior_std_list.append(np.std(posterior, axis = 0)[-1])
		np.save('%s/flatmeanposterior' %(directory), posterior_list)
		np.save('%s/flatstdposterior' %(directory), posterior_std_list)
		print(np.mean(posterior, axis = 0)[0], np.mean(posterior, axis = 0)[-1])

		if i>202:
			print('posterior -1, -100, -200',np.mean(posterior, axis = 0)[-1], np.mean(posterior, axis = 0)[-100], np.mean(posterior, axis = 0)[-200])
			print('posterior 0, 100, 200',np.mean(posterior, axis = 0)[0], np.mean(posterior, axis = 0)[100], np.mean(posterior, axis = 0)[200])
		#print("Mean acceptance fraction:", sampler.acceptance_fraction)
		elapsed1 = (time.time() - start1)
		print('calculation so far took', elapsed1, ' seconds')
		if i>a.min_mcmc_iterations and np.abs(np.mean(posterior, axis = 0)[-1] - np.mean(posterior, axis = 0)[-100]) < a.mcmc_tolerance and np.abs(np.mean(posterior, axis = 0)[-1] - np.mean(posterior, axis = 0)[-200]) < a.mcmc_tolerance:
			break
	if a.send_email:
		send_email(nthreads, i, np.mean(posterior, axis = 0)[0], np.mean(posterior, axis = 0)[-1], a, elapsed1)

def send_email(thread_count, iteration_count, posterior_beginning, posterior_end, parameters, time):
	from email.MIMEMultipart import MIMEMultipart
	from email.MIMEText import MIMEText
	import smtplib

	fromaddr = "pythonspeaking@gmail.com"
	toaddr = "rybizki@mpia.de"
	msg = MIMEMultipart()
	msg['From'] = fromaddr
	msg['To'] = toaddr
	msg['Subject'] = "Threads = %d, Run finished after %.2f hours" %(thread_count, time/3600.)
	body = "After %.1f hours %d threads produced %d iterations.\n The posterior at beginning was: %.2f. The posterior now is: %.2f.\n The stellar identifier list = %s.\n The error marginalization is %s \n  The yields are: %s %s %s \n " %(time/3600., thread_count, iteration_count, posterior_beginning, posterior_end, str(parameters.stellar_identifier_list), str(parameters.error_marginalization), parameters.yield_table_name_sn2, parameters.yield_table_name_agb, parameters.yield_table_name_1a)
	msg.attach(MIMEText(body, 'plain'))

	server = smtplib.SMTP('smtp.gmail.com', 587)
	server.ehlo()
	server.starttls()
	server.ehlo()
	server.login("pythonspeaking@gmail.com", "MPIA_Server_runs")
	text = msg.as_string()
	server.sendmail(fromaddr, toaddr, text)


#!/usr/bin/env python
# -*- coding: utf-8 -*-

"""
sRGB Colourspace
================

Defines the *sRGB* colourspace:

-   :attr:`sRGB_COLOURSPACE`.

See Also
--------
`RGB Colourspaces IPython Notebook
<http://nbviewer.ipython.org/github/colour-science/colour-ipython/blob/master/notebooks/models/rgb.ipynb>`_  # noqa

References
----------
.. [1]  `Recommendation ITU-R BT.709-5 - Parameter values for the HDTV
        standards for production and international programme exchange
        <http://www.itu.int/dms_pubrec/itu-r/rec/bt/R-REC-BT.709-5-200204-I!!PDF-E.pdf>`_  # noqa
        (Last accessed 24 February 2014)
"""

from __future__ import division, unicode_literals

import numpy as np

from colour.colorimetry import ILLUMINANTS
from colour.models import RGB_Colourspace

__author__ = 'Colour Developers'
__copyright__ = 'Copyright (C) 2013 - 2014 - Colour Developers'
__license__ = 'New BSD License - http://opensource.org/licenses/BSD-3-Clause'
__maintainer__ = 'Colour Developers'
__email__ = 'colour-science@googlegroups.com'
__status__ = 'Production'

__all__ = ['sRGB_PRIMARIES',
           'sRGB_WHITEPOINT',
           'sRGB_TO_XYZ_MATRIX',
           'XYZ_TO_sRGB_MATRIX',
           'sRGB_TRANSFER_FUNCTION',
           'sRGB_INVERSE_TRANSFER_FUNCTION',
           'sRGB_COLOURSPACE']

sRGB_PRIMARIES = np.array(
    [[0.6400, 0.3300],
     [0.3000, 0.6000],
     [0.1500, 0.0600]])
"""
*sRGB* colourspace primaries.

sRGB_PRIMARIES : ndarray, (3, 2)
"""

sRGB_WHITEPOINT = ILLUMINANTS.get(
    'CIE 1931 2 Degree Standard Observer').get('D65')
"""
*sRGB* colourspace whitepoint.

sRGB_WHITEPOINT : tuple
"""

sRGB_TO_XYZ_MATRIX = np.array(
    [[0.41238656, 0.35759149, 0.18045049],
     [0.21263682, 0.71518298, 0.0721802],
     [0.01933062, 0.11919716, 0.95037259]])
"""
*sRGB* colourspace to *CIE XYZ* colourspace matrix.

sRGB_TO_XYZ_MATRIX : array_like, (3, 3)
"""

XYZ_TO_sRGB_MATRIX = np.linalg.inv(sRGB_TO_XYZ_MATRIX)
"""
*CIE XYZ* colourspace to *sRGB* colourspace matrix.

XYZ_TO_sRGB_MATRIX : array_like, (3, 3)
"""

sRGB_TRANSFER_FUNCTION = lambda x: (
    x * 12.92 if x <= 0.0031308 else 1.055 * (x ** (1 / 2.4)) - 0.055)
"""
Transfer function from linear to *sRGB* colourspace.

sRGB_TRANSFER_FUNCTION : object
"""

sRGB_INVERSE_TRANSFER_FUNCTION = lambda x: (
    x / 12.92 if x <= 0.0031308 else ((x + 0.055) / 1.055) ** 2.4)
"""
Inverse transfer function from *sRGB* colourspace to linear.

sRGB_INVERSE_TRANSFER_FUNCTION : object
"""

sRGB_COLOURSPACE = RGB_Colourspace(
    'sRGB',
    sRGB_PRIMARIES,
    sRGB_WHITEPOINT,
    sRGB_TO_XYZ_MATRIX,
    XYZ_TO_sRGB_MATRIX,
    sRGB_TRANSFER_FUNCTION,
    sRGB_INVERSE_TRANSFER_FUNCTION)
"""
*sRGB* colourspace.

sRGB_COLOURSPACE : RGB_Colourspace
"""


from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import numpy as np

import libs.boxes.cython_bbox as cython_bbox
import libs.configs.config_v1 as cfg
from libs.boxes.bbox_transform import bbox_transform, bbox_transform_inv, clip_boxes
from libs.boxes.anchor import anchors_plane
from libs.logs.log import LOG
# FLAGS = tf.app.flags.FLAGS

_DEBUG = False

def encode(gt_boxes, all_anchors, height, width, stride):
  """Matching and Encoding groundtruth into learning targets
  Sampling

  Parameters
  ---------
  gt_boxes: an array of shape (G x 5), [x1, y1, x2, y2, class]
  all_anchors: an array of shape (h, w, A, 4),
  width: width of feature
  height: height of feature
  stride: downscale factor w.r.t the input size, e.g., [4, 8, 16, 32]
  Returns
  --------
  labels:   Nx1 array in [0, num_classes]
  bbox_targets: N x (4) regression targets
  bbox_inside_weights: N x (4), in {0, 1} indicating to which class is assigned.
  """
  # TODO: speedup this module
  if all_anchors is None:
    all_anchors = anchors_plane(height, width, stride=stride)

  # anchors, inds_inside, total_anchors
  border = cfg.FLAGS.allow_border
  all_anchors = all_anchors.reshape((-1, 4))
  inds_inside = np.where(
    (all_anchors[:, 0] >= -border) &
    (all_anchors[:, 1] >= -border) &
    (all_anchors[:, 2] < (width * stride) + border) &
    (all_anchors[:, 3] < (height * stride) + border))[0]
  anchors = all_anchors[inds_inside, :]
  total_anchors = all_anchors.shape[0]

  # choose boxes to assign to this stride
  # TODO gt assignment outside
  if False:
    areas = (gt_boxes[:, 3] - gt_boxes[:, 1] + 1) * (gt_boxes[:, 2] - gt_boxes[:, 0] + 1)
    ks = np.floor(4 + np.log2(np.sqrt(areas) / 224.0))
    K = int(np.log2(stride))
    inds = np.where((K == ks + 4))[0]
    if inds.size > 0:
      gt_boxes = gt_boxes[inds]
    else:
      labels = np.zeros((total_anchors), dtype=np.float32)
      bbox_targets = np.zeros((total_anchors, 4), dtype=np.float32)
      bbox_inside_weights = np.zeros((total_anchors, 4), dtype=np.float32)
      return labels, bbox_targets, bbox_inside_weights

  labels = np.zeros((anchors.shape[0], ), dtype=np.float32)

  if gt_boxes.size > 0:
      overlaps = cython_bbox.bbox_overlaps(
                 np.ascontiguousarray(anchors, dtype=np.float),
                 np.ascontiguousarray(gt_boxes[:, :4], dtype=np.float))

      if _DEBUG:
          print ('gt_boxes shape: ', gt_boxes.shape)
          print ('anchors shape: ', anchors.shape)
          print ('overlaps shape: ', overlaps.shape)

      gt_assignment = overlaps.argmax(axis=1)  # (A)
      max_overlaps = overlaps[np.arange(len(inds_inside)), gt_assignment]
      gt_argmax_overlaps = overlaps.argmax(axis=0)  # G
      gt_max_overlaps = overlaps[gt_argmax_overlaps,
                                 np.arange(overlaps.shape[1])]

      if False:
        # this is sentive to boxes of little overlaps, no need!
        gt_argmax_overlaps = np.where(overlaps == gt_max_overlaps)[0]

      # fg label: for each gt, assign anchor with highest overlap despite its overlaps
      labels[gt_argmax_overlaps] = 1
      # fg label: above threshold IOU
      labels[max_overlaps >= cfg.FLAGS.fg_threshold] = 1
      # print (np.min(labels), np.max(labels))

      # subsample positive labels if there are too many
      num_fg = int(cfg.FLAGS.fg_rpn_fraction * cfg.FLAGS.rpn_batch_size)
      fg_inds = np.where(labels == 1)[0]
      if len(fg_inds) > num_fg:
        disable_inds = np.random.choice(fg_inds, size=(len(fg_inds) - num_fg), replace=False)
        labels[disable_inds] = -1
  else:
      # if there is no gt
      labels[:] = 0

  # TODO: mild hard negative mining
  # subsample negative labels if there are too many
  num_bg = cfg.FLAGS.rpn_batch_size - np.sum(labels == 1)
  bg_inds = np.where(labels == 0)[0]
  if len(bg_inds) > num_bg:
    disable_inds = np.random.choice(bg_inds, size=(len(bg_inds) - num_bg), replace=False)
    labels[disable_inds] = -1

  bbox_targets = np.zeros((len(inds_inside), 4), dtype=np.float32)
  if gt_boxes.size > 0:
    bbox_targets = _compute_targets(anchors, gt_boxes[gt_assignment, :])
  bbox_inside_weights = np.zeros((len(inds_inside), 4), dtype=np.float32)
  bbox_inside_weights[labels == 1, :] = 1

  # mapping to whole outputs
  labels = _unmap(labels, total_anchors, inds_inside, fill=-1)
  bbox_targets = _unmap(bbox_targets, total_anchors, inds_inside, fill=0)
  bbox_inside_weights = _unmap(bbox_inside_weights, total_anchors, inds_inside, fill=0)

  labels = labels.reshape((1, height, width, -1))
  bbox_targets = bbox_targets.reshape((1, height, width, -1))
  bbox_inside_weights = bbox_inside_weights.reshape((1, height, width, -1))

  return labels, bbox_targets, bbox_inside_weights

def decode(boxes, scores, all_anchors, ih, iw):
  """Decode outputs into boxes
  Parameters
  ---------
  boxes: an array of shape (1, h, w, Ax4)
  scores: an array of shape (1, h, w, Ax2),
  all_anchors: an array of shape (1, h, w, Ax4), [x1, y1, x2, y2]

  Returns
  --------
  final_boxes: of shape (R x 4)
  classes: of shape (R) in {0,1,2,3... K-1}
  scores: of shape (R) in [0 ~ 1]
  """
  h, w = boxes.shape[1], boxes.shape[2]
  if all_anchors is  None:
    stride = 2 ** int(round(np.log2((iw + 0.0) / w)))
    all_anchors = anchors_plane(h, w, stride=stride)
  all_anchors = all_anchors.reshape((-1, 4))
  boxes = boxes.reshape((-1, 4))
  scores = scores.reshape((-1, 2))
  assert scores.shape[0] == boxes.shape[0] == all_anchors.shape[0], \
    'Anchor layer shape error %d vs %d vs %d' % (scores.shape[0],boxes.shape[0],all_anchors.reshape[0])
  boxes = bbox_transform_inv(all_anchors, boxes)
  classes = np.argmax(scores, axis=1)
  scores = scores[:, 1]
  final_boxes = boxes
  final_boxes = clip_boxes(final_boxes, (ih, iw))
  classes = classes.astype(np.int32)
  return final_boxes, classes, scores

def sample(boxes, scores, ih, iw, is_training):
  """
  Sampling the anchor layer outputs for next stage, mask or roi prediction or roi

  Params
  ----------
  boxes:  of shape (? ,4)
  scores: foreground prob
  ih:     image height
  iw:     image width
  is_training:  'test' or 'train'

  Returns
  ----------
  rois: of shape (N, 4)
  scores: of shape (N, 1)
  batch_ids:
  """
  return

def _unmap(data, count, inds, fill=0):
  """ Unmap a subset of item (data) back to the original set of items (of
  size count) """
  if len(data.shape) == 1:
    ret = np.empty((count,), dtype=np.float32)
    ret.fill(fill)
    ret[inds] = data
  else:
    ret = np.empty((count,) + data.shape[1:], dtype=np.float32)
    ret.fill(fill)
    ret[inds, :] = data
  return ret

def _compute_targets(ex_rois, gt_rois):
  """Compute bounding-box regression targets for an image."""

  assert ex_rois.shape[0] == gt_rois.shape[0]
  assert ex_rois.shape[1] == 4
  assert gt_rois.shape[1] == 5

  return bbox_transform(ex_rois, gt_rois[:, :4]).astype(np.float32, copy=False)

if __name__ == '__main__':

  import time
  t = time.time()

  for i in range(10):
    cfg.FLAGS.fg_threshold = 0.1
    classes = np.random.randint(0, 3, (50, 1))
    boxes = np.random.randint(10, 50, (50, 2))
    s = np.random.randint(20, 50, (50, 2))
    s = boxes + s
    boxes = np.concatenate((boxes, s), axis=1)
    gt_boxes = np.hstack((boxes, classes))
    # gt_boxes = boxes
    rois = np.random.randint(10, 50, (20, 2))
    s = np.random.randint(0, 20, (20, 2))
    s = rois + s
    rois = np.concatenate((rois, s), axis=1)
    labels, bbox_targets, bbox_inside_weights = encode(gt_boxes, all_anchors=None, height=200, width=300, stride=4)
    labels, bbox_targets, bbox_inside_weights = encode(gt_boxes, all_anchors=None, height=100, width=150, stride=8)
    labels, bbox_targets, bbox_inside_weights = encode(gt_boxes, all_anchors=None, height=50, width=75, stride=16)
    labels, bbox_targets, bbox_inside_weights = encode(gt_boxes, all_anchors=None, height=25, width=37, stride=32)
    # anchors, _, _ = anchors_plane(200, 300, stride=4, boarder=0)

  print('average time: %f' % ((time.time() - t)/10.0))


# This script repeatedly estimates the relative binding free energy of a single edge, along with the gradient of the
# estimate with respect to force field parameters, and adjusts the force field parameters to improve tha accuracy
# of the free energy prediction.

import argparse
import numpy as np
import jax
from jax import numpy as jnp

from fe.free_energy import construct_lambda_schedule
from fe.utils import convert_uIC50_to_kJ_per_mole
from fe import model
from md import builders

from testsystems.relative import hif2a_ligand_pair

from ff.handlers.serialize import serialize_handlers
from ff.handlers.nonbonded import AM1CCCHandler, LennardJonesHandler
from parallel.client import CUDAPoolClient
from parallel.utils import get_gpu_count

from typing import Union, Optional, Iterable, Any, Tuple, Dict

from optimize.step import truncated_step

array = Union[np.array, jnp.array]
Handler = Union[AM1CCCHandler, LennardJonesHandler] # TODO: do these all inherit from a Handler class already?

if __name__ == "__main__":

    parser = argparse.ArgumentParser(
        description="Relative Binding Free Energy Testing",
        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
    )

    parser.add_argument(
        "--num_gpus",
        type=int,
        help="number of gpus",
        default=get_gpu_count()
    )

    parser.add_argument(
        "--num_complex_windows",
        type=int,
        help="number of vacuum lambda windows",
        required=True
    )

    parser.add_argument(
        "--num_solvent_windows",
        type=int,
        help="number of solvent lambda windows",
        required=True
    )

    parser.add_argument(
        "--num_equil_steps",
        type=int,
        help="number of equilibration steps for each lambda window",
        required=True
    )

    parser.add_argument(
        "--num_prod_steps",
        type=int,
        help="number of production steps for each lambda window",
        required=True
    )

    cmd_args = parser.parse_args()

    client = CUDAPoolClient(max_workers=cmd_args.num_gpus)

    # fetch mol_a, mol_b, core, forcefield from testsystem
    mol_a, mol_b, core = hif2a_ligand_pair.mol_a, hif2a_ligand_pair.mol_b, hif2a_ligand_pair.core
    forcefield = hif2a_ligand_pair.ff

    # compute ddG label from mol_a, mol_b
    # TODO: add label upon testsystem construction
    # (ytz): these are *binding* free energies, i.e. values that are less than zero.
    label_dG_a = convert_uIC50_to_kJ_per_mole(float(mol_a.GetProp("IC50[uM](SPA)")))
    label_dG_b = convert_uIC50_to_kJ_per_mole(float(mol_b.GetProp("IC50[uM](SPA)")))
    label_ddG = label_dG_b - label_dG_a  # complex - solvent

    print("binding dG_a", label_dG_a)
    print("binding dG_b", label_dG_b)

    hif2a_ligand_pair.label = label_ddG

    # construct lambda schedules for complex and solvent
    complex_schedule = construct_lambda_schedule(cmd_args.num_complex_windows)
    solvent_schedule = construct_lambda_schedule(cmd_args.num_solvent_windows)

    # build the protein system.
    complex_system, complex_coords, _, _, complex_box, _ = builders.build_protein_system(
        'tests/data/hif2a_nowater_min.pdb')
    complex_box += np.eye(3) * 0.1  # BFGS this later

    # build the water system.
    solvent_system, solvent_coords, solvent_box, _ = builders.build_water_system(4.0)
    solvent_box += np.eye(3) * 0.1  # BFGS this later

    binding_model = model.RBFEModel(
        client,
        forcefield,
        complex_system,
        complex_coords,
        complex_box,
        complex_schedule,
        solvent_system,
        solvent_coords,
        solvent_box,
        solvent_schedule,
        cmd_args.num_equil_steps,
        cmd_args.num_prod_steps
    )

    vg_fn = jax.value_and_grad(binding_model.loss, argnums=0)

    ordered_params = forcefield.get_ordered_params()
    ordered_handles = forcefield.get_ordered_handles()

    handle_types_being_optimized = [AM1CCCHandler, LennardJonesHandler]

    def flatten(params) -> Tuple[np.array, callable]:
        """Turn params dict into flat array, with an accompanying unflatten function

        TODO: note that the result is going to be in the order given by ordered_handles (filtered by presence in handle_types)
            rather than in the order they appear in handle_types_being_optimized

        TODO: maybe leave out the reference to handle_types_being optimized altogether

        TODO: does Jax have a pytree-based flatten / unflatten utility?
        """

        theta_list = []
        _shapes = dict()
        _handle_types = []

        for param, handle in zip(params, ordered_handles):
            assert handle.params.shape == param.shape
            key = type(handle)

            if key in handle_types_being_optimized:
                theta_list.append(param.flatten())
                _shapes[key] = param.shape
                _handle_types.append(key)

        theta = np.hstack(theta_list)

        def unflatten(theta: array) -> Dict[Handler, array]:
            params = dict()
            i = 0
            for key in _handle_types:
                shape = _shapes[key]
                num_params = int(np.prod(shape))
                params[key] = np.array(theta[i: i + num_params]).reshape(shape)
                i += num_params
            return params

        return theta, unflatten

    # in each optimization step, don't step so far that you think you're jumping to
    #   loss_next = relative_improvement_bound * loss_current
    relative_improvement_bound = 0.95

    flat_theta_traj = []
    flat_grad_traj = []
    loss_traj = []

    for epoch in range(1000):
        epoch_params = serialize_handlers(ordered_handles)
        (loss, aux), loss_grad = vg_fn(ordered_params, mol_a, mol_b, core, label_ddG)

        print("epoch", epoch, "loss", loss)

        # note: unflatten_grad and unflatten_theta have identical definitions for now
        flat_loss_grad, unflatten_grad = flatten(loss_grad)
        flat_theta, unflatten_theta = flatten(ordered_params)

        step_lower_bound = loss * relative_improvement_bound
        theta_increment = truncated_step(flat_theta, loss, flat_loss_grad, step_lower_bound=step_lower_bound)
        param_increments= unflatten_theta(theta_increment)

        # for any parameter handler types being updated, update in place
        for handle in ordered_handles:
            handle_type = type(handle)
            if handle_type in param_increments:
                print(f'updating {handle_type.__name__}')

                print(f'\tbefore update: {handle.params}')
                handle.params += param_increments[handle_type] # TODO: careful -- this must be a "+=" or "-=" not an "="!
                print(f'\tafter update:  {handle.params}')

                # useful for debugging to dump out the grads
                # for smirks, dp in zip(handle.smirks, loss_grad):
                    # if np.any(dp) > 0:
                        # print(smirks, dp)

        # checkpoint results to npz (overwrite
        flat_theta_traj.append(np.array(flat_theta))
        flat_grad_traj.append(flat_loss_grad)
        loss_traj.append(loss)

        path_to_npz = 'results_checkpoint.npz'
        print(f'saving theta, grad, loss trajs to {path_to_npz}')
        np.savez(
            path_to_npz,
            theta_traj=np.array(flat_theta_traj),
            grad_traj=np.array(flat_grad_traj),
            loss_traj=np.array(loss_traj)
        )

        # write ff parameters after each epoch
        path_to_ff_checkpoint = f"checkpoint_{epoch}.py"
        print(f'saving force field parameter checkpoint to {path_to_ff_checkpoint}')
        with open(path_to_ff_checkpoint, 'w') as fh:
            fh.write(epoch_params)


from __future__ import division
from build.el_propagator import el_run
from mqc.mqc import MQC
from misc import au_to_K, call_name, typewriter
import os, shutil, textwrap
import numpy as np
import pickle

class Eh(MQC):
    """ Class for Ehrenfest dynamics

        :param object molecule: Molecule object
        :param object thermostat: Thermostat object
        :param integer istate: Initial state
        :param double dt: Time interval
        :param integer nsteps: Total step of nuclear propation
        :param integer nesteps: Total step of electronic propagation
        :param string elec_object: Electronic equation of motions
        :param string propagator: Electronic propagator
        :param boolean l_print_dm: Logical to print BO population and coherence
        :param boolean l_adj_nac: Logical to adjust nonadiabatic coupling
        :param init_coef: Initial BO coefficient
        :type init_coef: Double, list or complex, list
        :param string unit_dt: Unit of time step (fs = femtosecond, au = atomic unit)
        :param integer out_freq: Frequency of printing output
        :param integer verbosity: Verbosity of output
    """
    def __init__(self, molecule, thermostat=None, istate=0, dt=0.5, nsteps=1000, nesteps=20, \
        elec_object="density", propagator="rk4", l_print_dm=True, l_adj_nac=True, \
        init_coef=None, unit_dt="fs", out_freq=1, verbosity=0):
        # Initialize input values
        super().__init__(molecule, thermostat, istate, dt, nsteps, nesteps, \
            elec_object, propagator, l_print_dm, l_adj_nac, init_coef, unit_dt, out_freq, verbosity)

        # Debug variables
        self.dotpopnac = np.zeros(self.mol.nst)

    def run(self, qm, mm=None, output_dir="./", l_save_qm_log=False, l_save_mm_log=False, l_save_scr=True, restart=None):
        """ Run MQC dynamics according to Ehrenfest dynamics

            :param object qm: QM object containing on-the-fly calculation infomation
            :param object mm: MM object containing MM calculation infomation
            :param string output_dir: Name of directory where outputs to be saved.
            :param boolean l_save_qm_log: Logical for saving QM calculation log
            :param boolean l_save_mm_log: Logical for saving MM calculation log
            :param boolean l_save_scr: Logical for saving scratch directory
            :param string restart: Option for controlling dynamics restarting
        """
        # Initialize PyUNIxMD
        base_dir, unixmd_dir, qm_log_dir, mm_log_dir =\
             self.run_init(qm, mm, output_dir, l_save_qm_log, l_save_mm_log, l_save_scr, restart)
        bo_list = [ist for ist in range(self.mol.nst)]
        qm.calc_coupling = True
        self.print_init(qm, mm, restart)

        if (restart == None):
            # Calculate initial input geometry at t = 0.0 s
            self.istep = -1
            self.mol.reset_bo(qm.calc_coupling)
            qm.get_data(self.mol, base_dir, bo_list, self.dt, self.istep, calc_force_only=False)
            if (self.mol.l_qmmm and mm != None):
                mm.get_data(self.mol, base_dir, bo_list, self.istep, calc_force_only=False)
            self.mol.get_nacme()

            self.update_energy()

            self.write_md_output(unixmd_dir, self.istep)
            self.print_step(self.istep)

        elif (restart == "write"):
            # Reset initial time step to t = 0.0 s
            self.istep = -1
            self.write_md_output(unixmd_dir, self.istep)
            self.print_step(self.istep)

        elif (restart == "append"):
            # Set initial time step to last successful step of previous dynamics
            self.istep = self.fstep

        self.istep += 1

        # Main MD loop
        for istep in range(self.istep, self.nsteps):

            self.calculate_force()
            self.cl_update_position()

            self.mol.backup_bo()
            self.mol.reset_bo(qm.calc_coupling)
            qm.get_data(self.mol, base_dir, bo_list, self.dt, istep, calc_force_only=False)
            if (self.mol.l_qmmm and mm != None):
                mm.get_data(self.mol, base_dir, bo_list, istep, calc_force_only=False)

            if (self.l_adj_nac):
                self.mol.adjust_nac()

            self.calculate_force()
            self.cl_update_velocity()

            self.mol.get_nacme()

            el_run(self)

            if (self.thermo != None):
                self.thermo.run(self)

            self.update_energy()

            if ((istep + 1) % self.out_freq == 0):
                self.write_md_output(unixmd_dir, istep)
                self.print_step(istep)
            if (istep == self.nsteps - 1):
                self.write_final_xyz(unixmd_dir, istep)

            self.fstep = istep
            restart_file = os.path.join(base_dir, "RESTART.bin")
            with open(restart_file, 'wb') as f:
                pickle.dump({'qm':qm, 'md':self}, f)

        # Delete scratch directory
        if (not l_save_scr):
            tmp_dir = os.path.join(unixmd_dir, "scr_qm")
            if (os.path.exists(tmp_dir)):
                shutil.rmtree(tmp_dir)

            if (self.mol.l_qmmm and mm != None):
                tmp_dir = os.path.join(unixmd_dir, "scr_mm")
                if (os.path.exists(tmp_dir)):
                    shutil.rmtree(tmp_dir)

    def calculate_force(self):
        """ Calculate the Ehrenfest force
        """
        self.rforce = np.zeros((self.mol.nat, self.mol.ndim))

        for ist, istate in enumerate(self.mol.states):
            self.rforce += istate.force * self.mol.rho.real[ist, ist]

        for ist in range(self.mol.nst):
            for jst in range(ist + 1, self.mol.nst):
                self.rforce += 2. * self.mol.nac[ist, jst] * self.mol.rho.real[ist, jst] \
                    * (self.mol.states[ist].energy - self.mol.states[jst].energy)

    def update_energy(self):
        """ Routine to update the energy of molecules in Ehrenfest dynamics
        """
        # Update kinetic energy
        self.mol.update_kinetic()
        self.mol.epot = 0.
        for ist, istate in enumerate(self.mol.states):
            self.mol.epot += self.mol.rho.real[ist, ist] * self.mol.states[ist].energy
        self.mol.etot = self.mol.epot + self.mol.ekin

    def write_md_output(self, unixmd_dir, istep):
        """ Write output files

            :param string unixmd_dir: PyUNIxMD directory
            :param integer istep: Current MD step
        """
        # Write the common part
        super().write_md_output(unixmd_dir, istep)

        # Write time-derivative BO population
        self.write_dotpop(unixmd_dir, istep)

    def write_dotpop(self, unixmd_dir, istep):
        """ Write time-derivative BO population

            :param string unixmd_dir: PyUNIxMD directory
            :param integer istep: Current MD step
        """
        # Write NAC term in DOTPOPNAC
        if (self.verbosity >= 1):
            tmp = f'{istep + 1:9d}' + "".join([f'{pop:15.8f}' for pop in self.dotpopnac])
            typewriter(tmp, unixmd_dir, "DOTPOPNAC", "a")

    def print_init(self, qm, mm, restart):
        """ Routine to print the initial information of dynamics

            :param object qm: QM object containing on-the-fly calculation infomation
            :param object mm: MM object containing MM calculation infomation
            :param string restart: Option for controlling dynamics restarting
        """
        # Print initial information about molecule, qm, mm and thermostat
        super().print_init(qm, mm, restart)

        # Print dynamics information for start line
        dynamics_step_info = textwrap.dedent(f"""\

        {"-" * 118}
        {"Start Dynamics":>65s}
        {"-" * 118}
        """)

        # Print INIT for each step
        INIT = f" #INFO{'STEP':>8s}{'Kinetic(H)':>15s}{'Potential(H)':>15s}{'Total(H)':>13s}{'Temperature(K)':>17s}{'norm':>8s}"
        dynamics_step_info += INIT

        print (dynamics_step_info, flush=True)

    def print_step(self, istep):
        """ Routine to print each steps infomation about dynamics

            :param integer istep: Current MD step
        """
        ctemp = self.mol.ekin * 2. / float(self.mol.ndof) * au_to_K
        norm = 0.
        for ist in range(self.mol.nst):
            norm += self.mol.rho.real[ist, ist]

        # Print INFO for each step
        INFO = f" INFO{istep + 1:>9d} "
        INFO += f"{self.mol.ekin:14.8f}{self.mol.epot:15.8f}{self.mol.etot:15.8f}"
        INFO += f"{ctemp:13.6f}"
        INFO += f"{norm:11.5f}"
        print (INFO, flush=True)