"""Typed CLI configuration for train.py.

One adapter (lora2r: rank-2r Gaussian-init LoRA, A+B trainable, SGTM-style
three-way hard block masking; see src/vgrout/lora2r.py) and three arms:

  none    gate pinned clean (0,0): quarantine never trains -- the capacity- and
          structure-matched vanilla control.
  routeV  per-rollout three-way gate from the c-probe gradient vs v_grad.
  absorb  gate pinned mid (1,0): both blocks train on everything, no gate --
          isolates the value of the gate + hard masks vs absorption alone.
"""
from __future__ import annotations

from dataclasses import dataclass
from pathlib import Path
from typing import Literal

from .rewards import EnvMode


@dataclass(kw_only=True)
class Config:
    intervention: Literal["none", "routeV", "absorb"] = "routeV"
    lora_r: int = 32
    lora_init_seed: int = 0

    model: str = "Qwen/Qwen3-4B"
    steps: int = 100
    group: int = 6
    max_new: int = 1024
    n_problems: int = 992
    prompts_per_step: int = 8
    lr: float = 1e-4
    adam_beta1: float = 0.9
    adam_beta2: float = 0.99
    clip: float = 0.2
    # AdamW decay pulls raw A/B toward 0, not toward the init, which would drive
    # the net delta to -B0@A0 -- must stay 0 for this adapter.
    weight_decay: float = 0.0
    warmup_frac: float = 0.1
    grad_clip: float = 10.0
    seed: int = 41
    unbiased: bool = True

    vhack_refresh_every: int = 5
    vhack_pairs_path: Path = Path("out/pairsets/pairs_authored.json")
    # Haar-random direction control (placebo): same routing machinery, no pair signal.
    routeV_random_v_seed: int | None = None
    rollout_ablate_frac: float = 0.0

    env_mode: EnvMode = "run_tests"
    # Rotating per-step chance a problem is shown hint-free (graded honest-only);
    # keeps solve pressure alive. 10% learned solve too slowly; 25% default since 2026-06-10.
    unhackable_frac: float = 0.25
    teacher_pool_dir: Path | None = None
    mix_ratio: float = 0.125
    teacher_off_step: int | None = 30
    teacher_modes: tuple[str, ...] | None = None

    eval_ablate_every: int = 0
    eval_n_prompts: int = 32
    eval_batch_size: int = 2
    save_ckpt_every: int = 10
    out_tag: str = ""

    @property
    def preset_name(self) -> str:
        return type(self).__name__.removesuffix("Config").lower() or "base"

    @property
    def arm(self) -> str:
        # _lora2r suffix kept so these runs never conflate with the retired
        # PiSSA-substrate runs of the same intervention (rename-on-logic-change).
        return {"none": "vanilla_lora2r", "routeV": "routingV_lora2r",
                "absorb": "absorb_lora2r"}[self.intervention]


@dataclass(kw_only=True)
class SmokeConfig(Config):
    model: str = "llamafactory/tiny-random-qwen3"
    lora_r: int = 4   # tiny model min Linear dim is 16; 2r=8 fits
    steps: int = 30
    group: int = 4
    max_new: int = 32
    n_problems: int = 100
    prompts_per_step: int = 1


@dataclass(kw_only=True)
class FastConfig(Config):
    model: str = "Qwen/Qwen3-4B"
    steps: int = 100
    teacher_pool_dir: Path | None = Path("out/pools/teacher_pool_runtests_dense")
    group: int = 8
    max_new: int = 512
    n_problems: int = 200
    prompts_per_step: int = 4
    adam_beta1: float = 0.5
    adam_beta2: float = 0.9
