from functools import cache
from typing import Tuple
import numpy as np
import const_utils
import matplotlib.pyplot as plt

class ConstellationPoints():

    def __init__(self, length:int=None, constellation_dict:dict=None, radius:Tuple[int, float]=1):
        self._radius = radius
        if constellation_dict:
            self._constellation = constellation_dict
        elif length:
            self._length = length


    def _generate_from_length(self, length):
        const_utils.validate_intpow2(length, 'length')
        object.__setattr__(self, '_length', length)
        object.__setattr__(self, '_n', int(self._length/2+0.5))
        object.__setattr__(self, '_m', self._length // 2)
        object.__setattr__(self, '_constellation', const_utils.bidict(generate_rectangular_constellation(self._length)))


    def _generate_from_dict(self, constellation_dict):
        const_utils.validate_coords(constellation_dict)
        const_utils.validate_int(next(iter(constellation_dict.keys())), 'labels')

        # check if constellation is a one-to-one mapping
        if len(set(constellation_dict.values())) != len(constellation_dict):
            raise ValueError('constellation must be a one-to-one mapping of labels and coordinates')

        object.__setattr__(self, '_constellation', const_utils.bidict(constellation_dict))

        object.__setattr__(self, '_length', len(constellation_dict))
        object.__setattr__(self, '_n', int(self._length/2+0.5))
        object.__setattr__(self, '_m', self._length // 2)


    def plot(self):
        # STUB ConstellationPots.plot()
        raise NotImplementedError()


    def __len__(self):
        return self._length


    @property
    def symbols(self):
        return dict(self._constellation).values()


    @property
    def labels(self):
        return dict(self._constellation.inverse).values()


    @property
    def constellation(self):
        return dict(self._constellation)


    @property
    def lookup(self):
        return dict(self._constellation.inverse)


    def get_symbol(self, label):
        return self._constellation[label]


    def get_label(self, symbol):
        return self._constellation.inverse[symbol]


    def __setattr__(self, name: str, value) -> None:
        if name == '_constellation':
            self._generate_from_dict(value)
        elif name == '_length':
            self._generate_from_length(value)
        elif name == '_radius':
            self._set_radius(value)
        else:
            object.__setattr__(self, name, value)


    def _set_radius(self, value):
        if isinstance(value, (int, float)):
            object.__setattr__(self, '_radius',value)
        else:
            raise TypeError('radius must be an integer or a float')


@cache
def gray_1d(k: int, label: int) -> int:
    const_utils.gray_1d_input_validation(k, label)

    if k == 1:
        return 1 if label==0 else -1

    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: int, m: int, label: int) -> tuple:
    const_utils.gray_2d_input_validation(n, m, label)

    if m == 0:
        return (gray_1d(n, label), 0) # it's a 1d case in disguise!
    
    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)

    return nearest_symbols


def gray_penalty(constellation):
    raise NotImplementedError
    # 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 = constellation['meta']['len']
    inverted_constellation = {tuple(symbol):label for label,symbol in constellation.items() if label != 'meta'} # -> invert constellation dict
    syms = [symbol for _, symbol in constellation.items()]

    const_utils.validate_intpow2()

    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 generate_rectangular_constellation(length: int):
    # const_utils.validate_int(length, 'length')
    const_utils.validate_intpow2(length, 'length')
    lengthexp = int(np.log2(length))
    n = int(lengthexp/2+0.5) # ceil
    m = lengthexp // 2  # floor
    return {label:gray_2d(n, m, label) for label in range(length)}


def transform_rectangular_mapping(constellation: ConstellationPoints):
    n, m = constellation['meta'] # TODO def generate_rectangular_constellation(n)
    # r, c = 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 r == 1: # 1D-constellation
    #     return constellation

    # # get n and m for one quadrant
    # n = c/2
    # m = r/2

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

    # [ ] set transformed flag in constellation?

    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: # REVIEW m+2 correct? p. 7
        new_const = {}
        s = 2**(n-1)
        for label, symbol in constellation.items():
            # STUB transfrom_rectangular_mapping(constellation) -> generalized non-square-QAM
            raise NotImplementedError()
    
    else:
        # TODO define what should happen here
        return constellation

    # 2^(2n+1)
    # for 32-QAM: 2^5 -> n = 2
    # rectangular grid of 4*8 -> 2*4 per quadrant

    # for 128-QAM: 2^7 -> n = 3
    # rectangular grid of 


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__':
    const0 = ConstellationPoints()
    const128 = ConstellationPoints(length=128)
    const_ext = ConstellationPoints(constellation_dict=generate_rectangular_constellation(64))

    print(vars(const0))
    print(vars(const128))
    print(vars(const_ext))

    # constellation_128 = {label:gray_2d(3, 4, label) for label in range(128)}