code/downstream/features.py

"""Frozen transition-model feature extraction for candidate plans."""

from __future__ import annotations

import json
import math
from dataclasses import dataclass
from pathlib import Path

import numpy as np

from code.downstream.families import SourceFamily
from code.downstream.plan_utils import ProblemContext, load_problem_context, rollout_plan
from code.tokenization.factory import create_tokenizer
from code.tokenization.multidomain import MultiDomainUnionTokenizer, load_tokenizer_from_manifest


SERIES_NAMES = [
    "residual_l2",
    "residual_l1_mean",
    "cosine_distance",
    "pred_norm",
    "pred_delta_norm",
    "candidate_delta_norm",
    "current_goal_l2",
    "pred_goal_l2",
    "goal_progress_l2",
]
SUMMARY_STATS = ["mean", "std", "min", "max", "final"]
SCALAR_FEATURES = [
    "plan_len",
    "log_plan_len",
    "plan_to_budget_ratio",
    "initial_goal_l2",
    "final_goal_l2",
    "final_minus_initial_goal_l2",
    "hidden_h_final_norm",
    "hidden_c_final_norm",
    "hidden_h_abs_mean",
    "hidden_c_abs_mean",
]
FEATURE_NAMES = [
    f"{series}_{stat}"
    for series in SERIES_NAMES
    for stat in SUMMARY_STATS
] + SCALAR_FEATURES


@dataclass
class LoadedSource:
    """Loaded frozen model/tokenizer artifacts for one family/domain/seed."""

    family: SourceFamily
    domain: str
    model: object
    tokenizer: object
    mode: str
    model_kind: str
    input_dim: int
    device: object | None = None


