"""
PyTorch absorber — inject knowledge into a torch.nn.Module.

Two strategies:
1. Direct weight transfer (if source is also PyTorch and architectures overlap)
2. Knowledge distillation (universal — works for ANY source model type)
"""

from __future__ import annotations

import copy
from typing import Any, Dict, Optional

import numpy as np

from model_osmosis.absorbers.base import BaseAbsorber
from model_osmosis.core.config import OsmosisConfig
from model_osmosis.core.exceptions import AbsorptionError
from model_osmosis.knowledge.representation import UniversalKnowledge


class PyTorchAbsorber(BaseAbsorber):
    """Absorb knowledge into a PyTorch model."""

    def can_handle(self, model: Any) -> bool:
        try:
            import torch.nn as nn

            return isinstance(model, nn.Module)
        except ImportError:
            return False

    def absorb(
        self,
        target_model: Any,
        knowledge: UniversalKnowledge,
        config: OsmosisConfig,
    ) -> Any:
        """Inject knowledge into a PyTorch model.

        Strategy selection:
        - If source is PyTorch and has compatible layers: direct weight transfer
          followed by distillation fine-tuning.
        - Otherwise: pure knowledge distillation using I/O pairs.
        """
        import torch

        # Attempt direct weight transfer first
        if knowledge.source_type == "pytorch":
            transferred = self._direct_transfer(target_model, knowledge, config)
            if transferred > 0 and config.verbose:
                print(
                    f"[Absorber] Direct weight transfer: {transferred} "
                    f"parameters copied"
                )

        # Always do distillation to refine
        self._distill(target_model, knowledge, config)

        return target_model

    def _direct_transfer(
        self,
        target_model: Any,
        knowledge: UniversalKnowledge,
        config: OsmosisConfig,
    ) -> int:
        """Copy compatible weight tensors directly from source to target.

        Returns the number of parameters successfully transferred.
        """
        import torch

        transferred = 0
        target_state = target_model.state_dict()
        source_params = knowledge.structural.parameters

        for name, target_param in target_state.items():
            if name in source_params:
                source_np = source_params[name]
                if source_np.shape == target_param.shape:
                    target_state[name] = torch.from_numpy(source_np).to(
                        target_param.dtype
                    )
                    transferred += source_np.size
                elif _can_partial_copy(source_np.shape, target_param.shape):
                    # Partial copy: source is smaller, copy into top-left corner
                    slices = tuple(
                        slice(0, min(s, t))
                        for s, t in zip(source_np.shape, target_param.shape)
                    )
                    new_param = target_param.clone()
                    source_slice = tuple(
                        slice(0, min(s, t))
                        for s, t in zip(source_np.shape, target_param.shape)
                    )
                    new_param[slices] = torch.from_numpy(
                        source_np[source_slice]
                    ).to(target_param.dtype)
                    target_state[name] = new_param
                    transferred += np.prod(
                        [min(s, t) for s, t in zip(source_np.shape, target_param.shape)]
                    )

        target_model.load_state_dict(target_state)
        return int(transferred)

    def _distill(
        self,
        target_model: Any,
        knowledge: UniversalKnowledge,
        config: OsmosisConfig,
    ) -> None:
        """Train the target model to replicate the source's I/O behavior.

        Uses functional knowledge (input/output pairs) as training data.
        """
        import torch
        import torch.nn as nn
        from torch.utils.data import DataLoader, TensorDataset

        device = torch.device(config.device)
        target_model = target_model.to(device)
        target_model.train()

        # Prepare data
        x = torch.from_numpy(knowledge.functional.inputs).float().to(device)
        y = torch.from_numpy(knowledge.functional.outputs).float().to(device)

        dataset = TensorDataset(x, y)
        loader = DataLoader(
            dataset,
            batch_size=config.distillation_batch_size,
            shuffle=True,
        )

        # Loss function
        if knowledge.functional.is_classifier:
            # Use KL divergence for classification
            loss_fn = nn.KLDivLoss(reduction="batchmean")
        else:
            loss_fn = nn.MSELoss()

        optimizer = torch.optim.Adam(
            target_model.parameters(), lr=config.distillation_lr
        )

        # Training loop
        for epoch in range(config.distillation_epochs):
            epoch_loss = 0.0
            n_batches = 0

            for batch_x, batch_y in loader:
                optimizer.zero_grad()
                pred = target_model(batch_x)

                if knowledge.functional.is_classifier:
                    # Apply log_softmax for KL divergence
                    pred = torch.log_softmax(pred, dim=-1)
                    target_soft = torch.softmax(batch_y, dim=-1)
                    loss = loss_fn(pred, target_soft)
                else:
                    loss = loss_fn(pred, batch_y)

                loss.backward()
                optimizer.step()
                epoch_loss += loss.item()
                n_batches += 1

            if config.verbose and (epoch + 1) % max(1, config.distillation_epochs // 5) == 0:
                avg_loss = epoch_loss / max(n_batches, 1)
                print(
                    f"[Absorber] Distillation epoch {epoch + 1}/"
                    f"{config.distillation_epochs} — loss: {avg_loss:.6f}"
                )

        target_model.eval()


def _can_partial_copy(source_shape, target_shape) -> bool:
    """Check if a partial weight copy is possible (same number of dims)."""
    return len(source_shape) == len(target_shape)
