experiment/knowledge_editing/run_baselines.py

"""
Run knowledge-editing baselines on the hallucination suppression task.

Methods supported:
  - lora:     LoRA fine-tuning via EasyEdit
  - dualedit: DualEdit (vision + text adapters, custom implementation)

Usage:
    python -m experiment.knowledge_editing.run_baselines \
        --edit_set experiment/knowledge_editing/edit_set.json \
        --methods lora dualedit \
        --output_dir step4_ke_outputs
"""

import argparse
import json
import os
import sys
import copy
from datetime import datetime
from pathlib import Path

import torch
from PIL import Image
from tqdm import tqdm

# Ensure project root is on path
sys.path.insert(0, str(Path(__file__).resolve().parents[2]))

from experiment.knowledge_editing.llava15_compat import (
    LLaVA15ProcessorWrapper,
    LLaVA15ImageProcessor,
)


def load_edit_set(path: str) -> dict:
    with open(path) as f:
        return json.load(f)


def _load_image(inst: dict, hf_images: dict | None) -> Image.Image | None:
    """Load an image from file path or HF dataset cache."""
    image_path = inst.get("image_path")
    image_id = inst.get("image_id")

    if image_path and os.path.isfile(image_path):
        return Image.open(image_path).convert("RGB")

    if hf_images and image_id in hf_images:
        return hf_images[image_id].convert("RGB")

    return None


def _build_hf_image_cache(dataset_id: str) -> dict:
    """Build an {image_id: PIL.Image} lookup from the HF dataset."""
    from experiment.data.hf_loader import load_hf_dataset

    print(f"  Loading images from HuggingFace dataset ({dataset_id})...")
    ds = load_hf_dataset(dataset_id)
    if hasattr(ds, "keys"):
        from datasets import concatenate_datasets
        ds = concatenate_datasets([ds[s] for s in ds])

    return {item["image_id"]: item["image"] for item in ds}


def build_requests(edit_set: dict, dataset_id: str = "pbcong/bathroom-toilet",
                   use_eval_instances: bool = False):
    """Convert edit_set instances into request format.

    Args:
        use_eval_instances: If True, use eval_instances for the efficacy category
            (HF val split, the fixed 50-image eval set). Used by DualEdit so
            editing and evaluation are on the same images. If False, use
            edit_instances.train (HF train split) for methods like LoRA.
    """
    # Resolve efficacy category name from edit_descriptor
    relation_key = edit_set.get("edit_descriptor", {}).get("relation", "bathroom_toilet")
    efficacy_cat = edit_set.get("edit_descriptor", {}).get("concept", "bathroom_no_toilet")

    if use_eval_instances and "eval_instances" in edit_set:
        instances = edit_set["eval_instances"].get(efficacy_cat, [])
        print(f"  Using eval_instances.{efficacy_cat} ({len(instances)} images, HF val split)")
    else:
        instances = edit_set["edit_instances"]["train"]
        print(f"  Using edit_instances.train ({len(instances)} images, HF train split)")

    locality = edit_set["locality_instances"]

    edit_prompt = edit_set["prompts"]["edit_prompt"]

    generality_prompts = edit_set["prompts"]["generality_prompts"]
    rephrase_prompt = generality_prompts[0] if generality_prompts else edit_prompt

    # Use first available locality category (scene_with_object)
    loc_cat_name = next(iter(locality), None)
    loc_bwt = locality[loc_cat_name] if loc_cat_name else []

    # Filter for usable instances first, then apply n_edits cap so we don't
    # waste the budget on instances that have no target (non-hallucinating images).
    usable_all = [
        inst for inst in instances
        if inst.get("target") is not None and inst.get("is_usable", True)
    ]
    skipped = len(instances) - len(usable_all)
    if skipped:
        print(f"  {skipped} instances skipped (no target or degenerate after cleaning)")
    usable = usable_all

    hf_images = None
    all_instances = list(usable) + list(loc_bwt)
    needs_hf = any(
        not inst.get("image_path") or not os.path.isfile(inst.get("image_path", ""))
        for inst in all_instances
    )
    if needs_hf:
        hf_images = _build_hf_image_cache(dataset_id)

    requests = []
    for i, inst in enumerate(usable):
        edit_image = _load_image(inst, hf_images)
        if edit_image is None:
            print(f"  Skipping {inst['image_id']}: image not found")
            continue

        rephrase_idx = (i + 1) % len(usable)
        rephrase_inst = usable[rephrase_idx]
        rephrase_image = _load_image(rephrase_inst, hf_images) or edit_image

        text_loc_prompt = "What is the capital of France?"
        text_loc_answer = "Paris"

        loc_inst = loc_bwt[i % len(loc_bwt)]
        loc_image = _load_image(loc_inst, hf_images) or edit_image
        loc_vis_prompt = edit_prompt
        loc_vis_answer = loc_inst.get("original_caption") or "A room with various objects."

        request = {
            "prompt": edit_prompt,
            "target": inst["target"],
            "image": edit_image,
            "file_type": "image",
            "rephrase_prompt": rephrase_prompt,
            "image_rephrase": rephrase_image,
            "locality_prompt": text_loc_prompt,
            "locality_ground_truth": text_loc_answer,
            "multimodal_locality_image": loc_image,
            "multimodal_locality_prompt": loc_vis_prompt,
            "multimodal_locality_ground_truth": loc_vis_answer,
            "_image_id": inst["image_id"],
        }
        requests.append(request)

    print(f"Built {len(requests)} edit requests (completion formulation)")
    return requests


