from __future__ import annotations

import argparse
import csv
from pathlib import Path

ROOT = Path(__file__).resolve().parents[1]
SOURCE = ROOT / "data/input/tunisia_source_aligned_quarterly_panel.csv"
CANONICAL = ROOT / "data/input/tunisia_aligned_quarterly_panel.csv"
CORE = ["trust_obs", "shock_obs", "policy_proxy"]


def expected_quarters() -> list[str]:
    return [f"{year}Q{quarter}" for year in range(2000, 2020) for quarter in range(1, 5)]


def read_source() -> list[dict[str, str]]:
    with SOURCE.open(newline="", encoding="utf-8-sig") as handle:
        rows = list(csv.DictReader(handle))
    required = {"quarter", "year_decimal", *CORE}
    missing = required.difference(rows[0].keys() if rows else [])
    if missing:
        raise ValueError(f"Source panel is missing columns: {sorted(missing)}")
    if len(rows) != 80:
        raise ValueError(f"Expected 80 rows, found {len(rows)}")
    quarters = [row["quarter"] for row in rows]
    if quarters != expected_quarters():
        raise ValueError("Quarter sequence is not exactly 2000Q1--2019Q4")
    if len(set(quarters)) != len(quarters):
        raise ValueError("Duplicate quarter labels found")
    for row in rows:
        for coordinate in CORE:
            value = float(row[coordinate])
            if not (0.0 <= value <= 1.0):
                raise ValueError(f"{coordinate} outside [0,1] at {row['quarter']}: {value}")
    return rows


def projection(rows: list[dict[str, str]]) -> list[dict[str, str | int]]:
    return [
        {
            "quarter": row["quarter"],
            "year": row["year_decimal"],
            "quarter_index": index,
            "trust_obs": row["trust_obs"],
            "shock_obs": row["shock_obs"],
            "policy_proxy": row["policy_proxy"],
        }
        for index, row in enumerate(rows)
    ]


def write_projection(rows: list[dict[str, str]]) -> None:
    fields = ["quarter", "year", "quarter_index", *CORE]
    with CANONICAL.open("w", newline="", encoding="utf-8") as handle:
        writer = csv.DictWriter(handle, fieldnames=fields)
        writer.writeheader()
        writer.writerows(projection(rows))


def check_projection(rows: list[dict[str, str]]) -> None:
    if not CANONICAL.exists():
        raise FileNotFoundError(f"Missing canonical analysis input: {CANONICAL}")
    with CANONICAL.open(newline="", encoding="utf-8-sig") as handle:
        deposited = list(csv.DictReader(handle))
    expected = projection(rows)
    if len(deposited) != len(expected):
        raise ValueError("Canonical projection row count differs from the source panel")
    fields = ["quarter", "year", "quarter_index", *CORE]
    for index, (actual, target) in enumerate(zip(deposited, expected)):
        for field in fields:
            if str(actual[field]) != str(target[field]):
                raise ValueError(
                    f"Projection mismatch at row {index + 2}, field {field}: "
                    f"deposited={actual[field]!r}, expected={target[field]!r}"
                )
    print("Source projection PASS: canonical analysis input is an exact field projection of the author-supplied panel.")


def main(check: bool) -> None:
    rows = read_source()
    if check:
        check_projection(rows)
    else:
        write_projection(rows)
        check_projection(rows)


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--check", action="store_true", help="Check without rewriting the canonical CSV")
    arguments = parser.parse_args()
    main(arguments.check)
