from __future__ import annotations
import types
import importlib
import torch


def _sib(mod, *names):
    try:
        m = importlib.import_module("." + mod, __package__ or None)
    except (ImportError, TypeError, ValueError):
        m = importlib.import_module(mod)
    return [getattr(m, n) for n in names]


(CPUGatedDeltaNet,) = _sib("cpu_delta", "CPUGatedDeltaNet")
(CPULandmark,) = _sib("cpu_landmark", "CPULandmark")
CPUSlidingWindow, swa_fused_forward = _sib(
    "cpu_swa", "CPUSlidingWindow", "swa_fused_forward"
)


def _is_delta_block(block) -> bool:
    return (
        getattr(block, "is_gla", False)
        and getattr(block.mixer, "gla_delta", False)
        and hasattr(block.mixer, "gdn")
    )


def _is_landmark_block(block) -> bool:
    return getattr(block, "is_landmark", False)


def _has_swa_fused(block) -> bool:
    return int(getattr(block.mixer, "swa_fused_window", 0)) > 0


def _make_block_methods(cpu_gdn, cpu_swa):

    def step(self, x, bstate):
        h = self.ln1(x)
        if cpu_swa is not None:
            h, bstate["swa"] = cpu_swa.step(h, bstate.get("swa"))
            if h.dim() == 2:
                h = h.unsqueeze(1)
        o, bstate["gdn"] = cpu_gdn.step(h, bstate.get("gdn"))
        x = x + o
        x = x + self.ffn(self.ln2(x))
        return (x, bstate)

    def forward_chunk(self, x, bstate):
        h = self.ln1(x)
        if cpu_swa is not None:
            h, bstate["swa"] = cpu_swa.forward_chunk(h, bstate.get("swa"))
        o, bstate["gdn"] = cpu_gdn.forward_chunk(h, bstate.get("gdn"))
        x = x + o
        x = x + self.ffn(self.ln2(x))
        return (x, bstate)

    return (step, forward_chunk)


def _make_landmark_block_methods(cpu_land):

    def step(self, x, bstate):
        h = self.ln1(x)
        o, bstate["land"] = cpu_land.step(h, bstate.get("land"))
        x = x + o
        x = x + self.ffn(self.ln2(x))
        return (x, bstate)

    def forward_chunk(self, x, bstate):
        h = self.ln1(x)
        o, bstate["land"] = cpu_land.forward_chunk(h, bstate.get("land"))
        x = x + o
        x = x + self.ffn(self.ln2(x))
        return (x, bstate)

    return (step, forward_chunk)


def _make_mixer_forward(cpu_gdn, cpu_swa, orig_forward):

    def forward(self, x):
        if x.is_cuda:
            return orig_forward(x)
        if cpu_swa is not None:
            x = swa_fused_forward(self, x)
        return cpu_gdn.forward(x)

    return forward


def _make_init_state(model):
    cfg = model.cfg
    H = cfg.n_head
    D = cfg.n_embd // H
    M = cfg.fno_modes
    C = cfg.n_embd

    def init_state(self, batch_size: int = 1, device=None):
        if device is None:
            device = next(self.parameters()).device
        states = []
        for block in self.blocks:
            if _is_delta_block(block):
                states.append(
                    {"swa": None, "gdn": None}
                    if _has_swa_fused(block)
                    else {"gdn": None}
                )
            elif _is_landmark_block(block):
                states.append({"land": None})
            elif block.is_gla:
                states.append(
                    {
                        "gla_state": torch.zeros(batch_size, H, D, D, device=device),
                        "z_norm": torch.zeros(batch_size, H, D, device=device),
                    }
                )
            else:
                states.append(
                    {"buf": torch.zeros(batch_size, M, C, device=device), "pos": 0}
                )
        return states

    return init_state


def enable_cpu_delta(model) -> int:
    n = 0
    for block in model.blocks:
        if _is_delta_block(block):
            cpu_gdn = CPUGatedDeltaNet(block.mixer.gdn)
            cpu_swa = CPUSlidingWindow(block.mixer) if _has_swa_fused(block) else None
            step, forward_chunk = _make_block_methods(cpu_gdn, cpu_swa)
            block.step = types.MethodType(step, block)
            block.forward_chunk = types.MethodType(forward_chunk, block)
            block.mixer.forward = types.MethodType(
                _make_mixer_forward(cpu_gdn, cpu_swa, block.mixer.forward), block.mixer
            )
            block.mixer._cpu_gdn = cpu_gdn
            block.mixer._cpu_swa = cpu_swa
            n += 1
        elif _is_landmark_block(block):
            cpu_land = CPULandmark(block.mixer)
            step, forward_chunk = _make_landmark_block_methods(cpu_land)
            block.step = types.MethodType(step, block)
            block.forward_chunk = types.MethodType(forward_chunk, block)
            block.mixer._cpu_land = cpu_land
            if not hasattr(block.mixer, "prepare_inference"):
                block.mixer.prepare_inference = types.MethodType(
                    lambda self: None, block.mixer
                )
            n += 1
    model.init_state = types.MethodType(_make_init_state(model), model)
    return n