def load_model_and_processor(model_name: str, device: str, dtype: torch.dtype):
    """Load LLaVA-1.5 model and wrap processor."""
    from transformers import AutoTokenizer, LlavaForConditionalGeneration, LlavaProcessor
    from transformers import CLIPImageProcessor

    print(f"Loading {model_name}...")
    model = LlavaForConditionalGeneration.from_pretrained(
        model_name, torch_dtype=dtype, device_map={"": device},
    )
    tokenizer = AutoTokenizer.from_pretrained(model_name, use_fast=False)
    image_processor = CLIPImageProcessor.from_pretrained(model_name)
    raw_processor = LlavaProcessor(tokenizer=tokenizer, image_processor=image_processor)
    processor = LLaVA15ProcessorWrapper(raw_processor)

    return model, processor


def load_hparams(method: str, hparams_dir: str):
    """Load hparams for a given method."""
    yaml_path = os.path.join(hparams_dir, f"{method}.yaml")
    if not os.path.exists(yaml_path):
        raise FileNotFoundError(f"Hparams not found: {yaml_path}")

    if method == "lora":
        from easyeditor.models.lora import LoRAMultimodalHyperParams
        return LoRAMultimodalHyperParams.from_hparams(yaml_path)
    elif method == "dualedit":
        from experiment.knowledge_editing.dualedit.dualedit_hparams import DualEditHyperParams
        return DualEditHyperParams.from_hparams(yaml_path)
    else:
        raise ValueError(f"Unknown method: {method}")


def get_apply_algo(method: str):
    """Get the algorithm function for a method."""
    if method == "lora":
        from easyeditor.models.lora.lora_main import apply_lora_to_multimodal_model
        return apply_lora_to_multimodal_model
    elif method == "dualedit":
        from experiment.knowledge_editing.dualedit import apply_dualedit_to_multimodal_model
        return apply_dualedit_to_multimodal_model
    else:
        raise ValueError(f"Unknown method: {method}")


def save_edited_model(model, processor, output_dir: str, method: str):
    """Save the edited model for later evaluation."""
    save_dir = os.path.join(output_dir, f"{method}_edited")
    merged_dir = os.path.join(save_dir, "merged_for_eval")
    os.makedirs(merged_dir, exist_ok=True)

    from transformers import PreTrainedModel
    base_model = model
    while not isinstance(base_model, PreTrainedModel) and hasattr(base_model, "model"):
        base_model = base_model.model

    if method == "dualedit":
        if hasattr(base_model, '_dualedit_state'):
            dualedit_state = base_model._dualedit_state
            save_state = {k: v for k, v in dualedit_state.items()
                          if k not in ("vision_hook_handle", "text_hook_handle")}
            adapter_path = os.path.join(merged_dir, "dualedit_state.pt")
            torch.save(save_state, adapter_path)
            print(f"  Saved DualEdit adapter state to {adapter_path}")
        else:
            print("  WARNING: DualEdit state not found on model — nothing saved")
    else:
        # LoRA: edits are merged into the weights
        base_model.save_pretrained(merged_dir)
        processor._processor.save_pretrained(merged_dir)

    print(f"  Saved edited model to {merged_dir}")
    return save_dir


def run_single_method(
    method: str,
    model,
    processor,
    requests: list[dict],
    hparams_dir: str,
    output_dir: str,
):
    """Run one method on the edit requests."""
    print(f"\n{'='*60}")
    print(f"Running: {method.upper()}")
    print(f"  {len(requests)} edit instances")
    print(f"{'='*60}")

    hparams = load_hparams(method, hparams_dir)
    apply_algo = get_apply_algo(method)

    edited_model = model

    try:
        checkpoint_dir = os.path.join(output_dir, "checkpoints")
        edited_model, weights_copy = apply_algo(
            model,
            processor,
            requests,
            hparams,
            copy=False,
            return_orig_weights=True,
            keep_original_weight=False,
            checkpoint_dir=checkpoint_dir,
        )
    except Exception as e:
        print(f"  {method} apply failed: {e}")
        import traceback
        traceback.print_exc()

    save_dir = save_edited_model(edited_model, processor, output_dir, method)

    return edited_model, save_dir


