import functools
from typing import Any, Callable

import torch
import torch.nn as nn

from .context import Frame, Mode, pop_frame, push_frame


def _clean_key(key: str) -> str:
    return key.replace("__", ".")


class BlazeModule(nn.Module):
    """A PyTorch ``nn.Module`` produced by :func:`transform`.

    Call :meth:`init` once with a sample input to initialise all sub-modules,
    then use the model normally (``model(x)``).
    """

    def __init__(self):
        super().__init__()
        self._inited = False

    @torch.jit.ignore
    def init(self, *args: Any, **kwargs: Any) -> "BlazeModule":
        """Run the wrapped function in INIT mode to discover and create all modules."""
        frame = Frame(mode=Mode.INIT)
        push_frame(frame)
        try:
            self._fn(*args, **kwargs)
        finally:
            pop_frame()

        self._registry = frame.registry
        self._inited = True

        # Register discovered modules so parameters() / buffers() work.
        for key, mod in self._registry.items():
            self.add_module(key.replace(".", "__"), mod)

        return self

    @torch.jit.unused
    def _forward_dynamic(self, *args: Any, **kwargs: Any) -> Any:
        """Full dynamic dispatch — used in eager mode."""
        if not self._inited:
            raise RuntimeError(
                "Model not compiled. Call .init(sample_input) first."
            )
        frame = Frame(mode=Mode.APPLY)
        frame.registry = dict(self._registry)
        push_frame(frame)
        try:
            return self._fn(*args, **kwargs)
        finally:
            pop_frame()

    @torch.jit.ignore
    def train(self, mode: bool = True):
        super().train(mode)
        if hasattr(self, "_registry"):
            for mod in self._registry.values():
                mod.train(mode)
        return self

    @torch.jit.ignore
    def named_parameters(self, prefix="", recurse=True):
        seen = set()
        for name, p in nn.Module.named_parameters(self, prefix=prefix, recurse=recurse):
            clean = _clean_key(name)
            if p.data_ptr() not in seen:
                seen.add(p.data_ptr())
                yield clean, p

    @torch.jit.ignore
    def named_buffers(self, prefix="", recurse=True):
        seen = set()
        for name, b in nn.Module.named_buffers(self, prefix=prefix, recurse=recurse):
            clean = _clean_key(name)
            if b.data_ptr() not in seen:
                seen.add(b.data_ptr())
                yield clean, b

    @torch.jit.ignore
    def state_dict(self, *args, **kwargs):
        sd = nn.Module.state_dict(self, *args, **kwargs)
        out = {}
        for k, v in sd.items():
            clean = _clean_key(k)
            if clean not in out:
                out[clean] = v
        return out

    @torch.jit.ignore
    def load_state_dict(self, state_dict, strict=True, **kwargs):
        internal_sd = nn.Module.state_dict(self)
        clean_to_internal = {_clean_key(k): k for k in internal_sd}
        mapped = {clean_to_internal.get(k, k): v for k, v in state_dict.items()}
        nn.Module.load_state_dict(self, mapped, strict=strict, **kwargs)

    def forward(self, *args: Any, **kwargs: Any) -> Any:
        return self._forward_dynamic(*args, **kwargs)

    @torch.jit.ignore
    def __repr__(self) -> str:
        keys = list(self._registry.keys()) if hasattr(self, "_registry") else []
        return f"BlazeModule(modules={keys})"

@torch.jit.ignore
def transform(fn: Callable, *args: Any, **kwargs: Any) -> BlazeModule:
    """Wrap a forward-pass function into a :class:`BlazeModule`.

    Any extra ``*args`` / ``**kwargs`` are partially applied to *fn* so that
    :meth:`~BlazeModule.init` and the forward call only need the
    tensor inputs.

    Parameters:
        fn: A function that constructs and applies sub-modules. The first time
            it's called, Blaze will run it in INIT mode to discover all modules
            created and build the internal registry. Subsequent calls run in
            APPLY mode, where the registry is fixed and module creation is an
            error.
        *args, **kwargs: Optional extra arguments to partially apply to *fn*.
    Example::

        model = blaze.transform(forward, in_size=10, out_size=20)
        model.init(torch.randn(5, 10))
        output = model(torch.randn(5, 10))
    """
    if args or kwargs:
        fn = functools.partial(fn, *args, **kwargs)
    model = BlazeModule()
    model._fn = fn  # stored outside __init__ to stay invisible to TorchScript
    return model
