#!/usr/bin/env python3
"""Read-only Cohort001 entry-to-frozen-oracle failure mining.

Associations in this report are observational and must not be interpreted as
causal feature effects.  The runner never replays a strategy or alters inputs.
"""

from __future__ import annotations

import argparse
import hashlib
import json
from collections import Counter, defaultdict
from pathlib import Path
from typing import Any

import numpy as np


ARTIFACT = "COHORT001_FEATURE_FAILURE_MINING_V1"
GOOD_LABEL = "GOOD_ENTRY"
FAILURE_LABELS = ("WRONG_DIRECTION", "HIGH_MAE_ENTRY", "LATE_SIGNAL", "RECOVERY_DEPENDENT")
ENTRY_FIELDS = ("entry_bar", "entry_index", "entry_bar_index", "entry_idx", "bar_index")


def file_hash(path: Path) -> str:
    digest = hashlib.sha256()
    with path.open("rb") as source:
        for block in iter(lambda: source.read(1024 * 1024), b""):
            digest.update(block)
    return digest.hexdigest()


def directory_hash(path: Path) -> str:
    """Hash names and contents without loading the oracle into memory."""
    digest = hashlib.sha256()
    for item in sorted(path.glob("*.npy")):
        digest.update(item.name.encode("utf-8"))
        digest.update(b"\0")
        digest.update(bytes.fromhex(file_hash(item)))
    return digest.hexdigest()


def json_compatible(value: Any) -> Any:
    """Convert NumPy values to strict JSON-compatible Python values."""
    if isinstance(value, np.ndarray):
        return json_compatible(value.tolist())
    if isinstance(value, np.bool_):
        return bool(value)
    if isinstance(value, np.integer):
        return int(value)
    if isinstance(value, np.floating):
        value = float(value)
    if isinstance(value, float):
        return value if np.isfinite(value) else None
    if isinstance(value, dict):
        return {json_compatible(key): json_compatible(item) for key, item in value.items()}
    if isinstance(value, (list, tuple)):
        return [json_compatible(item) for item in value]
    return value


def load_engineering_map(path: Path) -> dict[str, dict[str, Any]]:
    payload = json.loads(path.read_text(encoding="utf-8"))
    rows = payload.get("primitives", payload.get("requests", []))
    if not isinstance(rows, list):
        raise ValueError("engineering map must contain primitives or requests")
    result = {}
    for row in rows:
        try:
            key = feature_key(row)
        except (KeyError, TypeError, ValueError):
            continue
        if key in result:
            raise ValueError(f"duplicate engineering primitive: {key}")
        result[key] = dict(row, feature_key=key)
    if not result:
        raise ValueError("engineering map contains no usable primitives")
    return result


def feature_key(row: dict[str, Any]) -> str:
    return f"{int(row['indicator_id'])}:{int(row['period'])}:{float(row['p1']):g}"


def load_semantics(path: Path, engineering: dict[str, dict[str, Any]]) -> dict[str, dict[str, str]]:
    try:
        import pyarrow.parquet as pq
    except ImportError as error:
        raise ValueError("pyarrow is required to read semantic-map parquet") from error
    result = {key: {"domain": "unclassified", "output_type": "continuous"} for key in engineering}
    for row in pq.read_table(path).to_pylist():
        try:
            key = str(row.get("feature_key") or feature_key(row))
        except (KeyError, TypeError, ValueError):
            continue
        if key not in result:
            continue
        raw = str(row.get("output_type", row.get("semantic_type", row.get("value_type", "continuous")))).lower()
        kind = "event" if raw in {"event", "detection", "binary_event"} else "state" if raw in {"state", "categorical", "boolean"} else "continuous"
        result[key] = {"domain": str(row.get("domain", row.get("semantic_domain", row.get("engineering_family", row.get("family", "unclassified"))))), "output_type": kind}
    return result


def feature_keys_from_row(row: dict[str, Any]) -> list[str]:
    """Read explicit triples first, then the persisted flat five-triple genome."""
    for field in ("feature_triples_json", "feature_triples"):
        value = row.get(field)
        if isinstance(value, str):
            try:
                value = json.loads(value)
            except json.JSONDecodeError:
                continue
        if isinstance(value, list):
            keys = []
            for triple in value:
                if isinstance(triple, dict):
                    try:
                        keys.append(feature_key(triple))
                    except (KeyError, TypeError, ValueError):
                        pass
            if keys:
                return keys
    value = row.get("genome_json", row.get("genome"))
    if isinstance(value, str):
        try:
            value = json.loads(value)
        except json.JSONDecodeError:
            return []
    combo = value.get("combo") if isinstance(value, dict) else None
    if isinstance(combo, str):
        try:
            combo = json.loads(combo)
        except json.JSONDecodeError:
            return []
    if not isinstance(combo, list) or len(combo) < 15:
        return []
    try:
        return [f"{int(combo[index])}:{int(combo[index + 1])}:{float(combo[index + 2]):g}" for index in range(0, 15, 3)]
    except (TypeError, ValueError):
        return []


