"""
⚡ ULTRON Systems // Multi-Task Financial Loss Function
Combines:
1. Huber Loss (Smooth L1) on forward return magnitudes (robust against fat-tailed black swan spikes).
2. Binary Cross-Entropy with Logits on directional movement sign.
3. Directional Accuracy evaluation metric.
"""
import torch
import torch.nn as nn
import torch.nn.functional as F
from typing import Tuple, Dict

class MultiTaskFinancialLoss(nn.Module):
    def __init__(self, dir_weight: float = 0.5, huber_delta: float = 0.01):
        super().__init__()
        self.dir_weight = dir_weight
        self.huber = nn.HuberLoss(delta=huber_delta)
        self.bce = nn.BCEWithLogitsLoss()

    def forward(
        self, 
        pred_returns: torch.Tensor, 
        pred_logits: torch.Tensor, 
        target_returns: torch.Tensor
    ) -> Tuple[torch.Tensor, Dict[str, float]]:
        """
        pred_returns: [Batch, Horiz]
        pred_logits:  [Batch, Horiz]
        target_returns: [Batch, Horiz]
        """
        # 1. Huber regression loss
        reg_loss = self.huber(pred_returns, target_returns)
        
        # 2. Binary target: 1.0 if return > 0 else 0.0
        target_dir = (target_returns > 0.0).float()
        dir_loss = self.bce(pred_logits, target_dir)
        
        # Total combined loss
        total_loss = reg_loss + (self.dir_weight * dir_loss)
        
        # Compute real-time directional accuracy (%)
        with torch.no_grad():
            pred_dir_bool = (pred_logits > 0.0)
            target_dir_bool = (target_returns > 0.0)
            dir_acc = (pred_dir_bool == target_dir_bool).float().mean().item() * 100.0
            
            # Sign accuracy of regression output
            reg_dir_acc = ((pred_returns > 0.0) == target_dir_bool).float().mean().item() * 100.0
            
        metrics = {
            "total_loss": total_loss.item(),
            "reg_loss": reg_loss.item(),
            "dir_loss": dir_loss.item(),
            "dir_acc_bce": dir_acc,
            "dir_acc_reg": reg_dir_acc
        }
        
        return total_loss, metrics