from ctypes import cdll, CDLL, RTLD_GLOBAL
from ctypes import POINTER, byref, c_int, c_int64, c_int32, c_bool, c_char_p, c_double, c_void_p, CFUNCTYPE, py_object, cast, byref, Structure
import ctypes

from scipy.linalg import block_diag
from ase.data import chemical_symbols
from ase.geometry import get_distances
import ase, os, warnings
import numpy as np
import numpy.typing as npt
from typing import List, Any
from scipy.io import netcdf_file
from asi4py.pyasi import matrix_asi_to_numpy
from scipy.optimize import LinearConstraint, minimize


def free_energy_fun(x, E, beta):
  '''
    Free energy
  '''
  return np.sum(x * (E*beta - np.log(x)))

def free_energy_jac(x, E, beta):
  '''
    Free energy gradient
  '''
  return E*beta - np.log(x)-1

def get_elems_ordered(atoms):
  '''
    List elements in the order of tiers.
    https://gitlab.com/ase/ase/-/commit/7b70eddd026154d636faf404cc2f8c7b08d89667
    https://mail.python.org/pipermail/python-dev/2017-December/151283.html
  '''
  return list(dict.fromkeys(atoms.symbols))

def split_dm(dm, magmom:float, iS:int):
  '''
    Split spin-paired density matrix into up/down spin-polarizes DM
  '''
  w,v = np.linalg.eigh(dm)
  w /= 2
  w_positive_indices = np.where(w>0)[0]
  E = -np.log(w[w_positive_indices])
  n_avg = np.sum(w)
  n_up   = n_avg + magmom/2
  n_down = n_avg - magmom/2
  b_eq = n_up if iS==1 else n_down
  w_new = minimize(lambda x: free_energy_fun(x, E, 0.1) , w[w_positive_indices], jac=lambda x: free_energy_jac(x, E, 0.1), bounds=[(1e-10, 1.0) for _ in E], constraints=[LinearConstraint(np.ones((1, len(E))), b_eq, b_eq), ] ).x
   
  w[w_positive_indices] = w_new
  return v @ np.diag(w) @ v.T

def build_dm_free_atoms(atoms, elem_dms, n_spin:int, iS:int):
  '''
    Build atoms density matrix from single-atomic matrices for each element
    :param n_spin: number of spin channels. 1 or 2.
    :param iS: index of the spin channel. 1 or 2.
  '''
  assert n_spin in [1,2], n_spin
  
  atomic_dms = [elem_dms[s] for s in atoms.symbols]
  if n_spin==2:
    atomic_dms = [split_dm(dm,magmom,iS) for dm,magmom in zip(atomic_dms, atoms.get_initial_magnetic_moments())]

  return block_diag(*atomic_dms).T

def load_atomic_dm(elem:str, path:str=None):
  if path is None:
    path = os.environ['ASI_FREE_ATOM_DMS']
  elemZ = chemical_symbols.index(elem)
  with netcdf_file(f'{path}/{elemZ:03d}_{elem}.ncdf', mmap=False) as f:
    return f.variables['DM'].data

def save_atomic_dm(dm, elem:str, path:str=None):
  '''
    Save density matrix for element `elem` into file `{path}/{elemZ:03d}_{elem}.ncdf`
    If `path` is `None` then it defaults to environment variable `ASI_FREE_ATOM_DMS`
    This files maybe used by :class:`PredictFreeAtoms` predictor
  '''
  if path is None:
    path = os.environ['ASI_FREE_ATOM_DMS']
  elemZ = chemical_symbols.index(elem)
  with netcdf_file(f'{path}/{elemZ:03d}_{elem}.ncdf','w', mmap=False) as f:
    f.createDimension('n_basis', dm.shape[0])
    f.createVariable('DM', dm.dtype, ('n_basis', 'n_basis'))[:,:] = dm

def bool2int_selector(atoms_selector):
  atoms_selector = np.array(atoms_selector)
  if np.issubdtype(atoms_selector.dtype, bool):
    atoms_selector = np.where(atoms_selector)[0]
  assert np.issubdtype(atoms_selector.dtype, np.integer)
  return atoms_selector

