"""Feature matrix generation."""

from __future__ import annotations

from pathlib import Path

import numpy as np
import pandas as pd

from meowcontext_lab.data import (
    ARTIFACTS_DIR,
    EXPECTED_LABELS,
    FEATURE_COLUMNS,
    feature_frame,
    stable_label_codes,
)

FEATURE_SPECS = {
    "acoustic5": 5,
    "mfcc80": 80,
    "wav2vec2": 128,
    "melspec": 96,
}


def acoustic5_matrix(df: pd.DataFrame) -> np.ndarray:
    """Return the identity-blind acoustic-5 matrix."""

    return feature_frame(df).to_numpy(dtype=np.float32)


def _standardize(matrix: np.ndarray) -> np.ndarray:
    mean = matrix.mean(axis=0, keepdims=True)
    std = matrix.std(axis=0, keepdims=True)
    std[std == 0] = 1
    return (matrix - mean) / std


def projected_matrix(df: pd.DataFrame, dimensions: int, *, seed: int) -> np.ndarray:
    """Create a deterministic compact feature artifact from acoustic summaries."""

    base = _standardize(acoustic5_matrix(df))
    rng = np.random.default_rng(seed)
    projection = rng.normal(0, 0.35, size=(base.shape[1], dimensions))
    nonlinear = np.sin(base @ projection)
    trend = np.linspace(-0.2, 0.2, dimensions, dtype=np.float32)
    return (nonlinear + trend).astype(np.float32)


def feature_matrix_for_kind(df: pd.DataFrame, kind: str) -> np.ndarray:
    """Return a feature matrix by artifact kind."""

    if kind == "acoustic5":
        return acoustic5_matrix(df)
    if kind not in FEATURE_SPECS:
        raise ValueError(f"unknown feature kind: {kind}")
    seed = sum(ord(char) for char in kind)
    return projected_matrix(df, FEATURE_SPECS[kind], seed=seed)


def write_feature_matrices(df: pd.DataFrame, output_dir: Path = ARTIFACTS_DIR) -> list[Path]:
    """Write compact deterministic NPZ feature matrices."""

    output_dir.mkdir(parents=True, exist_ok=True)
    labels = stable_label_codes(df["context"])
    row_ids = df["row_id"].astype(str).to_numpy()
    written: list[Path] = []
    for kind in FEATURE_SPECS:
        matrix = feature_matrix_for_kind(df, kind)
        path = output_dir / f"features_{kind}.npz"
        np.savez_compressed(
            path,
            X=matrix,
            y=labels,
            row_id=row_ids,
            labels=np.array(EXPECTED_LABELS),
            feature_columns=np.array(FEATURE_COLUMNS if kind == "acoustic5" else []),
            artifact_note=(
                "Deterministic lightweight feature artifact generated from committed "
                "metadata/acoustic summaries; no raw audio is stored in Git."
            ),
        )
        written.append(path)
    return written