"""
Debug script — inspect mfaytin/mask2former-satellite predictions class by class.

Usage:
    python debug_hardscape.py --address "17531 Madison St, Omaha, NE 68135"
    python debug_hardscape.py --address "..." --imagery google

This fetches the satellite image for the address and shows what each
semantic class looks like spatially, so we can see what the model
is calling 'building', 'pavement', 'road', etc.
"""

from __future__ import annotations

import argparse
import re

import numpy as np
import torch
import requests
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
from PIL import Image
from transformers import AutoImageProcessor, AutoModelForUniversalSegmentation
from dotenv import load_dotenv

load_dotenv()

from lawn_estimator.config import Config, OUTPUT_DIR  # noqa: E402
from lawn_estimator.regions import resolve_region  # noqa: E402
from lawn_estimator.sources.google import GoogleStaticMaps  # noqa: E402
from lawn_estimator.sources.imagery import CountyOrthoImagery, NaipImagery  # noqa: E402

MODEL_ID = "mfaytin/mask2former-satellite"
MODEL_LABELS = {
    0: "background",
    1: "bareland",
    2: "grass",
    3: "pavement",
    4: "road",
    5: "tree",
    6: "water",
    7: "cropland",
    8: "building",
}


def fetch_image(address: str, imagery: str = "naip") -> Image.Image:
    config = Config()
    session = requests.Session()

    region, geocode = resolve_region(address, config, session)
    print(f"  lat={geocode.latitude:.6f}, lon={geocode.longitude:.6f}")

    if imagery == "county":
        client = region.display_imagery or CountyOrthoImagery(session=session)
    elif imagery == "naip":
        client = NaipImagery(session)
    else:
        client = GoogleStaticMaps(config.google_maps_api_key, session)
    return client.fetch_satellite_image(
        latitude=geocode.latitude,
        longitude=geocode.longitude,
        zoom=config.zoom,
        image_size=config.image_size,
        scale=config.image_scale,
    )


def run_model(image: Image.Image) -> np.ndarray:
    print(f"  Loading {MODEL_ID}...")
    processor = AutoImageProcessor.from_pretrained(MODEL_ID)
    model = AutoModelForUniversalSegmentation.from_pretrained(MODEL_ID)
    model.eval()

    inputs = processor(images=image, return_tensors="pt")
    with torch.inference_mode():
        outputs = model(**inputs)

    maps = processor.post_process_semantic_segmentation(
        outputs, target_sizes=[image.size[::-1]]
    )
    return maps[0].detach().cpu().numpy()


def debug(address: str, imagery: str = "naip") -> None:
    print(f"\nDebug hardscape model for: {address} (imagery: {imagery})\n")

    print("[1/3] Fetching satellite image...")
    image = fetch_image(address, imagery)
    print(f"  Image size: {image.size}")

    print("\n[2/3] Running model...")
    prediction = run_model(image)

    unique = np.unique(prediction)
    print("\n  Classes detected:")
    for code in unique:
        label = MODEL_LABELS.get(int(code), f"unknown_{code}")
        count = (prediction == code).sum()
        print(f"    {code}: {label:<15} {count:>8,} px  ({count / prediction.size * 100:.1f}%)")

    print("\n[3/3] Rendering per-class visualization...")
    fig, axes = plt.subplots(3, 3, figsize=(16, 16))
    fig.suptitle(f"mask2former-satellite — per-class predictions\n{address}", fontsize=13)

    for idx, (code, label) in enumerate(MODEL_LABELS.items()):
        ax = axes[idx // 3][idx % 3]
        mask = prediction == code
        ax.imshow(image, alpha=0.6)
        if mask.any():
            ax.imshow(
                np.ma.masked_where(~mask, np.ones_like(mask, dtype=float)),
                cmap="Reds", vmin=0, vmax=1, alpha=0.65,
            )
        ax.set_title(
            f"[{code}] {label}\n{mask.sum():,} px ({mask.sum()/mask.size*100:.1f}%)",
            fontsize=10,
        )
        ax.axis("off")

    plt.tight_layout()

    out_dir = OUTPUT_DIR
    out_dir.mkdir(parents=True, exist_ok=True)
    sanitized = re.sub(r"[^\w\s-]", "", address).strip().replace(" ", "_")[:60]
    out_path = out_dir / f"debug_hardscape_{sanitized}.png"
    plt.savefig(out_path, dpi=120, bbox_inches="tight")
    plt.close(fig)
    print(f"\n  Saved to: {out_path}")


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--address", type=str,
        default="17531 Madison St, Omaha, NE 68135",
    )
    parser.add_argument("--imagery", choices=["naip", "google", "county"], default="naip")
    args = parser.parse_args()
    debug(args.address, args.imagery)