import argparse
import csv
from itertools import accumulate, chain, repeat
from pathlib import Path
import random
from typing import Any, Generator, TypeVar

import matplotlib.pyplot as plt
from matplotlib.widgets import Slider, RadioButtons
from matplotlib.animation import FuncAnimation
import numpy as np
from numpy.typing import NDArray

T = TypeVar("T")


def to_base(b: int, x: int) -> list[int]:
    """Convert x to an integer base b."""
    if x == 0:
        return [0]

    ret = []
    while x > 0:
        ret.append(x % b)
        x = x // b
    return ret


def from_base(b: float, xs: list[int]) -> float:
    """
    Interpret xs as a string in base `b`.
    For genericness, the base and return are both floats.
    """
    ret = 0
    plval = 1
    for i in xs:
        ret += i * plval
        plval *= b
    return ret


class Expansions:
    """List of base-b expansions of a number"""

    def __init__(self, base: float, name: str, xs: list[list[int]]):
        self.xs = xs
        self.name = name
        self.base = base

    @classmethod
    def from_file(cls, base: float, expansions_file: Path, name: str | None = None):
        ret = []
        with open(expansions_file) as f:
            for row in csv.reader(f.readlines()):
                row_even = [int(i) for i in row]
                row_odd = row_even.copy()
                row_odd[0] = 1
                ret.extend([row_even, row_odd])

        # Extract name of base from filename
        if name is None:
            name_from_file = expansions_file.name.removesuffix(expansions_file.suffix)
            second_underscore = name_from_file.find("_", name_from_file.find("_") + 1)
            name = (
                name_from_file[:second_underscore]
                if second_underscore >= 0
                else name_from_file
            )

        return cls(base, name, ret)

    @classmethod
    def from_integer(cls, base: int, count: int, name: str | None = None):
        return cls(
            float(base),
            str(base) if name is None else name,
            [to_base(base, i) for i in range(count)],
        )

    @classmethod
    def from_random(cls, base: int, count: int, digits_count: int):
        return cls(
            float(base),
            f"{base}_(random)",
            [[random.randint(0, 1) for _ in range(digits_count)] for _ in range(count)]
        )


cendree_divmod = Expansions.from_file(
    1 + 3**0.5, Path("./cendree_DivMod_count_1024_256_digits.csv"), "divMod_cendree"
)
cendree_quotrem = Expansions.from_file(
    1 + 3**0.5, Path("./cendree_QuotRem_count_1024_256_digits.csv"), "quotRem_cendree"
)


def slices(xs):
    for i in range(1, len(xs) + 1):
        yield xs[:i]


def extend_last(xs: Generator[T, Any, Any]) -> Generator[T, Any, Any]:
    i = next(xs)
    yield i
    for i in xs:
        yield i
    while True:
        yield i


# More efficient embedding
def to_angle(xs: list[int], len_: int, b: float) -> complex:
    return np.exp(2j * np.pi * from_base(b, xs) / b ** (len_ + 1))


def _adic_angles(b: float, expansion: list[int], maxpow: int = 10) -> list[complex]:
    return [
        to_angle(x, i, b) for i, x in zip(range(maxpow), extend_last(slices(expansion)))
    ]


def geosum(c, phis: list[complex]) -> complex:
    return sum(
        i * j
        for i, j in zip(phis, accumulate(chain([1], repeat(c)), lambda x, y: x * y))
    )


