"""
Test run, including ss.MultiSim and ss.parallel()
"""

# %% Imports and settings
import starsim as ss
import sciris as sc
import numpy as np
import matplotlib.pyplot as plt

n_agents = 1_000
do_plot = False
sc.options(interactive=False) # Assume not running interactively


def make_sim_pars(beta=0.1, **kwargs):
    pars = sc.objdict(
        n_agents = n_agents,
        demographics = True,
        networks = 'random',
        diseases = sc.objdict(type='sis', beta=beta), # To allow for modification later
        **kwargs,
    )
    return pars


@sc.timer()
def test_parallel():
    """ Test running two identical sims in parallel """
    sc.heading('Testing parallel...')
    pars = make_sim_pars()

    # Check that two identical sims match
    msim = ss.MultiSim([ss.Sim(pars, label='Sim1'), ss.Sim(pars, label='Sim2')])
    msim.run(shrink=False)
    s1, s2 = msim.sims
    assert np.allclose(s1.summary[:], s2.summary[:], rtol=0, atol=0, equal_nan=True), "Sims don't match and should"

    # Check that two non-identical sims don't match
    pars2 = sc.dcp(pars)
    pars2.diseases.beta *= 2
    sims = ss.MultiSim([ss.Sim(pars, label='Sim1'), ss.Sim(pars2, label='Sim2')])
    sims.run(shrink=False)
    s1, s2 = sims.sims
    assert not np.allclose(s1.summary[:], s2.summary[:], rtol=0, atol=0, equal_nan=True), "Sims do match and shouldn't"

    return s1, s2


@sc.timer()
def test_multisim():
    """ Check MultiSim methods """
    sc.heading('Testing MultiSim')
    pars = make_sim_pars()
    msim = ss.MultiSim(ss.Sim(pars))
    msim.init_sims()
    msim.run(n_runs=4)

    # Plot individual sims
    msim.plot()

    # Reduce and plot mean
    msim.mean()
    msim.plot()

    # Export results
    res_df = msim.results.to_df(resample='2YE')
    assert res_df.sis_n_susceptible_low.loc['2030-12-31'] == msim.results.sis_n_susceptible.low[29:31].mean()

    # Reduce and plot median
    msim.median()
    msim.plot()

    # Other methods
    msim.show()
    msim.summarize()
    msim.reset()

    return msim


@sc.timer()
def test_multisim_construction(do_plot=False):
    """ Test other MultiSim construction methods """
    sc.heading('Testing MultiSim construction')
    
    # Make sims
    def make_sims(n=2):
        sims = sc.autolist()
        for beta in 0.1*np.random.rand(n):
            pars = make_sim_pars(beta=beta, dur=20, rand_seed=np.random.randint(999))
            sims += ss.Sim(pars=pars, label=f'{beta=}')
        return sims
    
    # Construction options
    m1 = ss.MultiSim(make_sims(), label='list') # As list
    m2 = ss.MultiSim(tuple(make_sims()), label='tuple') # As tuple
    m3 = ss.MultiSim(m1.copy(), label='msim') # As MultiSim
    m4 = ss.MultiSim([m1.copy(), m2.copy()], label='multi-msim') # As list of MultiSims
    
    # Check that run works
    for msim in [m1, m2, m3, m4]:
        msim.run()
        if do_plot:
            msim.plot()
    
    return m4


@sc.timer()
def test_other():
    """ Check other run options """
    sc.heading('Testing other run options')
    # Check parallel
    pars = make_sim_pars()
    pars2 = sc.dcp(pars)
    pars2.diseases.beta *= 2
    s1 = ss.Sim(pars)
    s2 = ss.Sim(pars2)
    ss.parallel(s1, s2, inplace=False)
    assert not s1.initialized
    ss.parallel(s1, s2, inplace=True)
    assert not np.allclose(s1.summary[:], s2.summary[:], rtol=0, atol=0, equal_nan=True), "Sims do match and shouldn't"

    # Check single run
    s3 = ss.Sim(pars)
    s4 = ss.Sim(pars)
    s3 = ss.single_run(s3)
    s4 = ss.single_run(s4, n_agents=2000)
    assert not np.allclose(s3.summary[:], s4.summary[:], rtol=0, atol=0, equal_nan=True), "Sims do match and shouldn't"

    return s3,s4


if __name__ == '__main__':
    do_plot = True
    sc.options(interactive=do_plot)
    T = sc.timer()

    s1, s2 = test_parallel()
    msim = test_multisim()
    msim2 = test_multisim_construction(do_plot)
    s3,s4 = test_other()

    T.toc()

    if do_plot:
        plt.show()
