from __future__ import annotations

import logging

import awkward as ak
import hist
import numpy as np
from lgdo import lh5
from numpy.typing import ArrayLike
from scipy.spatial.transform import Rotation as rot

from revertex import utils
from revertex.core import _get_chunks, convert_output, sample_histogram

log = logging.getLogger(__name__)


def save_beta_spectrum(
    n_gen: int, in_file: str, out_file: str, seed: int | None = None, eunit: str = "keV"
) -> None:
    """Save positions generated by the function to a file.

    Parameters
    ----------
    func
        function which generates the positions, must return an `ak.Array` of positions.
    n_gen
        number of positions to generate
    in_file
        path to the CSV theory file.
    out_file
        path to the output file.
    seed
        random seed, for each chunk will be multiplied by 7.
    lunit
        The length unit returned by the function.
    **kwargs
        keyword arguments to be passed to the function.
    """
    # read input
    energies, phase_space = utils.read_input_beta_csv(in_file, delimiter=",")

    msg = f"Read beta spectrum from {in_file} E = {energies}, phase space = {phase_space} "
    log.info(msg)

    chunks = _get_chunks(n_gen, 1000_000)

    for idx, chunk in enumerate(chunks):
        # generate kinematics
        kin_ak = generate_beta_spectrum(
            chunk, energies=energies, phase_space=phase_space, seed=seed, eunit=eunit
        )
        msg = f"Generated beta kinematics {kin_ak}"
        log.debug(msg)

        # update the seed
        seed = seed * 7 if seed is not None else None

        # convert
        kin_lh5 = convert_output(kin_ak, mode="kin", eunit=eunit)

        # write
        mode = "of" if idx == 0 else "append"
        lh5.write(kin_lh5, "vtx/kin", out_file, wo_mode=mode)


def generate_beta_spectrum(
    size: int,
    *,
    energies: ArrayLike,
    phase_space: ArrayLike,
    seed: int | None = None,
    eunit: str = "keV",
) -> ak.Array:
    """Generate samples from a beta spectrum defined by a list of energies and phase space
    values.

    This function interprets the energies and phase_space as a histogram and samples
    energies from this. These are then converted into momenta. The energies should be
    the left edge of the bin of a histogram

    Parameters
    ----------
    energies
        the energy values
    spectrum
        the phase space values
    size
        number of events to generate
    seed
        random seed.
    eunit
        the unit for energy in the input file, default keV.

    Returns
    -------
    An awkward array with the sampled kinematics, in keV.
    """
    # convert energies to bin edges
    if eunit == "MeV":
        factor = 1 / 1000.0
    elif eunit == "keV":
        factor = 1
    elif eunit == "eV":
        factor = 1000
    else:
        msg = f"Only eunits keV, MeV or eV are supported not {eunit}"
        raise ValueError(msg)

    histo = hist.Hist(hist.axis.Variable(energies * factor))

    for b in range(histo.size - 2):
        histo[b] = phase_space[b]

    energy_samples = sample_histogram(histo, size, seed=seed)
    matrix = np.vstack(
        [[energy_samples, np.zeros_like(energy_samples), np.zeros_like(energy_samples)]]
    )

    mass = 511.0  # keV
    momentum = np.sqrt(energy_samples**2 + 2 * mass * energy_samples)

    rng = np.random.default_rng(seed)  # Set seed
    rand_rot = rot.random(random_state=rng)  # Random rotation

    matrix = np.vstack([[momentum, np.zeros_like(momentum), np.zeros_like(momentum)]])
    momenta = rand_rot.as_matrix() @ matrix

    return ak.Array(
        {
            "px": momenta[0, :],
            "py": momenta[1, :],
            "pz": momenta[2, :],
            "ekin": energy_samples,
            "g4_pid": np.full_like(energy_samples, 11),
        }
    )
