"""Canonical reproduction of the controlled RCH toy comparison.

Outputs
-------
table2_data.csv
    Canonical numerical results used in Table 2 and Figure 1.
results.xlsx
    Excel workbook containing the same canonical results, run metadata,
    and a native Excel chart linked to the canonical results.
figures/fig-000.png and figures/fig-000.pdf
    Publication Figure 1 regenerated from the canonical Table 2 values.
figures/fig-001.png and figures/fig-001.pdf
    Publication Figure 2 regenerated from the canonical Excel sample-size sensitivity sheet.

The Python and Excel outputs are generated from the same deterministic
simulation seed. The Excel chart is therefore a second presentation of the
same numerical results rather than an independently estimated result.
"""

import csv
import math
from pathlib import Path

import matplotlib.pyplot as plt
import numpy as np
from openpyxl import Workbook, load_workbook
from openpyxl.chart import LineChart, Reference
from openpyxl.styles import Font, Alignment

SEED = 20260827
A_VALUES = [0.0, 0.2, 0.4, 0.6, 0.8, 0.95]
REPETITIONS = 1000
N_CAL = 1000
N_TEST = 1000

OUT = Path(__file__).resolve().parent
FIG_DIR = OUT / "figures"
FIG_DIR.mkdir(exist_ok=True)


def simulate(a, repetitions, n_cal, n_test, rng):
    values = []
    for _ in range(repetitions):
        xcal = np.empty(n_cal)
        xtest = np.empty(n_test)
        xcal[0] = rng.normal()
        xtest[0] = rng.normal()
        scale = math.sqrt(1.0 - a * a)
        ec = rng.normal(size=n_cal - 1) * scale
        et = rng.normal(size=n_test - 1) * scale
        for i in range(1, n_cal):
            xcal[i] = a * xcal[i - 1] + ec[i - 1]
        for i in range(1, n_test):
            xtest[i] = a * xtest[i - 1] + et[i - 1]

        ahat = np.dot(xcal[:-1], xcal[1:]) / np.dot(xcal[:-1], xcal[:-1])
        sigma2_r = 1.0 - ahat * ahat
        mu = np.mean(xcal)
        sigma2_n = np.mean((xcal - mu) ** 2)

        y = xtest[1:]
        prev = xtest[:-1]
        ll_r = (
            -0.5 * (n_test - 1) * math.log(2 * math.pi * sigma2_r)
            -0.5 * np.sum((y - ahat * prev) ** 2) / sigma2_r
        )
        ll_n = (
            -0.5 * (n_test - 1) * math.log(2 * math.pi * sigma2_n)
            -0.5 * np.sum((y - mu) ** 2) / sigma2_n
        )
        values.append(ll_r - ll_n)

    values = np.asarray(values)
    return {
        "a": a,
        "mean_Acomp": float(values.mean()),
        "p5_Acomp": float(np.percentile(values, 5)),
        "p95_Acomp": float(np.percentile(values, 95)),
        "positive_percent": float(np.mean(values > 0) * 100),
    }


def write_csv(rows):
    path = OUT / "table2_data.csv"
    header = ["a", "mean_Acomp", "p5_Acomp", "p95_Acomp", "positive_percent"]
    with path.open("w", encoding="utf-8", newline="") as f:
        f.write(",".join(header) + "\n")
        for r in rows:
            f.write(
                f"{r['a']:.2f},{r['mean_Acomp']:.2f},{r['p5_Acomp']:.2f},"
                f"{r['p95_Acomp']:.2f},{r['positive_percent']:.1f}\n"
            )
    return path


def read_canonical_csv(path):
    rows = []
    with path.open("r", encoding="utf-8", newline="") as f:
        for r in csv.DictReader(f):
            rows.append({
                "a": float(r["a"]),
                "mean_Acomp": float(r["mean_Acomp"]),
                "p5_Acomp": float(r["p5_Acomp"]),
                "p95_Acomp": float(r["p95_Acomp"]),
                "positive_percent": float(r["positive_percent"]),
            })
    return rows


