import copy as cp
import jax.numpy as jnp
import jax.random as jrnd
import jax.scipy.special as jspecial
import joblib as jl
import numpy as np
import numpyro
import numpyro.contrib.control_flow as noflow
import numpyro.diagnostics as ndiag
import numpyro.distributions as ndist
import numpyro.infer as ninfer
import pandas as pd
import time as tm
import utils


class ALMetaRisk:
    """
    The Additive Linear MetaRisk Model
    """

    @staticmethod
    def load(path, random_seed: int = None, show_progress: bool = False):
        """
        Loads from File

        :param path: Joblib file to load
        :param random_seed: Seed for model
        :param show_progress: Progress Bar
        :return: Initialised ALMetaRisk object
        """
        state = jl.load(path)
        mdl = ALMetaRisk(**state['args'], random_seed=random_seed, show_progress=show_progress)
        mdl.__posterior = state['posterior']
        return mdl

    def __init__(
            self,
            # ==== Data ==== #
            risk_model: str,
            # ==== Priors ==== #
            lambda_mv: tuple,
            sex_init_ab: dict,
            init_var: float,
            tran_mv: dict,
            cov_mvs: dict,
            phi_ab: tuple = (0.01, 10),
            nu_ab: tuple = (0.01, 2),
            # ==== MCMC Parameters ==== #
            inferer: tuple = ('NUTS', {}),
            num_warmup: int = 1000,
            num_samples: int = 1000,
            num_chains: int = 1,
            # ==== Other ==== #
            random_seed: int = None,
            show_progress: bool = False,
    ):
        """
        Initialises the ALMetaRisk Model

        :param risk_model: str -- Name of the risk model
        :param lambda_mv: (mu, sigma) -- Mean and Variance for the effect of the exposure (Lambda)
        :param sex_init_ab: {Sex: (a, b)} -- Lower and Upper bounds for the uniform over the starting risk per Sex
        :param init_var: float -- Variance for the Gaussian distribution over the initial R^*
        :param tran_mv: {Sex: (mu, sigma)} -- Mean and Variance for the Sex-specific disease progression dynamics.
        :param cov_mvs: {Cov: ((mu_1, sigma_1), (mu_2, sigma_2)} -- Mean and Variance for the Stage 1 and 2 observations per covariate
        :param phi_ab: (a, b) -- Lower and Upper bounds for the variance of the noise in the risk model
        :param nu_ab: (a, b) -- Lower and Upper bounds for the variance of the disease transition model
        :param inferer: two-tuple initialiser -- Initialiser for the Inference algorithm
        :param num_warmup: int -- Number of warmup steps in each MCMC Chain
        :param num_samples: int -- Number of samples to use in each MCMC Chain
        :param num_chains: int -- Total number of chains to run.
        :param random_seed: int -- Random seed
        :param show_progress: bool -- If True, show bars during model fitting
        """
        # Do some Checks
        if not utils.ensure_positive(lambda_mv, (1,)):
            raise ValueError('Variance over Exposure effect must be positive')
        if not utils.ensure_positive(sex_init_ab):
            raise ValueError('It does not make sense to have negative values for sex-specific initialisation.')
        if init_var < 0:
            raise ValueError('Initialisation variance must be positive')
        if not utils.ensure_positive(phi_ab):
            raise ValueError('Variance over Risk Model noise must be positive')
        if not utils.ensure_positive(tran_mv, (1,)):
            raise ValueError('Transition Variance must be positive')
        if not utils.ensure_positive(cov_mvs, (1,)):
            raise ValueError('Covariate-effect Variances must be positive')
        if not utils.ensure_positive(nu_ab):
            raise ValueError('Variance over Transition model must be positive')

        # Risk Model
        self.__rm = risk_model

        # Get Sex Ordering as this is crucial in defining one-hot-encoding and correct ordering of dicts.
        self.__sex_order = tuple(sex_init_ab.keys())

        # Get Set of Covariates
        self.__covariates = tuple(cov_mvs.keys())

        # Build Prior
        self.__prior = (
            cp.deepcopy(sex_init_ab),
            init_var,
            cp.deepcopy(phi_ab),
            cp.deepcopy(lambda_mv),
            cp.copy(tran_mv),
            cp.deepcopy(cov_mvs),
            cp.deepcopy(nu_ab)
        )

        # Setup MCMC Sampler
        self.__mcmc = ninfer.MCMC(
            ninfer.__dict__[inferer[0]](model=self.__define_model, **inferer[1]),
            num_warmup=num_warmup,
            num_samples=num_samples,
            num_chains=num_chains,
            progress_bar=show_progress,
        )
        self.__key = jrnd.key(int(tm.time()) if random_seed is None else random_seed)

        # Currently not fit
        self.__posterior = None

    def fit(self, Xy):
        """
        Fits the ALMetaRisk model

        :param Xy: DataFrame containing the Data.
                * Index: Double-level by 'Subject' and 'Time'
                * Columns: [cov_Sex, cov_Age, cov_Exposure, cov_{COV}, risk_{MDL}]
        :return: self, for chaining
        """
        # Re-Initialise
        self.__posterior = None

        # Run Model
        self.__mcmc.run(
            self.__key,
            data=self.prepare_data(Xy),
            priors=cp.deepcopy(self.__prior),
            sex_order=self.__sex_order,
            extra_fields=("~z.r*_0", "~z.r*_t")
        )

        # Update status and return
        self.__posterior = {
            'Samples': cp.deepcopy(self.__mcmc.get_samples(False)),
            'Diagnostics': ndiag.summary(self.__mcmc.get_samples(True))
        }
        return self

    def predict(self, X):
        """
        Predict the Latent State

        :param X: DataFrame containing the Data.
                * Index: Double-level by 'Subject' and 'Time'
                * Columns: [cov_Sex, cov_Age, cov_Exposure, cov_{COV}]
        :return: Numpy Array of Latent State over time [N x T x Samples] as Prob Distribution: T is the largest time-span,
        """
        # Format Data
        data = self.prepare_data(X)

        # Generate Samples over Latent State
        predictor = ninfer.Predictive(self.__define_model, self.get_posteriors())
        l = predictor(self.__key, data=data, priors=self.__prior, sex_order=self.__sex_order)

        # Format and Invalidate beyond valid
        l = np.array(jnp.concatenate((l['r*_0'][:, jnp.newaxis, :], l['r*_t']), axis=1).swapaxes(0, 2))
        for i, v in enumerate(data['v']):
            l[i, v:, :] = np.nan

        # Return Array
        return l

    def log_likelihood(self, X):
        """
        Returns the (evidence) log-likelihood for the data

        :param X: DataFrame containing the Data.
                * Index: Double-level by 'Subject' and 'Time'
                * Columns: [cov_Sex, cov_Age, cov_Exposure, cov_{COV}]
        :return: Per-Sample (Normalised) log-likelihood
        """
        # Format Data
        d = self.prepare_data(X)
        p = self.get_posteriors()

        # Get Latent Posterior on this Data
        inferer = ninfer.Predictive(self.__define_model, posterior_samples=p, return_sites=[*p.keys(), 'r*_0', 'r*_t'])
        preds = inferer(jrnd.key(10), data=d, priors=self.__prior, sex_order=self.__sex_order)

        # Compute individual Log-Likelihoods and invalidate observations beyond valid
        ll = ninfer.log_likelihood( self.__define_model, preds, data=d, priors=self.__prior, sex_order=self.__sex_order)
        r_t = np.array(ll['r_t'])
        for i, v in enumerate(d['v']): r_t[i, v-1:, :] = 0

        # Resolve as single normalised value
        ll = r_t.sum(axis=1) + ll['r_0']  # Resolve probabilities over time
        ll = (jspecial.logsumexp(ll, axis=0) - np.log(len(ll)))  # Resolve average over posterior samples
        return ll.mean()  # Return per-subject mean

    def save(self, path: str):
        """
        Persist the model to disk as a joblib file

        :param path: Path of file to save
        :return: self, for chaining
        """
        arguments = {
            'risk_model': self.__rm,
            'sex_init_ab': self.__prior[0],
            'init_var': self.__prior[1],
            'phi_ab': self.__prior[2],
            'lambda_mv': self.__prior[3],
            'tran_mv': self.__prior[4],
            'cov_mvs': self.__prior[5],
            'nu_ab': self.__prior[6],
            'num_warmup': self.__mcmc.num_warmup,
            'num_samples': self.__mcmc.num_samples,
            'num_chains': self.__mcmc.num_chains,
        }
        # Save Model
        jl.dump({'args': arguments, 'posterior': self.__posterior}, path, compress=True)
        # Return self
        return self

    def get_posteriors(self, key=None):
        """
        Get the Posterior distribution over the parameters

        :param key: if specified, index directly into parameter
        :return:
        """
        if self.__posterior is not None:
            return cp.copy(self.__posterior['Samples'] if key is None else self.__posterior['Samples'][key])
        else:
            raise RuntimeError('Data must be fit successfully first')

    def get_convergence_diagnostics(self):
        """
        Returns Convergence Statistics for the MCMC run

        :return: Pandas DataFrame
        """
        return pd.DataFrame(self.__posterior['Diagnostics']).T

    @property
    def exposure_effect(self):
        """
        Access to posterior over lambda parameter
        """
        return self.get_posteriors('lambda')

    def prepare_data(self, df):
        """
        Prepares the dataframe for injesting by the model

        @param df: DataFrame containing the Time-Series
        """
        return {
            's': jnp.asarray(pd.get_dummies(df.cov_Sex.groupby('Subject').first())[list(self.__sex_order)].astype(int).T),
            'a': jnp.asarray(df.cov_Age.groupby('Subject').first()),
            'v': jnp.asarray(df.groupby('Subject').size()),
            'e': jnp.asarray(df.cov_Exposure.unstack(0).fillna(0)),
            'c': {k: self.__extract_covariate(df[f'cov_{k}']) for k in self.__covariates},
            'r': jnp.asarray(df[f'risk_{self.__rm}'].unstack(0).fillna(0)),
        }

    @staticmethod
    def __extract_covariate(cov):
        # First Setup Historical Observations
        cov = pd.concat([cov * 2, (cov.groupby('Subject').cumsum() > 0).astype(int)], axis=1).max(axis=1)
        # Expand 1-Hot Encoding and get common lengths by unstacking
        cov = pd.Series(pd.Categorical(cov, categories=(0, 1, 2)), index=cov.index)
        cov = pd.get_dummies(cov)[[0, 1, 2]].astype(int).unstack(1).fillna(0)
        # Convert To Numpy and Reshape
        return jnp.asarray(cov.to_numpy().reshape(len(cov), 3, -1)).T

    @staticmethod
    def __define_model(data: dict, priors: tuple, sex_order: tuple):
        """
        Static method that defines the model DAG

        :param data: Dictionary of data parameters (all Jax.numpy arrays)
                * s: [2, N] Sex of the subjects (one-hot encoded)
                * a: [N] Age of the subject at initialisation (before first observation)
                * v: [N] Number of valid observations (T_n) for this subject
                * e: [T, N] Per-subject exposure to the event over time-points T
                * c: {Cov: [T, 3, N]} Per-Covariate per-subject value over time (one-hot encoded)
                * r: [T, N] Per-subject risk model over time-points T
        :param priors: Tuple of prior parameters. These are a type-checked ordering of the passed priors to the
                    initializer, in the following order: ab_psi, sigma_0, ab_phi, mv_lambda, mv_omega_s, mv_omega_i and
                    ab_nu.
        :param sex_order: crucial for one-hot encoding
        :return: None
        """
        # Get Sizes and common data values and resolve priors
        T, N = data['e'].shape
        sex, covs = data['s'], data['c']
        ab_psi, sigma_0, ab_phi, mv_lambda, mv_omega_s, mv_omega_i, ab_nu = priors

        # Start Defining Prior distributions over parameters
        psi_p = jnp.array([numpyro.sample(f'psi_{i}', ndist.Uniform(*ab_psi[i])) for i in sex_order])
        phi_p = numpyro.sample('phi', ndist.Uniform(*ab_phi))
        lambda_p = numpyro.sample('lambda', ndist.Normal(*mv_lambda))
        omega_s_p = jnp.array([numpyro.sample(f'omega_s_{i}', ndist.Normal(*mv_omega_s[i])) for i in sex_order])
        omega_c_p = {
            c: jnp.array([0, *[numpyro.sample(f'omega_{c}_{i+1}', ndist.Normal(*v)) for i, v in enumerate(cc)]])
            for c, cc in mv_omega_i.items()
        }  # Note that we add a 0 for when there is no observation (state 0)
        nu_p = numpyro.sample('nu', ndist.Uniform(*ab_nu))

        # Now define the Initial Distribution and observations therein
        with numpyro.plate('N', N):
            r_init = numpyro.sample('r*_0', ndist.TruncatedNormal(psi_p @ sex - data['a'], sigma_0, low=0))
            # r_init = numpyro.sample('r*_0', ndist.Normal(psi_p @ sex - data['a'], sigma_0))
            _ = numpyro.sample('r_0', ndist.Normal(r_init, phi_p), obs=data['r'][0, :])

        # Define Transition Dynamics & run in Time
        def _dynamics(state, t):
            r_prev = state # Keep Track of Previous State
            with numpyro.plate('N', N):
                with numpyro.handlers.mask(mask=t<data['v']):
                    state_mean = r_prev + omega_s_p @ sex + lambda_p * data['e'][t, :]
                    for cov, obs in covs.items():
                        state_mean += omega_c_p[cov] @ obs[t,:,:]
                    r_curr = numpyro.sample('r*_t', ndist.Normal(state_mean, nu_p))
                    _ = numpyro.sample('r_t', ndist.Normal(r_curr, phi_p), obs=data['r'][t, :])
            return r_curr, None
        _, _ = noflow.scan(_dynamics, r_init, jnp.arange(1, T))


