import logging
import os
from functools import partial
from typing import Literal

import cryojax.simulator as cxs
import equinox as eqx
import jax
import jax.numpy as jnp
from cryojax.dataset import (
    RelionParticleDataset,
    RelionParticleParameterFile,
    simulate_particle_stack,
)
from cryojax.ndimage import CircularCosineMask, FourierGaussian
from cryojax.rotations import SO3
from jaxtyping import Array, Float, PRNGKeyArray

from ..internal._config_validators import DatasetSimulatorConfig
from ..io import load_gmm_volume_parametrization
from ..simulator._image_simulation import simulate_image_with_white_gaussian_noise


def make_relion_parameter_file(
    key: PRNGKeyArray, config: DatasetSimulatorConfig
) -> RelionParticleParameterFile:
    """

    This functions writes to disk a starfile containing the
    generated particle parameters. The starfile is saved to the path
    specified in the configuration. The function also returns a
    `cryojax.dataset.RelionParticleParameterFile` object that can be used to read the
    starfile and access the particle parameters.

    **Arguments:**
        key: JAX random key for generating random numbers.
        config: Configuration object containing the parameters for generating the dataset.
    **Returns:**
        parameter_file: A RelionParticleParameterFile object containing
        the generated particle parameters.
    """

    config_dict = dict(config.model_dump())
    logging.info("Generating Starfile...")
    key, *subkeys = jax.random.split(key, config_dict["number_of_images"] + 1)
    particle_parameters = _make_particle_parameters(jnp.array(subkeys), config_dict)

    logging.info(
        "Starfile generated. Saved to {}".format(config_dict["path_to_starfile"])
    )

    parameter_file = RelionParticleParameterFile(
        path_to_starfile=config_dict["path_to_starfile"],
        mode="w",
        exists_ok=config_dict["overwrite"],
    )
    parameter_file.append(particle_parameters)

    return parameter_file


def simulate_relion_dataset(config: DatasetSimulatorConfig) -> RelionParticleDataset:
    os.makedirs(config.path_to_relion_project, exist_ok=True)

    seed = config.rng_seed
    key = jax.random.key(seed)

    key_param, key_stack = jax.random.split(key)
    parameter_file = make_relion_parameter_file(key_param, config)

    # dumping so serialization happens
    config_dict = dict(config.model_dump())

    volumes = load_gmm_volume_parametrization(
        config_dict["atomic_models_params"]["path_to_atomic_models"],
        selection_string=config_dict["atomic_models_params"]["atom_selection"],
        loads_b_factors=config_dict["atomic_models_params"]["loads_b_factors"],
    )

    mask = CircularCosineMask(
        coordinate_grid=parameter_file[0]["image_config"].coordinate_grid_in_pixels,
        radius=config_dict["mask_radius"],
        rolloff_width=config_dict["mask_rolloff_width"],
    )

    return _simulate_relion_dataset(
        key=key_stack,
        parameter_file=parameter_file,
        path_to_relion_project=config_dict["path_to_relion_project"],
        images_per_file=config_dict["images_per_file"],
        volumes=volumes,
        ensemble_probabilities=config_dict["atomic_models_params"][
            "atomic_models_probabilities"
        ],
        mask=mask,
        noise_snr_range=config_dict["noise_snr"],
        data_sign=config_dict["data_sign"],
        overwrite=config_dict["overwrite"],
        batch_size=config_dict["batch_size_for_generation"],
    )


def _simulate_relion_dataset(
    key: PRNGKeyArray,
    parameter_file: RelionParticleParameterFile,
    path_to_relion_project: str,
    images_per_file: int,
    volumes: tuple[cxs.AbstractVolumeRepresentation],
    ensemble_probabilities: Float[Array, " n_volumes"],
    mask: CircularCosineMask,
    noise_snr_range: list[Float],
    data_sign: Literal["dark-on-light", "light-on-dark"],
    *,
    overwrite: bool = False,
    batch_size: int = 1,
) -> RelionParticleDataset:
    assert len(volumes) == len(ensemble_probabilities), (
        "The number of volumes must be equal to the number of ensemble probabilities."
        f" {len(volumes)} != {len(ensemble_probabilities)}"
    )

    # generate random keys
    key_snr, key_noise, key_volumes = jax.random.split(key, 3)

    # Make sure the ensemble probabilities sum to 1
    ensemble_probabilities = jnp.array(ensemble_probabilities)
    ensemble_probabilities /= jnp.sum(ensemble_probabilities)

    # Generate parameters for each image
    snr_per_image = jax.random.uniform(
        key_snr,
        (len(parameter_file),),
        minval=noise_snr_range[0],
        maxval=noise_snr_range[1],
    )

    keys_per_image = jax.random.split(key_noise, len(parameter_file))

    ensemble_indices_per_image = jax.random.choice(
        key_volumes,
        a=len(ensemble_probabilities),
        shape=(len(parameter_file),),
        p=ensemble_probabilities,
        replace=True,
    )

    # Write metadata to disk
    jnp.savez(
        os.path.join(path_to_relion_project, "metadata.npz"),
        snr_per_image=snr_per_image,
        ensemble_indices_per_image=ensemble_indices_per_image,
    )
    logging.info("Metadata: noise variance and ensemble indices saved to:")
    logging.info(f"  {os.path.join(path_to_relion_project, 'metadata.npz')}")

    # Bundle arguments and write images
    data_sign_factor = -1.0 if data_sign == "dark-on-light" else 1.0
    constant_args = (volumes, mask, data_sign_factor)
    per_particle_args = (
        keys_per_image,
        ensemble_indices_per_image,
        snr_per_image,
    )

    relion_dataset = RelionParticleDataset(
        parameter_file=parameter_file,
        path_to_relion_project=path_to_relion_project,
        mode="w",
        mrcfile_settings={"overwrite": overwrite},
    )

    simulate_particle_stack(
        dataset=relion_dataset,
        simulate_fn=simulate_image_with_white_gaussian_noise,
        constant_args=constant_args,
        per_particle_args=per_particle_args,
        images_per_file=images_per_file,
        batch_size=batch_size,
        overwrite=overwrite,
    )
    logging.info(f"Images generated. Saved to {path_to_relion_project}")
    logging.info("Simulated dataset generation complete.")
    return relion_dataset


