from functools import cache
import numpy as np
import const_utils

# https://www.ieee802.org/3/bn/public/nov13/prodan_3bn_02_1113.pdf

@cache
def gray_1d(k, label):
    const_utils.gray_1d_input_validation(k, label)
    # special case
    if k == 1:
        return 1 if label==0 else -1

    # all other cases -> recurse
    b0, new_symbol = const_utils.next_symbol(k, label)
    return (1-2*b0)*(2**(k-1)+gray_1d(k-1, new_symbol))


def gray_2d(n, m, label):
    const_utils.gray_2d_input_validation(n, m, label)
    # n or m is 0
    if (coord:=const_utils.gray_2d_handle_1d(n, m, label)) is not None:
        return coord
    
    # all other cases
    symbol_i, symbol_q = const_utils.split_symbol(n, m, label)
    return (gray_1d(n, symbol_i), gray_1d(m, symbol_q))

def hamming_dist(a, b):
    if not isinstance(a, int):
        raise ValueError('a must be an integer')
    if not isinstance(b, int):
        raise ValueError('b must be an integer')

def euclidean_distance(coord1, coord2):
    if isinstance(coord1, int):
        return abs(coord1 - coord2)
    return np.sqrt((coord1[0] - coord2[0])**2 + (coord1[1] - coord2[1])**2)

def find_nearest(coord, coords):
    min_distance = float('inf')
    nearest_symbols = []

    for c in coords:
        dist = euclidean_distance(coord, c)
        if dist == 0:
            continue
        elif dist < min_distance:
            min_distance = dist
            nearest_symbols = [c]
        elif dist == min_distance:
            nearest_symbols.append(c)
        # else:
        #     pass

    return nearest_symbols

def gray_penalty(constellation):
    # constellation: {label_0:coordinate_0, label_1:coordinate_1, .., label_2^n-1:coordinate_2^n-1}

    # 2^n-QAM -> 2^n symbols S_i, where i=0,1,..2^n-1, ex. S_0 = (-3,-3) or S_0 = -2
    # N(S_i): set of (euclidean) nearest symbols S_j -> N((-3,-3)) = {(-3,-2), (-3,-4), (-2,-3), (-4,-3)}
    # |N(S_i)|: size of set N(S_i)
    # l(S): label given by mapping -> inverse of gray_Qd -> generate all symbols/labels for given constellation
    # wt(l_1, l_2), hamming distance btw. two labels

    t = len(constellation)
    inverted_constellation = {tuple(symbol):label for label,symbol in constellation.items() if label != 'meta'} # -> invert constellation dict
    syms = [symbol for _, symbol in constellation.items()]

    if (n:=np.log2(t)) != int(n):
        raise ValueError('only constellations with 2^n points supported')

    G = 0
    for li, si in constellation.items():
        N = find_nearest(si, syms)
        size_N = len(N)
        wt = sum(hamming_dist(inverted_constellation[tuple(sj)], li) for sj in N)
        G += wt/size_N
    G /= t

    return G

def find_rows_columns(coordinates):
    if not coordinates:
        return 0, 0

    min_row = min(coord[0] for coord in coordinates.values())
    max_row = max(coord[0] for coord in coordinates.values())
    row_spacing = abs(coordinates[next(iter(coordinates))][0] - coordinates[next(iter(coordinates))][0])

    min_col = min(coord[1] for coord in coordinates.values())
    max_col = max(coord[1] for coord in coordinates.values())
    col_spacing = abs(coordinates[next(iter(coordinates))][1] - coordinates[next(iter(coordinates))][1])

    num_rows = (max_row - min_row) // row_spacing + 1
    num_cols = (max_col - min_col) // col_spacing + 1

    return num_rows, num_cols

def transform_rectangular_mapping(constellation):
    n, m = find_rows_columns(constellation)

    # example: 32-qam -> 2^(2n+1) -> n = 2

    two_n1 = np.log2(len(constellation))
    if int(two_n1) != two_n1:
        raise ValueError('only constellations with 2^m points allowed')

    if n == 1 or m == 1: # 1D-constellation
        return constellation

    n = c/2
    m = r/2

    const_utils._validate_integer(n, 'n')
    const_utils._validate_integer(m, 'm')

    if n == m:  # square 2^(2n)-QAM
        return constellation

    if n == 2 and m == 1: # rectangular 8-QAM (4*2)
        return transform_8QAM(constellation)
    elif n == m+2:
        new_const = {}
        s = 2**(n-1)
        for label, symbol in constellation.items():



def transform_8QAM(constellation):
    new_const = {}
    for label, symbol in constellation.items():
        if symbol[0] < 3:
            new_const[label] = symbol
        else:
            i_rct, q_rct = symbol
            i_cr = -np.sign(i_rct)*(4-np.abs(i_rct))
            q_cr = np.sign(q_rct)*(np.abs(q_rct)+2)
            new_const[label] = [i_cr, q_cr]
    return new_const# rectangular 2^(m+n)-QAM



if __name__ == '__main__':
    # print(gray_1d(2, 0))
    print(gray_2d(2, 3, 4))
    print(gray_2d(0, 2, 4))