#!/usr/bin/env python3
"""
EvoRM Comprehensive Experiment Runner
=======================================
Covers all 5 Research Questions from the TKDE paper:
  RQ1: Main Results (different datasets)
  RQ2: Cross-Model Analysis
  RQ3: Ablation Study
  RQ4: Efficiency Analysis
  RQ5: Cold-Start Scaling

Usage:
  python run_all_experiments.py --all          # Run everything
  python run_all_experiments.py --ablation     # RQ3 only
  python run_all_experiments.py --coldstart    # RQ5 only
  python run_all_experiments.py --main         # RQ1 only
  python run_all_experiments.py --efficiency   # RQ4 only
"""

import os
import sys
import json
import time
import shutil
import argparse
import subprocess
from datetime import datetime
from collections import defaultdict

# Add parent dir to path
sys.path.insert(0, '/root/autodl-tmp/AdaCoAgentEA')
sys.path.insert(0, '/root/autodl-tmp/AdaCoAgentEA/Area2')

os.environ["OPENAI_API_BASE"] = os.environ.get("OPENAI_BASE_URL", "https://api.openai.com/v1")
# OPENAI_API_KEY set via environment variable

RESULTS_DIR = "/root/autodl-tmp/AdaCoAgentEA/results"
DATA_BASE = "/root/autodl-tmp/AdaCoAgentEA/data"

# Dataset configurations
DATASETS = {
    'icews_wiki': {
        'path': f'{DATA_BASE}/icews_wiki',
        'ref_pairs': f'{DATA_BASE}/icews_wiki/ref_pairs',
        'n_pairs': 3540,
        'type': 'HHEA',
    },
    # ICEWS-YAGO and BETA require full pipeline (Area1→Area2→Area3)
    # They only have index/role_assignments, not raw KG triples
}

def ensure_dir(d):
    os.makedirs(d, exist_ok=True)

def load_ref_pairs(path):
    """Load reference entity pairs."""
    pairs = []
    with open(path, 'r') as f:
        for line in f:
            parts = line.strip().split('\t')
            if len(parts) == 2:
                pairs.append((int(parts[0]), int(parts[1])))
    return pairs

def compute_metrics(aligned_file, ref_pairs_file):
    """Compute Hits@1, MRR, Precision, Recall."""
    # Load aligned pairs
    aligned = set()
    if os.path.exists(aligned_file):
        with open(aligned_file, 'r') as f:
            for line in f:
                parts = line.strip().split('\t')
                if len(parts) == 2:
                    try:
                        aligned.add((int(parts[0]), int(parts[1])))
                    except ValueError:
                        continue

    # Load reference pairs
    ref = set()
    with open(ref_pairs_file, 'r') as f:
        for line in f:
            parts = line.strip().split('\t')
            if len(parts) == 2:
                try:
                    ref.add((int(parts[0]), int(parts[1])))
                except ValueError:
                    continue

    if not ref:
        return {'error': 'No reference pairs found'}

    # Compute metrics
    true_positives = aligned & ref
    false_positives = aligned - ref
    false_negatives = ref - aligned

    tp = len(true_positives)
    fp = len(false_positives)
    fn = len(false_negatives)

    precision = tp / max(1, tp + fp)
    recall = tp / max(1, tp + fn)
    hits1 = tp / len(ref)  # Hits@1 since we have 1-1 alignment
    mrr = hits1  # MRR = Hits@1 for 1-1 alignment

    return {
        'Hits@1': round(hits1, 4),
        'MRR': round(mrr, 4),
        'Precision': round(precision, 4),
        'Recall': round(recall, 4),
        'TP': tp, 'FP': fp, 'FN': fn,
        'Total_Ref': len(ref),
        'Total_Aligned': len(aligned),
    }

