from typing import List

import numpy
from numpy.typing import NDArray

from fracted.types import Point, TransformationLike


class IFS:
    """This class implements an Iterated Function System (IFS) fractal.

    `IFS` class holds all the properties of the fractal. It can be used to generate
    the fractal step-by-step with control of the result or to generate it all at once.

    Attributes
    ----------
    transfs : List[TransformationLike]
        A list of transformations used to generate the fractal
    probs : List[float] | None
        A list of probabilities asociated with `transformations`
        If None, the sample assumes a uniform distribution over all transformations.
    point : Point
        Current position of the transformed point
    resolution : float
        Pixels per unit in `array`
    min_x : float
        Left edge of the image
    max_x : float
        Right edge of the image
    min_y : float
        Top edge of the image
    max_y : float
        Bottom edge of the image
    array : NDArray[numpy.uint32]
        An array representing the final fractal image

    Methods
    -------
    step(draw=False)
        One iteration in fractal generating process
    draw_point(point=None)
        Draws current point to `array`.
    draw(start_iter, n_iter)
        Draws the fractal and returns the `array`.

    Notes
    -----
    IFS fractals are generated by repeatedly applying a set of affine transformations.
    For more background, see:
    https://en.wikipedia.org/wiki/Iterated_function_system
    """

    transfs: List[TransformationLike]
    probs: List[float] | None
    rng: numpy.random.Generator = numpy.random.default_rng()
    point: Point
    resolution: float
    min_x: float
    max_x: float
    min_y: float
    max_y: float
    array: NDArray[numpy.uint32]

    def __init__(
        self,
        transfs: List[TransformationLike],
        probs: List[float] | None = None,
        resolution: float = 1,
        min_x: float = -100,
        max_x: float = 100,
        min_y: float = -100,
        max_y: float = 100,
        start_point: Point = (0, 0),
    ) -> None:
        """
        Parameters
        ----------
        transfs : List[TransformationLike]
            a list of transformations used to generate the fractal
        probs : List[float] | None
            A list of probabilities asociated with the transformations.
            If not given, the sample assumes a uniform distribution over all transformations.
            If not sum to 1, they will be divided by their sum.
        resolution : float
            pixels per unit in `array`
        min_x : float
            left edge of the image
        max_x : float
            right edge of the image
        min_y : float
            top edge of the image
        max_y : float
            bottom edge of the image
        array : NDArray[numpy.uint32]
            an array representing the final fractal image
        start_point : Point
            initial position of the transformed point

        Raises
        ------
        ValueError
            If `transfs` and `probs` don't have the same length
        """
        if not (probs is None):
            if len(probs) != len(transfs):
                raise ValueError("Tranfs and probs must have the same length.")
            s = sum(probs)
            if s != 1:
                for i in range(len(probs)):
                    probs[i] /= s
        self.transfs = transfs
        self.probs = probs
        self.min_x = min_x
        self.max_x = max_x
        self.min_y = min_y
        self.max_y = max_y
        self.array = numpy.zeros(
            (int((max_x - min_x) * resolution), int((max_y - min_y) * resolution)),
            dtype=numpy.uint32,
        )
        self.point = start_point
        self.resolution = resolution

    def step(self, draw: bool = False) -> None:
        """Applies random transformation to the point.

        Parameters
        ----------
        draw: bool, optional
            If True, new point will be drawn to `array`. Default is False.
        """
        i: int = self.rng.choice(len(self.transfs), p=self.probs)
        t: TransformationLike = self.transfs[i]
        self.point = t(self.point)
        if draw:
            self.draw_point(self.point)

    def draw_point(self, point: Point | None = None) -> None:
        """Draws the point to `array`.

        Parameters
        ----------
        point: Point, optional
            The point to be drawn. If not given, the obect `point` attribute
            will be used.
        """
        if point is None:
            point = self.point
        x, y = point
        if self.min_x < x < self.max_x and self.min_y < y < self.max_y:
            self.array[
                int(self.resolution * (x + self.min_x)),
                int(self.resolution * (y + self.min_y)),
            ] += 1

    def draw(self, start_iter: int, n_iter: int) -> NDArray[numpy.uint32]:
        """Draws the fractal.

        Applies `n_iter` iterations (`step()` method), draws point attractor
        to `array` and returns the array.

        Parameters
        ----------
        start_iter: int
            Number of iterations performed before drawing to get te point
            to the fractal attractor.

        n_iter: int
            Number of iterations performed after `start_iter`. After each iteration,
            the point will be drawn to `array`.

        Returns
        -------
        array
            The object's `array` after the iterations
        """
        for _ in range(start_iter):
            self.step(draw=False)
        for _ in range(n_iter):
            self.step(draw=True)
        return self.array
