"""Internal support module.

Conceptually, this module turns irregular scalar, vector, and object payloads
into stable NumPy-friendly shapes.

It exists as a separate unit so higher-level loaders can rely on one shared
coercion layer instead of repeating shape repair and alias-handling logic.

It connects generic core helpers to both NWB-specific and processed-data
loading modules.
"""

from __future__ import annotations

from typing import Any, Iterable, List, Mapping, Optional

import numpy as np

from .core import to_1d, unwrap_scalar_obj


__all__ = [
    "_entry_get_any",
    "_entry_scalar_int",
    "_normalize_vector_length",
    "_infer_entry_n_trials",
    "_infer_entry_n_units",
    "align_to_len",
    "make_picklable",
    "sanitize_time_axis",
    "to_2d_xyz_rows",
    "to_numpy_compact_series",
]


def _entry_get_any(
    entry: Mapping[str, Any],
    keys: Iterable[str],
    default: Any = None,
) -> Any:
    """Return the first non-``None`` entry value for a list of aliases."""

    for key in keys:
        if key in entry and entry[key] is not None:
            return entry[key]
    return default



def _entry_scalar_int(
    entry: Mapping[str, Any],
    keys: Iterable[str],
    default: int = 0,
) -> int:
    """Read a scalar integer from an entry mapping with alias support."""

    raw = _entry_get_any(entry, keys, default)
    arr = np.asarray(unwrap_scalar_obj(raw)).reshape(-1)
    if arr.size == 0:
        return int(default)
    try:
        return int(arr[0])
    except Exception:
        return int(default)



def _normalize_vector_length(
    values: Any,
    expected_size: int,
    dtype: Optional[np.dtype] = None,
    fill_value: Any = 0,
) -> np.ndarray:
    """Coerce a vector-like payload to one target length.

    Short vectors are padded with `fill_value` and long vectors are truncated.
    The result is always one-dimensional with length `expected_size`.
    """

    if expected_size < 0:
        raise ValueError("expected_size must be >= 0")

    arr = to_1d(values, dtype=dtype) if dtype is not None else to_1d(values)
    if arr.size == expected_size:
        return arr

    if dtype is None:
        arr = np.asarray(arr)
        out = np.empty(expected_size, dtype=arr.dtype if arr.size else np.asarray([fill_value]).dtype)
    else:
        out = np.empty(expected_size, dtype=dtype)
    out[:] = fill_value
    keep = min(expected_size, arr.size)
    if keep > 0:
        out[:keep] = arr[:keep]
    return out


def align_to_len(v: Any, n: int, dtype: Any) -> np.ndarray:
    """Pad or truncate a vector-like payload to a target length.

    Parameters
    ----------
    v:
        Scalar or vector-like payload to coerce.
    n:
        Target output length.
    dtype:
        Output dtype.

    Returns
    -------
    np.ndarray
        One-dimensional array with length `n`. Scalar payloads are broadcast
        when `n > 1`.
    """

    a = np.asarray(unwrap_scalar_obj(v), dtype=dtype).reshape(-1)
    if a.size == int(n):
        return a
    if a.size == 0:
        return np.zeros(int(n), dtype=dtype)
    if a.size == 1 and int(n) > 1:
        return np.repeat(a, int(n)).astype(dtype)
    if a.size < int(n):
        out = np.zeros(int(n), dtype=dtype)
        out[: a.size] = a
        return out
    return a[: int(n)]



def _infer_entry_n_trials(entry: Mapping[str, Any]) -> int:
    """Infer the number of trials from common entry payload aliases."""

    n_trials = _entry_scalar_int(entry, ("n_trials",), default=0)
    if n_trials > 0:
        return n_trials

    for keys in (
        ("trial", "trial_type"),
        ("lick", "lick_flag"),
        ("stim", "whisker_stim"),
        ("early_lick",),
        ("quiet_whisker", "quiet_trial_whisker_speed"),
        ("quiet_jaw", "quiet_trial_jaw_movement"),
        ("lick_time",),
        ("start_time",),
    ):
        raw = _entry_get_any(entry, keys, None)
        if raw is not None:
            size = to_1d(raw).size
            if size > 0:
                return int(size)

    spike_counts = _entry_get_any(entry, ("spike_counts",), None)
    if spike_counts is not None:
        arr = np.asarray(unwrap_scalar_obj(spike_counts))
        if arr.ndim == 3:
            return int(arr.shape[1])

    return 0



def _infer_entry_n_units(entry: Mapping[str, Any]) -> int:
    """Infer the number of units from common entry payload aliases."""

    n_units = _entry_scalar_int(entry, ("n_units",), default=0)
    if n_units > 0:
        return n_units

    for keys in (
        ("unit_rs", "unit_rsUnits"),
        ("unit_fs", "unit_fsUnits"),
        ("unit_ccf_location",),
        ("GlobalclusterID",),
    ):
        raw = _entry_get_any(entry, keys, None)
        if raw is not None:
            size = to_1d(raw).size
            if size > 0:
                return int(size)

    spike_counts = _entry_get_any(entry, ("spike_counts",), None)
    if spike_counts is not None:
        arr = np.asarray(unwrap_scalar_obj(spike_counts))
        if arr.ndim == 3:
            return int(arr.shape[2])

    return 0