def run_evaluation(
    model_type: str,
    model_dir: str,
    base_model_name: str,
    output_dir: str,
    run_name: str,
    edit_set_path: str,
):
    """Run evaluation pipeline on the edited model."""
    import subprocess

    cmd = [
        sys.executable, "-m", "experiment.evaluation.validate",
        "--model_type", model_type,
        "--model_dir", model_dir,
        "--base_model_name", base_model_name,
        "--inference_backend", "transformers",
        "--mention_method", "keyword",
        "--output_dir", output_dir,
        "--num_per_category", "50",
    ]

    print(f"\n  Running evaluation: {run_name}")
    print(f"  Command: {' '.join(cmd)}")

    result = subprocess.run(cmd, capture_output=True, text=True)
    if result.returncode != 0:
        print(f"  Evaluation failed:\n{result.stderr}")
    else:
        print(f"  Evaluation complete")
        for line in result.stdout.split("\n"):
            if any(k in line for k in ["efficacy", "generality", "locality",
                                        "consistency", "Efficacy", "Generality",
                                        "Locality", "Consistency"]):
                print(f"    {line}")

    return result.returncode == 0


def main():
    parser = argparse.ArgumentParser(
        description="Run LoRA and DualEdit baselines for hallucination suppression"
    )
    parser.add_argument("--edit_set", type=str,
                        default="experiment/knowledge_editing/edit_set.json")
    parser.add_argument("--methods", nargs="+",
                        default=["lora", "dualedit"],
                        choices=["lora", "dualedit"])
    parser.add_argument("--model_name", type=str,
                        default="llava-hf/llava-1.5-7b-hf")
    parser.add_argument("--hparams_dir", type=str,
                        default=os.path.join(os.path.dirname(os.path.abspath(__file__)), "hparams"))
    parser.add_argument("--output_dir", type=str,
                        default="step4_ke_outputs")
    parser.add_argument("--dataset_id", type=str, default=None,
                        help="HuggingFace dataset ID for loading images")
    parser.add_argument("--device", type=str, default="cuda")
    parser.add_argument("--batch", action="store_true",
                        help="Ignored (kept for backwards compatibility)")
    parser.add_argument("--skip_eval", action="store_true",
                        help="Skip evaluation (just run edits and save)")
    args = parser.parse_args()

    run_id = datetime.now().strftime("%Y%m%d_%H%M%S")
    run_dir = os.path.join(args.output_dir, f"ke_run_{run_id}")
    os.makedirs(run_dir, exist_ok=True)

    with open(os.path.join(run_dir, "run_config.json"), "w") as f:
        json.dump(vars(args), f, indent=2)

    print("Loading edit set...")
    edit_set = load_edit_set(args.edit_set)
    print(f"  Stats: {edit_set['stats']}")

    # Resolve dataset_id: CLI > edit_set > default
    dataset_id = args.dataset_id
    if not dataset_id:
        dataset_id = edit_set.get("data_config", {}).get("dataset_id", "pbcong/bathroom-toilet")

    # DualEdit edits on the eval set (HF val, 50 images) so that editing and
    # evaluation use the exact same images. Other methods (LoRA) train on the
    # full HF train split and are evaluated on the separate eval set.
    use_eval = set(args.methods) == {"dualedit"}
    requests = build_requests(edit_set, dataset_id=dataset_id,
                              use_eval_instances=use_eval)
    if not requests:
        print("ERROR: No valid requests built. Check edit_set.json.")
        return

    # Save the exact image IDs used for editing so eval can pin to the same images.
    edit_image_ids = [r["_image_id"] for r in requests]
    edit_image_ids_path = os.path.join(run_dir, "edit_image_ids.json")
    with open(edit_image_ids_path, "w") as f:
        json.dump(edit_image_ids, f)
    print(f"  Saved {len(edit_image_ids)} edit image IDs to {edit_image_ids_path}")

    results_summary = {}

    for method in args.methods:
        print(f"\nLoading fresh model for {method}...")
        model, processor = load_model_and_processor(
            args.model_name, args.device, torch.float16,
        )

        try:
            edited_model, save_dir = run_single_method(
                method=method,
                model=model,
                processor=processor,
                requests=requests,
                hparams_dir=args.hparams_dir,
                output_dir=run_dir,
            )

            results_summary[method] = {
                "status": "edited",
                "save_dir": save_dir,
                "n_edits": len(requests),
            }

            if not args.skip_eval:
                merged_dir = os.path.join(save_dir, "merged_for_eval")
                eval_model_type = "dualedit" if method == "dualedit" else "merged"
                success = run_evaluation(
                    model_type=eval_model_type,
                    model_dir=merged_dir,
                    base_model_name=args.model_name,
                    output_dir=run_dir,
                    run_name=f"{method}_n{len(requests)}",
                    edit_set_path=args.edit_set,
                )
                results_summary[method]["eval_success"] = success

        except Exception as e:
            print(f"  {method} FAILED: {e}")
            import traceback
            traceback.print_exc()
            results_summary[method] = {"status": "failed", "error": str(e)}

    summary_path = os.path.join(run_dir, "results_summary.json")
    with open(summary_path, "w") as f:
        json.dump(results_summary, f, indent=2)

    print(f"\n{'='*60}")
    print("All methods complete.")
    print(f"Results saved to: {run_dir}")
    print(f"Summary: {summary_path}")
    for method, result in results_summary.items():
        print(f"  {method}: {result['status']}")
    print(f"{'='*60}")


if __name__ == "__main__":
    main()