#!/usr/bin/env python3
import csv
import json
from pathlib import Path

import matplotlib

matplotlib.use("Agg")
import matplotlib.pyplot as plt
from matplotlib.colors import BoundaryNorm, ListedColormap
from matplotlib.patches import Patch
import numpy as np


ROOT = Path(__file__).resolve().parents[1]
DATA = ROOT / "data" / "F3_marriott_documentary_state_r2_data.csv"
OUT_PNG = ROOT / "figures" / "F3_marriott_documentary_state_r2.png"
OUT_SVG = ROOT / "figures" / "F3_marriott_documentary_state_r2.svg"
OUT_VALIDATION = ROOT / "validation" / "F3_marriott_documentary_state_r2_validation.json"
FIGSIZE = (10.4, 6.4)
CELL_FONT_SIZE = 8.5
FIGURE_BADGE = "NOT_GOLD · DOCUMENTARY SUPPORT ONLY"
SHOW_INTERNAL_TITLE = False
STATES = [
    "BASELINE_ESTABLISHED",
    "REMAINS_SUPPORTED",
    "REQUIRES_REASSESSMENT",
    "UNRESOLVED",
    "CHANGED",
    "SUPERSEDED",
    "NOT_APPLICABLE",
]
DISPLAY_LABELS = {
    "BASELINE_ESTABLISHED": "BASELINE",
    "REMAINS_SUPPORTED": "SUPPORTED",
    "REQUIRES_REASSESSMENT": "REASSESS",
    "UNRESOLVED": "UNRESOLVED",
    "CHANGED": "CHANGED",
    "SUPERSEDED": "SUPERSEDED",
    "NOT_APPLICABLE": "N/A",
}
COLORS = {
    "BASELINE_ESTABLISHED": "#d6dde5",
    "REMAINS_SUPPORTED": "#dce8e2",
    "REQUIRES_REASSESSMENT": "#e8e2d7",
    "UNRESOLVED": "#eee9d8",
    "CHANGED": "#cbd8e8",
    "SUPERSEDED": "#e5d7dc",
    "NOT_APPLICABLE": "#f2f2f2",
}


def cell_label(state, source_id=None):
    del source_id
    return DISPLAY_LABELS[state]


def validate_rows(rows):
    required = {
        "claim_id", "claim_label", "event_id", "event_label", "event_date",
        "publication_date", "state", "source_id", "passage_locator", "basis",
        "evidence_class", "unresolved_limitation",
    }
    checked = []
    for row in rows:
        if not required.issubset(row):
            raise ValueError(f"missing columns: {sorted(required - set(row))}")
        item = dict(row)
        if item["state"] not in STATES:
            raise ValueError(f"invalid documentary state: {item['state']}")
        for field in required - {"state"}:
            if not item.get(field):
                raise ValueError(f"blank {field} in {item.get('claim_id')}/{item.get('event_id')}")
        if item["evidence_class"] != "documentary_not_gold":
            raise ValueError("every cell must remain documentary_not_gold")
        checked.append(item)
    return checked


def read_rows(path=DATA):
    with Path(path).open(newline="", encoding="utf-8") as handle:
        rows = validate_rows(list(csv.DictReader(handle)))
    claims = list(dict.fromkeys(row["claim_id"] for row in rows))
    events = list(dict.fromkeys(row["event_id"] for row in rows))
    if len(rows) != len(claims) * len(events):
        raise ValueError("documentary matrix is incomplete")
    if len({(row["claim_id"], row["event_id"]) for row in rows}) != len(rows):
        raise ValueError("duplicate claim/event cell")
    return rows, claims, events


