"""Street furniture detection from nDSM and point cloud data.

Detects and classifies small urban objects like lamps, benches, signs, and bollards
by analyzing the normalized DSM (DOM1 - DGM1) and optionally refining with LPO points.
"""

from __future__ import annotations

import csv
import os
from dataclasses import dataclass
from enum import Enum
from typing import List, Optional, Tuple

import numpy as np
from osgeo import gdal, ogr

from .citygml_utils import find_citygml_lod2
from .config import Config
from .gdal_utils import open_dataset
from .pointcloud import (
    PointCloud,
    find_pointcloud_file,
    has_lpo_data,
    read_pointcloud_file,
)


class FurnitureType(Enum):
    """Types of detectable street furniture."""
    LAMP = "lamp"
    BENCH = "bench"
    SIGN = "sign"
    BOLLARD = "bollard"
    UNKNOWN = "unknown"


@dataclass
class FurnitureDetection:
    """A detected street furniture object."""
    tile_id: str
    x_local: float  # X relative to tile origin
    y_local: float  # Y relative to tile origin
    z_ground: float  # Ground elevation
    height: float   # Object height above ground
    width: float    # Approximate object width
    furniture_type: FurnitureType
    confidence: float

    def to_csv_row(self) -> list:
        """Convert to CSV row."""
        return [
            self.tile_id,
            f"{self.x_local:.2f}",
            f"{self.y_local:.2f}",
            f"{self.z_ground:.2f}",
            f"{self.height:.2f}",
            self.furniture_type.value,
            f"{self.confidence:.3f}",
        ]


def classify_furniture(
    height: float,
    width: float,
    cfg: Config,
) -> Tuple[FurnitureType, float]:
    """Classify object type based on height and width.

    Args:
        height: Object height in meters
        width: Object width in meters
        cfg: Configuration with height ranges

    Returns:
        Tuple of (FurnitureType, base_confidence)
    """
    sf_cfg = cfg.street_furniture

    # Check each type's height range
    lamp_min, lamp_max = sf_cfg.lamp_height_range
    bench_min, bench_max = sf_cfg.bench_height_range
    sign_min, sign_max = sf_cfg.sign_height_range
    bollard_min, bollard_max = sf_cfg.bollard_height_range

    # Lamp: tall and narrow
    if lamp_min <= height <= lamp_max and width < 0.5:
        return FurnitureType.LAMP, 0.7

    # Bench: low and wide
    if bench_min <= height <= bench_max and 0.5 <= width <= 2.0:
        return FurnitureType.BENCH, 0.65

    # Sign: medium height, narrow
    if sign_min <= height <= sign_max and width < 0.5:
        return FurnitureType.SIGN, 0.6

    # Bollard: short, very narrow
    if bollard_min <= height <= bollard_max and width < 0.3:
        return FurnitureType.BOLLARD, 0.55

    return FurnitureType.UNKNOWN, 0.3


def _connected_components_2d(
    binary_mask: np.ndarray,
    min_pixels: int = 4,
) -> List[Tuple[np.ndarray, int, int, int, int]]:
    """Find connected components in binary mask.

    Args:
        binary_mask: 2D boolean array
        min_pixels: Minimum component size

    Returns:
        List of (component_mask, row_min, row_max, col_min, col_max)
    """
    from scipy import ndimage

    labeled, num_features = ndimage.label(binary_mask)
    components = []

    for i in range(1, num_features + 1):
        component_mask = labeled == i
        pixel_count = np.sum(component_mask)

        if pixel_count < min_pixels:
            continue

        rows, cols = np.where(component_mask)
        row_min, row_max = rows.min(), rows.max()
        col_min, col_max = cols.min(), cols.max()

        components.append((component_mask, row_min, row_max, col_min, col_max))

    return components


