from __future__ import annotations

import json
import platform
from pathlib import Path
from typing import Any

import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
from matplotlib.patches import FancyArrowPatch, FancyBboxPatch
import numpy as np

from .configuration import load_study_configuration
from .io import sha256_file, sha256_tree, write_csv, write_json
from .mission import plan_virtual_mission
from .simulator import (
    analyse_virtual_mosaic,
    coverage_grid,
    simulate_capture_qc,
    simulate_dn_and_calibrate,
    virtual_canopy,
)
from .specs import MAVIC_3M


ROOT_FIELDS = [
    "captureId", "line", "sequenceOnLine", "xM", "yM", "altitudeAglM", "elapsedS",
    "fileCountExpected", "rtkFlag", "rtkInterpretation", "sunAltitudeDeg",
    "sunlightSensorRelative", "exposureTimeUs", "iso", "blurScore", "saturationFraction",
    "missingBand", "filenames", "fileCountObserved", "qcPassed", "qcFailures",
]


def _save_figures(root: Path, mission: dict[str, Any], records: list[dict[str, Any]], arrays: dict[str, np.ndarray], canopy: dict[str, np.ndarray], result: dict[str, Any]) -> None:
    out = root / "outputs" / "figures"
    out.mkdir(parents=True, exist_ok=True)
    plt.rcParams.update({"font.family": "sans-serif", "axes.titlesize": 11, "axes.labelsize": 9})

    fig, ax = plt.subplots(figsize=(10.2, 6.2), constrained_layout=True)
    accepted = [row for row in records if row["qcPassed"]]
    rejected = [row for row in records if not row["qcPassed"]]
    for line in range(1, mission["scanLineCount"] + 1):
        row = [event for event in records if event["line"] == line]
        ax.plot([event["xM"] for event in row], [event["yM"] for event in row], color="#8d6e63", lw=0.9)
    ax.scatter([row["xM"] for row in accepted], [row["yM"] for row in accepted], s=8, color="#2e7d32", label="QC accepted")
    ax.scatter([row["xM"] for row in rejected], [row["yM"] for row in rejected], s=30, marker="x", color="#c62828", label="QC rejected")
    ax.set(xlim=(0, 400), ylim=(0, 300), xlabel="local east-west distance (m)", ylabel="local north-south distance (m)")
    ax.set_title("Virtual 12 ha route and synchronized five-camera capture events")
    ax.set_aspect("equal")
    ax.legend(loc="upper right")
    ax.text(3, 4, "LOCAL_SYNTHETIC_METRES · NOT FOR FLIGHT", color="#c62828", fontsize=9, weight="bold")
    fig.savefig(out / "figure-1-mission-route.png", dpi=190)
    plt.close(fig)

    fig, ax = plt.subplots(figsize=(10.4, 4.6), constrained_layout=True)
    ax.axis("off")
    labels = [
        (0.02, "1  Virtual AOI\n400 × 300 m"), (0.20, "2  Mission plan\n60 m · 80/70 overlap"),
        (0.40, "3  Five cameras\nRGB + G/R/RE/NIR"), (0.60, "4  Fail-closed QC\nRTK · blur · saturation"),
        (0.80, "5  Relative map\nindices + review zones"),
    ]
    for x, label in labels:
        box = FancyBboxPatch((x, 0.35), 0.16, 0.30, boxstyle="round,pad=0.02", fc="#f7f1e8", ec="#805233", lw=1.2)
        ax.add_patch(box)
        ax.text(x + 0.08, 0.50, label, ha="center", va="center", fontsize=9)
    for x in (0.18, 0.38, 0.58, 0.78):
        ax.add_patch(FancyArrowPatch((x, 0.50), (x + 0.02, 0.50), arrowstyle="-|>", mutation_scale=13, color="#476a50"))
    ax.text(0.5, 0.80, "End-to-end digital-twin demo (no aircraft connection, no field flight)", ha="center", fontsize=13, weight="bold")
    ax.text(0.5, 0.16, "Every stage emits a machine-readable receipt; rejected captures cannot enter the mosaic.", ha="center", color="#555")
    fig.savefig(out / "figure-2-system-flow.png", dpi=190)
    plt.close(fig)

    fig, axes = plt.subplots(2, 3, figsize=(11.2, 7.0), constrained_layout=True)
    panels = [
        (canopy["rgb"], "synthetic RGB", None, None, None),
        (arrays["ndvi"], "NDVI", "YlGn", 0.25, 0.9),
        (arrays["ndre"], "NDRE", "YlGn", 0.15, 0.7),
        (arrays["gndvi"], "GNDVI", "YlGn", 0.25, 0.9),
        (arrays["anomaly_score"], "relative spectral-low score", "magma", 0, 1),
        (arrays["review_mask"], "ground-review priority", "Reds", 0, 1),
    ]
    for ax, (value, title, cmap, low, high) in zip(axes.ravel(), panels):
        image = ax.imshow(value, cmap=cmap, vmin=low, vmax=high, origin="lower")
        ax.set_title(title)
        ax.set_axis_off()
        if cmap is not None:
            fig.colorbar(image, ax=ax, fraction=0.044, pad=0.02)
    fig.suptitle("Pose-assisted synthetic mosaic: analysis output is review priority, not diagnosis", fontsize=13)
    fig.savefig(out / "figure-3-analysis-panel.png", dpi=190)
    plt.close(fig)

    fig, axes = plt.subplots(1, 2, figsize=(10.6, 4.3), constrained_layout=True)
    reasons = result.get("captureQcSummary", {}).get("failureReasonCounts", {})
    axes[0].bar(list(reasons), list(reasons.values()), color="#9c6644")
    axes[0].set_title("Synthetic capture rejections by gate")
    axes[0].tick_params(axis="x", rotation=25)
    axes[0].set_ylabel("capture events")
    coverage = arrays["coverage"]
    values, counts = np.unique(coverage, return_counts=True)
    axes[1].bar(values, counts / counts.sum() * 100, color="#47704f")
    axes[1].set_title("Accepted-frame coverage multiplicity")
    axes[1].set(xlabel="accepted frames covering each 1 m² cell", ylabel="field area (%)")
    fig.savefig(out / "figure-4-qc-coverage.png", dpi=190)
    plt.close(fig)


