"""Read and derive Q1 V2 feature views from the native-source artifacts."""
from __future__ import annotations

import json
import re
from pathlib import Path
from typing import Any

import numpy as np
from scipy import sparse

Q1_DIR = Path(__file__).resolve().parent
FEATURE_DIR = Q1_DIR / "features_v2"
MODALITIES = ("text", "audio", "vision")
FEATURE_DIMS = {"text": 768, "audio": 74, "vision": 35, "speech": 768}


def load_sample(sample_id: str, feature_dir: Path = FEATURE_DIR) -> dict[str, Any]:
    """Load one sample by `video_id/clip_id` and return arrays plus metadata."""
    manifest = feature_dir / "manifest_q1.jsonl"
    found = None
    for line in manifest.read_text(encoding="utf-8").splitlines():
        row = json.loads(line)
        if row["sample_id"] == sample_id:
            found = row
            break
    if found is None:
        raise KeyError(f"sample_id not found: {sample_id}")
    path = Q1_DIR / found["feature_path"]
    with np.load(path, allow_pickle=False) as archive:
        result = {name: archive[name] for name in archive.files}
    result["_path"] = path
    result["_manifest"] = found
    result["_meta"] = json.loads(str(result["meta_json"]))
    return result


def csr_from_sample(sample: dict[str, Any], prefix: str) -> sparse.csr_matrix:
    """Reconstruct a stored sparse alignment/query map."""
    shape = tuple(int(v) for v in sample[f"{prefix}_shape"])
    return sparse.csr_matrix(
        (
            sample[f"{prefix}_data"].astype(np.float32),
            sample[f"{prefix}_indices"].astype(np.int32),
            sample[f"{prefix}_indptr"].astype(np.int32),
        ),
        shape=shape,
    )


def _native(sample: dict[str, Any], modality: str) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
    if modality == "text":
        return (
            sample["native_text_features"].astype(np.float32),
            np.broadcast_to(sample["native_text_observed"][:, None], sample["native_text_features"].shape),
            sample["native_text_intervals"].astype(np.float32),
            sample["native_text_quality_effective"].astype(np.float32),
            sample["native_text_quality_available"].astype(bool),
        )
    if modality in ("audio", "vision"):
        return (
            sample[f"native_{modality}_features"].astype(np.float32),
            sample[f"native_{modality}_mask"].astype(bool),
            sample[f"native_{modality}_intervals"].astype(np.float32),
            sample[f"native_{modality}_quality"].astype(np.float32),
            np.asarray(sample[f"native_{modality}_quality_available"], dtype=bool),
        )
    if modality == "speech":
        meta = sample["_meta"]
        safe = lambda value: re.sub(r"[^A-Za-z0-9_.-]+", "_", str(value))
        path = Q1_DIR / "cache" / "native" / f"{safe(meta['video_id'])}__{safe(meta['clip_id'])}.npz"
        if not path.is_file():
            raise FileNotFoundError(f"B4 speech source cache is unavailable: {path}")
        with np.load(path, allow_pickle=False) as archive:
            values = archive["speech_features"].astype(np.float32)
        times = sample["native_ctc_times"].astype(np.float32)
        if len(values) != len(times):
            raise ValueError("B4 speech rows do not match the stored CTC timestamps")
        step = 320.0 / 16000.0
        intervals = np.column_stack((np.maximum(0.0, times - step / 2), np.minimum(float(meta["duration_s"]), times + step / 2))).astype(np.float32)
        return values, np.isfinite(values), intervals, np.ones(len(values), np.float32), np.zeros(len(values), bool)
    raise ValueError(f"unknown modality: {modality}")


def _aggregate(
    values: np.ndarray,
    observed: np.ndarray,
    source_intervals: np.ndarray,
    quality: np.ndarray,
    target_intervals: np.ndarray,
    quality_available: np.ndarray | None = None,
) -> dict[str, np.ndarray]:
    """Quality-weight B0 means while keeping physical coverage and availability separate."""
    n, d = len(target_intervals), values.shape[1]
    mean = np.zeros((n, d), np.float32)
    var = np.zeros((n, d), np.float32)
    count = np.zeros((n, d), np.uint8)
    coverage = np.zeros((n, d), np.float32)
    qbar = np.ones((n, d), np.float32)
    qavail = np.zeros((n, d), np.float32)
    available = np.zeros(len(values), bool) if quality_available is None else np.asarray(quality_available, bool)
    mask = np.zeros((n, d), bool)
    for i, (left, right) in enumerate(target_intervals):
        width = float(right - left)
        if width <= 0 or not len(source_intervals):
            continue
        overlap = np.maximum(
            0.0,
            np.minimum(source_intervals[:, 1], right) - np.maximum(source_intervals[:, 0], left),
        )
        physical = overlap[:, None] * observed
        weighted = physical * quality[:, None]
        total = weighted.sum(axis=0)
        valid = total > 0
        if valid.any():
            weighted_sum = (values * weighted).sum(axis=0)
            weighted_sq = (np.square(values) * weighted).sum(axis=0)
            mean[i, valid] = (weighted_sum[valid] / total[valid]).astype(np.float32)
            var[i, valid] = np.maximum(0.0, weighted_sq[valid] / total[valid] - mean[i, valid] ** 2)
            mask[i, valid] = True
        physical_total = physical.sum(axis=0)
        physical_valid = physical_total > 0
        qbar[i, physical_valid] = weighted.sum(axis=0)[physical_valid] / physical_total[physical_valid]
        qavail[i, physical_valid] = (physical * available[:, None]).sum(axis=0)[physical_valid] / physical_total[physical_valid]
        coverage[i] = np.minimum(1.0, physical_total / width)
        count[i] = np.minimum(255, ((overlap[:, None] > 0) & observed).sum(axis=0)).astype(np.uint8)
    return {
        "x": mean.astype(np.float16),
        "var": var.astype(np.float32),
        "count": count,
        "coverage": coverage,
        "mask": mask,
        "quality_mean": qbar.astype(np.float32),
        "quality_available_fraction": qavail.astype(np.float32),
    }


