"""
    biological_showcases.jl

Shared biological data-generating processes for CDCS package-showcase chapters
(Ch. 05b discovery, Ch. 20 TMLE, Ch. 28 CDM, Ch. 28b case studies) and the
book integration test suite.

Structural narratives are informed by host–parasite immunology and vector-control
research (nutrition confounding, treatment, worm burden / biomarkers, lagged
immune–parasite coupling) with **generic variable names** in code.
"""

# Default sizes / seeds (keep in sync with package tests when possible)
const DEFAULT_COHORT_N = 2500
const DEFAULT_COHORT_SEED = 42
const DEFAULT_TS_LENGTH = 800
const DEFAULT_TS_SEED = 11
const DEFAULT_CDM_T = 100
const DEFAULT_CDM_SEED = 34

"""
    confounded_cohort_dgp(; n=2500, rng=Random.default_rng()) -> DataFrame

Simulate an observational cohort with nutrition confounding treatment assignment.

Structural DAG: `nutrition → treatment → worm_burden`, `nutrition → worm_burden`.

Columns: `:nutrition`, `:treatment`, `:worm_burden`.

Identification: adjust `:nutrition` for the effect of `:treatment` on `:worm_burden`.
"""
function confounded_cohort_dgp(; n::Integer = DEFAULT_COHORT_N, rng = Random.default_rng())
    nutrition = randn(rng, n)
    treatment = 0.8 .* nutrition .+ 0.3 .* randn(rng, n)
    worm_burden = 0.7 .* treatment .+ 0.5 .* nutrition .+ 0.3 .* randn(rng, n)
    return DataFrames.DataFrame(
        nutrition = nutrition,
        treatment = treatment,
        worm_burden = worm_burden,
    )
end

"""
    immune_parasite_ts_dgp(; T=800, rng=Random.default_rng())

Bivariate lagged host–parasite time series for OCE / temporal identification.

- `:immune` — immune marker (e.g. antibody / cytokine proxy)
- `:parasite_load` — parasite burden proxy

Coupling: immune at `t-1` affects parasite load at `t`, and vice versa (feedback).
"""
function immune_parasite_ts_dgp(; T::Integer = DEFAULT_TS_LENGTH, rng = Random.default_rng())
    immune = zeros(T)
    parasite_load = zeros(T)
    immune[1] = randn(rng)
    parasite_load[1] = randn(rng)
    for t in 2:T
        immune[t] = 0.4 * immune[t - 1] + 0.45 * parasite_load[t - 1] + 0.1 * randn(rng)
        parasite_load[t] = 0.5 * immune[t - 1] + 0.35 * parasite_load[t - 1] + 0.1 * randn(rng)
    end
    return (
        series = [immune, parasite_load],
        variables = [:immune, :parasite_load],
    )
end

"""
    host_immunity_cdm(x₀, c₀, A, B, D, ρ, α, γ, σ_w, σ_v, σ_c, σ_a; μ_c=0.1)

Build a `DiscreteTimeCDM` for a latent host immune marker under confounded treatment.

- `x` — latent immune marker
- `y` — noisy assay
- `a` — treatment / vaccine channel
- `c` — nutrition or infection-intensity confounder

`protein_treatment_cdm` is a backward-compatible alias.
"""
function host_immunity_cdm(x₀, c₀, A, B, D, ρ, α, γ, σ_w, σ_v, σ_c, σ_a; μ_c = 0.1)
    α_c = μ_c^2 / (σ_c^2)
    θ_c = σ_c^2 / μ_c
    return CausalDynamics.DiscreteTimeCDM(
        [:x, :y, :a, :c];
        initialise = (rng) -> begin
            x = max(0.0, x₀)
            c = max(0.0, c₀)
            u_a = rand(rng, Distributions.Normal(0, σ_a))
            u_y = rand(rng, Distributions.Normal(0, σ_v))
            a = γ * c + u_a
            y = max(0.0, x + u_y)
            (x = x, y = y, a = a, c = c)
        end,
        sample_noise = (rng, state, t) -> begin
            if t == 1
                (u_x = 0.0, u_y = 0.0, u_c = 0.0, u_a = 0.0)
            else
                (
                    u_x = rand(rng, Distributions.Normal(0, σ_w)),
                    u_y = rand(rng, Distributions.Normal(0, σ_v)),
                    u_c = rand(rng, Distributions.Gamma(α_c, θ_c)),
                    u_a = rand(rng, Distributions.Normal(0, σ_a)),
                )
            end
        end,
        step = (state, t, noise, intervention) -> begin
            c = max(0.0, ρ * state.c + α * state.a + noise.u_c)
            a_obs = γ * c + noise.u_a
            a = CausalDynamics.intervention_value(intervention, :a, t, a_obs, state)
            x = max(0.0, A * state.x + B * state.a + D * state.c + noise.u_x)
            y = max(0.0, x + noise.u_y)
            (x = x, y = y, a = a, c = c)
        end,
    )
end

"""Backward-compatible alias for [`host_immunity_cdm`](@ref)."""
const protein_treatment_cdm = host_immunity_cdm

"""
    default_host_immunity_parameters()

Default structural parameters for showcase simulations (Ch. 28).
"""
function default_host_immunity_parameters()
    return (
        x₀ = 1.0,
        c₀ = 0.5,
        A = 0.9,
        B = 0.5,
        D = 0.3,
        ρ = 0.8,
        α = -0.2,
        γ = 0.6,
        σ_w = sqrt(0.1),
        σ_v = sqrt(0.2),
        σ_c = sqrt(0.05),
        σ_a = sqrt(0.1),
    )
end

"""
    simulate_host_immunity_scenario(cdm; T=100, rng=...)

Run factual, interventional (`do(a=1)` / `do(a=0)`), and shared-`U` counterfactual trajectories.
"""
function simulate_host_immunity_scenario(
    cdm;
    T::Integer = DEFAULT_CDM_T,
    factual_seed::Integer = DEFAULT_CDM_SEED,
    treat_seed::Integer = 35,
    control_seed::Integer = 36,
)
    factual = CausalDynamics.simulate(cdm, T; rng = StableRNGs.StableRNG(factual_seed))
    treated = CausalDynamics.simulate(
        cdm, T;
        rng = StableRNGs.StableRNG(treat_seed),
        intervention = CausalDynamics.do_sequence(:a, ones(T)),
    )
    control = CausalDynamics.simulate(
        cdm, T;
        rng = StableRNGs.StableRNG(control_seed),
        intervention = CausalDynamics.do_sequence(:a, zeros(T)),
    )
    counterfactual = CausalDynamics.counterfactual(
        cdm,
        factual.noise;
        intervention = CausalDynamics.do_sequence(:a, ones(T)),
        initial = (
            x = factual.series[:x][1],
            y = factual.series[:y][1],
            a = factual.series[:a][1],
            c = factual.series[:c][1],
        ),
    )
    return (factual = factual, treated = treated, control = control, counterfactual = counterfactual)
end

"""
    cohort_for_tmle(df) -> DataFrame

Prepare cohort data for TMLE.jl (categorical treatment column).
Requires `DataFrames` and `CategoricalArrays` in the calling module.
"""
function cohort_for_tmle(df)
    treatment_binary = Float64[df.treatment[i] > median(df.treatment) ? 1.0 : 0.0 for i in 1:DataFrames.nrow(df)]
    return DataFrames.DataFrame(
        nutrition = df.nutrition,
        treatment = CategoricalArrays.categorical(treatment_binary),
        worm_burden = df.worm_burden,
    )
end