from abc import ABCMeta
from ase.calculators.calculator import Calculator
import numpy as np
from scipy.linalg import solve_triangular

try:
    from mpi4py import MPI
    has_mpi = True

    COMM_WORLD = MPI.COMM_WORLD

except ImportError:
    has_mpi = False
    COMM_WORLD = None


class BaseCommitteeCalculator(Calculator, metaclass=ABCMeta):
    implemented_properties = ['energy', 'forces', 'stress', 'desc_energy', 'desc_forces', 'desc_stress', 
                              'comm_energy', 'comm_forces', 'comm_stress', 'bias_energy', 'bias_forces', 'bias_stress']
    default_parameters = {}
    name = 'BaseCommitteeCalculator'

    def __init__(self, committee_size, descriptor_size, prior_weight, energy_weight=None, forces_weight=None, stress_weight=None, 
                 sqrt_prior=None, lowmem=False, random_seed=None, regularisation=1e-4, mpi_comm=COMM_WORLD, **kwargs):
        '''
        Parameters
        ----------
        committee_size : int
            size of the committee used to approximate the posterior
        
        descriptor_size : int
            length of the descriptor vector. Should be set by child classes, rather than the end user
        
        prior_weight: float
            Weight given to the prior, relative to the likelihood, in forming the posterior
        
        energy_weight: float or None, optional
            Weight given to energy observations in forming the likelihood. If None, energy observations will be ommitted
        
        forces_weight: float or None, optional
            Weight given to force observations in forming the likelihood. If None, force observations will be ommitted
        
        stress_weight: float or None, optional
            Weight given to stress observations in forming the likelihood. If None, stress observations will be ommitted
        
        sqrt_prior : np.array of float, optional
            Square root of the prior matrix, used to form the posterior. If None, an identity matrix is used.

        lowmem : bool, optional
            Whether to use the low memory variant of solving for the posterior. Cannot be changed after initialisation.
            Defaults to False
        
        random_seed : int or np.random.RandomState object
            Seed or random state to use for all random processes.

        regularisation: float, optional
            Regularisation strength used to ensure likelihood is positive definite in the low memory variant.
            Used in a cholesky decomposition cholesky(likelihood + regularisation * np.eye(self.n_desc)) to obtain
            a square root of the likelihood. Default is 1e-4

        mpi_comm: mpi4py Comm or IntraComm object, or None
            MPI communicator. Defaults to COMM_WORLD if mpi4py detected, otherwise MPI is not used

        **kwargs: Keyword Arguments
            Extra keyword arguments sent to ase.calculators.calculator.Calculator

        '''

        self.comm = mpi_comm

        if self.comm is not None:
            if not has_mpi or type(self.comm) not in  [MPI.Comm, MPI.Intracomm]:
                raise RuntimeError("mpi_comm argument passed without a valid MPI4Py Communicator")
            
            self.rank = self.comm.Get_rank()
            self.comm_size = self.comm.Get_size()
        else:
            self.rank = 0
            self.comm_size = 1

        super().__init__(**kwargs)

        self.regularisation = regularisation

        self.n_comm = committee_size
        self.n_desc = descriptor_size

        self.selected_structures = []

        self.weights = {
             "energy": None,
             "forces": None, 
             "stress": None
        }
        self.energy_weight = energy_weight
        self.forces_weight = forces_weight
        self.stress_weight = stress_weight

        self.prior_weight = prior_weight

        if sqrt_prior is None:
            self.sqrt_prior = np.eye(self.n_desc, self.n_desc)
        else:
            self.sqrt_prior = sqrt_prior
        
        self._lowmem = lowmem

        if self._lowmem:
            self.likelihood = {key : np.zeros_like(self.sqrt_prior) for key in ["energy", "forces", "stress"]}
        else:
            self.likelihood = {key : [] for key in ["energy", "forces", "stress"]}

        if random_seed is not None:
            if type(random_seed) == np.random.RandomState:
                self.rng = random_seed
            elif type(random_seed) == int:
                self.rng = np.random.RandomState(seed=random_seed)
            else:
                raise RuntimeError("random_seed must be an integer, or a np.random.RandomState object")
        else:
            self.rng = np.random.RandomState()

        self.committee_weights = None


    @property
    def energy_weight(self):
        return self.weights["energy"]
    
    @property
    def forces_weight(self):
        return self.weights["forces"]
    
    @property
    def stress_weight(self):
        return self.weights["stress"]
    
    @energy_weight.setter
    def energy_weight(self, weight):
        if weight is not None:
            assert weight > 0

        self.weights["energy"] = weight

    @forces_weight.setter
    def forces_weight(self, weight):
        if weight is not None:
            assert weight > 0

        self.weights["forces"] = weight

    @stress_weight.setter
    def stress_weight(self, weight):
        if weight is not None:
            assert weight > 0
    
        self.weights["stress"] = weight

    def get_descriptor_energy(self, atoms=None):
        '''
        Get "descriptor energy", which is the average of descriptor vectors in the structure

        Returns an array of length self.n_desc
        '''
        return self.get_property("desc_energy", atoms)

    def get_descriptor_forces(self, atoms=None):
        '''
        Get "descriptor forces", which are the derivatives w.r.t atomic positions of the total descriptor energy

        Returns an array of shape (self.n_desc, Nats, 3)
        '''
        return self.get_property("desc_forces", atoms)

    def get_descriptor_stress(self, atoms=None):
        '''
        Get "descriptor stresses", which are the stress analogues to the total descriptor energy

        Returns an array of shape (self.n_desc, 3, 3)
        '''
        return self.get_property("desc_stress", atoms)

    def get_committee_energies(self, atoms=None):
        '''
        Get energy predictions by each member of the committee

        Returns an array of length self.n_comm
        
        '''
        return self.get_property("comm_energy", atoms)
    
    def get_committee_forces(self, atoms=None):
        '''
        Get force predictions by each member of the committee

        Returns an array of shape (self.n_comm, Nats, 3)
        
        '''
        
        return self.get_property("comm_forces", atoms)
    
    def get_committee_stresses(self, atoms=None):
        '''
        Get stress predictions by each member of the committee

        Returns an array of shape (self.n_comm, 3, 3)

        '''

        return self.get_property("comm_stress", atoms)
    
    def get_bias_energy(self, atoms=None):
        '''

        '''
        return self.get_property("bias_energy", atoms)
    
    def get_bias_forces(self, atoms=None):
        '''
        
        '''
        return self.get_property("bias_forces", atoms)
    
    def get_bias_stress(self, atoms=None):
        '''
        
        '''
        return self.get_property("bias_stress", atoms)

    def __update_likelihood_core(self, atoms, weights):
        l = {}

        props = ["desc_" + key for key in weights.keys() if weights[key] is not None]

        self.calculate(atoms, props, self.check_state(atoms))

        for key in weights.keys():
            if weights[key] is not None:
                l = np.sqrt(weights[key]) * self.results["desc_" + key].reshape(self.n_desc, -1).T

                if self._lowmem:
                    # Low memory variant
                    # Setup problem as Phi^T Phi + Prior
                    self.likelihood[key] += l.T @ l
                else:
                    # Normal variant
                    # Assemble list of all observations
                    # Maintain as list to not shift results around in memory until needed.
                    self.likelihood[key].append(l)

    def _update_likelihood(self, atoms):
        '''
        Update the likelihood based on observations obtained from atoms

        Parameters
        ----------
        atoms: ase.atoms.Atoms object
            Atoms object to derive energy, force, and stress observations from
        
        '''
        self.__update_likelihood_core(atoms, self.weights)

    def _MPI_broadcast_selection(self, atoms):
        '''
        Use MPI4Py to send a selected atoms object to all other processes, along with a 
        snapshot of the current energy, force, and stress weights
        
        Skipped if not enabled
        '''

        data = [atoms.copy(), self.weights]

        if has_mpi and self.comm is not None:
            for i in range(self.comm_size):
                if i != self.rank:
                   self.comm.isend(data, dest=i) # Non-blocking broadcast to each other process
    
    def _MPI_receive_all_selections(self):
        '''
        Use MPI4Py to receive all structures selected by other processes, along with the energy, force
        and stress weights at the time the structures were selected

        Skipped if MPI not enabled
        '''
        if not (has_mpi and self.comm is not None):
            return # MPI not enabled, so skip this step
        
        req = self.comm.irecv()

        while req.get_status(): # True when message is pending
            st, msg = req.test() # Recieve status and message via request


            self.selected_structures.append(msg[0].copy()) # Add selected structure to list
            self.__update_likelihood_core(*msg)

            req.free() # Close and open a new request, for the new message
            req = self.comm.irecv()

        req.free()
    
    def sync(self):
         if not (has_mpi and self.comm is not None):
            return # MPI not enabled, so skip this step
         
         self.comm.barrier() # MPI barrier to make sure all ranks are together
         self._MPI_receive_all_selections() # Process all MPI messages
         self.comm.barrier() # 2nd barrier to ensure all messages are recieved on all ranks
    
    def select_structure(self, atoms):
        self.selected_structures.append(atoms.copy())
        self._update_likelihood(atoms)

        self._MPI_broadcast_selection(atoms) # Send selected structure to all other processes, to be picked up later

    def resample_committee(self, committee_size=None):
        '''
        Resample the committee, based on the states of self.likelihood and self.sqrt_prior
        Populates self.committee_weights based on the newly sampled committee

        Parameters
        ----------
        committee_size : int, optional
            New size of the committee, if supplied.
            By default, a committee of size self.n_comm is drawn

        '''
        self._MPI_receive_all_selections() # Sync up with selections from other processes


        if committee_size is not None:
            self.n_comm = committee_size


        if self._lowmem:
            L_likelihood = np.linalg.cholesky(np.sum([self.likelihood[key] for key in ["energy", "forces", "stress"]]) 
                                              + self.regularisation * np.eye(self.n_desc))

            sqrt_posterior = L_likelihood + np.sqrt(self.prior_weight) * self.sqrt_prior

            Q, R = np.linalg.qr(sqrt_posterior)
        
        else:
            l_list = []

            for key in ["energy", "forces", "stress"]:
                l_key = self.likelihood[key]
                if len(l_key):
                    l_list.extend(l_key)
            sqrt_posterior = np.vstack(l_list + [self.sqrt_prior])
            Q, R = np.linalg.qr(sqrt_posterior)

        
        z = self.rng.normal(loc=0, scale=1, size=(self.n_desc, self.n_comm))

        self.committee_weights = solve_triangular(R, z, lower=False).T # zero mean committee, so no mean term

    def calculate(self, atoms=None, properties=..., system_changes=...):
        super().calculate(atoms, properties, system_changes)

        if len(system_changes):
            self.results = {}

class HALBiasPotential(BaseCommitteeCalculator, metaclass=ABCMeta):
    def _bias_energy(self, comm_energy):
        return np.std(comm_energy)
    
    def _bias_forces(self, comm_forces, comm_energy):
        Es = comm_energy - np.mean(comm_energy)
        Fs = comm_forces - np.mean(comm_forces, axis=0)

        return np.mean([E * F for E, F in zip(Es, Fs)], axis=0) / self._bias_energy(comm_energy)

    def _bias_stress(self, comm_stress, comm_energy):
        Es = comm_energy - np.mean(comm_energy)
        Ss = comm_stress - np.mean(comm_stress, axis=0)

        return np.mean([E * S for E, S in zip(Es, Ss)], axis=0) / self._bias_energy(comm_energy)