"""Simulated trials in the review's data format, to test the analysis before any outcome data exist.

    python3 -I code/simulate.py <out dir> [masked_g] [awake_g]
        90 invented trials with a known masked and awake effect, a few with means no whole-number totals
        could give (as fabricated data might), and a registry, so a run shows whether the code recovers
        the effects and flags those trials.
    python3 -I code/simulate.py --designs <trials.json> <out dir>
        Invented numbers for the real trials' designs (their arms, measure names, and days), to test that
        the analysis reads real designs. The numbers say nothing about any trial's results.

Nothing here comes from any trial's results.
"""

import json
import pathlib
import random
import shutil
import sys

here = pathlib.Path(__file__).resolve().parent
FORMS = ["mean_sd"] * 6 + ["median_iqr", "mean_se", "count", "difference"]


def simulated(out, masked_g, awake_g):
    out.mkdir(parents=True, exist_ok=True)
    rng = random.Random(7)
    shutil.copy(here.parent / "data" / "parameters.json", out / "parameters.json")
    trials, extraction = [], []
    for i in range(90):
        group = "none" if i < 40 else "some"
        setting = "surgery" if (i < 38 or 40 <= i < 65) else "cesarean"
        selected = "selected" if rng.random() < 0.2 else "unselected"
        g = (masked_g if group == "none" else awake_g) + rng.gauss(0, 0.25)
        n = rng.choice([40, 60, 80, 100, 120, 200])
        tid = f"T{i + 1:03d}"
        code = rng.choice(["HAMD", "EPDS", "PHQ9", "HADS_D", "MADRS"])
        design = {
            "eligible": True,
            "arms": [{"name": "K", "role": "ketamine"}, {"name": "C", "role": "control"}],
            "conscious_dosing": group, "late_dosing": rng.random() < 0.3, "setting": setting, "depression_selection": selected,
            "depression_measures": [{"measure": "D", "code": code, "days": [1, 3, 7, 30], "primary": True}],
            "pain_measures": [{"measure": "VAS rest", "code": "rest", "days": [1], "primary": False}],
            "registration": {"registry": "ChiCTR", "id": f"ChiCTR22000{i:05d}", "date": "2022-01-15" if rng.random() < 0.6 else "2022-09-01"},
            "enrollment": {"start": "2022-03", "end": "2022-12"},
        }
        trials.append({"trial": tid, "source": rng.choice(["fulltext", "fulltext", "abstract", "preprint"]), "design": design,
                       "notices": [], "registration_year": 2022, "pump_only": False, "drug": rng.choice(["esketamine", "ketamine"]),
                       "dose_ketamine_equivalent_mg_per_kg": round(rng.uniform(0.2, 1.0), 2)})
        sd, half = 4.0, n // 2
        entries = []
        for day, shift in [(-1, 0.0), (1, 1.0), (3, 0.9), (7, 0.6), (30, 0.3)]:
            form = rng.choice(FORMS) if day > 0 else "mean_sd"
            mk = 10 + (g * sd * shift if day > 0 else 0) + rng.gauss(0, sd / half ** 0.5)
            mc = 10 + rng.gauss(0, sd / half ** 0.5)
            if i % 15 != 7:
                # Honest trials: each mean is a whole-number total divided by the arm's n (GRIM-consistent).
                mk, mc = round(mk * half) / half, round(mc * half) / half
            if form == "mean_sd":
                vals = {"K": {"mean": f"{mk:.2f}", "sd": f"{sd:.2f}", "n": str(half)}, "C": {"mean": f"{mc:.2f}", "sd": f"{sd:.2f}", "n": str(half)}}
            elif form == "mean_se":
                se = sd / half ** 0.5
                vals = {"K": {"mean": f"{mk:.2f}", "se": f"{se:.3f}", "n": str(half)}, "C": {"mean": f"{mc:.2f}", "se": f"{se:.3f}", "n": str(half)}}
            elif form == "median_iqr":
                vals = {"K": {"median": f"{mk:.0f}", "q1": f"{mk - 2.7:.0f}", "q3": f"{mk + 2.7:.0f}", "n": str(half)},
                        "C": {"median": f"{mc:.0f}", "q1": f"{mc - 2.7:.0f}", "q3": f"{mc + 2.7:.0f}", "n": str(half)}}
            elif form == "count":
                pk = min(0.95, max(0.02, 0.3 + 0.12 * g * shift))
                vals = {"K": {"events": str(round(pk * half)), "n": str(half)}, "C": {"events": str(round(0.3 * half)), "n": str(half)}}
            else:
                md, se = mk - mc, sd * (4 / n) ** 0.5
                vals = {"estimate": f"{md:.2f}", "lower": f"{md - 1.96 * se:.2f}", "upper": f"{md + 1.96 * se:.2f}", "level": "0.95", "scale": "mean_difference"}
            entries.append({"measure": "D", "day": day, "form": form, "location": "sim", "values": vals})
        pain_g = -0.5 + rng.gauss(0, 0.2)
        pain = [{"measure": "VAS rest", "day": 1, "form": "mean_sd", "location": "sim",
                 "values": {"K": {"mean": f"{3 + pain_g * 1.5:.1f}", "sd": "1.5", "n": str(half)}, "C": {"mean": "3.0", "sd": "1.5", "n": str(half)}}}]
        hall_k = rng.randint(0, 6 if group == "some" else 2)
        ae = [{"window": "postoperative", "values": {"K": {"hallucination": str(hall_k), "dizziness": str(rng.randint(2, 8))},
                                                     "C": {"hallucination": str(rng.randint(0, 1)), "dizziness": str(rng.randint(1, 5))}}}]
        extraction.append({"trial": tid, "source": "fulltext", "arms": {"K": {"randomized": str(half), "analysed": str(half)}, "C": {"randomized": str(half), "analysed": str(half)}},
                           "depression": entries, "pain": pain, "adverse_events": ae, "enrollment": {"start": "2022-03", "end": "2022-12"}})
    registry = {"registrations": [{"nct": f"NCT0{i:07d}", "eligible": True, "due": i < 30, "published": i < 18, "results_posted": i in (20, 21)} for i in range(40)]}
    (out / "trials.json").write_text(json.dumps(trials, indent=1))
    (out / "extraction.json").write_text(json.dumps(extraction, indent=1))
    (out / "registry.json").write_text(json.dumps(registry, indent=1))
    print("simulated", len(trials), "trials: masked g", masked_g, "awake g", awake_g)


