import numpy as np
import scipy.io as sio
import os

# ------------------ LOAD LUT ------------------

def load_LUT(mat_filename=None):
    """
    Loads the list of Huffman Codebooks (LUTs)

    Returns:
        huffLUT : list (index 1..11 used, index 0 unused)
    """
    if mat_filename is None:
        current_dir = os.path.dirname(os.path.abspath(__file__))
        mat_filename = os.path.join(current_dir, "huffCodebooks.mat")

    mat = sio.loadmat(mat_filename)


    huffCodebooks_raw = mat['huffCodebooks'].squeeze()

    huffCodebooks = []
    for i in range(11):
        huffCodebooks.append(np.array(huffCodebooks_raw[i]))

    # Build inverse VLC tables
    invTable = [None] * 11

    for i in range(11):
        h = huffCodebooks[i][:, 2].astype(int)       # column 3
        hlength = huffCodebooks[i][:, 1].astype(int) # column 2

        hbin = []
        for j in range(len(h)):
            hbin.append(format(h[j], f'0{hlength[j]}b'))

        invTable[i] = vlc_table(hbin)

    # Build Huffman LUT dicts
    huffLUT = [None] * 12  # index 0 unused
    params = [
        (4, 1, True),
        (4, 1, True),
        (4, 2, False),
        (4, 2, False),
        (2, 4, True),
        (2, 4, True),
        (2, 7, False),
        (2, 7, False),
        (2, 12, False),
        (2, 12, False),
        (2, 16, False),
    ]

    for i, (nTupleSize, maxAbs, signed) in enumerate(params, start=1):
        huffLUT[i] = {
            'LUT': huffCodebooks[i-1],
            'invTable': invTable[i-1],
            'codebook': i,
            'nTupleSize': nTupleSize,
            'maxAbsCodeVal': maxAbs,
            'signedValues': signed
        }

    return huffLUT

def vlc_table(code_array):
    """
    codeArray: list of strings, each string is a Huffman codeword (e.g. '0101')
    returns:
        h : NumPy array of shape (num_nodes, 3)
            columns:
              [ next_if_0 , next_if_1 , symbol_index ]
    """
    h = np.zeros((1, 3), dtype=int)

    for code_index, code in enumerate(code_array, start=1):
        word = [int(bit) for bit in code]
        h_index = 0

        for bit in word:
            k = bit
            next_node = h[h_index, k]
            if next_node == 0:
                h = np.vstack([h, [0, 0, 0]])
                new_index = h.shape[0] - 1
                h[h_index, k] = new_index
                h_index = new_index
            else:
                h_index = next_node

        h[h_index, 2] = code_index

    return h

# ------------------ ENCODE ------------------