def run_demo(
    study_root: str | Path,
    seed: int = 20260804,
    *,
    output_root: str | Path | None = None,
    mission_config: str | Path | None = None,
    field_config: str | Path | None = None,
) -> dict[str, Any]:
    source_root = Path(study_root)
    artifact_root = Path(output_root) if output_root is not None else source_root
    demo = artifact_root / "outputs" / "demo"
    demo.mkdir(parents=True, exist_ok=True)
    field, parameters, configuration_paths = load_study_configuration(
        source_root,
        mission_config=mission_config,
        field_config=field_config,
    )
    mission = plan_virtual_mission(field, parameters)
    records, qc_summary = simulate_capture_qc(
        mission["captureEvents"],
        sun_altitude_deg=parameters.sun_altitude_deg,
    )
    geometry = mission["cameraGeometry"]
    coverage = coverage_grid(
        records,
        field,
        float(geometry["multispectralFootprintAlongTrackM"]),
        float(geometry["multispectralFootprintCrossTrackM"]),
    )
    canopy = virtual_canopy(field, seed)
    source_bands = {name: canopy[name] for name in ("green", "red", "red_edge", "nir")}
    raw_dn, calibrated, calibration = simulate_dn_and_calibrate(source_bands, seed)
    arrays, analysis = analyse_virtual_mosaic(canopy, calibrated, coverage)
    analysis["captureQcSummary"] = qc_summary

    mission_public = {key: value for key, value in mission.items() if key != "captureEvents"}
    write_json(demo / "mavic-3m-specification-receipt.json", MAVIC_3M.as_receipt())
    write_json(demo / "mission-summary.json", mission_public)
    write_json(demo / "capture-qc-summary.json", qc_summary)
    write_csv(demo / "capture-manifest.csv", records, ROOT_FIELDS)
    write_json(demo / "radiometric-processing-receipt.json", calibration)
    write_json(demo / "system-analysis.json", analysis)
    np.savez_compressed(
        demo / "synthetic-analysis-arrays.npz",
        **{f"dn_{name}": value for name, value in raw_dn.items()},
        **{f"reflectance_{name}": value for name, value in calibrated.items()},
        rgb=canopy["rgb"],
        truth_lower_vigor=canopy["lower_vigor_truth"].astype(np.uint8),
        illumination_confounder=canopy["illumination_confounder"].astype(np.uint8),
        **arrays,
    )
    _save_figures(artifact_root, mission, records, arrays, canopy, analysis)

    receipt_path = demo / "engineering-demo-receipt.json"
    artifacts = [
        demo / "capture-manifest.csv",
        demo / "capture-qc-summary.json",
        demo / "mavic-3m-specification-receipt.json",
        demo / "mission-summary.json",
        demo / "radiometric-processing-receipt.json",
        demo / "synthetic-analysis-arrays.npz",
        demo / "system-analysis.json",
        artifact_root / "outputs" / "figures" / "figure-1-mission-route.png",
        artifact_root / "outputs" / "figures" / "figure-2-system-flow.png",
        artifact_root / "outputs" / "figures" / "figure-3-analysis-panel.png",
        artifact_root / "outputs" / "figures" / "figure-4-qc-coverage.png",
    ]
    receipt = {
        "schemaVersion": 1,
        "reportNumber": "SHWRP-2026-034",
        "runStatus": "PASSED",
        "seed": seed,
        "dataProvenance": "deterministic_synthetic_digital_twin",
        "fieldFlightExecuted": False,
        "aircraftConnected": False,
        "realMavic3MFilesProcessed": False,
        "flightExecutionAuthorized": False,
        "captureEventCount": mission["captureEventCount"],
        "expectedImageFileCount": mission["expectedImageFileCount"],
        "acceptedCaptureEventCount": qc_summary["acceptedEventCount"],
        "rejectedCaptureEventCount": qc_summary["rejectedEventCount"],
        "reviewZoneCount": analysis["reviewZoneCount"],
        "syntheticFunctionalTest": analysis["syntheticFunctionalTest"],
        "python": platform.python_version(),
        "sourceTreeSha256": sha256_tree(source_root / "m3m_twin"),
        "configuration": {
            name: {
                "path": path.relative_to(source_root.resolve()).as_posix()
                if path.is_relative_to(source_root.resolve())
                else str(path),
                "sha256": sha256_file(path),
            }
            for name, path in configuration_paths.items()
        },
        "artifacts": [
            {
                "path": path.relative_to(artifact_root).as_posix(),
                "bytes": path.stat().st_size,
                "sha256": sha256_file(path),
            }
            for path in artifacts
        ],
        "claimBoundary": (
            "Passing means the deterministic software demo completed and its internal checks were met. It is not "
            "evidence of a real flight, real crop condition, disease accuracy, agronomic validity, legal approval, or ROI."
        ),
    }
    write_json(receipt_path, receipt)
    return receipt


def print_demo(
    study_root: str | Path,
    seed: int = 20260804,
    *,
    output_root: str | Path | None = None,
    mission_config: str | Path | None = None,
    field_config: str | Path | None = None,
) -> None:
    print(json.dumps(run_demo(
        study_root,
        seed,
        output_root=output_root,
        mission_config=mission_config,
        field_config=field_config,
    ), ensure_ascii=False, indent=2))
