experiments/edge_weight_ablation.py

"""
Edge Weight Ablation Study for QAFD-RAG
========================================

Runs all experiments from a single script to avoid reloading the KG each time.

Experiments:
1. Query-aware vs query-agnostic
2. Weight scheme comparison (original/multiply/add)
3. (a, b) sensitivity sweep for Hybrid
4. Qualitative per-query diagnostics
"""

import json
import logging
import os
import sys
import time
from typing import Dict, List, Tuple

import numpy as np

# ---------------------------------------------------------------------------
# Setup paths
# ---------------------------------------------------------------------------
_project_root = os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))
sys.path.insert(0, _project_root)

import types as _types
for _pkg_path in ["src", "src.retrievers", "src.passage_entity"]:
    if _pkg_path not in sys.modules:
        _m = _types.ModuleType(_pkg_path)
        _m.__path__ = [os.path.join(_project_root, *_pkg_path.split("."))]
        _m.__package__ = _pkg_path
        sys.modules[_pkg_path] = _m

import importlib.util as _ilu
def _load_mod(fqn, filepath):
    spec = _ilu.spec_from_file_location(fqn, filepath)
    mod = _ilu.module_from_spec(spec)
    sys.modules[fqn] = mod
    spec.loader.exec_module(mod)
    return mod

_src = os.path.join(_project_root, "src")
_load_mod("src.retrievers.base", os.path.join(_src, "retrievers", "base.py"))
_load_mod("src.retrievers.flow_diffusion", os.path.join(_src, "retrievers", "flow_diffusion.py"))

from src.passage_entity.config import PassageEntityConfig
from src.passage_entity.embedding_store import EmbeddingModelWrapper
from src.passage_entity.kg_builder import KGBuilder
from src.passage_entity.openie import OpenIE
from src.passage_entity.reranker import FactReranker
from src.passage_entity.retriever import PassageEntityRetriever
from src.passage_entity.graph_adapter import run_igraph_qafd
from src.passage_entity.benchmark_runner import (
    get_gold_docs, get_gold_answers, recall_at_k,
    exact_match, f1_score, run_qa, _openai_embed, _openai_complete,
)

logging.basicConfig(level=logging.WARNING)
logger = logging.getLogger(__name__)

# ---------------------------------------------------------------------------
# Load KG once
# ---------------------------------------------------------------------------

def load_everything(dataset="musique", num_queries=10):
    """Load KG, embeddings, data — once for all experiments."""
    api_key = os.environ.get("OPENAI_API_KEY", "")

    config = PassageEntityConfig(
        llm_model="gpt-4o-mini",
        embedding_model_key="openai-small",
        dataset=dataset,
        save_dir="outputs",
    )

    import asyncio

    async def llm_func(prompt, system_prompt=None, history_messages=[], **kwargs):
        return await _openai_complete(
            model="gpt-4o-mini", prompt=prompt,
            system_prompt=system_prompt, history_messages=history_messages,
            api_key=api_key, **kwargs
        )

    async def embed_func(texts):
        return await _openai_embed(texts, model="text-embedding-3-small", api_key=api_key)

    embedding_model = EmbeddingModelWrapper(embed_func, batch_size=16)
    openie = OpenIE(llm_func)
    builder = KGBuilder(config, embedding_model, openie)

    # Load data
    data_dir = os.path.join(_project_root, "data", "multihop")
    with open(os.path.join(data_dir, f"{dataset}_corpus.json")) as f:
        corpus = json.load(f)
    docs = [f"{d['title']}\n{d['text']}" for d in corpus]

    with open(os.path.join(data_dir, f"{dataset}.json")) as f:
        samples = json.load(f)[:num_queries]

    queries = [s["question"] for s in samples]
    gold_answers = get_gold_answers(samples)
    gold_docs = get_gold_docs(samples, dataset)

    # Build / load KG
    builder.index(docs)

    # Build retriever
    reranker = FactReranker(llm_func)
    retriever = PassageEntityRetriever(
        config=config,
        embedding_model=embedding_model,
        reranker=reranker,
        graph=builder.graph,
        chunk_embedding_store=builder.chunk_embedding_store,
        entity_embedding_store=builder.entity_embedding_store,
        fact_embedding_store=builder.fact_embedding_store,
        openie_results_path=builder.openie_results_path,
    )
    retriever.prepare()

    return config, retriever, queries, gold_answers, gold_docs, llm_func


