from __future__ import annotations

import torch
import torch.nn as nn

from chimera.quantization import BitLinear


class MeZOOptimizer:
    """Memory-Efficient Zeroth-Order optimiser (Princeton MeZO)."""

    def __init__(
        self,
        model: nn.Module,
        lr: float = 1e-4,
        eps: float = 1e-3,
        weight_decay: float = 0.0,
        momentum: float = 0.0,
        direction: str = "rademacher",
    ):
        self.model = model
        self.lr = float(lr)
        self.eps = float(eps)
        self.wd = float(weight_decay)
        self.momentum = float(momentum)
        if direction not in ("rademacher", "gaussian"):
            raise ValueError(f"unknown direction: {direction!r}")
        self.direction = direction

        self._bitlinear_modules: list[tuple[str, BitLinear]] = []
        self._dense_params: list[tuple[str, torch.Tensor]] = []
        seen: set[int] = set()

        for name, module in model.named_modules():
            if isinstance(module, BitLinear):
                self._bitlinear_modules.append((name, module))
                seen.add(id(module.weight))
                if module.bias is not None:
                    seen.add(id(module.bias))

        for name, param in model.named_parameters():
            if param.requires_grad and id(param) not in seen:
                self._dense_params.append((name, param))
                seen.add(id(param))

        self._momentum: dict[int, torch.Tensor] = {}
        if self.momentum > 0:
            for _, param in self._dense_params:
                self._momentum[id(param)] = torch.zeros_like(param.data)
            for _, module in self._bitlinear_modules:
                self._momentum[id(module.weight)] = torch.zeros_like(module.weight.data)

        self._step_masks: dict[int, torch.Tensor] = {}

    def _direction(self, p: torch.Tensor, seed: int) -> torch.Tensor:
        gen = torch.Generator(device="cpu")
        gen.manual_seed(int(seed) & 0x7FFF_FFFF_FFFF_FFFF)
        if self.direction == "gaussian":
            return torch.randn(p.shape, dtype=p.dtype, device="cpu", generator=gen).to(p.device)
        z = torch.empty(p.shape, dtype=p.dtype, device="cpu")
        z.bernoulli_(0.5, generator=gen).mul_(2).sub_(1)
        return z.to(p.device)

    def _walk_params(self):
        offset = 0
        for _, module in self._bitlinear_modules:
            yield offset, module.weight.data, self._step_masks.get(id(module.weight))
            offset += 1
            if module.bias is not None:
                yield offset, module.bias.data, None
                offset += 1
        for _, param in self._dense_params:
            yield offset, param.data, None
            offset += 1

    def _perturb(self, base_seed: int, scale: float) -> None:
        for off, param, mask in self._walk_params():
            z = self._direction(param, base_seed + off * 1_000_003)
            if mask is not None:
                z = z * mask.to(dtype=z.dtype, device=z.device)
            param.add_(z, alpha=scale)
        for _, module in self._bitlinear_modules:
            module.invalidate_packed()

    def _update(self, base_seed: int, projected_grad: float) -> None:
        for off, param, mask in self._walk_params():
            z = self._direction(param, base_seed + off * 1_000_003)
            if mask is not None:
                z = z * mask.to(dtype=z.dtype, device=z.device)
            buf = self._momentum.get(id(param))
            if buf is not None:
                buf.mul_(self.momentum).add_(z, alpha=projected_grad)
                param.add_(buf, alpha=-self.lr)
            else:
                param.add_(z, alpha=-self.lr * projected_grad)
            if self.wd > 0:
                param.mul_(1 - self.lr * self.wd)
        for _, module in self._bitlinear_modules:
            module.invalidate_packed()

    @torch.no_grad()
    def step(self, loss_fn, batch) -> float:
        seed = int(torch.randint(0, 2**31, (1,)).item())
        self._step_masks = {id(m.weight): m.ternary_nonzero_mask().detach() for _, m in self._bitlinear_modules}
        self._perturb(seed, +self.eps)
        loss_pos = float(loss_fn(batch).item())
        self._perturb(seed, -2.0 * self.eps)
        loss_neg = float(loss_fn(batch).item())
        self._perturb(seed, +self.eps)
        projected_grad = (loss_pos - loss_neg) / (2.0 * self.eps)
        self._update(seed, projected_grad)
        self._step_masks = {}
        return 0.5 * (loss_pos + loss_neg)