"""Run a small, inspectable CPU semantic-mapping baseline.

The first capability version deliberately uses deterministic RGB/HSV rules so it
can run without downloading model weights.  GeoAI's ``masks_to_vector`` is used
for class-wise polygon extraction, while the metadata keeps the model boundary
explicit: this is a B capability, not a built-in GeoAI segmentation model.
"""

from __future__ import annotations

import argparse
import json
import time
from datetime import UTC, datetime
from pathlib import Path
from typing import Any

import cv2
import geopandas as gpd
import numpy as np
import pandas as pd
import rasterio
from PIL import Image
from rasterio.features import shapes
from rasterio.transform import Affine
from shapely.geometry import shape


CLASS_INFO = {
    0: {"key": "other", "label": "其他", "color": (80, 80, 80)},
    1: {"key": "vegetation", "label": "植被", "color": (40, 180, 70)},
    2: {"key": "water", "label": "水体", "color": (35, 130, 220)},
    3: {"key": "impervious", "label": "不透水面", "color": (220, 150, 55)},
}
SUPPORTED_SUFFIXES = {".jpg", ".jpeg", ".png", ".tif", ".tiff"}


def classify_rgb(rgb: np.ndarray) -> np.ndarray:
    """Classify RGB pixels with conservative, explainable colour rules."""
    rgb_f = rgb.astype(np.float32)
    r, g, b = [rgb_f[..., i] for i in range(3)]
    hsv = cv2.cvtColor(rgb, cv2.COLOR_RGB2HSV)
    saturation = hsv[..., 1].astype(np.float32)
    value = hsv[..., 2].astype(np.float32)
    result = np.zeros(rgb.shape[:2], dtype=np.uint8)

    green = (g > r * 1.04) & (g > b * 1.04) & (g > 45) & (saturation > 45)
    blue = (b > r * 1.08) & (b > g * 1.02) & (b > 45) & (saturation > 35)
    gray = (np.max(rgb_f, axis=2) - np.min(rgb_f, axis=2) < 38) & (value > 45)
    result[green] = 1
    result[blue & ~green] = 2
    result[gray & ~green & ~blue] = 3
    return result


def _raster_profile(input_path: Path, width: int, height: int) -> tuple[Affine, str | None, bool]:
    if input_path.suffix.lower() not in {".tif", ".tiff"}:
        return Affine.identity(), None, False
    try:
        with rasterio.open(input_path) as src:
            transform = src.transform
            crs = src.crs.to_string() if src.crs else None
            return transform, crs, bool(src.crs and src.transform != Affine.identity())
    except rasterio.errors.RasterioIOError:
        return Affine.identity(), None, False


def _geojson_from_mask(mask: np.ndarray, transform: Affine, crs: str | None, class_id: int) -> gpd.GeoDataFrame:
    records: list[dict[str, Any]] = []
    binary = (mask == class_id).astype(np.uint8)
    for geometry, value in shapes(binary, mask=binary.astype(bool), transform=transform):
        if value != 1:
            continue
        polygon = shape(geometry)
        if polygon.area < 100:
            continue
        records.append({"geometry": polygon, "class_id": class_id, "class_key": CLASS_INFO[class_id]["key"], "class_label": CLASS_INFO[class_id]["label"], "pixel_area": float(polygon.area)})
    return gpd.GeoDataFrame(records, geometry="geometry", crs=crs)


def vectorize_with_geoai(mask: np.ndarray, transform: Affine, crs: str | None, output_path: Path) -> bool:
    """Use GeoAI's class-wise mask utility, then preserve class properties."""
    try:
        from geoai import masks_to_vector

        frames: list[gpd.GeoDataFrame] = []
        output_path.parent.mkdir(parents=True, exist_ok=True)
        for class_id in (1, 2, 3):
            temp_path = output_path.with_name(f".{output_path.stem}-{class_id}.tif")
            binary = np.where(mask == class_id, 255, 0).astype(np.uint8)
            with rasterio.open(temp_path, "w", driver="GTiff", height=binary.shape[0], width=binary.shape[1], count=1, dtype="uint8", transform=transform, crs=crs, nodata=0) as dst:
                dst.write(binary, 1)
            try:
                frame = masks_to_vector(str(temp_path), min_object_area=100, simplify_tolerance=1.0)
                if not frame.empty:
                    frame = frame.assign(class_id=class_id, class_key=CLASS_INFO[class_id]["key"], class_label=CLASS_INFO[class_id]["label"])
                    frames.append(frame[["geometry", "class_id", "class_key", "class_label", "confidence"]])
            finally:
                temp_path.unlink(missing_ok=True)
        merged = gpd.GeoDataFrame(pd_concat(frames), geometry="geometry", crs=crs)
        merged.to_file(output_path, driver="GeoJSON")
        return True
    except Exception:
        fallback = gpd.GeoDataFrame(pd_concat([_geojson_from_mask(mask, transform, crs, class_id) for class_id in (1, 2, 3)]), geometry="geometry", crs=crs)
        fallback.to_file(output_path, driver="GeoJSON")
        return False


