"""GPU-accelerated syndrome decoder hooks.
Copyright 2026 Abir Maheshwari
"""
import numpy as np
import time
from typing import List, Optional

def decode_syndrome_gpu(
    syndrome: List[int],
    H: Optional[np.ndarray] = None,
    use_gpu: bool = False,
    max_iterations: int = 50,
    error_rate: float = 0.05
) -> List[int]:
    """
    Decode syndrome using a vectorized Belief Propagation decoder.
    Supports GPU acceleration via CuPy if use_gpu=True and CuPy is installed.
    """
    synd = np.array(syndrome, dtype=int)
    m = len(synd)
    
    if H is None:
        # Default fallback: construct a simple code where n = 2 * m
        # Systematic form H = [P^T | I_m]
        n = 2 * m
        P = np.zeros((m, m), dtype=int)
        for j in range(m):
            P[j, j] = 1
            P[(j + 1) % m, j] = 1
        H = np.zeros((m, n), dtype=int)
        H[:, :m] = P.T
        H[:, m:] = np.eye(m, dtype=int)
    else:
        H = np.array(H, dtype=int)
        
    m_shape, n_shape = H.shape
    if len(synd) != m_shape:
        raise ValueError(f"Syndrome length {len(synd)} does not match parity matrix rows {m_shape}")
        
    use_cp = False
    if use_gpu:
        try:
            import cupy as cp
            use_cp = True
        except ImportError:
            use_cp = False
            
    xp = cp if use_cp else np
    
    # Send to GPU/CPU arrays
    H_xp = xp.array(H, dtype=int)
    synd_xp = xp.array(synd, dtype=float)
    
    # Prior LLRs for each qubit
    llr = xp.full(n_shape, xp.log((1.0 - error_rate) / error_rate), dtype=float)
    
    # Variable-to-check messages
    v2c = xp.tile(llr, (m_shape, 1)) * H_xp
    c2v = xp.zeros((m_shape, n_shape), dtype=float)
    
    # Syndrome signs: syndrome=1 maps to -1, syndrome=0 maps to +1
    synd_sgn = xp.where(synd_xp == 1, -1.0, 1.0)
    
    for _ in range(max_iterations):
        # Check-to-Variable updates
        v2c_abs = xp.abs(v2c)
        v2c_abs = xp.clip(v2c_abs, 1e-12, 50.0)
        
        A = -xp.log(xp.tanh(v2c_abs / 2.0))
        A[H_xp == 0] = 0.0
        
        row_sums = xp.sum(A, axis=1, keepdims=True)
        A_ex = row_sums - A
        
        sgn = xp.sign(v2c)
        sgn[H_xp == 0] = 1.0
        row_sgn_prod = xp.prod(sgn, axis=1, keepdims=True) * synd_sgn[:, xp.newaxis]
        sgn_ex = row_sgn_prod * sgn
        
        A_ex = xp.clip(A_ex, 1e-12, 50.0)
        c2v = sgn_ex * 2.0 * xp.arctanh(xp.exp(-A_ex))
        c2v[H_xp == 0] = 0.0
        
        # Variable-to-Check updates
        col_sums = xp.sum(c2v, axis=0, keepdims=True)
        v2c = (llr + col_sums) - c2v
        v2c[H_xp == 0] = 0.0
        
        # Posterior and hard decision
        posterior = llr + xp.sum(c2v, axis=0)
        decoded = (posterior < 0).astype(int)
        
        # Check if syndrome is satisfied
        current_synd = (H_xp @ decoded) % 2
        if xp.all(current_synd == synd_xp.astype(int)):
            if use_cp:
                return decoded.get().tolist()
            return decoded.tolist()
            
    # Final decoding result
    posterior = llr + xp.sum(c2v, axis=0)
    decoded = (posterior < 0).astype(int)
    if use_cp:
        return decoded.get().tolist()
    return decoded.tolist()