def run_llm1_alignment(data_dir, from_m3=False, ablation_config=None, 
                        no_optimization_tool=False, n_entities=None):
    """Run LLM1 label selector alignment."""
    from LLM1_label_selector import align_entities
    
    # Backup original important_entities if we're limiting entities
    important_file = os.path.join(data_dir, "message_pool", "important_entities.txt")
    
    if n_entities:
        # Create subset
        bak = important_file + ".bak_full"
        if not os.path.exists(bak):
            shutil.copy(important_file, bak)
        # Read first n_entities
        with open(bak, 'r') as f:
            lines = f.readlines()
        with open(important_file, 'w') as f:
            f.writelines(lines[:n_entities])
    
    try:
        start_time = time.time()
        aligned_pairs = align_entities(
            data_dir, 
            from_m3=from_m3,
            ablation_config=ablation_config,
            no_optimization_tool=no_optimization_tool
        )
        elapsed = time.time() - start_time
        
        # Get token count
        import tokens_cal
        tokens = tokens_cal.global_tokens
        tokens_cal.global_tokens = 0  # Reset
        
        return {
            'aligned_pairs': len(aligned_pairs),
            'time_seconds': elapsed,
            'tokens': tokens,
        }
    finally:
        # Restore
        if n_entities and os.path.exists(bak):
            shutil.copy(bak, important_file)

