"""
PyTorch model extractor.

Extracts structural (weights, layer configs), functional (I/O pairs, activations),
and meta knowledge from any torch.nn.Module.
"""

from __future__ import annotations

from collections import OrderedDict
from typing import Any, Dict, List, Optional, Tuple

import numpy as np

from model_osmosis.core.exceptions import ExtractionError
from model_osmosis.extractors.base import BaseExtractor
from model_osmosis.knowledge.representation import (
    FunctionalKnowledge,
    MetaKnowledge,
    StructuralKnowledge,
)


class PyTorchExtractor(BaseExtractor):
    """Extract knowledge from PyTorch ``torch.nn.Module`` models."""

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

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

    def _extract_structural(self, model: Any) -> StructuralKnowledge:
        import torch

        parameters: Dict[str, np.ndarray] = {}
        layer_info: List[Dict[str, Any]] = []
        total_params = 0

        # Extract all named parameters as numpy arrays
        for name, param in model.named_parameters():
            parameters[name] = param.detach().cpu().numpy()
            total_params += param.numel()

        # Extract layer-by-layer info
        for name, module in model.named_modules():
            if name == "":
                continue  # skip the top-level module itself

            info: Dict[str, Any] = {
                "name": name,
                "type": type(module).__name__,
            }

            # Capture layer-specific details
            if hasattr(module, "in_features"):
                info["in_features"] = module.in_features
            if hasattr(module, "out_features"):
                info["out_features"] = module.out_features
            if hasattr(module, "in_channels"):
                info["in_channels"] = module.in_channels
            if hasattr(module, "out_channels"):
                info["out_channels"] = module.out_channels
            if hasattr(module, "kernel_size"):
                info["kernel_size"] = module.kernel_size
            if hasattr(module, "weight") and module.weight is not None:
                info["weight_shape"] = list(module.weight.shape)
            if hasattr(module, "bias") and module.bias is not None:
                info["bias_shape"] = list(module.bias.shape)

            layer_info.append(info)

        # Build architecture description
        arch_desc = _describe_pytorch_architecture(model)

        return StructuralKnowledge(
            model_type="pytorch",
            architecture_desc=arch_desc,
            layer_info=layer_info,
            parameters=parameters,
            total_params=total_params,
        )

    def _extract_functional(
        self, model: Any, sample_data: np.ndarray
    ) -> FunctionalKnowledge:
        import torch

        model.eval()
        device = next(model.parameters()).device

        x_tensor = torch.from_numpy(sample_data).float().to(device)

        # Capture intermediate activations via hooks
        activations: Dict[str, np.ndarray] = {}
        hooks = []

        def _make_hook(layer_name: str):
            def hook_fn(module, inp, output):
                if isinstance(output, torch.Tensor):
                    activations[layer_name] = output.detach().cpu().numpy()
            return hook_fn

        for name, module in model.named_modules():
            if name == "":
                continue
            h = module.register_forward_hook(_make_hook(name))
            hooks.append(h)

        # Forward pass
        with torch.no_grad():
            outputs = model(x_tensor)

        # Remove hooks
        for h in hooks:
            h.remove()

        if isinstance(outputs, torch.Tensor):
            out_np = outputs.detach().cpu().numpy()
        else:
            # Handle tuple/list outputs — take the first element
            out_np = outputs[0].detach().cpu().numpy()

        # Determine if classifier
        is_classifier = False
        num_classes = 0
        if out_np.ndim == 2 and out_np.shape[1] > 1:
            # Check if outputs look like probabilities/logits
            is_classifier = True
            num_classes = out_np.shape[1]

        return FunctionalKnowledge(
            inputs=sample_data,
            outputs=out_np,
            intermediate_activations=activations,
            is_classifier=is_classifier,
            num_classes=num_classes,
        )

    def _extract_meta(self, model: Any) -> MetaKnowledge:
        total_params = sum(p.numel() for p in model.parameters())

        # Try to infer input shape from first layer
        input_shape: Tuple[int, ...] = ()
        output_shape: Tuple[int, ...] = ()

        for module in model.modules():
            if hasattr(module, "in_features"):
                input_shape = (module.in_features,)
                break
            if hasattr(module, "in_channels"):
                input_shape = (module.in_channels,)
                break

        # Try to infer output shape from last layer
        last_linear = None
        for module in model.modules():
            if hasattr(module, "out_features"):
                last_linear = module
        if last_linear is not None:
            output_shape = (last_linear.out_features,)

        task_type = "unknown"
        if output_shape and output_shape[0] == 1:
            task_type = "regression"
        elif output_shape and output_shape[0] > 1:
            task_type = "classification"

        return MetaKnowledge(
            source_framework="pytorch",
            source_class_name=type(model).__name__,
            input_shape=input_shape,
            output_shape=output_shape,
            task_type=task_type,
            extra={"total_params": total_params},
        )

    def _generate_sample_data(
        self, model: Any, meta: MetaKnowledge, num_samples: int
    ) -> np.ndarray:
        """Generate sample data matching the model's input expectations."""
        if meta.input_shape and len(meta.input_shape) > 0:
            shape = (num_samples, *meta.input_shape)
        else:
            shape = (num_samples, 10)
        return np.random.rand(*shape).astype(np.float32)


def _describe_pytorch_architecture(model: Any) -> str:
    """Build a human-readable architecture description."""
    lines = [f"PyTorch {type(model).__name__}"]
    for name, module in model.named_modules():
        if name == "":
            continue
        lines.append(f"  {name}: {type(module).__name__}")
        if hasattr(module, "in_features") and hasattr(module, "out_features"):
            lines.append(
                f"    shape: ({module.in_features}) -> ({module.out_features})"
            )
    return "\n".join(lines)