class PadicPlot:
    EXPANSIONS_LENGTH = 1024

    def __init__(self, expansions: Expansions, axes=None, plot=False, interactive=False, **kwargs):
        self._expansions: Expansions = expansions
        self._angles: NDArray[np.complex64]
        self._embedding: NDArray[np.complex64]

        self._list_base: int = 2 if isinstance(expansions.base, float) else expansions.base
        self._embedding_base: float = 2 if isinstance(expansions.base, float) else expansions.base
        self._geometric: float = 0.9
        self._update(set_data=False)

        self._axes = {}
        self._inputs = {}
        self._plot = None

        axes = axes if axes is not None else plt.gca()
        if interactive:
            self.plot(axes=axes, **kwargs)
            self._interactive(axes=axes)
        elif plot:
            self.plot(axes=axes, **kwargs)

    def plot(self, axes, linestyle="None", marker=".", **kwargs):
        self._axes["plot"] = axes
        self._plot = axes.plot(
            self._embedding.real,
            self._embedding.imag,
            linestyle=linestyle,
            marker=marker,
            **kwargs,
        )[0]

        plt.tight_layout()

    def _interactive(self, axes):
        # adjust the main plot to make room for the sliders
        fig = axes.get_figure()
        fig.subplots_adjust(bottom=0.25)

        self._axes["adic"] = fig.add_axes([0.65, 0.025, 0.3, 0.175])  # type: ignore
        self._inputs["adic"] = RadioButtons(
            labels=["b-adics", "$\\kappa$-adic (balanced)", "$\\kappa$-adic (binary)", "Random Binary"],
            ax=self._axes["adic"],
        )
        self._inputs["adic"].on_clicked(self._update_adic)

        self._axes["list_base"] = fig.add_axes([0.1, 0.15, 0.45, 0.03])  # type: ignore
        self._inputs["list_base"] = Slider(
            ax=self._axes["list_base"],
            label="$b$",
            valmin=2,
            valmax=10,
            valstep=1,
            valinit=2,
        )
        self._inputs["list_base"].on_changed(self._update_list_base)

        self._axes["embedding_base"] = fig.add_axes([0.1, 0.1, 0.45, 0.03])  # type: ignore
        self._inputs["embedding_base"] = Slider(
            ax=self._axes["embedding_base"],
            label="$p$",
            valmin=1.1,
            valmax=10,
            valstep=0.1,
            valinit=2,
        )
        self._inputs["embedding_base"].on_changed(self._update_embedding_base)

        self._axes["geometric"] = fig.add_axes([0.1, 0.05, 0.45, 0.03])  # type: ignore
        self._inputs["geometric"] = Slider(
            ax=self._axes["geometric"],
            label="$c$",
            valmin=0.001,
            valmax=0.999,
            valinit=0.9,
        )
        self._inputs["geometric"].on_changed(self._update_geometric)

    def _update_adic(self, val):
        if val == "b-adics":
            expansions = Expansions.from_integer(int(self._list_base), self.EXPANSIONS_LENGTH)
        elif val == "$\\kappa$-adic (balanced)":
            expansions = cendree_quotrem
        elif val == "$\\kappa$-adic (binary)":
            expansions = cendree_divmod
        elif val == "Random Binary":
            expansions = Expansions.from_random(2, self.EXPANSIONS_LENGTH, self.EXPANSIONS_LENGTH // 10)
        else:
            raise ValueError("Impossible")

        self._expansions = expansions
        self._update()

    def _update_list_base(self, val):
        self._list_base = val
        if val == "b-adics":
            self._expansions = Expansions.from_integer(int(val), self.EXPANSIONS_LENGTH)
            self._update()

    def _update_embedding_base(self, val):
        self._embedding_base = val
        self._update()

    def _update_geometric(self, val):
        self._geometric = val
        self._update(False)

    def _update(self, regen_angles: bool = True, set_data: bool = True):
        if regen_angles:
            self._angles = np.array([
                _adic_angles(self._embedding_base, expansion, 15) for expansion in self._expansions.xs
            ])
        # TODO: vectorize
        self._embedding: NDArray[np.complex64] = np.array(
            [geosum(self._geometric, phis) for phis in self._angles]
        )

        if set_data:
            self._set_data()

    def _set_data(self):
        self._plot.set_data(self._embedding.real, self._embedding.imag)  # type: ignore
        self._axes["plot"].set_xlim(
            min(self._embedding.real),
            max(self._embedding.real)
        )
        self._axes["plot"].set_ylim(
            min(self._embedding.imag),
            max(self._embedding.imag)
        )
        self._axes["plot"].get_figure().canvas.draw_idle()


def interactive(expansions: Expansions):
    PadicPlot(expansions=expansions, interactive=True)
    plt.show()


def anim(expansions: Expansions):
    phiss = [
        [_adic_angles(expansions.base, expansion, 200) for expansion in expansions.xs]
    ]
    embedding: NDArray[np.complex64] = np.array(
        [geosum(0.9, phis) for phis in phiss[0]]
    )

    # Create the figure and the line that we will manipulate
    fig, ax = plt.subplots()

    line = ax.plot(embedding.real, embedding.imag, linestyle="None", marker=".")[0]

    vals = np.linspace(0.9, 0.1, 100)
    i = [0]

    # The function to be called anytime a slider's value changes
    def update(_):
        embedding: NDArray[np.complex64] = np.array(
            [geosum(vals[i[0]], phis) for phis in phiss[0]]
        )
        line.set_data(embedding.real, embedding.imag)  # type: ignore
        ax.set_xlim(min(embedding.real), max(embedding.real))
        ax.set_ylim(min(embedding.imag), max(embedding.imag))
        ax.set_title(f"{vals[i[0]]}")
        i[0] = i[0] + 1

    FuncAnimation(fig, update, frames=99).save(  # type: ignore
        f"{expansions.name}-adics_count_{len(expansions.xs)}_{len(expansions.xs[-1])}_digits.mp4"
    )



def main():
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--interactive",
        "-i",
        action="store_true",
        help="Interactively set the geometric constant of the Fourier series.",
    )
    parser.add_argument(
        "--base",
        "-p",
        default=2,
        type=int,
        help="Use p-adic integers"
    )
    parser.add_argument(
        "--count",
        "-c",
        default=1024,
        type=int,
        help="Number of numbers to use. Only applies when using -p.",
    )
    parser.add_argument("--cendree", "-k", type=Path, help="Use cendree-adic integers")

    args = parser.parse_args()

    expansions: Expansions
    if args.cendree:
        if not args.cendree.exists():
            raise ValueError(f"Received cendree-adic expansion file: {args.cendree}")
        expansions = Expansions.from_file(1 + 3**0.5, args.cendree)
    else:
        expansions = Expansions.from_integer(args.base, args.count)

    if args.interactive:
        interactive(expansions)
    else:
        anim(expansions)


if __name__ == "__main__":
    main()
