import arviz as az
import pymc3 as pm
import numpy as np
import matplotlib.pyplot as plt

"""
Explanation: Refitting PyMC3 models with ArviZ
ArviZ is backend agnostic and therefore does not sample directly. In order to take advantage of algorithms that require refitting models several times, ArviZ uses {class}~arviz.SamplingWrapper to convert the API of the sampling backend to a common set of functions. Hence, functions like Leave Future Out Cross Validation can be used in ArviZ independently of the sampling backend used.
Below there is an example of SamplingWrapper usage for PyMC3.
Before starting, it is important to note that PyMC3 cannot modify the shapes of the input data using the same compiled model. Thus, each refitting will require a recompilation of the model.
End of explanation
"""


np.random.seed(26)

xdata = np.linspace(0, 50, 100)
b0, b1, sigma = -2, 1, 3
ydata = np.random.normal(loc=b1 * xdata + b0, scale=sigma)

plt.plot(xdata, ydata);

"""
Explanation: For the example we will use a linear regression model.
End of explanation
"""


def compile_linreg_model(xdata, ydata):
    with pm.Model() as model:
        b0 = pm.Normal("b0", 0, 10)
        b1 = pm.Normal("b1", 0, 10)
        sigma_e = pm.HalfNormal("sigma_e", 10)

        y = pm.Normal("y", b0 + b1 * xdata, sigma_e, observed=ydata)
    return model

sample_kwargs = {"draws": 500, "tune": 500, "chains": 4}
with compile_linreg_model(xdata, ydata) as linreg_model:
    trace = pm.sample(**sample_kwargs)

"""
Explanation: Now we will write the PyMC3 model, keeping in mind the following two points:
1. Data must be modifiable (both x and y).
2. The model must be recompiled in order to be refitted with the modified data.
We, therefore, have to create a function that recompiles the model when it's called.
Note how we would generally use pm.Data to populate xyz_data groups in the {ref}InferenceData &lt;xarray_for_arviz&gt; result, however, given that we need to recompile the model every time, we are consciously skipping this step to speed up the recompilation.
End of explanation
"""


dims = {"y": ["time"]}
idata_kwargs = {
    "dims": dims,
}
idata = az.from_pymc3(trace, model=linreg_model, **idata_kwargs)

"""
Explanation: We have defined a dictionary sample_kwargs that will be passed to the SamplingWrapper in order to make sure that all refits use the same sampler parameters. 
We follow the same pattern with {func}az.from_pymc3 &lt;arviz.from_pymc3&gt;. 
Note, however, how coords are not set. This is done to prevent errors due to coordinates and values shapes being incompatible during refits. Otherwise, we'd have to handle subsetting of the coordinate values even though the refits are never used outside the refitting functions such as {func}~arviz.reloo.
We also exclude the model because the model, like the trace, is different for every refit. This may seem counterintuitive or even plain wrong, but we have to remember that the pm.Model object contains information like the observed data.
End of explanation
"""


class PyMC3LinRegWrapper(az.SamplingWrapper):
    def __init__(self, xdata, ydata, **kwargs):
        super(PyMC3LinRegWrapper, self).__init__(**kwargs)

        self.xdata = xdata
        self.ydata = ydata
        
    def sample(self, modified_observed_data):
        with self.model(**modified_observed_data) as linreg_model:
            trace = pm.sample(
                **self.sample_kwargs, 
                return_inferencedata=False, 
                idata_kwargs={"log_likelihood": False}
            )
        self.pymc3_model = linreg_model
        return trace
    
    def get_inference_data(self, trace):
        idata = az.from_pymc3(trace, model=self.pymc3_model, **self.idata_kwargs)
        idata.pymc3_trace = trace
        return idata
        
    def log_likelihood__i(self, excluded_observed_data, idata__i):
        model_ex = compile_linreg_model(**excluded_observed_data)
        log_lik__i = az.from_pymc3(idata__i.pymc3_trace, model=model_ex).log_likelihood["y"]
        return log_lik__i
        
    def sel_observations(self, idx):
        mask = np.isin(np.arange(len(self.xdata)), idx)
        data_dict = {"xdata": self.xdata, "ydata": self.ydata}
        data__i = {key: value[~mask] for key, value in data_dict.items()}
        data_ex = {key: value[mask] for key, value in data_dict.items()}
        return data__i, data_ex

loo_orig = az.loo(idata, pointwise=True)
loo_orig

"""
Explanation: We will create a subclass of az.SamplingWrapper.
End of explanation
"""


loo_orig.pareto_k[[13, 42, 56, 73]] = np.array([0.8, 1.2, 2.6, 0.9])

"""
Explanation: In this case, the Leave-One-Out Cross Validation (LOO-CV) approximation using Pareto Smoothed Importance Sampling (PSIS) works for all observations, so we will use modify loo_orig in order to make az.reloo believe that PSIS failed for some observations. This will also serve as a validation of our wrapper, as the PSIS LOO-CV already returned the correct value.
End of explanation
"""


pymc3_wrapper = PyMC3LinRegWrapper(
    model=compile_linreg_model, xdata=xdata, ydata=ydata, sample_kwargs=sample_kwargs, idata_kwargs=idata_kwargs
)

"""
Explanation: We initialize our sampling wrapper
End of explanation
"""


loo_relooed = az.reloo(pymc3_wrapper, loo_orig=loo_orig)

loo_relooed

loo_orig

"""
Explanation: And eventually, we can use this wrapper to call az.reloo, and compare the results with the PSIS LOO-CV results.
End of explanation
"""


