import numpy as np, torch
import json
import torch.nn.functional as F

from .utils import fastKNN_batched

def getLossFnc(loss_name, **kwargs):
    if loss_name == "huber": return torch.nn.HuberLoss(**kwargs)
    elif loss_name == "l1": return torch.nn.L1Loss(**kwargs)
    elif loss_name == "l2": return torch.nn.MSELoss(**kwargs)
    elif loss_name == "cross_entropy": return torch.nn.CrossEntropyLoss(**kwargs) # logits vs target (labels)
    else: raise NotImplementedError

def robust_segment_triangle_loss(p1, p2, v0, v1, v2, tolerance=2e-3):
    """
    tolerance: float (meters). Penetrations smaller than this are ignored.
               Default 2e-3 = 2mm.
    """

    edge1 = v1 - v0
    edge2 = v2 - v0

    # 1. Geometry
    with torch.no_grad():
        ray = p2 - p1

        # 2. Möller–Trumbore
        p_vec = torch.cross(ray, edge2, dim=-1)
        det = (edge1 * p_vec).sum(dim=-1, keepdim=True)
        not_parallel = (det.abs() > 1e-8).float()
        
        det_sign = torch.sign(det)
        det_sign[det_sign == 0] = 1.0 
        det_safe = torch.where(det.abs() < 1e-8, det_sign * 1e-8, det)
        inv_det = 1.0 / det_safe
        
        t_vec = p1 - v0
        u = (t_vec * p_vec).sum(dim=-1, keepdim=True) * inv_det
        q_vec = torch.cross(t_vec, edge1, dim=-1)
        v = (ray * q_vec).sum(dim=-1, keepdim=True) * inv_det
        t = (edge2 * q_vec).sum(dim=-1, keepdim=True) * inv_det

        # 3. Soft Intersection Probability
        k = 100.0 
        def sigmoid(x): return torch.sigmoid(k * x)

        is_intersecting_geo = (
            sigmoid(u) * sigmoid(1.0 - u) * sigmoid(v) * sigmoid(1.0 - v) * sigmoid(1.0 - (u + v)) *
            sigmoid(t) * sigmoid(1.0 - t)
        ) * not_parallel

    normal = torch.cross(edge1, edge2, dim=-1)
    normal_len = torch.linalg.norm(normal, dim=-1, keepdim=True) + 1e-8
    unit_normal = normal / normal_len

    # --- 4. Penalties with Tolerance (The Fix) ---
    
    # Signed distance to plane
    d1 = ((p1 - v0) * unit_normal).sum(dim=-1, keepdim=True)
    d2 = ((p2 - v0) * unit_normal).sum(dim=-1, keepdim=True)
    
    # Calculate raw depth (positive value = how deep inside we are)
    # If d1 is negative (inside), -d1 is positive depth.
    depth1 = -d1
    depth2 = -d2
    
    # Apply Tolerance:
    # If depth is 1mm (0.001) and tolerance is 2mm (0.002):
    #   0.001 - 0.002 = -0.001 -> ReLU -> 0.0 (Ignored)
    # If depth is 5mm (0.005):
    #   0.005 - 0.002 = 0.003 -> ReLU -> 0.003 (Penalized)
    pen_p1 = F.relu(depth1 - tolerance)
    pen_p2 = F.relu(depth2 - tolerance)
    
    loss = is_intersecting_geo * (pen_p1 + pen_p2)
    
    # 5. Updated Visualization Logic
    # We define "penetrating" only if it exceeds the tolerance too, 
    # to match what the loss sees.
    is_penetrating = is_intersecting_geo * ( (pen_p1 > 0).float() + (pen_p2 > 0).float() ).clamp(max=1.0)
    
    return is_penetrating.squeeze(), loss.squeeze()


