import os

import equinox as eqx
import jax
import jax.numpy as jnp
import jax.random as jr
import pytest
from jaxtyping import install_import_hook


with install_import_hook("cryojax", "typeguard.typechecked"):
    import cryojax as cryojax
    import cryojax.simulator as cs
    from cryojax.image import operators as op, rfftn
    from cryojax.io import read_array_with_spacing_from_mrc


# jax.config.update("jax_numpy_dtype_promotion", "strict")
# jax.config.update("jax_numpy_rank_promotion", "raise")
jax.config.update("jax_enable_x64", True)


@pytest.fixture
def sample_mrc_path():
    return os.path.join(os.path.dirname(__file__), "data", "3j9g_potential_ps4_4.mrc")


@pytest.fixture
def sample_subunit_mrc_path():
    return os.path.join(
        os.path.dirname(__file__), "data", "3j9g_subunit_potential_ps4_4.mrc"
    )


@pytest.fixture
def sample_pdb_path():
    return os.path.join(os.path.dirname(__file__), "data", "1uao.pdb")


@pytest.fixture
def toy_gaussian_cloud():
    atom_positions = jnp.array(
        [
            [0.0, 0.0, 0.0],
            [0.0, 0.0, 1.0],
            [1.0, 0.0, 0.0],
            [0.0, 1.0, 0.0],
        ]
    )
    num_atoms = atom_positions.shape[0]
    ff_a = jnp.array(
        num_atoms
        * [
            [1.0, 0.5],
        ]
    )

    ff_b = jnp.array(
        num_atoms
        * [
            [0.3, 0.2],
        ]
    )

    n_voxels_per_side = (128, 128, 128)
    voxel_size = 0.05
    return (atom_positions, ff_a, ff_b, n_voxels_per_side, voxel_size)


@pytest.fixture
def pixel_size():
    return 4.4


@pytest.fixture
def config(pixel_size):
    return cs.ImageConfig((65, 66), pixel_size, pad_scale=1.1)


@pytest.fixture
def integrator():
    return cs.FourierSliceExtract(interpolation_order=1)


@pytest.fixture
def potential(sample_mrc_path):
    real_voxel_grid, voxel_size = read_array_with_spacing_from_mrc(sample_mrc_path)
    return cs.FourierVoxelGridPotential.from_real_voxel_grid(
        real_voxel_grid, voxel_size, pad_scale=1.3
    )


@pytest.fixture
def filters(config):
    return op.LowpassFilter(config.wrapped_padded_frequency_grid_in_pixels.get())


@pytest.fixture
def masks(config):
    return op.CircularMask(
        config.wrapped_padded_coordinate_grid_in_angstroms.get(),
        radius=20 * float(config.pixel_size),
    )


@pytest.fixture
def instrument():
    voltage_in_kilovolts = 300.0
    return cs.Instrument(
        voltage_in_kilovolts,
        optics=cs.WeakPhaseOptics(cs.CTF()),
        dose=cs.ElectronDose(electrons_per_angstrom_squared=1000.0),
        detector=cs.GaussianDetector(cs.IdealDQE(fraction_detected_electrons=1.0)),
    )


@pytest.fixture
def pose():
    return cs.EulerAnglePose(
        view_phi=30.0,
        view_theta=100.0,
        view_psi=-10.0,
        offset_x_in_angstroms=10.0,
        offset_y_in_angstroms=-5.0,
    )


@pytest.fixture
def specimen(potential, integrator, pose):
    return cs.Specimen(potential, integrator, pose)


@pytest.fixture
def solvent():
    return cs.GaussianIce(op.Constant(0.001**2))


@pytest.fixture
def noiseless_model(config, specimen, instrument):
    instrument = eqx.tree_at(lambda ins: ins.detector, instrument, None)
    return cs.ImagePipeline(config=config, specimen=specimen, instrument=instrument)


@pytest.fixture
def noisy_model(config, specimen, instrument, solvent):
    return cs.ImagePipeline(
        config=config,
        specimen=specimen,
        instrument=instrument,
        solvent=solvent,
    )


@pytest.fixture
def filtered_model(config, specimen, instrument, solvent, filters):
    return cs.ImagePipeline(
        config=config,
        specimen=specimen,
        instrument=instrument,
        solvent=solvent,
        filter=filters,
    )


@pytest.fixture
def filtered_and_masked_model(config, specimen, instrument, solvent, filters, masks):
    return cs.ImagePipeline(
        config=config,
        specimen=specimen,
        instrument=instrument,
        solvent=solvent,
        filter=filters,
        mask=masks,
    )


@pytest.fixture
def test_image(noisy_model):
    image = noisy_model.sample(jr.PRNGKey(1234))
    return rfftn(image)
