"""Loop-only LoRA for band layers: deltas active ONLY during band re-runs.

The initial capture forward runs the pristine base model (k=0 stays
bit-exact); BandLooper.band() re-runs flip LOOP_ACTIVE on, so the LoRA only
shapes the recurrence. Entrance-weighted per the rung-2 design note.
"""

import torch
import torch.nn as nn

_STATE = {"on": False}


class LoopLoRA(nn.Module):
    """Wraps an nn.Linear; adds scaled low-rank delta when loop is active."""

    def __init__(self, base: nn.Linear, rank=8, scale=1.0):
        super().__init__()
        self.base = base
        self.rank = rank
        self.scale = scale
        self.A = nn.Parameter(
            torch.randn(rank, base.in_features, dtype=torch.float32) * 0.01)
        self.B = nn.Parameter(
            torch.zeros(base.out_features, rank, dtype=torch.float32))

    def forward(self, x):
        out = self.base(x)
        if _STATE["on"]:
            delta = (x.float() @ self.A.T) @ self.B.T
            out = out + (self.scale * delta).to(out.dtype)
        return out


def loop_active(on: bool):
    _STATE["on"] = on


def inject_band_lora(tm, l0, layers_scales, rank=8, targets=("q_proj", "v_proj",
                                                             "down_proj")):
    """Wrap target Linears in given band layers. layers_scales: {layer: scale}.
    Returns list of LoRA params."""
    params = []
    for li, scale in layers_scales.items():
        layer = tm.layers[li]
        for name, mod in list(layer.named_modules()):
            leaf = name.split(".")[-1]
            if leaf in targets and isinstance(mod, nn.Linear):
                parent = layer
                parts = name.split(".")
                for p in parts[:-1]:
                    parent = getattr(parent, p)
                wrapped = LoopLoRA(mod, rank=rank, scale=scale)
                setattr(parent, parts[-1], wrapped)
                params += [wrapped.A, wrapped.B]
    return params


def entrance_faded_scales(l0, fade_to, full_until=None):
    """scale 1.0 at l0, linear fade to 0 by fade_to (exclusive)."""
    full_until = full_until if full_until is not None else l0
    out = {}
    for li in range(l0, fade_to):
        if li <= full_until:
            out[li] = 1.0
        else:
            out[li] = max(0.0, 1.0 - (li - full_until) / (fade_to - full_until))
    return {li: s for li, s in out.items() if s > 0}