%load_ext autoreload
%autoreload 2

import sys
sys.path.append("../src/")
sys.path.append("..")

from src.examples.example_setups import setup_stat_scm
from src.utils.sequential_sampling import sequentially_sample_model
from src.utils.sem_utils.toy_sems import StationaryDependentSEM as StatSEM
from src.utils.plotting import plot_outcome
from src.utils.sem_utils.sem_estimate import build_sem_hat

from numpy.random import seed

# Models
from src.methods.bo import BO
from src.methods.cbo import CBO
from src.methods.abo import ABO
from src.methods.dcbo import DCBO

seed(seed=0)

"""
Explanation: Stationary SCM and DAG (STAT.) from figure 1 in paper
End of explanation
"""


T = 3 # In this example, as in the paper, we consider three time-slices
init_sem, sem, dag_view, G, exploration_sets, intervention_domain, true_objective_values  = setup_stat_scm(T=T)
# This is just a snapshot of the DAG at hand
dag_view

"""
Explanation: Sample SEM to get observational samples
End of explanation
"""


type(G)

"""
Explanation: The above DAG is the graphical structure we will be working with and is faithful to the one used in figure one in the paper. dag is a networkx object i.e.
End of explanation
"""


print(dag_view)

"""
Explanation: and is the one we will use for exploring the optimization methods further down. dag_view is simply a convenient way to visualise the structure at hand and is just a string that we visulise using pygraphviz:
End of explanation
"""


# Number of independent samples (samples here are the time-series on the horizontal) per time-index
D_O = sequentially_sample_model(init_sem,
                                sem,
                                total_timesteps=T,
                                sample_count=5, # How many samples we take per node in each time-slice
                                epsilon=None) # If we employ a noise model or not

D_O

# Simulate some interventional data

from src.utils.utilities import powerset, get_monte_carlo_expectation
from src.utils.sequential_intervention_functions import make_sequential_intervention_dict
from numpy import zeros

D_I = {k: None for k in powerset(["X", "Z"])}
"""
do(Z_0)
"""
interv = make_sequential_intervention_dict(G,T)
# Univariate intervention at time 0
interv["Z"][0] = 1.0
static_noise_model = {k: zeros(T) for k in ["X", "Z", "Y"]}
# Sample this model with one intervention, and no noise
intervention_samples = sequentially_sample_model(
    init_sem,
    sem,
    total_timesteps=T,
    interventions=interv,
    sample_count=1,
    epsilon=static_noise_model
)

# Note: we obvs don't need to take MC samples if noise-model is zero
D_I[("Z",)] = get_monte_carlo_expectation(intervention_samples)


"""
Explanation: Notice that we make use of two types of structural equations models: init_sem and sem. The latter concerns interactions with the first time-slice in the DBN, which has no incoming edges from the previous time-slices, and is only active at $t=0$. For all other time-slices i.e. when $t>0$ the sem model is used.
End of explanation
"""


# Contains the exploration sets we will be investigating
print("Exploration sets:", exploration_sets)
# The intervention domains for the manipulative variables
print("Intervention domains:", intervention_domain) 
# The true outcome values of Y given an optimal intervention on the three time-slices
print("True optimal outcome values:", [r"y^*_{} = {}".format(t,val.round(3)) for t,val in enumerate(true_objective_values)])
# Number of trials
N = 10

"""
Explanation: The above observation samples contained in D_obs are sequential in the sense that the sample function generates each column per variable, conditional on the column samples at $t-1$.
Other setup parameters
End of explanation
"""


# Setup the input parameters to be used with the Bayesian optimization
BO_input_params = {
    "G":G,
    "sem": StatSEM,
    "base_target_variable": "Y",
    "observation_samples": D_O, # Observational samples
    "intervention_domain": intervention_domain,
    "intervention_samples":None,   
    "number_of_trials": N,
    "optimal_assigned_blankets": None,
    "sample_anchor_points": True,
    "seed_anchor_points": 1,
    "change_points": None
}
# Run Bayesian optimization
bo = BO(**BO_input_params)
bo.run()

plot_outcome(T,N,outcomes=[bo.optimal_outcome_values_during_trials],labels=['BO'],true_objective_values=true_objective_values)

"""
Explanation: Explore optimization methods
In all these examples we do not employ any interventional data, just observational.
Vanilla Bayesian Optimization
End of explanation
"""


abo = ABO(**BO_input_params)
abo.run()

plot_outcome(T,N,outcomes=[bo.optimal_outcome_values_during_trials,abo.optimal_outcome_values_during_trials],labels=['BO', 'ABO'],true_objective_values=true_objective_values)

"""
Explanation: Adaptive Bayesian Optimization
The setup for ABO is the same as for BO.
End of explanation
"""


CBO_input_params = {
    "G": G,
    "sem": StatSEM,
    "base_target_variable": "Y",
    "make_sem_estimator": build_sem_hat,
    "exploration_sets":exploration_sets,
    "observation_samples": D_O,
    "intervention_samples":D_I, # Interventional data as well
    "intervention_domain": intervention_domain,
    "intervention_samples":None,
    "number_of_trials": N,
    "sample_anchor_points": True,
    "seed_anchor_points": 1,
}
cbo = CBO(**CBO_input_params)
cbo.run()

plot_outcome(T,N,outcomes=[bo.optimal_outcome_values_during_trials, abo.optimal_outcome_values_during_trials, cbo.optimal_outcome_values_during_trials],labels=['BO', 'ABO', 'CBO'],true_objective_values=true_objective_values)

"""
Explanation: Causal Bayesian Optimization
CBO has one extra parameter (though there are many others which we are not demonstrating in this demo) which is the SEM estimator method build_sem_hat. It estimates a SEM model from the observational data.
End of explanation
"""


dcbo = DCBO(**CBO_input_params)
dcbo.run()

plot_outcome(T,N,outcomes=[bo.optimal_outcome_values_during_trials, abo.optimal_outcome_values_during_trials, cbo.optimal_outcome_values_during_trials, dcbo.optimal_outcome_values_during_trials],labels=['BO', 'ABO', 'CBO', 'DCBO'],true_objective_values=true_objective_values)

"""
Explanation: Dynamic Causal Bayesian Optimization
DCBO has the same input parameters as CBO. But they work very differently under the hood. DCBO takes 'horizontal' information into account i.e. is able to transfer information between time-slices. CBO has no notion of time.
End of explanation
"""


