# cython: linetrace=False
# cython: profile=False
# cython: binding=False
# cython: boundscheck=False
# cython: wraparound=False
# cython: initializedcheck=False
# cython: nonecheck=False
# cython: language_level=3
# cython: infer_types=True

import numpy as np
cimport numpy as np
import pandas as pd
import warnings
from tqdm import tqdm
from libc.math cimport log

from .PRSModel cimport PRSModel
from magenpy.stats.h2.ldsc import simple_ldsc
from viprs.utils.exceptions import OptimizationDivergence
from viprs.utils.c_utils cimport dot, vec_sum, elementwise_add_mult, sigmoid, clip
from viprs.utils.compute_utils import dict_mean, dict_sum, dict_concat, fits_in_memory


cdef class VIPRS(PRSModel):
    """
    The base class for performing Variational Inference of Polygenic Risk Scores (VIPRS).
    """

    def __init__(self, gdl, fix_params=None, load_ld='auto', tracked_theta=None, verbose=True, threads=1):
        """
        :param gdl: An instance of GWADataLoader
        :param fix_params: A dictionary of hyperparameters with their fixed values.
        :param load_ld: A flag that specifies whether to load the LD matrix to memory (Default: `auto`).
        :param tracked_theta: A list of hyperparameters to track throughout the optimization procedure. Useful
        for debugging/model checking. Currently, we allow the user to track the following:
            - The proportion of causal variants (`pi`).
            - The heritability ('heritability').
            - The residual variance (`sigma_epsilon`).
        :param verbose: Verbosity of the information printed to standard output
        :param threads: The number of threads to use (experimental)
        """

        super().__init__(gdl)

        # Variational parameters:
        self.var_gamma = {}
        self.var_mu = {}
        self.var_sigma = {}

        # Properties of proposed distribution:
        self.eta = {}  # The posterior mean, E[B] = \gamma*\mu_beta
        self.zeta = {}  # The expectation of B^2 under the posterior, E[B^2] = \gamma*(\mu_beta^2 + \sigma_beta^2)

        # q-factor (keeps track of LD-related terms)
        self.q = {}

        # ---------- Inputs to the model: ----------

        # LD-related quantities:
        self.ld = self.gdl.get_ld_matrices()
        self.ld_bounds = self.gdl.get_ld_boundaries()

        # If load_ld is set to `auto`, then determine whether to load
        # the LD matrices by examining the available memory resources:
        if load_ld == 'auto':
            self.load_ld = fits_in_memory(sum([ld.estimate_uncompressed_size() for ld in self.ld.values()]))
        else:
            self.load_ld = load_ld

        # Standardized betas:
        self.std_beta = {c: ss.get_snp_pseudo_corr() for c, ss in self.gdl.sumstats_table.items()}

        # ---------- General properties: ----------

        self.threads = threads
        self.fix_params = fix_params or {}

        self.verbose = verbose
        self.history = {}
        self.tracked_theta = tracked_theta or []

    cpdef initialize(self, theta_0=None, param_0=None):
        """
        A convenience method to initialize all the objects associated with the model.
        :param theta_0: A dictionary of initial values for the hyperparameters theta
        :param param_0: A dictionary of initial values for the variational parameters
        """

        if self.verbose:
            print("> Initializing model parameters")

        self.initialize_theta(theta_0)
        self.initialize_variational_params(param_0)
        self.init_history()

    cpdef init_history(self):
        """
        Initialize the history object to track various quantities.
        """

        self.history = {
            'ELBO': []
        }

        for tt in self.tracked_theta:
            self.history[tt] = []

    cpdef initialize_theta(self, theta_0=None):
        """
        Initialize the global hyper-parameters
        :param theta_0: A dictionary of initial values for the hyperparameters theta
        """

        if theta_0 is not None and self.fix_params is not None:
            theta_0.update(self.fix_params)
        elif self.fix_params is not None:
            theta_0 = self.fix_params
        elif theta_0 is None:
            theta_0 = {}

        # ----------------------------------------------
        # (1) If 'pi' is not set, initialize from a uniform
        if 'pi' not in theta_0:
            self.pi = np.random.uniform(low=max(0.005, 1. / self.M), high=.1)
        else:
            self.pi = theta_0['pi']

        # ----------------------------------------------
        # (2) Initialize sigma_epsilon and sigma_beta
        # Assuming that the genotype and phenotype are normalized,
        # these two quantities are conceptually linked.
        # The initialization routine here assumes that:
        # Var(y) = h2 + sigma_epsilon
        # Where, by assumption, Var(y) = 1,
        # And h2 ~= pi*M*sigma_beta

        if 'sigma_epsilon' not in theta_0:
            if 'sigma_beta' not in theta_0:

                # If neither sigma_beta nor sigma_epsilon are given,
                # then initialize using the SNP heritability estimate

                try:
                    naive_h2g = clip(simple_ldsc(self.gdl), 1e-3, 1. - 1e-3)
                except Exception as e:
                    naive_h2g = np.random.uniform(low=.001, high=.999)

                self.sigma_epsilon = 1. - naive_h2g
                self.sigma_beta = naive_h2g / (self.pi * self.M)
            else:

                # If sigma_beta is given, use it to initialize sigma_epsilon

                self.sigma_beta = theta_0['sigma_beta']
                self.sigma_epsilon = np.clip(1. - self.sigma_beta*(self.pi * self.M),
                                             a_min=1e-6, a_max=1. - 1e-6)
        else:

            # If sigma_epsilon is given, use it in the initialization

            self.sigma_epsilon = theta_0['sigma_epsilon']

            if 'sigma_beta' in theta_0:
                self.sigma_beta = theta_0['sigma_beta']
            else:
                self.sigma_beta = (1. - self.sigma_epsilon) / (self.pi * self.M)

    cpdef initialize_variational_params(self, param_0=None):
        """
        Initialize the variational parameters.
        :param param_0: A dictionary of initial values for the variational parameters
        """

        param_0 = param_0 or {}

        self.var_mu = {}
        self.var_sigma = {}
        self.var_gamma = {}

        for c, c_size in self.shapes.items():

            # Initialize the variational parameters according to the derived update equations,
            # ignoring correlations between SNPs.
            if 'sigma' in param_0:
                self.var_sigma[c] = param_0['sigma'][c]
            else:
                self.var_sigma[c] = self.sigma_epsilon / (
                        self.Nj[c] + self.sigma_epsilon / self.get_sigma_beta(c)
                )

            if 'mu' in param_0:
                self.var_mu[c] = param_0['mu'][c]
            else:
                self.var_mu[c] = np.zeros(c_size)

            if 'gamma' in param_0:
                self.var_gamma[c] = param_0['gamma'][c]
            else:
                pi = self.get_pi(c)
                if isinstance(self.pi, dict):
                    self.var_gamma[c] = pi.copy()
                else:
                    self.var_gamma[c] = pi*np.ones(c_size)

        self.eta = self.compute_eta()
        self.zeta = self.compute_zeta()

    cpdef e_step(self):
        """
        In the E-step, update the variational parameters for each SNP 
        in a coordinate-wise fashion.
        """

        # Define memoryviews objects for fast access
        cdef:
            unsigned int j, start, end, j_idx
            double u_j
            double[::1] std_beta, Dj  # Inputs
            double[::1] var_gamma, var_mu, var_sigma  # Variational parameters
            double[::1] mu_mult, u_logs, recip_sigma  # Helpers + other quantities that we need inside the for loop
            double[::1] eta, q  # Properties of proposed distribution
            long[:, ::1] ld_bound

        for c, c_size in self.shapes.items():

            # Get the priors:
            sigma_beta = self.get_sigma_beta(c)
            pi = self.get_pi(c)

            # Updates for sigma_beta variational parameters:
            self.var_sigma[c] = self.sigma_epsilon / (
                    self.Nj[c] + self.sigma_epsilon / sigma_beta
            )

            # Compute some quantities that are needed for the per-SNP updates:
            mu_mult = self.Nj[c]*self.var_sigma[c]/self.sigma_epsilon
            u_logs = np.log(pi / (1. - pi)) + .5*np.log(self.var_sigma[c] / sigma_beta)
            recip_sigma = .5/self.var_sigma[c]

            # Set the numpy vectors into memoryviews for fast access:
            std_beta = self.std_beta[c]
            var_gamma = self.var_gamma[c]
            var_mu = self.var_mu[c]
            eta = self.eta[c]
            ld_bound = self.ld_bounds[c]
            q = np.zeros(shape=c_size)

            for j, Dj in enumerate(self.ld[c]):

                start, end = ld_bound[:, j]
                j_idx = j - start

                # Compute the variational mu beta:
                var_mu[j] = mu_mult[j]*(
                        std_beta[j] - dot(Dj, eta[start: end]) + Dj[j_idx]*eta[j]
                )

                # Compute the variational gamma:
                u_j = u_logs[j] + recip_sigma[j]*var_mu[j]*var_mu[j]
                var_gamma[j] = clip(sigmoid(u_j), 1e-6, 1. - 1e-6)

                # Compute the posterior mean:
                eta[j] = var_gamma[j]*var_mu[j]

                # Update the q factors:
                if j_idx > 0:
                    # Update the q factor for snp j by adding the contribution of previous SNPs.
                    q[j] = dot(Dj[:j_idx], eta[start: j])
                    # Update the q factors for all previously updated SNPs that are in LD with SNP j
                    q[start: j] = elementwise_add_mult(q[start: j], Dj[:j_idx], eta[j])

            # Convert memoryviews back to numpy arrays:
            self.var_gamma[c] = np.asarray(var_gamma)
            self.var_mu[c] = np.asarray(var_mu)
            self.q[c] = np.asarray(q)

            self.eta[c] = np.asarray(eta)
            self.zeta[c] = self.var_gamma[c]*(self.var_mu[c]**2 + self.var_sigma[c])

    cpdef update_pi(self):
        """
        Update the prior probability of a variant being causal, pi
        """

        if 'pi' not in self.fix_params:

            # Get the average of the gammas:
            pi_estimate = dict_mean(self.var_gamma)
            # Clip value:
            self.pi = clip(pi_estimate, 1./self.M, 1. - 1./self.M)

    cpdef update_sigma_beta(self):
        """
        Update the prior variance on the effect size, sigma_beta
        """

        if 'sigma_beta' not in self.fix_params:

            # Sigma_beta estimate:
            sigma_beta_estimate = dict_sum(self.zeta) / dict_sum(self.var_gamma)
            # Clip value:
            self.sigma_beta = clip(sigma_beta_estimate, 1e-12, 1. - 1e-12)

    cpdef update_sigma_epsilon(self):
        """
        Update the residual variance, sigma_epsilon.
        """

        if 'sigma_epsilon' not in self.fix_params:

            sig_eps = 0.

            for c, c_size in self.shapes.items():

                sig_eps += (
                        - 2. * dot(self.eta[c], self.std_beta[c]) +
                        vec_sum(self.zeta[c]) +
                        dot(self.eta[c], self.q[c])
                )

            self.sigma_epsilon = clip(1. + sig_eps, 1e-12, 1. - 1e-12)

    cpdef m_step(self):
        """
        In the M-step, update the global hyperparameters of the model.
        """

        self.update_pi()
        self.update_sigma_beta()
        self.update_sigma_epsilon()

    cpdef objective(self):
        """
        Compute the variational objective, the Evidence Lower-BOund (ELBO),
        from GWAS summary statistics.
        """

        # Concatenate the dictionary items for easy computation:
        var_gamma = dict_concat(self.var_gamma)
        null_gamma = 1. - dict_concat(self.compute_pip())  # The gamma for the null component
        var_mu = dict_concat(self.var_mu)
        var_sigma = dict_concat(self.var_sigma)

        if isinstance(self.pi, dict):
            pi = dict_concat(self.pi)
        else:
            pi = self.pi

        # The pi for the null component
        try:
            null_pi = 1. - np.sum(pi, axis=1)
        except np.AxisError:
            null_pi = 1. - np.sum(pi)


        if isinstance(self.sigma_beta, dict):
            sigma_beta = dict_concat(self.sigma_beta)
        else:
            sigma_beta = self.sigma_beta

        q = dict_concat(self.q)
        eta = dict_concat(self.eta)
        zeta = dict_concat(self.zeta)

        std_beta = dict_concat(self.std_beta)

        # The ELBO is made of two components:
        elbo = 0.  # (1) log of joint density

        # -----------------------------------------------
        # (1) Compute the log of the joint density:

        #
        # (1.1) The following terms are an expansion of ||Y - X\beta||^2
        #
        # -N/2log(2pi*sigma_epsilon)
        elbo -= .5 * self.N * log(2 * np.pi * self.sigma_epsilon)

        # -Y'Y/(2*sigma_epsilon), where we assume Y'Y = N
        elbo -= .5 * (self.N / self.sigma_epsilon)

        # + (1./sigma_epsilon)*\beta*(XY), where we assume XY = N\hat{\beta}
        elbo += (self.N / self.sigma_epsilon) * dot(eta, std_beta)

        # (-1/2sigma_epsilon)\beta'X'X\beta, where we assume X_j'X_j = N
        # Note that the q factor is equivalent to X'X\beta (excluding diagonal)
        elbo -= .5 * (self.N / self.sigma_epsilon) * (dot(eta, q) + vec_sum(zeta))

        elbo -= (var_gamma * np.log(var_gamma / pi)).sum()
        elbo -= (null_gamma * np.log(null_gamma / null_pi)).sum()

        elbo += .5 * (var_gamma * (1. + np.log(var_sigma / sigma_beta) -
                                   (var_mu ** 2 + var_sigma) / sigma_beta)).sum()

        return elbo

    cpdef get_sigma_epsilon(self):
        """
        Get the value of the hyperparameter sigma_epsilon
        """
        return self.sigma_epsilon

    cpdef get_sigma_beta(self, chrom=None):
        """
        Get the value of the hyperparameter sigma_beta
        :param chrom: Get the value of `sigma_beta` for a given chromosome.
        """
        if chrom is None:
            return self.sigma_beta
        else:
            if isinstance(self.sigma_beta, dict):
                return self.sigma_beta[chrom]
            else:
                return self.sigma_beta

    cpdef get_pi(self, chrom=None):
        """
        Get the value of the hyperparameter pi
        :param chrom: Get the value of `pi` for a given chromosome.
        """

        if chrom is None:
            return self.pi
        else:
            if isinstance(self.pi, dict):
                return self.pi[chrom]
            else:
                return self.pi

    cpdef get_proportion_causal(self):
        """
        Get the proportion of causal variants for the trait.
        """
        if isinstance(self.pi, dict):
            return dict_mean(self.pi)
        else:
            return self.pi

    cpdef get_average_effect_size_variance(self):
        """
        Get the average per-SNP variance for the prior mixture components
        """
        if isinstance(self.sigma_beta, dict):
            return dict_mean(self.sigma_beta)
        else:
            return self.sigma_beta

    cpdef get_heritability(self):
        """
        Estimate the heritability, or proportion of variance explained by SNPs.
        """

        sigma_g = np.sum([
            vec_sum(self.zeta[c]) + dot(self.q[c], self.eta[c])
            for c in self.shapes
        ])

        h2g = sigma_g / (sigma_g + self.sigma_epsilon)

        return h2g

    cpdef to_theta_table(self):
        """
        Output the values for the hyperparameters (theta) to a pandas table.
        """

        theta_table = [
            {'Parameter': 'Residual_variance', 'Value': self.sigma_epsilon},
            {'Parameter': 'Heritability', 'Value':self.get_heritability()},
            {'Parameter': 'Proportion_causal', 'Value': self.get_proportion_causal()},
        ]

        sigmas = self.get_average_effect_size_variance()
        try:
            sigmas = list(sigmas)
            for i in range(len(sigmas)):
                theta_table.append({'Parameter': f'sigma_beta_{i+1}', 'Value': sigmas[i]})
        except TypeError:
            theta_table.append({'Parameter': 'sigma_beta', 'Value': sigmas})

        return pd.DataFrame(theta_table)

    cpdef write_inferred_theta(self, f_name, sep="\t"):
        """
        Write the inferred (and fixed) hyperparameters to file.
        :param f_name: The file name
        :param sep: The separator for the hyperparameter file.
        """

        # Write the table to file:
        try:
            self.to_theta_table().to_csv(f_name, sep=sep, index=False)
        except Exception as e:
            raise e

    cpdef update_theta_history(self):
        """
        For all the tracked hyperparameters set in `tracked_theta`, append 
        their current values to the history object.
        """

        for tt in self.tracked_theta:
            if tt == 'pi':
                self.history['pi'].append(self.get_proportion_causal())
            if tt == 'heritability':
                self.history['heritability'].append(self.get_heritability())
            if tt == 'sigma_epsilon':
                self.history['sigma_epsilon'].append(self.sigma_epsilon)

    cpdef compute_pip(self):
        """
        Compute the posterior inclusion probability
        """
        try:
            return {c: v.sum(axis=1) for c, v in self.var_gamma.items()}
        except np.AxisError:
            return self.var_gamma

    cpdef compute_eta(self):
        """
        Compute the mean for the effect size under the variational posterior.
        """
        try:
            return {c: (v*self.var_mu[c]).sum(axis=1) for c, v in self.var_gamma.items()}
        except np.AxisError:
            return {c: v*self.var_mu[c] for c, v in self.var_gamma.items()}

    cpdef compute_zeta(self):
        """
        Compute the expectation of the squared effect size under the variational posterior.
        """

        try:
            return {c: (v * (self.var_mu[c] ** 2 + self.var_sigma[c])).sum(axis=1)
                    for c, v in self.var_gamma.items()}
        except np.AxisError:
            return {c: (v * (self.var_mu[c] ** 2 + self.var_sigma[c]))
                    for c, v in self.var_gamma.items()}

    cpdef update_posterior_moments(self):
        """
        Update the posterior objects containing the posterior moments,
        including the PIP and posterior mean and variance for the effect size.
        """

        self.pip = self.compute_pip()
        self.post_mean_beta = {c: eta.copy() for c, eta in self.eta.items()}
        self.post_var_beta = {c: zeta - self.eta[c]**2 for c, zeta in self.zeta.items()}

    cpdef fit(self, max_iter=1000, theta_0=None, param_0=None,
              continued=False, f_abs_tol=1e-3, x_rel_tol=1e-3, max_elbo_drops=10):
        """
        Fit the model parameters to data.

        :param max_iter: Maximum number of iterations. 
        :param theta_0: A dictionary of values to initialize the hyperparameters
        :param theta_0: A dictionary of values to initialize the variational parameters
        :param continued: If true, continue the model fitting for more iterations.
        :param f_abs_tol: The absolute tolerance threshold for the objective (ELBO)
        :param x_rel_tol: The relative tolerance threshold for the parameters (mean beta)
        :param max_elbo_drops: The maximum number of times the objective is allowed to drop before termination.
        """

        if not continued:
            self.initialize(theta_0, param_0)
            start_idx = 1
        else:
            start_idx = len(self.history['ELBO']) + 1

        if self.load_ld:

            if self.verbose:
                print("> Loading LD matrices into memory...")

            self.gdl.load_ld()

        elbo_dropped_count = 0
        converged = False

        if self.verbose:
            print("> Performing model fit...")
            print(f"> Using up to {self.threads} threads.")

        for i in tqdm(range(start_idx, start_idx + max_iter), disable=not self.verbose):

            self.update_theta_history()

            self.e_step()
            self.m_step()

            self.history['ELBO'].append(self.objective())

            if i > 1:

                curr_elbo = self.history['ELBO'][i - 1]
                prev_elbo = self.history['ELBO'][i - 2]

                if curr_elbo < prev_elbo:
                    elbo_dropped_count += 1
                    warnings.warn(f"Iteration {i}: ELBO dropped from {prev_elbo:.6f} "
                                  f"to {curr_elbo:.6f}.")

                    if elbo_dropped_count > max_elbo_drops:
                        warnings.warn("The optimization is halted due to numerical instabilities!")
                        break

                    continue

                if np.isclose(prev_elbo, curr_elbo, atol=f_abs_tol, rtol=0.):
                    print(f"Converged at iteration {i} || ELBO: {curr_elbo:.6f}")
                    converged = True
                    break
                elif all([np.allclose(v, self.post_mean_beta[c], atol=0., rtol=x_rel_tol)
                          for c, v in self.eta.items()]):
                    print(f"Converged at iteration {i} | ELBO: {curr_elbo:.6f}")
                    converged = True
                    break

                if abs((curr_elbo - prev_elbo) / prev_elbo) > 1. and abs(curr_elbo - prev_elbo) > 1e3:
                    raise OptimizationDivergence(f"Stopping at iteration {i}: "
                                                 f"The optimization algorithm is not converging!\n"
                                                 f"Previous ELBO: {prev_elbo:.6f} | "
                                                 f"Current ELBO: {curr_elbo:.6f}")
                elif self.get_heritability() >= 1.:
                    raise OptimizationDivergence(f"Stopping at iteration {i}: "
                                                 f"The optimization algorithm is not converging!\n"
                                                 f"Value of estimated heritability exceeded 1.")
            self.update_posterior_moments()

        if converged:
            self.update_posterior_moments()
        elif i - start_idx == max_iter - 1:
            warnings.warn("Maximum iterations reached without convergence. "
                          "You may need to run the model for more iterations.")

        if self.verbose:
            print(f"> Final ELBO: {self.history['ELBO'][len(self.history['ELBO'])-1]:.6f}")
            print(f"> Estimated heritability: {self.get_heritability():.6f}")
            print(f"> Estimated proportion of causal variants: {self.get_proportion_causal():.6f}")

        return self
