import numpy as np
from autumn.calibration.proposal_tuning import perform_all_params_proposal_tuning
from autumn.core.project import Project, ParameterSet, load_timeseries, build_rel_path, get_all_available_scenario_paths, \
    use_tuned_proposal_sds
from autumn.calibration import Calibration
from autumn.calibration.priors import UniformPrior, BetaPrior,TruncNormalPrior
from autumn.calibration.targets import (
    NormalTarget,
    get_dispersion_priors_for_gaussian_targets,
)
from autumn.models.covid_19 import base_params, build_model
from autumn.settings import Region, Models
from autumn.projects.covid_19.sri_lanka.sri_lanka.scenario_builder import get_all_scenario_dicts

# Load and configure model parameters.
default_path = build_rel_path("params/default.yml")
#scenario_paths = [build_rel_path(f"params/scenario-{i}.yml") for i in range(7, 9)]
mle_path = build_rel_path("params/mle-params.yml")
baseline_params = base_params.update(default_path).update(mle_path, calibration_format=True)
all_scenario_dicts = get_all_scenario_dicts("LKA")
#scenario_params = [baseline_params.update(p) for p in scenario_paths]
scenario_params = [baseline_params.update(sc_dict) for sc_dict in all_scenario_dicts]
param_set = ParameterSet(baseline=baseline_params, scenarios=scenario_params)

ts_set = load_timeseries(build_rel_path("timeseries.json"))
notifications_ts = ts_set["notifications"].rolling(7).mean().loc[350::7]
death_ts = ts_set["infection_deaths"].loc[350:]
targets = [
    NormalTarget(notifications_ts),
    NormalTarget(death_ts),
]

priors = [
    # Dispersion parameters based on targets
    *get_dispersion_priors_for_gaussian_targets(targets),
    *get_dispersion_priors_for_gaussian_targets(targets),
    # Regional parameters
    UniformPrior("contact_rate", [0.024, 0.027]),
    UniformPrior("infectious_seed", [275.0, 450.0]),
    # Detection
    UniformPrior("testing_to_detection.assumed_cdr_parameter", [0.009, 0.025]),
    UniformPrior("infection_fatality.multiplier", [0.09, 0.13]),
    #VoC
    UniformPrior("voc_emergence.alpha_beta.start_time", [370, 410]),
    UniformPrior("voc_emergence.alpha_beta.contact_rate_multiplier", [3.2, 4.5]),
    UniformPrior("voc_emergence.delta.start_time", [475, 530]),
    UniformPrior("voc_emergence.delta.contact_rate_multiplier", [8.5, 11.5]),
]

# Load proposal sds from yml file
# use_tuned_proposal_sds(priors, build_rel_path("proposal_sds.yml"))

calibration = Calibration(priors, targets)

# FIXME: Replace with flexible Python plot request API.
import json

plot_spec_filepath = build_rel_path("timeseries.json")
with open(plot_spec_filepath) as f:
    plot_spec = json.load(f)


project = Project(
    Region.SRI_LANKA, Models.COVID_19, build_model, param_set, calibration, plots=plot_spec
)