def write_excel(rows, sensitivity_rows):
    path = OUT / "results.xlsx"
    wb = Workbook()
    wb.properties.title = "RCH canonical numerical results and reproducibility materials"
    wb.properties.subject = "Controlled recursive-versus-non-recursive toy comparison"
    wb.properties.creator = "Anonymous Author"

    info = wb.active
    info.title = "Article metadata"
    metadata_rows = [
        ("Article title", "Testing Recursive World-Generating Processes: An Abductive and Information-Theoretic Framework for the Recursive Cosmological Hypothesis"),
        ("Target journal", "Foundations of Science"),
        ("Author", "Anonymous Author"),
        ("Affiliation", "Anonymous Researcher, Australia, Australia"),
        ("Corresponding email", ""),
        ("Workbook purpose", "Canonical numerical results, run metadata, and sample-size sensitivity supporting the manuscript"),
    ]
    info.append(["Item", "Value"])
    for k, v in metadata_rows:
        info.append([k, v])
    for cell in info[1]:
        cell.font = Font(bold=True)
    info.column_dimensions["A"].width = 24
    info.column_dimensions["B"].width = 110
    info.freeze_panes = "A2"

    ws = wb.create_sheet("Table 2")
    ws.title = "Table 2"

    headers = ["a", "Mean A_comp", "5th percentile", "95th percentile", "A_comp > 0 (%)"]
    ws.append(headers)
    for cell in ws[1]:
        cell.font = Font(bold=True)
    for r in rows:
        ws.append([
            r["a"], r["mean_Acomp"], r["p5_Acomp"], r["p95_Acomp"], r["positive_percent"]
        ])

    for row in ws.iter_rows(min_row=2, max_col=5):
        row[0].number_format = "0.00"
        for cell in row[1:4]:
            cell.number_format = "0.00"
        row[4].number_format = "0.0"

    ws.freeze_panes = "A2"
    widths = [10, 18, 18, 18, 18]
    for i, width in enumerate(widths, start=1):
        ws.column_dimensions[chr(64 + i)].width = width

    chart = LineChart()
    chart.title = "Held-out compression advantage"
    chart.y_axis.title = "Mean A_comp (nats)"
    chart.x_axis.title = "Recursive strength a"
    data = Reference(ws, min_col=2, min_row=1, max_row=1 + len(rows))
    cats = Reference(ws, min_col=1, min_row=2, max_row=1 + len(rows))
    chart.add_data(data, titles_from_data=True)
    chart.set_categories(cats)
    chart.height = 9
    chart.width = 17
    ws.add_chart(chart, "G2")

    meta = wb.create_sheet("Run metadata")
    metadata = [
        ("Random seed", SEED),
        ("Canonical repetitions", REPETITIONS),
        ("Calibration observations", N_CAL),
        ("Held-out test observations", N_TEST),
        ("Generator", "Stationary Gaussian AR(1)"),
        ("Null at a=0", "Recursive and independent Gaussian generators coincide"),
        ("Statistic", "A_comp = -ln P(D_test|theta_N,M_N) + ln P(D_test|theta_R,M_R)"),
        ("Figure source", "table2_data.csv regenerated by reproduction.py"),
        ("Excel role", "Independent presentation/check of the same canonical numerical results"),
    ]
    meta.append(["Item", "Value"])
    for k, v in metadata:
        meta.append([k, v])
    for cell in meta[1]:
        cell.font = Font(bold=True)
    meta.column_dimensions["A"].width = 32
    meta.column_dimensions["B"].width = 90

    sens = wb.create_sheet("Sample-size sensitivity")
    sens_headers = ["Calibration/Test N", "a", "Mean A_comp", "5th percentile", "95th percentile", "A_comp > 0 (%)"]
    sens.append(sens_headers)
    for cell in sens[1]:
        cell.font = Font(bold=True)
    for r in sensitivity_rows:
        sens.append([r["n"], r["a"], r["mean_Acomp"], r["p5_Acomp"], r["p95_Acomp"], r["positive_percent"]])
    for i, width in enumerate([20, 10, 18, 18, 18, 18], start=1):
        sens.column_dimensions[chr(64 + i)].width = width

    # Native Excel presentation of the same sensitivity data used for Figure 2.
    # The publication PNG is still regenerated by Python after reading this sheet.
    chart2 = LineChart()
    chart2.title = "Sample-size sensitivity"
    chart2.y_axis.title = "Mean A_comp (nats)"
    chart2.x_axis.title = "Recursive strength a"
    for n, col in [(500, 3), (1000, 3), (2000, 3)]:
        rows_for_n = [i + 2 for i, r in enumerate(sensitivity_rows) if r["n"] == n]
        if not rows_for_n:
            continue
        start_row, end_row = min(rows_for_n), max(rows_for_n)
        data = Reference(sens, min_col=col, min_row=start_row - 1, max_row=end_row)
        chart2.add_data(data, titles_from_data=True)
        cats = Reference(sens, min_col=2, min_row=start_row, max_row=end_row)
        chart2.set_categories(cats)
    chart2.height = 9
    chart2.width = 17
    sens.add_chart(chart2, "H2")

    wb.save(path)
    return path


