"""Generate a read-only semantic map for the 711 historical primitives.

The generator consumes engineering/source metadata only.  It never imports or
evaluates a historical formula, and optional checkpoints are inspected solely
to count already-materialized finite observations.
"""

from __future__ import annotations

import argparse
import hashlib
import json
from pathlib import Path
from typing import Any

import numpy as np
import pyarrow as pa
import pyarrow.parquet as pq


EXPECTED_PRIMITIVES = 711


def digest(path: Path) -> str:
    if path.is_file():
        return hashlib.sha256(path.read_bytes()).hexdigest()
    hasher = hashlib.sha256()
    for item in sorted(path.glob("*.npy")):
        hasher.update(item.name.encode("utf-8"))
        hasher.update(b"\0")
        hasher.update(bytes.fromhex(digest(item)))
    return hasher.hexdigest()


def key_for(row: dict[str, Any]) -> tuple[int, int, float]:
    return int(row["indicator_id"]), int(row["period"]), float(row["p1"])


def key_text(key: tuple[int, int, float]) -> str:
    return f"{key[0]}:{key[1]}:{key[2]:g}"


def composite_key(value: Any) -> tuple[int, int, float]:
    """Decode canonical primitive keys stored as strings or Parquet structs."""
    if isinstance(value, dict):
        return key_for(value)
    if isinstance(value, (list, tuple)) and len(value) == 3:
        return int(value[0]), int(value[1]), float(value[2])
    if not isinstance(value, str):
        raise ValueError(f"invalid composite primitive key: {value!r}")
    parts = value.strip().split(":")
    if len(parts) != 3:
        parts = value.strip().split("_")
        if len(parts) == 4 and parts[0] == "hs22":
            parts = parts[1:]
    if len(parts) != 3:
        raise ValueError(f"invalid composite primitive key: {value!r}")
    try:
        return int(parts[0]), int(parts[1]), float(parts[2])
    except ValueError as error:
        raise ValueError(f"invalid composite primitive key: {value!r}") from error


def source_name(row: dict[str, Any]) -> str:
    return str(row.get("source_function", "")).lower()


def semantic_mapping(row: dict[str, Any]) -> tuple[str, str, str, str, str]:
    """Return domain, subdomain, type, provenance, and deterministic rule id."""
    name = source_name(row)
    family = str(row.get("engineering_family", ""))
    rules = (
        (("cross_", "crossover", "crossunder", "breakout", "divergence"), "price_action", "crossing_or_breakout", "event", "source_derived", "explicit_event_name"),
        (("regime_", "entropy_", "chop", "fractal_dimension"), "market_regime", "regime_detection", "state", "source_derived", "explicit_regime_name"),
        (("time_", "session", "weekday", "hour"), "time_context", "session_or_calendar", "state", "source_derived", "explicit_time_name"),
        (("micro_", "orderflow", "vwap", "obv", "mfi", "cmf", "volume"), "volume_and_flow", "volume_flow", "continuous", "source_derived", "explicit_volume_name"),
        (("osc_", "rsi", "stoch", "cci", "williams", "roc", "momentum"), "momentum", "oscillator_or_rate", "continuous", "source_derived", "explicit_momentum_name"),
        (("bb_", "kelt_", "donch_", "ichimoku_", "supertrend", "psar", "channel"), "volatility", "bands_and_channels", "state", "source_derived", "explicit_band_name"),
        (("vol_", "atr", "true_range", "variance", "std", "range"), "volatility", "range_and_dispersion", "continuous", "source_derived", "explicit_volatility_name"),
        (("ma_", "ema", "sma", "wma", "tema", "dema"), "trend", "moving_average", "continuous", "source_derived", "explicit_average_name"),
        (("adx", "aroon_", "linreg_", "trend_"), "trend", "direction_and_strength", "continuous", "source_derived", "explicit_trend_name"),
        (("pivot_", "prev_", "rolling_", "quantile_"), "price_structure", "reference_level", "state", "source_derived", "explicit_reference_name"),
        (("norm_", "zscore", "percentile"), "normalization", "scaled_price_or_signal", "continuous", "source_derived", "explicit_normalization_name"),
    )
    for prefixes, domain, subdomain, value_type, provenance, rule in rules:
        if any(prefix in name for prefix in prefixes):
            return domain, subdomain, value_type, provenance, rule
    family_fallbacks = {
        "moving_average": ("trend", "moving_average", "continuous"),
        "oscillator": ("momentum", "oscillator_or_rate", "continuous"),
        "volatility": ("volatility", "range_and_dispersion", "continuous"),
        "regime": ("market_regime", "regime_detection", "state"),
        "microstructure": ("market_microstructure", "price_volume_structure", "continuous"),
        "momentum": ("momentum", "oscillator_or_rate", "continuous"),
        "trend_structure": ("trend", "direction_and_strength", "continuous"),
        "normalization": ("normalization", "scaled_price_or_signal", "continuous"),
        "cross_indicator": ("price_action", "crossing_or_breakout", "event"),
        "time_session": ("time_context", "session_or_calendar", "state"),
        "reference_level": ("price_structure", "reference_level", "state"),
        "band_channel": ("volatility", "bands_and_channels", "state"),
    }
    domain, subdomain, value_type = family_fallbacks.get(
        family, ("specialized", "unclassified_source_function", "continuous")
    )
    return domain, subdomain, value_type, "heuristic", "engineering_family_fallback"


