"""CPU ChangeStar demo for two-date imagery with inspectable raster/vector outputs."""

from __future__ import annotations

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

import cv2
import numpy as np
import rasterio
from affine import Affine
from PIL import Image
from rasterio.features import shapes
from rasterio.transform import rowcol
from shapely.geometry import shape as shapely_shape


MODEL_NAME = "s1_s1c1_vitb"
DEFAULT_THRESHOLD = 0.5
DEFAULT_TILE_SIZE = 512
DEFAULT_OVERLAP = 64
DEFAULT_MAX_DIMENSION = 1024
DEFAULT_PROCESSING_MODE = "auto"
AUTO_MAX_DIMENSION = 0
MIN_THRESHOLD = 0.01
MAX_THRESHOLD = 0.99
MIN_MAX_DIMENSION = 512
MAX_MAX_DIMENSION = 4096
PROCESSING_MODES = {"auto", "image", "geotiff"}


def _to_uint8(data: np.ndarray) -> np.ndarray:
    if data.dtype == np.uint8:
        return data
    if np.issubdtype(data.dtype, np.integer):
        info = np.iinfo(data.dtype)
        return np.clip(data.astype(np.float32) / max(1, info.max) * 255.0, 0, 255).astype(np.uint8)
    values = data.astype(np.float32)
    if float(np.nanmax(values)) <= 1.0:
        values = values * 255.0
    return np.clip(values, 0, 255).astype(np.uint8)


def _read_rgb(path: Path, processing_mode: str = DEFAULT_PROCESSING_MODE) -> tuple[np.ndarray, dict[str, Any]]:
    if processing_mode not in PROCESSING_MODES:
        raise ValueError(f"processing_mode must be one of {sorted(PROCESSING_MODES)}")
    if processing_mode in {"auto", "geotiff"} and path.suffix.lower() in {".tif", ".tiff"}:
        try:
            with rasterio.open(path) as dataset:
                has_georeference = dataset.crs is not None and not dataset.transform.is_identity
                if has_georeference:
                    if dataset.count < 3:
                        raise ValueError(f"GeoTIFF must contain at least 3 bands: {path.name}")
                    data = _to_uint8(dataset.read([1, 2, 3]))
                    image = np.transpose(data, (1, 2, 0))
                    return image, {
                        "processing_mode": "geotiff",
                        "georeferenced": True,
                        "crs": dataset.crs.to_string(),
                        "transform": dataset.transform,
                        "pixel_size": [abs(float(dataset.transform.a)), abs(float(dataset.transform.e))],
                        "grid_basis": [float(dataset.transform.a), float(dataset.transform.b), float(dataset.transform.d), float(dataset.transform.e)],
                    }
                if processing_mode == "geotiff":
                    raise ValueError(f"GeoTIFF has no valid CRS and affine transform: {path.name}")
        except rasterio.errors.RasterioIOError as exc:
            if processing_mode == "geotiff":
                raise ValueError(f"Unable to read GeoTIFF {path.name}: {exc}") from exc
    with Image.open(path) as image:
        return np.asarray(image.convert("RGB")), {
            "processing_mode": "image",
            "georeferenced": False,
            "crs": None,
            "transform": None,
            "pixel_size": None,
        }