class PredictDMByAtoms:
  '''
    This is a base class for density matrix predictors. It is meant to be subclassed.
    Subclasses must overwrite the method :func:`PredictDMByAtoms.__call__` that takes :class:`ase.Atoms` object as argument
    and returns :class:`numpy.ndarray` with a density matrix guess.
  '''
  def __init__(self):
    pass
  
  def register_DM_init(self, asi):
    '''
      Register this density matrix predictor as an ASI callback for density matrix initialization
      
      :param asi: an instance of :class:`asi4py.pyasi.ASILib`
      
    '''
    self.asi = asi
    asi.register_DM_init(PredictDMByAtoms.dm_init_callback, self)

  def dm_init_callback(self, iK, iS, blacs_descr, data, matrix_descr_ptr):
    '''
      ASI DM init callback. Not to be invoked directly.
    '''
    self = cast(self, py_object).value
    n_basis = self.asi.n_basis
    m = self(self.asi.atoms, iK, iS) if self.asi.scalapack.is_root(blacs_descr) else None
    
    assert m is None or (m.shape == (n_basis, n_basis)), \
                     f"m.shape=={m.shape} != n_basis=={n_basis}"
    self.asi.scalapack.scatter_numpy(m, blacs_descr, data)
    return 1

  def __call__(self, atoms, iK:int=1, iS:int=1):
    '''
      This method is meant to be overwritten by derived classes.
    '''
    raise RuntimeError("Not implemented in base class")
    #return build_dm_free_atoms(atoms, self.elem_dms)

class PredictFreeAtoms(PredictDMByAtoms):
  '''
    Density matrix predictor that uses single-atomic density matrices for initialization.
    See :class:`PredictFreeAtoms.__init__` for details
  '''
  def __init__(self, elem_dms:dict[str, Any]=None, elem_dm_path:str=None,  n_spin:int=1):
    '''
      If ``elem_dms`` parameter is not ``None``, then ``elem_dm_path`` should be ``None`.
      If both ``elem_dms`` and ``elem_dm_path`` are ``None``, then  ``elem_dm_path` value
      defaults to the ``ASI_FREE_ATOM_DMS`` environment variable.
      
      :param elem_dms: a dictionary that maps chemical elements to its density matrices
      :param elem_dm_path: a path to a folder that contains `*.npz` files with density matrices of chemical elements
      :param n_spin:  number of spin channels, should be 2 if initial magnetic moment is specified
    '''
    super().__init__()
    self.n_spin = n_spin
    assert (elem_dms is None) or (elem_dm_path is None)
    self.elem_dms = elem_dms
    if self.elem_dms is None:
      self.elem_dm_path = elem_dm_path if elem_dm_path is not None else os.environ['ASI_FREE_ATOM_DMS']

  def __call__(self, atoms, iK:int=1, iS:int=1):
    if iK != 1:
      return np.zeros((self.asi.n_basis, self.asi.n_basis))
    if self.elem_dms is None:
      assert self.elem_dm_path is not None
      self.elem_dms = {elem:load_atomic_dm(elem, self.elem_dm_path) for elem in get_elems_ordered(atoms)}
    return build_dm_free_atoms(atoms, self.elem_dms, self.n_spin, iS)

class PredictConstAtoms(PredictDMByAtoms):
  '''
    Density matrix predictor that returns a density matrix passed to its constructor, independently from
    atomic coordinates.
  '''
  def __init__(self, const_atoms:ase.Atoms, const_dm:Any):
    '''
      :param const_atoms: :class:`ase.Atoms` object that is only used for checking order of chemical elements on density matrix prediction
      :param const_dm: a density matrix that will be returned as initial guess on SCF loop initialization
    '''
    super().__init__()
    self.const_atoms = const_atoms
    self.const_dm = const_dm

  def __call__(self, atoms, iK:int=1, iS:int=1):
    np.testing.assert_allclose(atoms.numbers, self.const_atoms.numbers)
    
    np.testing.assert_allclose(
      atoms.positions - atoms.get_center_of_mass(), 
      self.const_atoms.positions - self.const_atoms.get_center_of_mass(), 
      atol=1e-1, rtol=1e-1)
    
    return self.const_dm

