import collections
import openpathsampling as paths
from openpathsampling.netcdfplus import StorableNamedObject
from openpathsampling.numerics import LookupFunction
import pandas as pd
import numpy as np

from .core import EnsembleHistogrammer, MultiEnsembleSamplingAnalyzer

class FullHistogramMaxLambdas(EnsembleHistogrammer):
    """Histogramming the full max-lambda function (one way of getting TCP)

    This histograms the maximum value of lambda for each ensemble. One of
    these objects is made per transition.

    Parameters
    ----------
    transition: :class:`.TISTransition`
        the transition to be analyzed
    hist_parameters: dict
        Histogram parameters to use with this collective variable: allowed
        keys are 'bin_width' and 'bin_range'; value for 'bin_width' is a
        float; for 'bin_range' is a tuple with `(left_edge, right_edge)`
        (only left edge is used)
    max_lambda_func: callable
        function to use to map the trajectories to a histogram; default is
        `None`, which uses the maximum value of the order parameter
        associated with the interface set. Overriding this can be used if
        either (a) the interface set does not have an order parameter
        associated with it, or (b) you want to calculate the values along
        some other order parameter
    """
    def __init__(self, transition, hist_parameters, max_lambda_func=None):
        self.transition = transition
        if max_lambda_func is None:
            try:
                max_lambda_func = transition.interfaces.cv_max
            except AttributeError:
                pass  # leave max_lambda_func as None

        if max_lambda_func is None:
            raise RuntimeError("Can't identify function to determine max "
                               + "value of order parameter.")

        # TODO: is this used?
        self.lambdas = {e: l for (e, l) in zip(transition.ensembles,
                                               transition.interfaces.lambdas)}
        super(FullHistogramMaxLambdas, self).__init__(
            ensembles=transition.ensembles,
            f=max_lambda_func,
            hist_parameters=hist_parameters
        )


#class PerEnsembleMaxLambdas(EnsembleHistogrammer):
    # TODO: this just maps the count to the ensemble, not the full histogram
    #def __init__(self, transition):
        #interfaces_lambdas = transition.interfaces.lambdas


class TotalCrossingProbability(MultiEnsembleSamplingAnalyzer):
    """
    Calculate the total crossing probability function.

    The total crossing probability function is generated by calculating the
    individual ensemble crossing probability functions (using, e.g.,
    :class:`.FullHistogramMaxLambdas`, and combining them using some
    combining method (default is :class:`.WHAM`). One of these objects is
    instantiated per transition.

    Parameters
    ----------
    max_lambda_calc: :class:`.EnsembleHistogrammer`
        usually :class:`.FullHistogramMaxLambdas`; object that creates the
        max lambda histograms for the ensembles associated with this
        transition.
    combiner: TODO
        class that combines multiple histograms (with restricted sampling)
        into a single result. If `None` (default), uses :class:`.WHAM`
    """
    def __init__(self, max_lambda_calc, combiner=None):
        transition = max_lambda_calc.transition
        super(TotalCrossingProbability, self).__init__(transition.ensembles)
        self.max_lambda_calc = max_lambda_calc
        self.transition = transition
        if combiner is None:
            lambdas = self.transition.interfaces.lambdas
            combiner = paths.numerics.WHAM(interfaces=lambdas)
        self.combiner = combiner

    def from_weighted_trajectories(self, input_dict):
        """Calculate results from a weighted trajectories dictionary.

        Parameters
        ----------
        input_dict : dict of {:class:`.Ensemble`: collections.Counter}
            ensemble as key, and a counter mapping each trajectory
            associated with that ensemble to its counter of time spent in
            the ensemble (output of ``steps_to_weighted_trajectories``)

        Returns
        -------
        :class:`.LookupFunction`
            the total crossing probability function
        """
        hists = self.max_lambda_calc.from_weighted_trajectories(input_dict)
        return self.from_ensemble_histograms(hists)

    def from_ensemble_histograms(self, hists):
        """Calculate results from a dict of ensemble histograms.

        Parameters
        ----------
        hists : dict of {:class:`.Ensemble`: :class:`.numerics.Histogram`}
            histogram for each ensemble (from ``self.max_lambda_calc``)

        Returns
        -------
        :class:`.LookupFunction`
            the total crossing probability function
        """
        tcp_results = {}
        input_hists = [hists[ens] for ens in self.transition.ensembles]
        df = paths.numerics.histograms_to_pandas_dataframe(
            input_hists,
            fcn="reverse_cumulative"
        ).sort_index(axis=1)
        # TODO: remove WHAM-specific name here
        tcp = self.combiner.wham_bam_histogram(df).to_dict()
        return LookupFunction(tcp.keys(), tcp.values())
