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)
[((16, 15, 16, 51), 'autumn.core.project.build_rel_path', 'build_rel_path', ({(16, 30, 16, 50): '"""params/default.yml"""'}, {}), "('params/default.yml')", False, 'from autumn.core.project import Project, ParameterSet, load_timeseries, build_rel_path, get_all_available_scenario_paths, use_tuned_proposal_sds\n'), ((18, 11, 18, 50), 'autumn.core.project.build_rel_path', 'build_rel_path', ({(18, 26, 18, 49): '"""params/mle-params.yml"""'}, {}), "('params/mle-params.yml')", False, 'from autumn.core.project import Project, ParameterSet, load_timeseries, build_rel_path, get_all_available_scenario_paths, use_tuned_proposal_sds\n'), ((20, 21, 20, 50), 'autumn.projects.covid_19.sri_lanka.sri_lanka.scenario_builder.get_all_scenario_dicts', 'get_all_scenario_dicts', ({(20, 44, 20, 49): '"""LKA"""'}, {}), "('LKA')", False, 'from autumn.projects.covid_19.sri_lanka.sri_lanka.scenario_builder import get_all_scenario_dicts\n'), ((23, 12, 23, 77), 'autumn.core.project.ParameterSet', 'ParameterSet', (), '', False, 'from autumn.core.project import Project, ParameterSet, load_timeseries, build_rel_path, get_all_available_scenario_paths, use_tuned_proposal_sds\n'), ((53, 14, 53, 42), 'autumn.calibration.Calibration', 'Calibration', ({(53, 26, 53, 32): 'priors', (53, 34, 53, 41): 'targets'}, {}), '(priors, targets)', False, 'from autumn.calibration import Calibration\n'), ((58, 21, 58, 54), 'autumn.core.project.build_rel_path', 'build_rel_path', ({(58, 36, 58, 53): '"""timeseries.json"""'}, {}), "('timeseries.json')", False, 'from autumn.core.project import Project, ParameterSet, load_timeseries, build_rel_path, get_all_available_scenario_paths, use_tuned_proposal_sds\n'), ((63, 10, 65, 1), 'autumn.core.project.Project', 'Project', (), '', False, 'from autumn.core.project import Project, ParameterSet, load_timeseries, build_rel_path, get_all_available_scenario_paths, use_tuned_proposal_sds\n'), ((25, 25, 25, 58), 'autumn.core.project.build_rel_path', 'build_rel_path', ({(25, 40, 25, 57): '"""timeseries.json"""'}, {}), "('timeseries.json')", False, 'from autumn.core.project import Project, ParameterSet, load_timeseries, build_rel_path, get_all_available_scenario_paths, use_tuned_proposal_sds\n'), ((29, 4, 29, 34), 'autumn.calibration.targets.NormalTarget', 'NormalTarget', ({(29, 17, 29, 33): 'notifications_ts'}, {}), '(notifications_ts)', False, 'from autumn.calibration.targets import NormalTarget, get_dispersion_priors_for_gaussian_targets\n'), ((30, 4, 30, 26), 'autumn.calibration.targets.NormalTarget', 'NormalTarget', ({(30, 17, 30, 25): 'death_ts'}, {}), '(death_ts)', False, 'from autumn.calibration.targets import NormalTarget, get_dispersion_priors_for_gaussian_targets\n'), ((38, 4, 38, 48), 'autumn.calibration.priors.UniformPrior', 'UniformPrior', ({(38, 17, 38, 31): '"""contact_rate"""', (38, 33, 38, 47): '[0.024, 0.027]'}, {}), "('contact_rate', [0.024, 0.027])", False, 'from autumn.calibration.priors import UniformPrior, BetaPrior, TruncNormalPrior\n'), ((39, 4, 39, 51), 'autumn.calibration.priors.UniformPrior', 'UniformPrior', ({(39, 17, 39, 34): '"""infectious_seed"""', (39, 36, 39, 50): '[275.0, 450.0]'}, {}), "('infectious_seed', [275.0, 450.0])", False, 'from autumn.calibration.priors import UniformPrior, BetaPrior, TruncNormalPrior\n'), ((41, 4, 41, 78), 'autumn.calibration.priors.UniformPrior', 'UniformPrior', ({(41, 17, 41, 61): '"""testing_to_detection.assumed_cdr_parameter"""', (41, 63, 41, 77): '[0.009, 0.025]'}, {}), "('testing_to_detection.assumed_cdr_parameter', [0.009, 0.025])", False, 'from autumn.calibration.priors import UniformPrior, BetaPrior, TruncNormalPrior\n'), ((42, 4, 42, 63), 'autumn.calibration.priors.UniformPrior', 'UniformPrior', ({(42, 17, 42, 48): '"""infection_fatality.multiplier"""', (42, 50, 42, 62): '[0.09, 0.13]'}, {}), "('infection_fatality.multiplier', [0.09, 0.13])", False, 'from autumn.calibration.priors import UniformPrior, BetaPrior, TruncNormalPrior\n'), ((44, 4, 44, 67), 'autumn.calibration.priors.UniformPrior', 'UniformPrior', ({(44, 17, 44, 54): '"""voc_emergence.alpha_beta.start_time"""', (44, 56, 44, 66): '[370, 410]'}, {}), "('voc_emergence.alpha_beta.start_time', [370, 410])", False, 'from autumn.calibration.priors import UniformPrior, BetaPrior, TruncNormalPrior\n'), ((45, 4, 45, 80), 'autumn.calibration.priors.UniformPrior', 'UniformPrior', ({(45, 17, 45, 67): '"""voc_emergence.alpha_beta.contact_rate_multiplier"""', (45, 69, 45, 79): '[3.2, 4.5]'}, {}), "('voc_emergence.alpha_beta.contact_rate_multiplier', [3.2, 4.5])", False, 'from autumn.calibration.priors import UniformPrior, BetaPrior, TruncNormalPrior\n'), ((46, 4, 46, 62), 'autumn.calibration.priors.UniformPrior', 'UniformPrior', ({(46, 17, 46, 49): '"""voc_emergence.delta.start_time"""', (46, 51, 46, 61): '[475, 530]'}, {}), "('voc_emergence.delta.start_time', [475, 530])", False, 'from autumn.calibration.priors import UniformPrior, BetaPrior, TruncNormalPrior\n'), ((47, 4, 47, 76), 'autumn.calibration.priors.UniformPrior', 'UniformPrior', ({(47, 17, 47, 62): '"""voc_emergence.delta.contact_rate_multiplier"""', (47, 64, 47, 75): '[8.5, 11.5]'}, {}), "('voc_emergence.delta.contact_rate_multiplier', [8.5, 11.5])", False, 'from autumn.calibration.priors import UniformPrior, BetaPrior, TruncNormalPrior\n'), ((60, 16, 60, 28), 'json.load', 'json.load', ({(60, 26, 60, 27): 'f'}, {}), '(f)', False, 'import json\n'), ((19, 18, 19, 50), 'autumn.models.covid_19.base_params.update', 'base_params.update', ({(19, 37, 19, 49): 'default_path'}, {}), '(default_path)', False, 'from autumn.models.covid_19 import base_params, build_model\n'), ((35, 5, 35, 56), 'autumn.calibration.targets.get_dispersion_priors_for_gaussian_targets', 'get_dispersion_priors_for_gaussian_targets', ({(35, 48, 35, 55): 'targets'}, {}), '(targets)', False, 'from autumn.calibration.targets import NormalTarget, get_dispersion_priors_for_gaussian_targets\n'), ((36, 5, 36, 56), 'autumn.calibration.targets.get_dispersion_priors_for_gaussian_targets', 'get_dispersion_priors_for_gaussian_targets', ({(36, 48, 36, 55): 'targets'}, {}), '(targets)', False, 'from autumn.calibration.targets import NormalTarget, get_dispersion_priors_for_gaussian_targets\n')]