from __future__ import annotations

from dataclasses import asdict

import numpy as np
from scipy.ndimage import gaussian_filter1d

from .structures import FrameHypothesis


def normalize(vector: np.ndarray) -> np.ndarray:
    norm = float(np.linalg.norm(vector))
    if norm == 0:
        return vector.copy()
    return vector / norm


def orthonormalize_axes(x_axis: np.ndarray, z_axis_seed: np.ndarray) -> np.ndarray:
    x_axis = normalize(x_axis)
    z_axis_seed = normalize(z_axis_seed)
    y_axis = normalize(np.cross(z_axis_seed, x_axis))
    z_axis = normalize(np.cross(x_axis, y_axis))
    return np.vstack([x_axis, y_axis, z_axis])


def transform_points(points: np.ndarray, axes: np.ndarray, origin: np.ndarray) -> np.ndarray:
    return (points - origin) @ axes.T


def transform_vectors(vectors: np.ndarray, axes: np.ndarray) -> np.ndarray:
    return vectors @ axes.T


def generate_frame_hypotheses(points: np.ndarray) -> tuple[np.ndarray, list[np.ndarray]]:
    centered = points - points.mean(axis=0, keepdims=True)
    covariance = centered.T @ centered / max(len(points), 1)
    eigenvalues, eigenvectors = np.linalg.eigh(covariance)
    order = np.argsort(eigenvalues)[::-1]
    eigenvectors = eigenvectors[:, order]
    primary = eigenvectors[:, 0]
    residual_axes = [eigenvectors[:, 1], eigenvectors[:, 2]]
    return primary, residual_axes


def _contiguous_ranges(mask: np.ndarray) -> list[tuple[int, int]]:
    ranges: list[tuple[int, int]] = []
    start = None
    for idx, value in enumerate(mask):
        if value and start is None:
            start = idx
        elif not value and start is not None:
            ranges.append((start, idx - 1))
            start = None
    if start is not None:
        ranges.append((start, len(mask) - 1))
    return ranges


