from typing import Union, Iterable, List
from collections import namedtuple
from scipy.stats import gaussian_kde
import numpy as np
import pandas as pd

import dpdata

default_kB = 0.0083144621  # Boltzmann constant in kJ/mol/K
default_ev_to_kjmol = 96.4853365  # conversion factor from eV to kJ/mol


_FesResult = namedtuple('FesResult', ['fes', 'grid', 'extend'])

def compute_fes(cvs: np.ndarray, bias: np.ndarray, temp: float, grid=None,
                w=1, kB=default_kB, grid_size=100j):
    """
    Compute free energy surface from biased sampling data using Gaussian KDE
    Support 1D and 2D collective variables

    :param cv: collective variable
    :param bias: bias potential
    :param temp: temperature
    :param grid: grid points, if None, will use grid_size to generate automatically
    :param w: weights
    :param kB: Boltzmann constant
    :param grid_size: grid size in np.mgrid style

    :return: free energy
    """
    if cvs.ndim == 2:
        if cvs.shape[0] == 1:
            cvs = cvs[0]  # flatten
        elif cvs.shape[0] != 2:
            raise ValueError(f'Invalid cvs shape {cvs.shape}, only support 1D or 2D cvs')
    elif cvs.ndim != 1:
        raise ValueError(f'Invalid cvs shape {cvs.shape}, only support 1D or 2D cvs')

    kBT = kB * temp
    beta = 1 / kBT
    weights = w * np.exp(beta * bias)
    kde = gaussian_kde(cvs, weights=weights)

    extend = None
    if cvs.ndim == 1:  # handle 1d
        if grid is None:
            x_min, x_max = cvs.min(), cvs.max()
            grid = np.mgrid[x_min:x_max:grid_size]
            extend = [x_min, x_max]
        pdf = kde.evaluate(grid)
    else:  # handle 2d
        if grid is None:
            x_min, x_max = cvs[0].min(), cvs[0].max()
            y_min, y_max = cvs[1].min(), cvs[1].max()
            grid = np.mgrid[x_min:x_max:grid_size, y_min:y_max:grid_size]
            extend = [x_min, x_max, y_min, y_max]
        X, Y = grid
        pos = np.vstack([X.ravel(), Y.ravel()])
        # TODO: generated by copilot, I am not sure if it is correct
        pdf = kde.evaluate(pos).T.reshape(X.shape)
    fes = -kBT * np.log(pdf)
    # return fes with grid and extend for plotting
    return _FesResult(fes=fes, grid=grid, extend=extend)




def compute_kde_weight(baseline_energy: np.ndarray, target_energy: np.ndarray,
                       temp: float, kB=default_kB,
                       ev_to_kjmol=default_ev_to_kjmol,):
    """
    Compute KDE weights for reweighting

    :param baseline_energy: baseline energy in eV, it's the output of DeepMD evaluation
    :param target_energy: target energy in eV, it's the output of DeepMD evaluation
    """
    kBT = kB * temp
    beta = 1 / kBT

    # use relative energy and convert from eV to kJ/mol
    be = (baseline_energy- np.min(baseline_energy)) * ev_to_kjmol
    te = (target_energy- np.min(target_energy)) * ev_to_kjmol
    return np.exp( -beta * (te - be))



class ReweightTool:

    @classmethod
    def from_dpdata_system(cls, dp_sys):
        ...


    def __init__(self):
        ...