"""
Test demographic consistency
"""

import numpy as np
import pandas as pd
import sciris as sc
import matplotlib.pyplot as plt
import starsim as ss
import starsim_examples as sse
import pytest

sc.options(interactive=False) # Assume not running interactively
datadir = ss.root / 'tests/test_data'


@sc.timer()
def test_nigeria(which='births', dt=1, start=1995, dur=15, do_plot=False):
    """
    Make a Nigeria sim with demographic modules
    Switch between which='births' or 'pregnancy' to determine which demographic module to use
    """
    sc.heading('Testing Nigeria demographics')

    # Make demographic modules
    demographics = sc.autolist()

    if which == 'births':
        birth_rates = pd.read_csv(datadir/'nigeria_births.csv')
        births = ss.Births(birth_rate=birth_rates)
        demographics += births

    elif which == 'pregnancy':
        fertility_rates = pd.read_csv(datadir/'nigeria_asfr.csv')
        pregnancy = ss.Pregnancy(fertility_rate=fertility_rates, rel_fertility=1, burnin=False)  # 4/3
        demographics += pregnancy

    death_rates = pd.read_csv(datadir/'nigeria_deaths.csv')
    death = ss.Deaths(death_rate=death_rates, rate_units=1)
    demographics += death

    # Make people
    n_agents = 5_000
    nga_pop_1995 = 106819805
    age_data = pd.read_csv(datadir/'nigeria_age.csv')
    ppl = ss.People(n_agents, age_data=age_data)

    sim = ss.Sim(
        dt=dt,
        total_pop=nga_pop_1995,
        start=start,
        dur=dur,
        people=ppl,
        demographics=demographics,
    )

    if do_plot:
        sim.init()
        # Plot histograms of the age distributions - simulated vs data
        bins = np.arange(0, 101, 1)
        init_scale = nga_pop_1995 / n_agents
        counts, bins = np.histogram(sim.people.age, bins)
        plt.bar(bins[:-1], counts * init_scale, alpha=0.5, label='Simulated')
        plt.bar(bins, age_data.value.values * 1000, alpha=0.5, color='r', label='Data')
        plt.legend(loc='upper right')

    sim.run()

    stop = start + dur
    nigeria_popsize = pd.read_csv(datadir/'nigeria_popsize.csv')
    data = nigeria_popsize[(nigeria_popsize.year >= start) & (nigeria_popsize.year <= stop)]

    nigeria_cbr = pd.read_csv(datadir/'nigeria_births.csv')
    cbr_data = nigeria_cbr[(nigeria_cbr.Year >= start) & (nigeria_cbr.Year <= stop)]

    nigeria_cmr = pd.read_csv(datadir/'nigeria_cmr.csv')
    cmr_data = nigeria_cmr[(nigeria_cmr.Year >= start) & (nigeria_cmr.Year <= stop)]

    # Tests
    if which == 'pregnancy':

        print("Check we don't have more births than pregnancies")
        assert sum(sim.results.pregnancy.births) <= sum(sim.results.pregnancy.pregnancies)
        print('✓ (births <= pregnancies)')

        if dt == 1:
            print("Checking that births equal pregnancies with dt=1")
            assert np.array_equal(sim.results.pregnancy.pregnancies, sim.results.pregnancy.births)
            print('✓ (births == pregnancies)')

    rtol = 0.05
    print(f'Check final pop size within {rtol*100:n}% of data')
    assert np.isclose(data.n_alive.values[-1], sim.results.n_alive[-1], rtol=rtol), f'Final population size not within {rtol*100:n}% of data'
    print(f'✓ (simulated/data={sim.results.n_alive[-1] / data.n_alive.values[-1]:.2f})')

    # Plots
    if do_plot:
        tvec = sim.timevec
        fig, ax = plt.subplots(2, 2)
        ax = ax.ravel()
        ax[0].scatter(data.year, data.n_alive, alpha=0.5)
        ax[0].plot(tvec, sim.results.n_alive, color='k')
        ax[0].set_title('Population')

        ax[1].plot(tvec, 1000 * sim.results.deaths.cmr, label='Simulated CMR')
        ax[1].scatter(cmr_data.Year, cmr_data.CMR, label='Data CMR')
        ax[1].set_title('CMR')
        ax[1].legend()

        if which == 'births':
            ax[2].plot(tvec, sim.results.births.cbr, label='Simulated CBR')
        elif which == 'pregnancy':
            ax[2].plot(tvec, sim.results.pregnancy.cbr, label='Simulated CBR')
        ax[2].scatter(cbr_data.Year, cbr_data.CBR, label='Data CBR')
        ax[2].set_title('CBR')
        ax[2].legend()

        if which == 'pregnancy':
            ax[3].plot(tvec, sim.results.pregnancy.pregnancies, label='Pregnancies')
            ax[3].plot(tvec, sim.results.pregnancy.births, label='Births')
            ax[3].set_title('Pregnancies and births')
            ax[3].legend()

        fig.tight_layout()

    return sim


