from __future__ import annotations

import numpy as np
import torch


def quat_mul_wxyz(q1: torch.Tensor, q2: torch.Tensor) -> torch.Tensor:
    w1, x1, y1, z1 = q1.unbind(dim=-1)
    w2, x2, y2, z2 = q2.unbind(dim=-1)
    w = w1 * w2 - x1 * x2 - y1 * y2 - z1 * z2
    x = w1 * x2 + x1 * w2 + y1 * z2 - z1 * y2
    y = w1 * y2 - x1 * z2 + y1 * w2 + z1 * x2
    z = w1 * z2 + x1 * y2 - y1 * x2 + z1 * w2
    return torch.stack([w, x, y, z], dim=-1)


def rotmat_to_quat_wxyz(Rm: torch.Tensor) -> torch.Tensor:
    m00, m01, m02 = Rm[0, 0], Rm[0, 1], Rm[0, 2]
    m10, m11, m12 = Rm[1, 0], Rm[1, 1], Rm[1, 2]
    m20, m21, m22 = Rm[2, 0], Rm[2, 1], Rm[2, 2]
    tr = m00 + m11 + m22
    if tr > 0.0:
        s = torch.sqrt(tr + 1.0) * 2.0
        w = 0.25 * s
        x = (m21 - m12) / s
        y = (m02 - m20) / s
        z = (m10 - m01) / s
    elif (m00 > m11) and (m00 > m22):
        s = torch.sqrt(1.0 + m00 - m11 - m22) * 2.0
        w = (m21 - m12) / s
        x = 0.25 * s
        y = (m01 + m10) / s
        z = (m02 + m20) / s
    elif m11 > m22:
        s = torch.sqrt(1.0 + m11 - m00 - m22) * 2.0
        w = (m02 - m20) / s
        x = (m01 + m10) / s
        y = 0.25 * s
        z = (m12 + m21) / s
    else:
        s = torch.sqrt(1.0 + m22 - m00 - m11) * 2.0
        w = (m10 - m01) / s
        x = (m02 + m20) / s
        y = (m12 + m21) / s
        z = 0.25 * s
    q = torch.stack([w, x, y, z])
    return q / q.norm().clamp(min=1e-8)


def to_k4(k3: torch.Tensor) -> torch.Tensor:
    b = k3.shape[0]
    out = torch.eye(4, dtype=k3.dtype, device=k3.device).unsqueeze(0).repeat(b, 1, 1)
    out[:, :3, :3] = k3
    return out


def warmup_cosine_lr(step: int, warmup: int, total: int, lr0: float, lr1: float) -> float:
    if step <= warmup:
        return lr0 * float(step) / float(max(1, warmup))
    t = (step - warmup) / float(max(1, total - warmup))
    cos = 0.5 * (1 + np.cos(np.pi * t))
    return lr1 + (lr0 - lr1) * cos


@torch.no_grad()
def compute_frustum_mask(
    depth: torch.Tensor,
    tgt_w2c: torch.Tensor,
    src_w2c: torch.Tensor,
    src_k3: torch.Tensor,
    tgt_k3: torch.Tensor,
    img_h: int,
    img_w: int,
    source_img_h: int | None = None,
    source_img_w: int | None = None,
    depth_min: float = 0.05,
    margin: float = 0.05,
) -> torch.Tensor:
    dev = depth.device
    f32 = torch.float32
    src_h = int(img_h if source_img_h is None else source_img_h)
    src_w = int(img_w if source_img_w is None else source_img_w)

    d = depth[0, 0].to(f32)
    valid = d > depth_min

    vy, vx = torch.meshgrid(
        torch.arange(img_h, device=dev, dtype=f32),
        torch.arange(img_w, device=dev, dtype=f32),
        indexing="ij",
    )

    fx_t = tgt_k3[0, 0, 0].to(f32)
    fy_t = tgt_k3[0, 1, 1].to(f32)
    cx_t = tgt_k3[0, 0, 2].to(f32)
    cy_t = tgt_k3[0, 1, 2].to(f32)
    X_t = (vx - cx_t) / fx_t * d
    Y_t = (vy - cy_t) / fy_t * d
    Z_t = d
    pts_t = torch.stack([X_t, Y_t, Z_t], dim=-1).reshape(-1, 3)

    c2w_t = torch.linalg.inv(tgt_w2c[0].to(f32))
    pts_w = pts_t @ c2w_t[:3, :3].T + c2w_t[:3, 3][None, :]

    w2c_s = src_w2c[0].to(f32)
    pts_s = pts_w @ w2c_s[:3, :3].T + w2c_s[:3, 3][None, :]

    Z_s = pts_s[:, 2].clamp(min=1e-4)
    fx_s = src_k3[0, 0, 0].to(f32)
    fy_s = src_k3[0, 1, 1].to(f32)
    cx_s = src_k3[0, 0, 2].to(f32)
    cy_s = src_k3[0, 1, 2].to(f32)
    u_s = pts_s[:, 0] / Z_s * fx_s + cx_s
    v_s = pts_s[:, 1] / Z_s * fy_s + cy_s

    half_w = (src_w - 1) * 0.5
    half_h = (src_h - 1) * 0.5
    x_ndc = (u_s - half_w) / half_w
    y_ndc = (v_s - half_h) / half_h

    in_frust = (
        (x_ndc.abs() <= 1.0 + margin)
        & (y_ndc.abs() <= 1.0 + margin)
        & (pts_s[:, 2] > 0)
    )

    mask = in_frust.reshape(img_h, img_w).float()
    mask = mask * valid.float()
    return mask[None, None]