@partial(eqx.filter_vmap, in_axes=(0, None))
def _make_particle_parameters(key: PRNGKeyArray, config: dict) -> dict:
    """
    WARNING: this function assumes the `config` has been validated
    by `cryojax_ensemble_refinement.internal.GeneratorConfig`.
    Skipping this step could lead to unexpected behavior.
    """
    image_config = cxs.BasicImageConfig(
        shape=(config["box_size"], config["box_size"]),
        pixel_size=config["pixel_size"],
        voltage_in_kilovolts=config["voltage_in_kilovolts"],
    )
    # Generate random parameters

    # Pose
    # ... instantiate rotations
    key, subkey = jax.random.split(key)  # split the key to use for the next random number

    rotation = SO3.sample_uniform(subkey)
    key, subkey = jax.random.split(key)  # do this everytime you use a key!!

    # ... now in-plane translation
    offset_x_in_angstroms = jax.random.uniform(
        subkey,
        (1,),
        minval=config["offset_x_in_angstroms"][0],
        maxval=config["offset_x_in_angstroms"][1],
    )
    key, subkey = jax.random.split(key)

    offset_y_in_angstroms = jax.random.uniform(
        subkey,
        (1,),
        minval=config["offset_y_in_angstroms"][0],
        maxval=config["offset_y_in_angstroms"][1],
    )
    key, subkey = jax.random.split(key)

    in_plane_offset_in_angstroms = jnp.concatenate(
        [offset_x_in_angstroms, offset_y_in_angstroms]
    )
    # ... convert 2D in-plane translation to 3D, setting the out-of-plane translation to
    # zero
    offset_in_angstroms = jnp.pad(in_plane_offset_in_angstroms, ((0, 1),))
    # ... build the pose
    pose = cxs.EulerAnglePose.from_rotation_and_translation(rotation, offset_in_angstroms)

    # CTF Parameters
    # ... defocus
    defocus_in_angstroms = jax.random.uniform(
        subkey,
        (),
        minval=config["defocus_in_angstroms"][0],
        maxval=config["defocus_in_angstroms"][1],
    )
    key, subkey = jax.random.split(key)

    astigmatism_in_angstroms = jax.random.uniform(
        subkey,
        (),
        minval=config["astigmatism_in_angstroms"][0],
        maxval=config["astigmatism_in_angstroms"][1],
    )
    key, subkey = jax.random.split(key)

    astigmatism_angle = jax.random.uniform(
        subkey,
        (),
        minval=config["astigmatism_angle_in_degrees"][0],
        maxval=config["astigmatism_angle_in_degrees"][1],
    )
    key, subkey = jax.random.split(key)

    phase_shift = jax.random.uniform(
        subkey, (), minval=config["phase_shift"][0], maxval=config["phase_shift"][1]
    )
    key, subkey = jax.random.split(key)

    b_factor = jax.random.uniform(
        subkey,
        (),
        minval=config["envelope_b_factor"][0],
        maxval=config["envelope_b_factor"][1],
    )
    # no more random numbers needed

    # now generate your non-random values
    spherical_aberration_in_mm = config["spherical_aberration_in_mm"]
    amplitude_contrast_ratio = config["amplitude_contrast_ratio"]
    ctf_scale_factor = config["ctf_scale_factor"]

    b_factor = jnp.clip(b_factor, 1e-16, None)
    envelope = FourierGaussian(b_factor=b_factor, amplitude=ctf_scale_factor)

    # ... build the CTF
    transfer_theory = cxs.ContrastTransferTheory(
        ctf=cxs.AstigmaticCTF(
            defocus_in_angstroms=defocus_in_angstroms,
            astigmatism_in_angstroms=astigmatism_in_angstroms,
            astigmatism_angle=astigmatism_angle,
            spherical_aberration_in_mm=spherical_aberration_in_mm,
        ),
        amplitude_contrast_ratio=amplitude_contrast_ratio,
        phase_shift=phase_shift,
        envelope=envelope,
    )

    relion_particle_parameters = {
        "image_config": image_config,
        "pose": pose,
        "transfer_theory": transfer_theory,
    }
    return relion_particle_parameters
