import numpy as np
import matplotlib.pyplot as plt
from typing import Self
from functools import cache
import sys


class constellation_point:
    def __init__(self, num: int = 0, n: int = 0, m: int = 0):
        self.number = num
        self.n = n
        self.m = m
        self.symbol = self.encode(self.number)
        self.constellation = self._gray(self.symbol, n, m)
        self.transformed = self.transform()
        self.neighbours = []

    @staticmethod
    @cache
    def encode(number):
        if number <= 1:
            return number
        symbol = constellation_point.encode(number - 1)
        j = constellation_point.determine_bit_position(number)
        symbol = symbol ^ (1 << j)
        return symbol

    @staticmethod
    def determine_bit_position(i):
        @cache
        def find_lowest_bit(x):
            ld = np.log2(x)
            if ld == int(ld):
                return ld
            else:
                return find_lowest_bit(x - 2 ** int(ld))

        if i % 2:
            j = 0
        else:
            j = find_lowest_bit(i)
        return int(j)

    @staticmethod
    def _gray_1d(symbol: int, k: int):
        if symbol >= 2**k or symbol < 0:
            raise ValueError(f"symbol {symbol} out of range (0..{2**k})")

        # symbol += 2**(k-1)
        # symbol %= 2**k

        if k < 0:
            raise ValueError(f"k must >= 0, is {k}")
        if k == 0:
            return 0
        else:
            b0 = (symbol & (2 ** (k - 1))) >> (k - 1)  # highest value bit, shifted to the right
            b1k = symbol & ((2 ** (k - 1)) - 1)  # remaining bits
            return (1 - 2 * b0) * (2 ** (k - 1) + constellation_point._gray_1d(b1k, k=k - 1))

    @staticmethod
    def _gray_2d(symbol: int, n: int, m: int):
        #          a           b
        # symbol: [0 1 .. n-1][0 1 .. m-1]
        # mask:    1 1 .. 1    0 0 0 0
        maskb = 2 ** (m) - 1
        b = symbol & maskb
        maska = 2 ** (m + n) - 1 & ~maskb
        a = (symbol & maska) >> m
        return np.array((constellation_point._gray_1d(a, n), constellation_point._gray_1d(b, m)))

    @staticmethod
    def _gray(symbol: int, k: int = 0, m: int = 0):
        if m == 0:
            return np.array((constellation_point._gray_1d(symbol, k), 0)).flatten()
        return constellation_point._gray_2d(symbol, k, m).flatten()

    def set_neighbours(self, points: tuple[Self]):
        self.neighbours = tuple(points)
        self.hamming_distance = sum(
            (self.symbol ^ neighbour.symbol).bit_count() for neighbour in self.neighbours
        ) / len(self.neighbours)
        return self.hamming_distance

    # @staticmethod
    # def calc_hamming_distance(a:Self, b:Self):
    #     return (a.symbol ^ b.symbol).bit_count()

    def transform(self):
        if self.n == self.m:
            return self.constellation
        if self.n == 2:
            return self._transform_8()
        return self._transform_2n1()

    def _transform_2n1(self):
        ir, qr = self.constellation
        # ic, qc = self.transformed

        s = 2 ** (self.m - 1)

        # 3, -4, 2 for n = 2, m = 1
        if np.abs(ir) < 3 * s:
            return self.constellation

        if np.abs(qr) > s:
            ic = np.sign(ir) * (np.abs(ir) - 2 * s)
            qc = np.sign(qr) * (4 * s - np.abs(qr))
        else:
            ic = np.sign(ir) * (4 * s - np.abs(ir))
            qc = np.sign(qr) * (np.abs(qr) + 2 * s)

        return np.array((ic, qc))

    def _transform_8(self):
        ir, qr = self.constellation

        # 3, -4, 2 for n = 2, m = 1
        if ir < 3:
            return self.constellation

        ic = np.sign(ir) * (np.abs(ir) - 4)
        qc = np.sign(qr) * (2 + np.abs(qr))

        return np.array((ic, qc))

    def __repr__(self):
        _digits_num = int(np.log10(2 ** (self.n + self.m))) + 1
        _digits_x = int(np.log10(2 ** (self.n) - 1)) + 1
        _digits_y = int(np.log10(2 ** (self.m) - 1)) + 1
        sym = f"{self.number:>{_digits_num}d}"
        sym_b = f"{self.symbol:0{self.n + self.m}b}"
        const = f"({self.constellation[0]:+0{_digits_x}d}, {self.constellation[1]:+0{_digits_y}d})"
        return f"{sym} [{sym_b}] - {const} | n:{self.n:d}, m:{self.m:d}"


