from numpy import zeros
from numpy.typing import NDArray
from typing import NamedTuple, Any


def compute_matrix(number: int):
    """
    Compute the index order (from 1 because it will be used by MATLAB) of measurmenet taken with a "snail" shape.
    Used for the backscattermeter analysis MATLAB script

    Parameters
    ----------
    number: int
        Number of block from center to the extremity of the map
    """
    matrix = zeros((number * 2 + 1, number * 2 + 1), dtype=int)
    cursor = [number, number]
    counter = 1
    matrix[*cursor] = counter
    for i in range(number * 2):
        for _ in range(i + 1):
            if i % 2 == 0:
                cursor[1] += 1
            else:
                cursor[1] -= 1
            counter += 1
            matrix[*cursor] = counter
        for _ in range(i + 1):
            if i % 2 == 0:
                cursor[0] -= 1
            else:
                cursor[0] += 1
            counter += 1
            matrix[*cursor] = counter
    for i in range(number):
        cursor[1] += 1
        counter += 1
        matrix[*cursor] = counter
    for i in range(number):
        cursor[1] += 1
        counter += 1
        matrix[*cursor] = counter
    return matrix


def diff_matrix(number: int) -> list[list[int]]:
    posX, negX, posY, negY = [], [], [], []
    index = 0
    for i in range(number * 2):
        for _ in range(i + 1):
            if i % 2 == 0:
                posX.append(index)
            else:
                negX.append(index)
            index += 1
        for _ in range(i + 1):
            if i % 2 == 0:
                posY.append(index)
            else:
                negY.append(index)
            index += 1
    for _ in range(number):
        posX.append(index)
        index += 1
    for _ in range(number):
        posY.append(index)
        index += 1
    return [posX, negX, posY, negY]


def python2matlab(array: NDArray[Any]):
    """
    Return the MATLAB representation of a 2d python numpy array

    Parameters
    ----------
    array: NDArray[Any]
        2d array
    """
    if len(array.shape) != 2:
        raise ValueError("Can only show the representation of 2D numpy array")
    result = "matrix = [ "
    for column_index in range(array.shape[1]):
        for row_index in range(array.shape[0]):
            result += "{} ".format(array[row_index, column_index])
        if column_index != array.shape[1] - 1:
            result += "; "
    result += "]';"
    return result


class Vector(NamedTuple):
    x: float
    y: float

    def __truediv__(self, other: "Vector | float | int") -> "Vector":
        if type(other) is type(self):
            return Vector(x=self.x / other.x, y=self.y / other.y)
        return Vector(x=self.x / other, y=self.y / other)

    def __mul__(self, other: "Vector | float | int") -> "Vector":
        if type(other) is type(self):
            return Vector(x=self.x * other.x, y=self.y * other.y)
        return Vector(x=self.x * other, y=self.y * other)

    def add(self, other: "Point | Vector | float | int") -> "Vector":
        if type(other) == type(self) or type(other) is Point:
            return Vector(x=self.x + other.x, y=self.y + other.y)
        return Vector(x=self.x + other, y=self.y + other)


class Point(NamedTuple):
    x: float
    y: float

    def add(self, other: "Point | Vector | float | int") -> "Point":
        if type(other) == type(self) or type(other) is Vector:
            return Point(x=self.x + other.x, y=self.y + other.y)
        return Point(x=self.x + other, y=self.y + other)


class Axis(NamedTuple):
    pos: Vector
    neg: Vector


class Directions(NamedTuple):
    x: Axis
    y: Axis


def convert_to_base(base: tuple[Vector, Vector], vector: Vector) -> Vector:
    y = (vector.y - (base[0].y * vector.x / base[0].x)) / (
        base[1].y - (base[0].y * base[1].x / base[0].x)
    )
    x = vector.x / base[0].x - y * (base[1].x / base[0].x)
    return Vector(x=x, y=y)


def move_snail(
    n_grid: int, point: Point, directions: Directions
) -> tuple[list[Vector], Point]:
    movement: list[Vector] = []
    for i in range(2 * n_grid):
        for _ in range(i + 1):
            if (i % 2) == 0:
                movement.append(directions.x.pos)
            else:
                movement.append(directions.x.neg)
            point = point.add(movement[-1])
        for _ in range(i + 1):
            if (i % 2) == 0:
                movement.append(directions.y.pos)
            else:
                movement.append(directions.y.neg)
            point = point.add(movement[-1])
    # go back
    for i in range(n_grid):
        movement.append(directions.x.pos)
        point = point.add(movement[-1])
    for i in range(n_grid):
        movement.append(directions.y.pos)
        point = point.add(movement[-1])
    return movement, point