code/modeling/inference_lstm.py

import argparse
from concurrent.futures import ThreadPoolExecutor
import json
import os
import pickle
import sys
from pathlib import Path
from code.common.fsf_wrapper import FSFEncoder
from code.common.utils import set_seed, validate_plan
from code.modeling.models import StateCentricLSTM, StateCentricLSTM_Delta
from code.tokenization.factory import create_tokenizer
from code.tokenization.multidomain import MultiDomainUnionTokenizer, load_tokenizer_from_manifest

import numpy as np
import torch
import torch.nn.functional as F
from pyperplan.grounding import ground
from pyperplan.pddl.parser import Parser
from tqdm import tqdm


DEFAULT_SPLITS = ["validation", "test-interpolation", "test-extrapolation"]


def progress_enabled() -> bool:
    return bool(sys.stdout.isatty())


def resolve_device(device_arg: str) -> torch.device:
    """Resolve runtime device from CLI preference."""
    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":
        if not torch.cuda.is_available():
            raise RuntimeError("CUDA requested but not available.")
        return torch.device("cuda")

    if device_arg == "mps":
        if not torch.backends.mps.is_available():
            raise RuntimeError("MPS requested but not available.")
        return torch.device("mps")

    return torch.device("cpu")


def normalize_encoding_name(name: str) -> str:
    """Normalize encoding aliases to canonical tokenizer names."""
    return "wl" if name == "graphs" else name


def score_distances(
    pred_batch: torch.Tensor,
    candidate_batch: torch.Tensor,
    *,
    score_metric: str,
    delta: bool,
) -> torch.Tensor:
    """
    Compute successor distances under the requested metric.

    `native` preserves the original behavior from the main study:
      - cosine distance for state prediction
      - L2 distance for delta prediction
    """
    metric = score_metric
    if metric == "native":
        metric = "l2" if delta else "cosine"

    if metric == "l2":
        return torch.norm(pred_batch - candidate_batch, p=2, dim=-1).reshape(-1)
    if metric == "cosine":
        return 1.0 - F.cosine_similarity(pred_batch, candidate_batch, dim=-1).reshape(-1)
    raise ValueError(f"Unsupported score metric: {score_metric}")


def resolve_model_artifact_path(data_dir: str, filename: str) -> str:
    """
    Resolve files from data/encodings/models robustly.
    Supports callers passing either `data` or an encoding subdir.
    """
    p = Path(data_dir).resolve()
    search_dirs = [
        p / "encodings" / "models",
        p / "models",
        p.parent / "models",
        p.parent.parent / "models",
        Path("data") / "encodings" / "models",
    ]

    seen = set()
    for d in search_dirs:
        d_str = str(d)
        if d_str in seen:
            continue
        seen.add(d_str)
        candidate = d / filename
        if candidate.exists():
            return str(candidate)

    # Default fallback path for clear error messages.
    return str(search_dirs[0] / filename)


def resolve_vocab_path(data_dir: str, domain: str, raw_encoding: str) -> str | None:
    """Find the most likely vocabulary file for a tokenizer."""
    normalized = normalize_encoding_name(raw_encoding)
    names = [f"{domain}_{normalized}.json"]

    if normalized == "wl":
        names.extend(
            [
                f"{domain}_wl_tok.json",
                f"{domain}_wl.json",
                f"{domain}_graphs.json",
            ]
        )

    if raw_encoding != normalized:
        names.append(f"{domain}_{raw_encoding}.json")

    for name in names:
        path = resolve_model_artifact_path(data_dir, name)
        if os.path.exists(path):
            return path

    return None


def get_fsf_tensor(atoms_set, encoder, objects, obj_map, device):
    """Helper for FSF Inference embedding"""
    # Convert set of strings to list of tuples
    atom_tuples = []
    for a in atoms_set:
        content = a.replace("(", "").replace(")", "").lower()
        atom_tuples.append(tuple(content.split()))

    vec = encoder._state_to_vector(atom_tuples, objects, obj_map)
    # [1, 1, D]
    return torch.tensor(vec).float().to(device).unsqueeze(0).unsqueeze(0)