def run_ablation_study():
    """RQ3: Ablation Study - Full System vs w/o Stage-1 vs w/o Maintenance."""
    print("\n" + "="*70)
    print("RQ3: ABLATION STUDY")
    print("="*70)
    
    data_dir = f"{DATA_BASE}/icews_wiki"
    n_entities = 200  # Subset for efficiency
    aligned_file = os.path.join(data_dir, "message_pool", "aligned_entities.txt")
    ref_pairs = os.path.join(data_dir, "ref_pairs")
    
    # Set up important_entities subset
    important_file = os.path.join(data_dir, "message_pool", "important_entities.txt")
    bak = important_file + ".bak_full"
    if not os.path.exists(bak):
        shutil.copy(important_file, bak)
    
    # Read reference pairs to get entity IDs
    ref_set = load_ref_pairs(ref_pairs)
    kg1_entities = list(set(p[0] for p in ref_set))
    
    # Create important_entities with first n_entities
    with open(important_file, 'w') as f:
        for eid in kg1_entities[:n_entities]:
            f.write(f"{eid}\n")
    
    # Also update aligned_entities to only contain relevant ref pairs
    with open(important_file, 'r') as f:
        selected_entities = set(int(line.strip()) for line in f if line.strip())
    
    results = {}
    
    configs = [
        ('Full System (EvoRM)', None),
        ('w/o Stage-1 (Symbolic)', 'no_stage1'),
        ('w/o Maintenance', 'no_maintenance'),
    ]
    
    for config_name, ablation_mode in configs:
        print(f"\n{'='*50}")
        print(f"Config: {config_name}")
        print(f"{'='*50}")
        
        # Clear evorm state
        evorm_state = os.path.join(data_dir, "evorm_state")
        if os.path.exists(evorm_state):
            shutil.rmtree(evorm_state)
        
        # Clear aligned_entities
        if os.path.exists(aligned_file):
            os.remove(aligned_file)
        
        # Reset tokens
        import tokens_cal
        tokens_cal.global_tokens = 0
        
        # Run alignment with ablation config
        # We need to patch the module to accept ablation_mode
        import evorm_plugin
        
        # Monkey-patch: create a new EvoRMPlugin with ablation_mode
        original_init = evorm_plugin.EvoRMPlugin.__init__
        
        # Run using LLM1_label_selector directly
        from LLM1_label_selector import align_entities
        
        # We need to inject ablation_mode into the EvoRMPlugin initialization
        # The simplest way is to modify the module temporarily
        import LLM1_label_selector as llm1_module
        
        # Save original
        orig_align = llm1_module.align_entities
        
        # Create a wrapper
        def make_align_with_ablation(ablation):
            def wrapped_align(data_dir, from_m3=False, ablation_config=None, no_optimization_tool=False):
                # Override evorm_plugin to use ablation
                original_plugin = evorm_plugin.EvoRMPlugin
                class AblationPlugin(original_plugin):
                    def __init__(self, **kwargs):
                        kwargs['ablation_mode'] = ablation
                        super().__init__(**kwargs)
                evorm_plugin.EvoRMPlugin = AblationPlugin
                try:
                    return orig_align(data_dir, from_m3, ablation_config, no_optimization_tool)
                finally:
                    evorm_plugin.EvoRMPlugin = original_plugin
            return wrapped_align
        
        if ablation_mode:
            llm1_module.align_entities = make_align_with_ablation(ablation_mode)
        
        try:
            start_time = time.time()
            aligned_pairs = llm1_module.align_entities(data_dir, from_m3=False)
            elapsed = time.time() - start_time
            
            tokens = tokens_cal.global_tokens
            tokens_cal.global_tokens = 0
            
            # Compute metrics
            metrics = compute_metrics(aligned_file, ref_pairs)
            
            results[config_name] = {
                'aligned_pairs': len(aligned_pairs),
                'time_seconds': elapsed,
                'tokens': tokens,
                'metrics': metrics,
            }
            
            print(f"  Aligned: {len(aligned_pairs)}, Time: {elapsed:.1f}s, Tokens: {tokens}")
            print(f"  Hits@1: {metrics.get('Hits@1', 'N/A')}, MRR: {metrics.get('MRR', 'N/A')}")
            print(f"  Precision: {metrics.get('Precision', 'N/A')}, Recall: {metrics.get('Recall', 'N/A')}")
            
        finally:
            if ablation_mode:
                llm1_module.align_entities = orig_align
    
    # Restore full important_entities
    if os.path.exists(bak):
        shutil.copy(bak, important_file)
    
    # Save results
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    result_file = os.path.join(RESULTS_DIR, f"ablation_{timestamp}.json")
    ensure_dir(RESULTS_DIR)
    with open(result_file, 'w') as f:
        json.dump(results, f, indent=2)
    
    print(f"\nAblation results saved to: {result_file}")
    
    # Print summary table
    print("\n" + "="*70)
    print("ABLATION STUDY SUMMARY")
    print("="*70)
    print(f"{'Configuration':<30} {'Hits@1':<10} {'MRR':<10} {'Precision':<10} {'Recall':<10} {'Tokens':<10}")
    print("-"*70)
    for config_name in ['Full System (EvoRM)', 'w/o Stage-1 (Symbolic)', 'w/o Maintenance']:
        if config_name in results:
            r = results[config_name]
            m = r['metrics']
            print(f"{config_name:<30} {m.get('Hits@1', 0):<10.4f} {m.get('MRR', 0):<10.4f} {m.get('Precision', 0):<10.4f} {m.get('Recall', 0):<10.4f} {r['tokens']:<10}")
    
    return results