class FrozenTransitionFeatureExtractor:
    """Extract fixed-size consistency summaries from frozen transition models."""

    def __init__(
        self,
        *,
        run_root: str | Path,
        source_data_dir: str | Path,
        family: SourceFamily,
        seed: int,
        device: str = "cpu",
        xgb_n_jobs: int = 1,
    ) -> None:
        self.run_root = Path(run_root)
        self.source_data_dir = Path(source_data_dir)
        self.family = family
        self.seed = seed
        self.device_arg = device
        self.xgb_n_jobs = xgb_n_jobs
        self._source_cache: dict[str, LoadedSource] = {}
        self._problem_cache: dict[tuple[str, str, str], ProblemContext] = {}
        self._wl_problem_cache: dict[tuple[str, str, str], object] = {}

    @property
    def feature_names(self) -> list[str]:
        return list(FEATURE_NAMES)

    def extract(self, candidate: dict) -> np.ndarray:
        """Extract one feature vector for a candidate-plan record."""
        domain = candidate["domain"]
        split = candidate["split"]
        problem = candidate["problem"]
        plan = [str(action) for action in candidate.get("plan", [])]

        source = self._load_source(domain)
        problem_context = self._load_problem_context(domain, split, problem)
        states = rollout_plan(problem_context, plan)

        goal_vec = self._embed_goal(source, problem_context)
        initial_vec = self._embed_state(source, problem_context, states[0])
        final_vec = self._embed_state(source, problem_context, states[-1])
        initial_goal_l2 = _l2(initial_vec - goal_vec)
        final_goal_l2 = _l2(final_vec - goal_vec)

        series_values = {name: [] for name in SERIES_NAMES}
        hidden_summary = {
            "hidden_h_final_norm": 0.0,
            "hidden_c_final_norm": 0.0,
            "hidden_h_abs_mean": 0.0,
            "hidden_c_abs_mean": 0.0,
        }

        hidden = None
        for current_atoms, next_atoms in zip(states, states[1:]):
            current_vec = self._embed_state(source, problem_context, current_atoms)
            next_vec = self._embed_state(source, problem_context, next_atoms)
            pred_next_vec, pred_delta_vec, hidden = self._predict_next(
                source,
                current_vec,
                goal_vec,
                hidden,
            )

            candidate_delta = next_vec - current_vec
            residual = pred_next_vec - next_vec
            current_goal_l2 = _l2(current_vec - goal_vec)
            next_goal_l2 = _l2(next_vec - goal_vec)
            pred_goal_l2 = _l2(pred_next_vec - goal_vec)

            series_values["residual_l2"].append(_l2(residual))
            series_values["residual_l1_mean"].append(float(np.mean(np.abs(residual))))
            series_values["cosine_distance"].append(_cosine_distance(pred_next_vec, next_vec))
            series_values["pred_norm"].append(_l2(pred_next_vec))
            series_values["pred_delta_norm"].append(_l2(pred_delta_vec))
            series_values["candidate_delta_norm"].append(_l2(candidate_delta))
            series_values["current_goal_l2"].append(current_goal_l2)
            series_values["pred_goal_l2"].append(pred_goal_l2)
            series_values["goal_progress_l2"].append(current_goal_l2 - next_goal_l2)

        if source.model_kind == "lstm" and hidden is not None:
            hidden_summary.update(_summarize_lstm_hidden(hidden))

        values: list[float] = []
        for series in SERIES_NAMES:
            values.extend(_summarize_series(series_values[series]))

        plan_len = float(len(plan))
        transition_budget = float(max(100, 10 * max(1, len(problem_context.objects))))
        values.extend(
            [
                plan_len,
                math.log1p(plan_len),
                plan_len / transition_budget,
                initial_goal_l2,
                final_goal_l2,
                final_goal_l2 - initial_goal_l2,
                hidden_summary["hidden_h_final_norm"],
                hidden_summary["hidden_c_final_norm"],
                hidden_summary["hidden_h_abs_mean"],
                hidden_summary["hidden_c_abs_mean"],
            ]
        )

        return np.nan_to_num(np.asarray(values, dtype=np.float32))

    def _load_problem_context(self, domain: str, split: str, problem: str) -> ProblemContext:
        key = (domain, split, problem)
        cached = self._problem_cache.get(key)
        if cached is None:
            cached = load_problem_context(
                data_dir=self.source_data_dir,
                domain=domain,
                split=split,
                problem=problem,
            )
            self._problem_cache[key] = cached
        return cached

    def _load_source(self, domain: str) -> LoadedSource:
        cached = self._source_cache.get(domain)
        if cached is not None:
            return cached

        if self.family.model == "xgboost":
            source = self._load_xgb_source(domain)
        elif self.family.model == "lstm":
            source = self._load_lstm_source(domain)
        else:
            raise ValueError(f"Unsupported source model: {self.family.model}")

        self._source_cache[domain] = source
        return source

    def _load_xgb_source(self, domain: str) -> LoadedSource:
        import xgboost as xgb

        checkpoint_dir = self._checkpoint_dir(domain)
        model_name = domain if self.family.is_domain_dependent else "all_domains"
        model_path = checkpoint_dir / f"{model_name}_xgb.json"
        meta_path = checkpoint_dir / f"{model_name}_xgb_meta.json"
        if not model_path.exists():
            raise FileNotFoundError(f"Missing XGBoost source model: {model_path}")

        model = xgb.XGBRegressor(device="cpu", n_jobs=self.xgb_n_jobs)
        model.load_model(str(model_path))
        meta = _read_json(meta_path)
        tokenizer = self._load_tokenizer(domain)
        input_dim = int(meta.get("output_dim", tokenizer.get_embedding_dim()))

        return LoadedSource(
            family=self.family,
            domain=domain,
            model=model,
            tokenizer=tokenizer,
            mode=str(meta.get("delta", True) and "delta" or "state"),
            model_kind="xgboost",
            input_dim=input_dim,
        )

    def _load_lstm_source(self, domain: str) -> LoadedSource:
        import torch

        from code.modeling.models import StateCentricLSTM, StateCentricLSTM_Delta

        checkpoint_dir = self._checkpoint_dir(domain)
        model_name = domain if self.family.is_domain_dependent else "all_domains"
        checkpoint_path = checkpoint_dir / f"{model_name}_lstm_best.pt"
        meta_path = checkpoint_dir / f"{model_name}_lstm_meta.json"
        if not checkpoint_path.exists():
            raise FileNotFoundError(f"Missing LSTM source model: {checkpoint_path}")

        meta = _read_json(meta_path)
        tokenizer = self._load_tokenizer(domain)
        input_dim = int(meta.get("input_dim", tokenizer.get_embedding_dim()))
        hidden_dim = int(meta.get("hidden_dim", 256))
        use_projection = not bool(meta.get("no_projection", False))
        mode = str(meta.get("mode", self.family.mode))

        device = _resolve_torch_device(self.device_arg)
        if mode == "delta":
            model = StateCentricLSTM_Delta(
                input_dim,
                hidden_dim=hidden_dim,
                use_projection=use_projection,
            ).to(device)
        else:
            model = StateCentricLSTM(
                input_dim,
                hidden_dim=hidden_dim,
                use_projection=use_projection,
            ).to(device)
        model.load_state_dict(torch.load(checkpoint_path, map_location=device))
        model.eval()

        return LoadedSource(
            family=self.family,
            domain=domain,
            model=model,
            tokenizer=tokenizer,
            mode=mode,
            model_kind="lstm",
            input_dim=input_dim,
            device=device,
        )

    def _checkpoint_dir(self, domain: str) -> Path:
        root = self.run_root / self.family.regime / "checkpoints" / f"seed_{self.seed}"
        root = root / self.family.tokenizer / self.family.model_mode_dir
        if self.family.is_domain_dependent:
            root = root / domain
        return root

    def _load_tokenizer(self, domain: str):
        domain_pddl = self.source_data_dir / "pddl" / domain / "domain.pddl"
        if self.family.is_all_domains:
            manifest = (
                self.run_root
                / "all_domains"
                / "tokenizers"
                / self.family.tokenizer
                / f"all_domains_{self.family.tokenizer}.json"
            )
            tokenizer = load_tokenizer_from_manifest(str(manifest))
            if isinstance(tokenizer, MultiDomainUnionTokenizer):
                tokenizer.set_active_domain(domain, str(domain_pddl))
            elif hasattr(tokenizer, "set_domain"):
                tokenizer.set_domain(str(domain_pddl))
            return tokenizer

        tokenizer = create_tokenizer(self.family.tokenizer)
        vocab_path = self._domain_vocab_path(domain)
        if not vocab_path.exists():
            raise FileNotFoundError(f"Missing tokenizer vocabulary: {vocab_path}")
        tokenizer.load_vocabulary(str(vocab_path))
        if hasattr(tokenizer, "set_domain"):
            tokenizer.set_domain(str(domain_pddl))
        return tokenizer

    def _domain_vocab_path(self, domain: str) -> Path:
        model_dir = (
            self.run_root
            / "domain_dependent"
            / "data"
            / "encodings"
            / "models"
        )
        if self.family.tokenizer == "wl":
            return model_dir / f"{domain}_wl_tok.json"
        return model_dir / f"{domain}_{self.family.tokenizer}.json"

    def _embed_state(
        self,
        source: LoadedSource,
        context: ProblemContext,
        atoms: frozenset[str],
    ) -> np.ndarray:
        tokenizer = source.tokenizer
        wl_prob = self._wl_problem(source, context)
        state_atoms = list(atoms)
        try:
            vec = tokenizer.transform_state(
                state_atoms,
                list(context.goal_atoms),
                list(context.objects),
                problem_pddl_path=str(context.problem_path),
                _wl_prob=wl_prob,
            )
        except TypeError:
            try:
                vec = tokenizer.transform_state(
                    state_atoms,
                    list(context.goal_atoms),
                    list(context.objects),
                    problem_pddl_path=str(context.problem_path),
                )
            except TypeError:
                vec = tokenizer.transform_state(
                    state_atoms,
                    list(context.goal_atoms),
                    list(context.objects),
                )
        return np.asarray(vec, dtype=np.float32).reshape(-1)

    def _embed_goal(self, source: LoadedSource, context: ProblemContext) -> np.ndarray:
        tokenizer = source.tokenizer
        wl_prob = self._wl_problem(source, context)
        try:
            vec = tokenizer.transform_goal(
                list(context.goal_atoms),
                list(context.objects),
                problem_pddl_path=str(context.problem_path),
                _wl_prob=wl_prob,
            )
        except TypeError:
            try:
                vec = tokenizer.transform_goal(
                    list(context.goal_atoms),
                    list(context.objects),
                    problem_pddl_path=str(context.problem_path),
                )
            except TypeError:
                vec = tokenizer.transform_goal(list(context.goal_atoms), list(context.objects))
        return np.asarray(vec, dtype=np.float32).reshape(-1)

    def _predict_next(
        self,
        source: LoadedSource,
        current_vec: np.ndarray,
        goal_vec: np.ndarray,
        hidden,
    ) -> tuple[np.ndarray, np.ndarray, object | None]:
        if source.model_kind == "xgboost":
            x = np.hstack([current_vec.reshape(1, -1), goal_vec.reshape(1, -1)])
            pred = np.asarray(source.model.predict(x), dtype=np.float32).reshape(-1)
            if source.mode == "delta":
                return current_vec + pred, pred, hidden
            return pred, pred - current_vec, hidden

        return self._predict_lstm(source, current_vec, goal_vec, hidden)

    def _predict_lstm(
        self,
        source: LoadedSource,
        current_vec: np.ndarray,
        goal_vec: np.ndarray,
        hidden,
    ) -> tuple[np.ndarray, np.ndarray, object | None]:
        import torch

        device = source.device
        state_tensor = (
            torch.tensor(current_vec, dtype=torch.float32, device=device)
            .reshape(1, 1, -1)
        )
        goal_tensor = torch.tensor(goal_vec, dtype=torch.float32, device=device).reshape(1, -1)
        with torch.inference_mode():
            pred, next_hidden = source.model(state_tensor, goal_tensor, hidden=hidden)
        pred_vec = pred.detach().cpu().numpy().reshape(-1).astype(np.float32)
        if source.mode == "delta":
            return current_vec + pred_vec, pred_vec, next_hidden
        return pred_vec, pred_vec - current_vec, next_hidden

    def _wl_problem(self, source: LoadedSource, context: ProblemContext):
        if source.family.tokenizer != "wl" and not isinstance(source.tokenizer, MultiDomainUnionTokenizer):
            return None
        key = (context.domain, context.split, context.problem)
        if key not in self._wl_problem_cache:
            try:
                from wlplan.planning import parse_problem as wl_parse_problem

                self._wl_problem_cache[key] = wl_parse_problem(
                    str(context.domain_path),
                    str(context.problem_path),
                )
            except Exception:
                self._wl_problem_cache[key] = None
        return self._wl_problem_cache[key]


