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.ndimage import CircularCosineMask, FourierGaussian
from cryojax.rotations import SO3
from cryospax import (
    RelionParticleDataset,
    RelionParticleParameterFile,
    simulate_particle_stack,
)
from jaxtyping import Array, Float, PRNGKeyArray

from ..internal._config_validators import DatasetSimulatorConfig
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
    `cryospax.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",
        exist_ok=config_dict["overwrite"],
        options=dict(
            broadcasts_image_config=True
        ),  # Change to False when cryospax is updated
    )
    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 = tuple(
        [
            cxs.load_tabulated_volume(
                filename,
                output_type=cxs.GaussianMixtureVolume,
                include_b_factors=config_dict["atomic_models_params"]["loads_b_factors"],
                selection_string=config_dict["atomic_models_params"]["atom_selection"],
            )
            for filename in config_dict["atomic_models_params"]["path_to_atomic_models"]
        ]
    )

    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.GaussianMixtureVolume, ...],
    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
