# This file is part of PyEMMA.
#
# Copyright (c) 2016, 2017 Computational Molecular Biology Group, Freie Universitaet Berlin (GER)
#
# PyEMMA is free software: you can redistribute it and/or modify
# it under the terms of the GNU Lesser General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU Lesser General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.

import numpy as _np
from pyemma.util import types

__all__ = [
    'get_averaged_bias_matrix',
    'get_umbrella_sampling_data',
    'get_multi_temperature_data',
    'assign_unbiased_state_label']

# ==================================================================================================
# helpers for discrete estimations
# ==================================================================================================

def get_averaged_bias_matrix(bias_sequences, dtrajs, nstates=None):
    r"""
    Computes a bias matrix via an exponential average of the observed frame wise bias energies.

    Parameters
    ----------
    bias_sequences : list of numpy.ndarray(T_i, num_therm_states)
        A single reduced bias energy trajectory or a list of reduced bias energy trajectories.
        For every simulation frame seen in trajectory i and time step t, btrajs[i][t, k] is the
        reduced bias energy of that frame evaluated in the k'th thermodynamic state (i.e. at
        the k'th Umbrella/Hamiltonian/temperature)
    dtrajs : list of numpy.ndarray(T_i) of int
        A single discrete trajectory or a list of discrete trajectories. The integers are indexes
        in 0,...,num_conf_states-1 enumerating the num_conf_states Markov states or the bins the
        trajectory is in at any time.
    nstates : int, optional, default=None
        Number of configuration states.

    Returns
    -------
    bias_matrix : numpy.ndarray(shape=(num_therm_states, num_conf_states)) object
        bias_energies_full[j, i] is the bias energy in units of kT for each discrete state i
        at thermodynamic state j.
    """
    from pyemma.thermo.extensions.util import (logsumexp as _logsumexp, logsumexp_pair as _logsumexp_pair)

    nmax = int(_np.max([dtraj.max() for dtraj in dt
rajs]))
    if nstates is None:
        nstates = nmax + 1
    elif nstates < nmax + 1:
        raise ValueError("nstates is smaller than the number of observed microstates")
    nthermo = bias_sequences[0].shape[1]
    bias_matrix = -_np.ones(shape=(nthermo, nstates), dtype=_np.float64) * _np.inf
    counts = _np.zeros(shape=(nstates,), dtype=_np.intc)
    for s in range(len(bias_sequences)):
        for i in range(nstates):
            idx = (dtrajs[s] == i)
            nidx = idx.sum()
   
         if nidx == 0:
                continue
            counts[i] += nidx
            selected_bias_sequence = bias_sequences[s][idx, :]
            for k in range(nthermo):
                bias_matrix[k, i] = _logsumexp_pair(
                    bias_matrix[k, i],
                    _logsumexp(
                        _np.ascontiguousarray(-selected_bias_sequence[:, k]),
                        inplace=False))
    idx = counts.nonzero()
    log_counts = _np.log(counts[idx])
    bias_matrix *= -1.0
    bias_matrix[:, idx] += log_counts[_np.newaxis, :]
    return bias_matrix

# ==================================================================================================
# helpers for umbrella sampling simulations
# ==================================================================================================

def _ensure_umbrella_center(candidate, dimension):
    try:
        candidate = _np.asarray(candidate).astype(_np.float64).reshape((-1,))
    except ValueError:
        raise ValueError("Umbrella center " + str(candidate) + " cannot be cast as numpy.ndarray")
    if candidate.shape[0] == 1 and dimension > 1:
        return candidate[0] * _np.ones(shape=(dimension,), dtype=_np.float64)
    elif candidate.shape[0] != dimension:
        raise ValueError("Unmatching dimensions: umbrella center " + str(candidate) + \
            " is not compatible with dimension %d" % dimension)
    return candidate

def _ensure_force_constant(candidate, dimension):
    try:
        candidate = _np.asarray(candidate).astype(_np.float64)
    except ValueError:
        raise ValueError("Force constant " + str(candidate) + " cannot be cast as numpy.ndarray")
    if candidate.ndim == 0:
        candidate = candidate * _np.ones(shape=(dimension,), dtype=_np.float64)
    if candidate.shape[0] != dimension:
        raise ValueError("Force constant " + str(candidate) + \
            " has shape[0]=%d instead of %d" % (candidate.shape[0], dimension))
    if candidate.ndim == 1:
        _candidate = _np.zeros(shape=(dimension, dimension), dtype=_np.float64)
        for i, x in enumerate(candidate):
            _candidate[i, i] = x
        candidate = _candidate
    elif candidate.ndim == 2:
        if candidate.shape[0] != dimension:
            raise ValueError("Force constant " + str(candidate) + \
                " has shape[1]=%d instead of %d" % (candidate.shape[1], dimension))
    else:
        raise ValueError("Force constant " + str(candidate) + " must be a 2d numpy.ndarray")
    return candidate

def _get_umbrella_sampling_parameters(
    us_trajs, us_centers, us_force_constants, md_trajs=None, kT=None):
    umbrella_centers = []
    force_constants = []
    ttrajs = []
    nthermo = 0
    unbiased_state = None
    dimension = None
    for i, traj in enumerate(us_trajs):
        state = None
        try:
            _dimension = traj.shape[1]
        except IndexError:
            _dimension = 1
        if dimension is None:
            dimension = _dimension
        else:
            if dimension != _dimension:
                raise ValueError(
                    "Trajectory %i has unmatching dimension %d!=%d" % (i, _dimension, dimension))
        this_center = _ensure_umbrella_center(
            us_centers[i], dimension)
        this_force_constant = _ensure_force_constant(
            us_force_constants[i], dimension)
        if _np.all(this_force_constant == 0.0):
            this_center *= 0.0
        for j in range(nthermo):
            if _np.all(umbrella_centers[j] == this_center) and \
                _np.all(force_constants[j] == this_force_constant):
                state = j
                break
        if state is None:
            umbrella_centers.append(this_center.copy())
            force_constants.append(this_force_constant.copy())
            ttrajs.append(nthermo * _np.ones(shape=(traj.shape[0],), dtype=_np.intc))
            nthermo += 1
        else:
            ttrajs.append(state * _np.ones(shape=(traj.shape[0],), dtype=_np.intc))
    if md_trajs is not None:
        if not isinstance(md_trajs, (list, tuple)):
            md_trajs = [md_trajs]
        this_center = umbrella_centers[-1] * 0.0
        this_force_constant = force_constants[-1] * 0.0
        for j in range(nthermo):
            if _np.all(force_constants[j] == this_force_constant):
                unbiased_state = j
                break
        if unbiased_state is None:
            umbrella_centers.append(this_center.copy())
            force_constants.append(this_force_constant.copy())
            unbiased_state = nthermo
            nthermo += 1
        for traj in md_trajs:
            ttrajs.append(unbiased_state * _np.ones(shape=(traj.shape[0],), dtype=_np.intc))
    umbrella_centers = _np.array(umbrella_centers, dtype=_np.float64)
    force_constants = _np.array(force_constants, dtype=_np.float64)
    if kT is not None:
        if not isinstance(kT, (int, float)):
            raise ValueError("kT has wrong type:" + str(type(kT)))
        if kT <= 0.0:
            raise ValueError("non-positive kT: %f" % kT)
      