"""Public helper module.

Conceptually, this module centralizes trial and unit selection rules built from
normalized entry payloads.

It exists as a separate unit so decoding, movement, coding-direction, and
dropout pipelines can share the same masking semantics instead of reimplementing
them locally.

It connects entry-normalization helpers from `loading` to downstream analysis
pipelines that consume trial and unit masks.
"""

from __future__ import annotations

from typing import Any, Mapping, Optional, Set

import numpy as np

from .loading import _entry_get_any, _infer_entry_n_trials, _infer_entry_n_units, _normalize_vector_length


def get_completion_mask(e: Mapping[str, Any], completion_state: str) -> np.ndarray:
    """Return a trial-level completion mask for one entry payload.

    Parameters
    ----------
    e:
        Entry mapping containing trial-aligned payloads such as `early_lick`,
        `lick_time`, and `start_time`.
    completion_state:
        Completion mode. Supported values are `completed_trial`,
        `completed_trials`, and `early_licks`.

    Returns
    -------
    np.ndarray
        Boolean mask with shape `(n_trials,)`.
    """

    n_trials = _infer_entry_n_trials(e)
    early_lick = _normalize_vector_length(
        _entry_get_any(e, ("early_lick",), np.zeros(n_trials, dtype=bool)),
        expected_size=n_trials,
        dtype=bool,
        fill_value=False,
    )
    if completion_state in ("completed_trial", "completed_trials"):
        return ~early_lick
    if completion_state == "early_licks":
        # Use the positive lick-to-start delay to distinguish actual early licks
        # from missing or placeholder lick timestamps.
        # Absolute session timestamps (up to ~10^4 s): keep float64 so the
        # lick-after-start test matches Code_M double precision.
        lick_time = _normalize_vector_length(
            _entry_get_any(e, ("lick_time",), np.zeros(n_trials, dtype=np.float64)),
            expected_size=n_trials,
            dtype=np.float64,
            fill_value=0.0,
        )
        start_time = _normalize_vector_length(
            _entry_get_any(e, ("start_time",), np.zeros(n_trials, dtype=np.float64)),
            expected_size=n_trials,
            dtype=np.float64,
            fill_value=0.0,
        )
        return (0.0 < (lick_time - start_time)) & early_lick
    raise ValueError("completion_state must be one of: completed_trial, completed_trials, early_licks")


def get_quiet_mask(e: Mapping[str, Any], quiet_state: str) -> np.ndarray:
    """Return a quiet-state trial mask for one entry payload.

    Parameters
    ----------
    e:
        Entry mapping containing quiet-trial signals or their aliases.
    quiet_state:
        Quiet-state selector such as `Quiet_(jaw & whisker)` or `Non_quiet`.

    Returns
    -------
    np.ndarray
        Boolean mask with shape `(n_trials,)`. Missing quiet-trial signals fall
        back to all-True vectors.
    """

    n_trials = _infer_entry_n_trials(e)
    quiet_whisker = _normalize_vector_length(
        _entry_get_any(e, ("quiet_whisker", "quiet_trial_whisker_speed"), np.ones(n_trials, dtype=bool)),
        expected_size=n_trials,
        dtype=bool,
        fill_value=True,
    )
    quiet_jaw = _normalize_vector_length(
        _entry_get_any(e, ("quiet_jaw", "quiet_trial_jaw_movement"), np.ones(n_trials, dtype=bool)),
        expected_size=n_trials,
        dtype=bool,
        fill_value=True,
    )

    if quiet_state == "Quiet_(whisker_speed)":
        return quiet_whisker
    if quiet_state == "Quiet_(jaw_movement)":
        return quiet_jaw
    if quiet_state == "Quiet_(jaw & whisker)":
        return quiet_jaw & quiet_whisker
    if quiet_state == "Non_quiet":
        return ~(quiet_jaw & quiet_whisker)
    if quiet_state in ("All_trial", "All_trials"):
        return np.ones(n_trials, dtype=bool)
    raise ValueError("Invalid quietstate")


