import os
import re

import numpy as np
import torch
from PIL import Image


def _safe_name(name):
    return re.sub(r"[^A-Za-z0-9_.-]+", "_", str(name))


def _prepare_uv(point_uv):
    uv = point_uv.detach().cpu()
    if uv.dim() == 4 and uv.shape[1] == 1:
        uv = uv[:, 0]
    if uv.dim() != 3 or uv.shape[-1] != 2:
        raise ValueError(f"Expected point_uv shape [V, 1, N, 2] or [V, N, 2], got {tuple(point_uv.shape)}")
    return uv


def _prepare_depth(point_z, num_views):
    if point_z is None:
        return None

    depth = point_z.detach().cpu().squeeze()
    if depth.dim() != 2:
        return None
    if depth.shape[0] == num_views:
        return depth
    if depth.shape[1] == num_views:
        return depth.t()
    return None


def save_projection_debug(input_img, point_uv, point_z=None, out_dir="train_imgs/projection_debug", prefix="debug", view_names=None, max_points=30000):
    os.makedirs(out_dir, exist_ok=True)

    images = input_img.detach().clamp(0.0, 1.0).cpu()
    uv = _prepare_uv(point_uv)
    depth = _prepare_depth(point_z, uv.shape[0])

    num_views = min(images.shape[0], uv.shape[0])
    summary = []

    for view_idx in range(num_views):
        image = images[view_idx]
        if image.shape[0] == 1:
            image = image.repeat(3, 1, 1)
        image_np = (image[:3].permute(1, 2, 0).numpy() * 255.0).astype(np.uint8)
        height, width = image_np.shape[:2]

        view_uv = uv[view_idx]
        finite = torch.isfinite(view_uv).all(dim=-1)
        in_frame = (
            finite
            & (view_uv[:, 0] >= -1.0)
            & (view_uv[:, 0] <= 1.0)
            & (view_uv[:, 1] >= -1.0)
            & (view_uv[:, 1] <= 1.0)
        )
        valid_indices = torch.nonzero(in_frame, as_tuple=False).flatten()
        total_points = view_uv.shape[0]
        valid_points = valid_indices.numel()

        if valid_points > max_points:
            sample_ids = torch.linspace(0, valid_points - 1, max_points).long()
            valid_indices = valid_indices[sample_ids]

        sampled_uv = view_uv[valid_indices]
        x = torch.round((sampled_uv[:, 0] + 1.0) * 0.5 * (width - 1)).long().numpy()
        y = torch.round((sampled_uv[:, 1] + 1.0) * 0.5 * (height - 1)).long().numpy()

        overlay = image_np.copy()
        if depth is not None:
            depth_values = depth[view_idx][valid_indices].clamp(0.0, 1.0).numpy()
            colors = np.stack(
                [
                    (255.0 * depth_values).astype(np.uint8),
                    np.full_like(depth_values, 64.0, dtype=np.float32).astype(np.uint8),
                    (255.0 * (1.0 - depth_values)).astype(np.uint8),
                ],
                axis=1,
            )
        else:
            colors = np.tile(np.array([[255, 32, 32]], dtype=np.uint8), (len(x), 1))

        for dy in (-1, 0, 1):
            yy = np.clip(y + dy, 0, height - 1)
            for dx in (-1, 0, 1):
                xx = np.clip(x + dx, 0, width - 1)
                overlay[yy, xx] = colors

        blended = (0.65 * image_np + 0.35 * overlay).astype(np.uint8)
        view_name = view_names[view_idx] if view_names and view_idx < len(view_names) else f"view{view_idx}"
        filename = f"{prefix}_{view_idx:02d}_{_safe_name(view_name)}_projection.png"
        Image.fromarray(blended).save(os.path.join(out_dir, filename))
        summary.append(f"{view_idx} {view_name}: {valid_points}/{total_points} points inside image")

    with open(os.path.join(out_dir, f"{prefix}_projection_summary.txt"), "w") as f:
        f.write("\n".join(summary))