"""Shared fixtures: a small, fast configuration of the Part B model solved once per session."""
from __future__ import annotations

import os
import sys

import numpy as np
import pytest

ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
sys.path.insert(0, ROOT)
sys.path.insert(0, os.path.join(ROOT, "src"))

from src import distributions as D  # noqa: E402
from src import equilibrium as E  # noqa: E402
from src import household_solver as HS  # noqa: E402
from src import model as M  # noqa: E402

SMALL = {"simulation.n_dynasties": 12000, "simulation.burn_periods": 36, "grids.n_a": 40, "grids.n_x": 100,
         "income.sigma_transitory": 0.0}   # exact identities; the noisy-income identity has its own test


@pytest.fixture(scope="session")
def P_small():
    cfg = M.load_config(overrides=SMALL)
    return M.build_params(cfg)


@pytest.fixture(scope="session")
def lottery(P_small):
    K = P_small.K_estate
    return (np.array([0.0, 0.0, 0.05, 0.4, 1.0])[:K], np.full(K, 1.0 / K))


@pytest.fixture(scope="session")
def prices(P_small):
    return M.stationary_prices(P_small, 1.0, 1.0)


@pytest.fixture(scope="session")
def pols(P_small, prices, lottery):
    return HS.solve_lifecycle(P_small, prices, lottery)


@pytest.fixture(scope="session")
def sim(P_small, pols, prices):
    recs, panel = D.simulate_stationary(P_small, pols, prices, n_record=6)
    return D.pool_records(recs), panel


@pytest.fixture(scope="session")
def ge_small(P_small):
    """A small general equilibrium: supply scales from the calibration, then a GE solve from a
    perturbed price pair. Takes ~1 minute."""
    P = M.build_params(M.load_config(overrides=SMALL))
    log = E.RunLog("pytest_ge", echo=False)
    E.calibrate_supply(P, log)
    res = E.solve_stationary(P, p0=1.03, ptil0=0.98, log=log, name="pytest_ge", diagnostics=True)
    return P, res