"""Bytes/hex distance (Hamming/popcount) and timestamp FFT phase distance."""

import numpy as np
import polars as pl


# ---------------------------------------------------------------------------
#  Bytes Hamming distance
# ---------------------------------------------------------------------------

_HEX_POPCOUNT: dict[str, int] = {
    '0': 0, '1': 1, '2': 1, '3': 2,
    '4': 1, '5': 2, '6': 2, '7': 3,
    '8': 1, '9': 2,
    'a': 2, 'b': 3, 'c': 2, 'd': 3, 'e': 3, 'f': 4,
    'A': 2, 'B': 3, 'C': 2, 'D': 3, 'E': 3, 'F': 4,
}


def _hex_popcount_normalized(val: str | None) -> float:
    """Return the proportion of 1-bits in a hex byte string (0.0–1.0).

    This is the Hamming distance from the all-zero bit string,
    normalised by total bits.  Returns ``NaN`` for null or
    non-hex input.
    """
    if val is None:
        return float("nan")
    s = str(val).replace(' ', '').replace('\t', '').replace('\n', '')
    if not s:
        return 0.0

    total_ones = 0
    for ch in s:
        ones = _HEX_POPCOUNT.get(ch)
        if ones is None:
            return float("nan")  # non-hex character
        total_ones += ones

    total_bits = len(s) * 4
    return total_ones / total_bits


def bytes_distance(
    lf: pl.LazyFrame,
    byte_columns: list[str],
) -> pl.LazyFrame:
    """Append normalised Hamming-distance columns for hex byte columns.

    Each byte column is converted from hex to binary, and the
    proportion of 1-bits (Hamming distance from the all-zero
    bit vector) is computed as a ``_bytes_dist_N`` Float64 column
    normalised to [0, 1].  Non-hex / null values produce ``NaN``.

    Parameters
    ----------
    lf:
        LazyFrame containing the byte columns.
    byte_columns:
        Column names whose values are hex byte strings
        (with optional whitespace separators, e.g. ``"A3 F2"``).

    Returns
    -------
    pl.LazyFrame
        Input frame with ``_bytes_dist_0``, ``_bytes_dist_1``, …
        appended.  Columns not present in *lf* are silently skipped.

    Examples
    --------
    >>> df = pl.DataFrame({"payload": ["A3F2", "FFFF", None]})
    >>> result = bytes_distance(df.lazy(), ["payload"])
    >>> result.collect()
    shape: (3, 2)
    ┌─────────┬────────────────┐
    │ payload ┆ _bytes_dist_0  │
    │ ---     ┆ ---            │
    │ str     ┆ f64            │
    ╞═════════╪════════════════╡
    │ A3F2    ┆ 0.625          │
    │ FFFF    ┆ 1.0            │
    │ null    ┆ NaN            │
    └─────────┴────────────────┘
    """
    existing = set(lf.collect_schema().names())
    result = lf

    for i, col_name in enumerate(byte_columns):
        if col_name not in existing:
            continue
        result = result.with_columns(
            pl.col(col_name)
            .cast(pl.Utf8)
            .map_elements(_hex_popcount_normalized, return_dtype=pl.Float64)
            .alias(f'_bytes_dist_{i}')
        )

    return result


# ---------------------------------------------------------------------------
#  Timestamp FFT phase distance
# ---------------------------------------------------------------------------

# ── threshold: minimum fraction of repeated intervals to use FFT path ──
_FFT_PERIODICITY_THRESHOLD = 0.05
# ── minimum samples before FFT is attempted ──
_FFT_MIN_SAMPLES = 4


def _parse_timestamp_to_epoch(series: pl.Series) -> pl.Series:
    """Convert a timestamp series to Unix epoch seconds (float64).

    Handles int/float (assumed already epoch), pl.Datetime/pl.Date
    (ns→seconds), and str/pl.Utf8 (strptime→ns→seconds, with direct
    float-cast fallback on parse failure).
    """
    dtype = series.dtype
    if dtype in (pl.Utf8, pl.String):
        try:
            parsed = series.str.strptime(pl.Datetime, format=None, strict=False)
        except Exception:
            return series.cast(pl.Float64)
        return parsed.cast(pl.Int64).cast(pl.Float64) / 1_000_000_000.0
    elif dtype in (pl.Datetime, pl.Date):
        return series.cast(pl.Int64).cast(pl.Float64) / 1_000_000_000.0
    else:
        return series.cast(pl.Float64)