def load_engineering_map(path: Path) -> list[dict[str, Any]]:
    payload = json.loads(path.read_text(encoding="utf-8"))
    if payload.get("artifact") != "ENGINEERING_FAMILY_MAP_V1":
        raise ValueError("--engineering-map is not ENGINEERING_FAMILY_MAP_V1")
    rows = payload.get("primitives")
    if not isinstance(rows, list):
        raise ValueError("engineering map has no primitives list")
    keys = [key_for(row) for row in rows]
    if len(rows) != EXPECTED_PRIMITIVES or len(set(keys)) != EXPECTED_PRIMITIVES:
        raise ValueError("engineering map must contain exactly 711 unique primitives")
    return rows


def usage_counts(path: Path, expected_indicator_ids: set[int]) -> dict[int, int]:
    """Load indicator-level usage; it cannot distinguish parameterized primitives."""
    counts: dict[int, int] = {}
    source = pq.ParquetFile(path)
    columns = set(source.schema.names)
    required = {"entity_type", "indicator_id", "usage_count"}
    if missing := required - columns:
        raise ValueError(f"usage parquet is missing required columns: {', '.join(sorted(missing))}")
    for batch in source.iter_batches(columns=["entity_type", "indicator_id", "usage_count"]):
        for row in batch.to_pylist():
            if row["entity_type"] != "indicator_id":
                continue
            indicator_id = int(row["indicator_id"])
            if indicator_id not in expected_indicator_ids:
                continue
            if indicator_id in counts:
                raise ValueError(f"usage parquet has duplicate indicator row: {indicator_id}")
            counts[indicator_id] = int(row["usage_count"])
    if set(counts) != expected_indicator_ids:
        raise ValueError("usage parquet does not cover every engineering indicator ID")
    return counts


