from __future__ import annotations

import logging

import awkward as ak
import numpy as np
from lgdo import VectorOfVectors

from ..optmap import convolve
from ..units import units_conv_ak

log = logging.getLogger(__name__)


def load_optmap_all(map_file: str) -> convolve.OptmapForConvolve:
    """Load an optical map file for later use with :py:func:`detected_photoelectrons`."""
    return convolve.open_optmap(map_file)


def load_optmap(map_file: str, spm_det_uid: int) -> convolve.OptmapForConvolve:
    """Load an optical map file for later use with :py:func:`detected_photoelectrons`."""
    return convolve.open_optmap_single(map_file, spm_det_uid)


def _nested_unflatten(data: ak.Array, lengths: ak.Array):
    return ak.unflatten(ak.unflatten(ak.flatten(data), ak.flatten(lengths)), ak.num(lengths))


def corrected_photoelectrons(
    simulated_pe: ak.Array,
    simulated_uids: ak.Array,
    data_pe: ak.Array,
    data_uids: ak.Array,
    *,
    seed: int | None = None,
) -> tuple[ak.Array, ak.Array]:
    r"""Add a correction to the observed number of photoelectrons (p.e.) using forced trigger data.

    For every simulated event a corresponding forced trigger event in data is chosen
    and the resulting number of p.e. for each channel (i) is:

     .. math::

        n_i = n_{\text{sim},i} + n_{\text{data},i}

    .. warning::
       The number of supplied forced trigger events in data should ideally be
       more than that in the simulations. If this is not the case and "allow_data_reuse"
       is True then some data events will be used multiple times. This introduces
       a small amount of correlation between the simulated events, but is probably acceptable
       in most circumstances.

    Parameters
    ----------
    simulated_pe
        The number of number of detected pe per sipm channel.
    simulated_uids
        The unique identifier (uid) for each sipm hit.
    data_pe
        The collection of forced trigger pe.
    data_uids
        The uids for each forced trigger event.
    seed
        Seed for random number generator

    Returns
    -------
    a tuple of the corrected pe and sipm uids.
    """
    rand = np.random.default_rng(seed=seed)
    rand_ints = rand.integers(0, len(data_pe), size=len(simulated_pe))

    selected_data_pe = data_pe[rand_ints]
    selected_data_uids = data_uids[rand_ints]

    # combine sims with data
    pe_tot = ak.concatenate([simulated_pe, selected_data_pe], axis=1)
    uid_tot = ak.concatenate([simulated_uids, selected_data_uids], axis=1)

    # sort by uid
    order = ak.argsort(uid_tot)
    pe_tot = pe_tot[order]
    uid_tot = uid_tot[order]

    # add an extra axis
    n = ak.run_lengths(uid_tot)

    # add another dimension
    pe_tot = _nested_unflatten(pe_tot, n)
    uid_tot = _nested_unflatten(uid_tot, n)

    # sum pe and take the first uid (should all be the same)
    corrected_pe = ak.sum(pe_tot, axis=-1)
    uid_tot = ak.fill_none(ak.firsts(uid_tot, axis=-1), np.nan)

    return corrected_pe, uid_tot