def run_coldstart_experiment():
    """RQ5: Cold-Start Scaling - different warmup sizes."""
    print("\n" + "="*70)
    print("RQ5: COLD-START SCALING EXPERIMENT")
    print("="*70)
    
    data_dir = f"{DATA_BASE}/icews_wiki"
    aligned_file = os.path.join(data_dir, "message_pool", "aligned_entities.txt")
    ref_pairs = os.path.join(data_dir, "ref_pairs")
    
    # Read all reference pairs for ground truth
    ref_set = load_ref_pairs(ref_pairs)
    kg1_entities = list(set(p[0] for p in ref_set))
    
    # Important entities file
    important_file = os.path.join(data_dir, "message_pool", "important_entities.txt")
    bak = important_file + ".bak_full"
    if not os.path.exists(bak):
        shutil.copy(important_file, bak)
    
    results = {}
    
    warmup_sizes = [100, 200, 500, 1000, 2000]
    
    for n_warmup in warmup_sizes:
        print(f"\n--- N_warmup = {n_warmup} ---")
        
        # Clear evorm state
        evorm_state = os.path.join(data_dir, "evorm_state")
        if os.path.exists(evorm_state):
            shutil.rmtree(evorm_state)
        
        # Clear aligned_entities
        if os.path.exists(aligned_file):
            os.remove(aligned_file)
        
        # Create subset
        with open(important_file, 'w') as f:
            for eid in kg1_entities[:n_warmup]:
                f.write(f"{eid}\n")
        
        import tokens_cal
        tokens_cal.global_tokens = 0
        
        from LLM1_label_selector import align_entities
        
        start_time = time.time()
        aligned_pairs = align_entities(data_dir, from_m3=False)
        elapsed = time.time() - start_time
        
        tokens = tokens_cal.global_tokens
        tokens_cal.global_tokens = 0
        
        metrics = compute_metrics(aligned_file, ref_pairs)
        
        results[f'n_warmup_{n_warmup}'] = {
            'aligned_pairs': len(aligned_pairs),
            'time_seconds': elapsed,
            'tokens': tokens,
            'tokens_per_pair': tokens / max(1, len(aligned_pairs)),
            'metrics': metrics,
        }
        
        print(f"  Aligned: {len(aligned_pairs)}, Tokens: {tokens}, "
              f"Tokens/Pair: {tokens/max(1,len(aligned_pairs)):.1f}")
        print(f"  Hits@1: {metrics.get('Hits@1', 'N/A')}")
    
    # Restore
    if os.path.exists(bak):
        shutil.copy(bak, important_file)
    
    # Save results
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    result_file = os.path.join(RESULTS_DIR, f"coldstart_{timestamp}.json")
    ensure_dir(RESULTS_DIR)
    with open(result_file, 'w') as f:
        json.dump(results, f, indent=2)
    
    print(f"\nCold-start results saved to: {result_file}")
    
    # Print summary
    print("\n" + "="*70)
    print("COLD-START SCALING SUMMARY")
    print("="*70)
    print(f"{'N_warmup':<12} {'Hits@1':<10} {'Tokens':<12} {'Tokens/Pair':<12} {'Time(s)':<10}")
    print("-"*56)
    for n_warmup in warmup_sizes:
        key = f'n_warmup_{n_warmup}'
        if key in results:
            r = results[key]
            m = r['metrics']
            print(f"{n_warmup:<12} {m.get('Hits@1', 0):<10.4f} {r['tokens']:<12} {r['tokens_per_pair']:<12.1f} {r['time_seconds']:<10.1f}")
    
    return results