def _score_frame_hypothesis(
    coords: np.ndarray,
    normals: np.ndarray,
    curvature: np.ndarray,
    config: dict,
) -> tuple[float, float, float, float, float, float]:
    frame_cfg = config["frame"]
    coarse_slice_count = frame_cfg["coarse_slice_count"]
    x_min, x_max = float(coords[:, 0].min()), float(coords[:, 0].max())
    y_min, y_max = float(coords[:, 1].min()), float(coords[:, 1].max())
    width = max(y_max - y_min, 1e-6)
    slice_edges = np.linspace(x_min, x_max, coarse_slice_count + 1)
    flat_thresh = frame_cfg["flatness_mm"]

    bottom_scores: list[float] = []
    symmetry_scores: list[float] = []
    flatness_scores: list[float] = []
    continuity_flags: list[float] = []
    hole_scores: list[float] = []

    global_curvature_scale = float(np.quantile(curvature, 0.75) + 1e-6)

    for slice_index in range(coarse_slice_count):
        mask = (coords[:, 0] >= slice_edges[slice_index]) & (
            coords[:, 0] < slice_edges[slice_index + 1]
        )
        slice_points = coords[mask]
        if len(slice_points) < frame_cfg["candidate_bin_min_points"] * 6:
            continue
        slice_normals = normals[mask]
        slice_curvature = curvature[mask]
        trim_q = 0.02
        y_lo, y_hi = np.quantile(slice_points[:, 1], [trim_q, 1.0 - trim_q])
        y_mid = 0.5 * (y_lo + y_hi)
        y_edges = np.linspace(y_lo, y_hi, frame_cfg["coarse_y_bins"] + 1)

        low_z = np.full(frame_cfg["coarse_y_bins"], np.nan)
        flatness = np.full(frame_cfg["coarse_y_bins"], np.nan)
        normal_alignment = np.full(frame_cfg["coarse_y_bins"], np.nan)
        counts = np.zeros(frame_cfg["coarse_y_bins"], dtype=np.int32)
        curvature_score = np.zeros(frame_cfg["coarse_y_bins"], dtype=np.float64)

        for bin_index in range(frame_cfg["coarse_y_bins"]):
            bin_mask = (slice_points[:, 1] >= y_edges[bin_index]) & (
                slice_points[:, 1] < y_edges[bin_index + 1]
            )
            counts[bin_index] = int(np.count_nonzero(bin_mask))
            if counts[bin_index] < frame_cfg["candidate_bin_min_points"]:
                continue
            bin_points = slice_points[bin_mask]
            z_threshold = np.quantile(bin_points[:, 2], frame_cfg["bottom_quantile"])
            low_mask = bin_points[:, 2] <= z_threshold
            low_points = bin_points[low_mask]
            low_z[bin_index] = float(np.median(low_points[:, 2]))
            flatness[bin_index] = float(
                np.median(np.abs(low_points[:, 2] - low_z[bin_index]))
            )
            normal_alignment[bin_index] = float(
                np.mean(np.abs(slice_normals[bin_mask][low_mask, 2]))
            )
            curvature_score[bin_index] = float(
                np.mean(slice_curvature[bin_mask][low_mask]) / global_curvature_scale
            )

        valid = ~np.isnan(low_z)
        if valid.sum() < 4:
            continue
        smoothed_low_z = low_z.copy()
        smoothed_low_z[valid] = gaussian_filter1d(low_z[valid], sigma=1.0)
        z_cut = float(
            np.quantile(smoothed_low_z[valid], 0.25) + frame_cfg["low_band_margin_mm"]
        )
        candidate_mask = (
            valid
            & (smoothed_low_z <= z_cut)
            & (flatness <= flat_thresh)
            & (normal_alignment >= frame_cfg["normal_alignment_min"])
        )
        groups = _contiguous_ranges(candidate_mask)
        if not groups:
            continue

        group_summaries = []
        for start, end in groups:
            y_center = 0.5 * (y_edges[start] + y_edges[end + 1])
            group_summaries.append(
                {
                    "start": start,
                    "end": end,
                    "y_center": y_center,
                    "z_low": float(np.nanmedian(smoothed_low_z[start : end + 1])),
                    "flatness": float(np.nanmean(flatness[start : end + 1])),
                    "normal": float(np.nanmean(normal_alignment[start : end + 1])),
                    "count": int(np.sum(counts[start : end + 1])),
                    "curvature": float(np.nanmean(curvature_score[start : end + 1])),
                }
            )

        left = [group for group in group_summaries if group["y_center"] < y_mid]
        right = [group for group in group_summaries if group["y_center"] >= y_mid]
        if not left or not right:
            continuity_flags.append(0.35)
            continue
        left_best = min(left, key=lambda item: item["z_low"])
        right_best = min(right, key=lambda item: item["z_low"])

        slice_bottom_ref = float(np.quantile(slice_points[:, 2], 0.15))
        low_scale = max(np.std(slice_points[:, 2]), 1.0)
        bottom_scores.append(
            np.clip(
                ((slice_bottom_ref - left_best["z_low"]) + (slice_bottom_ref - right_best["z_low"]))
                / (2.0 * low_scale)
                + 1.0,
                0.0,
                2.0,
            )
        )
        sym_error = abs(
            abs(left_best["y_center"] - y_mid) - abs(right_best["y_center"] - y_mid)
        ) / max(width * 0.5, 1.0)
        symmetry_scores.append(float(np.clip(1.0 - sym_error, 0.0, 1.0)))
        flatness_scores.append(
            float(
                np.clip(
                    (
                        (1.0 - left_best["flatness"] / max(flat_thresh, 1e-6))
                        + (1.0 - right_best["flatness"] / max(flat_thresh, 1e-6))
                    )
                    / 2.0
                    * (left_best["normal"] + right_best["normal"])
                    / 2.0,
                    0.0,
                    1.5,
                )
            )
        )
        continuity_flags.append(1.0)
        hole_scores.append(
            float(
                np.clip(
                    (left_best["curvature"] + right_best["curvature"]) / 2.0,
                    0.0,
                    1.5,
                )
            )
        )

    if not continuity_flags:
        return (0.0, 0.0, 0.0, 0.0, 0.0, 0.0)

    bottom_score = float(np.mean(bottom_scores)) if bottom_scores else 0.0
    symmetry_score = float(np.mean(symmetry_scores)) if symmetry_scores else 0.0
    flatness_score = float(np.mean(flatness_scores)) if flatness_scores else 0.0
    hole_score = float(np.mean(hole_scores)) if hole_scores else 0.0
    continuity_score = float(np.mean(continuity_flags))

    total_score = (
        frame_cfg["bottom_weight"] * bottom_score
        + frame_cfg["symmetry_weight"] * symmetry_score
        + frame_cfg["flatness_weight"] * flatness_score
        + frame_cfg["hole_weight"] * hole_score
        + continuity_score
    )
    return (
        total_score,
        bottom_score,
        symmetry_score,
        flatness_score,
        hole_score,
        continuity_score,
    )


