"""Experiments for world3 simulator."""
import numpy as np

from whynot.framework import DynamicsExperiment, parameter
from whynot.simulators import world3


def get_experiments():
    """Return all experiments for world3."""
    return [PollutionRCT,
            PollutionConfounding,
            PollutionUnobservedConfounding,
            PollutionMediation]


def sample_initial_states(rng):
    """Sample initial state by randomly perturbing the default state."""
    def random_scale(scale, base=10):
        """Make 10**(-scale) to 10**scale smaller/bigger uniformly."""
        return rng.choice(np.logspace(-scale, scale, base=base, num=50))

    state = world3.State()
    state.population_0_to_14 *= random_scale(0.33)
    state.population_15_to_44 *= random_scale(0.33)
    state.population_45_to_64 *= random_scale(0.33)
    state.population_65_and_over *= random_scale(0.33)
    state.industrial_capital *= random_scale(0.33)
    state.service_capital *= random_scale(0.33)
    state.arable_land *= random_scale(0.33)
    state.potentially_arable_land *= random_scale(0.33)
    state.urban_industrial_land *= random_scale(0.33)
    state.nonrenewable_resources *= random_scale(0.33)
    state.persistent_pollution *= random_scale(0.33)
    state.land_fertility *= random_scale(0.26)

    return state


##################
# RCT Experiment
##################
# pylint: disable-msg=invalid-name
#: An RCT experiment to study the effect of decreases in pollution generation on population.
PollutionRCT = DynamicsExperiment(
    name="world3_rct",
    description="Study effect of intervening in 1975 to decrease pollution generation.",
    simulator=world3,
    simulator_config=world3.Config(persistent_pollution_generation_factor=1.),
    intervention=world3.Intervention(year=1975, persistent_pollution_generation_factor=0.85),
    state_sampler=sample_initial_states,
    propensity_scorer=0.5,
    outcome_extractor=lambda run: run[2050].total_population,
    covariate_builder=lambda run: run.initial_state.values())


##########################
# Confounding Experiments
##########################
@parameter(name="treatment_bias", default=0.8, values=np.linspace(0.5, 0.05, 1),
           description="Bias of probability of treatment between top and bottom pollution runs.")
def pollution_confounded_propensity(intervention, control_runs, treatment_bias):
    """Probability of treating each unit.

    To generate confounding, we are more likely to treat worlds with high pollution.
    """
    def persistent_pollution(run):
        return run[intervention.year].persistent_pollution

    pollution = [persistent_pollution(run) for run in control_runs]
    upper_quantile = np.quantile(pollution, 0.9)

    def treatment_prob(idx):
        if pollution[idx] > upper_quantile:
            return treatment_bias
        return 1. - treatment_bias

    return np.array([treatment_prob(idx) for idx in range(len(control_runs))])


# pylint: disable-msg=invalid-name
#: An observational experiment with confounding. Polluted states are more likely to be treated.
PollutionConfounding = DynamicsExperiment(
    name="world3_pollution_confounding",
    description=("Study effect of intervening to decrease pollution.  Confounding "
                 "arises becauses states with high pollution are more likely "
                 "to receive treatment."),
    simulator=world3,
    simulator_config=world3.Config(persistent_pollution_generation_factor=1.),
    intervention=world3.Intervention(year=1975, persistent_pollution_generation_factor=0.85),
    state_sampler=sample_initial_states,
    propensity_scorer=pollution_confounded_propensity,
    outcome_extractor=lambda run: run[2050].total_population,
    covariate_builder=lambda intervention, run: run[intervention.year].values())


# pylint: disable-msg=invalid-name
#: An observational experiment with unobserved confounding.
PollutionUnobservedConfounding = DynamicsExperiment(
    name="world3_pollution_unobserved_confounding",
    description=("Study effect of intervening to decrease pollution.  Confounding "
                 "arises becauses states with high pollution are more likely "
                 "to receive treatment. However, only variables that directly affect "
                 "treatment assignment are observed."),
    simulator=world3,
    simulator_config=world3.Config(persistent_pollution_generation_factor=1.),
    intervention=world3.Intervention(year=1975, persistent_pollution_generation_factor=0.85),
    state_sampler=sample_initial_states,
    propensity_scorer=pollution_confounded_propensity,
    outcome_extractor=lambda run: run[2050].total_population,
    covariate_builder=lambda intervention, run: np.array([
        run[intervention.year].persistent_pollution]))


########################
# Mediation Experiments
########################
@parameter(name="mediation_year", default=2015, values=[1980, 2000, 2020, 2045],
           description="What year to get mediating variables")
@parameter(name="num_mediators", default=11, values=[0, 3, 5, 7, 10],
           description="number of mediators to include.")
def mediation_covariates(intervention, run, mediation_year, num_mediators):
    """Build the causal dataset."""
    # Just use all of the states at the moment of treatment assignment
    confounders = run[intervention.year].values()
    mediators = run[mediation_year].values()[:num_mediators]
    return np.concatenate([confounders, mediators])


# pylint: disable-msg=invalid-name
#: An observational experiment with mediation from states after intervention.
PollutionMediation = DynamicsExperiment(
    name="world3_pollution_mediation",
    description=("Study effect of intervening to decrease pollution.  Confounding "
                 "arises becauses states with high pollution are more likely "
                 "to receive treatment. Mediation arises since future states "
                 "are also observed."),
    simulator=world3,
    simulator_config=world3.Config(persistent_pollution_generation_factor=1.),
    intervention=world3.Intervention(year=1975, persistent_pollution_generation_factor=0.85),
    state_sampler=sample_initial_states,
    propensity_scorer=pollution_confounded_propensity,
    outcome_extractor=lambda run: run[2050].total_population,
    covariate_builder=mediation_covariates)
