import torch
import torch.nn as nn
import numpy as np
from typing import List, Optional
from pimenet.logger import get_logger

logger = get_logger(__name__)

class AlignedMTL(nn.Module):
    """
    Aligned Multi-Task Learning (Aligned-MTL) implementation.
    
    Based on "Independent Component Alignment for Multi-Task Learning (CVPR 2023)":
    - Computes gradient alignment matrix M = G @ G^T
    - Performs eigenvalue decomposition to find independent components
    - Uses alignment-based weighting for gradient combination
    
    Args:
        task_num: Number of tasks
        device: Device to run on
    """
    def __init__(self, task_num: Optional[int] = None, device: Optional[torch.device] = None):
        super(AlignedMTL, self).__init__()
        self.task_num = task_num
        self.device = device
        self.share_params: Optional[List[nn.Parameter]] = None
        self._grad_index: Optional[List[int]] = None
        self._grad_dim: Optional[int] = None
        self._prefix_sum: Optional[torch.Tensor] = None
        
    def _get_grads(self, losses: List[torch.Tensor], mode: str = 'backward') -> torch.Tensor:
        """
        Get gradients for all tasks.
        
        Args:
            losses: List of task losses
            mode: 'backward' or 'forward'
            
        Returns:
            Gradient matrix [num_tasks, grad_dim]
        """
        if self.share_params is None:
            raise ValueError("share_params not set. Call set_share_params first.")
        
        grads = []
        expected_grad_dim = sum(p.numel() for p in self.share_params)
        
        for i, loss in enumerate(losses):
            grad = torch.autograd.grad(
                loss, 
                self.share_params, 
                retain_graph=True, 
                create_graph=False,
                allow_unused=True
            )
            grad_list = []
            for j, g in enumerate(grad):
                if g is not None:
                    grad_list.append(g.flatten())
                else:
                    grad_list.append(torch.zeros(self.share_params[j].numel(), device=loss.device))
            
            grad_flat = torch.cat(grad_list)
            if grad_flat.shape[0] != expected_grad_dim:
                raise ValueError(f"Gradient dimension mismatch for task {i}: expected {expected_grad_dim}, got {grad_flat.shape[0]}")
            grads.append(grad_flat)
        
        if len(grads) == 0:
            raise ValueError("No valid gradients computed")
        
        grads = torch.stack(grads, dim=0)
        return grads
    
    def _backward_new_grads(self, alpha: torch.Tensor, grads: torch.Tensor):
        """
        Apply weighted gradients to shared parameters (memory-efficient).
        
        Args:
            alpha: Task weights [num_tasks]
            grads: Gradient matrix [num_tasks, grad_dim]
        """
        weighted_grad = torch.zeros_like(grads[0])
        for i in range(grads.shape[0]):
            weighted_grad += alpha[i] * grads[i]
        
        if self._prefix_sum is None:
            start_idx = 0
            for param in self.share_params:
                end_idx = start_idx + param.numel()
                param_grad = weighted_grad[start_idx:end_idx].view(param.shape)
                if param.grad is None:
                    param.grad = param_grad
                else:
                    param.grad += param_grad
                start_idx = end_idx
        else:
            for i, param in enumerate(self.share_params):
                start_idx = self._prefix_sum[i]
                end_idx = self._prefix_sum[i + 1]
                param_grad = weighted_grad[start_idx:end_idx].view(param.shape)
                if param.grad is None:
                    param.grad = param_grad
                else:
                    param.grad += param_grad
    
    def backward(self, losses: List[torch.Tensor], return_numpy: bool = True, apply_grads: bool = True, **kwargs):
        """
        Aligned-MTL backward pass (memory-efficient).
        
        Args:
            losses: List of task losses
            return_numpy: Whether to return weights as numpy array
            apply_grads: Whether to apply gradients to parameters
            **kwargs: Extra parameters
            
        Returns:
            Task weights as numpy array or torch.Tensor
        """
        grads = self._get_grads(losses, mode='backward')
        M = torch.matmul(grads, grads.t())
        lmbda, V = torch.linalg.eigh(M, UPLO='U')
        tol = (
            torch.max(lmbda) 
            * max(M.shape[-2:]) 
            * torch.finfo(lmbda.dtype).eps
        )
        rank = sum(lmbda > tol)
        order = torch.argsort(lmbda, dim=-1, descending=True)
        lmbda = lmbda[order][:rank]
        V = V[:, order][:, :rank]
        sigma = torch.diag(1 / lmbda.sqrt())
        B = lmbda[-1].sqrt() * ((V @ sigma) @ V.t())
        alpha = B.sum(0)
        
        if apply_grads:
            self._backward_new_grads(alpha, grads)
        
        # Cache alpha for task weight lookup
        self._last_alpha = alpha.detach().cpu().numpy()
        
        del M, lmbda, V, sigma, B, grads
        
        if return_numpy:
            return self._last_alpha
        return alpha
    
    def set_share_params(self, share_params: List[nn.Parameter]):
        """
        Set shared parameters for gradient computation.
        
        Args:
            share_params: List of shared parameters
        """
        self.share_params = share_params
        self._grad_index = []
        self._grad_dim = 0
        for param in share_params:
            self._grad_index.append(self._grad_dim)
            self._grad_dim += param.numel()
        
        self._prefix_sum = torch.zeros(len(share_params) + 1, dtype=torch.long)
        for i, param in enumerate(share_params):
            self._prefix_sum[i + 1] = self._prefix_sum[i] + param.numel()
    
    def get_task_weights(self) -> Optional[np.ndarray]:
        """
        Get current task weights.
        
        Returns:
            Task weights as numpy array, or None if not computed
        """
        return getattr(self, '_last_alpha', None)
    
    def forward(self, losses: List[torch.Tensor]) -> torch.Tensor:
        """
        Forward pass - just return weighted sum of losses.
        
        Args:
            losses: List of task losses
            
        Returns:
            Weighted sum of losses
        """
        if len(losses) != self.task_num:
            raise ValueError(f"Expected {self.task_num} losses, got {len(losses)}")
        return sum(losses)