from collections.abc import Sequence

from nested_mapping import NestedMapping
from numba import njit
from numpy import empty, isnan
from numpy.typing import NDArray


@njit
def weekly_to_daily(array: NDArray) -> NDArray:
    ret = empty(array.shape[0] * 7)

    for i in range(7):
        ret[i::7] = array

    return ret


@njit
def inperiod_to_daily(array: NDArray, ndays: NDArray) -> NDArray:
    ndays_total = ndays.sum()
    ret = empty(ndays_total)

    i = 0
    for ndays_i, arr in zip(ndays, array):
        ret[i : i + ndays_i] = arr
        i += ndays_i

    return ret


@njit
def periods_to_days(array: NDArray, ndays: NDArray) -> NDArray:
    ndays_total = ndays.sum()
    ret = empty(ndays_total, dtype="i")

    i = 0
    for arr, ndays_p in zip(array, ndays):
        for day in range(ndays_p):
            ret[i] = arr + day
            i += 1

    return ret


@njit
def weeks_to_days(array: NDArray) -> NDArray:
    ret = empty(array.shape[0] * 7)

    for i in range(7):
        ret[i::7] = array
        ret[i::7] += i

    return ret


def refine_neutrino_rate_data(
    source: NestedMapping,
    target: NestedMapping,
    *,
    reactors: Sequence[str],
    periods: Sequence[int] = (6, 8, 7),
    clean_source: bool = True,
) -> None:
    """Refine neutrino rate data.

    Does the following items:
        - checks input consistency
        - split arrays based on isotope number and period name
        - provides similar data on a daily basis
        - splits into multiple arrays based on the reactor index

    Source should contain the following fields:
        - period : int
              period number
        - day : int
              number of the first day of the period, relative to some common day.
        - ndays : int
              length of the period in days. This function can deal only with periods of
              equal length and with no gaps.
        - core : int
              number of the current reactor (1 to 6).
        - ndet : int
              number of active detectors during the period (6, 8 or 7).
        - neutrino_rate_per_s : dict[str, NDArray]
              total neutrino rate per second for each reactor core.

    The target dictionary will be populated as follows:
    ```python
    target = {
        "days": array(),
        "neutrino_rate": {
            "period_name": {
                "reactor_name": array(),
                ...
            },
        }
    }
    ```
    with all the arrays of same lengths.

    Parameters
    ----------
    source : NestedMapping
        storage/record/mapping with source data.
    target : NestedMapping
        storage to write target arrays to.
    reactors : Sequence[str]
        reactor names to use when writing data.
    isotopes : Sequence[str]
        names of the isotopes to read and write neutrino rate.
    periods : Sequence[int], default=(6, 8, 7)
        periods (number of active detectors) to select.
    clean_source : bool, default=True
        if True, remove used data from source.
    """
    for corename in reactors:
        period = source["period", corename]
        day = source["day", corename]
        ndays = source["n_days", corename]
        n_det_mask = source["n_det_mask", corename]

        neutrino_rate_per_s = source["neutrino_rate_per_s"][corename]

        step = period[1:] - period[:-1]
        assert (step == 1).all(), "Expect reactor data for with distinct period, no gaps"

        target["days"] = (days_storage := {})
        for period in periods:
            period_bit = {
                    6: 0b001,
                    8: 0b010,
                    7: 0b100,
                    }[period]
            mask_period = (n_det_mask & period_bit)>0
            periodname = f"{period}AD"

            mask = mask_period
            key = (
                periodname,
                corename,
            )
            ndays_p = ndays[mask]
            target[("neutrino_rate_per_s",) + key] = inperiod_to_daily(neutrino_rate_per_s[mask], ndays_p)

            days = periods_to_days(day[mask], ndays_p)
            if isnan(days).any():
                raise RuntimeError()
            days_stored = days_storage.setdefault(periodname, days)
            if days is not days_stored:
                assert all(days == days_stored)

    for key, array in target.walkjoineditems():
        if isnan(array).any():
            raise ValueError(f"Invalid refined reactor data for {key}")

    if clean_source:
        for key in tuple(source.walkkeys()):
            source.delete_with_parents(key)
