# coding: utf-8
import torch
import torch.nn as nn
from typing import Optional

from helpers import getSkeletalModelStructure, getSkeletalParentsDict
from pinn_losses import PINNLoss, PINNConfig

class Loss(nn.Module):

    def __init__(self, cfg, target_pad=0.0):
        super(Loss, self).__init__()

        self.loss = cfg["training"]["loss"].lower()
        self.bone_loss = cfg["training"]["bone_loss"].lower()

        if self.loss == "l1":
            self.criterion = nn.L1Loss()
        elif self.loss == "mse":
            self.criterion = nn.MSELoss()
        else:
            print("Loss not found - revert to default L1 loss")
            self.criterion = nn.L1Loss()

        if self.bone_loss == "l1":
            self.criterion_bone = nn.L1Loss()
        elif self.bone_loss == "mse":
            self.criterion_bone = nn.MSELoss()
        else:
            print("Loss not found - revert to default MSE loss")
            self.criterion_bone = nn.MSELoss()

        model_cfg = cfg["model"]
        training_cfg = cfg["training"]

        self.target_pad = target_pad
        self.loss_scale = model_cfg.get("loss_scale", 1.0)

        # PINN Configuration
        self.use_pinn = training_cfg.get("use_pinn", False)
        self.lambda_pinn = float(training_cfg.get("lambda_pinn", 0.5))
        
        if self.use_pinn:
            # Create PINN configuration from training config with proper type conversion
            pinn_cfg = PINNConfig(
                lambda_bone=float(training_cfg.get("pinn_lambda_bone", 1.0)),
                lambda_vel=float(training_cfg.get("pinn_lambda_vel", 0.1)),
                lambda_acc=float(training_cfg.get("pinn_lambda_acc", 0.05)),
                lambda_fk=float(training_cfg.get("pinn_lambda_fk", 0.5)),
                eps=float(training_cfg.get("pinn_eps", 1e-8)),
                dt=float(training_cfg.get("pinn_dt", 1.0)),
                rest_from=str(training_cfg.get("pinn_rest_from", "first_valid")),
                detach_rest=bool(training_cfg.get("pinn_detach_rest", True)),
                use_huber=bool(training_cfg.get("pinn_use_huber", True)),
                huber_delta=float(training_cfg.get("pinn_huber_delta", 1.0))
            )
            
            # Get parents dictionary from skeletal structure
            parents_dict = getSkeletalParentsDict()
            
            # Initialize PINN loss module
            self.pinn_loss = PINNLoss(
                parents=parents_dict,
                num_joints=50,
                cfg=pinn_cfg
            )
            
            print(f"PINN Loss initialized with {len(parents_dict)} parent-child pairs")
            print(f"  - lambda_pinn: {self.lambda_pinn}")
            print(f"  - lambda_bone: {pinn_cfg.lambda_bone}")
            print(f"  - lambda_vel: {pinn_cfg.lambda_vel}")
            print(f"  - lambda_acc: {pinn_cfg.lambda_acc}")
            print(f"  - lambda_fk: {pinn_cfg.lambda_fk}")
        else:
            self.pinn_loss = None
            print("PINN Loss is disabled")

    def forward(self, preds, targets, mask: Optional[torch.Tensor] = None):
        """
        Compute loss with optional PINN regularization
        
        Args:
            preds: predicted skeleton (B, T, 150)
            targets: target skeleton (B, T, 150)
            mask: optional mask (B, T) for valid frames
        """
        # Create loss mask from target padding
        loss_mask = (targets != self.target_pad)

        # Find the masked predictions and targets using loss mask
        preds_masked = preds * loss_mask
        targets_masked = targets * loss_mask

        # Compute bone length and direction features
        preds_masked_length, preds_masked_direct = get_length_direct(preds_masked)
        targets_masked_length, targets_masked_direct = get_length_direct(targets_masked)

        preds_masked_length = preds_masked_length * loss_mask[:, :, :50]
        targets_masked_length = targets_masked_length * loss_mask[:, :, :50]
        preds_masked_direct = preds_masked_direct * loss_mask[:, :, :150]
        targets_masked_direct = targets_masked_direct * loss_mask[:, :, :150]

        # Calculate base reconstruction loss
        recon_loss = self.criterion(preds_masked, targets_masked) + \
                     0.1 * self.criterion_bone(preds_masked_direct, targets_masked_direct)

        # Add PINN loss if enabled
        if self.use_pinn and self.pinn_loss is not None:
            # Create mask for PINN (B, T) - True for valid frames
            if mask is None:
                # Infer mask from targets: a frame is valid if it's not all padding
                # Check if any coordinate in the frame is non-pad
                pinn_mask = (targets[:, :, 0] != self.target_pad)  # (B, T)
            else:
                pinn_mask = mask
            
            # Compute PINN losses on predictions
            pinn_out = self.pinn_loss(preds, pinn_mask)
            
            # Total PINN loss
            pinn_total = pinn_out["total"]
            
            # Combined loss
            loss = recon_loss + self.lambda_pinn * pinn_total
        else:
            loss = recon_loss

        # Multiply loss by the loss scale
        if self.loss_scale != 1.0:
            loss = loss * self.loss_scale

        return loss


def get_length_direct(trg):
    """
    Compute bone lengths and unit directions from skeleton coordinates
    
    Args:
        trg: skeleton tensor (B, T, 150) - 50 joints x 3 coordinates
    
    Returns:
        lengths: (B, T, num_bones) - length of each bone
        directs: (B, T, 3*num_bones) - unit direction vectors (3D) for each bone
    """
    trg_reshaped = trg.view(trg.shape[0], trg.shape[1], 50, 3)
    trg_list = trg_reshaped.split(1, dim=2)
    trg_list_squeeze = [t.squeeze(dim=2) for t in trg_list]
    skeletons = getSkeletalModelStructure()

    length = []
    direct = []
    for skeleton in skeletons:
        Skeleton_length = torch.norm(trg_list_squeeze[skeleton[0]]-trg_list_squeeze[skeleton[1]], p=2, dim=2, keepdim=True)
        result_length = Skeleton_length
        result_direct = (trg_list_squeeze[skeleton[0]]-trg_list_squeeze[skeleton[1]]) / (Skeleton_length+torch.finfo(Skeleton_length.dtype).tiny)
        direct.append(result_direct)
        length.append(result_length)
    lengths = torch.stack(length, dim=-1).squeeze()
    directs = torch.stack(direct, dim=2).view(trg.shape[0], trg.shape[1], -1)

    return lengths, directs