def _rasterize_citygml_footprints(
    tile_id: str,
    cfg: Config,
    shape: Tuple[int, int],
    geotransform: Tuple[float, ...],
) -> np.ndarray:
    """Rasterize CityGML building footprints to create exclusion mask.

    Args:
        tile_id: Tile identifier
        cfg: Configuration
        shape: Output raster shape (rows, cols)
        geotransform: GDAL geotransform

    Returns:
        Boolean mask where True = building footprint
    """
    # Find CityGML file for tile
    gml_path = find_citygml_lod2(tile_id, cfg)
    if not gml_path:
        return np.zeros(shape, dtype=bool)

    # Create memory raster for rasterization
    driver = gdal.GetDriverByName("MEM")
    target_ds = driver.Create("", shape[1], shape[0], 1, gdal.GDT_Byte)
    target_ds.SetGeoTransform(geotransform)

    # Open GML and rasterize
    try:
        gml_ds = ogr.Open(gml_path)
        if gml_ds is None:
            return np.zeros(shape, dtype=bool)

        layer = gml_ds.GetLayer(0)
        gdal.RasterizeLayer(target_ds, [1], layer, burn_values=[1])

        mask = target_ds.GetRasterBand(1).ReadAsArray()
        return mask > 0

    except Exception as e:
        print(f"[street_furniture] Warning: Could not rasterize CityGML for {tile_id}: {e}")
        return np.zeros(shape, dtype=bool)


def detect_furniture_in_tile(
    tile_id: str,
    xmin: float,
    ymin: float,
    xmax: float,
    ymax: float,
    cfg: Config,
) -> List[FurnitureDetection]:
    """Detect street furniture in a single tile.

    Args:
        tile_id: Tile identifier (e.g., '32_328_5511')
        xmin, ymin, xmax, ymax: Tile bounds in meters
        cfg: Configuration

    Returns:
        List of detected furniture objects
    """
    from scipy import ndimage

    sf_cfg = cfg.street_furniture
    pc_cfg = cfg.pointcloud

    # Load DOM1 (surface model)
    dom1_pattern = os.path.join(pc_cfg.dom1_dir, f"*{tile_id}*.tif")
    import glob
    dom1_files = glob.glob(dom1_pattern)
    if not dom1_files:
        print(f"[street_furniture] No DOM1 found for {tile_id}")
        return []

    dom1_ds = open_dataset(dom1_files[0], required=False)
    if dom1_ds is None:
        return []

    # Load DGM1 (terrain model)
    dgm1_pattern = os.path.join(cfg.raw.dgm1_dir, f"*{tile_id}*.tif")
    dgm1_files = glob.glob(dgm1_pattern)
    if not dgm1_files:
        print(f"[street_furniture] No DGM1 found for {tile_id}")
        return []

    dgm1_ds = open_dataset(dgm1_files[0], required=False)
    if dgm1_ds is None:
        return []

    # Read rasters
    dom1 = dom1_ds.GetRasterBand(1).ReadAsArray().astype(np.float32)
    dgm1 = dgm1_ds.GetRasterBand(1).ReadAsArray().astype(np.float32)
    gt = dom1_ds.GetGeoTransform()

    # Handle nodata
    dom1_nodata = dom1_ds.GetRasterBand(1).GetNoDataValue() or -9999
    dgm1_nodata = dgm1_ds.GetRasterBand(1).GetNoDataValue() or -9999
    valid_mask = (dom1 > dom1_nodata + 1) & (dgm1 > dgm1_nodata + 1)

    # Compute nDSM (normalized DSM)
    ndsm = np.where(valid_mask, dom1 - dgm1, 0)
    ndsm = np.clip(ndsm, 0, None)  # Remove negative values

    # Create building exclusion mask
    building_mask = _rasterize_citygml_footprints(
        tile_id, cfg, dom1.shape, gt
    )

    # Create tree exclusion mask (objects > 2m that are likely trees)
    tree_mask = ndsm > 2.0

    # Exclude buildings from tree mask
    tree_mask = tree_mask & ~building_mask

    # Detect small objects (between min and max height, not buildings/trees)
    object_mask = (
        (ndsm >= sf_cfg.min_height_m) &
        (ndsm <= sf_cfg.max_height_m) &
        ~building_mask &
        ~tree_mask
    )

    # Find connected components
    components = _connected_components_2d(object_mask, min_pixels=2)

    detections = []
    pixel_size = abs(gt[1])  # Assume square pixels

    for comp_mask, row_min, row_max, col_min, col_max in components:
        # Compute object properties
        rows, cols = np.where(comp_mask)
        heights = ndsm[rows, cols]

        max_height = float(np.max(heights))
        mean_height = float(np.mean(heights))

        # Compute width from bounding box
        width_pixels = max(col_max - col_min + 1, row_max - row_min + 1)
        width = width_pixels * pixel_size

        # Compute centroid
        centroid_row = (row_min + row_max) / 2
        centroid_col = (col_min + col_max) / 2

        # Convert to world coordinates
        world_x = gt[0] + centroid_col * gt[1]
        world_y = gt[3] + centroid_row * gt[5]

        # Get ground elevation at centroid
        dgm_row = int(centroid_row)
        dgm_col = int(centroid_col)
        if 0 <= dgm_row < dgm1.shape[0] and 0 <= dgm_col < dgm1.shape[1]:
            z_ground = float(dgm1[dgm_row, dgm_col])
        else:
            z_ground = 0.0

        # Convert to local coordinates (relative to tile origin)
        x_local = world_x - xmin
        y_local = world_y - ymin

        # Classify object
        furniture_type, base_confidence = classify_furniture(max_height, width, cfg)

        # Skip unknown types below threshold
        if furniture_type == FurnitureType.UNKNOWN and base_confidence < sf_cfg.min_confidence:
            continue

        detection = FurnitureDetection(
            tile_id=tile_id,
            x_local=x_local,
            y_local=y_local,
            z_ground=z_ground,
            height=max_height,
            width=width,
            furniture_type=furniture_type,
            confidence=base_confidence,
        )
        detections.append(detection)

    # Optionally refine with LPO points
    if sf_cfg.use_lpo_refinement and has_lpo_data(tile_id, pc_cfg.lpo_dir):
        detections = _refine_with_lpo(
            detections,
            tile_id,
            xmin, ymin, xmax, ymax,
            cfg,
        )

    # Filter by minimum confidence
    detections = [d for d in detections if d.confidence >= sf_cfg.min_confidence]

    return detections


