from sympy.abc import x, y
import sympy
from sympy.polys.polytools import degree
import numpy as np

from .carry import Carry


def add(
    expansion: np.ndarray, carry: Carry, val: int = 1, center: tuple[int, int] = (0, 0)
):
    """
    Add `val` to a position in an `expansion` designated the `center`, then apply `carry`.
    """
    expansion[center] += val
    carry.apply(expansion)
    return expansion


def times(expansion: np.ndarray, carry: Carry, val: int):
    """
    Multiply all digits in an `expansion` by `val`, then apply `carry`.
    """
    expansion *= val
    carry.apply(expansion)
    return expansion


def poly_from_array(
    array: np.ndarray, center: tuple[int, int] = (0, 0), x_var=x, y_var=y
):
    """
    Convert a numpy array to a polynomial in two variables.
    Note that indices are row-major in Numpy, so center is in the form (y, x).
    """
    ret = 0
    for i, row in enumerate(array):
        for j, val in enumerate(row):
            ret += val * (x_var ** (i - center[1]) * y_var ** (j - center[0]))
    return ret


def poly_to_array(
    poly: sympy.Expr,
    alloc: np.ndarray | None = None,
    x_var=x,
    y_var=y
) -> np.ndarray:
    '''Convert a polynomial of 2 variables to a 2D array'''
    if isinstance(poly, (int, sympy.core.numbers.Zero)):
        if alloc is None:
            alloc = np.zeros((1, 1), dtype=np.int32)
            alloc[0,0] = int(poly)
            return alloc
        alloc.fill(0)
        alloc[0,0] = int(poly)
        return alloc

    coeffs = list(sorted(
        map(
            lambda z: (z, (degree(z[0], x_var), degree(z[0], y_var))),  # type: ignore
            poly.as_coefficients_dict().items()
        ),
        key=lambda z: sum(z[1])
    ))

    max_x, max_y = coeffs[-1][1]
    if alloc is None:
        alloc = np.zeros((max_x + max_y + 1, max_x + max_y + 1), dtype=np.int32)
    elif alloc.shape[0] < (max_x + 1) or alloc.shape[1] < (max_y + 1):
        raise ValueError("return array has insufficient size")
    else:
        alloc.fill(0)

    for (_, coeff), (x_deg, y_deg) in coeffs:
        alloc[y_deg, x_deg] = coeff
    return alloc


def _first_nonzero(arr: np.ndarray, axis: int):
    """
    Find the index of the first nonzero entry of `arr` along `axis`.
    Modified from https://stackoverflow.com/questions/47269390
    """
    mask = arr != 0
    return min(
        np.where(mask.any(axis=axis), mask.argmax(axis=axis), arr.shape[axis] + 1)
    )


def latex_polynumber(
    arr: np.ndarray, center: tuple[int, int] | None = None, show_zero: bool = False
):
    """
    Convert a 2D Numpy array into a LaTeX array representing a polynomial in two variables.
    Vertical lines and `\\hline`s are placed to indicate the transition between negative
        and nonnegative powers, similar to a typical fractional separator (decimal point).

    By default, this assumes the constant term is in position (0, 0) of the array.
    It may be alternatively supplied by `center`, which designates `arr[center]` as the constant term.
    Note that indices are row-major in Numpy, so `center` is in the form (y, x).

    `show_zero` defaults to False, which causes entries in `arr` which are 0 to not be drawn.
    When true, all 0 entries will be drawn.
    """
    upper_left = _first_nonzero(arr, 0), _first_nonzero(arr, 1)
    lower_right = (
        len(arr) - _first_nonzero(np.flip(arr, 0), 0) - 1,
        len(arr[1]) - _first_nonzero(np.flip(arr, 1), 1) - 1,
    )

    center_left, center_top = center or (0, 0)
    if center is not None:
        # this does not need offset, since we iterate over this range
        center_top = center[0]
        # but this does for the array environment argument
        center_left = center[1] - upper_left[1]

    num_columns = lower_right[1] - upper_left[1]
    column_layout = ("c" * center_left) + "|" + ("c" * (num_columns - center_left))

    # build array output
    ret = "\\begin{array}{" + column_layout + "}"
    for i in range(upper_left[0], lower_right[0] + 1):
        if i == center_top:
            ret += " \\hline "
        ret += " & ".join(
            [
                str(arr[i, j] if arr[i, j] != 0 or show_zero else "")
                for j in range(upper_left[1], lower_right[1] + 1)
            ]
        )
        # next row
        ret += " \\\\ "
    return ret + "\\end{array}"