import numpy as np
from core import geodesics
from core import utils
from bayesquad.bq_wrapper import BQWrapper
import time
import copy
import dill

class LandQuadrature():
    
    def __init__(self, land):
        self.land = land
        self.integration_params = land.integration_params
        
    def estimate_norm_constant(self, K):
        pass
    

class MCQuadrature(LandQuadrature):
    # This function estimates the normalization constant using Monte Carlo sampling on the tangent space
    def estimate_norm_constant(self, k, consider_failed=False):
        if self.integration_params['method'] != 'MC':
            print("oops, this shouldn't happen. Integration parameters are probably misspecified.")
        mu = self.land.means[k]
        Sigma = self.land.sigmas[k]
        S = self.land.model_params["S"]
        start = time.time()

        D = Sigma.shape[0]
        Z_eucl = np.sqrt(((2 * np.pi) ** D) * np.linalg.det(Sigma))  # The Euclidean normalization constant

        # Initialize the matrices for the samples and the matrix A
        L, U = np.linalg.eigh(Sigma)
        A = U @ np.diag(np.sqrt(L))
        V_samples = np.zeros((S, D))
        V_samples[:] = np.nan
        X_samples = np.zeros((S, D))
        X_samples[:] = np.nan

        s = 0
        while True:
            try:
                v = A @ np.random.randn(D, 1)  # D x 1, get white noise to sample from N(0, Sigma)
                curve, failed = geodesics.expmap(self.land.manifold, x=mu.reshape(-1, 1), v=v.reshape(-1, 1))
                if not failed:
                    X_samples[s, :] = curve(1)[0].flatten()
                    V_samples[s, :] = v.flatten()
                    s = s + 1
                else:
                    print('Expmap failed for v %s' % v)
            except Exception:
                print('Expmap failed for v %s' % v)

            if s == S:  # We have collected all the samples we need
                break

        inds = np.isnan(X_samples[:, 0])  # The failed exponential maps
        X_samples = X_samples[~inds, :]  # Keep the non-failed ones
        V_samples = V_samples[~inds, :]

        volM_samples = self.land.manifold.measure(X_samples.T).flatten()  # Compute the volume element sqrt(det(M))
        norm_constant = np.mean(volM_samples) * Z_eucl  # Estimate the normalization constant

        end = time.time()
        time_mc = end - start
        print("Const: %s" % norm_constant)
        print("Runtime for MC: %s" % time_mc)


        # save some debug information
        if self.integration_params["logger"]:
            savedir = self.integration_params["savedir"]

            _means = self.land.means.copy()
            _sigmas = self.land.sigmas.copy()
            _consts = self.land.consts
            _means[k] = mu
            _sigmas[k] = Sigma
            _consts[k] = norm_constant

            res_dict = {'mu' : mu, 'Sigma' : Sigma, 'Z' : Z_eucl, 'X' : None, 'Y' : None, 'm_int' : norm_constant / Z_eucl, \
                        'x_preds' : None, 'y_preds' : None, 'pred_means' : None, 'pred_variances' : None,\
                        'lengthscale' : None, 'variance' : None,\
                        'V_samples' : V_samples, 'volM_samples' : volM_samples, 'X_samples' : X_samples,\
                        'runtime' : time_mc, 'logmaps' : self.land.logmaps[k], 'k' : k,\
                        'all_mus' : _means,\
                        'all_sigmas' : _sigmas,\
                        'all_logmaps' : self.land.logmaps,\
                        'all_consts' : _consts,\
                        'weights' : self.land.weights}
            utils.logger(res_dict, savedir)

        self.land.consts[k] = norm_constant
        self.land.V_samples[k] = V_samples
        self.land.volM_samples[:,k] = volM_samples.flatten()
        self.land.Z_eucl[k] = Z_eucl

        # we can compute the "true" manifold sigma = int v v^t exp(-0.5 <v, Gamma v>) dv
        try:
            self.land.manifold_sigmas[k] = 1/norm_constant * 1/V_samples.shape[0] * V_samples.T @ np.diag(volM_samples.flatten()) @ V_samples
            #print("true Sigma is:")
            #print(self.land.manifold_sigmas[k])
        except:
            pass

        if self.integration_params.get("save_lands"):
            # save the land object
            landcopy = self.land.getDumpableLand()
            utils.logger(landcopy, self.integration_params["savedir"] + "lands/", startswith="land_")


        # TODO: variance for MC estimator
        return norm_constant, 0 
        
        
        
        
