from __future__ import annotations

import argparse
import hashlib
import json
import math
import statistics
from datetime import datetime, timezone
from pathlib import Path
from typing import Any


AREA_BUCKETS = (
    ("tiny_lt_25_m2", 0.0, 25.0),
    ("small_25_100_m2", 25.0, 100.0),
    ("medium_100_500_m2", 100.0, 500.0),
    ("large_gte_500_m2", 500.0, math.inf),
)


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Audit persisted GeoIntel false-negative QA evidence across model portfolios."
    )
    parser.add_argument(
        "--portfolio",
        action="append",
        required=True,
        help="Label and calibration_evidence_portfolio.json path as label=/path/file.json.",
    )
    parser.add_argument("--output-dir", required=True, type=Path)
    return parser.parse_args()


def load_json(path: Path) -> dict[str, Any]:
    if not path.is_file():
        raise SystemExit(f"JSON input is not readable: {path}")
    payload = json.loads(path.read_text(encoding="utf-8-sig"))
    if not isinstance(payload, dict):
        raise SystemExit(f"JSON input must be an object: {path}")
    return payload


def parse_portfolio_arg(raw: str) -> tuple[str, Path]:
    label, separator, path_raw = raw.partition("=")
    label = label.strip()
    if not separator or not label or not path_raw.strip():
        raise SystemExit("--portfolio must use label=/path/to/calibration_evidence_portfolio.json")
    return label, Path(path_raw.strip()).expanduser().resolve()


def resolve_evidence_path(portfolio_path: Path, sample: dict[str, Any]) -> Path:
    raw = str(sample.get("evidence_geojson_path") or "")
    configured = Path(raw).expanduser()
    candidates = [configured]
    if raw and not configured.is_absolute():
        candidates.append(portfolio_path.parent / configured)
    sample_slug = str(sample.get("sample_slug") or "").strip().lower()
    candidates.append(
        portfolio_path.parent
        / "samples"
        / sample_slug
        / "evidence"
        / "calibration_evidence.geojson"
    )
    for candidate in candidates:
        if candidate.is_file():
            return candidate.resolve()
    raise SystemExit(f"Evidence GeoJSON is not readable for {sample_slug}: {raw}")


def percentile(values: list[float], fraction: float) -> float | None:
    if not values:
        return None
    ordered = sorted(values)
    position = (len(ordered) - 1) * fraction
    lower = math.floor(position)
    upper = math.ceil(position)
    if lower == upper:
        return ordered[lower]
    return ordered[lower] + (ordered[upper] - ordered[lower]) * (position - lower)


def area_stats(values: list[float]) -> dict[str, float | int | None]:
    return {
        "count": len(values),
        "min": min(values) if values else None,
        "p10": percentile(values, 0.10),
        "median": statistics.median(values) if values else None,
        "p90": percentile(values, 0.90),
        "max": max(values) if values else None,
    }


def area_bucket(area_m2: float) -> str:
    for label, minimum, maximum in AREA_BUCKETS:
        if minimum <= area_m2 < maximum:
            return label
    raise AssertionError(f"No area bucket for {area_m2}")


def stable_reference_id(feature: dict[str, Any], geometry: Any) -> str:
    properties = feature.get("properties") or {}
    source_feature_id = properties.get("source_feature_id")
    if source_feature_id not in (None, ""):
        return f"source:{source_feature_id}"
    reference_feature_id = properties.get("reference_feature_id")
    if reference_feature_id not in (None, ""):
        return f"reference:{reference_feature_id}"
    normalized_wkb = getattr(geometry.normalize(), "wkb", geometry.wkb)
    return f"geometry:{hashlib.sha256(normalized_wkb).hexdigest()}"


