"""
JAXAdapter — Load and convert JAX/Flax model parameters for ModelOsmosis.

Lazy-imports JAX so it's not a required dependency.

Usage::

    from model_osmosis.adapters import JAXAdapter

    adapter = JAXAdapter()
    expert = adapter.from_params(jax_params)
    osmo = adapter.to_osmo(jax_params)
    jax_params = adapter.from_osmo(osmo_model)
"""

from __future__ import annotations

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

import numpy as np


def _flatten_pytree(
    tree: Any, prefix: str = "", separator: str = "."
) -> Dict[str, np.ndarray]:
    """Flatten a JAX pytree (nested dict) into a flat dict of name → array.

    Args:
        tree: A nested dict/FrozenDict of arrays.
        prefix: Current key prefix for recursion.
        separator: Separator between nested keys.

    Returns:
        Flat dict of dotted-key → numpy array.
    """
    flat = {}
    if isinstance(tree, dict):
        items = tree.items()
    elif hasattr(tree, "items"):
        # Support FrozenDict, flax.core.frozen_dict, etc.
        items = tree.items()
    else:
        # Leaf node — convert to numpy
        arr = np.asarray(tree)
        flat[prefix] = arr
        return flat

    for key, value in items:
        full_key = f"{prefix}{separator}{key}" if prefix else key
        if isinstance(value, (dict,)) or hasattr(value, "items"):
            flat.update(_flatten_pytree(value, full_key, separator))
        else:
            flat[full_key] = np.asarray(value)

    return flat


def _unflatten_pytree(
    flat: Dict[str, np.ndarray], separator: str = "."
) -> Dict[str, Any]:
    """Unflatten a flat dict back into a nested dict (pytree).

    Args:
        flat: Flat dict of dotted-key → numpy array.
        separator: Separator used during flattening.

    Returns:
        Nested dict matching the original pytree structure.
    """
    tree: Dict[str, Any] = {}
    for key, value in flat.items():
        parts = key.split(separator)
        current = tree
        for part in parts[:-1]:
            if part not in current:
                current[part] = {}
            current = current[part]
        current[parts[-1]] = value
    return tree


class JAXExpert:
    """Wrapper around JAX/Flax model params for use in ModelOsmosis pipelines."""

    def __init__(
        self,
        params: Dict[str, np.ndarray],
        apply_fn: Any = None,
        name: str = "jax_model",
    ):
        self.params = params
        self.apply_fn = apply_fn
        self.name = name

    @property
    def flat_params(self) -> Dict[str, np.ndarray]:
        """Get parameters as a flat dict."""
        if isinstance(self.params, dict) and any(
            isinstance(v, dict) or hasattr(v, "items")
            for v in self.params.values()
        ):
            return _flatten_pytree(self.params)
        return self.params

    def predict(self, data: np.ndarray) -> Optional[np.ndarray]:
        """Run inference if apply_fn is available."""
        if self.apply_fn is None:
            return None
        try:
            import jax.numpy as jnp

            jax_input = jnp.array(data)
            output = self.apply_fn(self.params, jax_input)
            return np.asarray(output)
        except Exception:
            return None


class JAXAdapter:
    """Adapter for JAX/Flax model parameters.

    Supports:
    - Flattening/unflattening JAX pytrees
    - Converting pytree params to/from OsmoFormat
    - Working with Flax model params directly

    Requires: jax, jaxlib (optional: flax)
    """

    def __init__(self):
        self._jax = None

    def _lazy_import(self):
        """Lazy-import JAX to avoid hard dependency."""
        if self._jax is None:
            try:
                import jax

                self._jax = jax
            except ImportError:
                raise ImportError(
                    "JAX is required for JAXAdapter. "
                    "Install it with: pip install jax jaxlib"
                )

    def from_params(
        self, params: Any, apply_fn: Any = None, name: str = "jax_model"
    ) -> JAXExpert:
        """Create a JAXExpert from raw parameters.

        Args:
            params: JAX pytree of model parameters (nested dict of arrays).
            apply_fn: Optional function(params, input) -> output for inference.
            name: Human-readable name.

        Returns:
            JAXExpert wrapping the parameters.
        """
        return JAXExpert(params=params, apply_fn=apply_fn, name=name)

    def extract_weights(self, model: Any) -> Dict[str, np.ndarray]:
        """Extract all weights from JAX params as flat numpy arrays.

        Args:
            model: A JAXExpert, dict (pytree), or Flax TrainState.

        Returns:
            Flat dict of parameter name → numpy array.
        """
        if isinstance(model, JAXExpert):
            return model.flat_params

        # Handle Flax TrainState
        if hasattr(model, "params"):
            params = model.params
        else:
            params = model

        return _flatten_pytree(params)

    def load_weights(
        self, target_params: Any, weights: Dict[str, np.ndarray]
    ) -> Dict[str, Any]:
        """Load flat numpy weights back into a JAX pytree structure.

        Args:
            target_params: Original pytree to match structure of.
            weights: Flat dict of weight name → numpy array.

        Returns:
            New pytree with updated weights.
        """
        return _unflatten_pytree(weights)

    def to_osmo(self, model: Any) -> Any:
        """Convert JAX/Flax model params to OsmoFormat.

        Args:
            model: JAXExpert, pytree dict, or Flax TrainState.

        Returns:
            OsmoModel with the model's weights and architecture.
        """
        from model_osmosis.format.osmo_format import (
            ArchitectureSpec,
            LayerSpec,
            LayerType,
            OsmoManifest,
            OsmoModel,
        )

        weights = self.extract_weights(model)
        name = model.name if isinstance(model, JAXExpert) else "jax_model"

        # Build architecture spec from parameter shapes
        layers = []
        for key, arr in weights.items():
            if arr.ndim >= 2:
                layer_type = LayerType.LINEAR
            elif "norm" in key.lower():
                layer_type = LayerType.NORM
            elif "embed" in key.lower():
                layer_type = LayerType.EMBEDDING
            else:
                layer_type = LayerType.LINEAR

            layers.append(
                LayerSpec(
                    name=key,
                    layer_type=layer_type,
                    shape=list(arr.shape),
                    dtype=str(arr.dtype),
                )
            )

        hidden_size = 0
        for arr in weights.values():
            if arr.ndim == 2:
                hidden_size = max(arr.shape)
                break

        architecture = ArchitectureSpec(
            layers=layers,
            hidden_size=hidden_size,
            num_layers=len(layers),
            extra={"framework": "jax"},
        )

        manifest = OsmoManifest(
            model_name=name,
            description="JAX/Flax model converted via JAXAdapter",
            tags=["jax", "flax"],
        )

        return OsmoModel(
            manifest=manifest,
            architecture=architecture,
            parameters=weights,
        )

    def from_osmo(
        self, osmo_model: Any, target_params: Any = None
    ) -> Dict[str, Any]:
        """Convert an OsmoModel back to JAX pytree format.

        Args:
            osmo_model: The OsmoModel to convert.
            target_params: Optional original pytree to match structure.
                If None, returns a flat dict.

        Returns:
            Nested dict (pytree) of parameters, or flat dict if no target.
        """
        flat_params = osmo_model.parameters

        if target_params is not None:
            return _unflatten_pytree(flat_params)

        # Try to reconstruct tree structure from dotted keys
        has_dots = any("." in k for k in flat_params.keys())
        if has_dots:
            return _unflatten_pytree(flat_params)

        return flat_params