#perform_all_params_proposal_tuning(project, calibration, priors, n_points=50, relative_likelihood_reduction=0.2)
[((738, 774), 'autumn.core.project.build_rel_path', 'build_rel_path', (['"""params/default.yml"""'], {}), "('params/default.yml')\n", (752, 774), False, 'from autumn.core.project import Project, ParameterSet, load_timeseries, build_rel_path, get_all_available_scenario_paths, use_tuned_proposal_sds\n'), ((870, 909), 'autumn.core.project.build_rel_path', 'build_rel_path', (['"""params/mle-params.yml"""'], {}), "('params/mle-params.yml')\n", (884, 909), False, 'from autumn.core.project import Project, ParameterSet, load_timeseries, build_rel_path, get_all_available_scenario_paths, use_tuned_proposal_sds\n'), ((1024, 1053), 'autumn.projects.covid_19.sri_lanka.sri_lanka.scenario_builder.get_all_scenario_dicts', 'get_all_scenario_dicts', (['"""LKA"""'], {}), "('LKA')\n", (1046, 1053), False, 'from autumn.projects.covid_19.sri_lanka.sri_lanka.scenario_builder import get_all_scenario_dicts\n'), ((1223, 1288), 'autumn.core.project.ParameterSet', 'ParameterSet', ([], {'baseline': 'baseline_params', 'scenarios': 'scenario_params'}), '(baseline=baseline_params, scenarios=scenario_params)\n', (1235, 1288), False, 'from autumn.core.project import Project, ParameterSet, load_timeseries, build_rel_path, get_all_available_scenario_paths, use_tuned_proposal_sds\n'), ((2435, 2463), 'autumn.calibration.Calibration', 'Calibration', (['priors', 'targets'], {}), '(priors, targets)\n', (2446, 2463), False, 'from autumn.calibration import Calibration\n'), ((2555, 2588), 'autumn.core.project.build_rel_path', 'build_rel_path', (['"""timeseries.json"""'], {}), "('timeseries.json')\n", (2569, 2588), False, 'from autumn.core.project import Project, ParameterSet, load_timeseries, build_rel_path, get_all_available_scenario_paths, use_tuned_proposal_sds\n'), ((2666, 2766), 'autumn.core.project.Project', 'Project', (['Region.SRI_LANKA', 'Models.COVID_19', 'build_model', 'param_set', 'calibration'], {'plots': 'plot_spec'}), '(Region.SRI_LANKA, Models.COVID_19, build_model, param_set,\n    calibration, plots=plot_spec)\n', (2673, 2766), False, 'from autumn.core.project import Project, ParameterSet, load_timeseries, build_rel_path, get_all_available_scenario_paths, use_tuned_proposal_sds\n'), ((1315, 1348), 'autumn.core.project.build_rel_path', 'build_rel_path', (['"""timeseries.json"""'], {}), "('timeseries.json')\n", (1329, 1348), False, 'from autumn.core.project import Project, ParameterSet, load_timeseries, build_rel_path, get_all_available_scenario_paths, use_tuned_proposal_sds\n'), ((1487, 1517), 'autumn.calibration.targets.NormalTarget', 'NormalTarget', (['notifications_ts'], {}), '(notifications_ts)\n', (1499, 1517), False, 'from autumn.calibration.targets import NormalTarget, get_dispersion_priors_for_gaussian_targets\n'), ((1523, 1545), 'autumn.calibration.targets.NormalTarget', 'NormalTarget', (['death_ts'], {}), '(death_ts)\n', (1535, 1545), False, 'from autumn.calibration.targets import NormalTarget, get_dispersion_priors_for_gaussian_targets\n'), ((1752, 1796), 'autumn.calibration.priors.UniformPrior', 'UniformPrior', (['"""contact_rate"""', '[0.024, 0.027]'], {}), "('contact_rate', [0.024, 0.027])\n", (1764, 1796), False, 'from autumn.calibration.priors import UniformPrior, BetaPrior, TruncNormalPrior\n'), ((1802, 1849), 'autumn.calibration.priors.UniformPrior', 'UniformPrior', (['"""infectious_seed"""', '[275.0, 450.0]'], {}), "('infectious_seed', [275.0, 450.0])\n", (1814, 1849), False, 'from autumn.calibration.priors import UniformPrior, BetaPrior, TruncNormalPrior\n'), ((1871, 1945), 'autumn.calibration.priors.UniformPrior', 'UniformPrior', (['"""testing_to_detection.assumed_cdr_parameter"""', '[0.009, 0.025]'], {}), "('testing_to_detection.assumed_cdr_parameter', [0.009, 0.025])\n", (1883, 1945), False, 'from autumn.calibration.priors import UniformPrior, BetaPrior, TruncNormalPrior\n'), ((1951, 2010), 'autumn.calibration.priors.UniformPrior', 'UniformPrior', (['"""infection_fatality.multiplier"""', '[0.09, 0.13]'], {}), "('infection_fatality.multiplier', [0.09, 0.13])\n", (1963, 2010), False, 'from autumn.calibration.priors import UniformPrior, BetaPrior, TruncNormalPrior\n'), ((2025, 2088), 'autumn.calibration.priors.UniformPrior', 'UniformPrior', (['"""voc_emergence.alpha_beta.start_time"""', '[370, 410]'], {}), "('voc_emergence.alpha_beta.start_time', [370, 410])\n", (2037, 2088), False, 'from autumn.calibration.priors import UniformPrior, BetaPrior, TruncNormalPrior\n'), ((2094, 2170), 'autumn.calibration.priors.UniformPrior', 'UniformPrior', (['"""voc_emergence.alpha_beta.contact_rate_multiplier"""', '[3.2, 4.5]'], {}), "('voc_emergence.alpha_beta.contact_rate_multiplier', [3.2, 4.5])\n", (2106, 2170), False, 'from autumn.calibration.priors import UniformPrior, BetaPrior, TruncNormalPrior\n'), ((2176, 2234), 'autumn.calibration.priors.UniformPrior', 'UniformPrior', (['"""voc_emergence.delta.start_time"""', '[475, 530]'], {}), "('voc_emergence.delta.start_time', [475, 530])\n", (2188, 2234), False, 'from autumn.calibration.priors import UniformPrior, BetaPrior, TruncNormalPrior\n'), ((2240, 2312), 'autumn.calibration.priors.UniformPrior', 'UniformPrior', (['"""voc_emergence.delta.contact_rate_multiplier"""', '[8.5, 11.5]'], {}), "('voc_emergence.delta.contact_rate_multiplier', [8.5, 11.5])\n", (2252, 2312), False, 'from autumn.calibration.priors import UniformPrior, BetaPrior, TruncNormalPrior\n'), ((2641, 2653), 'json.load', 'json.load', (['f'], {}), '(f)\n', (2650, 2653), False, 'import json\n'), ((928, 960), 'autumn.models.covid_19.base_params.update', 'base_params.update', (['default_path'], {}), '(default_path)\n', (946, 960), False, 'from autumn.models.covid_19 import base_params, build_model\n'), ((1611, 1662), 'autumn.calibration.targets.get_dispersion_priors_for_gaussian_targets', 'get_dispersion_priors_for_gaussian_targets', (['targets'], {}), '(targets)\n', (1653, 1662), False, 'from autumn.calibration.targets import NormalTarget, get_dispersion_priors_for_gaussian_targets\n'), ((1669, 1720), 'autumn.calibration.targets.get_dispersion_priors_for_gaussian_targets', 'get_dispersion_priors_for_gaussian_targets', (['targets'], {}), '(targets)\n', (1711, 1720), False, 'from autumn.calibration.targets import NormalTarget, get_dispersion_priors_for_gaussian_targets\n')]