def save_feature_matrix(
    *,
    path: str | Path,
    candidates: list[dict],
    features: np.ndarray,
    feature_names: list[str],
) -> None:
    """Save feature arrays and row metadata in one compressed artifact."""
    path = Path(path)
    path.parent.mkdir(parents=True, exist_ok=True)
    np.savez_compressed(
        path,
        X=features.astype(np.float32),
        y=np.asarray([int(row["label_valid"]) for row in candidates], dtype=np.int64),
        candidate_ids=np.asarray([row["candidate_id"] for row in candidates], dtype=object),
        domains=np.asarray([row["domain"] for row in candidates], dtype=object),
        splits=np.asarray([row["split"] for row in candidates], dtype=object),
        problems=np.asarray([row["problem"] for row in candidates], dtype=object),
        corruption_types=np.asarray(
            [row["corruption_type"] for row in candidates],
            dtype=object,
        ),
        feature_names=np.asarray(feature_names, dtype=object),
    )


def load_feature_matrix(path: str | Path) -> dict:
    """Load a feature matrix artifact with object metadata enabled."""
    return dict(np.load(path, allow_pickle=True))


def _read_json(path: Path) -> dict:
    if not path.exists():
        return {}
    with open(path, "r", encoding="utf-8") as f:
        return json.load(f)


