"""Task-aware expansion backends for strict `osmosis expand` flows."""

from __future__ import annotations

import copy
from dataclasses import dataclass
from typing import Any, Dict, Optional, Tuple

import numpy as np


def count_params(model: Any) -> int:
    """Count parameters for torch modules or dict-like weight maps."""
    try:
        import torch.nn as nn

        if isinstance(model, nn.Module):
            return int(sum(p.numel() for p in model.parameters()))
    except Exception:
        pass

    if isinstance(model, dict):
        total = 0
        for value in model.values():
            total += int(np.asarray(value).size)
        return int(total)
    return 0


def _to_tensor_output(value: Any) -> Optional[np.ndarray]:
    """Best-effort conversion of framework output object to numpy."""
    if value is None:
        return None
    if hasattr(value, "detach"):
        try:
            return value.detach().cpu().numpy()
        except Exception:
            pass
    if hasattr(value, "logits"):
        return _to_tensor_output(getattr(value, "logits"))
    if isinstance(value, (list, tuple)) and value:
        return _to_tensor_output(value[0])
    try:
        return np.asarray(value)
    except Exception:
        return None


def _extract_tabular_features(sample_data: Any) -> Optional[np.ndarray]:
    if sample_data is None:
        return None
    if isinstance(sample_data, np.ndarray):
        arr = np.asarray(sample_data)
        if arr.ndim >= 2 and arr.size:
            return arr.astype("float32", copy=False)
        return None
    if isinstance(sample_data, dict):
        kind = str(sample_data.get("kind", "")).strip().lower()
        if kind == "tabular":
            features = sample_data.get("features")
            if features is None:
                return None
            arr = np.asarray(features)
            if arr.ndim >= 2 and arr.size:
                return arr.astype("float32", copy=False)
    return None


def evaluate_metric(model: Any, sample_data: Any, task: str, *, seed: int = 42) -> float:
    """Task-aware cheap deterministic quality proxy."""
    base = float(count_params(model))
    try:
        import torch
        import torch.nn as nn
    except Exception:
        return base

    if not isinstance(model, nn.Module):
        return base

    model = model.eval()
    with torch.no_grad():
        # text/code path: deterministic token-id probe
        if task in {"text", "code"}:
            vocab = 32000
            if hasattr(model, "config"):
                try:
                    vocab = int(getattr(model.config, "vocab_size", vocab) or vocab)
                except Exception:
                    pass
            batch = 2
            seq = 16
            rng = np.random.default_rng(int(seed))
            token_ids = rng.integers(0, max(2, vocab - 1), size=(batch, seq), dtype=np.int64)
            input_ids = torch.from_numpy(token_ids)
            try:
                out = model(input_ids=input_ids)
            except Exception:
                try:
                    out = model(input_ids)
                except Exception:
                    return base
            out_np = _to_tensor_output(out)
            if out_np is None or out_np.size == 0:
                return base
            return base + float(np.mean(np.abs(out_np)))

        # vision/audio/tabular path: tensor batch probe
        if isinstance(sample_data, dict):
            kind = str(sample_data.get("kind", "")).strip().lower()
            if kind == "vision":
                arr = np.asarray(sample_data.get("images", np.empty((0,))))
            elif kind == "audio":
                arr = np.asarray(sample_data.get("waveforms", np.empty((0,))))
            elif kind == "tabular":
                arr = np.asarray(sample_data.get("features", np.empty((0,))))
            else:
                arr = np.asarray(sample_data.get("features", np.empty((0,))))
        else:
            arr = np.asarray(sample_data) if sample_data is not None else np.empty((0,))

        if arr.size == 0:
            return base

        x = torch.from_numpy(arr.astype("float32", copy=False))
        try:
            out = model(x)
        except Exception:
            return base
        out_np = _to_tensor_output(out)
        if out_np is None or out_np.size == 0:
            return base
        return base + float(np.mean(np.abs(out_np)))


def quality_gate(task: str, baseline: float, post: float) -> Dict[str, Any]:
    """Task-specific non-regression gate."""
    tolerances = {
        "text": -1e-6,
        "code": -1e-6,
        "vision": -1e-6,
        "audio": -1e-6,
        "classification": -1e-6,
        "regression": -1e-6,
        "auto": -1e-6,
    }
    min_delta = float(tolerances.get(str(task).lower(), -1e-6))
    delta = float(post - baseline)
    return {
        "task": str(task),
        "baseline": float(baseline),
        "post": float(post),
        "delta": float(delta),
        "min_delta": float(min_delta),
        "passed": bool(delta >= min_delta),
    }