def make_figure(rows, claims, events):
    claim_labels = {row["claim_id"]: row["claim_label"] for row in rows}
    by_key = {(row["claim_id"], row["event_id"]): row for row in rows}
    state_index = {state: index for index, state in enumerate(STATES)}
    matrix = np.array([[state_index[by_key[(claim, event)]["state"]] for event in events] for claim in claims])

    plt.rcParams.update({
        "font.family": "DejaVu Sans",
        "font.size": 9,
        "axes.titlesize": 13,
        "svg.hashsalt": "eoji-v07-r2-f3",
    })
    fig, ax = plt.subplots(figsize=FIGSIZE)
    cmap = ListedColormap([COLORS[state] for state in STATES])
    norm = BoundaryNorm(np.arange(-0.5, len(STATES) + 0.5, 1), cmap.N)
    ax.imshow(matrix, cmap=cmap, norm=norm, aspect="auto")

    for i, claim in enumerate(claims):
        for j, event in enumerate(events):
            row = by_key[(claim, event)]
            ax.text(j, i, cell_label(row["state"], row["source_id"]), ha="center", va="center", fontsize=CELL_FONT_SIZE, fontweight="bold")

    wrapped_events = {
        "E1": "2019 agreement\nMA01",
        "E2": "Apr-2020\namendment\nMA02/MA03",
        "E3": "Q1 report\n+ events\nMA04",
        "E4": "Q2 report\nMA05",
        "E5": "Jan-2021\nSecond amdt\nMA07",
        "E6": "Jan-2021\nThird amdt\nMA08",
        "E7": "2020 10-K\nFeb-2021 · MA06",
    }
    ax.set_xticks(range(len(events)), [wrapped_events[event] for event in events])
    ax.set_yticks(range(len(claims)), [claim_labels[claim] for claim in claims])
    ax.tick_params(axis="x", top=True, bottom=False, labeltop=True, labelbottom=False, length=0, pad=8)
    ax.tick_params(axis="y", length=0, pad=8)
    for label in ax.get_xticklabels():
        label.set_fontsize(8.6)
    for label in ax.get_yticklabels():
        label.set_fontsize(10.0)
    ax.set_xticks(np.arange(-0.5, len(events), 1), minor=True)
    ax.set_yticks(np.arange(-0.5, len(claims), 1), minor=True)
    ax.grid(which="minor", color="white", linestyle="-", linewidth=2)
    ax.tick_params(which="minor", bottom=False, left=False)
    for spine in ax.spines.values():
        spine.set_visible(False)

    if SHOW_INTERNAL_TITLE:
        fig.suptitle("Selective documentary state change across one public facility history", y=0.97, fontsize=14, fontweight="bold")
    fig.text(0.5, 0.968, FIGURE_BADGE, ha="center", fontsize=10, fontweight="bold", color="#3f4854")
    legend = [Patch(facecolor=COLORS[state], edgecolor="#8b9299", label=DISPLAY_LABELS[state]) for state in STATES]
    ax.legend(handles=legend, ncol=4, loc="lower center", bbox_to_anchor=(0.5, -0.25), frameon=False, fontsize=8.2)
    fig.subplots_adjust(left=0.285, right=0.99, top=0.80, bottom=0.24)
    fig.text(
        0.5,
        0.035,
        "Event time and publication time remain distinct. Source IDs are column-level; exact locators and limitations are in the source map.",
        ha="center",
        fontsize=8.3,
        color="#3f4854",
    )
    metadata = {"Creator": "EOJI v0.7 R2 readability correction", "Date": None}
    fig.savefig(OUT_PNG, dpi=220, metadata=metadata)
    fig.savefig(OUT_SVG, metadata=metadata)
    plt.close(fig)


def build_validation(rows, claims, events):
    by_key = {(row["claim_id"], row["event_id"]): row for row in rows}
    expected_corrections = {
        "C1/E1": "BASELINE_ESTABLISHED",
        "C5/E1": "BASELINE_ESTABLISHED",
        "C3/E5": "UNRESOLVED",
        "C3/E6": "UNRESOLVED",
    }
    actual = {key: by_key[tuple(key.split("/"))]["state"] for key in expected_corrections}
    if actual != expected_corrections:
        raise ValueError(f"mandatory corrections failed: {actual}")
    if any("does not establish facility balance" not in by_key[("C3", event)]["basis"] for event in ("E5", "E6")):
        raise ValueError("January 2021 balance basis is not explicit")
    return {
        "status": "PASS_WITH_SOURCE_CAPTURE_ADJUDICATION_AND_MAIN_TEXT_PLACEMENT_GAPS",
        "evidence_class": "documentary_not_gold",
        "cells": len(rows),
        "claims": len(claims),
        "events": len(events),
        "valid_state_taxonomy": True,
        "complete_claim_event_matrix": True,
        "all_cells_source_bound": all(row["source_id"] and row["passage_locator"] for row in rows),
        "all_cells_have_limitations": all(row["unresolved_limitation"] for row in rows),
        "event_and_publication_time_separate": True,
        "mandatory_corrections": actual,
        "r2_readability_change": {
            "data_changed": False,
            "internal_title_removed": True,
            "not_gold_badge_visible": True,
            "cell_font_points": CELL_FONT_SIZE,
            "figure_size_inches": list(FIGSIZE),
            "legend_columns": 4
        },
        "source_ids": sorted({source for row in rows for source in row["source_id"].split("|")}),
        "raw_source_capture": {
            "status": "BLOCKED",
            "requests": 9,
            "accepted_raw_bytes": 0,
            "reason": "official SEC endpoints returned 403 undeclared-automated-tool policy pages",
        },
        "outputs": [str(OUT_PNG.relative_to(ROOT)), str(OUT_SVG.relative_to(ROOT))],
        "open_gates": [
            "no accepted local raw SEC bytes/hashes",
            "no qualified legal/financial adjudication",
            "single selected public case",
            "exact publication dates unresolved for columns E1-E6",
        ],
        "exclusions": [
            "not legal default",
            "not compliance ground truth",
            "not action authority",
            "not current creditworthiness",
            "not prevalence",
            "not EOJI performance",
        ],
    }


def main():
    rows, claims, events = read_rows()
    make_figure(rows, claims, events)
    OUT_VALIDATION.write_text(json.dumps(build_validation(rows, claims, events), indent=2) + "\n", encoding="utf-8")


if __name__ == "__main__":
    main()