def _resolve_torch_device(device_arg: str):
    import torch

    if device_arg == "auto":
        if torch.cuda.is_available():
            return torch.device("cuda")
        if torch.backends.mps.is_available():
            return torch.device("mps")
        return torch.device("cpu")
    if device_arg == "cuda" and not torch.cuda.is_available():
        return torch.device("cpu")
    if device_arg == "mps" and not torch.backends.mps.is_available():
        return torch.device("cpu")
    return torch.device(device_arg)


def _summarize_series(values: list[float]) -> list[float]:
    if not values:
        return [0.0 for _ in SUMMARY_STATS]
    arr = np.asarray(values, dtype=np.float32)
    return [
        float(np.mean(arr)),
        float(np.std(arr)),
        float(np.min(arr)),
        float(np.max(arr)),
        float(arr[-1]),
    ]


def _summarize_lstm_hidden(hidden) -> dict[str, float]:
    h, c = hidden
    h_cpu = h.detach().float().cpu()
    c_cpu = c.detach().float().cpu()
    return {
        "hidden_h_final_norm": float(h_cpu[-1].norm().item()),
        "hidden_c_final_norm": float(c_cpu[-1].norm().item()),
        "hidden_h_abs_mean": float(h_cpu.abs().mean().item()),
        "hidden_c_abs_mean": float(c_cpu.abs().mean().item()),
    }


def _l2(vec: np.ndarray) -> float:
    return float(np.linalg.norm(np.asarray(vec, dtype=np.float32).reshape(-1)))


def _cosine_distance(left: np.ndarray, right: np.ndarray) -> float:
    u = np.asarray(left, dtype=np.float32).reshape(-1)
    v = np.asarray(right, dtype=np.float32).reshape(-1)
    denom = float(np.linalg.norm(u) * np.linalg.norm(v))
    if denom == 0.0:
        return 1.0
    return float(1.0 - np.dot(u, v) / denom)