"""
Memory consolidation - compressing recent learning into stable patterns.

Like sleep for neural memory: strengthens important patterns,
prunes noise, and prevents catastrophic forgetting.
"""

from typing import Dict, List, Union

import torch
import torch.nn as nn
from torch import Tensor


class MemoryConsolidator:
    """
    Consolidates memory by compressing recent learning into stable patterns.

    Uses Elastic Weight Consolidation (EWC) to protect important weights.
    """

    def __init__(self, ewc_lambda: float = 0.1):
        """
        Initialize consolidator.

        Args:
            ewc_lambda: Weight for EWC penalty term
        """
        self.ewc_lambda = ewc_lambda
        self.fisher: Dict[str, Tensor] = {}
        self.optimal: Dict[str, Tensor] = {}

    def compute_fisher(self, model: nn.Module, data_loader: List[Tensor]) -> None:
        """
        Compute Fisher information matrix for EWC.

        The Fisher matrix measures how important each weight is
        for the current task.
        """
        self.fisher = {}
        self.optimal = {}

        # Store current optimal weights
        for name, param in model.named_parameters():
            self.optimal[name] = param.clone().detach()
            self.fisher[name] = torch.zeros_like(param)

        model.eval()
        for x in data_loader:
            model.zero_grad()
            output = model(x)
            # Use reconstruction loss as proxy for importance
            loss = output.sum()
            loss.backward()

            for name, param in model.named_parameters():
                if param.grad is not None:
                    self.fisher[name] += param.grad.data.pow(2)

        # Normalize Fisher values
        for name in self.fisher:
            self.fisher[name] /= len(data_loader)

    def ewc_loss(self, model: nn.Module, current_loss: Tensor) -> Tensor:
        """
        Add EWC penalty to protect important weights from past tasks.

        Args:
            model: The neural memory model
            current_loss: Current task loss

        Returns:
            Loss with EWC penalty added
        """
        if not self.fisher:
            return current_loss

        ewc_penalty = torch.tensor(0.0, device=current_loss.device)
        for name, param in model.named_parameters():
            if name in self.fisher:
                ewc_penalty += (self.fisher[name] * (param - self.optimal[name]).pow(2)).sum()

        return current_loss + self.ewc_lambda * ewc_penalty

    def consolidate(
        self, model: nn.Module, recent_observations: List[Tensor]
    ) -> Dict[str, Union[int, float]]:
        """
        Perform consolidation pass.

        Args:
            model: Neural memory model
            recent_observations: Recent data to identify important patterns

        Returns:
            Consolidation metrics
        """
        # Create simple data loader from observations
        data_loader = recent_observations

        # Compute new Fisher information
        old_fisher = self.fisher.copy() if self.fisher else {}
        self.compute_fisher(model, data_loader)

        # Merge with existing Fisher (weighted average)
        patterns_merged = 0
        if old_fisher:
            for name in self.fisher:
                if name in old_fisher:
                    self.fisher[name] = 0.5 * (self.fisher[name] + old_fisher[name])
                    patterns_merged += 1

        # Compute compression metric (how much the Fisher changed)
        compression = 0.0
        if old_fisher:
            for name in self.fisher:
                if name in old_fisher:
                    diff = (self.fisher[name] - old_fisher[name]).abs().mean()
                    compression += diff.item()

        return {
            "patterns_merged": patterns_merged,
            "memory_compressed_by": compression,
            "stability_score": 1.0 / (1.0 + compression),
        }