def entry_bar(row: dict[str, Any]) -> tuple[int | None, str | None]:
    """Resolve only an integer bar offset; timestamps are intentionally not guessed."""
    containers = [("column", row)]
    for field in ("raw_ledger_json", "context_json"):
        value = row.get(field)
        if isinstance(value, str):
            try:
                value = json.loads(value)
            except json.JSONDecodeError:
                continue
        if isinstance(value, dict):
            containers.append((field, value))
    for source, values in containers:
        for field in ENTRY_FIELDS:
            value = values.get(field)
            if isinstance(value, bool) or value is None:
                continue
            try:
                numeric = float(value)
            except (TypeError, ValueError):
                continue
            if np.isfinite(numeric) and numeric.is_integer() and numeric >= 0:
                return int(numeric), f"{source}.{field}"
    return None, None


def checkpoint_paths(directory: Path, engineering: dict[str, dict[str, Any]]) -> dict[str, Path]:
    paths = {}
    for path in directory.glob("*.npy"):
        parts = path.stem.split("_")
        if len(parts) != 4 or parts[0] != "hs22":
            continue
        try:
            key = f"{int(parts[1])}:{int(parts[2])}:{float(parts[3]):g}"
        except ValueError:
            continue
        if key in engineering and key not in paths:
            paths[key] = path
    return paths


def ks_statistic(good: np.ndarray, failure: np.ndarray) -> tuple[float, str]:
    try:
        from scipy.stats import ks_2samp

        return float(ks_2samp(good, failure, method="auto").statistic), "scipy_ks_2samp"
    except ImportError:
        # Exact empirical CDF distance evaluated at every observed rank boundary.
        values = np.sort(np.concatenate((good, failure)))
        left = np.searchsorted(np.sort(good), values, side="right") / good.size
        right = np.searchsorted(np.sort(failure), values, side="right") / failure.size
        return float(np.max(np.abs(left - right))), "exact_rank_approx"


def odds_ratio(good: np.ndarray, failure: np.ndarray) -> float | None:
    if not np.all(np.isin(np.concatenate((good, failure)), (0, 1))):
        return None
    # Haldane-Anscombe correction keeps complete separation reportable.
    good_on, failure_on = good.sum(), failure.sum()
    return float(((failure_on + 0.5) / (failure.size - failure_on + 0.5)) / ((good_on + 0.5) / (good.size - good_on + 0.5)))


def comparison(good: list[tuple[float, str]], failure: list[tuple[float, str]]) -> dict[str, Any]:
    good_values, failure_values = np.array([item[0] for item in good]), np.array([item[0] for item in failure])
    median_difference = float(np.median(failure_values) - np.median(good_values))
    pooled = np.concatenate((good_values, failure_values))
    good_variance = np.var(good_values, ddof=1) if good_values.size > 1 else 0.0
    failure_variance = np.var(failure_values, ddof=1) if failure_values.size > 1 else 0.0
    degrees_of_freedom = good_values.size + failure_values.size - 2
    scale = float(np.sqrt(((good_values.size - 1) * good_variance + (failure_values.size - 1) * failure_variance) / degrees_of_freedom)) if degrees_of_freedom else 0.0
    effect = median_difference / scale if scale else None
    pooled_sorted = np.sort(pooled)
    quantile_separation = float(
        np.searchsorted(pooled_sorted, np.median(failure_values), side="right") / pooled.size
        - np.searchsorted(pooled_sorted, np.median(good_values), side="right") / pooled.size
    )
    ks, ks_method = ks_statistic(good_values, failure_values)
    direction = np.sign(median_difference)
    folds = []
    for fold in sorted(set(item[1] for item in good) & set(item[1] for item in failure)):
        fold_difference = np.median([item[0] for item in failure if item[1] == fold]) - np.median([item[0] for item in good if item[1] == fold])
        folds.append(float(fold_difference))
    consistent = sum(np.sign(item) == direction for item in folds) if direction else sum(item == 0 for item in folds)
    return {
        "good_samples": int(good_values.size), "failure_samples": int(failure_values.size),
        "median_difference_failure_minus_good": median_difference,
        "standardized_effect_size": effect, "standardized_effect_size_method": "median_difference_over_pooled_within_group_standard_deviation", "quantile_separation": quantile_separation,
        "ks_statistic": ks, "ks_method": ks_method,
        "folds_compared": len(folds), "fold_direction_consistent": consistent,
        "fold_consistency": consistent / len(folds) if folds else None,
    }


