from __future__ import annotations

import torch
from torch.nn import functional as F

LAPLACIAN = torch.tensor(
    [
        [0.05, 0.20, 0.05],
        [0.20, -1.00, 0.20],
        [0.05, 0.20, 0.05],
    ],
    dtype=torch.float32,
).reshape(1, 1, 3, 3)


def initial_state(
    batch: int,
    size: int,
    seed: int,
) -> torch.Tensor:
    generator = torch.Generator().manual_seed(seed)
    u = torch.ones(batch, 1, size, size)
    v = torch.zeros(batch, 1, size, size)
    radius = max(2, size // 10)
    center = size // 2
    u[
        :,
        :,
        center - radius : center + radius,
        center - radius : center + radius,
    ] = 0.50
    v[
        :,
        :,
        center - radius : center + radius,
        center - radius : center + radius,
    ] = 0.25
    noise = torch.rand((batch, 1, size, size), generator=generator) * 0.04 - 0.02
    u = torch.clamp(u + noise, 0, 1)
    v = torch.clamp(v - noise, 0, 1)
    return torch.cat([u, v], dim=1)


def gray_scott_step(
    state: torch.Tensor,
    feed: torch.Tensor,
    kill: torch.Tensor,
    diffusion_u: float = 0.16,
    diffusion_v: float = 0.08,
) -> torch.Tensor:
    kernel = LAPLACIAN.to(state)
    u = state[:, 0:1]
    v = state[:, 1:2]
    padded_u = F.pad(u, (1, 1, 1, 1), mode="circular")
    padded_v = F.pad(v, (1, 1, 1, 1), mode="circular")
    laplacian_u = F.conv2d(padded_u, kernel)
    laplacian_v = F.conv2d(padded_v, kernel)
    reaction = u * v.square()
    feed_field = feed.reshape(-1, 1, 1, 1)
    kill_field = kill.reshape(-1, 1, 1, 1)
    delta_u = diffusion_u * laplacian_u - reaction + feed_field * (1 - u)
    delta_v = diffusion_v * laplacian_v + reaction - (feed_field + kill_field) * v
    return torch.clamp(
        torch.cat([u + delta_u, v + delta_v], dim=1),
        0,
        1,
    )