def pd_concat(frames: list[gpd.GeoDataFrame]) -> gpd.GeoDataFrame:
    if not frames:
        return gpd.GeoDataFrame({"geometry": []}, geometry="geometry")
    return gpd.GeoDataFrame(pd.concat(frames, ignore_index=True), geometry="geometry", crs=frames[0].crs)


def process_image(input_path: Path, output_dir: Path) -> dict[str, Any]:
    started = time.perf_counter()
    with Image.open(input_path) as image:
        rgb = np.asarray(image.convert("RGB"))
    height, width = rgb.shape[:2]
    mask = classify_rgb(rgb)
    transform, crs, georeferenced = _raster_profile(input_path, width, height)
    stem = input_path.stem
    mask_path = output_dir / f"{stem}.mask.png"
    overlay_path = output_dir / f"{stem}.overlay.png"
    raster_path = output_dir / f"{stem}.mask.tif"
    vector_path = output_dir / f"{stem}.segments.geojson"
    Image.fromarray(mask).save(mask_path)
    palette = np.zeros((256, 3), dtype=np.uint8)
    for class_id, info in CLASS_INFO.items():
        palette[class_id] = info["color"]
    color_mask = palette[mask]
    overlay = (rgb.astype(np.float32) * 0.52 + color_mask.astype(np.float32) * 0.48).clip(0, 255).astype(np.uint8)
    Image.fromarray(overlay).save(overlay_path, quality=92)
    with rasterio.open(raster_path, "w", driver="GTiff", height=height, width=width, count=1, dtype="uint8", transform=transform, crs=crs, nodata=0) as dst:
        dst.write(mask, 1)
    used_geoai = vectorize_with_geoai(mask, transform, crs, vector_path)
    counts = {info["key"]: int(np.count_nonzero(mask == class_id)) for class_id, info in CLASS_INFO.items()}
    return {"file": input_path.name, "width": width, "height": height, "mask_file": mask_path.name, "overlay_file": overlay_path.name, "raster_file": raster_path.name, "vector_file": vector_path.name, "class_pixel_counts": counts, "georeferenced": georeferenced, "vectorizer": "geoai.masks_to_vector" if used_geoai else "rasterio.features.shapes fallback", "elapsed_seconds": round(time.perf_counter() - started, 3)}


def collect_inputs(input_path: Path) -> list[Path]:
    if input_path.is_file():
        candidates = [input_path]
    elif input_path.is_dir():
        candidates = sorted(item for item in input_path.iterdir() if item.is_file())
    else:
        raise SystemExit(f"Input path does not exist: {input_path}")
    inputs = [item for item in candidates if item.suffix.lower() in SUPPORTED_SUFFIXES]
    if not inputs:
        raise SystemExit("No JPG, PNG, or GeoTIFF inputs were found.")
    return inputs


def main() -> int:
    parser = argparse.ArgumentParser(description="Run the CPU semantic-mapping baseline.")
    parser.add_argument("--input", type=Path, default=Path("shared/data/raw/01-object-detection"))
    parser.add_argument("--output", type=Path, default=Path("shared/outputs/02-semantic-mapping"))
    args = parser.parse_args()
    inputs = collect_inputs(args.input)
    if args.output.exists() and any(args.output.iterdir()):
        raise SystemExit(f"Output directory is not empty: {args.output}. Use a new run directory.")
    args.output.mkdir(parents=True, exist_ok=True)
    started = time.perf_counter()
    images = [process_image(item, args.output) for item in inputs]
    metadata = {"capability": "02-semantic-mapping", "classification": "B", "created_at": datetime.now(UTC).isoformat(), "geoai_version": __import__("importlib.metadata").metadata.version("geoai-py"), "method": "deterministic RGB/HSV semantic baseline", "model": "classical-color-baseline", "device": "CPU", "thresholds": {"vegetation_green_ratio": 1.04, "water_blue_ratio": 1.08, "gray_channel_range": 38, "minimum_vector_area_pixels": 100}, "input_dir": args.input.as_posix(), "raw_input_dir": args.input.as_posix(), "input_count": len(images), "processed_images": len(images), "elapsed_seconds": round(time.perf_counter() - started, 3), "images": images, "classes": [info | {"id": class_id} for class_id, info in CLASS_INFO.items()], "limitations": ["这是可解释的颜色规则基线，不是经过训练的通用分割模型。", "普通 JPG/PNG 没有 CRS 时，GeoJSON 使用像素坐标；只有有效 GeoTIFF 地理参考才会保留地理坐标。", "复杂光照、阴影、材质相近区域可能误分；当前没有人工标注，因此不报告精度。"]}
    (args.output / "run_metadata.json").write_text(json.dumps(metadata, ensure_ascii=False, indent=2), encoding="utf-8")
    print(json.dumps(metadata, ensure_ascii=False, indent=2))
    return 0


if __name__ == "__main__":
    raise SystemExit(main())