"""Configuration for the HPC-Quantize V2 Q2_K engine.

Every tunable that was a compile-time macro, environment variable or hidden
global in the C engine is collected here as an explicit, validated field.
The defaults are the C engine's production defaults unless marked
``V2`` in the docstring, in which case they are deliberate improvements.
"""

from __future__ import annotations

from dataclasses import dataclass, field


@dataclass(slots=True)
class HPCConfig:
    # ---- block geometry -------------------------------------------------
    qk_k: int = 256

    # ---- candidate field ------------------------------------------------
    # V2 improvement: the default is the 6x6 symbolic lattice (36 exact
    # factored solves) instead of the 24x24 physical grid (576 solves).
    # Set to 24/24 to reproduce the full C candidate field.
    n_cand_d: int = 6
    n_cand_m: int = 6
    factor_pin: int = 3
    factor_iters: int = 1
    factor_1d_iters: int = 6
    factor_1d_refine_sweeps: int = 1
    factor_standalone_iters: int = 8
    factor_starts: int = 2
    factor_ulp_radius: int = 1
    factor_refine_1d: bool = True
    factor_portfolio: bool = True
    factor_spectral_refine: bool = True
    factor_exact_emit: bool = True
    use_weight_analogs: bool = True
    factor_coupled_q2k: bool = True
    factor_joint_subscales: bool = True

    # ---- objective ------------------------------------------------------
    use_spectral: bool = True
    dc_lambda: float = 1.0
    vw_lambda: float = 1.0
    fold_depth: int = 1
    fold_gamma: float = 2.0
    dc_decay: float = 0.85
    carry_cumulative: bool = True
    sse_budget: float = 5.0e-4
    greedy_override_ratio: float = 0.995

    # ---- sieve ----------------------------------------------------------
    same_penalty: float = 0.85
    log_slack: float = 1.5
    quhits_per_block: int = 2
    graph_max_blocks: int = 8192

    # ---- Viterbi --------------------------------------------------------
    optimizer: str = "hpc"
    viterbi_beta: float = 0.25
    viterbi_alpha: float = 0.08
    temperature_sample_blocks: int = 128

    # ---- assembly -------------------------------------------------------
    refine_passes: int = 3
    ulp_radius: int = 4
    subblock_states: int = 6
    subblock_sieve: bool = True
    use_factor_floor: bool = True
    use_candidate_floor: bool = True
    apply_dc_carry: bool = True
    quant_format: str = "mixed"
    custom_quant_map: dict[str, str] = field(default_factory=dict)
    enable_q4_fallback: bool = True
    quantize_embd_q4: bool = True
    promote_attn_q4: bool = False

    # ---- thermal refinement (pair heat-bath annealer, opt-in) ------------
    # V2-only knob (0 = off, bit-identical to the deterministic pipeline).
    # Runs a seeded checkerboard Gibbs sampler over the vesica-pair lattice
    # after Phase 4; late passes are monotone so output never regresses.
    therm_passes: int = 0
    therm_sweeps: int = 2
    therm_t0: float = 0.25
    therm_t_decay: float = 0.35
    therm_seed: int = 271828
    # Zero-residual anneal: multiplier >= 1 applied to the spectral lambdas
    # (dc_lambda, vw_lambda) while annealing.  At large tilts (~1e2-1e3) the
    # exact-conditional Gibbs chain spends weighted SSE to drain the *coherent*
    # residual (row DC + vesica pairs) toward the lattice floor, leaving an
    # (almost) spectrally white error.  1.0 = off.
    therm_lambda_tilt: float = 1.0
    # Zero-DC pin (the "Q2_K free lunch"): after assembly (and any thermal
    # refinement) drive each block's residual DC toward 0.  Q2_K re-solves
    # (d, m) on the exact dc=0 constraint line — a joint least-squares
    # translation whose closed form *lowers* SSE (free).  Q4_0 / Q3_K carry
    # a single scale dof, so the pin there is a deterministic code-flip
    # exchange: it spends at most ``pin_budget`` of each block's SSE to buy
    # the coherent-residual drop (Q3_K lands at ~-85% DC RMS for ~0.9% SSE;
    # Q4_0 is bounded by the narrower [0, 15] grid's DC slack).
    pin_dc: bool = False

    # Fraction of per-block SSE the Q4_0 / Q3_K code-flip pin may spend
    # (the Q2_K closed-form pin never spends any).  0 disables the expense.
    pin_budget: float = 0.05

    # ---- seed clamp -----------------------------------------------------
    clamp_kurtosis_off: float = 6.0
    clamp_kurtosis_soft: float = 20.0
    clamp_sigma_soft: float = 6.0
    clamp_sigma_hard: float = 4.0
    temper_imatrix: bool = True

    # ---- reporting & parallelism ----------------------------------------
    device: str = "auto"
    workers: int = 0
    verbose: bool = False

    def validate(self) -> "HPCConfig":
        if self.qk_k != 256:
            raise ValueError("Q2_K requires qk_k == 256")
        if not (1 <= self.n_cand_d <= 24 and 1 <= self.n_cand_m <= 24):
            raise ValueError("candidate counts must lie in [1, 24]")
        if self.optimizer not in ("hpc", "mse"):
            raise ValueError("optimizer must be 'hpc' or 'mse'")
        if not 0 <= self.factor_pin <= 3:
            raise ValueError("factor_pin must be a 2-bit mask")
        if self.fold_depth < 0:
            raise ValueError("fold_depth must be non-negative")
        if not 0.0 <= self.dc_decay <= 1.0:
            raise ValueError("dc_decay must lie in [0, 1]")
        if self.ulp_radius < 0 or self.factor_ulp_radius < 0:
            raise ValueError("ULP radii must be non-negative")
        if self.refine_passes < 1:
            raise ValueError("refine_passes must be >= 1")
        if not 1 <= self.subblock_states <= 6:
            raise ValueError("subblock_states must lie in [1, 6]")
        if self.factor_1d_iters < 1 or self.factor_standalone_iters < 1:
            raise ValueError("factor iteration budgets must be >= 1")
        if self.factor_starts < 1:
            raise ValueError("factor_starts must be >= 1")
        if self.sse_budget < 0.0:
            raise ValueError("sse_budget must be non-negative")
        if self.pin_budget < 0.0:
            raise ValueError("pin_budget must be non-negative")
        if not 0.0 < self.same_penalty <= 1.0:
            raise ValueError("same_penalty must lie in (0, 1]")
        if self.log_slack < 0.0:
            raise ValueError("log_slack must be non-negative")
        if self.viterbi_alpha < 0.0 or self.viterbi_beta < 0.0:
            raise ValueError("Viterbi weights must be non-negative")
        if self.workers < 0:
            raise ValueError("workers must be >= 0 (0 = auto CPU count)")
        if self.therm_passes < 0:
            raise ValueError("therm_passes must be >= 0 (0 disables thermal refinement)")
        if self.therm_sweeps < 1:
            raise ValueError("therm_sweeps must be >= 1")
        if self.therm_t0 <= 0.0:
            raise ValueError("therm_t0 must be positive")
        if not 0.0 < self.therm_t_decay < 1.0:
            raise ValueError("therm_t_decay must lie in (0, 1)")
        if self.therm_lambda_tilt < 1.0:
            raise ValueError("therm_lambda_tilt must be >= 1.0 (1.0 = off)")
        if self.device not in ("auto", "cuda", "cpu") and not self.device.startswith("cuda:"):
            raise ValueError("device must be 'auto', 'cuda', 'cuda:N', or 'cpu'")
        if self.quant_format not in ("mixed", "mixed_q1", "custom", "q1_0", "q3_k", "q4_0", "q2_k"):
            raise ValueError(
                "quant_format must be 'mixed', 'mixed_q1', 'custom', 'q1_0', 'q3_k', 'q4_0', or 'q2_k'"
            )
        allowed_custom = {"q1_0", "q2_k", "q3_k", "q4_0", "copy", "auto"}
        for tensor_type, quant in self.custom_quant_map.items():
            if str(quant).strip().lower() not in allowed_custom:
                raise ValueError(
                    f"custom_quant_map[{tensor_type!r}] must be one of "
                    f"Q1_0, Q2_K, Q3_K, Q4_0, Copy, Auto, got {quant!r}"
                )
        return self