def get_trial_selection_mask(
    e: Mapping[str, Any],
    *,
    trial_type: Optional[float] = None,
    lick_state: Optional[bool] = None,
    completion_state: str = "completed_trials",
    quiet_state: str = "All_trials",
) -> np.ndarray:
    """Return a combined trial-selection mask for one entry payload.

    Parameters
    ----------
    e:
        Entry mapping containing trial-aligned fields such as `trial_type` and
        `lick_flag`.
    trial_type:
        Optional trial-type code required by the selection. When omitted, the
        mask keeps every trial type.
    lick_state:
        Optional lick-state flag applied after the completion and quiet masks.
        When omitted, the mask keeps both lick outcomes.
    completion_state:
        Completion selector forwarded to `get_completion_mask`.
    quiet_state:
        Quiet-trial selector forwarded to `get_quiet_mask`.

    Returns
    -------
    np.ndarray
        Boolean mask with shape `(n_trials,)`.
    """
    n_trials = _infer_entry_n_trials(e)
    trial = _normalize_vector_length(
        _entry_get_any(e, ("trial_type",), np.full(n_trials, np.nan, dtype=np.float32)),
        expected_size=n_trials,
        dtype=np.float32,
        fill_value=np.nan,
    )
    lick = _normalize_vector_length(
        _entry_get_any(e, ("lick_flag",), np.zeros(n_trials, dtype=bool)),
        expected_size=n_trials,
        dtype=bool,
        fill_value=False,
    )

    mask = get_completion_mask(e, completion_state) & get_quiet_mask(e, quiet_state)
    if trial_type is not None:
        mask &= trial == float(trial_type)
    if lick_state is not None:
        mask &= lick == bool(lick_state)
    return mask


def get_celltype_mask(e: Mapping[str, Any], celltype: str) -> np.ndarray:
    """Return a unit-level mask for the requested cell-type selection.

    The returned array has shape `(n_units,)`.
    """

    n_units = _infer_entry_n_units(e)
    unit_rs = _normalize_vector_length(
        _entry_get_any(e, ("unit_rs", "unit_rsUnits"), np.ones(n_units, dtype=bool)),
        expected_size=n_units,
        dtype=bool,
        fill_value=True,
    )
    unit_fs = _normalize_vector_length(
        _entry_get_any(e, ("unit_fs", "unit_fsUnits"), np.ones(n_units, dtype=bool)),
        expected_size=n_units,
        dtype=bool,
        fill_value=True,
    )

    if celltype == "RS":
        return unit_rs
    if celltype == "FS":
        return unit_fs
    if celltype in ("RS-FS", "RS_FS", "FS_RS", "FS-RS"):
        return unit_rs | unit_fs
    if celltype == "All":
        return np.ones(n_units, dtype=bool)
    raise ValueError("celltype must be one of: RS, FS, RS-FS, All")


def get_ccf_mask(
    e: Mapping[str, Any],
    area_name: str,
    area_list: Optional[Mapping[str, Set[str]]] = None,
    enable_ccf_filter: bool = True,
) -> np.ndarray:
    """Return a CCF-based unit-selection mask for one target area.

    Parameters
    ----------
    e:
        Entry mapping containing per-unit CCF locations.
    area_name:
        Target area whose allowed CCF labels should be retained.
    area_list:
        Mapping from area name to the set of accepted CCF labels.
    enable_ccf_filter:
        If False, the function returns an all-True mask.

    Returns
    -------
    np.ndarray
        Boolean mask with shape `(n_units,)`. If CCF filtering is disabled or
        no area map is available, the mask is all True.
    """

    n_units = _infer_entry_n_units(e)
    if (not enable_ccf_filter) or (area_list is None) or (area_name not in area_list):
        return np.ones(n_units, dtype=bool)

    allowed = area_list[area_name]
    unit_ccf_location = _normalize_vector_length(
        _entry_get_any(e, ("unit_ccf_location",), np.array([""] * n_units, dtype=object)),
        expected_size=n_units,
        dtype=object,
        fill_value="",
    )
    # Preserve the original unit order so the mask can be applied directly to
    # per-unit matrices without additional remapping.
    out = np.array([loc in allowed for loc in unit_ccf_location], dtype=bool)
    if out.size != n_units:
        return np.ones(n_units, dtype=bool)
    return out