def to_numpy_compact_series(series: Any) -> np.ndarray:
    """Convert a pandas-like series to a compact NumPy representation."""

    import pandas as pd  # local import to keep base module lightweight

    arr = series.to_numpy()

    if pd.api.types.is_bool_dtype(series):
        return np.asarray(arr, dtype=bool)

    if pd.api.types.is_float_dtype(series):
        return np.asarray(arr, dtype=np.float32)

    if pd.api.types.is_integer_dtype(series):
        arr64 = np.asarray(arr, dtype=np.int64)
        if arr64.size == 0:
            return arr64.astype(np.int16)
        mn = int(np.nanmin(arr64))
        mx = int(np.nanmax(arr64))
        if np.iinfo(np.int16).min <= mn and mx <= np.iinfo(np.int16).max:
            return arr64.astype(np.int16)
        if np.iinfo(np.int32).min <= mn and mx <= np.iinfo(np.int32).max:
            return arr64.astype(np.int32)
        return arr64

    out: List[Any] = []
    for value in arr:
        if isinstance(value, bytes):
            out.append(value.decode("utf-8", errors="ignore").strip())
        else:
            out.append(value)
    return np.asarray(out, dtype=object)



def to_2d_xyz_rows(values: Any, n_rows: int) -> np.ndarray:
    """Coerce mixed xyz payloads to a stable `(n_rows, 3)` float32 matrix.

    Parameters
    ----------
    values:
        Dense numeric payload or object payload containing xyz-like rows.
    n_rows:
        Expected number of rows in the repaired matrix.

    Returns
    -------
    np.ndarray
        Float32 array with shape `(n_rows, 3)`.
    """

    raw = unwrap_scalar_obj(values)
    try:
        raw_arr = np.asarray(raw)
    except Exception:
        raw_arr = np.asarray(raw, dtype=object)

    if raw_arr.dtype == object:
        # Some serialized payloads store one xyz-like row per object entry
        # rather than one dense numeric matrix.
        flat = raw_arr.reshape(-1)
        rows = []
        ok = True
        for value in flat:
            try:
                current = np.asarray(unwrap_scalar_obj(value), dtype=np.float32).reshape(-1)
            except Exception:
                ok = False
                break
            if current.size == 0:
                current = np.zeros(3, dtype=np.float32)
            if current.size < 3:
                current = np.pad(current, (0, 3 - current.size), mode="constant")
            rows.append(current[:3])
        if ok and rows:
            out = np.vstack(rows).astype(np.float32, copy=False)
        else:
            out = np.zeros((max(int(n_rows), 0), 3), dtype=np.float32)
    else:
        out = np.asarray(raw, dtype=np.float32)
        if out.ndim == 0:
            out = out.reshape(1, 1)
        elif out.ndim == 1:
            if out.size == int(n_rows) * 3:
                out = out.reshape(int(n_rows), 3)
            else:
                out = np.repeat(out.reshape(-1, 1), 3, axis=1)
        if out.ndim == 1:
            out = out.reshape(-1, 1)
        if out.shape[1] < 3:
            out = np.pad(out, ((0, 0), (0, 3 - out.shape[1])), mode="constant")
        out = out[:, :3]
        if out.shape[0] != int(n_rows):
            # Only use resize as a final shape repair when upstream payloads do
            # not match the declared session length.
            out = np.resize(out, (int(n_rows), 3)).astype(np.float32)

    return out



def sanitize_time_axis(values: Any) -> np.ndarray:
    """Return a finite 1D float32 time axis.

    Parameters
    ----------
    values:
        Time-axis payload to sanitize.

    Returns
    -------
    np.ndarray
        One-dimensional float32 axis. Rare non-finite gaps are filled
        conservatively from neighboring values.
    """

    arr = np.asarray(unwrap_scalar_obj(values), dtype=np.float32).reshape(-1)
    if arr.size == 0:
        return arr

    finite = np.isfinite(arr)
    if np.all(finite):
        return arr

    out = arr.copy()
    finite_idx = np.flatnonzero(finite)
    if finite_idx.size == 0:
        return np.arange(arr.size, dtype=np.float32)

    if finite_idx.size >= 2:
        # Use the median finite step as a conservative local spacing estimate
        # before repairing missing timestamps.
        step = np.float32(np.nanmedian(np.diff(out[finite_idx])))
        if (not np.isfinite(step)) or step == 0:
            step = np.float32(1.0)
    else:
        step = np.float32(1.0)

    first = int(finite_idx[0])
    for idx in range(first - 1, -1, -1):
        out[idx] = out[idx + 1] - step

    prev = first
    for idx in finite_idx[1:]:
        idx = int(idx)
        gap = idx - prev
        if gap > 1:
            for fill_idx in range(prev + 1, idx):
                out[fill_idx] = out[prev] + step * np.float32(fill_idx - prev)
        prev = idx

    last = int(finite_idx[-1])
    for idx in range(last + 1, out.size):
        out[idx] = out[idx - 1] + step

    return out.astype(np.float32, copy=False)



def make_picklable(value: Any) -> Any:
    """Convert nested arrays and scalars to pickle-friendly builtins.

    Parameters
    ----------
    value:
        Arbitrary nested payload that may include NumPy arrays, pandas objects,
        or Python containers.

    Returns
    -------
    Any
        Recursively converted payload composed of pickle-friendly builtins and
        NumPy arrays.
    """

    if value is None:
        return None
    if isinstance(value, (str, int, float, bool, np.number)):
        return value
    if isinstance(value, np.ndarray):
        return value

    try:
        import pandas as pd  # local import to keep base module lightweight

        if isinstance(value, (pd.Series, pd.Index)):
            return value.to_numpy()
        if isinstance(value, pd.DataFrame):
            return {k: make_picklable(v) for k, v in value.to_dict(orient="list").items()}
    except Exception:
        pass

    if isinstance(value, dict):
        return {k: make_picklable(v) for k, v in value.items()}
    if isinstance(value, (list, tuple)):
        return [make_picklable(v) for v in value]

    # Fall back to a string representation rather than leaving an opaque object
    # that may fail during notebook export or caching.
    return str(value)