def audit_feature_collection(payload: dict[str, Any], geod: Any, shape: Any) -> dict[str, Any]:
    if payload.get("type") != "FeatureCollection":
        raise SystemExit("Evidence GeoJSON must be a FeatureCollection")
    features = payload.get("features") or []
    false_negative_ids: set[str] = set()
    reference_ids: set[str] = set()
    false_negative_areas: list[float] = []
    matched_reference_areas: list[float] = []
    reference_records: dict[str, dict[str, Any]] = {}
    bucket_counts = {
        label: {"false_negative": 0, "matched_reference": 0, "total_reference": 0, "false_negative_rate": None}
        for label, _, _ in AREA_BUCKETS
    }

    for feature in features:
        if not isinstance(feature, dict):
            raise SystemExit("Evidence GeoJSON contains a non-object feature")
        role = str((feature.get("properties") or {}).get("qa_evidence_role") or "")
        if role not in {"false_negative", "match_reference"}:
            continue
        geometry_payload = feature.get("geometry")
        if not isinstance(geometry_payload, dict):
            raise SystemExit(f"Evidence feature {feature.get('id')} has no geometry")
        geometry = shape(geometry_payload)
        if geometry.is_empty or not geometry.is_valid:
            raise SystemExit(f"Evidence feature {feature.get('id')} has invalid geometry")
        if geometry.geom_type not in {"Polygon", "MultiPolygon"}:
            raise SystemExit(
                f"Evidence feature {feature.get('id')} must be Polygon or MultiPolygon"
            )
        area_m2 = abs(float(geod.geometry_area_perimeter(geometry)[0]))
        bucket = area_bucket(area_m2)
        reference_id = stable_reference_id(feature, geometry)
        reference_ids.add(reference_id)
        reference_records[reference_id] = {
            "area_m2": area_m2,
            "area_bucket": bucket,
            "feature": feature,
        }
        bucket_role = "false_negative" if role == "false_negative" else "matched_reference"
        bucket_counts[bucket][bucket_role] += 1
        bucket_counts[bucket]["total_reference"] += 1
        if role == "false_negative":
            false_negative_ids.add(reference_id)
            false_negative_areas.append(area_m2)
        else:
            matched_reference_areas.append(area_m2)

    for values in bucket_counts.values():
        total = values["total_reference"]
        values["false_negative_rate"] = values["false_negative"] / total if total else None

    total_reference = len(false_negative_areas) + len(matched_reference_areas)
    return {
        "false_negative_ids": false_negative_ids,
        "reference_ids": reference_ids,
        "reference_records": reference_records,
        "false_negative_count": len(false_negative_areas),
        "matched_reference_count": len(matched_reference_areas),
        "total_reference_count": total_reference,
        "false_negative_rate": len(false_negative_areas) / total_reference if total_reference else None,
        "false_negative_area_m2": area_stats(false_negative_areas),
        "matched_reference_area_m2": area_stats(matched_reference_areas),
        "area_buckets": bucket_counts,
    }


def build_recommendations(samples: list[dict[str, Any]]) -> list[str]:
    recommendations: list[str] = []
    persistent_total = sum(sample["persistent_false_negative_count"] for sample in samples)
    if persistent_total:
        recommendations.append(
            f"Prioritize targeted positive sampling for {persistent_total} reference buildings missed by every compared portfolio."
        )
    ranked = sorted(
        samples,
        key=lambda sample: max(
            (item.get("false_negative_rate") or 0.0 for item in sample["portfolios"]),
            default=0.0,
        ),
        reverse=True,
    )
    if ranked:
        recommendations.append(
            "Expand or oversample the weakest AOIs first: "
            + ", ".join(sample["sample_slug"] for sample in ranked[:3])
            + "."
        )
    small_bias_samples: list[str] = []
    for sample in samples:
        for portfolio in sample["portfolios"]:
            buckets = portfolio["area_buckets"]
            small_total = sum(
                buckets[key]["total_reference"]
                for key in ("tiny_lt_25_m2", "small_25_100_m2")
            )
            small_fn = sum(
                buckets[key]["false_negative"]
                for key in ("tiny_lt_25_m2", "small_25_100_m2")
            )
            larger_total = sum(
                buckets[key]["total_reference"]
                for key in ("medium_100_500_m2", "large_gte_500_m2")
            )
            larger_fn = sum(
                buckets[key]["false_negative"]
                for key in ("medium_100_500_m2", "large_gte_500_m2")
            )
            small_rate = small_fn / small_total if small_total else 0.0
            larger_rate = larger_fn / larger_total if larger_total else 0.0
            if small_total >= 5 and small_rate >= larger_rate + 0.15:
                small_bias_samples.append(sample["sample_slug"])
                break
    if small_bias_samples:
        recommendations.append(
            "Review small-building label retention and pixel-size gates for: "
            + ", ".join(sorted(set(small_bias_samples)))
            + "."
        )
    if not recommendations:
        recommendations.append(
            "No dominant false-negative pattern was found; preserve current data gates and collect more independent AOIs."
        )
    return recommendations