# if __name__=='__main__':
#     import runpy
#     import os
#
#     def run_inference(args):
#         cfg, pth, name, seed = args
#         # Prepare
#         os.makedirs(os.path.join(pth, name[0]), exist_ok=True)  # Create Directory
#         df = pd.read_pickle(os.path.join(pth, 'Data.bz2'))  # Read Data
#         # Infer Model
#         mdl = globals()[cfg[0]](**cfg[1], random_seed=seed)  # Initialise
#         mdl.fit(df).save(os.path.join(pth, name[0], name[1] + '.jl'))
#
#
#     # Paths
#     DATA_PATH = '../../../../Data/Simulations/Test'  # Data Path
#     MODEL_PATH = '../Configs/Models/'  # Model Configs
#     INFER_PATH = '../Configs/Inference/'  # Inference Configs
#
#     # Selection of Configurations
#     DATA_CONFIGS = ('LA_C1_Obs',)
#     PARAM_CONFIGS = (
#         'V_Exposure',
#         'V_MaxFollowup',
#         'V_DiseaseSD',
#         'V_RiskScale',
#         'V_RiskSD',
#         'V_ComorbidEffect',
#         'V_ComorbidHeal',
#         'V_ComorbidInit',
#         'V_ComorbidMiss',
#         'R_Comorbidity',
#         'R_RiskModel',
#     )
#     MODEL_CONFIGS = ('ALMR_Base',)
#     INFER_CONFIGS = ('NUTS_1_1_5',)
#
#     # Randomness
#     RANDOM_SEED = 18122025
#
#     # Runtime
#     NUM_JOBS = 10
#
#     for dcfg in DATA_CONFIGS:
#         for pcfg in PARAM_CONFIGS:
#             print('======================================================================')
#             print(f'  {dcfg}/{pcfg}')
#             # Load Simulator config
#             data_cfg = runpy.run_path(os.path.join(DATA_PATH, dcfg, pcfg, 'Experiment.cfg'))['CONFIG']['EXPERIMENT']
#             for mcfg in MODEL_CONFIGS:
#                 for icfg in INFER_CONFIGS:
#                     print(' -------------------------------------------------------------------')
#                     print(f'    {mcfg}/{icfg}')
#                     # Load Model Configuration
#                     mdl_cfg = runpy.run_path(os.path.join(MODEL_PATH, mcfg + '.py'))['MODEL']
#                     infer_cfg = runpy.run_path(os.path.join(INFER_PATH, icfg + '.py'))['INFERENCE']
#                     mdl_cfg[1].update(infer_cfg)
#                     # Build Job List
#                     jobs = [
#                         (mdl_cfg, os.path.join(DATA_PATH, dcfg, pcfg, f'ParameterSet_{p}', f'Replica_{r}'),
#                          (mcfg, icfg), RANDOM_SEED)
#                         for p in range(data_cfg['PARAMETER_SAMPLES'])
#                         for r in range(data_cfg['REPLICAS'])
#                     ]
#                     for job in jobs:
#                         run_inference(job)
#                     print(' -------------------------------------------------------------------')
#             print('======================================================================')