def _register_after(before: np.ndarray, after: np.ndarray) -> tuple[np.ndarray, np.ndarray, dict[str, Any]]:
    """Register the second image using ORB homography, returning a valid-pixel mask."""
    height, width = before.shape[:2]
    scale = min(1.0, 1200.0 / max(height, width))
    small_before = cv2.resize(before, None, fx=scale, fy=scale, interpolation=cv2.INTER_AREA)
    small_after = cv2.resize(after, None, fx=scale, fy=scale, interpolation=cv2.INTER_AREA)
    gray_before = cv2.cvtColor(small_before, cv2.COLOR_RGB2GRAY)
    gray_after = cv2.cvtColor(small_after, cv2.COLOR_RGB2GRAY)
    orb = cv2.ORB_create(nfeatures=4000, fastThreshold=7)
    key_before, desc_before = orb.detectAndCompute(gray_before, None)
    key_after, desc_after = orb.detectAndCompute(gray_after, None)
    details: dict[str, Any] = {"method": "identity", "matches": 0, "inliers": 0, "inlier_ratio": 0.0}
    matrix = np.eye(3, dtype=np.float32)
    if desc_before is not None and desc_after is not None and len(key_before) >= 8 and len(key_after) >= 8:
        matcher = cv2.BFMatcher(cv2.NORM_HAMMING)
        pairs = matcher.knnMatch(desc_after, desc_before, k=2)
        good = [a for a, b in pairs if a.distance < 0.75 * b.distance]
        if len(good) >= 8:
            src = np.float32([key_after[m.queryIdx].pt for m in good]).reshape(-1, 1, 2)
            dst = np.float32([key_before[m.trainIdx].pt for m in good]).reshape(-1, 1, 2)
            candidate, mask = cv2.findHomography(src, dst, cv2.RANSAC, 5.0)
            if candidate is not None and mask is not None and int(mask.sum()) >= 8:
                matrix = candidate.astype(np.float32)
                inliers = int(mask.sum())
                details = {
                    "method": "orb_homography",
                    "matches": len(good),
                    "inliers": inliers,
                    "inlier_ratio": round(inliers / len(good), 4),
                }
    registered = cv2.warpPerspective(after, matrix, (width, height), flags=cv2.INTER_LINEAR, borderMode=cv2.BORDER_CONSTANT)
    valid = cv2.warpPerspective(np.full((after.shape[0], after.shape[1]), 255, dtype=np.uint8), matrix, (width, height), flags=cv2.INTER_NEAREST, borderMode=cv2.BORDER_CONSTANT)
    details["valid_ratio"] = round(float((valid > 0).mean()), 6)
    details["matrix"] = [[round(float(value), 8) for value in row] for row in matrix]
    return registered, valid > 0, details


def _resize_pair(before: np.ndarray, after: np.ndarray, valid: np.ndarray, max_dimension: int, *, native_resolution: bool = False) -> tuple[np.ndarray, np.ndarray, np.ndarray, dict[str, Any]]:
    height, width = before.shape[:2]
    scale = 1.0 if native_resolution else min(1.0, float(max_dimension) / max(height, width))
    out_width = width if native_resolution else max(32, int(round(width * scale / 32) * 32))
    out_height = height if native_resolution else max(32, int(round(height * scale / 32) * 32))
    # Keep a stable 32-pixel multiple for the model while documenting the actual scale.
    resized_before = cv2.resize(before, (out_width, out_height), interpolation=cv2.INTER_AREA)
    resized_after = cv2.resize(after, (out_width, out_height), interpolation=cv2.INTER_AREA)
    resized_valid = cv2.resize(valid.astype(np.uint8), (out_width, out_height), interpolation=cv2.INTER_NEAREST) > 0
    return resized_before, resized_after, resized_valid, {
        "original_width": width,
        "original_height": height,
        "processed_width": out_width,
        "processed_height": out_height,
        "scale_x": round(width / out_width, 8),
        "scale_y": round(height / out_height, 8),
        "native_resolution": native_resolution,
    }


def _reproject_rgb_to_grid(path: Path, destination_shape: tuple[int, int], destination_transform: Affine, destination_crs: str) -> np.ndarray:
    """Read the second GeoTIFF onto the first image's grid without changing either source file."""
    from rasterio.warp import Resampling, reproject

    height, width = destination_shape
    aligned = np.zeros((3, height, width), dtype=np.uint8)
    with rasterio.open(path) as source:
        for band in range(3):
            source_band = _to_uint8(source.read(band + 1))
            reproject(
                source=source_band,
                destination=aligned[band],
                src_transform=source.transform,
                src_crs=source.crs,
                dst_transform=destination_transform,
                dst_crs=destination_crs,
                resampling=Resampling.bilinear,
            )
    return np.transpose(aligned, (1, 2, 0))


def _write_rgb_geotiff(path: Path, image: np.ndarray, transform: Affine, crs: str | None) -> None:
    with rasterio.open(path, "w", driver="GTiff", height=image.shape[0], width=image.shape[1], count=3, dtype="uint8", transform=transform, crs=crs, compress="lzw") as dst:
        dst.write(np.transpose(image, (2, 0, 1)))


def _write_raster(path: Path, data: np.ndarray, transform: Affine, dtype: str, crs: str | None = None) -> None:
    with rasterio.open(path, "w", driver="GTiff", height=data.shape[0], width=data.shape[1], count=1, dtype=dtype, transform=transform, crs=crs, compress="lzw") as dst:
        dst.write(data.astype(dtype), 1)


def _write_overlay(path: Path, before: np.ndarray, after: np.ndarray, mask: np.ndarray) -> None:
    left = before.copy()
    right = after.copy()
    red = np.zeros_like(right)
    red[..., 0] = 255
    right[mask] = (right[mask].astype(np.float32) * 0.45 + red[mask].astype(np.float32) * 0.55).astype(np.uint8)
    separator = np.full((before.shape[0], 8, 3), 235, dtype=np.uint8)
    Image.fromarray(np.concatenate([left, separator, right], axis=1)).save(path, quality=92)