class SegmentTriangleMeshIntersectionLossDualHands(torch.nn.Module):
    """
    Instead of SegmentTriangleMeshIntersectionLoss, which tests intersection in general case by removing adjacent faces, here are
    tested distinct subsets of r hand
    """
    """
    This implementation is more accurate than simple point-to-plane distance.
    Although point-to-plane is simpler and works well, it will fail in some scenarios.
    For example, imagine two faces intersecting but their closest points are still outside of each other, outside the mesh interior.
    """
    def __init__(self, faces, num_triangles_test=5, target_remove_unlikely_parts=True, far_centroid_thr=0.02, is_intersect_prob = 0.5, intersection_tolerance=3e-3, segmentation_path=None, smplx_model_path=None):
        super().__init__()
        # self.register_buffer('faces', faces)
        self.intersection_tolerance = intersection_tolerance
        self.far_centroid_thr = far_centroid_thr # in meters
        self.is_intersect_prob = is_intersect_prob

        if num_triangles_test < 5:
            print('WARNING: not too high to keep efficiency! Current value:', num_triangles_test)

        if far_centroid_thr > 0.03:
            print('WARNING: far_centroid_thr should be at most 2.5cm for efficiency! Current value:', far_centroid_thr)

        if segmentation_path is None:
            raise ValueError("segmentation_path must be provided for SegmentTriangleMeshIntersectionLossDualHands")
        segmentation_idxs = json.load(open(segmentation_path, 'r'))
        print('SegmentTriangleMeshIntersectionLoss: number of points', faces.max().item()+1)
        hand_verts_l = np.concatenate([segmentation_idxs[k] for k in ['leftHand', 'leftHandIndex1']])
        hand_verts_r = np.concatenate([segmentation_idxs[k] for k in ['rightHand', 'rightHandIndex1']])
        hand_mask_l = torch.zeros((faces.max().item()+1,), dtype=bool)
        hand_mask_r = torch.zeros((faces.max().item()+1,), dtype=bool)
        hand_mask_l[hand_verts_l] = True
        hand_mask_r[hand_verts_r] = True
        is_hand_face_l = hand_mask_l[faces].any(dim=-1)
        is_hand_face_r = hand_mask_r[faces].any(dim=-1)
        self.register_buffer('query_faces_l', faces[is_hand_face_l])
        self.register_buffer('query_faces_r', faces[is_hand_face_r])

        self.target_remove_unlikely_parts = target_remove_unlikely_parts
        self.likely_mask = None
        if target_remove_unlikely_parts:
            unlikely_verts_l = np.concatenate([segmentation_idxs[k] for k in ['leftLeg', 'leftToeBase', 'leftFoot', 'leftUpLeg', 'eyeballs', 'rightToeBase', 'rightUpLeg', 'rightLeg', 'rightFoot', 'leftHand', 'leftHandIndex1', 'leftArm', 'leftForeArm']])
            unlikely_verts_r = np.concatenate([segmentation_idxs[k] for k in ['leftLeg', 'leftToeBase', 'leftFoot', 'leftUpLeg', 'eyeballs', 'rightToeBase', 'rightUpLeg', 'rightLeg', 'rightFoot', 'rightHand', 'rightHandIndex1', 'rightArm', 'rightForeArm']])

            likely_verts_l = np.concatenate([segmentation_idxs[k] for k in ['rightHand', 'rightHandIndex1', 'rightArm', 'rightForeArm']])
            likely_verts_r = np.concatenate([segmentation_idxs[k] for k in ['leftHand', 'leftHandIndex1', 'leftArm', 'leftForeArm']])
            import smplx
            if smplx_model_path is None:
                raise ValueError("smplx_model_path must be provided when target_remove_unlikely_parts=True")
            smplx_model_test = smplx.create(**{'model_path': smplx_model_path, 'model_type': 'smplx'}).to('cuda')
            
            self.likely_mask_l = torch.ones((faces.max().item()+1,), dtype=bool)
            self.likely_mask_l[unlikely_verts_l] = False
            self.likely_mask_l[smplx_model_test.v_template[...,2] < 0.0] = False # positive
            self.likely_mask_l[likely_verts_l] = True

            self.likely_mask_r = torch.ones((faces.max().item()+1,), dtype=bool)
            self.likely_mask_r[unlikely_verts_r] = False
            self.likely_mask_r[smplx_model_test.v_template[...,2] < 0.0] = False # positive
            self.likely_mask_r[likely_verts_r] = True

            is_likely_face_l = self.likely_mask_l[faces].all(dim=-1)
            target_faces_l = faces[is_likely_face_l]

            is_likely_face_r = self.likely_mask_r[faces].all(dim=-1)
            target_faces_r = faces[is_likely_face_r]
        else:
            raise NotImplementedError()

        self.register_buffer('target_faces_r', target_faces_r)
        self.register_buffer('target_faces_l', target_faces_l)

        # 2. Setup Query Edges
        q_f_l = self.query_faces_l
        q_f_r = self.query_faces_r
        edges_stacked_l = torch.stack([torch.stack([q_f_l[:,0], q_f_l[:,1]], dim=1), torch.stack([q_f_l[:,1], q_f_l[:,2]], dim=1), torch.stack([q_f_l[:,2], q_f_l[:,0]], dim=1)], dim=1)
        edges_stacked_r = torch.stack([torch.stack([q_f_r[:,0], q_f_r[:,1]], dim=1), torch.stack([q_f_r[:,1], q_f_r[:,2]], dim=1), torch.stack([q_f_r[:,2], q_f_r[:,0]], dim=1)], dim=1)

        # Flatten to (F_q * 3, 2) -> Keeps order [F0_e0, F0_e1, F0_e2, F1_e0...]
        self.register_buffer('query_edges_l', edges_stacked_l.view(-1, 2))
        self.register_buffer('query_edges_r', edges_stacked_r.view(-1, 2))

        self.num_test = num_triangles_test # should be high enough to cover for ignored adjacent faces

    def forwardQueryTarget(self, vertices, query_faces, target_faces, query_edges, vis=False):
        with torch.no_grad():
            centroids_t = vertices[:, target_faces].mean(dim=2) # B x F_t x 3
            centroids_q = vertices[:, query_faces].mean(dim=2) # B x F_q x 3

            closest_idxs, closest_dists = fastKNN_batched(centroids_q, centroids_t, k=self.num_test+1, exclude_self=True)

        valid_mask = closest_dists < self.far_centroid_thr # (B, F_q, K_test)

        if not valid_mask.any():
            return torch.tensor(0.0, device=vertices.device, requires_grad=True), 0
        
        valid_indices = torch.nonzero(valid_mask)
        b_idxs = valid_indices[:, 0]
        f_q_idxs = valid_indices[:, 1]
        k_idxs = valid_indices[:, 2]

        valid_closest_idxs = closest_idxs[b_idxs, f_q_idxs, k_idxs] # (N_valid,)

        # Look up the actual vertex indices for these target faces
        # target_faces: (F_t, 3) -> valid_target_faces: (N_valid, 3)
        valid_target_faces = target_faces[valid_closest_idxs]
        
        # Gather V0, V1, V2 from vertices
        # We assume vertices is (B, N, 3). We need advanced indexing.
        t_v0 = vertices[b_idxs, valid_target_faces[:, 0]] # (N_valid, 3)
        t_v1 = vertices[b_idxs, valid_target_faces[:, 1]]
        t_v2 = vertices[b_idxs, valid_target_faces[:, 2]]

        # 3. Gather Query Geometry (Only for valid pairs)
        # Each query face has 3 edges. We need to expand our valid list by 3x.
        
        # Repeat Target geometry 3 times (interleaved) to match the 3 edges
        # (N_valid, 3) -> (N_valid * 3, 3)
        v0 = t_v0.repeat_interleave(3, dim=0)
        v1 = t_v1.repeat_interleave(3, dim=0)
        v2 = t_v2.repeat_interleave(3, dim=0)
        
        # Get Edge Indices
        # We need edges corresponding to f_q_idxs.
        # Edges for face i are at: 3*i, 3*i+1, 3*i+2
        base_edge_idxs = f_q_idxs * 3
        # We stack them: [3i, 3i+1, 3i+2]
        # Shape (N_valid, 3) -> Flatten to (N_valid * 3)
        edge_idxs = torch.stack([base_edge_idxs, base_edge_idxs+1, base_edge_idxs+2], dim=1).view(-1)
        
        # Look up the vertex indices for these edges
        # query_edges: (F_q*3, 2)
        valid_edge_verts = query_edges[edge_idxs] # (N_valid*3, 2)
        
        # We also need to repeat the batch indices 3 times to match the edges
        b_idxs_expanded = b_idxs.repeat_interleave(3)
        
        # Gather P1, P2
        p1 = vertices[b_idxs_expanded, valid_edge_verts[:, 0]]
        p2 = vertices[b_idxs_expanded, valid_edge_verts[:, 1]]

        is_intersect, final_loss = robust_segment_triangle_loss(p1, p2, v0, v1, v2, tolerance=self.intersection_tolerance)

        if vis:
        # if True:
            num_show = 10
            pen_idxs = torch.where(is_intersect)[0]
            sorted_vals, sorted_indices = torch.sort(final_loss[pen_idxs], descending=True)
            pen_idxs = pen_idxs[sorted_indices[:num_show]]
                
            from drawer import plotPoses
            
            for idx in pen_idxs:
                b_idx = b_idxs_expanded[idx]
                edge_idx1 = valid_edge_verts[idx, 0].item()
                edge_idx2 = valid_edge_verts[idx, 1].item()
                face_idx1 = valid_target_faces[idx//3, 0].item()
                face_idx2 = valid_target_faces[idx//3, 1].item()
                face_idx3 = valid_target_faces[idx//3, 2].item()
                segment = torch.stack([p1[idx], p2[idx]], dim=0)
                tri_verts = torch.stack([v0[idx], v1[idx], v2[idx]], dim=0)
                
                assert torch.allclose(segment[0], vertices[b_idx, edge_idx1], atol=1e-5) and torch.allclose(segment[1], vertices[b_idx, edge_idx2], atol=1e-5), 'Segment vertices do not match expected vertices from the mesh!'
                assert torch.allclose(tri_verts[0], vertices[b_idx, face_idx1], atol=1e-5) and torch.allclose(tri_verts[1], vertices[b_idx, face_idx2], atol=1e-5) and torch.allclose(tri_verts[2], vertices[b_idx, face_idx3], atol=1e-5), 'Triangle vertices do not match expected vertices from the mesh!'

                tri_faces = [[0, 1], [1, 2], [2, 0]]
                
                loss_v = final_loss[idx].item()

                # Helper to visualize mesh if available
                v_plot = vertices[b_idx]
                
                plotPoses([tri_verts, segment, v_plot], 
                        [tri_faces, [[0,1]], None], 
                        s=[10, 100, 1], 
                        colors_p=['blue', 'red', 'gray'], 
                        colors_l=['blue', 'red', 'gray'], 
                        legend=['Target Face', 'Penetrating Point', 'Pts'], 
                        title=f'B{b_idx} Collision: {loss_v:.4f}m', 
                        show=True)

        is_intersect_refined_mask = is_intersect > self.is_intersect_prob
        if is_intersect_refined_mask.any():
            return final_loss.sum() / len(vertices), is_intersect_refined_mask.sum().item()
        else:
            return torch.tensor(0.0, device=vertices.device), 0

    def forward(self, smplx_out, vis=False):
        loss_l, count_l = self.forwardQueryTarget(smplx_out.vertices, self.query_faces_l, self.target_faces_l, self.query_edges_l, vis)
        loss_r, count_r = self.forwardQueryTarget(smplx_out.vertices, self.query_faces_r, self.target_faces_r, self.query_edges_r, vis)
        return loss_l + loss_r, count_l + count_r


def compute_pearson_loss_batched(x, y, mask, eps=1e-8):
    """
    Computes Pearson Correlation depth loss.
    
    Robust to:
    - Scale differences (meters vs raw logits)
    - Shift differences (absolute distance offsets)
    - Camera focal length mismatches
    
    Args:
        x: B x D
        y: B x D
        mask: B x D (boolean mask for valid entries)
    Returns:
        loss (scalar): 1 - mean(Pearson Correlation)
    """
    assert x.shape == y.shape == mask.shape, "Input shapes must match"
    assert x.ndim == 2

    # Apply mask to x and y
    x = x * mask
    y = y * mask

    # Compute means over valid entries (excluding masked-out values)
    mask_sum = mask.sum(dim=1, keepdim=True) + eps # B x 1
    mu_x = (x * mask).sum(-1, keepdim=True) / mask_sum
    mu_y = (y * mask).sum(-1, keepdim=True) / mask_sum

    # Center the data (subtract mean)
    x_c = x - mu_x
    y_c = y - mu_y

    # Compute variance and covariance terms
    var_x = ((x_c * mask) ** 2).sum(dim=1, keepdim=True) / mask_sum
    var_y = ((y_c * mask) ** 2).sum(dim=1, keepdim=True) / mask_sum
    cov_xy = ((x_c * mask) * (y_c * mask)).sum(dim=1, keepdim=True) / mask_sum

    # Compute Pearson correlation coefficient
    std_product = torch.sqrt(var_x * var_y) + eps
    rho = cov_xy / std_product

    # Return mean loss over batch dimension
    losses = 1.0 - rho
    
    #### test ####
    # losses_test = torch.zeros_like(losses)
    # for i in range(x.shape[0]):
    #     losses_test[i] = compute_pearson_loss(x[i][mask[i]>0], y[i][mask[i]>0], eps)
    # assert torch.allclose(losses, losses_test, atol=1e-5), f"Batch Pearson loss does not match individual computation! Max diff: {(losses - losses_test).abs().max().item()}"
    # loss_single = compute_pearson_loss(x.flatten()[mask.flatten()>0], y.flatten()[mask.flatten()>0], eps)
    ##############

    return losses.mean()

def compute_pearson_depth_loss(smpl_metric_depth, vda_relative_depth, visibility_weight, eps=1e-8):
    """
    smpl_metric_depth obtain by using persepctive projection of the mesh with known pin-hole camera parameters.
    vda_relative_depth: Video Depth Anything relative depth output. 

    Computes Pearson Correlation depth loss.
    
    Robust to:
    - Scale differences (meters vs raw logits)
    - Shift differences (absolute distance offsets)
    - Camera focal length mismatches
    
    B - batch size
    C - is number of cameras
    N - is number of joints
    Args:
        smpl_metric_depth (B, C, N): Metric Z depth (meters).
        vda_relative_depth (B, C, N): Raw disparity/inverse depth (any range).
        visibility_weight (B, C, N): Mask [0, 1] for valid joints.
        
    Returns:
        loss (scalar): 1 - mean(Pearson Correlation)
    """
    assert smpl_metric_depth.shape == vda_relative_depth.shape == visibility_weight.shape, "Input shapes must match"
    assert smpl_metric_depth.ndim == 3

    smpl_metric_depth = 1.0 / torch.clamp(smpl_metric_depth, min=0.01) # Convert to disparity space so that pearson correlates positively
    C = smpl_metric_depth.shape[1]
    loss = 0.0
    mask_flattened = visibility_weight > 0.5
    for c in range(C):
        loss += compute_pearson_loss_batched(smpl_metric_depth[:, c], vda_relative_depth[:, c], mask_flattened[:, c], eps)
    return loss / C