def encode_huff(coeff_sec, huff_LUT_list, force_codebook = None):
    """
    Huffman-encode a sequence of quantized coefficients.

    This function selects the appropriate Huffman codebook based on the
    maximum absolute value of the input coefficients, encodes the coefficients
    into a binary Huffman bitstream, and returns both the bitstream and the
    selected codebook index.

    This is the Python equivalent of the MATLAB `encodeHuff.m` function used
    in audio/image coding (e.g., scale factor band encoding). The input
    coefficient sequence is grouped into fixed-size tuples as defined by
    the chosen Huffman LUT. Zero-padding may be applied internally.

    Parameters
    ----------
    coeff_sec : array_like of int
        1-D array of quantized integer coefficients to encode.
        Typically corresponds to a "section" or scale-factor band.

    huff_LUT_list : list
        List of Huffman lookup-table dictionaries as returned by `loadLUT()`.
        Index 1..11 correspond to valid Huffman codebooks.
        Index 0 is unused.

    Returns
    -------
    huffSec : str
        Huffman-encoded bitstream represented as a string of '0' and '1'
        characters.

    huffCodebook : int
        Index (1..11) of the Huffman codebook used for encoding.
        A value of 0 indicates a special all-zero section.
    """
    if force_codebook is not None:
            return huff_LUT_code_1(huff_LUT_list[force_codebook], coeff_sec)
    
    maxAbsVal = np.max(np.abs(coeff_sec))

    if maxAbsVal == 0:
        huffCodebook = 0
        huffSec = huff_LUT_code_0()

    elif maxAbsVal == 1:
        candidates = [1, 2]
        huffSec1 = huff_LUT_code_1(huff_LUT_list[candidates[0]], coeff_sec)
        huffSec2 = huff_LUT_code_1(huff_LUT_list[candidates[1]], coeff_sec)
        if len(huffSec1) <= len(huffSec2):
            huffSec = huffSec1
            huffCodebook = candidates[0]
        else:
            huffSec = huffSec2
            huffCodebook = candidates[1]

    elif maxAbsVal == 2:
        candidates = [3, 4]
        huffSec1 = huff_LUT_code_1(huff_LUT_list[candidates[0]], coeff_sec)
        huffSec2 = huff_LUT_code_1(huff_LUT_list[candidates[1]], coeff_sec)
        if len(huffSec1) <= len(huffSec2):
            huffSec = huffSec1
            huffCodebook = candidates[0]
        else:
            huffSec = huffSec2
            huffCodebook = candidates[1]

    elif maxAbsVal in (3, 4):
        candidates = [5, 6]
        huffSec1 = huff_LUT_code_1(huff_LUT_list[candidates[0]], coeff_sec)
        huffSec2 = huff_LUT_code_1(huff_LUT_list[candidates[1]], coeff_sec)
        if len(huffSec1) <= len(huffSec2):
            huffSec = huffSec1
            huffCodebook = candidates[0]
        else:
            huffSec = huffSec2
            huffCodebook = candidates[1]

    elif maxAbsVal in (5, 6, 7):
        candidates = [7, 8]
        huffSec1 = huff_LUT_code_1(huff_LUT_list[candidates[0]], coeff_sec)
        huffSec2 = huff_LUT_code_1(huff_LUT_list[candidates[1]], coeff_sec)
        if len(huffSec1) <= len(huffSec2):
            huffSec = huffSec1
            huffCodebook = candidates[0]
        else:
            huffSec = huffSec2
            huffCodebook = candidates[1]

    elif maxAbsVal in (8, 9, 10, 11, 12):
        candidates = [9, 10]
        huffSec1 = huff_LUT_code_1(huff_LUT_list[candidates[0]], coeff_sec)
        huffSec2 = huff_LUT_code_1(huff_LUT_list[candidates[1]], coeff_sec)
        if len(huffSec1) <= len(huffSec2):
            huffSec = huffSec1
            huffCodebook = candidates[0]
        else:
            huffSec = huffSec2
            huffCodebook = candidates[1]

    elif maxAbsVal in (13, 14, 15):
        huffCodebook = 11
        huffSec = huff_LUT_code_1(huff_LUT_list[huffCodebook], coeff_sec)

    else:
        huffCodebook = 11
        huffSec = huff_LUT_code_ESC(huff_LUT_list[huffCodebook], coeff_sec)

    return huffSec, huffCodebook

def huff_LUT_code_1(huff_LUT, coeff_sec):
    LUT = huff_LUT['LUT']
    nTupleSize = huff_LUT['nTupleSize']
    maxAbsCodeVal = huff_LUT['maxAbsCodeVal']
    signedValues = huff_LUT['signedValues']

    numTuples = int(np.ceil(len(coeff_sec) / nTupleSize))

    if signedValues:
        coeff = coeff_sec + maxAbsCodeVal
        base = 2 * maxAbsCodeVal + 1
    else:
        coeff = coeff_sec
        base = maxAbsCodeVal + 1

    coeffPad = np.zeros(numTuples * nTupleSize, dtype=int)
    coeffPad[:len(coeff)] = coeff

    huffSec = []

    powers = base ** np.arange(nTupleSize - 1, -1, -1)

    for i in range(numTuples):
        nTuple = coeffPad[i*nTupleSize:(i+1)*nTupleSize]
        huffIndex = int(np.abs(nTuple) @ powers)

        hexVal = LUT[huffIndex, 2]
        huffLen = LUT[huffIndex, 1]

        bits = format(int(hexVal), f'0{int(huffLen)}b')

        if signedValues:
            huffSec.append(bits)
        else:
            signBits = ''.join('1' if v < 0 else '0' for v in nTuple)
            huffSec.append(bits + signBits)

    return ''.join(huffSec)

def huff_LUT_code_0():
    return ''