def run_audit(portfolio_args: list[str], output_dir: Path) -> tuple[Path, Path]:
    try:
        from pyproj import Geod
        from shapely.geometry import shape
    except ImportError as exc:
        raise SystemExit(
            "False-negative GIS audit requires the existing GeoIntel GIS extras (pyproj and shapely)."
        ) from exc

    parsed_portfolios = [parse_portfolio_arg(raw) for raw in portfolio_args]
    labels = [label for label, _ in parsed_portfolios]
    if len(labels) != len(set(labels)):
        raise SystemExit("Every --portfolio label must be unique")
    geod = Geod(ellps="WGS84")
    portfolio_samples: dict[str, dict[str, dict[str, Any]]] = {}
    portfolio_meta: list[dict[str, Any]] = []

    for label, portfolio_path in parsed_portfolios:
        portfolio = load_json(portfolio_path)
        samples = portfolio.get("samples") or []
        if not isinstance(samples, list) or not samples:
            raise SystemExit(f"Evidence portfolio has no samples: {portfolio_path}")
        sample_results: dict[str, dict[str, Any]] = {}
        for sample in samples:
            sample_slug = str(sample.get("sample_slug") or "").strip().lower()
            if not sample_slug:
                raise SystemExit(f"Evidence portfolio has a sample without sample_slug: {portfolio_path}")
            evidence_path = resolve_evidence_path(portfolio_path, sample)
            result = audit_feature_collection(load_json(evidence_path), geod, shape)
            result.update(
                {
                    "label": label,
                    "model_asset_id": portfolio.get("model_asset_id"),
                    "evidence_geojson_path": str(evidence_path),
                }
            )
            sample_results[sample_slug] = result
        portfolio_samples[label] = sample_results
        portfolio_meta.append(
            {
                "label": label,
                "path": str(portfolio_path),
                "model_asset_id": portfolio.get("model_asset_id"),
                "sample_count": len(sample_results),
            }
        )

    expected_slugs = set(next(iter(portfolio_samples.values())))
    for label, samples in portfolio_samples.items():
        if set(samples) != expected_slugs:
            raise SystemExit(
                f"Portfolio {label} has different AOIs; fixed-threshold comparisons require identical samples"
            )

    sample_reports: list[dict[str, Any]] = []
    persistent_evidence_features: list[dict[str, Any]] = []
    for sample_slug in sorted(expected_slugs):
        portfolio_rows = []
        false_negative_sets = []
        reference_sets = []
        for label, _ in parsed_portfolios:
            raw = portfolio_samples[label][sample_slug]
            false_negative_sets.append(raw["false_negative_ids"])
            reference_sets.append(raw["reference_ids"])
            portfolio_rows.append(
                {
                    key: value
                    for key, value in raw.items()
                    if key not in {"false_negative_ids", "reference_ids", "reference_records"}
                }
            )
        if any(reference_ids != reference_sets[0] for reference_ids in reference_sets[1:]):
            counts = ", ".join(
                f"{label}={len(reference_ids)}"
                for (label, _), reference_ids in zip(parsed_portfolios, reference_sets, strict=True)
            )
            raise SystemExit(
                f"Sample {sample_slug} has different reference populations across portfolios ({counts})"
            )
        persistent_ids = sorted(set.intersection(*false_negative_sets))
        reference_records = portfolio_samples[parsed_portfolios[0][0]][sample_slug]["reference_records"]
        persistent_areas = [float(reference_records[reference_id]["area_m2"]) for reference_id in persistent_ids]
        persistent_bucket_counts = {
            label: {"count": 0, "share": 0.0}
            for label, _, _ in AREA_BUCKETS
        }
        for reference_id in persistent_ids:
            record = reference_records[reference_id]
            persistent_bucket_counts[record["area_bucket"]]["count"] += 1
            source_feature = record["feature"]
            properties = dict(source_feature.get("properties") or {})
            properties.update(
                {
                    "qa_evidence_role": "persistent_false_negative",
                    "sample_slug": sample_slug,
                    "persistent_reference_id": reference_id,
                    "area_m2": record["area_m2"],
                    "area_bucket": record["area_bucket"],
                    "compared_portfolios": labels,
                }
            )
            persistent_evidence_features.append(
                {
                    "type": "Feature",
                    "id": f"{sample_slug}:{reference_id}",
                    "properties": properties,
                    "geometry": source_feature["geometry"],
                }
            )
        if persistent_ids:
            for values in persistent_bucket_counts.values():
                values["share"] = values["count"] / len(persistent_ids)
        sample_reports.append(
            {
                "sample_slug": sample_slug,
                "reference_population_count": len(reference_sets[0]),
                "persistent_false_negative_count": len(persistent_ids),
                "persistent_reference_ids": persistent_ids,
                "persistent_false_negative_area_m2": area_stats(persistent_areas),
                "persistent_area_buckets": persistent_bucket_counts,
                "portfolios": portfolio_rows,
            }
        )

    output_dir = output_dir.expanduser().resolve()
    output_dir.mkdir(parents=True, exist_ok=True)
    persistent_geojson_path = output_dir / "persistent_false_negatives.geojson"
    persistent_geojson_path.write_text(
        json.dumps(
            {"type": "FeatureCollection", "features": persistent_evidence_features},
            indent=2,
            sort_keys=True,
        ),
        encoding="utf-8",
    )
    report = {
        "generated_at": datetime.now(timezone.utc).isoformat(),
        "schema_version": 1,
        "input_crs": "EPSG:4326",
        "area_method": "WGS84 geodesic area via pyproj.Geod",
        "portfolio_count": len(portfolio_meta),
        "portfolios": portfolio_meta,
        "sample_count": len(sample_reports),
        "samples": sample_reports,
        "persistent_evidence_geojson_path": str(persistent_geojson_path),
        "recommendations": build_recommendations(sample_reports),
    }
    json_path = output_dir / "detection_false_negative_audit.json"
    json_path.write_text(json.dumps(report, indent=2, sort_keys=True), encoding="utf-8")

    lines = [
        "# Detection false-negative evidence audit",
        "",
        f"- Generated: {report['generated_at']}",
        f"- Portfolios: {report['portfolio_count']}",
        f"- AOIs: {report['sample_count']}",
        f"- Area method: {report['area_method']}",
        "",
        "## AOI comparison",
        "",
        "| AOI | Persistent misses | Persistent tiny/small | Persistent median m2 | "
        + " | ".join(f"{label} FN rate" for label, _ in parsed_portfolios)
        + " |",
        "|---|---:|---:|---:|" + "---:|" * len(parsed_portfolios),
    ]
    for sample in sample_reports:
        rates = [
            f"{(row['false_negative_rate'] or 0.0):.3f}"
            for row in sample["portfolios"]
        ]
        persistent_buckets = sample["persistent_area_buckets"]
        persistent_small_count = sum(
            persistent_buckets[key]["count"]
            for key in ("tiny_lt_25_m2", "small_25_100_m2")
        )
        persistent_median = sample["persistent_false_negative_area_m2"]["median"]
        persistent_median_text = f"{persistent_median:.1f}" if persistent_median is not None else "n/a"
        lines.append(
            f"| {sample['sample_slug']} | {sample['persistent_false_negative_count']} | "
            f"{persistent_small_count} | {persistent_median_text} | "
            + " | ".join(rates)
            + " |"
        )
    lines.extend(["", "## Recommended data actions", ""])
    lines.extend(f"- {item}" for item in report["recommendations"])
    markdown_path = output_dir / "detection_false_negative_audit.md"
    markdown_path.write_text("\n".join(lines) + "\n", encoding="utf-8")
    return json_path, markdown_path


def main() -> None:
    args = parse_args()
    json_path, markdown_path = run_audit(args.portfolio, args.output_dir)
    print(f"False-negative audit JSON: {json_path}")
    print(f"False-negative audit Markdown: {markdown_path}")


if __name__ == "__main__":
    main()