"""MoGe-2 depth sampling and Depth Balloon base-plane controls."""

from pathlib import Path

import bpy


PLANE_VARIANCE_FLOOR_RATIO = 0.025


def load_depth_balloon_data(path):
    import numpy as np

    try:
        with np.load(Path(path), allow_pickle=False) as archive:
            required = {
                "depth_signal",
                "baseline_vector",
                "valid",
                "origin",
                "reference_depth",
                "crop_size",
                "source_size",
            }
            missing = required.difference(archive.files)
            if missing:
                raise ValueError(
                    f"Depth Balloon data is missing: {', '.join(sorted(missing))}"
                )
            data = {
                "depth_signal": np.asarray(
                    archive["depth_signal"],
                    dtype=np.float32,
                ),
                "baseline_vector": np.asarray(
                    archive["baseline_vector"],
                    dtype=np.float32,
                ),
                "valid": np.asarray(archive["valid"], dtype=bool),
                "origin": tuple(int(value) for value in archive["origin"]),
                "reference_depth": float(archive["reference_depth"]),
                "crop_size": tuple(int(value) for value in archive["crop_size"]),
                "source_size": tuple(int(value) for value in archive["source_size"]),
            }
    except OSError as error:
        raise ValueError(f"Unable to read Depth Balloon data: {path}") from error
    if data["depth_signal"].ndim != 2:
        raise ValueError("Depth Balloon signal must be two-dimensional")
    if data["valid"].shape != data["depth_signal"].shape:
        raise ValueError("Depth Balloon mask dimensions do not match")
    if data["baseline_vector"].shape != (3,):
        raise ValueError("Depth Balloon base plane vector must have three components")
    if any(value <= 0 for value in (*data["crop_size"], *data["source_size"])):
        raise ValueError("Depth Balloon image dimensions must be positive")
    return data


def _sample_valid_bilinear(values, valid, x, y):
    import numpy as np

    height, width = values.shape
    x = np.clip(x, 0.0, max(width - 1, 0))
    y = np.clip(y, 0.0, max(height - 1, 0))
    x0 = np.floor(x).astype(np.int64)
    y0 = np.floor(y).astype(np.int64)
    x1 = np.minimum(x0 + 1, width - 1)
    y1 = np.minimum(y0 + 1, height - 1)
    fx = x - x0
    fy = y - y0
    samples = np.stack(
        (
            values[y0, x0],
            values[y0, x1],
            values[y1, x0],
            values[y1, x1],
        ),
        axis=1,
    )
    weights = np.stack(
        (
            (1.0 - fx) * (1.0 - fy),
            fx * (1.0 - fy),
            (1.0 - fx) * fy,
            fx * fy,
        ),
        axis=1,
    )
    sample_valid = np.stack(
        (
            valid[y0, x0],
            valid[y0, x1],
            valid[y1, x0],
            valid[y1, x1],
        ),
        axis=1,
    )
    weights *= sample_valid
    totals = weights.sum(axis=1)
    result = np.zeros(len(x), dtype=np.float32)
    usable = totals > 1e-8
    result[usable] = (samples[usable] * weights[usable]).sum(axis=1) / totals[usable]
    return result, usable


def _stable_plane(coordinates, values):
    import numpy as np

    coordinates = np.asarray(coordinates, dtype=np.float64)
    values = np.asarray(values, dtype=np.float64)
    if len(coordinates) < 2:
        return np.zeros(3, dtype=np.float32)
    center = coordinates.mean(axis=0)
    value_center = float(values.mean())
    centered = coordinates - center
    covariance = centered.T @ centered / len(centered)
    cross = centered.T @ (values - value_center) / len(centered)
    eigenvalues, eigenvectors = np.linalg.eigh(covariance)
    largest = max(float(eigenvalues[-1]), 0.0)
    if largest <= 1e-12:
        slopes = np.zeros(2, dtype=np.float64)
    else:
        supported = np.maximum(eigenvalues, largest * PLANE_VARIANCE_FLOOR_RATIO)
        slopes = eigenvectors @ ((eigenvectors.T @ cross) / supported)
    intercept = value_center - float(center @ slopes)
    return np.asarray((*slopes, intercept), dtype=np.float32)


