from __future__ import annotations

from dataclasses import dataclass
from typing import Optional, Tuple, Dict, Any, List
import logging

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.optim import AdamW
from torch.optim.lr_scheduler import CosineAnnealingLR

logger = logging.getLogger(__name__)


def _l2_normalize(x: torch.Tensor, dim: int = -1, eps: float = 1e-8) -> torch.Tensor:
    return x / (x.norm(dim=dim, keepdim=True) + eps)


def av_contrastive_loss(vision: torch.Tensor, audio: torch.Tensor, temperature: float = 0.07) -> torch.Tensor:
    """InfoNCE contrastive loss between vision and audio features.

    vision, audio: (B, D) — L2-normalised cosine similarity matrix.
    Trains the model to align vision and audio representations.
    """
    v = _l2_normalize(vision)
    a = _l2_normalize(audio)
    logits = v @ a.T / temperature
    labels = torch.arange(logits.size(0), device=logits.device)
    return (F.cross_entropy(logits, labels) + F.cross_entropy(logits.T, labels)) * 0.5


def av_sync_loss(vision_seq: torch.Tensor, audio_seq: torch.Tensor) -> torch.Tensor:
    """Audio-visual synchronisation margin loss.

    Inputs: (B, T, D) — positive pairs are in-sync, negatives are shifted.
    Encourages temporally aligned vision and audio sequences.
    """
    v = vision_seq.mean(dim=1)        # [B, D]
    a = audio_seq.mean(dim=1)         # [B, D]
    pos = (v * a).sum(dim=-1)
    neg = (v * a.roll(shifts=1, dims=0)).sum(dim=-1)
    return F.relu(1.0 - pos + neg).mean()


def masked_modeling_loss(
    seq: torch.Tensor,
    reconstruction_head: nn.Module,
    mask_ratio: float = 0.15,
) -> Tuple[torch.Tensor, torch.Tensor]:
    """Masked feature modeling loss with an actual reconstruction head.

    FIX: Previously zeroed masked tokens and computed MSE against original —
    the model never predicted anything, making the loss meaningless.
    Now the reconstruction_head predicts masked token values from context.

    seq:                 (B, T, D) — input feature sequence
    reconstruction_head: nn.Module — Linear(D, D) that maps context → prediction
    Returns: (loss, mask)
    """
    B, T, D = seq.shape
    num_mask = max(1, int(T * mask_ratio))
    mask = torch.zeros(B, T, dtype=torch.bool, device=seq.device)
    for b in range(B):
        idx = torch.randperm(T, device=seq.device)[:num_mask]
        mask[b, idx] = True

    target = seq.detach()                       # ground truth for masked positions
    corrupted = seq.clone()
    corrupted[mask] = 0.0                       # mask out input

    # Predict from corrupted context (reconstruction_head must be trained)
    pred = reconstruction_head(corrupted)       # [B, T, D]

    # Only compute loss on masked positions
    loss = F.mse_loss(pred[mask], target[mask])
    return loss, mask


def av_generation_loss(
    predicted_mel: torch.Tensor,
    target_mel: torch.Tensor,
    stop_logits: Optional[torch.Tensor] = None,
    stop_targets: Optional[torch.Tensor] = None,
) -> torch.Tensor:
    """Vision-to-speech mel generation loss.

    NEW: This is the primary training signal for the MelDecoder.
    Trains the model to produce correct mel spectrograms from vision input.

    predicted_mel: [B, T, n_mels]  — model output
    target_mel:    [B, T, n_mels]  — ground truth mel spectrogram
    stop_logits:   [B, T, 1]       — optional stop token predictions
    stop_targets:  [B, T, 1]       — optional binary stop targets (1=stop)
    """
    # L1 + L2 mel reconstruction for better perceptual quality
    mel_l1 = F.l1_loss(predicted_mel, target_mel)
    mel_l2 = F.mse_loss(predicted_mel, target_mel)
    mel_loss = mel_l1 + 0.5 * mel_l2

    if stop_logits is not None and stop_targets is not None:
        stop_loss = F.binary_cross_entropy_with_logits(stop_logits, stop_targets.float())
        return mel_loss + 0.1 * stop_loss

    return mel_loss


