import torch
import torch.nn.functional as F

from typing import Optional

# H36M17 skeleton edges
H36M17_EDGES = [
    (0, 1), (1, 2), (2, 3),             # left leg
    (0, 4), (4, 5), (5, 6),             # right leg
    (0, 7), (7, 8), (8, 9), (9, 10),    # pelvis to head
    (8, 11), (11, 12), (12, 13),        # right arm
    (8, 14), (14, 15), (15, 16),        # left arm
]

# Left-right pairs.
H36M17_LR_PAIRS = [
    (1, 4), (2, 5), (3, 6),
    (14, 11), (15, 12), (16, 13)
]

def return_edges():
    return H36M17_EDGES

def return_lr_edges():
    return H36M17_LR_PAIRS

def split_state(s_hat: torch.Tensor):
    """
        s_hat: (B, T, J, 12) = concat(p, v, a, j), each 3D
    returns
        p: (B, T, J, 3)
        v: (B, T, J, 3)
        a: (B, T, J, 3)
        j: (B, T, J, 3)
    """
    p, v, a, j = torch.split(s_hat, 3, dim=-1)
    return p, v, a, j

# Central difference.
def central_diff(x: torch.Tensor,
                 dt: float):
    """
        Central difference along time axis.
        x: (B, T, J, 3) -> dx/dt: (B, T, J, 3)
        Endpoints use forward/backward difference.
    """
    B, T, J, C = x.shape
    if T < 2:
        return torch.zeros_like(x)

    dx = torch.zeros_like(x)
    
    dx[:, 0]  = (x[:, 1] - x[:, 0]) / dt
    dx[:, -1] = (x[:, -1] - x[:, -2]) / dt

    if T > 2:
        dx[:, 1:-1] = (x[:, 2:] - x[:, :-2]) / (2.0 * dt)

    return dx


def masked_mean(x: torch.Tensor,
                mask: Optional[torch.Tensor] = None,
                eps: float = 1e-8):
    """
        Inputs:
            x: (...,) any shape
        
        Description:
            mask: same broadcastable shape as x without last dims, or exact shape of x expects float/bool with 1 for valid
    """
    if mask is None:
        return x.mean()

    m = mask
    if m.dtype != x.dtype:
        m = m.to(dtype=x.dtype)

    # broadcast
    while m.ndim < x.ndim:
        m = m.unsqueeze(-1)

    num = (x * m).sum()
    den = m.sum().clamp_min(eps)
    return num / den

def build_adj(J: int,
              edges):
    A = torch.zeros(J, J)
    for i, j in edges:
        A[i, j] = 1.0
        A[j, i] = 1.0
    A.fill_diagonal_(1.0)
    return A  # (J,J)

# Loss util
def variance_regularization(z: torch.Tensor, eps: float = 1e-4, target_std: float = 1.0):
    """
        z shape:
            (B,T,J,D) or (B,N,D) or (B,D)

        Treat the last dimension D as the feature dimension.
        Flatten all leading dimensions into samples before computing std.
    """
    if z.ndim < 2:
        raise ValueError(f"z.ndim must be >= 2, got {z.ndim}")

    z = z.reshape(-1, z.shape[-1])   # (N_all, D)
    std = torch.sqrt(z.var(dim=0, unbiased=False) + eps)  # (D,)
    loss_var = torch.mean(F.relu(target_std - std))
    return loss_var