from __future__ import annotations

import importlib.metadata
from typing import Any

import numpy as np

import pigreads as pig


def test_version():
    assert importlib.metadata.version("pigreads") == pig.__version__


def test_deltas():
    z, y, x = np.meshgrid(
        [0], np.arange(0, 4, 0.2), np.linspace(0, 1, 5), indexing="ij"
    )
    dz, dy, dx = pig.deltas(z, y, x)
    assert dz == 1
    assert dy == 0.2
    assert dx == 0.25
    assert pig.delta(y, ax=-2) == dy


def test_ithildin():
    z, y, x = np.mgrid[0:1, 0:3, 0:3]
    dz, dy, dx = pig.deltas(z, y, x)
    inhom = np.ones_like(x, dtype=int)
    models = pig.Models()
    models.add("marcotte2017dynamical")
    assert models[0]["variables"] == {"u": 0.0, "v": 0.0}
    states: np.ndarray[Any, Any] = models.resting_states(inhom, Nframes=2).astype(
        np.float32
    )
    weights = pig.weights(dz, dy, dx, inhom, diffusivity=0)
    Nt, dt = 3, 0.01
    states[1] = pig.run(models, inhom, weights, states[0], Nt=Nt, dt=dt)
    log, variables = pig.to_ithildin(Nt * dt, dt, dz, dy, dx, models, states, inhom)
    del log["Start date"]
    del log["Simulation parameters"]["Ithildin library version"]
    del log["Simulation parameters"]["Serial number"]
    assert (
        str(log)
        == r"{'Ithildin log version': 2, 'Simulation parameters': {'Timestep dt': 0.01, 'Frame duration': 0.03, 'Number of frames to take': 2, 'Name of simulation series': 'pigreads'}, 'Geometry parameters': {'Number of dimensions': 3, 'Voxel size': [1.0, 1.0, 1.0], 'Domain size': [3, 3, 1]}, 'Model parameters': {'Model type': 'Smoothed Karma (Marcotte & Grigoriev 2017)', 'Class': 'marcotte2017dynamical', 'Citation': 'https://doi.org/10.1063/1.5003259\nhttps://doi.org/10.1063/1.4915143\nhttps://doi.org/10.1103/PhysRevLett.71.1103\nhttps://doi.org/10.1063/1.166024', 'Parameters': {'diffusivity_u': 4.0062, 'diffusivity_v': 0.20031, 'beta': 1.389, 'eps': 0.01, 'ustar': 1.5415}, 'Initial values': {'u': 0.0, 'v': 0.0}, 'Variable names': ['u', 'v'], 'Number of vars': 2}}"
    )
    assert " ".join(variables.keys()) == "u v inhom"
    assert variables["u"].shape == (2, 1, 3, 3)
    assert variables["v"].shape == (2, 1, 3, 3)