@sc.timer()
def test_constant_pop(do_plot=False):
    """ Test pars for constant pop size """
    sc.heading('Testing constant population size')
    sim = ss.Sim(n_agents=10e3, birth_rate=ss.freqperyear(10), death_rate=ss.freqperyear(10/1010*1000), dur=ss.years(200), rand_seed=1).run()
    print("Check final pop size within 5% of starting pop")
    assert np.isclose(sim.results.n_alive[0], sim.results.n_alive[-1], rtol=0.05)
    print(f'✓ (final pop / starting pop={sim.results.n_alive[-1] / sim.results.n_alive[0]:.2f})')

    # Plots
    if do_plot:
        sim.plot()
    return sim


@sc.timer()
def test_module_adding():
    """ Test that modules can't be added twice """
    sc.heading('Testing module duplication')
    births = ss.Births(birth_rate=ss.freqperyear(10))
    deaths = ss.Deaths(death_rate=ss.freqperyear(10))
    demographics = [births, deaths]
    with pytest.raises(ValueError):
        ss.Sim(n_agents=1e3, demographics=demographics, birth_rate=ss.freqperyear(10), death_rate=ss.freqperyear(10)).run()
    return demographics


@sc.timer()
def test_aging():
    """ Test that aging is configured properly """
    sc.heading('Testing aging')
    n_agents = int(1e3)

    # With no demograhpics, people shouldn't age
    s1 = ss.Sim(n_agents=n_agents).init()
    orig_ages = s1.people.age.raw.copy()
    orig_age = orig_ages.mean()
    s1.run()
    end_age = s1.people.age.mean()
    assert orig_age == end_age, f'By default there should be no aging, but {orig_age} != {end_age}'

    # We should be able to manually turn on aging
    s2 = ss.Sim(n_agents=n_agents, use_aging=True).run()
    age2 = s2.people.age.mean()
    assert orig_age < age2, f'With aging, start age {orig_age} should be less than end age {age2}'

    # Aging should turn on automatically if we add demographics
    s3 = ss.Sim(n_agents=n_agents, demographics=True).run()
    agent = s3.people.auids[0] # Find first alive agent
    orig_agent_age = orig_ages[agent]
    age3 = s3.people.age[ss.uids(agent)]
    assert orig_agent_age < age3, f'With demographics, original agent age {orig_agent_age} should be less than end age {age3}'

    # ...but can be turned off manually
    s4 = ss.Sim(n_agents=n_agents, demographics=True, use_aging=False).run()
    agent = s4.people.auids[0] # Find first alive agent
    orig_agent_age = orig_ages[agent]
    age4 = s4.people.age[ss.uids(agent)]
    assert orig_agent_age == age4, f'With aging turned off, original agent age {orig_agent_age} should match end age {age4}'

    return s3 # Most interesting sim