def depth_balloon_fields(vertices, vertex_uv, depth_data, depth_scale):
    """Sample camera depth and fit its base plane at a fixed local scale."""
    import numpy as np

    vertices = np.asarray(vertices, dtype=np.float32)
    uv = np.asarray(vertex_uv, dtype=np.float32)
    if vertices.ndim != 2 or vertices.shape[1] != 3:
        raise ValueError("Depth Balloon vertices must have three components")
    if uv.shape != (len(vertices), 2):
        raise ValueError("Depth Balloon UV coordinates do not match its vertices")
    depth_signal = np.asarray(depth_data["depth_signal"], dtype=np.float32)
    valid = np.asarray(depth_data["valid"], dtype=bool)
    crop_width, crop_height = depth_data["crop_size"]
    depth_height, depth_width = depth_signal.shape

    crop_x = uv[:, 0] * max(crop_width - 1, 0)
    crop_y = (1.0 - uv[:, 1]) * max(crop_height - 1, 0)
    sample_x = crop_x * max(depth_width - 1, 0) / max(crop_width - 1, 1)
    sample_y = crop_y * max(depth_height - 1, 0) / max(crop_height - 1, 1)
    sampled, usable = _sample_valid_bilinear(
        depth_signal,
        valid,
        sample_x,
        sample_y,
    )
    inside_crop = (
        (crop_x >= 0.0)
        & (crop_x <= max(crop_width - 1, 0))
        & (crop_y >= 0.0)
        & (crop_y <= max(crop_height - 1, 0))
    )
    usable &= inside_crop
    sampled[~usable] = 0.0
    coordinates = np.stack(
        (
            sample_x / max(depth_width - 1, 1) * 2.0 - 1.0,
            sample_y / max(depth_height - 1, 1) * 2.0 - 1.0,
            np.ones(len(vertices), dtype=np.float32),
        ),
        axis=1,
    ).astype(np.float32)

    depth_scale = float(depth_scale)
    if not depth_scale > 0.0:
        raise ValueError("Depth Balloon reference scale must be positive")
    normalized_plane = np.asarray(
        depth_data["baseline_vector"],
        dtype=np.float32,
    )
    baseline_values = (
        coordinates @ normalized_plane * depth_scale * 0.5
    )
    local_coordinates = vertices[:, :2]
    if np.count_nonzero(usable) >= 3:
        local_plane = _stable_plane(
            local_coordinates[usable],
            baseline_values[usable],
        )
    else:
        local_plane = np.zeros(3, dtype=np.float32)
    depth_values = sampled * depth_scale
    if usable.any():
        plane_without_offset = local_coordinates @ local_plane[:2]
        residual = depth_values - plane_without_offset
        depth_values[usable] -= float(np.median(residual[usable]))
    return depth_values, local_plane


def orient_depth_balloon_plane(local_plane):
    slope_x, slope_y, offset = (float(value) for value in local_plane)
    return -slope_x, slope_y, offset


def create_depth_balloon_base_plane(context, mesh_object, vertices, local_plane):
    import numpy as np
    from mathutils import Vector

    vertices = np.asarray(vertices, dtype=np.float32)
    slope_y, slope_z, offset = (float(value) for value in local_plane)
    extent = vertices[:, 1:].max(axis=0) - vertices[:, 1:].min(axis=0)
    center_y = 0.0
    center_z = 0.0
    center_x = offset
    normal = Vector((-1.0, slope_y, slope_z)).normalized()

    controller = bpy.data.objects.new(f"{mesh_object.name} Base Plane", None)
    context.collection.objects.link(controller)
    controller.parent = mesh_object
    controller.location = (float(center_x), float(center_y), float(center_z))
    controller.rotation_euler = normal.to_track_quat("Z", "Y").to_euler()
    controller.empty_display_type = "SINGLE_ARROW"
    controller.empty_display_size = max(float(extent.max()) * 0.15, 0.01)
    controller.show_in_front = True
    controller.hide_render = True
    controller["anyimage_source_type"] = "DEPTH_BALLOON_BASE_PLANE"
    controller["anyimage_base_plane"] = (slope_y, slope_z, offset)
    return controller