class BQuadrature(LandQuadrature):
    # estimating the normalization constant using BQ
    def estimate_norm_constant(self, k, consider_failed=False):
        if self.integration_params['method'] != 'BQ':
            print("oops, this shouldn't happen. Integration parameters are probably misspecified.")

        mu = self.land.means[k]
        Sigma = self.land.sigmas[k]
        D = Sigma.shape[0]
        S = self.land.model_params['S']

        Z = np.sqrt(((2 * np.pi) ** D) * np.linalg.det(Sigma))  # The Euclidean normalization constant

        # these containers are needed for the gradients
        # but this requires that the BQ method uses MC inside to predict
        V_samples = np.zeros((S+1, D))
        V_samples[:] = np.nan

        volM_samples = np.zeros(S+1)
        volM_samples[:] = np.nan

        # the integrand
        # works directly on a tangent vector    
        # needs to be integrated against a gaussian with covariance matrix Sigma
        def fTangent(v): 
            curve, failed = geodesics.expmap(self.land.manifold, x=mu, v=v)  # Solve the IVP problem for expmap
            if failed:
                print("expmap failed for v: %s" % v)
                print("the failed result is %s" % curve(1)[0].reshape(-1,1))
            x = curve(1)[0].reshape(-1, 1)  # D x 1, the point on the manifold
            meas = self.land.manifold.measure(x) # The Riemannian measure |M(x)|^1/2
            if np.isnan(np.array(meas)):
                print("result is NaN")
                print("v:")
                print(v)
            return meas

        # for multiple tangent vectors at once
        def fTangent_multiple(vs): 
            return np.apply_along_axis(fTangent, 1, vs).reshape(-1,1,)


        def ExpmapCurve(v):
            curve, failed = geodesics.expmap(self.land.manifold, x=mu, v=v)  # Solve the IVP problem for expmap
            if failed:
                print("expmap failed for v: %s" % v)
                print("the failed result is %s" % curve(1)[0].reshape(-1,1))
            return curve

        start = time.time()
        fun = fTangent_multiple

        # reuse information from the last iteration
        last_mu = self.land.transfer_dict[k].get("last_mu")
        last_X = self.land.transfer_dict[k].get("last_X")
        last_Y = self.land.transfer_dict[k].get("last_Y")
        last_kernel = self.land.transfer_dict[k].get("last_kernel")

        reusing = False # indicates whether mu changed or not, we can reuse old observations if it didnt

        if self.integration_params.get("ever_reuse"):
            if last_mu is not None:
                print("can BQ reuse information?")
                print("last mu: %s " % last_mu)
                print("new mu: %s" % mu)
                if np.linalg.norm(last_mu - mu) < 1e-6:
                    print("reusing last " + str(last_X.shape[0]) + " observations.")
                    reusing = True
                    X = last_X
                    Y = last_Y
                else:
                    print("not reusing")
            else:
                print("last mu is None!")

        if True:#try:
            if not reusing:
                # if we don't reuse, then mu changed
                # but in this case, we can use the logmaps!
                # to do so, we have to evaluate the measure at the data points
                # but only use non-failed logmaps

                # if consider_failed=True, we have to check land.failed to see which
                # logmaps succeeded. this is only for the initialization of the LAND
                # afterwards, failed logmaps are NaN
                if self.integration_params.get('use_logmaps') is None or self.integration_params.get('use_logmaps'):
                    logmaps = []
                    measures = []
                    if consider_failed:
                        #print("considering failed logmaps..")
                        for il, l in enumerate(self.land.logmaps[k]):
                            if self.land.inducing_points[il] == 1 and not self.land.failed[il,k]:
                                logmaps.append(l)
                                measures.append(self.land.data_measures[il])

                    else:
                        for il, l in enumerate(self.land.logmaps[k]):
                            if self.land.inducing_points[il] == 1 and not np.isnan(l).any():
                                logmaps.append(l)
                                measures.append(self.land.data_measures[il])

                    measures = np.array([measures]).reshape(-1,1)
                    logmaps = np.vstack(logmaps).reshape(-1,D)
                    X = logmaps
                    Y = measures
                else:
                    X = np.array([])
                    Y = np.array([])

                if self.integration_params.get('verbose'):
                    print("y init: %s" % str(Y.shape))
                #last_kernel = None

            if True:
                print("last_mu is none!")
                logger = utils.logger if (self.integration_params.get('logger') is not None \
                    and self.integration_params.get('logger') is not False) else None
                w = BQWrapper(D, mu=mu, Z=Z, k=last_kernel, n_grad=S, plot=False, logger=logger, land=self.land)
                n_samples = self.integration_params["reusing_samples"] if reusing else self.integration_params["new_samples"]
                n_batches = self.integration_params["reusing_batches"] if reusing else self.integration_params["new_batches"]
                gp_mean = self.land.manifold.asymptotic_measure() if self.integration_params.get("asymptotic_mean") else 0.
                m_int, V_samples, volM_samples, int_variance = w.wsabi_integrate(fun, X, Y, n_batches, n_samples, \
                                                                variance=200., lengthscale=2.,\
                                                                prior_mean=np.zeros(D), prior_cov=Sigma, \
                                                                constant_mean = gp_mean,\
                                                                grad=True, \
                                                                integration_params=self.integration_params,\
                                                                component_k=k,
                                                                expmap=ExpmapCurve, f_manifold=self.land.manifold.measure)

            # update the last mu container
            last_mu = mu
            print("collected observations: %s" % (w.X).shape[0])
            last_X = w.X
            last_Y = w.Y

        else:#except Exception as e: 
            print("error in BQ")
            print(e)
            m_int = 0.0
            w = None

        # multiply with euclidean normalization constant
        Const = Z * m_int

        end = time.time()
        time_bq = end - start
        if self.integration_params.get("verbose"):
            print("Const: %s" % Const)
        print('Runtime for BQ : ' + str(time_bq))

        self.land.transfer_dict[k] = {'last_mu' : last_mu.copy(), 'last_X' : last_X.copy(), 'last_Y' : last_Y.copy(),\
                                      'last_kernel' : (w.k if w is not None else None)}

        # we do not have any X samples to return, thus we return None
        # but they are not used anyway
        self.land.V_samples[k,:,:] = V_samples
        self.land.volM_samples[:,k] = volM_samples.flatten()
        self.land.consts[k,:] = Const
        self.land.Z_eucl[k] = Z

        # we can compute the "true" manifold sigma = int g(v) v v^t exp(-0.5 <v, Gamma v>) dv
        #self.land.manifold_sigmas[k] = 1/Const * 1/V_samples.shape[0] * V_samples.T @ np.diag(volM_samples.flatten()) @ V_samples
        self.land.manifold_sigmas[k] = 1/Const * 1/V_samples.shape[0] * np.einsum('nd,ne,n->de', V_samples, V_samples, volM_samples.flatten())


        # save the land object
        if self.integration_params.get("save_lands"):
            landcopy = self.land.getDumpableLand()
            utils.logger(landcopy, self.integration_params["savedir"] + "lands/", startswith="land_")

        # return both the constant (mean) and its variance estimator
        return Const, int_variance