def select_basis_indices(all_basis_atoms, atoms_indices):
  '''
    Returns indices of basis functions for selected atoms.
    
    :param all_basis_atoms: atomic indices for all basis functions
    :param atoms_indices: indices of selected atoms
  '''
  return np.where(np.any(all_basis_atoms[None,:] == atoms_indices[:, None], axis=0))[0]

class PredictFrankensteinDM(PredictDMByAtoms):
  '''
    Density matrix predictor that "stitches" full density matrix prediction from 
    predictions of multiples predictors of smaller, possibly overlapping subsystems of the full system.
  '''
  def __init__(self, predictors_and_selectors:list[tuple[PredictDMByAtoms, list[int]]]):
    '''
      :param predictors_and_selectors: list of pairs. The first pair element is callable, that 
        should return predicted density matrix for a subsystem; it can be derived from :class:`PredictDMByAtoms`.
        The second pair element is the list of atomc that constitute the subsystem.
        
    '''
    # unzip https://stackoverflow.com/a/12974504/3213940
    self.predictors, atoms_selectors = list(zip(*predictors_and_selectors))
    self.atoms_groups_indices = list(map(bool2int_selector, atoms_selectors))

    if True: # extended assertion check
      all_selected_atoms_set = set().union(*self.atoms_groups_indices)
      max_range_atoms_set = set(range(max(all_selected_atoms_set) + 1))
      missed_atoms = max_range_atoms_set - all_selected_atoms_set
      # Heuristic check: doesn't guarantie full atoms coverage, because actual
      # number of atoms is not known here and max(all_selected_atoms_set) is 
      # just an heuristic
      assert len(missed_atoms) == 0, f"Missed atoms: {missed_atoms}"
  
  def register_DM_init(self, asi):
    super().register_DM_init(asi)
    self.init_basis_indices(asi)

  def init_basis_indices(self, asi):
    '''
      Loads indices of basis functions for subsystems from an ASI library.
      It should only be called if one is going to predict density matrix without registering 
      this predictor via :func:`PredictDMByAtoms.register_DM_init`
    '''
    all_basis_atoms = asi.basis_atoms
    self.basis_indices = [select_basis_indices(all_basis_atoms, atoms_group) for atoms_group in self.atoms_groups_indices]
    self.n_basis = asi.n_basis

    if True:
      all_selected_basis_indices = set().union(*self.basis_indices)
      total_basis_set = set(range(self.n_basis))
      missed_basis_functions = total_basis_set - all_selected_basis_indices
      assert len(missed_basis_functions) == 0, f'Basis functions missed from selection: {missed_basis_functions}'

  def __call__(self, atoms, iK:int=1, iS:int=1):
    assert len(self.predictors) == len(self.atoms_groups_indices)
    assert len(self.predictors) == len(self.basis_indices)

    total_dm = np.zeros((self.n_basis, self.n_basis), dtype=np.float64)
    total_dm_cnt = np.zeros(total_dm.shape, dtype=int)
    for predictor, atoms_group_indices, basis_group_indices in zip(self.predictors, self.atoms_groups_indices, self.basis_indices):
      total_dm[basis_group_indices[np.newaxis,:], basis_group_indices[:, np.newaxis]] += predictor(atoms[atoms_group_indices])
      total_dm_cnt[basis_group_indices[np.newaxis,:], basis_group_indices[:, np.newaxis]] += 1
    
    assert (total_dm[total_dm_cnt==0]==0).all()
    total_dm_cnt[total_dm_cnt==0] = 1 # to avoid division by zero

    return (total_dm / total_dm_cnt).T