def run_efficiency_analysis():
    """RQ4: Efficiency Analysis - token consumption and latency."""
    print("\n" + "="*70)
    print("RQ4: EFFICIENCY ANALYSIS")
    print("="*70)
    
    data_dir = f"{DATA_BASE}/icews_wiki"
    aligned_file = os.path.join(data_dir, "message_pool", "aligned_entities.txt")
    ref_pairs = os.path.join(data_dir, "ref_pairs")
    
    important_file = os.path.join(data_dir, "message_pool", "important_entities.txt")
    bak = important_file + ".bak_full"
    if not os.path.exists(bak):
        shutil.copy(important_file, bak)
    
    results = {}
    
    # Test with different entity counts
    entity_counts = [100, 200, 500]
    
    for n_entities in entity_counts:
        print(f"\n--- N_entities = {n_entities} ---")
        
        # Clear evorm state
        evorm_state = os.path.join(data_dir, "evorm_state")
        if os.path.exists(evorm_state):
            shutil.rmtree(evorm_state)
        
        if os.path.exists(aligned_file):
            os.remove(aligned_file)
        
        # Subset
        ref_set = load_ref_pairs(ref_pairs)
        kg1_entities = list(set(p[0] for p in ref_set))
        with open(important_file, 'w') as f:
            for eid in kg1_entities[:n_entities]:
                f.write(f"{eid}\n")
        
        import tokens_cal
        tokens_cal.global_tokens = 0
        
        from LLM1_label_selector import align_entities
        
        # Measure latency
        latencies = []
        original_task = None
        
        start_time = time.time()
        aligned_pairs = align_entities(data_dir, from_m3=False)
        total_time = time.time() - start_time
        
        tokens = tokens_cal.global_tokens
        tokens_cal.global_tokens = 0
        
        metrics = compute_metrics(aligned_file, ref_pairs)
        
        n_pairs = len(aligned_pairs)
        results[f'n_{n_entities}'] = {
            'n_pairs': n_pairs,
            'total_time_s': total_time,
            'time_per_pair_s': total_time / max(1, n_pairs),
            'tokens': tokens,
            'tokens_per_pair': tokens / max(1, n_pairs),
            'metrics': metrics,
        }
        
        print(f"  Pairs: {n_pairs}, Time: {total_time:.1f}s, "
              f"Time/Pair: {total_time/max(1,n_pairs):.2f}s, "
              f"Tokens/Pair: {tokens/max(1,n_pairs):.1f}")
    
    # Restore
    if os.path.exists(bak):
        shutil.copy(bak, important_file)
    
    # Save
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    result_file = os.path.join(RESULTS_DIR, f"efficiency_{timestamp}.json")
    ensure_dir(RESULTS_DIR)
    with open(result_file, 'w') as f:
        json.dump(results, f, indent=2)
    
    print(f"\nEfficiency results saved to: {result_file}")
    
    # Summary
    print("\n" + "="*70)
    print("EFFICIENCY SUMMARY")
    print("="*70)
    print(f"{'N_Entities':<12} {'Pairs':<8} {'Time(s)':<10} {'Time/Pair':<12} {'Tokens/Pair':<12} {'Hits@1':<10}")
    print("-"*64)
    for n_entities in entity_counts:
        key = f'n_{n_entities}'
        if key in results:
            r = results[key]
            m = r['metrics']
            print(f"{n_entities:<12} {r['n_pairs']:<8} {r['total_time_s']:<10.1f} {r['time_per_pair_s']:<12.2f} {r['tokens_per_pair']:<12.1f} {m.get('Hits@1', 0):<10.4f}")
    
    return results

def run_main_experiment():
    """RQ1: Main Results - full dataset run."""
    print("\n" + "="*70)
    print("RQ1: MAIN EXPERIMENT - ICEWS-WIKI (Full Dataset)")
    print("="*70)
    
    data_dir = f"{DATA_BASE}/icews_wiki"
    aligned_file = os.path.join(data_dir, "message_pool", "aligned_entities.txt")
    ref_pairs = os.path.join(data_dir, "ref_pairs")
    
    # Clear evorm state
    evorm_state = os.path.join(data_dir, "evorm_state")
    if os.path.exists(evorm_state):
        shutil.rmtree(evorm_state)
    
    if os.path.exists(aligned_file):
        os.remove(aligned_file)
    
    import tokens_cal
    tokens_cal.global_tokens = 0
    
    from LLM1_label_selector import align_entities
    
    start_time = time.time()
    aligned_pairs = align_entities(data_dir, from_m3=False)
    elapsed = time.time() - start_time
    
    tokens = tokens_cal.global_tokens
    tokens_cal.global_tokens = 0
    
    metrics = compute_metrics(aligned_file, ref_pairs)
    
    results = {
        'dataset': 'ICEWS-WIKI',
        'n_ref_pairs': 3540,
        'aligned_pairs': len(aligned_pairs),
        'time_seconds': elapsed,
        'tokens': tokens,
        'tokens_per_pair': tokens / max(1, len(aligned_pairs)),
        'metrics': metrics,
    }
    
    # Check evorm state
    if os.path.exists(evorm_state):
        rules_file = os.path.join(evorm_state, "rules.json")
        if os.path.exists(rules_file):
            with open(rules_file) as f:
                rules = json.load(f)
            results['n_rules'] = len(rules)
        hyperedges_file = os.path.join(evorm_state, "hyperedges.json")
        if os.path.exists(hyperedges_file):
            with open(hyperedges_file) as f:
                hedges = json.load(f)
            results['n_hyperedges'] = len(hedges)
    
    # Save
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    result_file = os.path.join(RESULTS_DIR, f"main_icews_wiki_{timestamp}.json")
    ensure_dir(RESULTS_DIR)
    with open(result_file, 'w') as f:
        json.dump(results, f, indent=2)
    
    print(f"Main results saved to: {result_file}")
    print(f"\n  Dataset: ICEWS-WIKI")
    print(f"  Aligned: {len(aligned_pairs)}, Time: {elapsed:.1f}s")
    print(f"  Tokens: {tokens}, Tokens/Pair: {tokens/max(1,len(aligned_pairs)):.1f}")
    print(f"  Hits@1: {metrics.get('Hits@1', 'N/A')}")
    print(f"  MRR: {metrics.get('MRR', 'N/A')}")
    print(f"  Precision: {metrics.get('Precision', 'N/A')}")
    print(f"  Recall: {metrics.get('Recall', 'N/A')}")
    if 'n_rules' in results:
        print(f"  Rules: {results['n_rules']}, Hyperedges: {results.get('n_hyperedges', 'N/A')}")
    
    return results