def from_designs(trials_path, out):
    r = random.Random(11)
    real = json.loads(pathlib.Path(trials_path).read_text())
    out.mkdir(parents=True, exist_ok=True)
    shutil.copy(here.parent / "data" / "parameters.json", out / "parameters.json")
    shutil.copy(pathlib.Path(trials_path).parent / "registry.json", out / "registry.json")
    ext = []
    for t in real:
        d = t["design"]
        if not d.get("eligible"):
            continue
        arms = [a for a in d.get("arms") or [] if a.get("role") in ("ketamine", "control")]
        n = max(int(d.get("planned_n") or 80) // max(len(arms), 1), 5)
        g = (-0.3 if d.get("conscious_dosing") == "none" else -0.5) + r.gauss(0, 0.25)
        entries = []
        for m in d.get("depression_measures") or []:
            days = [x for x in (m.get("days") or [1, 3]) if isinstance(x, (int, float)) and x > 0]
            for day in [-1] + days:
                vals = {}
                for a in arms:
                    mean = 10 + (g * 4 if a["role"] == "ketamine" and day > 0 else 0) + r.gauss(0, 0.5)
                    vals[a["name"]] = {"mean": f"{round(mean * n) / n:.2f}", "sd": "4.00", "n": str(n)}
                entries.append({"measure": m["measure"], "day": day, "form": "mean_sd", "location": "sim", "values": vals})
        pain = [{"measure": m["measure"], "day": 1, "form": "mean_sd", "location": "sim",
                 "values": {a["name"]: {"mean": "2.3" if a["role"] == "ketamine" else "3.0", "sd": "1.5", "n": str(n)} for a in arms}}
                for m in d.get("pain_measures") or []]
        ae = [{"window": "postoperative", "values": {a["name"]: {"hallucination": str(r.randint(0, 3))} for a in arms}}]
        ext.append({"trial": t["trial"], "source": t["source"], "arms": {a["name"]: {"randomized": str(n), "analysed": str(n)} for a in arms},
                    "depression": entries, "pain": pain, "adverse_events": ae, "enrollment": d.get("enrollment") or {}})
    (out / "trials.json").write_text(json.dumps(real, indent=1))
    (out / "extraction.json").write_text(json.dumps(ext, indent=1))
    print("invented numbers for", len(ext), "real designs")


if __name__ == "__main__":
    if sys.argv[1] == "--designs":
        from_designs(sys.argv[2], pathlib.Path(sys.argv[3]))
    else:
        simulated(pathlib.Path(sys.argv[1]), float(sys.argv[2]) if len(sys.argv) > 2 else -0.4, float(sys.argv[3]) if len(sys.argv) > 3 else -0.6)