def write_figure(excel_path):
    # The publication figure is regenerated from the canonical Excel Table 2
    # after the CSV has been written and checked. This keeps the displayed
    # figure, table, and workbook on one canonical numerical path.
    wb = load_workbook(excel_path, data_only=True, read_only=True)
    ws = wb["Table 2"]
    records = list(ws.iter_rows(min_row=2, values_only=True))
    wb.close()

    a = np.array([r[0] for r in records], dtype=float)
    mean = np.array([r[1] for r in records], dtype=float)
    lo = np.array([r[2] for r in records], dtype=float)
    hi = np.array([r[3] for r in records], dtype=float)

    fig, ax = plt.subplots(figsize=(8.6, 5.5), dpi=300)
    ax.errorbar(
        a,
        mean,
        yerr=[mean - lo, hi - mean],
        fmt="o-",
        capsize=4,
        linewidth=1.5,
        markersize=6,
    )
    ax.axhline(0, linewidth=1)
    ax.set_xlabel(r"Recursive strength $a$")
    ax.set_ylabel(r"Mean $A_{\mathrm{comp}}$ (nats)")
    ax.set_title("Held-out compression advantage")
    ax.grid(False)
    fig.tight_layout()
    path = FIG_DIR / "fig-000.png"
    vector_path = FIG_DIR / "fig-000.pdf"
    fig.savefig(path, bbox_inches="tight", dpi=300)
    fig.savefig(vector_path, bbox_inches="tight")
    plt.close(fig)
    return path, vector_path



def write_sensitivity_figure(excel_path):
    """Regenerate Figure 2 from the canonical Excel sensitivity sheet."""
    wb = load_workbook(excel_path, data_only=True, read_only=True)
    ws = wb["Sample-size sensitivity"]
    records = list(ws.iter_rows(min_row=2, values_only=True))
    wb.close()

    groups = {}
    for n, a, mean, p5, p95, positive in records:
        groups.setdefault(int(n), []).append((float(a), float(mean)))

    fig, ax = plt.subplots(figsize=(8.6, 5.5), dpi=300)
    for n in sorted(groups):
        pairs = sorted(groups[n])
        ax.plot(
            [p[0] for p in pairs],
            [p[1] for p in pairs],
            marker="o", linewidth=1.5, markersize=5, label=f"N = {n:,}"
        )
    ax.axhline(0, linewidth=1)
    ax.set_xlabel(r"Recursive strength $a$")
    ax.set_ylabel(r"Mean $A_{\mathrm{comp}}$ (nats)")
    ax.set_title("Sample-size sensitivity of the held-out comparison")
    ax.legend(frameon=False)
    ax.grid(False)
    fig.tight_layout()
    path = FIG_DIR / "fig-001.png"
    vector_path = FIG_DIR / "fig-001.pdf"
    fig.savefig(path, bbox_inches="tight", dpi=300)
    fig.savefig(vector_path, bbox_inches="tight")
    plt.close(fig)
    return path, vector_path

def main():
    rng = np.random.default_rng(SEED)
    rows = [simulate(a, REPETITIONS, N_CAL, N_TEST, rng) for a in A_VALUES]

    # Independent sample-size sensitivity check, also deterministic.
    sensitivity_rows = []
    for n in [500, 1000, 2000]:
        sensitivity_rng = np.random.default_rng(SEED + n)
        for a in A_VALUES:
            r = simulate(a, 1000, n, n, sensitivity_rng)
            sensitivity_rows.append({"n": n, **r})

    csv_path = write_csv(rows)
    canonical_rows = read_canonical_csv(csv_path)
    excel_path = write_excel(canonical_rows, sensitivity_rows)
    write_figure(excel_path)
    write_sensitivity_figure(excel_path)

    print("Canonical results")
    print("a,mean_Acomp,p5_Acomp,p95_Acomp,positive_percent")
    for r in rows:
        print(
            f"{r['a']:.2f},{r['mean_Acomp']:.2f},{r['p5_Acomp']:.2f},"
            f"{r['p95_Acomp']:.2f},{r['positive_percent']:.1f}"
        )


if __name__ == "__main__":
    main()