def transform_state_compat(tokenizer, state_atoms, goal_atoms, objects, problem_path):
    """
    Call tokenizer.transform_state with optional problem path when supported.
    WLTokenizer needs it; other tokenizers generally do not.
    """
    try:
        return tokenizer.transform_state(
            state_atoms,
            goal_atoms,
            objects,
            problem_pddl_path=problem_path,
        )
    except TypeError:
        return tokenizer.transform_state(state_atoms, goal_atoms, objects)


def transform_state_cached(
    tokenizer,
    state_atoms,
    goal_atoms,
    objects,
    problem_path,
    *,
    wl_prob=None,
):
    """Call tokenizer.transform_state while reusing a pre-parsed WL problem when available."""
    try:
        return tokenizer.transform_state(
            state_atoms,
            goal_atoms,
            objects,
            problem_pddl_path=problem_path,
            _wl_prob=wl_prob,
        )
    except TypeError:
        return transform_state_compat(tokenizer, state_atoms, goal_atoms, objects, problem_path)


def transform_goal_compat(tokenizer, goal_atoms, objects, problem_path):
    """
    Call tokenizer.transform_goal with optional problem path when supported.
    WLTokenizer needs it; other tokenizers generally do not.
    """
    try:
        return tokenizer.transform_goal(
            goal_atoms,
            objects,
            problem_pddl_path=problem_path,
        )
    except TypeError:
        return tokenizer.transform_goal(goal_atoms, objects)


def transform_goal_cached(tokenizer, goal_atoms, objects, problem_path, *, wl_prob=None):
    """Call tokenizer.transform_goal while reusing a pre-parsed WL problem when available."""
    try:
        return tokenizer.transform_goal(
            goal_atoms,
            objects,
            problem_pddl_path=problem_path,
            _wl_prob=wl_prob,
        )
    except TypeError:
        return transform_goal_compat(tokenizer, goal_atoms, objects, problem_path)


def get_generic_tensor(
    atoms_set,
    tokenizer,
    goal_atoms,
    objects,
    problem_path,
    device,
    *,
    wl_prob=None,
):
    """Helper for Generic Tokenizer Inference embedding"""
    state_atoms = list(atoms_set)
    vec = transform_state_cached(
        tokenizer,
        state_atoms,
        goal_atoms,
        objects,
        problem_path,
        wl_prob=wl_prob,
    )
    # [1, 1, D]
    return torch.tensor(vec).float().to(device).unsqueeze(0).unsqueeze(0)


def extract_objects_from_problem(prob, dom) -> list[str]:
    """Extract sorted object names from a parsed pyperplan problem/domain pair."""
    objs = set()

    prob_objects = getattr(prob, "objects", {})
    if isinstance(prob_objects, dict):
        objs.update(str(name) for name in prob_objects.keys())
    else:
        for obj in prob_objects:
            objs.add(obj.name if hasattr(obj, "name") else str(obj))

    dom_constants = getattr(dom, "constants", {})
    if isinstance(dom_constants, dict):
        objs.update(str(name) for name in dom_constants.keys())
    else:
        for obj in dom_constants:
            objs.add(obj.name if hasattr(obj, "name") else str(obj))

    return sorted(objs)


