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], s=format(i, f'0{N}b'))
    plt.xlabel('Real')
    plt.ylabel('Imaginary')
    plt.grid()
    plt.show()


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