"""
MultiSense-DF — Grad-CAM Explainability
Generates heatmap overlays for the visual branch and attention weight visualisations
"""

import torch
import torch.nn.functional as F
import numpy as np
import cv2
import matplotlib.pyplot as plt
from pathlib import Path


class GradCAM:
    """
    Gradient-weighted Class Activation Mapping for EfficientNet-B4 backbone.
    Hooks into the last convolutional block.
    """

    def __init__(self, model, target_layer_name='visual.backbone.blocks.6'):
        self.model = model
        self.gradients = None
        self.activations = None
        self._register_hooks(target_layer_name)

    def _register_hooks(self, layer_name):
        target = dict(self.model.named_modules()).get(layer_name)
        if target is None:
            # Fallback: last conv block of EfficientNet
            for name, module in self.model.named_modules():
                if 'blocks' in name and hasattr(module, 'conv_pwl'):
                    target = module
            if target is None:
                raise ValueError(f'Layer {layer_name} not found.')

        target.register_forward_hook(self._save_activation)
        target.register_full_backward_hook(self._save_gradient)

    def _save_activation(self, module, input, output):
        self.activations = output.detach()

    def _save_gradient(self, module, grad_input, grad_output):
        self.gradients = grad_output[0].detach()

    def generate(self, frame_tensor, target_class=1):
        """
        Generate Grad-CAM heatmap for a single frame.
        frame_tensor: (1, 3, 224, 224)
        Returns: numpy heatmap (224, 224) in [0, 1]
        """
        self.model.eval()
        frame_tensor.requires_grad_(True)

        # Forward through visual backbone only
        feat = self.model.visual.backbone(frame_tensor)
        score = feat.mean()   # simplified — use actual logit in full pipeline
        score.backward()

        # Global average pool gradients
        weights = self.gradients.mean(dim=(2, 3), keepdim=True)
        cam = (weights * self.activations).sum(dim=1, keepdim=True)
        cam = F.relu(cam)
        cam = F.interpolate(cam, size=(224, 224), mode='bilinear', align_corners=False)
        cam = cam.squeeze().cpu().numpy()
        cam = (cam - cam.min()) / (cam.max() - cam.min() + 1e-8)
        return cam

    def overlay(self, frame_np, cam, alpha=0.5):
        """
        Overlay Grad-CAM heatmap on original frame.
        frame_np: (H, W, 3) uint8 RGB
        Returns: (H, W, 3) uint8 RGB overlay
        """
        heatmap = cv2.applyColorMap(
            (cam * 255).astype(np.uint8), cv2.COLORMAP_JET
        )
        heatmap = cv2.cvtColor(heatmap, cv2.COLOR_BGR2RGB)
        heatmap = cv2.resize(heatmap, (frame_np.shape[1], frame_np.shape[0]))
        overlay = (alpha * heatmap + (1 - alpha) * frame_np).astype(np.uint8)
        return overlay


def visualise_attention_weights(attn_weights: dict, save_path: str = None):
    """
    Bar chart of per-modality contribution from fusion attention weights.
    attn_weights: dict with 'visual_weight', 'audio_weight', 'lipsync_weight'
    """
    labels = ['Visual', 'Audio', 'Lip-Sync']
    colors = ['#FF6B6B', '#4ECDC4', '#45B7D1']
    values = [
        attn_weights['visual_weight'].mean().item(),
        attn_weights['audio_weight'].mean().item(),
        attn_weights['lipsync_weight'].mean().item(),
    ]
    total = sum(values) + 1e-8
    percentages = [v / total * 100 for v in values]

    fig, ax = plt.subplots(figsize=(6, 3))
    bars = ax.barh(labels, percentages, color=colors, height=0.5)
    ax.set_xlim(0, 100)
    ax.set_xlabel('Relative Contribution (%)')
    ax.set_title('Per-Modality Attention Weights')
    for bar, pct in zip(bars, percentages):
        ax.text(bar.get_width() + 1, bar.get_y() + bar.get_height() / 2,
                f'{pct:.1f}%', va='center', fontsize=10, fontweight='bold')
    plt.tight_layout()

    if save_path:
        plt.savefig(save_path, dpi=150, bbox_inches='tight')
    plt.show()
    return percentages


def generate_full_explanation(model, sample, device='cuda', save_dir='results/explanation'):
    """
    Full explanation pipeline: Grad-CAM + attention weights + confidence scores.
    sample: dict with 'frames', 'waveform', 'mouth_crops', 'mel_specs'
    """
    save_dir = Path(save_dir)
    save_dir.mkdir(parents=True, exist_ok=True)

    model.eval()
    frames      = sample['frames'].unsqueeze(0).to(device)
    waveform    = sample['waveform'].unsqueeze(0).to(device)
    mouth_crops = sample['mouth_crops'].unsqueeze(0).to(device)
    mel_specs   = sample['mel_specs'].unsqueeze(0).to(device)

    with torch.no_grad():
        outputs = model(frames, waveform, mouth_crops, mel_specs)

    global_prob = torch.sigmoid(outputs['global_logit']).item()
    per_mod = {
        k: torch.sigmoid(v).item()
        for k, v in outputs['per_mod_logits'].items()
    }

    # Attention weights visualisation
    attn_path = str(save_dir / 'attention_weights.png')
    pcts = visualise_attention_weights(outputs['attn_weights'], save_path=attn_path)

    print('\n── MultiSense-DF Explanation ───────────────')
    print(f'  Verdict       : {"FAKE" if global_prob > 0.5 else "REAL"}')
    print(f'  Confidence    : {global_prob:.1%}')
    print(f'  Visual score  : {per_mod["visual"]:.1%}')
    print(f'  Audio score   : {per_mod["audio"]:.1%}')
    print(f'  Lip-sync score: {per_mod["lipsync"]:.1%}')
    print(f'  Contribution  : Visual={pcts[0]:.1f}%  Audio={pcts[1]:.1f}%  Lip-sync={pcts[2]:.1f}%')
    print('────────────────────────────────────────────\n')

    return {
        'verdict': 'FAKE' if global_prob > 0.5 else 'REAL',
        'confidence': global_prob,
        'visual_score': per_mod['visual'],
        'audio_score': per_mod['audio'],
        'lipsync_score': per_mod['lipsync'],
        'contributions': pcts
    }