import json
import subprocess
import shutil
from pathlib import Path

BASE_DIR = Path(__file__).resolve().parents[1]
DATASET = BASE_DIR / "datasets" / "cosmo_params.json"
BACKUP = BASE_DIR / "datasets" / "cosmo_params_backup.json"
RESULTS_DIR = BASE_DIR / "results"
RESULTS_DIR.mkdir(exist_ok=True)
RESULTS_FILE = RESULTS_DIR / "sn_slope_sweep_fine.txt"

# 🎯 FINE SWEEP SETTINGS
START = -4.70
END = -4.60
STEP = 0.01


def run_gate() -> str:
    result = subprocess.run(
        ["python", "-m", "harness.gate_row20_cosmology"],
        capture_output=True,
        text=True
    )
    return result.stdout


def extract_json_block(output: str) -> dict | None:
    lines = output.splitlines()
    start_idx = None
    for i, line in enumerate(lines):
        if line.strip().startswith("{"):
            start_idx = i
            break
    if start_idx is None:
        return None

    json_text = "\n".join(lines[start_idx:])
    try:
        return json.loads(json_text)
    except json.JSONDecodeError:
        return None


def extract_sn_chi2(parsed: dict) -> float | None:
    if not parsed:
        return None
    for part in parsed.get("parts", []):
        if part.get("tag") == "SN":
            return part.get("chi2")
    return None


def main():
    print("=== SN SLOPE FINE SWEEP START ===")

    shutil.copy(DATASET, BACKUP)

    with open(DATASET, "r", encoding="utf-8") as f:
        params = json.load(f)

    results = []

    slope = START
    while slope <= END + 1e-12:
        slope = round(slope, 5)
        print(f"\nRunning slope = {slope:.3f}")

        params["sn_mu_slope"] = slope
        with open(DATASET, "w", encoding="utf-8") as f:
            json.dump(params, f, indent=2)

        output = run_gate()
        parsed = extract_json_block(output)
        sn_chi2 = extract_sn_chi2(parsed)

        print(f"SN chi2 = {sn_chi2}")
        results.append((slope, sn_chi2))

        slope += STEP

    shutil.copy(BACKUP, DATASET)

    with open(RESULTS_FILE, "w", encoding="utf-8") as f:
        for slope, chi2 in results:
            f.write(f"{slope:.5f}, {chi2}\n")

    valid_results = [(s, c) for s, c in results if c is not None]
    if valid_results:
        best = min(valid_results, key=lambda x: x[1])
        print("\n=== BEST RESULT ===")
        print(f"slope = {best[0]:.5f}, chi2 = {best[1]}")
    else:
        print("\n=== NO VALID RESULTS FOUND ===")

    print("\nResults saved to:", RESULTS_FILE)


if __name__ == "__main__":
    main()