def run_retrieval_with_params(
    retriever: PassageEntityRetriever,
    queries: List[str],
    gold_docs: List[List[str]],
    gold_answers,
    weight_scheme: str = "original",
    hybrid_a: float = 1.0,
    hybrid_b: float = 0.5,
    query_aware: bool = True,
) -> Dict:
    """Run retrieval with specific edge weight params. Returns metrics dict."""
    from src.passage_entity.utils import compute_mdhash_id, min_max_normalize

    retriever._encode_queries(queries)
    k_list = [1, 2, 5, 10, 20, 50, 100, 200]
    all_retrieved = []
    qafd_times = []
    convergence_iters = []

    for q in queries:
        fact_scores = retriever._get_fact_scores(q)
        top_indices, top_facts, _ = retriever._rerank_facts(q, fact_scores)

        if len(top_facts) == 0:
            sorted_ids, sorted_scores = retriever._dense_passage_retrieval(q)
            all_retrieved.append([
                retriever.chunk_store.get_row(retriever.passage_node_keys[idx])["content"]
                for idx in sorted_ids[:200]
            ])
            qafd_times.append(0.0)
            convergence_iters.append(-1)  # DPR fallback
            continue

        # Compute seed weights (same as retriever._graph_search)
        n_nodes = retriever.graph.vcount()
        phrase_weights = np.zeros(n_nodes)
        passage_weights = np.zeros(n_nodes)
        number_of_occurs = np.zeros(n_nodes)

        for rank, f in enumerate(top_facts):
            subj, obj = f[0].lower(), f[2].lower()
            fs = fact_scores[top_indices[rank]] if fact_scores.ndim > 0 else float(fact_scores)
            for phrase in [subj, obj]:
                pk = compute_mdhash_id(phrase, prefix="entity-")
                pid = retriever.node_name_to_vertex_idx.get(pk)
                if pid is not None:
                    wfs = fs
                    num_chunks = len(retriever.ent_node_to_chunk_ids.get(pk, set()))
                    if num_chunks > 0:
                        wfs /= num_chunks
                    phrase_weights[pid] += wfs
                    number_of_occurs[pid] += 1

        nonzero = number_of_occurs > 0
        phrase_weights[nonzero] /= number_of_occurs[nonzero]

        dpr_ids, dpr_scores = retriever._dense_passage_retrieval(q)
        norm_dpr = min_max_normalize(dpr_scores)
        pw = retriever.config.passage_node_weight
        for i, did in enumerate(dpr_ids.tolist()):
            pk = retriever.passage_node_keys[did]
            pid = retriever.node_name_to_vertex_idx.get(pk)
            if pid is not None:
                passage_weights[pid] = norm_dpr[i] * pw

        node_weights = phrase_weights + passage_weights
        if np.sum(node_weights) == 0:
            sorted_ids, sorted_scores = dpr_ids, dpr_scores
            all_retrieved.append([
                retriever.chunk_store.get_row(retriever.passage_node_keys[idx])["content"]
                for idx in sorted_ids[:200]
            ])
            qafd_times.append(0.0)
            convergence_iters.append(-1)
            continue

        # Build node embeddings
        if not hasattr(retriever, '_node_emb_dict') or retriever._node_emb_dict is None:
            retriever._node_emb_dict = {}
            for i, nk in enumerate(retriever.entity_node_keys):
                if i < len(retriever.entity_embeddings):
                    retriever._node_emb_dict[nk] = retriever.entity_embeddings[i]
            for i, nk in enumerate(retriever.passage_node_keys):
                if i < len(retriever.passage_embeddings):
                    retriever._node_emb_dict[nk] = retriever.passage_embeddings[i]

        query_emb = retriever._query_emb_fact.get(q) if query_aware else None

        t0 = time.time()
        sorted_ids, sorted_scores = run_igraph_qafd(
            graph=retriever.graph,
            node_name_to_idx=retriever.node_name_to_vertex_idx,
            passage_node_idxs=retriever.passage_node_idxs,
            source_weights=node_weights,
            node_embeddings=retriever._node_emb_dict if query_aware else {},
            query_embedding=query_emb,
            alpha=retriever.config.qafd_alpha,
            epsilon=retriever.config.qafd_epsilon,
            max_iterations=retriever.config.qafd_max_iterations,
            step_size=retriever.config.qafd_step_size,
            weight_scheme=weight_scheme,
            hybrid_a=hybrid_a,
            hybrid_b=hybrid_b,
            use_node_degree=retriever.config.qafd_use_node_degree,
            random_seed=retriever.config.qafd_random_seed,
        )
        elapsed = time.time() - t0
        qafd_times.append(elapsed)

        top_docs = [
            retriever.chunk_store.get_row(retriever.passage_node_keys[idx])["content"]
            for idx in sorted_ids[:200]
        ]
        all_retrieved.append(top_docs)

    # Compute metrics
    recall_metrics = recall_at_k(gold_docs, all_retrieved, k_list) if gold_docs else {}

    return {
        "recall": recall_metrics,
        "avg_qafd_time": np.mean([t for t in qafd_times if t > 0]) if any(t > 0 for t in qafd_times) else 0,
        "dpr_fallback_count": sum(1 for t in convergence_iters if t == -1),
        "qafd_queries": sum(1 for t in qafd_times if t > 0),
    }


