"""The main search loop for the grid search."""

import math
from collections.abc import Callable
from typing import Any

import equinox as eqx
import equinox.internal as eqxi
import jax
import jax.numpy as jnp
import jax.tree_util as jtu
from jaxtyping import Array, PyTree

from ._internal import fori_loop_tqdm_decorator
from ._method import AbstractGridSearchMethod
from ._tree_grid import (
    tree_grid_shape,
    tree_grid_take,
    tree_grid_unravel_index,
)
from ._types import PyTreeGrid, PyTreeGridPoint


@eqx.filter_jit
def run_grid_search(
    fn: Callable[[PyTreeGridPoint, Any], Array],
    method: AbstractGridSearchMethod,
    tree_grid: PyTreeGrid,
    args: Any,
    *,
    is_leaf: Callable[[Any], bool] | None = None,
    progress_bar: bool = False,
    total_updates: int = 20,
) -> PyTree[Any]:
    """Run a grid search over the function `fn`.

    !!! question "What is a `tree_grid`?"

        This is the grid of inputs to `fn` to search over. It is
        represented as an arbitrary
        pytree, whose leaves are JAX arrays with a leading dimension.
        For a particular leaf, its leading dimension indexes a set
        grid points. The entire grid is then the cartesian product
        of the grid points of all of its leaves.

        From an arbitrary tree structure, create a `tree_grid`
        as `tree_grid, _ = equinox.partition(pytree, eqx.is_array)`.

    To learn more, see the `tree_grid` manipulation routines
    [`brutax.tree_grid_shape`][] and [`brutax.tree_grid_take`][].

    **Arguments:**

    - `fn`:
        The function to search over. This must have signature `fn(y, args)`,
        where `y` is a particular grid point of `tree_grid`. The value
        returned by `fn` must be compatible with the respective `method`.
    - `method`:
        An interface that specifies what we would like to do with
        each evaluation of `fn`. For example, [`brutax.MinimumSearchMethod`][].
    - `tree_grid`:
        The grid as an arbitrary pytree. Its leaves can only be JAX
        arrays, each with a leading dimension, and `None`.
    - `args`:
        Arguments passed to `fn`, i.e. `fn(y, args)`.
    - `is_leaf`:
        As [`jax.tree_util.tree_flatten`](https://jax.readthedocs.io/en/latest/_autosummary/jax.tree_util.tree_flatten.html).
        This specifies what is to be treated as a leaf in `tree_grid`. This is
        useful for wrapping multiple JAX arrays into the same grid point.
    - `progress_bar`:
        Add a [`tqdm`](https://github.com/tqdm/tqdm) progress bar to the
        search loop.
    - `total_updates`:
        The number of updates of the progress bar over the course of the search.
        By default, `20`. Ignored if `progress_bar = False`.

    **Returns:**

    The solution of the search. In general, this is the output of
    [`brutax.AbstractGridSearchMethod.postprocess`][]. For the case
    of [`brutax.MinimumSearchMethod`][], access the solution of the
    search as `sol = brutax.run_grid_search(...); sol.value`.
    ```
    """
    # Evaluate the shape and dtype of the output of `fn` using
    # eqx.filter_closure_convert.
    test_tree_grid_point = tree_grid_take(
        tree_grid,
        tree_grid_unravel_index(0, tree_grid, is_leaf=is_leaf),
    )
    fn = eqx.filter_closure_convert(fn, test_tree_grid_point, args)
    f_struct = jtu.tree_map(
        lambda x: x.value,
        jtu.tree_map(eqxi.Static, fn.out_struct),  # type: ignore
        is_leaf=lambda x: isinstance(x, eqxi.Static),
    )
    # Get the initial state of the search method
    init_state = method.init(tree_grid, f_struct, is_leaf=is_leaf)
    dynamic_init_state, static_state = eqx.partition(init_state, eqx.is_array)
    # Finally, build the loop
    init_carry = (dynamic_init_state, tree_grid)

    def brute_force_body_fun(iteration_index, carry):
        dynamic_state, tree_grid = carry
        state = eqx.combine(static_state, dynamic_state)
        tree_grid_point = tree_grid_take(
            tree_grid,
            tree_grid_unravel_index(iteration_index, tree_grid, is_leaf=is_leaf),
        )
        new_state = method.update(fn, tree_grid_point, args, state, iteration_index)
        new_dynamic_state, new_static_state = eqx.partition(new_state, eqx.is_array)
        assert eqx.tree_equal(static_state, new_static_state) is True
        return new_dynamic_state, tree_grid

    def batched_body_fun(iteration_index, carry):
        dynamic_state, tree_grid = carry
        state = eqx.combine(static_state, dynamic_state)
        raveled_grid_index_batch = jnp.linspace(
            iteration_index * method.batch_size,
            (iteration_index + 1) * method.batch_size - 1,
            method.batch_size,  # type: ignore
            dtype=jnp.int32,
        )
        tree_grid_points = tree_grid_take(
            tree_grid,
            tree_grid_unravel_index(raveled_grid_index_batch, tree_grid, is_leaf=is_leaf),
        )
        new_state = method.batch_update(
            fn, tree_grid_points, args, state, raveled_grid_index_batch
        )
        new_dynamic_state, new_static_state = eqx.partition(new_state, eqx.is_array)
        assert eqx.tree_equal(static_state, new_static_state) is True
        return new_dynamic_state, tree_grid

    # Get the number of iterations of the loop (the size of the grid)
    grid_size = math.prod(tree_grid_shape(tree_grid, is_leaf=is_leaf))
    if method.batch_size is None:
        n_iterations = grid_size
        body_fun = brute_force_body_fun
    else:
        n_iterations = grid_size // method.batch_size
        body_fun = batched_body_fun
    # Run and unpack results
    if progress_bar:
        print_every = n_iterations // total_updates
        if print_every == 0:
            raise ValueError(
                "The number of progress bar updates is greater than the "
                "number of iterations. Try decreasing `total_updates` "
                "or setting it to `None`."
            )
        body_fun = fori_loop_tqdm_decorator(n_iterations, print_every)(body_fun)
    final_carry = jax.lax.fori_loop(0, n_iterations, body_fun, init_carry)
    if method.batch_size is not None and grid_size % method.batch_size != 0:
        final_carry = jax.lax.fori_loop(
            n_iterations,
            n_iterations + (grid_size % method.batch_size),
            batched_body_fun,
            final_carry,
        )
    dynamic_final_state, _ = final_carry
    final_state = eqx.combine(static_state, dynamic_final_state)
    # Return the solution
    solution = method.postprocess(tree_grid, final_state, f_struct, is_leaf=is_leaf)
    return solution
