"""Baselines to compare the counter synapse against — in plain PyTorch.

Right now: a BitNet-b1.58-style ternary QAT linear (full FP32 master + Adam, ternary in the
forward) which isolates "cost of the 6-bit counter optimizer" vs "cost of ternarization".
The README roadmap lists the memory-efficient-training baselines worth adding next (8-bit
Adam / GaLore / LoMo), so the memory claim is measured against real competitors, not only
FP32+Adam.
"""
from __future__ import annotations

import math

import torch
import torch.nn as nn
import torch.nn.functional as F

__all__ = ["TernaryQATLinear", "make_linear"]


class TernaryQATLinear(nn.Module):
    """FP32 master weight, per-row absmean ternary in forward, straight-through gradient to
    the master; Adam updates the master. Same ternary inference as the counter layer, but a
    full FP32 optimizer -- so (counter vs this) isolates the 6-bit optimizer cost."""

    def __init__(self, fin: int, fout: int, init_gain: float = 1.0) -> None:
        super().__init__()
        self.weight = nn.Parameter(torch.empty(fout, fin))
        nn.init.normal_(self.weight, mean=0.0, std=init_gain * math.sqrt(1.0 / fin))

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        w = self.weight
        scale = w.abs().mean(dim=1, keepdim=True).clamp_min(1e-5)
        wq = (w / scale).round().clamp_(-1, 1) * scale
        wq = w + (wq - w).detach()  # straight-through estimator
        return F.linear(x, wq)


def make_linear(kind: str, fin: int, fout: int, init_gain: float = 1.0, **counter_kw):
    """Factory used by the GPT harness.
    kind in {dense, qat, counter, counter_rms, counter_packed, counter_triton, slowfast, group}.
    'slowfast' (M3) and 'group' (M2) are counter-linear drop-ins: SlowFastCounterLinear folds a
    low-rank fp residual into the base every K steps (cuts the base correlation 8-32x); the extra
    rank/merge_every knobs ride in counter_kw. GroupCounterLinear is a 2:4 structured-sparse counter
    linear (group/keep/hysteresis in counter_kw). Both replace any nn.Linear in the block."""
    from .counter import CompactCounterLinear, RMSCounterLinear
    from .packed import PackedRMSCounterLinear

    if kind == "counter":
        return CompactCounterLinear(fin, fout, init_gain=init_gain, **counter_kw)
    if kind == "counter_rms":
        return RMSCounterLinear(fin, fout, init_gain=init_gain, **counter_kw)
    if kind == "counter_packed":
        return PackedRMSCounterLinear(fin, fout, init_gain=init_gain, **counter_kw)
    if kind == "counter_triton":
        from .triton_counter import TritonCounterLinear
        return TritonCounterLinear(fin, fout, init_gain=init_gain, **counter_kw)
    if kind == "slowfast":                                  # M3
        from .slowfast import SlowFastCounterLinear
        return SlowFastCounterLinear(fin, fout, init_gain=init_gain, **counter_kw)
    if kind == "group":                                     # M2
        from .group_counter import GroupCounterLinear
        return GroupCounterLinear(fin, fout, init_gain=init_gain, **counter_kw)
    if kind == "qat":
        return TernaryQATLinear(fin, fout, init_gain)
    if kind == "dense":
        lin = nn.Linear(fin, fout, bias=False)
        nn.init.normal_(lin.weight, mean=0.0, std=init_gain * math.sqrt(1.0 / fin))
        return lin
    raise ValueError(f"unknown linear kind: {kind!r}")