"""36-state Viterbi selection over the Q2 symbolic lattice.

This is a direct port of Phase 3 of the C engine.  Each graph block owns a
``6 x 6`` coarse/fine lattice; candidate errors are reduced to a free-energy
per bin, the Sieve marginals become a log-prior bonus, and an exact DP finds
the minimum-cost state sequence with a Manhattan transition penalty.

Both the 576->36 bin reduction and the path->candidate mapping are fully
vectorised by reshaping the ``(24, 24)`` candidate grid into ``(6, 4, 6, 4)``.
"""

from __future__ import annotations

import numpy as np

from .config import HPCConfig
from .factored import CAND_TO_QUHIT, N_CAND_M, TOTAL_CANDIDATES

N_STATES = 36

_BIN_CAND_MAP = (
    np.arange(TOTAL_CANDIDATES, dtype=np.int64)
    .reshape(6, 4, 6, 4)
    .transpose(0, 2, 1, 3)
    .reshape(N_STATES, 16)
)


def _median_upper(values: np.ndarray) -> float:
    arr = np.asarray(values)
    if arr.size == 0:
        return 0.0
    k = arr.size // 2
    return float(np.partition(arr, k)[k])


def robust_temperature(
    candidate_errors: np.ndarray, cfg: HPCConfig
) -> float:
    n_blocks = candidate_errors.shape[0]
    if n_blocks <= 0:
        return 1.0e-10
    step = max(1, n_blocks // cfg.temperature_sample_blocks)
    medians = []
    for b in range(0, n_blocks, step):
        row = candidate_errors[b]
        finite = row[np.isfinite(row)]
        if finite.size == 0:
            continue
        medians.append(_median_upper(np.maximum(0.0, row - finite.min())))
        if len(medians) >= cfg.temperature_sample_blocks:
            break
    t = _median_upper(np.asarray(medians))
    return t if np.isfinite(t) and t > 1.0e-10 else 1.0e-10


def _bin_geometry():
    di = np.arange(TOTAL_CANDIDATES) // N_CAND_M
    mi = np.arange(TOTAL_CANDIDATES) % N_CAND_M
    qi_d = CAND_TO_QUHIT[di]
    qi_m = CAND_TO_QUHIT[mi]
    return qi_d, qi_m, qi_d * 6 + qi_m


def viterbi_select(
    candidate_errors: np.ndarray,
    coarse_marg: np.ndarray,
    fine_marg: np.ndarray,
    temperature: float,
    stride: int,
    graph_blocks: int,
    cfg: HPCConfig,
) -> tuple[np.ndarray, np.ndarray]:
    n_blocks = candidate_errors.shape[0]
    temp = max(float(temperature), 1.0e-12)

    # Reshape (n_blocks, 576) -> view (n_blocks, 6, 4, 6, 4) without copying
    ce_f32 = np.asarray(candidate_errors, dtype=np.float32)
    ce_view = ce_f32.reshape(n_blocks, 6, 4, 6, 4)
    blk_bin_min = np.min(ce_view, axis=(2, 4)).reshape(n_blocks, N_STATES).astype(np.float64)

    ce_bins = ce_view.transpose(0, 1, 3, 2, 4).reshape(n_blocks, N_STATES, 16)
    blk_bin_arg = np.argmin(ce_bins, axis=2)
    blk_bin_cand = _BIN_CAND_MAP[np.arange(N_STATES)[None, :], blk_bin_arg]
    blk_bin_cand = np.where(np.isfinite(blk_bin_min), blk_bin_cand, -1)

    pad = graph_blocks * stride - n_blocks
    if pad > 0:
        pad_arr = np.full((pad, 6, 4, 6, 4), np.inf, dtype=np.float32)
        grouped = np.concatenate([ce_view, pad_arr], axis=0).reshape(graph_blocks, stride, 6, 4, 6, 4)
    else:
        grouped = ce_view.reshape(graph_blocks, stride, 6, 4, 6, 4)

    group_min_4d = np.min(grouped, axis=(1, 3, 5))  # (graph_blocks, 6, 6)
    group_min = group_min_4d.reshape(graph_blocks, N_STATES).astype(np.float64)
    finite = np.isfinite(grouped)
    counts = np.sum(finite, axis=(1, 3, 5)).reshape(graph_blocks, N_STATES)
    valid = (counts > 0) & np.isfinite(group_min)
    gm_b = np.where( np.isfinite(group_min_4d), group_min_4d, 0.0)[:, None, :, None, :, None]
    shifted = np.where(finite, (grouped - gm_b) / np.float32(2.0 * temp), np.inf)
    sums = np.sum(np.exp(-shifted), axis=(1, 3, 5), dtype=np.float64).reshape(graph_blocks, N_STATES)
    mean_exp = np.maximum(sums / np.maximum(counts, 1), 1.0e-300)
    bin_err = np.where(valid, group_min - 2.0 * temp * np.log(mean_exp), np.inf)

    rep_idx = np.arange(graph_blocks) * stride
    bin_cand = blk_bin_cand[rep_idx]  # (graph_blocks, 36)

    c_tot = np.sum(coarse_marg, axis=1, keepdims=True)
    f_tot = np.sum(fine_marg, axis=1, keepdims=True)
    pc = np.where(c_tot > 1.0e-30, coarse_marg / np.maximum(c_tot, 1.0e-30), 1.0 / 6.0)
    pf = np.where(f_tot > 1.0e-30, fine_marg / np.maximum(f_tot, 1.0e-30), 1.0 / 6.0)
    log_prior = (
        np.log(pc[:, :, None] + 1.0e-30) + np.log(pf[:, None, :] + 1.0e-30)
    ).reshape(graph_blocks, N_STATES)

    finite_err = bin_err[np.isfinite(bin_err)]
    scale_err = float(finite_err.mean()) if finite_err.size else 1.0e-10
    scale_err = max(scale_err, 1.0e-20)

    local = np.where(
        np.isfinite(bin_err),
        bin_err - cfg.viterbi_beta * scale_err * log_prior,
        np.inf,
    )

    state_d = np.arange(N_STATES) // 6
    state_m = np.arange(N_STATES) % 6
    trans = (
        cfg.viterbi_alpha
        * scale_err
        * (np.abs(state_d[None, :] - state_d[:, None]) + np.abs(state_m[None, :] - state_m[:, None]))
    )

    dp = np.full((graph_blocks, N_STATES), np.inf, dtype=np.float64)
    back = np.zeros((graph_blocks, N_STATES), dtype=np.int64)
    dp[0] = local[0]
    s_range = np.arange(N_STATES)
    for gi in range(1, graph_blocks):
        totals = dp[gi - 1][:, None] + trans
        best_sp = np.argmin(totals, axis=0)
        best_val = totals[best_sp, s_range]
        dp[gi] = np.where(np.isfinite(best_val), best_val + local[gi], np.inf)
        back[gi] = best_sp

    if not np.any(np.isfinite(dp[-1])):
        path = np.zeros(graph_blocks, dtype=np.int64)
    else:
        path = np.zeros(graph_blocks, dtype=np.int64)
        path[-1] = int(np.argmin(dp[-1]))
        for gi in range(graph_blocks - 2, -1, -1):
            path[gi] = back[gi + 1, path[gi + 1]]

    # Vectorized path -> block candidate mapping:
    blk_gi = np.arange(n_blocks) // stride
    chosen_state = path[blk_gi]
    chosen_cand = blk_bin_cand[np.arange(n_blocks), chosen_state]
    global_best_cand = np.argmin(candidate_errors, axis=1).astype(np.int64)
    best_candidate = np.where(chosen_cand >= 0, chosen_cand, global_best_cand)

    cur_err = candidate_errors[np.arange(n_blocks), best_candidate]
    g_err = candidate_errors[np.arange(n_blocks), global_best_cand]
    best_candidate = np.where(
        g_err < cur_err * cfg.greedy_override_ratio, global_best_cand, best_candidate
    )

    return best_candidate, path