def test_pregnancy():
    """ Test pregnancy module"""
    sc.heading('Testing pregnancy module')

    sim = ss.Sim(
        n_agents=10e3,
        demographics=[
            ss.Pregnancy(fertility_rate=ss.freqperyear(10), burnin=True),
            ss.Deaths(death_rate=ss.freqperyear(10/1010*1000)),
        ],
        dur=ss.years(10),
        rand_seed=1,
        dt=ss.years(1/12),
        networks=ss.PrenatalNet(),
    )
    sim.run()

    # Tests
    # Check that everyone in their 2nd trimester conceived 3-6 months ago
    tri2 = sim.demographics.pregnancy.tri2_uids
    conception_time = sim.people.pregnancy.ti_pregnant[tri2]
    time_since_conception = ss.years((sim.ti - conception_time) * sim.t.dt_year).weeks
    assert np.all((time_since_conception >= 13) & (time_since_conception <= 26)), 'Some people in 2nd trimester did not conceive 3-6 months ago'

    # Check that gestational clock is present for all pregnant women, and that they all have a defined delivery time
    pregnant = sim.people.pregnancy.pregnant
    assert np.all(~np.isnan(sim.people.pregnancy.gestation[pregnant])), 'Some pregnant people are missing gestational clock'
    ti_delivery = sim.people.pregnancy.ti_delivery[pregnant]
    assert np.all(~np.isnan(ti_delivery)), 'Some people with ti_pregnant are missing ti_delivery'

    # Check that gestational clock is NaN for all non-pregnant women
    assert np.all(np.isnan(sim.people.pregnancy.gestation[~sim.people.pregnancy.pregnant])), 'Some non-pregnant people have gestational clock'
    sc.printgreen('✓ Pregnancy duration tests passed')

    # Check that everyone born during the sim has gestational age at birth recorded
    born_during_sim = (sim.people.age < sim.pars.dur.years) & (sim.people.age > sim.t.dt.years)
    aa = (sim.people.pregnancy.gestation_at_birth.isnan & born_during_sim).uids
    assert len(aa) == 0, f'Some babies born during the sim are missing gestational age at birth: {len(aa)} found'

    # Check that the mean gestational age at birth is approximately 40 weeks
    mean_gestation = np.nanmean(sim.people.pregnancy.gestation_at_birth[born_during_sim])
    assert np.isclose(mean_gestation, 40, atol=2), f'Mean gestational age at birth is {mean_gestation:.1f} weeks, expected ~40 weeks'

    # Check that the mean gestational age of babies born to women <18 or >35 is shorter than otherwise
    mothers = sim.people.parent[born_during_sim]
    age_of_mother = sim.people.age[mothers]
    age_of_mother_at_birth = age_of_mother - sim.people.age[born_during_sim]
    key_cohorts = {
        '<18': age_of_mother_at_birth < 18,
        '>35': age_of_mother_at_birth > 35,
        '18-35': (age_of_mother_at_birth >= 18) & (age_of_mother_at_birth <= 35),
    }
    mean_results = {k:None for k in key_cohorts.keys()}
    ptb_results = {k:None for k in key_cohorts.keys()}
    for k, v in key_cohorts.items():
        these_mothers_idx = v.nonzero()[-1]
        babies_of_these_mothers = born_during_sim.uids[these_mothers_idx]
        mean_gestation = sim.people.pregnancy.gestation_at_birth[babies_of_these_mothers].mean()
        print(f'Mean gestational age at birth for mothers aged {k} at birth: {mean_gestation:.1f} weeks')
        mean_results[k] = mean_gestation
        n = len(babies_of_these_mothers)
        print(f'  {n} babies born to mothers aged {k} at birth')
        # Share of babies born preterm (<37 weeks)
        n_preterm = np.sum(sim.people.pregnancy.gestation_at_birth[babies_of_these_mothers] < 37)
        preterm_share = n_preterm / n if n > 0 else np.nan
        print(f'  Preterm births (<37 weeks): {n_preterm} ({preterm_share:.1%})')
        ptb_results[k] = preterm_share

    # Assertions
    assert mean_results['<18'] < mean_results['18-35'], '<18 mothers do not have lower gestational age at birth than 18-35 mothers'
    assert mean_results['>35'] < mean_results['18-35'], '>35 mothers do not have lower gestational age at birth than 18-35 mothers'
    assert ptb_results['<18'] > ptb_results['18-35'], '<18 mothers do not have higher preterm birth share than 18-35 mothers'
    assert ptb_results['>35'] > ptb_results['18-35'], '>35 mothers do not have higher preterm birth share than 18-35 mothers'

    sc.printgreen('✓ Gestational age at birth by maternal age tests passed')

    return sim


def test_pregnancy_short_sim():
    """ Test that Pregnancy.finalize() works for sims shorter than 1 year """
    sc.heading('Testing pregnancy with short sim duration')

    sim = ss.Sim(
        n_agents=1e3,
        demographics=[
            ss.Pregnancy(fertility_rate=ss.freqperyear(10), burnin=False),
            ss.Deaths(death_rate=ss.freqperyear(10/1010*1000)),
        ],
        dur=ss.days(250),
        rand_seed=1,
        dt=ss.days(10),
        networks=ss.PrenatalNet(),
    )
    sim.run()  # Should not raise ValueError in finalize()

    sc.printgreen('✓ Short sim pregnancy finalize test passed')
    return sim