@dataclass
class ExpandExecutionResult:
    expanded_model: Any
    structural_applied: bool
    adaptation_applied: bool
    structural_reason: str
    adaptation_reason: str


class BaseExpandBackend:
    backend_id = "base"

    def can_expand(self, model: Any, sample_data: Any, task: str) -> Tuple[bool, str]:
        raise NotImplementedError

    def execute(
        self,
        model: Any,
        sample_data: Any,
        *,
        target_params: int,
        budget_minutes: int,
        seed: int,
    ) -> ExpandExecutionResult:
        raise NotImplementedError


def _replace_named_module(root: Any, name: str, replacement: Any) -> None:
    parts = str(name).split(".")
    parent = root
    for part in parts[:-1]:
        if part.isdigit():
            parent = parent[int(part)]
        else:
            parent = getattr(parent, part)
    leaf = parts[-1]
    if leaf.isdigit():
        parent[int(leaf)] = replacement
    else:
        setattr(parent, leaf, replacement)


def _first_linear_pair(model: Any) -> Tuple[Optional[Tuple[str, Any]], Optional[Tuple[str, Any]], str]:
    try:
        import torch.nn as nn
    except Exception:
        return None, None, "torch_not_available"

    # Strict safety: only support top-level Sequential for structural transforms.
    if not isinstance(model, nn.Sequential):
        return None, None, "strict_expand_requires_top_level_sequential"

    named = [(str(idx), mod) for idx, mod in enumerate(model)]
    for i in range(len(named) - 1):
        name_a, mod_a = named[i]
        name_b, mod_b = named[i + 1]
        if isinstance(mod_a, nn.Linear) and isinstance(mod_b, nn.Linear):
            if int(mod_b.in_features) != int(mod_a.out_features):
                continue
            return (name_a, mod_a), (name_b, mod_b), ""
    return None, None, "no_expandable_linear_pair_found"