def solve_problem(
    args,
    split,
    prob_file,
    model,
    device,
    encoder_type,
    feature_encoder,
    objects=None,
    obj_map=None,
    collect_search_stats: bool = False,
):
    """Unified Solver for Generic Tokenizers and FSF"""
    search_start_time = None
    if collect_search_stats:
        import time

        search_start_time = time.perf_counter()
        search_stats = {
            "beam_expansions": 0,
            "model_calls": 0,
            "successor_evals": 0,
            "outer_steps": 0,
            "terminated_reason": "max_steps",
        }

    domain_path = os.path.join(args.pddl_dir, args.domain, "domain.pddl")
    prob_path = os.path.join(args.pddl_dir, args.domain, split, prob_file)

    # 1. Pyperplan Parsing (for successors)
    try:
        parser = Parser(domain_path, prob_path)
        dom = parser.parse_domain()
        prob = parser.parse_problem(dom)
        task = ground(prob)
    except Exception as e:
        print(f"Pyperplan Parsing Error on {prob_file}: {e}")
        raise e

    # Match upstream inference behavior: allow longer searches on larger problems.
    num_objects = len(prob.objects) + len(dom.constants)
    effective_max_steps = max(args.max_steps, args.steps_per_object * num_objects)

    if objects is None and encoder_type != "fsf":
        objects = extract_objects_from_problem(prob, dom)

    initial_atoms = task.initial_state
    goal_set = set(task.goals)
    goal_atoms_list = list(goal_set)
    state_cache = {}
    successor_cache = {}
    sorted_operators = sorted(task.operators, key=lambda op: op.name)
    wl_prob = None

    if encoder_type == "wl":
        from wlplan.planning import parse_problem as wl_parse_problem

        wl_prob = wl_parse_problem(domain_path, prob_path)

    # 2. Embedding Setup
    if encoder_type == "fsf":
        encoder = feature_encoder
        objects = encoder._get_sorted_objects(prob_path)
        obj_map = encoder._get_object_indices(objects)

        def get_cached_tensor(atoms):
            key = frozenset(atoms)
            cached = state_cache.get(key)
            if cached is None:
                cached = get_fsf_tensor(atoms, encoder, objects, obj_map, device)
                state_cache[key] = cached
            return cached

        # Embed Goal
        goal_vec = encoder.embed_goal(prob_path)
        goal_tensor = torch.tensor(goal_vec).float().to(device).unsqueeze(0)  # [1, D]

        # Embed Init
        init_tensor = get_cached_tensor(initial_atoms)

    else:
        # Generic Tokenizer Logic
        tokenizer = feature_encoder

        def get_cached_tensor(atoms):
            key = frozenset(atoms)
            cached = state_cache.get(key)
            if cached is None:
                cached = get_generic_tensor(
                    atoms,
                    tokenizer,
                    goal_atoms_list,
                    objects,
                    prob_path,
                    device,
                    wl_prob=wl_prob,
                )
                state_cache[key] = cached
            return cached
        
        # Embed Goal
        goal_vec = transform_goal_cached(
            tokenizer,
            goal_atoms_list,
            objects,
            prob_path,
            wl_prob=wl_prob,
        )
        goal_tensor = torch.tensor(goal_vec).float().to(device).unsqueeze(0) # [1, D]

        # Embed Init
        init_tensor = get_cached_tensor(initial_atoms)

    # 3. Beam Search
    beam = [
        (0.0, None, init_tensor, initial_atoms, [], set())
    ]  # score, hidden, tensor, atoms, plan, visited

    def get_successors(atoms):
        state_hash = frozenset(atoms)
        cached = successor_cache.get(state_hash)
        if cached is None:
            cached = []
            for op in sorted_operators:
                if op.applicable(atoms):
                    next_atoms = op.apply(atoms)
                    cached.append((op.name, next_atoms, frozenset(next_atoms)))
            successor_cache[state_hash] = cached
        return cached

    for _ in range(effective_max_steps):
        if collect_search_stats:
            search_stats["outer_steps"] += 1
            search_stats["effective_max_steps"] = effective_max_steps
        candidates = []
        for score, hidden, last_tensor, current_atoms, plan, visited in beam:
            if collect_search_stats:
                search_stats["beam_expansions"] += 1
            # Check Goal (Internal Check)
            if goal_set.issubset(current_atoms):
                result = {
                    "problem": prob_file,
                    "search_solved": True,
                    "plan_len": len(plan),
                    "plan": plan,
                    "effective_max_steps": effective_max_steps,
                }
                if collect_search_stats:
                    search_stats["terminated_reason"] = "goal_reached"
                    search_stats["search_elapsed_sec"] = time.perf_counter() - search_start_time
                    result.update(search_stats)
                return result

            # Predict Next Latent State/Delta
            with torch.inference_mode():
                if collect_search_stats:
                    search_stats["model_calls"] += 1
                with torch.autocast(
                    device_type="cuda",
                    dtype=torch.float16,
                    enabled=getattr(args, "_use_amp", False),
                ):
                    # The model predicts the State directly
                    pred, next_hidden = model(last_tensor, goal_tensor, hidden=hidden)

                # reconstruct the next state (S_t + Delta) if delta
                # else, the model already predicts S_t+1 directly
                pred_next_emb = (last_tensor + pred) if args.delta else pred

            # Score Successors
            successor_items = []
            successor_tensors = []
            for op_name, next_atoms, next_hash in get_successors(current_atoms):
                if collect_search_stats:
                    search_stats["successor_evals"] += 1
                if next_hash in visited:
                    continue  # Skip cycles

                cand_tensor = get_cached_tensor(next_atoms)
                successor_items.append((op_name, next_atoms, next_hash, cand_tensor))
                successor_tensors.append(cand_tensor)

            if not successor_items:
                continue

            candidate_batch = torch.cat(successor_tensors, dim=0)
            pred_batch = pred_next_emb.expand(candidate_batch.shape[0], -1, -1)

            sims = score_distances(
                pred_batch,
                candidate_batch,
                score_metric=args.score_metric,
                delta=args.delta,
            )

            for (op_name, next_atoms, next_hash, cand_tensor), sim in zip(
                successor_items,
                sims.tolist(),
            ):

                # Update Score
                new_score = score + sim

                # Update Visited
                new_visited = visited.copy()
                new_visited.add(next_hash)

                # Append tuple
                candidates.append(
                    (
                        new_score,
                        next_hidden,
                        cand_tensor,
                        next_atoms,
                        plan + [op_name],
                        new_visited,
                    )
                )

            # Prune Beam
            if not candidates:
                if collect_search_stats:
                    search_stats["terminated_reason"] = "dead_end"
                break  # All paths led to dead ends

            # Stable Sort:
            candidates.sort(key=lambda x: (x[0], str(x[4])))

            beam = candidates[: args.beam_width]

    best = beam[0] if beam else (0, 0, 0, [], [], 0)

    result = {
        "problem": prob_file,
        "search_solved": False,
        "plan_len": len(best[4]),
        "plan": best[4],
        "effective_max_steps": effective_max_steps,
    }
    if collect_search_stats:
        if search_stats["terminated_reason"] == "max_steps":
            search_stats["terminated_reason"] = "max_steps"
        search_stats["search_elapsed_sec"] = time.perf_counter() - search_start_time
        result.update(search_stats)
    return result


