# decoder.py
import numpy as np
from encoder import ALPHABET, G


def decode_blocks(Y: np.ndarray, C: np.ndarray) -> str:
    """
    Decode received samples by maximum correlation score (state unknown, per-block).
    """
    n = C.shape[1]
    half = n // 2
    num = Y.size // n
    C1 = C[:, :half]
    C2 = C[:, half:]
    sqrtG = np.sqrt(G)
    recovered = []
    for k in range(num):
        Yb = Y[k * n:(k + 1) * n]
        Ye, Yo = Yb[0::2], Yb[1::2]
        s1 = sqrtG * (Ye @ C1.T) + (Yo @ C2.T)
        s2 = (Ye @ C1.T) + sqrtG * (Yo @ C2.T)
        score = np.maximum(s1, s2)
        best = int(np.argmax(score))
        recovered.append(ALPHABET[best])
    return "".join(recovered)


def decode_blocks_with_state(Y: np.ndarray, C: np.ndarray) -> (str, int):
    """
    1) Estimate the single channel state s∈{1,2} by comparing total energy
       on even vs odd positions across the entire Y.
    2) Decode **all** blocks using that one state’s scoring rule.

    Returns (decoded_string, estimated_state).
    """
    n = C.shape[1]
    half = n // 2
    num = Y.size // n
    C1 = C[:, :half]
    C2 = C[:, half:]
    sqrtG = np.sqrt(G)

    # 1) state estimate from full-length energies
    Ye_all = Y[0::2]
    Yo_all = Y[1::2]
    E_even = np.sum(Ye_all ** 2)
    E_odd = np.sum(Yo_all ** 2)
    s_est = 1 if E_even > E_odd else 2

    recovered = []
    for k in range(num):
        Yb = Y[k * n:(k + 1) * n]
        Ye, Yo = Yb[0::2], Yb[1::2]
        if s_est == 1:
            score = sqrtG * (Ye @ C1.T) + (Yo @ C2.T)
        else:
            score = (Ye @ C1.T) + sqrtG * (Yo @ C2.T)
        best = int(np.argmax(score))
        recovered.append(ALPHABET[best])

    return "".join(recovered), s_est