def get_view(
    sample: dict[str, Any],
    modality: str,
    view: str = "sec",
    context_s: float | None = None,
    geometry_aware: bool = False,
) -> dict[str, np.ndarray]:
    """Return `sec`, `phase50`, `word`, `multi`, or `posterior` features.

    `multi` is always aggregated from native rows around each 0.1 s bin midpoint.
    For it, set `context_s` to one of 0.1, 0.3, or 0.7.
    """
    if modality not in (*MODALITIES, "speech"):
        raise ValueError(f"unknown modality: {modality}")
    if view in ("sec", "phase50", "word"):
        bounds = sample[f"views_{view}_time_bounds_s"].astype(np.float32)
        if modality == "speech":
            values, observed, intervals, quality, available = _native(sample, modality)
            return {**_aggregate(values, observed, intervals, quality, bounds, available), "time_bounds_s": bounds}
        prefix = f"views_{view}_{modality}_"
        result = {
            "x": sample[prefix + "x"],
            "var": sample[prefix + "var"],
            "count": sample[prefix + "count"],
            "coverage": sample[prefix + "coverage_u8"].astype(np.float32) / 255.0,
            "mask": sample[prefix + "mask"].astype(bool),
            "quality_mean": sample[prefix + "qbar"].astype(np.float32),
            "quality_available_fraction": sample[prefix + "quality_available_fraction"].astype(np.float32),
            "time_bounds_s": bounds,
        }
        if view == "sec" and modality == "vision" and geometry_aware:
            result["x"] = sample["views_sec_vision_b1_x"]
            result["mask"] = sample["views_sec_vision_b1_mask"].astype(bool)
        return result
    if view == "multi":
        if context_s not in (0.1, 0.3, 0.7):
            raise ValueError("context_s must be one of 0.1, 0.3, or 0.7")
        centers = sample["views_sec_time_bounds_s"].astype(np.float32).mean(axis=1)
        duration = float(sample["_meta"]["duration_s"])
        bounds = np.column_stack((np.maximum(0.0, centers - context_s / 2), np.minimum(duration, centers + context_s / 2))).astype(np.float32)
        values, observed, intervals, quality, available = _native(sample, modality)
        return {**_aggregate(values, observed, intervals, quality, bounds, available), "time_bounds_s": bounds}
    if view == "posterior":
        if modality != "text":
            raise ValueError("posterior projection is defined for text only")
        return posterior_text_view(sample)
    raise ValueError(f"unknown view: {view}")


def posterior_text_view(sample: dict[str, Any]) -> dict[str, np.ndarray]:
    """Project CTC forward-backward word occupancy directly to the 0.1 s grid."""
    edges = sample["views_sec_time_bounds_s"].astype(np.float32)
    words = sample["native_text_features"].astype(np.float32)
    word_ok = sample["native_text_observed"].astype(bool)
    occupancy = sample["native_ctc_occupancy"].astype(np.float32)
    times = sample["native_ctc_times"].astype(np.float32)
    if occupancy.ndim != 2 or occupancy.shape[1] != len(words):
        raise ValueError("CTC occupancy word dimension does not match the native text rows")
    step = 320.0 / 16000.0
    half = step / 2
    duration = float(sample["_meta"]["duration_s"])
    source = np.column_stack((np.maximum(0.0, times - half), np.minimum(duration, times + half))).astype(np.float32)
    n, d = len(edges), words.shape[1]
    result = np.zeros((n, d), np.float32)
    var = np.zeros((n, d), np.float32)
    count = np.zeros(n, np.uint8)
    coverage = np.zeros(n, np.float32)
    mask = np.zeros(n, bool)
    mass_by_word = np.zeros((n, len(words)), np.float32)
    for i, (left, right) in enumerate(edges):
        overlap = np.maximum(0.0, np.minimum(source[:, 1], right) - np.maximum(source[:, 0], left))
        mass_by_word[i] = (occupancy * overlap[:, None]).sum(axis=0)
        w = mass_by_word[i] * word_ok
        total = float(w.sum())
        if total > 0:
            normalized = w / total
            result[i] = normalized @ words
            centered = words - result[i]
            var[i] = (normalized[:, None] * np.square(centered)).sum(axis=0)
            mask[i] = True
            count[i] = min(255, int(np.count_nonzero(w > 0)))
            coverage[i] = min(1.0, total / max(float(right - left), 1e-8))
    return {
        "x": result.astype(np.float16),
        "var": var,
        "count": np.broadcast_to(count[:, None], result.shape).copy(),
        "coverage": np.broadcast_to(coverage[:, None], result.shape).copy(),
        "mask": np.broadcast_to(mask[:, None], result.shape).copy(),
        "quality_mean": np.ones_like(result, dtype=np.float32),
        "quality_available_fraction": np.zeros_like(result, dtype=np.float32),
        "word_posterior_mass": mass_by_word,
        "time_bounds_s": edges,
    }