class constellation:
    def __init__(self, n_symbols: int = 1):
        self.k = int(np.ceil(np.log2(n_symbols)))
        self.n_symbols = n_symbols
        self.m = self.k // 2
        self.n = self.k - self.m

        # self.n, self.m = self.most_square_fact(n_symbols)
        # self.n = int(np.ceil(np.log2(self.n)))
        # self.m = int(np.ceil(np.log2(self.m)))
        self.points = [constellation_point(i, self.n, self.m) for i in range(self.n_symbols)]
        self.distances = self.calc_distances(transformed=False)
        self.set_nearest_neighbours()
        self.gray_penalty_original = self.calc_gray_penalty()

        self.distances = self.calc_distances(transformed=True)
        self.set_nearest_neighbours()
        self.gray_penalty_transformed = self.calc_gray_penalty()

        self.scale = 1

    def calc_distances(self, transformed=True):
        distances = np.zeros((len(self.points), len(self.points)))
        for i in range(len(self.points)):
            for j in range(len(self.points)):
                distances[i, j] = self.calc_distance(self.points[i], self.points[j], transformed=transformed)
        distances = np.ma.masked_array(distances, np.eye(*distances.shape))
        return distances

    def set_nearest_neighbours(self):
        # neighbours = []
        for i, point in enumerate(self.points):
            dists = self.distances[i]
            indices = np.where(dists == dists.min())[0]
            point.set_neighbours([self.points[index] for index in indices])
            # for j, point2 in enumerate(self.points):
            #     if i == j:
            #         continue
            #     neighbours.append(point2)
            # point.set_neighbours(neighbours)
        ...

    def calc_gray_penalty(self):
        return sum(pt.hamming_distance for pt in self.points) / self.n_symbols

    # def calc_gray_penalty(self):
    #     hamming = 0
    #     euclidean = 0
    #     for i, point in enumerate(self.points):
    #         for j, point2 in enumerate(self.points):
    #             if i == j:
    #                 continue
    #             hamming += (point.symbol ^ point2.symbol).bit_count()
    #             euclidean += np.linalg.norm(point.transformed - point2.transformed)
    #     penalty = hamming/euclidean
    #     return penalty

    def get_points_for_plot(self):
        pts = [point.constellation for point in self.points]
        xs_pre = [pt[0] * self.scale for pt in pts]
        ys_pre = [pt[1] * self.scale for pt in pts]
        labels = [point.symbol for point in self.points]
        pts = [point.transformed for point in self.points]
        xs = [pt[0] * self.scale for pt in pts]
        ys = [pt[1] * self.scale for pt in pts]
        return xs, ys, xs_pre, ys_pre, labels

    def plot(self, show=True, annotate=False, old=False):
        ax = plt.subplot()
        xs, ys, xs_pre, ys_pre, labels = self.get_points_for_plot()
        if old:
            ax.scatter(xs_pre, ys_pre, marker="o", facecolors="none", edgecolors="b")
            # circ = Circle.welzl([point.constellation for point in self.points])
            # circ.plot(ax, show=False)
        ax.scatter(xs, ys, color="b")
        circ = Circle.welzl(list(zip(xs, ys)))
        circ.plot(ax, show=False)
        if annotate:
            for i, label in enumerate(labels):
                ax.annotate(
                    f"{label:0{self.n + self.m}b}",
                    xy=(xs[i], ys[i] + self.scale / 4),
                    horizontalalignment="center",
                    verticalalignment="bottom",
                )
                # if old:
                #     ax.annotate(
                #         f"{label:0{self.n+self.m}b}",
                #         xy=(xs_pre[i], ys_pre[i] + self.scale/4),
                #         horizontalalignment='center',
                #         verticalalignment='bottom'
                #         )

        lims = (min(*xs, *ys, *xs_pre, *ys_pre) - self.scale, max(*xs, *ys, *xs_pre, *ys_pre) + self.scale)
        # ticks = tuple(range(lims[0], lims[1] + self.scale, 2 * self.scale))
        ticks = np.arange(lims[0], lims[1]+self.scale, 2*self.scale)
        ax.set_xticks(ticks)
        ax.set_yticks(ticks)
        lims = (min(circ.center[0] - circ.radius, circ.center[1] - circ.radius) - self.scale, max(circ.center[0] + circ.radius, circ.center[1] + circ.radius) + self.scale)
        ax.vlines(0, *lims, color="k", linewidth=1)
        ax.hlines(0, *lims, color="k", linewidth=1)
        ax.set_xlim(*lims)
        ax.set_ylim(*lims)
        ax.grid("major")
        ax.set_aspect("equal")
        ax.set_title(f"G_pt = {self.gray_penalty_transformed:.4f}\nG_po = {self.gray_penalty_original:.4f}")
        if show:
            plt.show()

    @staticmethod
    def most_square_fact(n):
        best_a = 1
        best_b = n
        for a in range(2, n // 2 + 1):
            b = n // a
            if b * a != n:
                continue
            if (a + b) < (best_a + best_b):
                best_a = a
                best_b = b
            elif (a + b) == (best_a + best_b) and (b - a) < (best_b - best_a):
                best_a = a
                best_b = b
        return (best_a, best_b)

    @staticmethod
    def calc_distance(a: constellation_point, b: constellation_point, transformed=True):
        if transformed:
            return np.linalg.norm(a.transformed - b.transformed)
        return np.linalg.norm(a.constellation - b.constellation)


class Circle:
    def __init__(self, center: tuple[float] | None = None, radius: float | None = None):
        self.center = center
        self.radius = radius
        self.points = None

    @staticmethod
    def welzl(P: list, R: list = []):
        if len(P) == 0 or len(R) == 3:
            return Circle.from_points(R)
        p = P.pop(np.random.randint(0, len(P)))
        D = Circle.welzl(P.copy(), R.copy())
        if p in D:
            return D

        R.append(p)
        return Circle.welzl(P.copy(), R.copy())

    @staticmethod
    def from_points(R: list | tuple):
        # if len(R) != 3:
        #     raise ValueError("R must have length 3")
        if len(R) == 0:
            return Circle((0, 0), 1)
        if len(R) == 1:
            return Circle(R[0], 1)
        if len(R) == 2:
            x = (R[0][0] + R[1][0]) / 2
            y = (R[0][1] + R[1][1]) / 2
            r = np.sqrt((R[0][0] - R[1][0]) ** 2 + (R[0][1] - R[1][1]) ** 2) / 2
            return Circle((x, y), r)
        if len(R) > 3:
            raise ValueError("more than 3 points given")
        z1, z2, z3 = (r[0] + 1j * r[1] for r in R)
        if z1 == z2 or z2 == z3 or z3 == z1:
            raise ValueError(f"Duplicate points: {z1}, {z2}, {z3}")
        w = (z3 - z1) / (z2 - z1)
        if w == w.real:
            raise ValueError(f"Points are collinear: {z1}, {z2}, {z3}")

        c = (z2 - z1) * (w - abs(w) ** 2) / (2j * w.imag) + z1
        r = abs(z1 - c)

        circ = Circle((c.real, c.imag), r)
        circ.points = R
        return circ

    def __contains__(self, pt):
        p = pt[0] + 1j * pt[1]
        c = self.center[0] + 1j * self.center[1]
        if np.linalg.norm(p - c) > self.radius:
            return False
        return True

    def __repr__(self):
        return f"Circle [c: ({self.center[0]:.2f}, {self.center[1]:.2f}), r:{self.radius:.2f}]"

    def xy_from_phase(self, phi):
        if isinstance(phi, float) or isinstance(phi, int):
            x = self.center[0] + np.cos(phi) * self.radius
            y = self.center[1] + np.sin(phi) * self.radius
            return (x, y)
        num = len(phi)
        # coords = np.zeros(num)
        center = np.array(self.center)
        coords = np.stack((np.cos(phi) * self.radius, np.sin(phi) * self.radius)) + np.expand_dims(
            center, axis=1
        ).repeat(num, axis=1)
        return coords

    def plot(self, ax=None, show=True, show_points=False):
        ax = plt.subplot() if ax is None else ax

        phi = np.arange(0, 2 * np.pi + np.pi / 50, np.pi / 50)
        xs, ys = self.xy_from_phase(phi)
        ax.plot(xs, ys, "k:")
        if self.points is not None and show_points:
            xs = [p[0] for p in self.points]
            ys = [p[1] for p in self.points]
            ax.scatter(xs, ys)
        ax.grid()
        ax.set_aspect("equal")

        if show:
            plt.show()


if __name__ == "__main__":
    # prevs = []
    # for i in range(32):
    #     gr = constellation_point.encode(i)
    #     if i > 0:
    #         prev_gr = constellation_point.encode(i-1)
    #         num_mark = i ^ (i -1)
    #         gr_mark = gr ^ (prev_gr)
    #         num_mark = f"{num_mark:05b}".replace("0", " ").replace("1", "|")
    #         gr_mark = f"{gr_mark:05b}".replace("0", " ").replace("1", "|")
    #         print(f"    {num_mark}         {gr_mark}")
    #     print(f"{i:>2d} [{i:05b}] -> {gr:>2d} [{gr:05b}] {"!" if gr in prevs else ""}")
    #     prevs.append(gr)
    # exit()
    N = 32
    if len(sys.argv) >= 2:
        N = int(sys.argv[1])

    cnst = constellation(N)
    cnst.plot(annotate=True)