import numpy as np
import math
import matplotlib.pyplot as plt


def display_alphabet(alphabet, values=None, a_vals=False, title="Alphabet constellation diagram"):
    rect = polar2rect(alphabet)
    if values is not None:
        rect2 = polar2rect(values)
        plt.plot(rect2[:, 0], rect2[:, 1], 'r.')
    plt.plot(rect[:, 0], rect[:, 1], 'b.', markersize=12)
    plt.title(title)
    N = math.ceil(math.log2(len(alphabet)))
    if a_vals:
        for i, value in enumerate(rect):
            plt.annotate(xy=value+[0.01, 0.01], text=format(i, f'0{N}b'))
    plt.xlabel('Real')
    plt.ylabel('Imaginary')
    plt.grid()
    plt.show()


def bit_matrix2one_hot(matrix: np.ndarray) -> np.ndarray:
    """
    Returns a copy of bit encoded matrix to one hot matrix. A row examples:
    [1010] (decimal 10) => [0000 0100 0000 0000]
    [0011] (decimal  3) => [0000 0000 0000 1000]
    each number represents true/false value in column
    """
    N = matrix.shape[1]
    encoder = 2**np.arange(N)
    values = np.dot(matrix, encoder)
    result = np.zeros((matrix.shape[0], 2**N), dtype=bool)
    result[np.arange(matrix.shape[0]), values] = True
    return result


def one_hot2bit_matrix(matrix: np.ndarray) -> np.ndarray:
    """
    Returns a copy of one hot matrix to bit encoded matrix. A row examples:
    [0000 0100 0000 0000] => [1010] (decimal 10)
    [0000 0000 0000 1000] => [0011] (decimal  3)
    each number represents true/false value in column
    """
    N = math.ceil(math.log2(matrix.shape[1]))
    values = np.dot(matrix, np.arange(2**N))
    return int2bit_array(values, N)


def int2bit_array(int_arr: np.ndarray, N: int) -> np.ndarray:
    x0 = np.array([int_arr], dtype=np.uint8)
    x1 = np.unpackbits(x0.T, bitorder='little', axis=1)
    result = x1[:, :N].astype(bool)  # , indices
    return result


def polar2rect(array, amp_column=0, phase_column=1) -> np.ndarray:
    """
    Return copy of array with amp_column and phase_column as polar coordinates replaced by rectangular coordinates
    """
    if len(array) == 2 or len(array.shape) == 1:
        array = np.array([array])

    if array.shape[1] < 2:
        raise ValueError('Matrix has less than two columns')

    result = array.copy()
    result[:, amp_column] = np.cos(array[:, phase_column]) * array[:, amp_column]
    result[:, phase_column] = np.sin(array[:, phase_column]) * array[:, amp_column]
    return result


def rect2polar(array, x_column=0, y_column=1) -> np.ndarray:
    """
    Return copy of array with x_column and y_column as rectangular coordinates replaced by polar coordinates
    """
    if len(array) == 2 or len(array.shape) == 1:
        array = np.array([array])

    if array.shape[1] < 2:
        raise ValueError('Matrix has less than two columns')

    x_arr = array[:, x_column]
    y_arr = array[:, y_column]
    result = array.copy()
    result[:, x_column] = np.sqrt(x_arr**2 + y_arr**2)
    result[x_arr != 0, y_column] = np.arctan(y_arr[x_arr != 0, ] / x_arr[x_arr != 0, ])
    result[np.bitwise_and(x_arr == 0, y_arr < 0), y_column] = -np.pi / 2
    result[np.bitwise_and(x_arr == 0, y_arr == 0), y_column] = 0
    result[np.bitwise_and(x_arr == 0, y_arr > 0), y_column] = np.pi / 2
    result[np.bitwise_and(x_arr < 0, y_arr < 0), y_column] -= np.pi
    result[np.bitwise_and(x_arr < 0, y_arr >= 0), y_column] += np.pi
    return result


def generate_random_bit_array(size):
    z, o = np.zeros(int(size // 2), dtype=bool), np.ones(int(size // 2), dtype=bool)
    if size % 2 == 1:
        p = (z, o, [np.random.choice([True, False])])
    else:
        p = (z, o)
    arr = np.concatenate(p)
    np.random.shuffle(arr)
    return arr