def checkpoint_counts(directory: Path, expected: set[tuple[int, int, float]]) -> tuple[dict[tuple[int, int, float], int], dict[tuple[int, int, float], str]]:
    counts: dict[tuple[int, int, float], int] = {}
    names: dict[tuple[int, int, float], str] = {}
    for path in sorted(directory.glob("*.npy")):
        parts = path.stem.split("_")
        if len(parts) != 4 or parts[0] != "hs22":
            raise ValueError(f"invalid checkpoint filename: {path.name}")
        try:
            key = int(parts[1]), int(parts[2]), float(parts[3])
        except ValueError as error:
            raise ValueError(f"invalid checkpoint filename: {path.name}") from error
        if key not in expected or key in counts:
            raise ValueError(f"checkpoint is not a unique engineering primitive: {path.name}")
        counts[key] = int(np.isfinite(np.load(path, allow_pickle=False)).sum())
        names[key] = path.name
    if set(counts) != expected:
        raise ValueError("checkpoint directory does not cover every engineering primitive")
    return counts, names


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--engineering-map", type=Path, required=True)
    parser.add_argument("--usage-parquet", type=Path)
    parser.add_argument("--checkpoint-dir", type=Path)
    parser.add_argument("--output-dir", type=Path, required=True)
    args = parser.parse_args()
    if args.usage_parquet and not args.usage_parquet.is_file():
        parser.error("--usage-parquet must be a file")
    if args.checkpoint_dir and not args.checkpoint_dir.is_dir():
        parser.error("--checkpoint-dir must be a directory")

    engineering = load_engineering_map(args.engineering_map)
    expected = {key_for(row) for row in engineering}
    expected_indicator_ids = {key[0] for key in expected}
    usage = usage_counts(args.usage_parquet, expected_indicator_ids) if args.usage_parquet else {}
    finite, checkpoint_names = checkpoint_counts(args.checkpoint_dir, expected) if args.checkpoint_dir else ({}, {})
    rows = []
    for engineering_row in engineering:
        key = key_for(engineering_row)
        domain, subdomain, value_type, provenance, rule = semantic_mapping(engineering_row)
        rows.append({
            "feature_key": key_text(key), "indicator_id": key[0], "period": key[1], "p1": key[2],
            "source_function": engineering_row.get("source_function"),
            "source_expression": engineering_row.get("source_expression"),
            "engineering_family": engineering_row.get("engineering_family"),
            "semantic_domain": domain, "semantic_subdomain": subdomain, "semantic_type": value_type,
            "semantic_mapping_provenance": provenance, "semantic_mapping_rule": rule,
            "usage_count": usage.get(key[0]),
            "usage_grain": "indicator_id_replicated_not_parameter_specific" if args.usage_parquet else None,
            "finite_observation_count": finite.get(key),
            "checkpoint_file": checkpoint_names.get(key),
        })

    args.output_dir.mkdir(parents=True, exist_ok=True)
    parquet_path = args.output_dir / "historical_feature_semantic_domain_map_v1.parquet"
    pq.write_table(pa.Table.from_pylist(rows), parquet_path, compression="zstd")
    sources: dict[str, dict[str, str]] = {"engineering_map": {"path": str(args.engineering_map), "sha256": digest(args.engineering_map)}}
    if args.usage_parquet:
        sources["usage_parquet"] = {"path": str(args.usage_parquet), "sha256": digest(args.usage_parquet)}
    if args.checkpoint_dir:
        sources["checkpoint_dir"] = {"path": str(args.checkpoint_dir), "sha256": digest(args.checkpoint_dir)}
    summary = {
        "schema_version": 1,
        "artifact": "HISTORICAL_FEATURE_SEMANTIC_DOMAIN_MAP_V1_SUMMARY",
        "method": "read-only source-function semantic classification; explicit name rules are source-derived and family fallbacks are heuristic; no historical formulas were imported or evaluated",
        "source_artifacts": sources,
        "counts": {
            "primitives": len(rows),
            "semantic_domains": len({row["semantic_domain"] for row in rows}),
            "semantic_subdomains": len({row["semantic_subdomain"] for row in rows}),
            "source_derived_mappings": sum(row["semantic_mapping_provenance"] == "source_derived" for row in rows),
            "heuristic_mappings": sum(row["semantic_mapping_provenance"] == "heuristic" for row in rows),
            "usage_count_total_replicated_across_primitives": sum(row["usage_count"] or 0 for row in rows),
            "finite_observation_count_total": sum(row["finite_observation_count"] or 0 for row in rows),
        },
        "semantic_type_counts": {kind: sum(row["semantic_type"] == kind for row in rows) for kind in ("continuous", "state", "event")},
        "usage_grain": "indicator_id_replicated_not_parameter_specific" if args.usage_parquet else None,
        "artifacts": {"semantic_domain_map_parquet": str(parquet_path), "semantic_domain_map_parquet_sha256": digest(parquet_path)},
    }
    (args.output_dir / "historical_feature_semantic_domain_map_v1_summary.json").write_text(
        json.dumps(summary, indent=2, sort_keys=True) + "\n", encoding="utf-8"
    )


if __name__ == "__main__":
    main()