@dataclass
class FineTuningConfig:
    """Configuration for fine-tuning LISA models."""
    learning_rate: float = 1e-4
    weight_decay: float = 1e-5
    batch_size: int = 8
    num_epochs: int = 10
    warmup_steps: int = 100
    save_every: int = 1000
    eval_every: int = 500
    gradient_clip: float = 1.0

    # Loss weights
    contrastive_weight: float = 1.0
    sync_weight: float = 0.5
    generation_weight: float = 1.0   # mel generation loss weight (replaces masked_weight)

    # Fine-tuning strategy
    freeze_vision_backbone: bool = False
    freeze_audio_backbone: bool = False
    fine_tune_fusion_only: bool = False
    fine_tune_decoder_only: bool = False  # freeze everything except MelDecoder


@dataclass
class FineTuner:
    """Fine-tuning harness for LISA Vision→Audio model.

    Supports:
    - Contrastive alignment of vision and audio features
    - Audio-visual sync loss
    - Mel generation loss (primary signal for speech output)
    """
    model: Any
    config: FineTuningConfig = None

    def __post_init__(self):
        if self.config is None:
            self.config = FineTuningConfig()

        self.optimizer = None
        self.scheduler = None
        self.step_count = 0
        self.epoch_count = 0

        # Reconstruction head for masked modeling (D→D)
        self._recon_head: Optional[nn.Module] = None

        logger.info(f"FineTuner initialized with config: {self.config}")

    def setup_training(self, output_dir: str = "./finetuned_model"):
        self.output_dir = output_dir
        params = self._get_trainable_parameters()

        # Build a shared reconstruction head matching audio embed dim
        audio_dim = getattr(self.model, '_detected_dims', {}).get('audio', 512)
        self._recon_head = nn.Linear(audio_dim, audio_dim)

        self.optimizer = AdamW(params, lr=self.config.learning_rate, weight_decay=self.config.weight_decay)
        self.scheduler = CosineAnnealingLR(
            self.optimizer,
            T_max=self.config.num_epochs * 1000,
            eta_min=self.config.learning_rate * 0.1,
        )
        logger.info("Training setup complete.")

    def _get_trainable_parameters(self) -> List:
        params = []

        if self.config.fine_tune_decoder_only:
            # Only train MelDecoder for vision→speech
            if hasattr(self.model, 'mel_decoder'):
                params.extend(self.model.mel_decoder.parameters())
            logger.info("Fine-tuning MelDecoder only")
            return params

        if self.config.fine_tune_fusion_only:
            if hasattr(self.model, 'fusion_model'):
                params.extend(self.model.fusion_model.parameters())
            if hasattr(self.model, 'mel_decoder'):
                params.extend(self.model.mel_decoder.parameters())
            logger.info("Fine-tuning fusion + decoder only")
            return params

        # Vision
        if hasattr(self.model, 'vision_vit') and not self.config.freeze_vision_backbone:
            params.extend(self.model.vision_vit.parameters())
            if hasattr(self.model, 'object_detector'):
                params.extend(self.model.object_detector.parameters())

        # Audio
        if hasattr(self.model, 'speech_recognizer') and not self.config.freeze_audio_backbone:
            params.extend(self.model.speech_recognizer.parameters())
            if hasattr(self.model, 'sound_classifier'):
                params.extend(self.model.sound_classifier.parameters())
            if hasattr(self.model, 'emotion_detector'):
                params.extend(self.model.emotion_detector.parameters())

        # Fusion + decoder always trained
        if hasattr(self.model, 'fusion_model'):
            params.extend(self.model.fusion_model.parameters())
        if hasattr(self.model, 'mel_decoder'):
            params.extend(self.model.mel_decoder.parameters())

        return params

    def train_step(
        self,
        vision_batch: torch.Tensor,
        audio_batch: torch.Tensor,
        target_mel: Optional[torch.Tensor] = None,
    ) -> Dict[str, float]:
        """One training step. target_mel required for mel generation loss."""
        self._set_train_mode()
        if self._recon_head is not None:
            self._recon_head.train()

        losses: Dict[str, float] = {}
        total_loss = torch.tensor(0.0, requires_grad=True)

        vision_features = self._extract_vision_features(vision_batch)
        audio_features  = self._extract_audio_features(audio_batch)

        # Contrastive loss
        if self.config.contrastive_weight > 0:
            cl = av_contrastive_loss(vision_features, audio_features)
            losses['contrastive'] = cl.item()
            total_loss = total_loss + self.config.contrastive_weight * cl

        # Sync loss (only when sequential data)
        if self.config.sync_weight > 0 and vision_features.dim() > 2:
            sl = av_sync_loss(vision_features, audio_features)
            losses['sync'] = sl.item()
            total_loss = total_loss + self.config.sync_weight * sl

        # Mel generation loss — the main training signal for vision→speech
        if self.config.generation_weight > 0 and target_mel is not None and hasattr(self.model, 'mel_decoder'):
            mel_result = self.model.mel_decoder(vision_features, target_mel)
            pred_mel   = mel_result['mel']
            stop_logits = mel_result.get('stop')

            # Align predicted and target lengths
            min_T = min(pred_mel.shape[1], target_mel.shape[1])
            gl = av_generation_loss(pred_mel[:, :min_T], target_mel[:, :min_T], stop_logits)
            losses['generation'] = gl.item()
            total_loss = total_loss + self.config.generation_weight * gl

        # Backprop
        self.optimizer.zero_grad()
        total_loss.backward()
        if self.config.gradient_clip > 0:
            torch.nn.utils.clip_grad_norm_(self._get_trainable_parameters(), self.config.gradient_clip)
        self.optimizer.step()
        self.scheduler.step()

        losses['total'] = total_loss.item()
        losses['lr']    = self.optimizer.param_groups[0]['lr']
        self.step_count += 1
        return losses

    def validate(
        self,
        val_vision_batch: torch.Tensor,
        val_audio_batch: torch.Tensor,
        target_mel: Optional[torch.Tensor] = None,
    ) -> Dict[str, float]:
        """Read-only validation — FIX: no optimizer calls inside validate."""
        self._set_eval_mode()

        with torch.no_grad():
            losses: Dict[str, float] = {}
            total_loss = 0.0

            vision_features = self._extract_vision_features(val_vision_batch)
            audio_features  = self._extract_audio_features(val_audio_batch)

            if self.config.contrastive_weight > 0:
                cl = av_contrastive_loss(vision_features, audio_features)
                losses['val_contrastive'] = cl.item()
                total_loss += self.config.contrastive_weight * cl.item()

            if self.config.generation_weight > 0 and target_mel is not None and hasattr(self.model, 'mel_decoder'):
                mel_result = self.model.mel_decoder(vision_features, target_mel)
                pred_mel   = mel_result['mel']
                stop_logits = mel_result.get('stop')
                min_T = min(pred_mel.shape[1], target_mel.shape[1])
                gl = av_generation_loss(pred_mel[:, :min_T], target_mel[:, :min_T], stop_logits)
                losses['val_generation'] = gl.item()
                total_loss += self.config.generation_weight * gl.item()

            losses['val_total'] = total_loss

        return losses

    # ---- internal helpers ----

    def _extract_vision_features(self, vision_batch: torch.Tensor) -> torch.Tensor:
        if hasattr(self.model, 'get_vision_features'):
            return self.model.get_vision_features(vision_batch)
        elif hasattr(self.model, 'vision_vit'):
            cls_token, _ = self.model.vision_vit(vision_batch)
            return cls_token
        raise ValueError("Cannot extract vision features from model")

    def _extract_audio_features(self, audio_batch: torch.Tensor) -> torch.Tensor:
        if hasattr(self.model, 'get_audio_features'):
            return self.model.get_audio_features(audio_batch)
        elif hasattr(self.model, 'speech_recognizer'):
            enc = self.model.speech_recognizer(audio_batch)   # [B, T, D]
            return enc.mean(dim=1)
        raise ValueError("Cannot extract audio features from model")

    def _set_train_mode(self):
        for attr in ['vision_vit', 'speech_recognizer', 'fusion_model', 'mel_decoder']:
            if hasattr(self.model, attr):
                getattr(self.model, attr).train()

    def _set_eval_mode(self):
        for attr in ['vision_vit', 'speech_recognizer', 'fusion_model', 'mel_decoder']:
            if hasattr(self.model, attr):
                getattr(self.model, attr).eval()

    def save_checkpoint(self, path: str, include_optimizer: bool = True):
        checkpoint = {
            'step_count': self.step_count,
            'epoch_count': self.epoch_count,
            'config': self.config,
        }
        if hasattr(self.model, 'state_dict'):
            checkpoint['model_state_dict'] = self.model.state_dict()
        else:
            checkpoint['model_state_dict'] = {}
            for attr in ['vision_vit', 'speech_recognizer', 'sound_classifier',
                         'emotion_detector', 'object_detector', 'vad_model',
                         'fusion_model', 'mel_decoder']:
                if hasattr(self.model, attr):
                    checkpoint['model_state_dict'][attr] = getattr(self.model, attr).state_dict()

        if include_optimizer and self.optimizer:
            checkpoint['optimizer_state_dict'] = self.optimizer.state_dict()
            checkpoint['scheduler_state_dict'] = self.scheduler.state_dict()

        torch.save(checkpoint, path)
        logger.info(f"Checkpoint saved to {path}")

    def load_checkpoint(self, path: str, load_optimizer: bool = True):
        checkpoint = torch.load(path, map_location='cpu')
        self.step_count  = checkpoint.get('step_count', 0)
        self.epoch_count = checkpoint.get('epoch_count', 0)

        if 'model_state_dict' in checkpoint:
            if hasattr(self.model, 'load_state_dict'):
                self.model.load_state_dict(checkpoint['model_state_dict'], strict=False)
            else:
                for attr, sd in checkpoint['model_state_dict'].items():
                    if hasattr(self.model, attr):
                        getattr(self.model, attr).load_state_dict(sd, strict=False)

        if load_optimizer and self.optimizer and 'optimizer_state_dict' in checkpoint:
            self.optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
            self.scheduler.load_state_dict(checkpoint['scheduler_state_dict'])

        logger.info(f"Checkpoint loaded from {path}")

    def step_av(self, vision_feat: torch.Tensor, audio_feat: torch.Tensor) -> torch.Tensor:
        """Legacy compatibility."""
        losses = self.train_step(vision_feat.unsqueeze(0), audio_feat.unsqueeze(0))
        return torch.tensor(losses['total'])


