from functools import partial
from typing import Callable, Literal

from matplotlib import animation, pyplot as plt
import numpy as np

from .carry import Carry

writer = animation.writers["ffmpeg"](fps=15, metadata={"artist": "Me"})


def times(base: np.ndarray, carry: Carry, val):
    base *= val
    carry.apply(base)
    return base


def add(base: np.ndarray, carry: Carry, val=1, center=(0, 0)):
    base[center] += val
    carry.apply(base)
    return base


def animate(
    fig, frames: list[int] | None, **kwargs
) -> Callable[..., animation.FuncAnimation]:
    # Why isn't this how the function is exposed by default?
    return lambda func: animation.FuncAnimation(fig, func, frames, **kwargs)


def animate_carry_count(
    carry: Carry,
    dims: int = 100,
    operation: Literal["add", "multiply"] = "add",
    op_val=1,
    center: tuple[int, int] | None = None,
    # Animation parameters
    frames: list[int] | None = None,
    interval=200,
) -> animation.FuncAnimation:
    if center is None:
        if carry.over_pos != (0, 0):
            center = (dims // 2, dims // 2)
        else:
            center = (0, 0)

    if operation == "add":
        next_func = partial(add, carry=carry, val=op_val, center=center)
    elif operation == "multiply":
        if op_val == 1:
            raise ValueError(f"too small value {op_val} for repeated multiplication")
        next_func = partial(times, carry=carry, val=op_val)
    else:
        raise ValueError(f"Cannot use {repr(operation)} for animation")

    expansion = np.zeros((dims, dims))
    expansion[center] = 1
    title = [1]

    fig = plt.gcf()

    image = plt.imshow(
        expansion,
        extent=[-center[0], dims - center[0], dims - center[1], -center[1]],  # type: ignore
    )
    def init():
        plt.title(f"{title[0]}")
        image.set_clim(0, carry.overflow - 1)
        fig.tight_layout()

    @animate(fig, frames, interval=interval, init_func=init)
    def ret(_):
        plt.title(f"{title[0]}")
        image.set_data(expansion)
        fig.tight_layout()

        next_func(expansion)
        if operation == "add":
            title[0] += op_val
        else:
            title[0] *= op_val

    return ret