def _refine_with_lpo(
    detections: List[FurnitureDetection],
    tile_id: str,
    xmin: float,
    ymin: float,
    xmax: float,
    ymax: float,
    cfg: Config,
) -> List[FurnitureDetection]:
    """Refine furniture detections using LPO point cloud.

    Args:
        detections: Initial detections from nDSM
        tile_id: Tile identifier
        xmin, ymin, xmax, ymax: Tile bounds
        cfg: Configuration

    Returns:
        Refined detections with updated heights and confidence
    """
    lpo_file = find_pointcloud_file(cfg.pointcloud.lpo_dir, tile_id)
    if not lpo_file:
        return detections

    # Load LPO points for tile
    lpo = read_pointcloud_file(lpo_file, bounds=(xmin, ymin, xmax, ymax))
    if len(lpo) == 0:
        return detections

    refined = []
    for det in detections:
        # Find LPO points near detection
        world_x = det.x_local + xmin
        world_y = det.y_local + ymin

        dist_sq = (lpo.x - world_x) ** 2 + (lpo.y - world_y) ** 2
        nearby_mask = dist_sq < (det.width + 1.0) ** 2

        if nearby_mask.any():
            nearby_z = lpo.z[nearby_mask]

            # Refine height using 95th percentile
            refined_height = float(np.percentile(nearby_z, 95)) - det.z_ground
            if refined_height > 0:
                det.height = refined_height

            # Boost confidence for LPO-confirmed detections
            if det.furniture_type == FurnitureType.LAMP:
                det.confidence += 0.2
            elif det.furniture_type == FurnitureType.SIGN:
                det.confidence += 0.15
            else:
                det.confidence += 0.1

            det.confidence = min(det.confidence, 1.0)

        refined.append(det)

    return refined