def huff_LUT_code_ESC(huff_LUT, coeff_sec):
    LUT = huff_LUT['LUT']
    nTupleSize = huff_LUT['nTupleSize']
    maxAbsCodeVal = huff_LUT['maxAbsCodeVal']

    numTuples = int(np.ceil(len(coeff_sec) / nTupleSize))
    base = maxAbsCodeVal + 1

    coeffPad = np.zeros(numTuples * nTupleSize, dtype=int)
    coeffPad[:len(coeff_sec)] = coeff_sec

    huffSec = []
    powers = base ** np.arange(nTupleSize - 1, -1, -1)

    for i in range(numTuples):
        nTuple = coeffPad[i*nTupleSize:(i+1)*nTupleSize]

        lnTuple = nTuple.astype(float)
        lnTuple[lnTuple == 0] = np.finfo(float).eps

        N4 = np.maximum(0, np.floor(np.log2(np.abs(lnTuple))).astype(int))
        N = np.maximum(0, N4 - 4)
        esc = np.abs(nTuple) > 15

        nTupleESC = nTuple.copy()
        nTupleESC[esc] = np.sign(nTupleESC[esc]) * 16

        huffIndex = int(np.abs(nTupleESC) @ powers)

        hexVal = LUT[huffIndex, 2]
        huffLen = LUT[huffIndex, 1]

        bits = format(int(hexVal), f'0{int(huffLen)}b')

        escSeq = ''
        for k in range(nTupleSize):
            if esc[k]:
                escSeq += '1' * N[k]
                escSeq += '0'
                escSeq += format(abs(nTuple[k]) - (1 << N4[k]), f'0{N4[k]}b')

        signBits = ''.join('1' if v < 0 else '0' for v in nTuple)
        huffSec.append(bits + signBits + escSeq)

    return ''.join(huffSec)

# ------------------ DECODE ------------------

def decode_huff(huff_sec, huff_LUT):
    """
    Decode a Huffman-encoded stream.

    Parameters
    ----------
    huff_sec : array-like of int or str
        Huffman encoded stream as a sequence of 0 and 1 (string or list/array).
    huff_LUT : dict
        Huffman lookup table with keys:
            - 'invTable': inverse table (numpy array)
            - 'codebook': codebook number
            - 'nTupleSize': tuple size
            - 'maxAbsCodeVal': maximum absolute code value
            - 'signedValues': True/False

    Returns
    -------
    decCoeffs : list of int
        Decoded quantized coefficients.
    """

    h = huff_LUT['invTable']
    huffCodebook = huff_LUT['codebook']
    nTupleSize = huff_LUT['nTupleSize']
    maxAbsCodeVal = huff_LUT['maxAbsCodeVal']
    signedValues = huff_LUT['signedValues']

    # Convert string to array of ints
    if isinstance(huff_sec, str):
        huff_sec = np.array([int(b) for b in huff_sec])

    eos = False
    decCoeffs = []
    streamIndex = 0

    while not eos:
        wordbit = 0
        r = 0  # start at root

        # Decode Huffman word using inverse table
        while True:
            b = huff_sec[streamIndex + wordbit]
            wordbit += 1
            rOld = r
            r = h[rOld, b]
            if h[r, 0] == 0 and h[r, 1] == 0:
                symbolIndex = h[r, 2] - 1  # zero-based
                streamIndex += wordbit
                break

        # Decode n-tuple magnitudes
        if signedValues:
            base = 2 * maxAbsCodeVal + 1
            nTupleDec = []
            tmp = symbolIndex
            for p in reversed(range(nTupleSize)):
                val = tmp // (base ** p)
                nTupleDec.append(val - maxAbsCodeVal)
                tmp = tmp % (base ** p)
            nTupleDec = np.array(nTupleDec)
        else:
            base = maxAbsCodeVal + 1
            nTupleDec = []
            tmp = symbolIndex
            for p in reversed(range(nTupleSize)):
                val = tmp // (base ** p)
                nTupleDec.append(val)
                tmp = tmp % (base ** p)
            nTupleDec = np.array(nTupleDec)

            # Apply sign bits
            nTupleSignBits = huff_sec[streamIndex:streamIndex + nTupleSize]
            nTupleSign = -(np.sign(nTupleSignBits - 0.5))
            streamIndex += nTupleSize
            nTupleDec = nTupleDec * nTupleSign

        # Handle escape sequences
        escIndex = np.where(np.abs(nTupleDec) == 16)[0]
        if huffCodebook == 11 and escIndex.size > 0:
            for idx in escIndex:
                N = 0
                b = huff_sec[streamIndex]
                while b:
                    N += 1
                    b = huff_sec[streamIndex + N]
                streamIndex += N +1
                N4 = N + 4
                escape_word = huff_sec[streamIndex:streamIndex + N4]
                escape_value = 2 ** N4 + int("".join(map(str, escape_word)), 2)
                nTupleDec[idx] = escape_value
                streamIndex += N4
            # Apply signs again
            nTupleDec[escIndex] *= nTupleSign[escIndex]

        decCoeffs.extend(nTupleDec.tolist())

        if streamIndex >= len(huff_sec):
            eos = True

    return decCoeffs