def _feature_summary(mask: np.ndarray, probability: np.ndarray, transform: Affine, georeferenced: bool) -> list[dict[str, Any]]:
    features: list[dict[str, Any]] = []
    for index, (geometry, value) in enumerate(shapes(mask.astype(np.uint8), transform=transform), start=1):
        if int(value) != 255:
            continue
        polygon = shapely_shape(geometry)
        coords = np.asarray(polygon.exterior.coords)
        rows, cols = rowcol(transform, coords[:, 0], coords[:, 1])
        x = np.clip(np.asarray(cols), 0, mask.shape[1] - 1)
        y = np.clip(np.asarray(rows), 0, mask.shape[0] - 1)
        sample = probability[y, x]
        item: dict[str, Any] = {
            "feature_id": index,
            "area_pixels": round(float(polygon.area / (abs(transform.a * transform.e))), 3),
            "mean_probability": round(float(sample.mean()) if sample.size else 0.0, 6),
            "max_probability": round(float(sample.max()) if sample.size else 0.0, 6),
        }
        if georeferenced:
            item["area_map_units"] = round(float(polygon.area), 6)
            item["bounds_map"] = [round(float(v), 6) for v in polygon.bounds]
        else:
            item["bounds_pixel"] = [round(float(v), 3) for v in polygon.bounds]
        features.append(item)
    return features


def _enrich_vector(vector_path: Path, features: list[dict[str, Any]], preserve_crs: bool = False) -> None:
    import geopandas as gpd

    if vector_path.is_file():
        vector = gpd.read_file(vector_path)
        for key in ("feature_id", "area_pixels", "area_map_units", "mean_probability", "max_probability", "bounds_pixel", "bounds_map"):
            vector[key] = [item.get(key) for item in features[: len(vector)]]
        vector.to_file(vector_path, driver="GeoJSON")
        payload = json.loads(vector_path.read_text(encoding="utf-8"))
        if not preserve_crs:
            payload.pop("crs", None)
        vector_path.write_text(json.dumps(payload, ensure_ascii=False, indent=2), encoding="utf-8")
    else:
        vector_path.write_text(json.dumps({"type": "FeatureCollection", "features": []}, indent=2), encoding="utf-8")