def export_furniture_csv(
    detections: List[FurnitureDetection],
    tile_id: str,
    cfg: Config,
) -> str:
    """Export furniture detections to CSV file.

    Args:
        detections: List of detections
        tile_id: Tile identifier
        cfg: Configuration

    Returns:
        Path to written CSV file
    """
    os.makedirs(cfg.street_furniture.csv_dir, exist_ok=True)
    csv_path = os.path.join(cfg.street_furniture.csv_dir, f"{tile_id}.csv")

    with open(csv_path, "w", newline="") as f:
        writer = csv.writer(f)
        writer.writerow([
            "tile_id", "x_local", "y_local", "z_ground",
            "height", "type", "confidence"
        ])
        for det in detections:
            writer.writerow(det.to_csv_row())

    return csv_path


def load_furniture_detections(
    tile_id: str,
    cfg: Config,
) -> List[FurnitureDetection]:
    """Load furniture detections from CSV.

    Args:
        tile_id: Tile identifier
        cfg: Configuration

    Returns:
        List of detections, empty if file not found
    """
    csv_path = os.path.join(cfg.street_furniture.csv_dir, f"{tile_id}.csv")
    if not os.path.exists(csv_path):
        return []

    detections = []
    with open(csv_path, "r") as f:
        reader = csv.DictReader(f)
        for row in reader:
            det = FurnitureDetection(
                tile_id=row["tile_id"],
                x_local=float(row["x_local"]),
                y_local=float(row["y_local"]),
                z_ground=float(row["z_ground"]),
                height=float(row["height"]),
                width=0.0,  # Not stored in CSV
                furniture_type=FurnitureType(row["type"]),
                confidence=float(row["confidence"]),
            )
            detections.append(det)

    return detections


def export_street_furniture(cfg: Config) -> int:
    """Export street furniture detections for all tiles.

    Args:
        cfg: Configuration

    Returns:
        Exit code (0 = success)
    """
    if not cfg.street_furniture.enabled:
        print("[street_furniture] Disabled in config, skipping.")
        return 0

    # Load tile manifest
    if not os.path.exists(cfg.export.manifest_path):
        print(f"[street_furniture] ERROR: Manifest not found: {cfg.export.manifest_path}")
        print("[street_furniture] Run heightmap export first.")
        return 1

    tiles = []
    with open(cfg.export.manifest_path, "r") as f:
        reader = csv.DictReader(f)
        for row in reader:
            tiles.append({
                "tile_id": row["tile_id"],
                "xmin": float(row["xmin"]),
                "ymin": float(row["ymin"]),
                "xmax": float(row["xmax"]),
                "ymax": float(row["ymax"]),
            })

    if not tiles:
        print("[street_furniture] No tiles found in manifest.")
        return 1

    os.makedirs(cfg.street_furniture.csv_dir, exist_ok=True)

    total_detections = 0
    for tile in tiles:
        tile_id = tile["tile_id"]
        print(f"[street_furniture] Processing {tile_id}...")

        detections = detect_furniture_in_tile(
            tile_id,
            tile["xmin"],
            tile["ymin"],
            tile["xmax"],
            tile["ymax"],
            cfg,
        )

        csv_path = export_furniture_csv(detections, tile_id, cfg)
        print(f"[street_furniture] {tile_id}: {len(detections)} objects -> {csv_path}")
        total_detections += len(detections)

    print(f"[street_furniture] DONE. Total detections: {total_detections}")
    return 0