def detected_photoelectrons(
    num_scint_ph: ak.Array,
    particle: ak.Array,
    time: ak.Array,
    xloc: ak.Array,
    yloc: ak.Array,
    zloc: ak.Array,
    optmap: convolve.OptmapForConvolve,
    material: str,
    spm_detector: str,
    map_scaling: float = 1,
    map_scaling_sigma: float = 0,
) -> VectorOfVectors:
    """Derive the number and arrival times of detected photoelectrons (p.e.) from scintillator hits using an optical map.

    .. deprecated :: 0.8.5
        Use the other, more fine-granular and split processors
        :func:`number_of_detected_photoelectrons` and :func:`photoelectron_times` to
        replace this legacy processor.

    Parameters
    ----------
    num_scint_ph
        array of emitted scintillation photons, as generated by
        :func:`emitted_scintillation_photons`.
    particle
        array of particle PDG IDs of scintillation events.
    time
        array of timestamps of scintillation events.
    xloc
        array of x coordinate position of scintillation events.
    yloc
        array of y coordinate position of scintillation events.
    zloc
        array of z coordinate position of scintillation events.
    optmap
        the optical map loaded via py:func:`load_optmap`.
    material
        scintillating material name.
    spm_detector
        SiPM detector name as used in the optical map.
    map_scaling
        scale the detection probability in the map for this detector by this factor.
    map_scaling_sigma
        if larger than zero, sample the used scaling factor for each (reshaped) event
        from a normal distribution with this standard deviation.
    """
    hits = ak.Array(
        {
            "num_scint_ph": units_conv_ak(num_scint_ph, "dimensionless"),
            "particle": units_conv_ak(particle, "dimensionless"),
            "time": units_conv_ak(time, "ns"),
            "xloc": units_conv_ak(xloc, "m"),
            "yloc": units_conv_ak(yloc, "m"),
            "zloc": units_conv_ak(zloc, "m"),
        }
    )

    scint_mat_params = convolve._get_scint_params(material)
    pe = convolve.iterate_stepwise_depositions_pois(
        hits, optmap, scint_mat_params, spm_detector, map_scaling, map_scaling_sigma
    )

    return VectorOfVectors(pe, attrs={"units": "ns"})


def emitted_scintillation_photons(
    edep: ak.Array, particle: ak.Array, material: str
) -> VectorOfVectors:
    """Derive the number of emitted scintillation photons from scintillator hits.

    Parameters
    ----------
    edep
        array of deposited energy in scintillation events.
    particle
        array of particle PDG IDs of scintillation events.
    material
        scintillating material name.
    """
    hits = ak.Array(
        {
            "edep": units_conv_ak(edep, "keV"),
            "particle": units_conv_ak(particle, "dimensionless"),
        }
    )

    scint_mat_params = convolve._get_scint_params(material)
    ph = convolve.iterate_stepwise_depositions_scintillate(hits, scint_mat_params)
    return VectorOfVectors(ph)


def number_of_detected_photoelectrons(
    xloc: ak.Array,
    yloc: ak.Array,
    zloc: ak.Array,
    num_scint_ph: ak.Array,
    optmap: convolve.OptmapForConvolve,
    spm_detector: str,
    map_scaling: float = 1,
    map_scaling_sigma: float = 0,
) -> VectorOfVectors:
    """Derive the number of detected photoelectrons (p.e.) from scintillator hits using an optical map.

    Parameters
    ----------
    xloc
        array of x coordinate position of scintillation events.
    yloc
        array of y coordinate position of scintillation events.
    zloc
        array of z coordinate position of scintillation events.
    num_scint_ph
        array of emitted scintillation photons, as generated by
        :func:`emitted_scintillation_photons`.
    detp
    """
    hits = ak.Array(
        {
            "xloc": units_conv_ak(xloc, "m"),
            "yloc": units_conv_ak(yloc, "m"),
            "zloc": units_conv_ak(zloc, "m"),
            "num_scint_ph": units_conv_ak(num_scint_ph, "dimensionless"),
        }
    )

    ph = convolve.iterate_stepwise_depositions_numdet(
        hits, optmap, spm_detector, map_scaling, map_scaling_sigma
    )
    return VectorOfVectors(ph)


def photoelectron_times(
    num_det_ph: ak.Array,
    particle: ak.Array,
    time: ak.Array,
    material: str,
) -> VectorOfVectors:
    """Derive the arrival times of scintillation photons.

    Parameters
    ----------
    num_det_ph
        array of detected scintillation photons, as generated by
        :func:`emitted_scintillation_photons`.
    particle
        array of particle PDG IDs of scintillation events.
    time
        array of timestamps of scintillation events.
    material
        scintillating material name.
    """
    hits = ak.Array(
        {
            "num_det_ph": units_conv_ak(num_det_ph, "dimensionless"),
            "particle": units_conv_ak(particle, "dimensionless"),
            "time": units_conv_ak(time, "ns"),
        }
    )

    scint_mat_params = convolve._get_scint_params(material)
    pe = convolve.iterate_stepwise_depositions_times(hits, scint_mat_params)

    return VectorOfVectors(pe, attrs={"units": "ns"})