class ModelRepairTrainer(FineTuner):
    """Specialized trainer for repairing and fine-tuning models."""

    def __init__(self, model, config: FineTuningConfig = None, repair_strategy: str = "adaptive"):
        super().__init__(model, config)
        self.repair_strategy = repair_strategy

    def repair_and_train(
        self,
        train_data: List[Tuple[torch.Tensor, torch.Tensor]],
        val_data: List[Tuple[torch.Tensor, torch.Tensor]] = None,
        output_dir: str = "./repaired_model",
    ) -> Dict[str, Any]:
        logger.info("Starting model repair and training process")
        self.setup_training(output_dir)

        train_losses = []
        val_losses = []

        for epoch in range(self.config.num_epochs):
            self.epoch_count = epoch
            epoch_losses = []

            for batch_idx, (vision_batch, audio_batch) in enumerate(train_data):
                losses = self.train_step(vision_batch, audio_batch)
                epoch_losses.append(losses)

                if batch_idx % 50 == 0:
                    logger.info(f"Epoch {epoch}, Batch {batch_idx}, Loss: {losses['total']:.4f}")

                if val_data and batch_idx % self.config.eval_every == 0:
                    vb = val_data[batch_idx % len(val_data)]
                    vl = self.validate(vb[0], vb[1])
                    val_losses.append(vl)

                if batch_idx % self.config.save_every == 0:
                    self.save_checkpoint(f"{output_dir}/checkpoint_epoch_{epoch}_batch_{batch_idx}.pt")

            avg = sum(l['total'] for l in epoch_losses) / len(epoch_losses)
            train_losses.append(avg)
            logger.info(f"Epoch {epoch} avg loss: {avg:.4f}")

        final = f"{output_dir}/final_model.pt"
        self.save_checkpoint(final)

        return {
            'train_losses': train_losses,
            'val_losses': val_losses,
            'final_model_path': final,
            'total_steps': self.step_count,
        }