class TabularExpandBackend(BaseExpandBackend):
    backend_id = "tabular"

    def can_expand(self, model: Any, sample_data: Any, task: str) -> Tuple[bool, str]:
        arr = _extract_tabular_features(sample_data)
        if arr is None:
            return False, "tabular_expand_requires_numeric_features"
        pair_a, pair_b, reason = _first_linear_pair(model)
        if pair_a is None or pair_b is None:
            return False, reason
        return True, ""

    def execute(
        self,
        model: Any,
        sample_data: Any,
        *,
        target_params: int,
        budget_minutes: int,
        seed: int,
    ) -> ExpandExecutionResult:
        try:
            import torch
            import torch.nn as nn
        except Exception:
            return ExpandExecutionResult(
                expanded_model=model,
                structural_applied=False,
                adaptation_applied=False,
                structural_reason="torch_not_available",
                adaptation_reason="torch_not_available",
            )

        pair_a, pair_b, reason = _first_linear_pair(model)
        if pair_a is None or pair_b is None:
            return ExpandExecutionResult(
                expanded_model=model,
                structural_applied=False,
                adaptation_applied=False,
                structural_reason=reason or "not_expandable",
                adaptation_reason="skipped_no_structural_expand",
            )

        teacher = model.eval()
        expanded = copy.deepcopy(model)
        (name_a, layer_a), (name_b, layer_b) = pair_a, pair_b

        old_out = int(layer_a.out_features)
        old_in = int(layer_a.in_features)
        next_out = int(layer_b.out_features)
        next_in = int(layer_b.in_features)
        if next_in != old_out:
            return ExpandExecutionResult(
                expanded_model=model,
                structural_applied=False,
                adaptation_applied=False,
                structural_reason="adjacent_linear_pair_shape_mismatch",
                adaptation_reason="skipped_no_structural_expand",
            )

        baseline_params = max(1, count_params(model))
        ratio = max(1.05, float(target_params) / float(baseline_params))
        new_out = max(old_out + 1, int(round(old_out * ratio)))

        rng = np.random.default_rng(int(seed))
        clone_indices = [int(x) for x in rng.integers(0, old_out, size=max(1, new_out - old_out))]

        new_a = nn.Linear(old_in, new_out, bias=layer_a.bias is not None)
        with torch.no_grad():
            new_a.weight[:old_out] = layer_a.weight.detach().clone()
            for i, src_idx in enumerate(clone_indices):
                dest = old_out + i
                if dest >= new_out:
                    break
                noise = torch.from_numpy(
                    rng.normal(0.0, 1e-3, size=tuple(layer_a.weight[src_idx].shape)).astype("float32")
                )
                new_a.weight[dest] = layer_a.weight[src_idx].detach().clone() + noise.to(new_a.weight.device)
            if layer_a.bias is not None and new_a.bias is not None:
                new_a.bias[:old_out] = layer_a.bias.detach().clone()
                for i, src_idx in enumerate(clone_indices):
                    dest = old_out + i
                    if dest >= new_out:
                        break
                    new_a.bias[dest] = layer_a.bias[src_idx].detach().clone()

        new_b = nn.Linear(new_out, next_out, bias=layer_b.bias is not None)
        with torch.no_grad():
            new_b.weight[:, :old_out] = layer_b.weight.detach().clone()
            base_col = torch.mean(layer_b.weight.detach(), dim=1)
            for i in range(new_out - old_out):
                new_b.weight[:, old_out + i] = base_col
            if layer_b.bias is not None and new_b.bias is not None:
                new_b.bias[:] = layer_b.bias.detach().clone()

        _replace_named_module(expanded, name_a, new_a)
        _replace_named_module(expanded, name_b, new_b)

        structural_applied = count_params(expanded) > count_params(model)
        if not structural_applied:
            return ExpandExecutionResult(
                expanded_model=model,
                structural_applied=False,
                adaptation_applied=False,
                structural_reason="param_count_not_increased",
                adaptation_reason="skipped_no_structural_expand",
            )

        features = _extract_tabular_features(sample_data)
        if features is None or features.size == 0:
            return ExpandExecutionResult(
                expanded_model=expanded.eval(),
                structural_applied=True,
                adaptation_applied=False,
                structural_reason="ok",
                adaptation_reason="no_features_for_adaptation",
            )

        x = torch.from_numpy(features.astype("float32", copy=False))
        with torch.no_grad():
            y = teacher(x)
            y_np = _to_tensor_output(y)
            if y_np is None:
                return ExpandExecutionResult(
                    expanded_model=expanded.eval(),
                    structural_applied=True,
                    adaptation_applied=False,
                    structural_reason="ok",
                    adaptation_reason="teacher_output_unavailable",
                )
            y = torch.from_numpy(y_np.astype("float32", copy=False))

        train_epochs = max(1, min(int(budget_minutes), 3))
        optimizer = torch.optim.Adam(expanded.parameters(), lr=1e-3)
        loss_fn = nn.MSELoss()
        expanded.train()
        for _ in range(train_epochs):
            optimizer.zero_grad()
            pred = expanded(x)
            if hasattr(pred, "logits"):
                pred = pred.logits
            if isinstance(pred, (list, tuple)):
                pred = pred[0]
            if not hasattr(pred, "dtype"):
                break
            if pred.dtype != y.dtype:
                pred = pred.to(y.dtype)
            loss = loss_fn(pred, y)
            loss.backward()
            optimizer.step()
        expanded.eval()
        return ExpandExecutionResult(
            expanded_model=expanded,
            structural_applied=True,
            adaptation_applied=True,
            structural_reason="ok",
            adaptation_reason="ok",
        )


class UnsupportedStrictBackend(BaseExpandBackend):
    def __init__(self, backend_id: str, reason: str):
        self.backend_id = backend_id
        self._reason = str(reason)

    def can_expand(self, model: Any, sample_data: Any, task: str) -> Tuple[bool, str]:
        return False, self._reason

    def execute(
        self,
        model: Any,
        sample_data: Any,
        *,
        target_params: int,
        budget_minutes: int,
        seed: int,
    ) -> ExpandExecutionResult:
        return ExpandExecutionResult(
            expanded_model=model,
            structural_applied=False,
            adaptation_applied=False,
            structural_reason=self._reason,
            adaptation_reason="skipped",
        )


class TextCodeExpandBackend(UnsupportedStrictBackend):
    def __init__(self):
        super().__init__(
            "text_code",
            "strict_text_code_expand_not_supported_for_this_architecture",
        )


class VisionExpandBackend(UnsupportedStrictBackend):
    def __init__(self):
        super().__init__(
            "vision",
            "strict_vision_expand_not_supported_for_this_architecture",
        )


class AudioExpandBackend(UnsupportedStrictBackend):
    def __init__(self):
        super().__init__(
            "audio",
            "strict_audio_expand_not_supported_for_this_architecture",
        )


def select_expand_backend(task: str) -> BaseExpandBackend:
    normalized = str(task or "auto").strip().lower()
    if normalized in {"classification", "regression", "auto"}:
        return TabularExpandBackend()
    if normalized in {"text", "code"}:
        return TextCodeExpandBackend()
    if normalized == "vision":
        return VisionExpandBackend()
    if normalized == "audio":
        return AudioExpandBackend()
    return UnsupportedStrictBackend("generic", "strict_expand_backend_unavailable")
