"""
Universal Knowledge Representation (UKR).

This module defines the data structures that serve as the universal "language"
for model knowledge. Any source model's knowledge gets translated into this
format, and any target model absorbs from it.
"""

from __future__ import annotations

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

import numpy as np


@dataclass
class StructuralKnowledge:
    """Architecture-level knowledge extracted from a model.

    This captures the structural elements: weight matrices, layer configurations,
    tree split structures, support vectors, etc.

    Attributes:
        model_type: String identifier like 'pytorch', 'sklearn_tree', 'sklearn_linear'.
        architecture_desc: Human-readable description of the architecture.
        layer_info: List of dicts describing each layer/component.
            For neural nets: {'name': str, 'type': str, 'in_features': int,
                              'out_features': int, 'weight': np.ndarray, 'bias': np.ndarray}
            For trees: {'tree_structure': dict with splits, thresholds, values}
            For linear: {'coef': np.ndarray, 'intercept': np.ndarray}
        parameters: Dict mapping parameter names to numpy arrays (raw weights).
        total_params: Total number of learnable parameters.
    """

    model_type: str
    architecture_desc: str = ""
    layer_info: List[Dict[str, Any]] = field(default_factory=list)
    parameters: Dict[str, np.ndarray] = field(default_factory=dict)
    total_params: int = 0


@dataclass
class FunctionalKnowledge:
    """Behavioral knowledge — captures what the model has learned via I/O pairs.

    This is the most universal form of knowledge: given inputs, what outputs
    does the model produce? Works for ANY model type.

    Attributes:
        inputs: Sample input data (N, *input_shape).
        outputs: Corresponding model outputs (N, *output_shape).
        intermediate_activations: Optional dict mapping layer names to their
            activations for the sample inputs. Only available for neural nets.
        is_classifier: Whether the model is a classifier (affects loss function
            choice during distillation).
        num_classes: Number of output classes (for classifiers).
        output_shape: Shape of a single output sample.
    """

    inputs: np.ndarray
    outputs: np.ndarray
    intermediate_activations: Dict[str, np.ndarray] = field(default_factory=dict)
    is_classifier: bool = False
    num_classes: int = 0
    output_shape: Tuple[int, ...] = ()

    def __post_init__(self) -> None:
        if len(self.output_shape) == 0 and self.outputs is not None:
            if self.outputs.ndim > 1:
                self.output_shape = tuple(self.outputs.shape[1:])
            else:
                self.output_shape = (1,)


@dataclass
class MetaKnowledge:
    """Metadata about the source model and its training.

    Attributes:
        source_framework: e.g., 'pytorch', 'sklearn', 'tensorflow', 'generic'.
        source_class_name: Class name of the source model.
        input_shape: Expected input shape (excluding batch dimension).
        output_shape: Expected output shape (excluding batch dimension).
        feature_importances: Per-feature importance scores (if available).
        task_type: 'classification', 'regression', or 'unknown'.
        extra: Arbitrary extra metadata.
    """

    source_framework: str = "unknown"
    source_class_name: str = "unknown"
    input_shape: Tuple[int, ...] = ()
    output_shape: Tuple[int, ...] = ()
    feature_importances: Optional[np.ndarray] = None
    task_type: str = "unknown"
    extra: Dict[str, Any] = field(default_factory=dict)


@dataclass
class UniversalKnowledge:
    """Complete knowledge bundle extracted from a source model.

    This is the central data structure of ModelOsmosis. It holds structural,
    functional, and meta knowledge in a framework-agnostic format.
    """

    structural: StructuralKnowledge
    functional: FunctionalKnowledge
    meta: MetaKnowledge

    @property
    def source_type(self) -> str:
        """Short identifier for the source model type."""
        return self.structural.model_type

    @property
    def param_count(self) -> int:
        """Total parameter count of the source model."""
        return self.structural.total_params

    @property
    def is_classifier(self) -> bool:
        """Whether the source model is a classifier."""
        return self.functional.is_classifier

    def summary(self) -> str:
        """Human-readable summary of the extracted knowledge."""
        lines = [
            "=== Universal Knowledge Summary ===",
            f"Source: {self.meta.source_framework} / {self.meta.source_class_name}",
            f"Model type: {self.structural.model_type}",
            f"Task: {self.meta.task_type}",
            f"Input shape: {self.meta.input_shape}",
            f"Output shape: {self.meta.output_shape}",
            f"Total parameters: {self.structural.total_params:,}",
            f"Functional samples: {len(self.functional.inputs):,}",
            f"Layers/components: {len(self.structural.layer_info)}",
        ]
        if self.functional.intermediate_activations:
            lines.append(
                f"Intermediate activation layers: "
                f"{len(self.functional.intermediate_activations)}"
            )
        return "\n".join(lines)