@sc.timer()
def test_loss_classification():
    """
    Test that prenatal deaths are correctly classified as miscarriage (<20w)
    or stillbirth (>=20w) based on gestational age.
    """
    sc.heading('Testing loss classification...')

    # Use a high loss rate so both early and late losses occur
    sim = ss.Sim(
        n_agents=5_000,
        demographics=[
            ss.Pregnancy(fertility_rate=ss.freqperyear(30), p_loss=ss.bernoulli(p=0.05), burnin=True),
            ss.Deaths(death_rate=ss.freqperyear(10/1010*1000)),
        ],
        dur=ss.years(5),
        dt=ss.years(1/12),
        rand_seed=1,
        networks=ss.PrenatalNet(),
    )
    sim.run()

    preg = sim.demographics.pregnancy
    mc = preg.results['miscarriages'].sum()
    sb = preg.results['stillbirths'].sum()
    births = preg.results['births'].sum()
    print(f'  Miscarriages: {mc}, Stillbirths: {sb}, Births: {births}')

    assert mc > 0, 'Expected some miscarriages with p_loss=0.05'
    assert sb > 0, 'Expected some stillbirths with p_loss=0.05'
    assert mc > sb, 'Expected more miscarriages than stillbirths (pregnancies spend more time at early GA)'
    assert births > 0, 'Expected some births despite losses'

    sc.printgreen('✓ Loss classification tests passed')
    return sim


@sc.timer()
def test_background_loss():
    """
    Test that background pregnancy loss (p_loss) produces losses and that
    higher p_loss produces more losses.
    """
    sc.heading('Testing background loss...')

    results = dict()
    for label, p_loss in [('none', 0), ('low', 0.01), ('high', 0.05)]:
        sim = ss.Sim(
            n_agents=5_000,
            demographics=[
                ss.Pregnancy(fertility_rate=ss.freqperyear(30), p_loss=ss.bernoulli(p=p_loss), burnin=True),
                ss.Deaths(death_rate=ss.freqperyear(10/1010*1000)),
            ],
            dur=ss.years(5),
            dt=ss.years(1/12),
            rand_seed=1,
            networks=ss.PrenatalNet(),
        )
        sim.run()
        preg = sim.demographics.pregnancy
        total_losses = preg.results['miscarriages'].sum() + preg.results['stillbirths'].sum()
        births = preg.results['births'].sum()
        results[label] = dict(losses=total_losses, births=births)
        print(f'  p_loss={p_loss}: losses={total_losses}, births={births}')

    assert results['low']['losses'] > results['none']['losses'], 'Expected more losses with p_loss=0.01 than p_loss=0'
    assert results['high']['losses'] > results['low']['losses'], 'Expected more losses with higher p_loss'
    assert results['high']['births'] < results['none']['births'], 'Expected fewer births with p_loss > 0'

    sc.printgreen('✓ Background loss tests passed')
    return sim


@sc.timer()
def test_nnd():
    """
    Test passive neonatal death detection using neonatal sepsis disease.
    Verify deaths of newborns within 28 days are classified as NNDs.
    """
    sc.heading('Testing neonatal death detection...')

    sim = ss.Sim(
        n_agents=5_000,
        demographics=[
            ss.Pregnancy(fertility_rate=ss.freqperyear(30), burnin=True),
            ss.Deaths(death_rate=ss.freqperyear(10/1010*1000)),
        ],
        diseases=sse.NeonatalSepsis(),
        dur=ss.years(3),
        dt=ss.years(1/12),
        rand_seed=1,
        networks=ss.PrenatalNet(),
    )
    sim.run()

    preg = sim.demographics.pregnancy
    neonatal_deaths = preg.results['neonatal_deaths'].sum()
    births = preg.results['births'].sum()
    print(f'  Neonatal deaths: {neonatal_deaths}, Births: {births}')

    assert neonatal_deaths > 0, 'Expected neonatal deaths from neonatal sepsis'
    assert neonatal_deaths < births, 'Neonatal deaths should be fewer than total births'

    sc.printgreen('✓ Neonatal death detection tests passed')
    return sim


if __name__ == '__main__':
    T = sc.timer()

    s1 = test_nigeria()
    s2 = test_nigeria(dt=1/12, which='pregnancy')
    s3 = test_constant_pop()
    s4 = test_module_adding()
    s5 = test_aging()
    s6 = test_pregnancy()
    s7 = test_loss_classification()
    s8 = test_background_loss()
    s9 = test_nnd()

    T.toc()
