farm.py

from collections import defaultdict
from typing import Annotated, Any, Literal, Self

import chex
import jax
import jax.numpy as jnp
import typer

app = typer.Typer(pretty_exceptions_enable=False)


class Farm:
    """The Farm problem is a problem where we want to place an amount of `cameras` on a
    given image. Each camera can have a unique viewing angle, orientation,
    benefitial viewing depth, on a specific location: A camera must be place on a
    given bin (horizontally or vertically), and anywere on the given bin, according
    to its respective size; horizontal beam -> width, vertical beam -> length.
    """

    def _check_angle(angle: chex.Array | chex.Scalar) -> None:
        """Checks if the angle given is valid (i.e. is < 180 degrees).

        Args:
            angle: The angle to be checked.

        Raises:
            ValueError: if The angle is >= 180 degrees.
        """

        if isinstance(angle, chex.Array):
            angle = angle.astype(float)

        if angle >= 180:
            raise ValueError(
                "An angle should be less than 180 degrees, as all triangle angles "
                f"should sum up to a total of 180 degrees. Given `triangles_angles` = "
                f"{angle.astype(float)}"
            )

    def __init__(
        self,
        cameras: chex.Scalar,
        width: chex.Scalar,
        length: chex.Scalar,
        horizontal_beam_locations: chex.Array,
        vertical_beam_locations: chex.Array,
        *,
        genotype: Literal["beams", "xy"] = "xy",
        resolution: chex.Scalar = 0.1,
        triangles_height: chex.Scalar = 3,
        triangles_angles: chex.Array = 102.0,
        bounds: chex.Array | None = None,
    ) -> None:
        """A Farm problem constructor.

        Args:
            cameras: The number of total cameras.
            width: The width of the photo.
            length: The length of the photo.
            horizontal_beam_locations: The possible horizontal beam locations, that a
              camera can be fixed upon.
            vertical_beam_locations: The possible vertical beam locations, that a camera
              can be fixed upon.
            genotype: The form of the genotype: {'beams', 'xy'}. Default is 'xy'.

                'beams' accepts a genotype with 4 genes: 1) camera orientation, 2) beam
                  orientation, 3) beam number, and 4) location on beam. This genotype
                  chooses the orientation of the camera, then which beams collection to
                  use (vertical or horizontal), which beam to use (from the respective
                  collection), and the location of that beam.

                'xy' accepts a genotype with 3 genes: 1) camera orientation, 2) x
                  location of camera in image, and 3) y location of camera in image.
                  This genotype chooses the orientation of the camera, and its location
                  on the image. Then, it finds which of the coordinates is close to on
                  of the beams, and clips it to that beam.
            resolution: The resolution of the image to convert Cartesian coordinates to
              `image` coordinates. Default is 0.1.
            triangles_height: The height of the triangle from the origin point to the
              base. Default is 3.
            triangles_angles: The angle of the `origin_points`. It can be a `float`, or
              a chex.Array with the matching `origin_points` The angle should not be
              >= 180 degrees. Default is 102.

            bounds: The bounds for each of the genes. Default is None.
              Example: [[min_1, max_1], [min_2, max_2], [min_3, max_3]].

        Raises:
            ValueError: If a mismatched type, shape, or an incorrect value is given for
              `triangles_angles`.
        """

        self.cameras: int = cameras
        self.width: chex.Scalar = width
        self.length: chex.Scalar = length
        self.horizontal_beam_locations: chex.Array = horizontal_beam_locations
        self.vertical_beam_locations: chex.Array = vertical_beam_locations
        self._genotype_structure: Literal["beams"] | Literal["xy"] = genotype
        self.resolution: chex.Array = resolution
        self.triangles_height: chex.Array = triangles_height
        self.bounds: chex.Array | None = bounds

        if isinstance(triangles_angles, chex.Array):
            if triangles_angles.shape[0] != cameras:
                raise ValueError(
                    f"`triangles_angles` shape {triangles_angles} does not match "
                    f"the total number of cameras {cameras}."
                )
            else:
                for angle in triangles_angles:
                    jax.debug.callback(Farm._check_angle, angle)
        else:
            jax.debug.callback(Farm._check_angle, triangles_angles)

        self.triangles_angles: chex.Array = triangles_angles

        if genotype == "beams":
            self._dims_per_gene: int = 4
        else:
            self._dims_per_gene: int = 3

        self.dims: int = cameras * self._dims_per_gene

    def _tree_flatten(self) -> tuple[tuple, dict[str, Any]]:
        """Required by JAX."""

        children: tuple = ()
        aux_data: dict[str, Any] = {
            "cameras": self.cameras,
            "width": self.width,
            "length": self.length,
            "horizontal_beam_locations": self.horizontal_beam_locations,
            "vertical_beam_locations": self.vertical_beam_locations,
            "genotype": self._genotype_structure,
            "resolution": self.resolution,
            "triangles_height": self.triangles_height,
            "bounds": self.bounds,
            "triangles_angles": self.triangles_angles,
        }
        return (children, aux_data)

    @classmethod
    def _tree_unflatten(cls, aux_data, children) -> Self:
        """Required by JAX."""

        return cls(*children, **aux_data)

    @jax.jit
    def draw_isosceles_triangle(
        image: chex.Array,
        origin: chex.Array,
        height: chex.Scalar,
        orientation: chex.Scalar,
        resolution: chex.Scalar,
        triangle_angle: chex.Scalar,
        /,
    ) -> chex.Array:
        """Draws an isosceles triangle on a given grey-scale image.

        Args:
            image: The image to draw the triangle on.
            origin: A jax.numpy.Array containing the coordinates of top side of the
              triangle, in Cartesian space.
            height: The height of the triangle.
            orientation: The orientation of the triangle on the Cartesian space in
              degrees.
            resolution: The resolution of the image to convert Cartesian coordinates to
              `image` coordinates.
            triangle_angle: The angle at the origin of the triangle.

        Returns:
            The image with the specified triangle drawned.
        """

        x1, y1 = origin
        orientation = jnp.radians(orientation)
        triangle_angle = jnp.radians(triangle_angle)

        half_base = height * jnp.tan(triangle_angle / 2)

        @jax.jit
        def rotate_coords(
            origin: tuple[chex.Scalar, chex.Scalar],
            angle: chex.Scalar,
            x: chex.Scalar,
            y: chex.Scalar,
            /,
        ) -> tuple[chex.Scalar, chex.Scalar]:
            origin_x, origin_y = origin
            return (
                origin_x + x * jnp.cos(angle) - y * jnp.sin(angle),
                origin_y + x * jnp.sin(angle) + y * jnp.cos(angle),
            )

        x2, y2 = rotate_coords(origin, orientation, height, half_base)  # Left point

        x3, y3 = rotate_coords(origin, orientation, height, -half_base)  # Right point

        @jax.jit
        def array_coords(x, y) -> tuple[int, int]:
            return jnp.round(x / resolution).astype(int), jnp.round(
                y / resolution
            ).astype(int)

        c1, r1 = array_coords(x1, y1)
        c2, r2 = array_coords(x2, y2)
        c3, r3 = array_coords(x3, y3)

        rows = jnp.array([r1, r2, r3])
        cols = jnp.array([c1, c2, c3])

        min_r = jnp.max(jnp.array([0, jnp.min(rows)])).astype(int)
        max_r = jnp.min(jnp.array([image.shape[0] - 1, jnp.max(rows)])).astype(int)
        min_c = jnp.max(jnp.array([0, jnp.min(cols)])).astype(int)
        max_c = jnp.min(jnp.array([image.shape[1] - 1, jnp.max(cols)])).astype(int)

        @jax.jit
        def is_point_in_triangle(x, y, x1, y1, x2, y2, x3, y3) -> bool:
            @jax.jit
            def area(x1, y1, x2, y2, x3, y3):
                return jnp.abs((x1 * (y2 - y3) + x2 * (y3 - y1) + x3 * (y1 - y2)) / 2.0)

            A = area(x1, y1, x2, y2, x3, y3)
            A1 = area(x, y, x1, y1, x2, y2)
            A2 = area(x, y, x2, y2, x3, y3)
            A3 = area(x, y, x1, y1, x3, y3)

            return jnp.round(A, 2) == jnp.round(A1 + A2 + A3, 2)

        @jax.jit
        def check_and_change_column(
            c_index: int, carry: tuple[int, chex.Array]
        ) -> tuple[int, chex.Array]:
            r_index, image = carry
            return (
                r_index,
                jax.lax.select(
                    is_point_in_triangle(
                        c_index * resolution,
                        r_index * resolution,
                        x1,
                        y1,
                        x2,
                        y2,
                        x3,
                        y3,
                    ),
                    image.at[r_index, c_index].set(1),
                    image,
                ),
            )

        @jax.jit
        def check_and_change_row(r_index: int, carry: chex.Array) -> chex.Array:
            image = carry
            (r_index, image) = jax.lax.fori_loop(
                min_c, max_c + 1, check_and_change_column, (r_index, image)
            )
            return image

        return jax.lax.fori_loop(min_r, max_r + 1, check_and_change_row, (image))

    def isosceles_triangle_coverage(
        origin_points: chex.Array,
        orientations: chex.Array,
        length: chex.Array,
        width: chex.Array,
        resolution: chex.Array,
        triangles_height: chex.Array,
        triangles_angles: chex.Array,
        /,
    ) -> tuple[chex.Scalar, chex.Array, chex.Array]:
        """Draws the given isosceles triangle points and orienations over a grey-scale
          image.

        Args:
            origin_points: A `jax.numpy.Array` of points containing the coordinates of
              the top side of the triangles in Cartesian space.
              Example: [[X1, X2], [Y1,Y2]].
            orientations: A `jax.numpy.Array` of the orientation of each triangle. The
              amount of orientations, should match the given `origin_points`.
            length: The length (Y axis) of the image to draw the triangles on.
            width: The width (X axis) of the image the triangles on.
            resolution: The resolution of the image. With the resolution, length and
              width are divided into cells/pixels.
            triangles_height: The height of the triangle from the origin point to the
              base.
            triangles_angles: The angle of the `origin_points`. It can be a `float`, or
              a chex.Array with the matching `origin_points`.

        Returns:
            The coverage of the image, the grey-scale resulting image, and the empty
              image.
        """

        blank_image = jnp.zeros((int(length // resolution), int(width // resolution)))
        covered_image = blank_image.copy()

        @jax.jit
        def draw(index: int, carry: chex.Array) -> chex.Array:
            image = carry
            return Farm.draw_isosceles_triangle(
                image,
                origin_points[:, index],
                triangles_height,
                orientations[index],
                resolution,
                (
                    triangles_angles
                    if not isinstance(triangles_angles, chex.Array)
                    or len(triangles_angles.shape) == 0
                    else triangles_angles[index]
                ),
            )

        covered_image = jax.lax.fori_loop(
            0, origin_points.shape[1], draw, covered_image
        )

        return (
            jnp.count_nonzero(covered_image) / covered_image.size,
            covered_image,
            blank_image,
        )

    @jax.jit
    def _calculate_xy_fitness(
        self, genotypes: chex.Array, /
    ) -> tuple[chex.Array, chex.Array, chex.Array, chex.Array]:
        """Calculates the fitness of the given genotype(s) with 3 genes.

        Each collection of genotypes should have 3 * n genes, where n is the number of
          total genotypes. For every genotype, its genes should have values for:
            1) camera orientation,

            2) x location of camera in image, and

            3) y location of camera in image.
        The must be in the same order as above.

        Args:
            genotypes: The genotype(s) to calculate the fitness.
              Expected shape = (n, 3 * m), where n and m are positive numbers.

        Returns:
            The coverage of the image, the grey-scale resulting image, the empty image,
              the origin points ([[X1, x2], [Y1, Y2]]), and their orientations.
        """

        @jax.jit
        def convert_genotype(genotype: chex.Array) -> chex.Array:
            result = genotype * jnp.array([360, self.width, self.length])
            v = jnp.abs(self.vertical_beam_locations - result[1])
            h = jnp.abs(self.horizontal_beam_locations - result[2])

            v_min_arg = jnp.argmin(v)
            h_min_arg = jnp.argmin(h)

            result = jax.lax.select(
                v[v_min_arg] < h[h_min_arg],
                result.at[1].set(self.vertical_beam_locations.at[v_min_arg].get()),
                result.at[2].set(self.horizontal_beam_locations.at[h_min_arg].get()),
            )

            return result

        result = jax.vmap(convert_genotype, (0))(genotypes)

        score, covered_image, blank_image = Farm.isosceles_triangle_coverage(
            result[:, 1:].T,
            result[:, 0],
            self.length,
            self.width,
            self.resolution,
            self.triangles_height,
            self.triangles_angles,
        )
        return score, covered_image, blank_image, result[:, 1:].T, result[:, 0]

    @jax.jit
    def _calculate_beams_fitness(
        self, genotypes: chex.Array, /
    ) -> tuple[chex.Array, chex.Array, chex.Array, chex.Array]:
        """Calculates the fitness of the given genotype(s) with 4 genes.

        Each collection of genotypes should have 4 * n genes, where n is the number of
        total genotypes. For every genotype, its genes should have values for:
            1) camera orientation,

            2) beam orientation,

            3) beam number, and

            4) location on beam.
        The must be in the same order as above.

        Args:
            genotypes: The genotype(s) to calculate the fitness.
              Expected shape = (n, 4 * m), where n and m are positive numbers.

        Returns:
            The coverage of the image, the grey-scale resulting image, the empty image,
              the origin points ([[X1, x2], [Y1, Y2]]), and their orientations.
        """

        @jax.jit
        def convert_genotype(genotype: chex.Array) -> chex.Array:
            result = jax.lax.select(
                genotype.at[1].get() < 0.5,
                jnp.array(
                    [
                        genotype.at[0].get() * 360,
                        genotype.at[3].get() * self.width,
                        self.horizontal_beam_locations.at[
                            (
                                genotype.at[2].get()
                                // (1 / self.horizontal_beam_locations.shape[0])
                            )
                            .clip(max=self.horizontal_beam_locations.shape[0])
                            .astype(int)
                        ].get(),
                    ]
                ),
                jnp.array(
                    [
                        genotype.at[0].get() * 360,
                        self.vertical_beam_locations.at[
                            (
                                genotype.at[2].get()
                                // (1 / self.vertical_beam_locations.shape[0])
                            )
                            .clip(max=self.vertical_beam_locations.shape[0])
                            .astype(int),
                        ].get(),
                        genotype.at[3].get() * self.length,
                    ]
                ),
            )

            return result

        result = jax.vmap(convert_genotype, (0))(genotypes)

        score, covered_image, blank_image = Farm.isosceles_triangle_coverage(
            result[:, 1:].T,
            result[:, 0],
            length=self.length,
            width=self.width,
            resolution=self.resolution,
            triangles_height=self.triangles_height,
            triangles_angles=self.triangles_angles,
        )
        return score, covered_image, blank_image, result[:, 1:].T, result[:, 0]

    @jax.jit
    def _normalize_genotypes(self, genotypes: chex.Array, /) -> chex.Array:
        """Normalizes the genotypes to between [0, 1].

        Args:
            genotypes: The genotypes to be normalized.

        Returns:
            The normalized genotypes.
        """

        genotypes = jnp.where(genotypes == jnp.inf, 1, genotypes)
        genotypes = jnp.where(genotypes == -jnp.inf, 0, genotypes)

        if self.bounds is None:
            min_bounds = jnp.nanmin(genotypes)
            max_bounds = jnp.nanmax(genotypes)
        else:
            min_bounds = self.bounds.at[:, 0].get()
            max_bounds = self.bounds.at[:, 1].get()

        return jax.lax.select(
            (max_bounds - min_bounds).sum() == 0,
            genotypes,
            (genotypes - min_bounds) / (max_bounds - min_bounds),
        )

    @jax.jit
    def _validate_genotypes(self, genotypes: chex.Array, /) -> None:
        """Validates the given `genotypes`.

        The `genotypes` should match the expected farm dimensions.

        Raises:
            ValueError: If `genotypes` size is not equal to the farm dimensions.
        """

        if genotypes.size != self.dims:
            raise ValueError(
                f"'genotypes' should have a dimensionality equal to {self.dims}, but "
                f"instead {genotypes.shape[0]} was given."
            )

    @jax.jit
    def fitness(
        self, genotypes: chex.Array, /
    ) -> tuple[chex.Array, chex.Array, chex.Array, chex.Array, chex.Array]:
        """Calculates the fitness of the given genotype(s).

        Each collection of genotypes should have a shape equal to the dimension of the
          farm = cameras * genes.

        Args:
            genotypes: The genotype(s) to calculate the fitness. Expected shape =
              (cameras * genes), where n and m are positive numbers.

        Returns:
            The coverage of the image, the grey-scale resulting image, the empty image,
              the origin points ([[X1, X2], [Y1, Y2]]), and their orientations.

        Raises:
            ValueError: If `genotypes` size is not equal to the farm dimensions.
        """

        func = self._calculate_xy_fitness
        total_genes = 3

        if self._genotype_structure == "beams":
            func = self._calculate_beams_fitness
            total_genes = 4

        self._validate_genotypes(genotypes)

        genotypes = genotypes.reshape(-1, total_genes)
        genotypes = self._normalize_genotypes(genotypes)

        return func(genotypes)

    @jax.jit
    def behaviour_descriptor_spread(
        self, genotypes: chex.Array, /, only_xy: bool = False
    ) -> chex.Scalar:
        """Calculates a spread behaviour descriptor.

        The spread behaviour descriptor, describes the well spreadness of the cameras.

        Args:
            genotypes: The genotype(s) to calculate the behaviour function. Expected
              shape = (cameras * genes), where n and m are positive numbers.
            only_xy: If True, it will only calculate the distance given X, and Y, not
              the orientation. Default is False.

        Returns:
            A Scalar in the range [0, 1], where 1 is well spread - not overlapping, and
              0 is not well spread - overlapping.
        """

        match self._genotype_structure:
            case "beams":
                genotypes = genotypes.reshape(-1, 4)
            case "xy":
                genotypes = genotypes.reshape(-1, 3)

        genotypes = self._normalize_genotypes(genotypes)

        @jax.jit
        def euclidean_distance(x: chex.Array, y: chex.Array) -> chex.Scalar:
            return jnp.linalg.norm(x - y, axis=1)

        distances = jax.lax.select(
            only_xy,
            jax.vmap(euclidean_distance, (0, None))(genotypes[:, 1:], genotypes[:, 1:]),
            jax.vmap(euclidean_distance, (0, None))(genotypes, genotypes),
        )

        min_distances = jnp.nanmin(
            jnp.fill_diagonal(distances, jnp.nan, inplace=False), 1
        )
        max_possible_distance = jnp.sqrt(
            jax.lax.select(
                only_xy,
                jnp.pow(jnp.ones(genotypes.shape[1])[1:], 2).sum(),
                jnp.pow(jnp.ones(genotypes.shape[1]), 2).sum(),
            )
        )

        return min_distances.sum() / (min_distances.size * max_possible_distance)

    @jax.jit
    def _count_beam_cameras(
        self,
        *,
        camera_points: chex.Array | None = None,
        genotypes: chex.Array | None = None,
    ) -> dict[str, chex.Array]:
        """Counts the cameras on each beam. At least of of the optional parameters must
        be given.

        Args:
            camera_points: The locations of each camera. Default is None.
            genotypes: The genotype(s) to calculate the behaviour function, to find the
              camera_points. Expected shape = (cameras * genes), where n and m are
              positive numbers. Default is None.

        Returns:
            A dict containing the number of cameras on each beam; horizontal and
              vertical. Each beam location and orientation are represented as the
              dictionary keys.

        Raises:
            ValueError: if both `genotypes` and `camera_points` are absent, or if
              `genotypes` does not match the farm dimensions, or `camera_points` do not
              match the farm cameras.
        """

        if camera_points is None:
            if genotypes is None:
                raise ValueError(
                    "One argument between 'genotypes' or 'camera_points' must be given."
                )
            else:
                _, _, _, camera_points, _ = self.fitness(genotypes)
        else:
            if camera_points.shape[1] != self.cameras:
                raise ValueError(
                    f"'camera_points' second dimension {camera_points.shape[0]} should "
                    f"match the farm cameras {self.cameras}."
                )

        beams: dict = defaultdict(int)
        for point in self.horizontal_beam_locations.tolist():
            beams[f"h_{point}"] = (camera_points.at[1].get() == point).astype(int).sum()

        for point in self.vertical_beam_locations.tolist():
            beams[f"v_{point}"] = (camera_points.at[0].get() == point).astype(int).sum()

        return beams

    @jax.jit
    def behaviour_descriptor_used_beams(
        self,
        *,
        camera_points: chex.Array | None = None,
        genotypes: chex.Array | None = None,
    ) -> dict[str, chex.Array]:
        """Calculates the 'used-beams' behaviour descriptor.  At least one of the
        optional parameters must be given.

        The 'used-beams' behaviour descriptor, describes each given beam as empty (no
        camera set on it) or not empty (at least one camera is set on it).

        Args:
            camera_points: The locations of each camera.
            genotypes: The genotype(s) to calculate the behaviour function. Expected
              shape = (cameras * genes), where n and m are positive numbers.

        Returns:
            A dict containing zeroes and ones; where 0 means not used, and 1 means
              in use, for every beam, horizontal, and vertical. Each beam location and
              orientation are represented as the dictionary keys.

        Raises:
            ValueError: if both `genotypes` and `camera_points` are absent, or if
              `genotypes` does not match the farm dimensions, or `camera_points` do not
              match the farm cameras.
        """

        beams: dict[str, chex.Array] = self._count_beam_cameras(
            camera_points=camera_points, genotypes=genotypes
        )

        for key, item in beams.items():
            beams[key] = jax.lax.select(item == 0, 0, 1).astype(int)

        return beams

    @jax.jit
    def behaviour_descriptor_beams_ratio(
        self,
        *,
        camera_points: chex.Array | None = None,
        genotypes: chex.Array | None = None,
    ) -> dict[str, chex.Array]:
        """Calculates the 'beams-ratio' behaviour descriptor.  At least one of the
        optional parameters must be given.

        The 'beams-ratio' behaviour descriptor, describes each given beam as 0%-33%
        of the cameras are on the beam, 33% - 67%, and 67% - 100%.

        Args:
            camera_points: The locations of each camera.
            genotypes: The genotype(s) to calculate the behaviour function. Expected
              shape = (cameras * genes), where n and m are positive numbers.

        Returns:
            A dict containing zeroes and ones; where 0 means not used, and 1 means
              in use, for every beam, horizontal, and vertical. Each beam location and
              orientation are represented as the dictionary keys.

        Raises:
            ValueError: if both `genotypes` and `camera_points` are absent, or if
              `genotypes` does not match the farm dimensions, or `camera_points` do not
              match the farm cameras.
        """

        beams: dict[str, chex.Array] = self._count_beam_cameras(
            camera_points=camera_points, genotypes=genotypes
        )

        for key, item in beams.items():
            beams[key] = item / self.cameras

        return beams


jax.tree_util.register_pytree_node(Farm, Farm._tree_flatten, Farm._tree_unflatten)


@app.command()
def main(
    length: Annotated[
        float, typer.Option("--length", "-l", help="The length of the image.")
    ] = 6.5,
    width: Annotated[
        float, typer.Option("--width", "-w", help="The width of the image.")
    ] = 20.5,
    resolution: Annotated[
        float, typer.Option("--resolution", "-r", help="The resolution of the image.")
    ] = 0.1,
) -> None:
    import time

    start = time.time()

    farm = Farm(
        2,
        width,
        length,
        jnp.linspace(0, length, 3),
        jnp.linspace(0, width, 5),
        genotype="xy",
        resolution=resolution,
        bounds=jnp.full((3, 2), jnp.array([0, 1])),
        triangles_angles=jnp.array([102, 105]),
    )

    score, covered_image, blank_image, _, _ = farm.fitness(
        jnp.array([1, 0.541641854, 0.0000112, 0.5, 0.54156, 0.85416]),
    )

    print(f"Score using 'xy' genotype: {score}")

    farm = Farm(
        2,
        width,
        length,
        jnp.linspace(0, length, 3),
        jnp.linspace(0, width, 5),
        bounds=jnp.full((4, 2), jnp.array([0, 1])),
        triangles_angles=102,
    )

    score, covered_image, blank_image, _, _ = farm.fitness(
        jnp.array([1, 0.25, 0.541641854, 0.0000112, 0.5, 0.75, 0.54156, 0.85416]),
    )

    print(f"Score using 'beams' genotype: {score}")
    print(time.time() - start)

    import matplotlib.pyplot as plt

    plt.imshow(blank_image, cmap="gray", origin="lower")
    plt.title("Initial (Blank) Image")
    plt.show()

    plt.imshow(covered_image, cmap="gray", origin="lower")
    plt.title("Coverage by cameras")
    plt.show()


if __name__ == "__main__":
    app()