def _compute_fft_distances(values: np.ndarray) -> np.ndarray:
    """Compute per-row phase distances for a 1-D array.

    Sorts *values*, checks for periodicity, then either:
    * FFT path: ``np.fft.fft()`` → ``np.angle()``, mapped back to
      original (unsorted) row order via rank lookup.
    * Statistical fallback: interval z-scores mapped to row order.

    Returns an ``(n,)`` float64 array, normalised to zero-mean
    unit-variance.
    """
    n = len(values)
    if n < _FFT_MIN_SAMPLES:
        return np.full(n, np.nan, dtype=np.float64)

    # Sort once; record rank→position mapping for unsorting later
    sorted_order = np.argsort(values)
    ranks = np.empty(n, dtype=np.int64)
    ranks[sorted_order] = np.arange(n)
    sorted_vals = values[sorted_order]

    # ── periodicity check ──────────────────────────────────────────
    intervals = np.diff(sorted_vals)
    unique_int, counts = np.unique(
        np.round(intervals, 6), return_counts=True
    )
    repeat_ratio = float(np.sum(counts[counts > 1])) / max(len(intervals), 1)

    if repeat_ratio >= _FFT_PERIODICITY_THRESHOLD:
        # ── FFT phase path ─────────────────────────────────────────
        centered = sorted_vals - np.mean(sorted_vals)
        fft_result = np.fft.fft(centered)
        phases = np.angle(fft_result)
        dist_values = phases[ranks]
    else:
        # ── statistical fallback ───────────────────────────────────
        mean_int = float(np.mean(intervals))
        std_int = float(np.std(intervals))
        if std_int > 1e-10:
            z_scores = (intervals - mean_int) / std_int
        else:
            z_scores = np.zeros_like(intervals)

        # Row at sorted position r gets the z-score of interval r-1
        # (preceding interval); row at position 0 reuses last interval.
        dist_values = np.empty(n, dtype=np.float64)
        for r in range(n):
            if r > 0:
                dist_values[ranks[r]] = z_scores[r - 1]
            else:
                dist_values[ranks[0]] = (
                    z_scores[-1] if n > 1 else 0.0
                )

    # Normalise to zero-mean unit-variance
    d_std = float(np.std(dist_values))
    d_mean = float(np.mean(dist_values))
    if d_std > 1e-10:
        dist_values = (dist_values - d_mean) / d_std
    else:
        dist_values = dist_values - d_mean

    return dist_values


def timestamp_fft_distance(
    lf: pl.LazyFrame, ts_columns: list[str]
) -> pl.LazyFrame:
    """Append FFT-based phase-distance columns for timestamp columns.

    For each column in *ts_columns* that exists in the schema:

    1. Parse values to Unix epoch float seconds (handles int, float,
       datetime strings, and ``pl.Datetime``).
    2. Collect the column to numpy, sort, then run ``np.fft.fft``
       followed by ``np.angle`` to extract per-row phase values.
    3. If fewer than 5 % of timestamp intervals repeat (weak
       periodicity), fall back to statistical features derived
       from interval mean, standard deviation, and z-scores.

    Columns are appended as ``_ts_dist_0``, ``_ts_dist_1``, … in the
    order *ts_columns* are supplied.  Rows whose timestamp is null
    receive ``NaN``.

    Parameters
    ----------
    lf : pl.LazyFrame
        The input LazyFrame with at least one timestamp column.
    ts_columns : list[str]
        Column names to process.

    Returns
    -------
    pl.LazyFrame
        A new LazyFrame with ``_ts_dist_N`` columns appended.

    Examples
    --------
    >>> import polars as pl
    >>> from analysis.distance import timestamp_fft_distance
    >>> df = pl.DataFrame({
    ...     'ts': [0, 60, 120, 180, 240, 300],
    ... })
    >>> result = timestamp_fft_distance(df.lazy(), ['ts'])
    >>> result.collect().columns
    ['ts', '_ts_dist_0']
    """
    schema = lf.collect_schema()
    available = [c for c in ts_columns if c in schema.names()]
    if not available:
        return lf

    # Collect all needed columns once
    df = lf.select(available).collect(streaming=True)
    n_rows = len(df)

    dist_arrays: list[np.ndarray] = []

    for col in available:
        series = df[col]
        ts_numeric = _parse_timestamp_to_epoch(series)

        # Identify valid (non-null) rows
        valid = ~ts_numeric.is_null()
        valid_indices = np.where(valid.to_numpy())[0]
        values = ts_numeric.filter(valid).to_numpy().astype(np.float64)

        if len(values) < _FFT_MIN_SAMPLES:
            continue

        phase_values = _compute_fft_distances(values)

        # Build full-length array; invalid rows stay NaN
        full_dist = np.full(n_rows, np.nan, dtype=np.float64)
        full_dist[valid_indices] = phase_values
        dist_arrays.append(full_dist)

    if not dist_arrays:
        return lf

    # Re-materialise the original frame and bolt on distance columns
    result_df = lf.collect(streaming=True)
    for i, arr in enumerate(dist_arrays):
        result_df = result_df.with_columns(
            pl.Series(f'_ts_dist_{i}', arr)
        )

    return result_df.lazy()
