import numpy as np
from numpy import logaddexp
from utils import index_to_char
from codebook import construct_codebook

def decode_message(Y, codebook):
    """
    ML decoding for unknown channel state s:
    p(Y|i) = 0.5*p(Y|i,s=1) + 0.5*p(Y|i,s=2)
    We use log-sum-exp to combine both branch metrics.
    """
    G = 10
    Y1, Y2 = Y[::2], Y[1::2]
    best_idx = None
    best_metric = -np.inf
    for i, c in enumerate(codebook):
        # Only consider indices that map to characters
        if i not in index_to_char:
            continue
        c1, c2 = c[::2], c[1::2]
        # Branch metrics (up to additive constants)
        s1 = np.sqrt(G) * np.dot(Y1, c1) + np.dot(Y2, c2)
        s2 = np.dot(Y1, c1) + np.sqrt(G) * np.dot(Y2, c2)
        # Combine via log-sum-exp
        metric = logaddexp(s1, s2)
        if metric > best_metric:
            best_metric = metric
            best_idx = i
    return best_idx


def signal_to_text(Y, codebook, r=6):
    # Reconstruct codebook length (seg_len)
    _, seg_len, _, _ = construct_codebook(r, 1)
    text = ''
    for i in range(40):
        seg = Y[i * seg_len:(i + 1) * seg_len]
        idx = decode_message(seg, codebook)
        text += index_to_char.get(idx, '?')
    return text