"""Exact stochastic Gillespie and deterministic ODE simulation of adaptive vs continuous therapy.

Evaluates competitive suppression of pre-existing resistant clones versus mutational
influx in de novo resistance under Gompertz/logistic crowding.
"""

import json
import math
import os
import random


def simulate_stochastic(f0, mu, regimen, seed, K=2000, N0=500, b_S=0.15, d_S=0.03, delta=0.22, s_cost=0.15, prog_threshold=1200, max_days=300.0):
    rng = random.Random(seed)
    N_S = int(round(N0 * (1.0 - f0)))
    N_R = int(round(N0 * f0))
    b_R = b_S * (1.0 - s_cost)
    d_R = d_S
    drug = True
    t = 0.0
    dose = 0.0

    while t < max_days:
        N = N_S + N_R
        if N >= prog_threshold:
            return {"ttp": t, "cure": False, "dose": dose}
        if N == 0:
            return {"ttp": max_days, "cure": True, "dose": dose}

        if regimen == "continuous":
            drug = True
        elif regimen == "adaptive":
            if drug and N <= 0.5 * N0:
                drug = False
            elif not drug and N >= N0:
                drug = True
        elif regimen == "metronomic":
            drug = ((int(t) // 7) % 2) == 0

        crowding = max(0.0, 1.0 - N / K)
        rate_div_S = b_S * crowding * N_S
        rate_birth_S = (1.0 - mu) * rate_div_S
        rate_mut_S = mu * rate_div_S
        rate_death_S = (d_S + (delta if drug else 0.0)) * N_S

        rate_birth_R = b_R * crowding * N_R
        rate_death_R = d_R * N_R

        total_rate = rate_birth_S + rate_mut_S + rate_death_S + rate_birth_R + rate_death_R
        if total_rate <= 0:
            break

        dt = rng.expovariate(total_rate)
        if drug:
            dose += dt
        t += dt

        r = rng.uniform(0, total_rate)
        if r < rate_birth_S:
            N_S += 1
        elif r < rate_birth_S + rate_mut_S:
            N_R += 1
        elif r < rate_birth_S + rate_mut_S + rate_death_S:
            N_S -= 1
        elif r < rate_birth_S + rate_mut_S + rate_death_S + rate_birth_R:
            N_R += 1
        else:
            N_R -= 1

    return {"ttp": max_days, "cure": (N_S + N_R == 0), "dose": dose}


def simulate_ode(f0=0.01, s_cost=0.10, regimen="continuous", dt=0.05, K=2000, N0=500, b_S=0.15, d_S=0.03, delta=0.22, mu=8e-5, prog_threshold=1200, max_days=500.0):
    N_S = N0 * (1.0 - f0)
    N_R = N0 * f0
    b_R_loc = b_S * (1.0 - s_cost)
    d_R = d_S
    drug = True
    t = 0.0
    dose = 0.0

    while t < max_days:
        N = N_S + N_R
        if N >= prog_threshold:
            return {"ttp": t, "dose": dose}
        if regimen == "continuous":
            drug = True
        elif regimen == "adaptive":
            if drug and N <= 0.5 * N0:
                drug = False
            elif not drug and N >= N0:
                drug = True

        if drug:
            dose += dt

        crowding = max(0.0, 1.0 - N / K)
        rate_div_S = b_S * crowding * N_S
        mut = mu * rate_div_S
        N_S += dt * ((1.0 - mu) * rate_div_S - (d_S + (delta if drug else 0.0)) * N_S)
        N_R += dt * (b_R_loc * crowding * N_R + mut - d_R * N_R)
        N_S = max(0.0, N_S)
        N_R = max(0.0, N_R)
        t += dt

    return {"ttp": max_days, "dose": dose}


def main():
    os.makedirs("results", exist_ok=True)
    n_trials = 500
    mu_denovo = 8e-5

    # --- Part 1: De Novo Resistance Arm (f0 = 0.0) ---
    cont_cures = 0
    adapt_cures = 0
    cont_ttps = []
    adapt_ttps = []
    cont_doses = []
    adapt_doses = []

    for i in range(n_trials):
        res_c = simulate_stochastic(f0=0.0, mu=mu_denovo, regimen="continuous", seed=10000 + i)
        if res_c["cure"]:
            cont_cures += 1
        else:
            cont_ttps.append(res_c["ttp"])
        cont_doses.append(res_c["dose"])

        res_a = simulate_stochastic(f0=0.0, mu=mu_denovo, regimen="adaptive", seed=20000 + i)
        if res_a["cure"]:
            adapt_cures += 1
        else:
            adapt_ttps.append(res_a["ttp"])
        adapt_doses.append(res_a["dose"])

    cure_rate_cont = round(cont_cures / n_trials * 100.0, 1)
    cure_rate_adapt = round(adapt_cures / n_trials * 100.0, 1)
    mean_ttp_adapt_denovo = round(sum(adapt_ttps) / len(adapt_ttps), 1) if adapt_ttps else 0.0
    mean_dose_adapt_denovo = round(sum(adapt_doses) / n_trials, 1)
    mean_dose_cont_denovo = round(sum(cont_doses) / n_trials, 1)

    r1 = {
        "n_trials": n_trials,
        "cure_rate_continuous_pct": cure_rate_cont,
        "cure_rate_adaptive_pct": cure_rate_adapt,
        "mean_ttp_adaptive_days": mean_ttp_adapt_denovo,
        "mean_dose_adaptive_days": mean_dose_adapt_denovo,
        "mean_dose_continuous_days": mean_dose_cont_denovo,
    }
    with open("results/R1.json", "w") as f:
        json.dump(r1, f, indent=2)
        f.write("\n")

    # --- Part 2: Pre-Existing Resistance Arm (f0 = 0.02) ---
    f0_pre = 0.02
    cont_pre_ttps = []
    adapt_pre_ttps = []
    cont_pre_doses = []
    adapt_pre_doses = []

    for i in range(n_trials):
        res_c = simulate_stochastic(f0=f0_pre, mu=mu_denovo, regimen="continuous", seed=30000 + i)
        cont_pre_ttps.append(res_c["ttp"])
        cont_pre_doses.append(res_c["dose"])

        res_a = simulate_stochastic(f0=f0_pre, mu=mu_denovo, regimen="adaptive", seed=40000 + i)
        adapt_pre_ttps.append(res_a["ttp"])
        adapt_pre_doses.append(res_a["dose"])

    ttp_cont_pre = round(sum(cont_pre_ttps) / n_trials, 1)
    ttp_adapt_pre = round(sum(adapt_pre_ttps) / n_trials, 1)
    dose_cont_pre = round(sum(cont_pre_doses) / n_trials, 1)
    dose_adapt_pre = round(sum(adapt_pre_doses) / n_trials, 1)
    dose_red_pct = round((dose_cont_pre - dose_adapt_pre) / dose_cont_pre * 100.0, 1)
    ttp_gain_days = round(ttp_adapt_pre - ttp_cont_pre, 1)

    r2 = {
        "n_trials": n_trials,
        "f0_pre_pct": 2.0,
        "ttp_continuous_days": ttp_cont_pre,
        "ttp_adaptive_days": ttp_adapt_pre,
        "ttp_gain_days": ttp_gain_days,
        "dose_continuous_days": dose_cont_pre,
        "dose_adaptive_days": dose_adapt_pre,
        "dose_reduction_pct": dose_red_pct,
    }
    with open("results/R2.json", "w") as f:
        json.dump(r2, f, indent=2)
        f.write("\n")

    # --- Part 3: Deterministic Fitness Cost Gradient ---
    costs = [0.20, 0.10, 0.05, 0.00]
    ode_out = {}
    for sc in costs:
        res_cont = simulate_ode(f0=0.01, s_cost=sc, regimen="continuous")
        res_adapt = simulate_ode(f0=0.01, s_cost=sc, regimen="adaptive")
        t_c = round(res_cont["ttp"], 1)
        t_a = round(res_adapt["ttp"], 1)
        ratio = round(t_a / t_c, 2)
        ode_out[str(sc)] = {"t_cont": t_c, "t_adapt": t_a, "ratio": ratio}

    r3 = {
        "cost_20_pct": 20.0,
        "cost_20": 0.20,
        "cost_10": 0.10,
        "cost_05": 0.05,
        "cost_00": 0.00,
        "f0_ode_pct": 1.0,
        "ttp_cont_cost_20": ode_out["0.2"]["t_cont"],
        "ttp_adapt_cost_20": ode_out["0.2"]["t_adapt"],
        "ttp_ratio_cost_20": ode_out["0.2"]["ratio"],
        "ttp_cont_cost_10": ode_out["0.1"]["t_cont"],
        "ttp_adapt_cost_10": ode_out["0.1"]["t_adapt"],
        "ttp_ratio_cost_10": ode_out["0.1"]["ratio"],
        "ttp_cont_cost_05": ode_out["0.05"]["t_cont"],
        "ttp_adapt_cost_05": ode_out["0.05"]["t_adapt"],
        "ttp_ratio_cost_05": ode_out["0.05"]["ratio"],
        "ttp_cont_cost_00": ode_out["0.0"]["t_cont"],
        "ttp_adapt_cost_00": ode_out["0.0"]["t_adapt"],
        "ttp_ratio_cost_00": ode_out["0.0"]["ratio"],
    }
    with open("results/R3.json", "w") as f:
        json.dump(r3, f, indent=2)
        f.write("\n")

    print("Computation complete. Wrote results/R1.json, results/R2.json, results/R3.json.")


if __name__ == "__main__":
    main()