def mine(rows: list[dict[str, Any]], engineering: dict[str, dict[str, Any]], semantics: dict[str, dict[str, str]], paths: dict[str, Path]) -> tuple[list[dict[str, Any]], dict[str, int], list[str]]:
    samples: dict[str, dict[str, list[tuple[float, str]]]] = defaultdict(lambda: defaultdict(list))
    skipped = Counter()
    schemas = set()
    arrays: dict[str, np.ndarray] = {}
    for row in rows:
        label = row.get("failure_label")
        if label not in (GOOD_LABEL, *FAILURE_LABELS):
            continue
        bar, schema = entry_bar(row)
        keys = feature_keys_from_row(row)
        if bar is None:
            skipped["unmappable_entry_bar"] += 1
            continue
        if not keys:
            skipped["unmappable_feature_triples"] += 1
            continue
        schemas.add(schema)
        fold = str(row.get("fold", "unknown"))
        for key in set(keys):
            if key not in engineering or key not in paths:
                skipped["feature_missing_from_oracle"] += 1
                continue
            values = arrays.setdefault(key, np.load(paths[key], allow_pickle=False, mmap_mode="r").reshape(-1))
            if bar >= values.size:
                skipped["entry_bar_outside_oracle"] += 1
                continue
            value = float(values[bar])
            if not np.isfinite(value):
                skipped["nonfinite_oracle_value"] += 1
                continue
            samples[key][label].append((value, fold))
    output = []
    for key in sorted(samples):
        good = samples[key][GOOD_LABEL]
        if not good:
            continue
        for label in FAILURE_LABELS:
            failure = samples[key][label]
            if not failure:
                continue
            result = {"feature_key": key, "failure_label": label, "attribution": "ASSOCIATIVE_NOT_CAUSAL", **engineering[key], **semantics[key], **comparison(good, failure)}
            if semantics[key]["output_type"] in {"state", "event"}:
                result["state_event_odds_ratio_failure_vs_good"] = odds_ratio(np.array([x[0] for x in good]), np.array([x[0] for x in failure]))
            output.append(result)
    return output, dict(sorted(skipped.items())), sorted(schema for schema in schemas if schema)


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--failure-context", type=Path, required=True)
    parser.add_argument("--semantic-map", type=Path, required=True)
    parser.add_argument("--oracle-checkpoint-dir", type=Path, required=True)
    parser.add_argument("--engineering-map", type=Path, required=True)
    parser.add_argument("--output-dir", type=Path, required=True)
    args = parser.parse_args()
    if not args.oracle_checkpoint_dir.is_dir():
        parser.error("--oracle-checkpoint-dir must be a directory")
    for path in (args.failure_context, args.semantic_map, args.engineering_map):
        if not path.is_file():
            parser.error(f"input is not a file: {path}")
    try:
        import pyarrow as pa
        import pyarrow.parquet as pq
    except ImportError as error:
        raise SystemExit("pyarrow is required for Cohort001 failure mining") from error
    engineering = load_engineering_map(args.engineering_map)
    semantics = load_semantics(args.semantic_map, engineering)
    paths = checkpoint_paths(args.oracle_checkpoint_dir, engineering)
    rows = pq.read_table(args.failure_context).to_pylist()
    features, skipped, entry_schemas = mine(rows, engineering, semantics, paths)
    domains: dict[tuple[str, str], list[dict[str, Any]]] = defaultdict(list)
    for item in features:
        domains[(item["domain"], item["failure_label"])].append(item)
    domain_rows = [{
        "domain": domain, "failure_label": label, "feature_comparisons": len(items),
        "median_standardized_effect_size": float(np.median([x["standardized_effect_size"] for x in items if x["standardized_effect_size"] is not None])) if any(x["standardized_effect_size"] is not None for x in items) else None,
        "median_ks_statistic": float(np.median([x["ks_statistic"] for x in items])),
        "median_fold_consistency": float(np.median([x["fold_consistency"] for x in items if x["fold_consistency"] is not None])) if any(x["fold_consistency"] is not None for x in items) else None,
        "attribution": "ASSOCIATIVE_NOT_CAUSAL",
    } for (domain, label), items in sorted(domains.items())]
    result = {
        "schema_version": 1, "artifact": ARTIFACT, "read_only": True,
        "attribution": "ASSOCIATIVE_NOT_CAUSAL",
        "attribution_note": "Entry-time associations do not establish causal feature effects.",
        "inputs": {
            str(path): file_hash(path) for path in (args.failure_context, args.semantic_map, args.engineering_map)
        } | {str(args.oracle_checkpoint_dir): directory_hash(args.oracle_checkpoint_dir)},
        "entry_bar_schema_detected": entry_schemas,
        "failure_context_rows": len(rows), "oracle_features_available": len(paths),
        "skipped": skipped, "feature_failure_comparisons": features, "domain_aggregates": domain_rows,
    }
    args.output_dir.mkdir(parents=True, exist_ok=True)
    pq.write_table(pa.Table.from_pylist(features), args.output_dir / "cohort001_feature_failure_mining_v1.parquet", compression="zstd")
    output = args.output_dir / "cohort001_feature_failure_mining_v1.json"
    output.write_text(json.dumps(json_compatible(result), indent=2, sort_keys=True, allow_nan=False) + "\n", encoding="utf-8")
    print(output)


if __name__ == "__main__":
    main()