#!/usr/bin/env python3
# Copyright (c) 2026 INNOVATIO SAS
# SPDX-License-Identifier: MIT
"""MSC-P-021: deterministic fixed, random and correlated-random-effects comparison.

The script uses only the Python standard library. By default it reads the published
synthetic panel CSV and prints the reference estimates. Use --emit-csv to recreate
the exact dataset from the declared data-generating process.
"""

from __future__ import annotations

import argparse
import csv
import math
from collections import defaultdict
from pathlib import Path


BASE_DEVIATIONS = (-3.5, -2.5, -1.5, -0.5, 0.5, 1.5, 2.5, 3.5)
BASE_SHOCKS = (1.0, -1.0, 1.0, -1.0, 1.0, -1.0, 1.0, -1.0)


def synthetic_rows() -> list[dict[str, float | int | str]]:
    rows: list[dict[str, float | int | str]] = []
    for store_index in range(12):
        store_mean = 15.0 + 2.0 * store_index
        stable_store_component = 2.0 * (1.0, -1.0, -1.0, 1.0)[store_index % 4]
        shock_sign = 1.0 if store_index % 2 == 0 else -1.0
        shift = store_index % len(BASE_DEVIATIONS)
        deviations = BASE_DEVIATIONS[shift:] + BASE_DEVIATIONS[:shift]
        shocks = BASE_SHOCKS[shift:] + BASE_SHOCKS[:shift]
        for week_index, (deviation, shock) in enumerate(zip(deviations, shocks), start=1):
            display_share = store_mean + deviation
            orders = (
                80.0
                - 1.2 * store_mean
                - 0.6 * deviation
                + stable_store_component
                + 0.8 * shock_sign * shock
            )
            rows.append(
                {
                    "store_id": f"S{store_index + 1:02d}",
                    "week": week_index,
                    "display_share_pct": round(display_share, 1),
                    "orders_per_1000": round(orders, 1),
                }
            )
    return rows


def read_rows(path: Path) -> list[dict[str, float | int | str]]:
    with path.open(newline="", encoding="utf-8") as handle:
        rows = []
        for raw in csv.DictReader(handle):
            rows.append(
                {
                    "store_id": raw["store_id"],
                    "week": int(raw["week"]),
                    "display_share_pct": float(raw["display_share_pct"]),
                    "orders_per_1000": float(raw["orders_per_1000"]),
                }
            )
    return rows


def write_rows(path: Path, rows: list[dict[str, float | int | str]]) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    with path.open("w", newline="", encoding="utf-8") as handle:
        writer = csv.DictWriter(
            handle,
            fieldnames=("store_id", "week", "display_share_pct", "orders_per_1000"),
            lineterminator="\n",
        )
        writer.writeheader()
        for row in rows:
            writer.writerow(row)


def slope_with_intercept(xs: list[float], ys: list[float]) -> tuple[float, float]:
    x_mean = sum(xs) / len(xs)
    y_mean = sum(ys) / len(ys)
    denominator = sum((x - x_mean) ** 2 for x in xs)
    slope = sum((x - x_mean) * (y - y_mean) for x, y in zip(xs, ys)) / denominator
    return y_mean - slope * x_mean, slope


def analyze(rows: list[dict[str, float | int | str]]) -> dict[str, float]:
    grouped: dict[str, list[dict[str, float | int | str]]] = defaultdict(list)
    for row in rows:
        grouped[str(row["store_id"])].append(row)

    n_entities = len(grouped)
    periods = {len(values) for values in grouped.values()}
    if periods != {8} or len(rows) != 96:
        raise ValueError("The reference calculation requires the declared balanced 12 x 8 panel.")

    means: dict[str, tuple[float, float]] = {}
    for store_id, values in grouped.items():
        means[store_id] = (
            sum(float(row["display_share_pct"]) for row in values) / len(values),
            sum(float(row["orders_per_1000"]) for row in values) / len(values),
        )

    xs = [float(row["display_share_pct"]) for row in rows]
    ys = [float(row["orders_per_1000"]) for row in rows]
    _, pooled = slope_with_intercept(xs, ys)

    dx = [float(row["display_share_pct"]) - means[str(row["store_id"])][0] for row in rows]
    dy = [float(row["orders_per_1000"]) - means[str(row["store_id"])][1] for row in rows]
    within_sxx = sum(value * value for value in dx)
    fixed = sum(x * y for x, y in zip(dx, dy)) / within_sxx
    fixed_residuals = [y - fixed * x for x, y in zip(dx, dy)]

    x_between = [value[0] for value in means.values()]
    y_between = [value[1] for value in means.values()]
    between_intercept, between = slope_with_intercept(x_between, y_between)
    between_residuals = [
        y - between_intercept - between * x for x, y in zip(x_between, y_between)
    ]

    sigma_e2 = sum(value * value for value in fixed_residuals) / (n_entities * (8 - 1) - 1)
    between_composite_variance = sum(value * value for value in between_residuals) / (n_entities - 2)
    sigma_u2 = max(0.0, between_composite_variance - sigma_e2 / 8)
    theta = 1.0 - math.sqrt(sigma_e2 / (sigma_e2 + 8 * sigma_u2))

    x_star = []
    y_star = []
    for row in rows:
        x_bar, y_bar = means[str(row["store_id"])]
        x_star.append(float(row["display_share_pct"]) - theta * x_bar)
        y_star.append(float(row["orders_per_1000"]) - theta * y_bar)
    _, random_effects = slope_with_intercept(x_star, y_star)

    cluster_scores = defaultdict(float)
    for row, x_dev, residual in zip(rows, dx, fixed_residuals):
        cluster_scores[str(row["store_id"])] += x_dev * residual
    finite_cluster_correction = n_entities / (n_entities - 1)
    fixed_cluster_variance = finite_cluster_correction * sum(score * score for score in cluster_scores.values()) / within_sxx**2
    fixed_cluster_se = math.sqrt(fixed_cluster_variance)
    # Two-sided t critical value for 11 cluster degrees of freedom. This is a
    # declared teaching convention. With only 12 clusters, the interval is a
    # worked example rather than a general small-sample inference guarantee.
    t975_df11 = 2.200985

    return {
        "observations": float(len(rows)),
        "stores": float(n_entities),
        "pooled_slope": pooled,
        "fixed_within_slope": fixed,
        "fixed_cluster_se": fixed_cluster_se,
        "fixed_ci_low": fixed - t975_df11 * fixed_cluster_se,
        "fixed_ci_high": fixed + t975_df11 * fixed_cluster_se,
        "cluster_degrees_of_freedom": float(n_entities - 1),
        "finite_cluster_correction": finite_cluster_correction,
        "between_slope": between,
        "cre_within_slope": fixed,
        "cre_between_slope": between,
        "cre_contextual_difference": between - fixed,
        "standard_re_slope": random_effects,
        "theta": theta,
        "sigma_e2": sigma_e2,
        "sigma_u2": sigma_u2,
        "icc": sigma_u2 / (sigma_u2 + sigma_e2),
    }


def main() -> None:
    parser = argparse.ArgumentParser()
    default_csv = Path(__file__).resolve().parents[1] / "datasets" / "msc-p021-panel.csv"
    parser.add_argument("--csv", type=Path, default=default_csv)
    parser.add_argument("--emit-csv", action="store_true")
    args = parser.parse_args()

    if args.emit_csv:
        write_rows(args.csv, synthetic_rows())

    results = analyze(read_rows(args.csv))
    for key, value in results.items():
        print(f"{key}={value:.6f}")


if __name__ == "__main__":
    main()