def generate_latex_table(all_results):
    """Generate paper-style LaTeX table."""
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    table_file = os.path.join(RESULTS_DIR, f"paper_table_{timestamp}.tex")
    
    latex = r"""\begin{table}[t]
\centering
\caption{Experimental Results of EvoRM on ICEWS-WIKI}
\label{tab:main_results}
\begin{tabular}{lcccc}
\hline
\textbf{Method} & \textbf{Hits@1} & \textbf{MRR} & \textbf{Precision} & \textbf{Recall} \\
\hline
"""
    
    # Main results
    if 'main' in all_results:
        m = all_results['main']['metrics']
        latex += f"EvoRM (Ours) & {m['Hits@1']:.4f} & {m['MRR']:.4f} & {m['Precision']:.4f} & {m['Recall']:.4f} \\\\\n"
    
    latex += r"""\hline
\end{tabular}
\end{table}
"""
    
    with open(table_file, 'w') as f:
        f.write(latex)
    
    print(f"LaTeX table saved to: {table_file}")
    return table_file

def main():
    parser = argparse.ArgumentParser(description='EvoRM Comprehensive Experiment Runner')
    parser.add_argument('--all', action='store_true', help='Run all experiments')
    parser.add_argument('--ablation', action='store_true', help='Run ablation study (RQ3)')
    parser.add_argument('--coldstart', action='store_true', help='Run cold-start scaling (RQ5)')
    parser.add_argument('--main', action='store_true', help='Run main experiment (RQ1)')
    parser.add_argument('--efficiency', action='store_true', help='Run efficiency analysis (RQ4)')
    args = parser.parse_args()
    
    # Default: run everything if no flag specified
    run_all = args.all or not any([args.ablation, args.coldstart, args.main, args.efficiency])
    
    all_results = {}
    
    if run_all or args.main:
        all_results['main'] = run_main_experiment()
    
    if run_all or args.ablation:
        all_results['ablation'] = run_ablation_study()
    
    if run_all or args.coldstart:
        all_results['coldstart'] = run_coldstart_experiment()
    
    if run_all or args.efficiency:
        all_results['efficiency'] = run_efficiency_analysis()
    
    # Generate paper table
    generate_latex_table(all_results)
    
    print("\n" + "="*70)
    print("ALL EXPERIMENTS COMPLETE!")
    print("="*70)
    print(f"Results saved in: {RESULTS_DIR}")
    for f in sorted(os.listdir(RESULTS_DIR)):
        print(f"  {f}")

if __name__ == '__main__':
    main()