import itertools
from collections.abc import Iterator

import equinox as eqx
import jax
import jax.numpy as jnp
import optax
import pytest
from equinox import nn
from jaxtyping import Array, Float, PRNGKeyArray, PyTree
from pytest_codspeed import BenchmarkFixture

from differt.geometry import Scene, TracedPaths
from differt.utils import sample_points_in_bounding_box


def random_tx_rx(
    base_scene: Scene, num_tx: int = 10, num_rx: int = 10, *, key: PRNGKeyArray
) -> Scene:
    key_tx, key_rx = jax.random.split(key, 2)
    bounding_box = base_scene.mesh.bounding_box
    scene = eqx.tree_at(
        lambda s: s.transmitters,
        base_scene,
        sample_points_in_bounding_box(bounding_box, (num_tx,), key=key_tx),
    )
    return eqx.tree_at(
        lambda s: s.receivers,
        scene,
        sample_points_in_bounding_box(bounding_box, (num_rx,), key=key_rx),
    )


@eqx.filter_jit
def random_scene(
    base_scene: Scene, num_tx: int = 4, num_rx: int = 8, *, key: PRNGKeyArray
) -> Scene:
    key_tx_rx, key_num_objects, key_sample_triangles = jax.random.split(key, 3)
    scene = random_tx_rx(base_scene, num_tx, num_rx, key=key_tx_rx)
    fill_factor = jax.random.uniform(key_num_objects, ())
    return eqx.tree_at(
        lambda s: s.mesh,
        scene,
        scene.mesh.sample(fill_factor, by_masking=True, key=key_sample_triangles),
    )


def train_dataloader(base_scene: Scene, *, key: PRNGKeyArray) -> Iterator[Scene]:
    while True:
        key, key_to_use = jax.random.split(key, 2)
        yield random_scene(base_scene, key=key_to_use)


class LOSModel(eqx.Module):
    embeds: nn.MLP
    logits: nn.Linear

    def __init__(
        self,
        num_embeds: int = 32,
        width: int = 64,
        depth: int = 3,
        *,
        key: PRNGKeyArray,
    ) -> None:
        key_embeds, key_logits = jax.random.split(key)
        self.embeds = nn.MLP(
            in_size=3,
            out_size=num_embeds,
            width_size=width,
            depth=depth,
            scan=True,
            key=key_embeds,
        )
        self.logits = nn.Linear(num_embeds * 3, "scalar", key=key_logits)

    def __call__(
        self,
        triangle_vertices: Float[Array, "num_triangles 3 3"],
        path_vertices: Float[Array, "2 3"],
    ) -> Float[Array, ""]:
        # [num_triangles 3 num_embeds] -> [num_triangles num_embeds] -> [num_embeds]
        scene_embeds = (
            jax.vmap(jax.vmap(self.embeds))(triangle_vertices).mean(axis=1).sum(axis=0)
        )

        # [2 num_embeds] -> [2*num_embeds]
        path_embeds = jax.vmap(self.embeds)(path_vertices).reshape(-1)

        logits = self.logits(jnp.concatenate([scene_embeds, path_embeds], axis=-1))
        return jax.nn.sigmoid(logits)


@eqx.filter_jit(donate="all-except-first")
def loss(model: LOSModel, scene: Scene) -> Float[Array, ""]:
    paths = scene.trace_paths(order=0, solver="exhaustive")
    assert isinstance(paths, TracedPaths)
    f = model

    for _ in range(paths.vertices.ndim - 2):
        f = jax.vmap(f, in_axes=(None, 0))

    pred = f(scene.mesh.triangle_vertices, paths.vertices)

    assert paths.mask is not None
    return jnp.mean((pred - paths.mask.astype(pred.dtype)) ** 2)


@pytest.mark.benchmark(group="training")
def test_train_step(
    simple_street_canyon_scene: Scene,
    benchmark: BenchmarkFixture,
    key: PRNGKeyArray,
) -> None:
    key_model, key_dataloader = jax.random.split(key)
    model = LOSModel(key=key_model)
    dataloader = train_dataloader(simple_street_canyon_scene, key=key_dataloader)
    optim = optax.adam(learning_rate=1e-3)
    opt_state = optim.init(eqx.filter(model, eqx.is_inexact_array))

    @eqx.filter_jit(donate="all")
    def make_step(
        model: LOSModel,
        opt_state: optax.OptState,
        scene: Scene,
    ) -> tuple[LOSModel, optax.OptState, Float[Array, ""]]:
        loss_value, grads = eqx.filter_value_and_grad(loss)(
            model,
            scene,
        )
        updates, opt_state = optim.update(
            grads, opt_state, eqx.filter(model, eqx.is_array)
        )
        model = eqx.apply_updates(model, updates)
        return model, opt_state, loss_value

    iterator = itertools.accumulate(
        dataloader,
        (lambda carry, scene: make_step(carry[0], carry[1], scene)),
        initial=(model, opt_state, jnp.array(0.0)),
    )

    @jax.block_until_ready
    def bench_fun() -> PyTree:
        return next(iterator)

    bench_fun()

    benchmark(bench_fun)