def run_change_detection(before_path: Path, after_path: Path, output_dir: Path, *, processed_dir: Path | None = None, model_name: str = MODEL_NAME, threshold: float = DEFAULT_THRESHOLD, tile_size: int = DEFAULT_TILE_SIZE, overlap: int = DEFAULT_OVERLAP, max_dimension: int = AUTO_MAX_DIMENSION, processing_mode: str = DEFAULT_PROCESSING_MODE) -> dict[str, Any]:
    started = time.perf_counter()
    if not math.isfinite(threshold) or not MIN_THRESHOLD <= threshold <= MAX_THRESHOLD:
        raise ValueError(f"threshold must be between {MIN_THRESHOLD} and {MAX_THRESHOLD}")
    if processing_mode not in PROCESSING_MODES:
        raise ValueError(f"processing_mode must be one of {sorted(PROCESSING_MODES)}")
    if not isinstance(max_dimension, int) or max_dimension != AUTO_MAX_DIMENSION and not MIN_MAX_DIMENSION <= max_dimension <= MAX_MAX_DIMENSION:
        raise ValueError(f"max_dimension must be 0 or between {MIN_MAX_DIMENSION} and {MAX_MAX_DIMENSION}")
    before, before_info = _read_rgb(before_path, processing_mode)
    after, after_info = _read_rgb(after_path, processing_mode)
    actual_mode = "geotiff" if before_info["georeferenced"] and after_info["georeferenced"] else "image"
    if before_info["georeferenced"] != after_info["georeferenced"]:
        raise ValueError("Both inputs must use the same processing mode and georeferencing.")
    if processing_mode == "geotiff" and actual_mode != "geotiff":
        raise ValueError("GeoTIFF mode requires both inputs to have a valid CRS and affine transform.")
    if actual_mode == "geotiff" and before_info["crs"] != after_info["crs"]:
        raise ValueError("GeoTIFF inputs must use the same CRS.")
    grid_alignment: dict[str, Any] = {"required": False, "method": "none"}
    if actual_mode == "geotiff":
        before_transform = before_info["transform"]
        after_transform = after_info["transform"]
        same_grid = before.shape[:2] == after.shape[:2] and np.allclose(
            [before_transform.a, before_transform.b, before_transform.c, before_transform.d, before_transform.e, before_transform.f],
            [after_transform.a, after_transform.b, after_transform.c, after_transform.d, after_transform.e, after_transform.f],
            rtol=1e-6,
            atol=1e-9,
        )
        if not same_grid:
            source_shape = [int(after.shape[0]), int(after.shape[1])]
            after = _reproject_rgb_to_grid(after_path, before.shape[:2], before_transform, before_info["crs"])
            grid_alignment = {
                "required": True,
                "method": "rasterio.reproject_bilinear",
                "source_shape": source_shape,
                "destination_shape": [int(before.shape[0]), int(before.shape[1])],
                "destination_crs": before_info["crs"],
            }
    if before.shape[:2] != after.shape[:2]:
        raise ValueError(f"Input dimensions must match before={before.shape[:2]} after={after.shape[:2]}")
    effective_max_dimension = max_dimension
    if effective_max_dimension == AUTO_MAX_DIMENSION:
        effective_max_dimension = 0 if actual_mode == "geotiff" else DEFAULT_MAX_DIMENSION
    output_dir.mkdir(parents=True, exist_ok=False)
    work_dir = processed_dir or output_dir
    if processed_dir is not None:
        processed_dir.mkdir(parents=True, exist_ok=False)
    registered, valid, registration = _register_after(before, after)
    native_resolution = actual_mode == "geotiff" and effective_max_dimension == AUTO_MAX_DIMENSION
    before_small, after_small, valid_small, resize_details = _resize_pair(before, registered, valid, effective_max_dimension, native_resolution=native_resolution)
    source_transform = before_info["transform"] if actual_mode == "geotiff" else Affine(1, 0, 0, 0, -1, before.shape[0])
    transform = source_transform if native_resolution else source_transform * Affine.scale(resize_details["scale_x"], resize_details["scale_y"])
    output_crs = before_info["crs"] if actual_mode == "geotiff" else None
    _write_rgb_geotiff(work_dir / "before.tif", before_small, transform, output_crs)
    _write_rgb_geotiff(work_dir / "after_registered.tif", after_small, transform, output_crs)

    from geoai import ChangeStarDetection, masks_to_vector

    detector = ChangeStarDetection(model_name=model_name, device="cpu")
    result = detector.predict(str(work_dir / "before.tif"), str(work_dir / "after_registered.tif"), tile_size=tile_size, overlap=overlap, threshold=threshold)
    probability = np.asarray(result["change_prob"], dtype=np.float32)
    raw_mask = ((probability >= threshold) & valid_small).astype(np.uint8) * 255
    # Keep model output intact in raw_mask, then apply a small component filter for presentation.
    cleaned = raw_mask.copy()
    count, labels, stats, _ = cv2.connectedComponentsWithStats((cleaned > 0).astype(np.uint8), connectivity=8)
    min_area = max(16, int(cleaned.size * 0.00002))
    for component in range(1, count):
        if int(stats[component, cv2.CC_STAT_AREA]) < min_area:
            cleaned[labels == component] = 0
    _write_raster(output_dir / "change_probability.tif", probability, transform, "float32", output_crs)
    _write_raster(output_dir / "change_mask_raw.tif", raw_mask, transform, "uint8", output_crs)
    _write_raster(output_dir / "change_mask.tif", cleaned, transform, "uint8", output_crs)
    _write_overlay(output_dir / "change_overlay.jpg", before_small, after_small, cleaned > 0)
    vector_path = output_dir / "changes.geojson"
    vector = masks_to_vector(str(output_dir / "change_mask.tif"), str(vector_path), simplify_tolerance=1.0, mask_threshold=0.5, min_object_area=min_area)
    features = _feature_summary(cleaned, probability, transform, actual_mode == "geotiff")
    _enrich_vector(vector_path, features, preserve_crs=actual_mode == "geotiff")
    (output_dir / "change_features.json").write_text(json.dumps({"features": features}, ensure_ascii=False, indent=2), encoding="utf-8")
    limitations = [
        "GeoTIFF 输入保留 CRS 和地图坐标；若两期网格不同，会在处理副本中将第二期双线性重投影到第一期网格，原始 TIFF 不会被改写。" if actual_mode == "geotiff" else "输入 JPG/PNG 或无 CRS TIFF 没有有效 CRS，GeoTIFF/GeoJSON 坐标是像素换算坐标，不是米或经纬度。",
        "ChangeStar 权重训练于 Changen2/S1 建筑变化数据；当前近景边坡照片不在其验证分布内。",
        "没有人工变化真值，不报告 precision、recall、IoU，也不输出变化类型或工程告警。",
        "ORB 配准只用于工作流演示；生产使用需要正射校正、同 GSD 和独立配准质量验收。",
    ]
    metadata: dict[str, Any] = {
        "schema_version": 1,
        "capability": "00-change-detection",
        "classification": "A",
        "geoai_version": "0.42.0",
        "method": "geoai.ChangeStarDetection + geoai.masks_to_vector",
        "model": model_name,
        "device": "cpu",
        "processing_mode": actual_mode,
        "requested_processing_mode": processing_mode,
        "thresholds": {"change_probability": threshold, "minimum_component_pixels": min_area},
        "tile_size": tile_size,
        "overlap": overlap,
        "max_dimension": max_dimension,
        "effective_max_dimension": effective_max_dimension,
        "input_count": 2,
        "processed_images": 2,
        "input_files": [before_path.name, after_path.name],
        "created_at": datetime.now(UTC).isoformat(),
        "input_shape": [int(before.shape[0]), int(before.shape[1])],
        "processed_shape": [int(before_small.shape[0]), int(before_small.shape[1])],
        "registration": registration,
        "grid_alignment": grid_alignment,
        "resize": resize_details,
        "valid_pixel_ratio": round(float(valid_small.mean()), 6),
        "raw_changed_pixels": int((raw_mask > 0).sum()),
        "changed_pixels": int((cleaned > 0).sum()),
        "changed_pixel_ratio": round(float((cleaned > 0).mean()), 6),
        "vector_feature_count": len(vector),
        "georeferenced": actual_mode == "geotiff",
        "coordinate_basis": "source_crs_map_coordinates" if actual_mode == "geotiff" else "pixel_coordinates_north_up_transform_y_from_image_bottom",
        "crs": output_crs,
        "elapsed_seconds": round(time.perf_counter() - started, 3),
        "python": platform.python_version(),
        "platform": platform.platform(),
        "artifacts": {
            "probability_raster": "change_probability.tif",
            "raw_mask_raster": "change_mask_raw.tif",
            "mask_raster": "change_mask.tif",
            "overlay": "change_overlay.jpg",
            "vector": "changes.geojson",
            "features": "change_features.json",
        },
        "limitations": limitations,
        "licenses": {
            "geoai-py": "MIT",
            "torchange": "Apache-2.0",
            "ever-beta": "PyPI metadata indicates rights reserved; review required",
            "model_weights": "CC BY-NC-SA 4.0 (non-commercial; source EVER-Z/Changen2-ChangeStar1x256)",
            "user_images": "user-provided; authorization not verified",
        },
    }
    (output_dir / "run_metadata.json").write_text(json.dumps(metadata, ensure_ascii=False, indent=2), encoding="utf-8")
    return metadata


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(description="Run CPU ChangeStar change detection on a pair of images.")
    parser.add_argument("--before", type=Path, required=True)
    parser.add_argument("--after", type=Path, required=True)
    parser.add_argument("--output", type=Path, required=True)
    parser.add_argument("--processed-output", type=Path)
    parser.add_argument("--model", default=MODEL_NAME)
    parser.add_argument("--threshold", type=float, default=DEFAULT_THRESHOLD)
    parser.add_argument("--tile-size", type=int, default=DEFAULT_TILE_SIZE)
    parser.add_argument("--overlap", type=int, default=DEFAULT_OVERLAP)
    parser.add_argument("--max-dimension", type=int, default=AUTO_MAX_DIMENSION, help="Long-edge cap in pixels; 0 keeps a valid GeoTIFF at native resolution and uses 1024 for ordinary images.")
    parser.add_argument("--processing-mode", choices=sorted(PROCESSING_MODES), default=DEFAULT_PROCESSING_MODE)
    return parser


if __name__ == "__main__":
    args = build_parser().parse_args()
    try:
        print(json.dumps(run_change_detection(args.before, args.after, args.output, processed_dir=args.processed_output, model_name=args.model, threshold=args.threshold, tile_size=args.tile_size, overlap=args.overlap, max_dimension=args.max_dimension, processing_mode=args.processing_mode), ensure_ascii=False, indent=2))
    except (FileNotFoundError, ValueError, RuntimeError) as exc:
        raise SystemExit(f"变化检测失败: {exc}") from exc