def infer_vehicle_frame(
    points: np.ndarray,
    normals: np.ndarray,
    curvature: np.ndarray,
    config: dict,
) -> tuple[list[FrameHypothesis], FrameHypothesis]:
    origin = points.mean(axis=0)
    primary_axis, residual_axes = generate_frame_hypotheses(points)
    hypotheses: list[FrameHypothesis] = []
    for x_sign in (-1.0, 1.0):
        for z_base_index, z_base in enumerate(residual_axes):
            for z_sign in (-1.0, 1.0):
                axes = orthonormalize_axes(x_sign * primary_axis, z_sign * z_base)
                coords = transform_points(points, axes, origin)
                transformed_normals = transform_vectors(normals, axes)
                score_tuple = _score_frame_hypothesis(
                    coords,
                    transformed_normals,
                    curvature,
                    config,
                )
                name = f"x{x_sign:+.0f}_z{z_base_index + 2}{z_sign:+.0f}"
                hypotheses.append(
                    FrameHypothesis(
                        name=name,
                        axes=axes,
                        origin=origin,
                        score=score_tuple[0],
                        bottom_score=score_tuple[1],
                        symmetry_score=score_tuple[2],
                        flatness_score=score_tuple[3],
                        hole_score=score_tuple[4],
                        continuity_score=score_tuple[5],
                    )
                )
    hypotheses.sort(key=lambda item: item.score, reverse=True)
    return hypotheses, hypotheses[0]


def build_bottom_roi_mask(points_vehicle: np.ndarray, config: dict) -> tuple[np.ndarray, dict]:
    roi_cfg = config["roi"]
    aabb = roi_cfg.get("aabb") if isinstance(roi_cfg, dict) else None
    z_floor = float(points_vehicle[:, 2].min())
    if aabb is not None:
        x_lo = float(aabb["x_min"])
        x_hi = float(aabb["x_max"])
        y_lo = float(aabb["y_min"])
        y_hi = float(aabb["y_max"])
        z_lo = float(aabb["z_min"])
        z_ceiling = float(aabb["z_max"])
        mask = (
            (points_vehicle[:, 0] >= x_lo)
            & (points_vehicle[:, 0] <= x_hi)
            & (points_vehicle[:, 1] >= y_lo)
            & (points_vehicle[:, 1] <= y_hi)
            & (points_vehicle[:, 2] >= z_lo)
            & (points_vehicle[:, 2] <= z_ceiling)
        )
        payload = {
            "source": "external_aabb",
            "x_min": x_lo,
            "x_max": x_hi,
            "y_min": y_lo,
            "y_max": y_hi,
            "z_min": z_lo,
            "z_floor": z_floor,
            "z_ceiling": z_ceiling,
            "point_count": int(np.count_nonzero(mask)),
        }
        return mask, payload

    x_lo = float(np.quantile(points_vehicle[:, 0], roi_cfg["x_trim_quantile"]))
    x_hi = float(np.quantile(points_vehicle[:, 0], 1.0 - roi_cfg["x_trim_quantile"]))
    z_quantile = float(np.quantile(points_vehicle[:, 2], roi_cfg["bottom_quantile"]))
    z_span = float(points_vehicle[:, 2].max() - z_floor)
    z_ceiling = min(
        z_quantile + roi_cfg["z_margin_mm"],
        z_floor + roi_cfg["bottom_height_ratio"] * z_span,
    )
    z_ceiling = max(z_ceiling, roi_cfg["min_z_ceiling_mm"])
    mask = (
        (points_vehicle[:, 0] >= x_lo)
        & (points_vehicle[:, 0] <= x_hi)
        & (points_vehicle[:, 2] <= z_ceiling)
    )
    payload = {
        "source": "auto",
        "x_min": x_lo,
        "x_max": x_hi,
        "z_floor": z_floor,
        "z_ceiling": z_ceiling,
        "point_count": int(np.count_nonzero(mask)),
    }
    return mask, payload


def frame_hypotheses_to_dict(hypotheses: list[FrameHypothesis]) -> list[dict]:
    payload = []
    for hypothesis in hypotheses:
        item = asdict(hypothesis)
        item["axes"] = hypothesis.axes.tolist()
        item["origin"] = hypothesis.origin.tolist()
        payload.append(item)
    return payload