class PredictSeqDM(PredictDMByAtoms):
  '''
    Extrapolates density matrix using kernel ridge regression from a few previous geometries.
    This predictor uses another predictor as a baseline and extrapolates only deviations of 
    baseline predictors from ground-state density matrices.
    The method :func:`PredictSeqDM.update_errs` should be invoked after SCF convergence to provide
    the predictor with ground-state density matrix, for extrapolation model adjustment.
    
  '''
  class PredictionSequence:
    def __init__(self):
      self.preds = []
      self.descrs = []
      self.errs = []

    def predict_err(self, predicted_dm, descr):
      from sklearn.gaussian_process.kernels import RBF
      from sklearn.kernel_ridge import KernelRidge
      k = len(self.preds)
      assert k == len(self.descrs), f'{k=} {len(self.descrs)=}'
      assert k == len(self.errs), f'{k=} {len(self.errs)=}'
      if k == 0:
        return np.zeros(predicted_dm.shape)

      X = np.hstack([np.array(self.preds).reshape((k, -1)), np.array(self.descrs).reshape((k, -1))])
      x = np.hstack([predicted_dm.ravel(), descr.ravel()]).reshape((1, -1))
      Y = np.array(self.errs).reshape((k, -1))
      

      X_mean = np.mean(X, axis=0)
      Y_mean = np.mean(Y, axis=0)
      
      X -= X_mean
      Y -= Y_mean
      x -= X_mean
      #------------------

      with warnings.catch_warnings():
        warnings.filterwarnings("ignore", message="Ill-conditioned matrix")
        reg = KernelRidge(alpha=1e-15, kernel='rbf')
        reg.fit(X, Y)
        y = reg.predict(x)

      #------------------
      y += Y_mean
      P = y.reshape(predicted_dm.shape)
      return P
      
  
  def __init__(self, base_predictor:PredictDMByAtoms, n_hist:int):
    '''
      :param base_predictor: Baseline predictor, for example :class:`PredictFreeAtoms` or :class:`PredictFrankensteinDM` 
      :param n_hist: history size for extrapolation
    '''
    self.base_predictor = base_predictor
    self.predseq = {} # dict (iK,iS):PredictionSequence
    self.n_hist = n_hist
  
  def __call__(self, atoms, iK:int=1, iS:int=1):
    predicted_dm = self.base_predictor(atoms, iK, iS)

    R,d = get_distances(atoms.positions)
    lowtri = np.tri(len(atoms), len(atoms), -1)==1
    invd = 1/d[lowtri]
    descr = invd

    if (iK, iS) not in self.predseq:
      self.predseq[(iK, iS)] = PredictSeqDM.PredictionSequence()

    predicted_err = self.predseq[(iK, iS)].predict_err(predicted_dm, descr)

    if len(self.predseq[(iK, iS)].preds) == self.n_hist:
      self.predseq[(iK, iS)].preds.pop(0)
      self.predseq[(iK, iS)].descrs.pop(0)
      self.predseq[(iK, iS)].errs.pop(0)
    self.predseq[(iK, iS)].preds.append(predicted_dm)
    self.predseq[(iK, iS)].descrs.append(descr)

    return (predicted_dm - predicted_err).T

  def register_DM_init(self, asi):
    '''
      Register this density matrix predictor as an ASI callback for density matrix initialization,
      and registers itself to process DM evaluation callback for errors update.
      
      :param asi: an instance of :class:`asi4py.pyasi.ASILib`
      
    '''
    super().register_DM_init(asi)
    asi.dm_calc_cnt = {}
    asi.dm_storage = {}
    asi.register_dm_callback(PredictSeqDM.dm_callback, self)

  def dm_callback(self, iK, iS, descr, data, matrix_descr_ptr):
    self = cast(self, py_object).value
    self.asi.dm_calc_cnt[(iK, iS)] = self.asi.dm_calc_cnt.get((iK, iS),0) + 1
    data = matrix_asi_to_numpy(self.asi, descr, data, matrix_descr_ptr)
    if data is not None:
      self.asi.dm_storage[(iK, iS)] = data.copy()
  
  def update_errs(self):
    '''
      This method should be called after SCF convergence, to adjust extrapolation model.
    '''
    for (iK, iS),pred_item in self.predseq.items():
      assert len(pred_item.preds) == len(pred_item.errs) + 1
      pred_item.errs.append(pred_item.preds[-1] - self.asi.dm_storage[(iK, iS)])
      self.asi.dm_calc_cnt[(iK, iS)] = 0
    
    self.register_DM_init(self.asi)
  