def run_inference(args):
    set_seed(args.seed)

    device = resolve_device(args.device)
    args._use_amp = bool(args.amp and device.type == "cuda")
    if args.fast and device.type == "cuda":
        torch.backends.cudnn.benchmark = True
        torch.backends.cudnn.deterministic = False
        torch.use_deterministic_algorithms(False)
        torch.set_float32_matmul_precision("high")

    print(f"Using device: {device}")
    if device.type == "cuda":
        print(f"CUDA device: {torch.cuda.get_device_name(0)}")
    print(f"AMP: {'enabled' if args._use_amp else 'disabled'}")
    
    # Common variables
    domain_pddl = os.path.join(args.pddl_dir, args.domain, "domain.pddl")

    # Load Encoder
    feature_encoder = None
    input_dim = 0
    
    normalized_encoding = normalize_encoding_name(args.encoding)

    if normalized_encoding == "fsf":
        # 1. Load Config
        config_path = resolve_model_artifact_path(
            args.data_dir, f"{args.domain}_fsf_config.json"
        )
        if not os.path.exists(config_path):
            print(f"Error: FSF Config not found at {config_path}")
            return

        with open(config_path, "r") as f:
            config = json.load(f)
            max_objects = config["max_objects"]

        # 2. Init Encoder
        feature_encoder = FSFEncoder(args.domain, domain_pddl, max_objects)

        # 3. Set Input Dim (Max Objects + 1 Global)
        input_dim = feature_encoder.vector_size
        print(f"FSF Input Dimension: {input_dim}")
    else:
        # Generic Tokenizer
        try:
            if args.tokenizer_manifest:
                tokenizer = load_tokenizer_from_manifest(args.tokenizer_manifest)
                print(f"Loaded tokenizer manifest from {args.tokenizer_manifest}")
            else:
                vocab_path = resolve_vocab_path(args.data_dir, args.domain, args.encoding)
                tokenizer = create_tokenizer(normalized_encoding)

                if vocab_path and os.path.exists(vocab_path):
                    tokenizer.load_vocabulary(vocab_path)
                    print(f"Loaded {normalized_encoding} vocabulary from {vocab_path}")
                else:
                    print(
                        f"Warning: Vocabulary file not found for '{args.encoding}'. Using default params."
                    )

            if isinstance(tokenizer, MultiDomainUnionTokenizer):
                tokenizer.set_active_domain(args.domain, domain_pddl)
            elif hasattr(tokenizer, "set_domain"):
                tokenizer.set_domain(domain_pddl)

            feature_encoder = tokenizer
            input_dim = tokenizer.get_embedding_dim()

        except Exception as e:
            print(f"Failed to initialize tokenizer '{args.encoding}': {e}")
            return

    # Load Model
    print(f"Loading LSTM from {args.checkpoint}...")

    # Determine projection usage
    use_projection = not args.no_projection

    if args.delta:
        model = StateCentricLSTM_Delta(
            input_dim, hidden_dim=args.hidden_dim, use_projection=use_projection
        ).to(device)
    else:
        model = StateCentricLSTM(
            input_dim, hidden_dim=args.hidden_dim, use_projection=use_projection
        ).to(device)

    model.load_state_dict(torch.load(args.checkpoint, map_location=device))
    model.eval()

    all_split_solved = []
    all_split_exec = []

    # 4. Run on Splits
    splits = args.splits
    try:
        validation_workers = max(1, args.validation_workers)
        for split in splits:
            print(f"\n*** Testing on {split} ***")
            split_dir = os.path.join(args.pddl_dir, args.domain, split)
            if not os.path.exists(split_dir):
                print(f"Skipping {split} (not found)")
                continue

            results = []
            solved_count = 0
            executable_count = 0
            prob_files = sorted([f for f in os.listdir(split_dir) if f.endswith(".pddl")])

            if args.problems:
                requested = set(args.problems)
                prob_files = [f for f in prob_files if f in requested]

            if args.max_problems is not None:
                prob_files = prob_files[: args.max_problems]

            print(f" Found {len(prob_files)} problems for {split}")

            pending_validations = []
            with ThreadPoolExecutor(max_workers=validation_workers) as validation_pool:
                for prob_file in tqdm(
                    prob_files,
                    desc=f"Solving {split}",
                    disable=(not progress_enabled()),
                ):
                    prob_path = os.path.join(split_dir, prob_file)

                    try:
                        # Generate Plan
                        res = solve_problem(
                            args,
                            split,
                            prob_file,
                            model,
                            device,
                            normalized_encoding,
                            feature_encoder,
                            objects=None,
                            collect_search_stats=args.collect_search_stats,
                        )

                        if args.skip_validation:
                            # Smoke-mode fallback: trust internal goal check.
                            is_solved = bool(res.get("search_solved", False))
                            is_executable = is_solved
                            res["val_skipped"] = True
                            res["val_solved"] = is_solved
                            res["val_executable"] = is_executable
                            res["solved"] = is_solved
                            results.append(res)
                            if is_solved:
                                solved_count += 1
                            if is_executable:
                                executable_count += 1
                        else:
                            future = validation_pool.submit(
                                validate_plan,
                                domain_pddl,
                                prob_path,
                                res["plan"],
                                args.val_path,
                            )
                            pending_validations.append((prob_file, res, future))

                    except Exception as e:
                        import traceback

                        traceback.print_exc()
                        print(f"Error processing {prob_file}: {e}")
                        results.append({"problem": prob_file, "solved": False, "error": str(e)})

                if not args.skip_validation:
                    for prob_file, res, future in pending_validations:
                        try:
                            is_solved, is_executable = future.result()
                        except Exception as e:
                            print(f"Validation failed for {prob_file}: {e}")
                            is_solved, is_executable = False, False

                        res["val_solved"] = is_solved
                        res["val_executable"] = is_executable
                        res["solved"] = is_solved
                        results.append(res)

                        if is_solved:
                            solved_count += 1
                        if is_executable:
                            executable_count += 1

            # Report
            total = len(prob_files)
            accuracy = solved_count / total if total else 0
            exec_rate = executable_count / total if total else 0
            avg_plan_len = (
                sum(r.get("plan_len", 0) for r in results if "plan_len" in r) / total
                if total
                else 0.0
            )

            print(
                f"Result {split}: Solved {solved_count}/{total} ({accuracy:.2%}) | Executable {executable_count}/{total} ({exec_rate:.2%})"
            )

            # Save
            os.makedirs(args.results_dir, exist_ok=True)
            tag_suffix = f"_{args.tag}" if getattr(args, "tag", "") else ""
            out_file = os.path.join(
                args.results_dir,
                f"{args.domain}_{args.encoding}_{split}{tag_suffix}_results.json",
            )
            with open(out_file, "w") as f:
                json.dump(results, f, indent=2)
            print(f"Saved results to {out_file}")

            all_split_solved.append(accuracy)
            all_split_exec.append(exec_rate)
    finally:
        pass


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--domain", required=True)
    parser.add_argument("--checkpoint", required=True)
    # Replaced choices with free text to allow all tokenizers
    parser.add_argument("--encoding", required=True)
    parser.add_argument("--results_dir", required=True)
    parser.add_argument("--pddl_dir", default="data/pddl")
    parser.add_argument("--data_dir", default="data")
    parser.add_argument(
        "--tokenizer_manifest",
        default=None,
        help="Optional explicit tokenizer manifest for pooled/all-domain runs",
    )
    parser.add_argument(
        "--device",
        choices=["auto", "cuda", "mps", "cpu"],
        default="auto",
        help="Device selection policy",
    )
    parser.add_argument("--hidden_dim", type=int, default=256)
    parser.add_argument("--max_steps", type=int, default=100)
    parser.add_argument(
        "--steps_per_object",
        type=int,
        default=10,
        help="Minimum search budget scales to max(max_steps, steps_per_object * num_objects).",
    )
    parser.add_argument("--beam_width", type=int, default=3, help="Search beam width")
    parser.add_argument(
        "--score_metric",
        choices=["native", "cosine", "l2"],
        default="native",
        help="Successor scoring metric. 'native' reproduces the original study setting.",
    )
    parser.add_argument(
        "--amp",
        dest="amp",
        action="store_true",
        help="Enable automatic mixed precision for CUDA",
    )
    parser.add_argument(
        "--no_amp",
        dest="amp",
        action="store_false",
        help="Disable automatic mixed precision",
    )
    parser.add_argument(
        "--fast",
        action="store_true",
        help="Enable fast CUDA settings (less deterministic, more throughput)",
    )
    parser.add_argument(
        "--splits",
        nargs="+",
        default=DEFAULT_SPLITS,
        help="Splits to run (default: validation/test-interpolation/test-extrapolation)",
    )
    parser.add_argument(
        "--problems",
        nargs="+",
        default=None,
        help="Optional explicit problem file names to run (e.g., probBLOCKS-8-0.pddl)",
    )
    parser.add_argument(
        "--max_problems",
        type=int,
        default=None,
        help="Optional cap on number of problems per split after filtering",
    )
    parser.add_argument(
        "--skip_validation",
        action="store_true",
        help="Skip VAL-based validation and use internal search goal check for solved status",
    )
    parser.add_argument(
        "--validation_workers",
        type=int,
        default=8,
        help="Number of background VAL validation workers to overlap with search.",
    )
    parser.add_argument(
        "--collect_search_stats",
        action="store_true",
        help="Record search-effort and termination statistics in the output JSON.",
    )
    parser.add_argument(
        "--delta",
        action="store_true",
        help="Flag to whether perform delta-based preds. Def. is False",
    )
    parser.add_argument(
        "--tag",
        default="state",
        help="Optional tag to disambiguate results, e.g., 'state' or 'delta'",
    )
    parser.add_argument(
        "--no_projection",
        action="store_true",
        help="If set, disables the input projection layer (must match training)",
    )
    parser.add_argument("--seed", type=int, default=13, help="Random seed")
    HOME = os.path.expanduser("~")
    ROOT_DIR = f"{HOME}/planning/"
    parser.add_argument(
        "--val_path",
        default=os.environ.get("VAL_PATH", f"{ROOT_DIR}VAL/bin/Validate"),
        help="Path to VAL binary",
    )
    parser.set_defaults(amp=True)

    args = parser.parse_args()

    run_inference(args)