# ===========================================================================
# Main experiments
# ===========================================================================

def main():
    import argparse
    parser = argparse.ArgumentParser()
    parser.add_argument("--dataset", default="musique")
    parser.add_argument("--num_queries", type=int, default=10)
    args = parser.parse_args()

    print("=" * 70)
    print("  QAFD-RAG Edge Weight Ablation Study")
    print(f"  Dataset: {args.dataset}, Queries: {args.num_queries}")
    print("=" * 70)

    config, retriever, queries, gold_answers, gold_docs, llm_func = load_everything(
        args.dataset, args.num_queries
    )

    results = {}

    # ── Experiment 1: Query-aware vs Query-agnostic ──────────────────
    print("\n[1/4] Query-aware vs Query-agnostic ablation")

    print("  Running: Hybrid (a=1, b=0.5) — query-aware ...")
    r = run_retrieval_with_params(
        retriever, queries, gold_docs, gold_answers,
        weight_scheme="original", hybrid_a=1.0, hybrid_b=0.5, query_aware=True,
    )
    results["hybrid_query_aware"] = r
    print(f"    R@10={r['recall'].get('Recall@10', 0):.4f}  R@100={r['recall'].get('Recall@100', 0):.4f}  QAFD={r['avg_qafd_time']:.3f}s  DPR_fallback={r['dpr_fallback_count']}")

    print("  Running: Query-agnostic (b=0) ...")
    r = run_retrieval_with_params(
        retriever, queries, gold_docs, gold_answers,
        weight_scheme="original", hybrid_a=1.0, hybrid_b=0.0, query_aware=False,
    )
    results["query_agnostic"] = r
    print(f"    R@10={r['recall'].get('Recall@10', 0):.4f}  R@100={r['recall'].get('Recall@100', 0):.4f}  QAFD={r['avg_qafd_time']:.3f}s  DPR_fallback={r['dpr_fallback_count']}")

    # ── Experiment 2: Weight scheme comparison ───────────────────────
    print("\n[2/4] Weight scheme comparison")

    for scheme in ["original", "multiply", "add"]:
        print(f"  Running: {scheme} ...")
        r = run_retrieval_with_params(
            retriever, queries, gold_docs, gold_answers,
            weight_scheme=scheme, query_aware=True,
        )
        results[f"scheme_{scheme}"] = r
        print(f"    R@10={r['recall'].get('Recall@10', 0):.4f}  R@100={r['recall'].get('Recall@100', 0):.4f}  QAFD={r['avg_qafd_time']:.3f}s")

    # ── Experiment 3: (a, b) sensitivity sweep ───────────────────────
    print("\n[3/4] (a, b) sensitivity sweep for Hybrid")

    a_values = [0.5, 1.0, 2.0]
    b_values = [0.0, 0.1, 0.25, 0.5, 1.0]

    sweep_results = {}
    for a in a_values:
        for b in b_values:
            label = f"a={a},b={b}"
            r = run_retrieval_with_params(
                retriever, queries, gold_docs, gold_answers,
                weight_scheme="original", hybrid_a=a, hybrid_b=b, query_aware=(b > 0),
            )
            sweep_results[label] = r
            r10 = r['recall'].get('Recall@10', 0)
            r100 = r['recall'].get('Recall@100', 0)
            print(f"  {label:>15}  R@10={r10:.4f}  R@100={r100:.4f}  QAFD={r['avg_qafd_time']:.3f}s")

    results["sweep"] = sweep_results

    # ── Experiment 4: Per-query diagnostics ──────────────────────────
    print("\n[4/4] Per-query diagnostics (first 5 queries)")

    # Compare query-aware vs agnostic per query
    diag_queries = queries[:5]
    diag_results = []

    for qi, q in enumerate(diag_queries):
        # Query-aware
        r_aware = run_retrieval_with_params(
            retriever, [q], [gold_docs[qi]], [gold_answers[qi]],
            weight_scheme="original", hybrid_a=1.0, hybrid_b=0.5, query_aware=True,
        )
        # Query-agnostic
        r_agnostic = run_retrieval_with_params(
            retriever, [q], [gold_docs[qi]], [gold_answers[qi]],
            weight_scheme="original", hybrid_a=1.0, hybrid_b=0.0, query_aware=False,
        )

        r10_aware = r_aware['recall'].get('Recall@10', 0)
        r10_agnostic = r_agnostic['recall'].get('Recall@10', 0)
        delta = r10_aware - r10_agnostic

        status = "HELPS" if delta > 0 else ("HURTS" if delta < 0 else "SAME")
        print(f"  Q{qi}: R@10 aware={r10_aware:.3f} agnostic={r10_agnostic:.3f}  delta={delta:+.3f} [{status}]")
        print(f"       Q: {q[:80]}...")

        diag_results.append({
            "query": q,
            "r10_aware": r10_aware,
            "r10_agnostic": r10_agnostic,
            "delta": delta,
            "status": status,
        })

    results["diagnostics"] = diag_results

    # ── Summary ──────────────────────────────────────────────────────
    print("\n" + "=" * 70)
    print("  SUMMARY")
    print("=" * 70)

    print("\n  Ablation: Query-aware vs Query-agnostic")
    print(f"    Query-aware (Hybrid a=1,b=0.5):  R@10={results['hybrid_query_aware']['recall'].get('Recall@10',0):.4f}  R@100={results['hybrid_query_aware']['recall'].get('Recall@100',0):.4f}")
    print(f"    Query-agnostic (b=0):             R@10={results['query_agnostic']['recall'].get('Recall@10',0):.4f}  R@100={results['query_agnostic']['recall'].get('Recall@100',0):.4f}")

    print("\n  Weight scheme comparison")
    for scheme in ["original", "multiply", "add"]:
        r = results[f"scheme_{scheme}"]
        print(f"    {scheme:>10}:  R@10={r['recall'].get('Recall@10',0):.4f}  R@100={r['recall'].get('Recall@100',0):.4f}")

    print("\n  Best (a,b) from sweep:")
    best_label = max(sweep_results, key=lambda k: sweep_results[k]['recall'].get('Recall@10', 0))
    best = sweep_results[best_label]
    print(f"    {best_label}:  R@10={best['recall'].get('Recall@10',0):.4f}  R@100={best['recall'].get('Recall@100',0):.4f}")

    print("\n  Per-query diagnostics:")
    helps = sum(1 for d in diag_results if d["status"] == "HELPS")
    hurts = sum(1 for d in diag_results if d["status"] == "HURTS")
    same = sum(1 for d in diag_results if d["status"] == "SAME")
    print(f"    Query-awareness HELPS: {helps}/{len(diag_results)}, HURTS: {hurts}/{len(diag_results)}, SAME: {same}/{len(diag_results)}")

    # Save
    out_dir = os.path.join(_project_root, "experiments", "results")
    os.makedirs(out_dir, exist_ok=True)
    out_path = os.path.join(out_dir, f"edge_weight_ablation_{args.dataset}.json")
    with open(out_path, "w") as f:
        json.dump(results, f, indent=2, default=str)
    print(f"\n  Results saved to {out_path}")
    print("=" * 70)


if __name__ == "__main__":
    main()