#!/usr/bin/env python3
"""
EvoRM 全面实验脚本
===================
覆盖论文全部 5 个 Research Questions:
  RQ1: 主实验 (ICEWS-WIKI + ICEWS-YAGO + BETA)
  RQ2: 跨模型分析 (GPT-3.5, GPT-4o-mini, DeepSeek)
  RQ3: 消融实验 (Full / w/o Stage-1 / w/o Maintenance)
  RQ4: 效率分析 (Token消耗 + 延迟)
  RQ5: 冷启动扩展 (不同 warmup 规模)

用法:
  python run_comprehensive.py --dataset icews_wiki --all
  python run_comprehensive.py --dataset icews_yago --main
  python run_comprehensive.py --dataset BETA --ablation
  python run_comprehensive.py --all_datasets --all
"""

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

# ===== 配置 =====
API_CONFIGS = {
    'gpt-3.5-turbo': {
        'base_url': 'https://hk.xty.app/v1',
        'api_key': 'sk-7a7Ev4VcVyysPLT5hqtqIVD6PybzJ1ZlEIVZddIR3NtZvPgK',
        'model': 'gpt-3.5-turbo-1106',
    },
    'gpt-4o-mini': {
        'base_url': 'https://hk.xty.app/v1',
        'api_key': 'sk-7a7Ev4VcVyysPLT5hqtqIVD6PybzJ1ZlEIVZddIR3NtZvPgK',
        'model': 'gpt-4o-mini',
    },
}

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

os.makedirs(RESULTS_DIR, exist_ok=True)

# ===== 工具函数 =====
def load_ref_pairs(path):
    pairs = set()
    if not os.path.exists(path):
        return pairs
    with open(path, 'r') as f:
        for line in f:
            parts = line.strip().split('\t')
            if len(parts) == 2:
                try:
                    pairs.add((int(parts[0]), int(parts[1])))
                except ValueError:
                    pass
    return pairs

def compute_metrics(aligned_file, ref_pairs_file, ref_subset=None):
    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:
                        continue

    if ref_subset:
        ref_set = ref_subset
    else:
        ref_set = load_ref_pairs(ref_pairs_file)

    if not ref_set:
        return {'error': 'No reference pairs'}

    tp = len(aligned & ref_set)
    fp = len(aligned - ref_set)
    fn = len(ref_set - aligned)

    precision = tp / max(1, tp + fp)
    recall = tp / max(1, tp + fn)
    hits1 = tp / len(ref_set)

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

def clear_state(data_dir):
    evorm_state = os.path.join(data_dir, "evorm_state")
    aligned_file = os.path.join(data_dir, "message_pool", "aligned_entities.txt")
    if os.path.exists(evorm_state):
        shutil.rmtree(evorm_state)
    if os.path.exists(aligned_file):
        os.remove(aligned_file)

def setup_api(model_name):
    config = API_CONFIGS[model_name]
    os.environ["OPENAI_API_BASE"] = config['base_url']
    os.environ["OPENAI_API_KEY"] = config['api_key']
    return config

# ===== RQ1: 主实验 =====
def run_main_experiment(data_dir, dataset_name, model_name='gpt-3.5-turbo'):
    print(f"\n{'='*70}")
    print(f"RQ1: MAIN EXPERIMENT - {dataset_name} ({model_name})")
    print(f"{'='*70}")

    config = setup_api(model_name)
    ref_pairs_file = os.path.join(data_dir, "ref_pairs")
    aligned_file = os.path.join(data_dir, "message_pool", "aligned_entities.txt")
    important_file = os.path.join(data_dir, "message_pool", "important_entities.txt")
    retriever_file = os.path.join(data_dir, "message_pool", "retriever_outputs.txt")

    # Check if dataset has standard KG structure
    has_kg = os.path.exists(os.path.join(data_dir, 'triples_1'))
    has_retriever = os.path.exists(retriever_file)

    if not has_retriever and not has_kg:
        print(f"  SKIP: {dataset_name} missing retriever_outputs.txt and triples")
        return None

    # Backup important_entities
    if os.path.exists(important_file):
        bak = important_file + ".bak_evorm"
        if not os.path.exists(bak):
            shutil.copy(important_file, bak)

    clear_state(data_dir)

    # If retriever outputs exist, use them as important_entities
    if has_retriever:
        shutil.copy(retriever_file, important_file)

    sys.path.insert(0, '/root/autodl-tmp/AdaCoAgentEA')
    sys.path.insert(0, '/root/autodl-tmp/AdaCoAgentEA/Area2')

    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
    n_pairs = len(aligned_pairs)

    metrics = {}
    if os.path.exists(ref_pairs_file):
        metrics = compute_metrics(aligned_file, ref_pairs_file)

    evorm_info = {}
    evorm_state = os.path.join(data_dir, "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:
                evorm_info['rules'] = len(json.load(f))
        hedges_file = os.path.join(evorm_state, "hyperedges.json")
        if os.path.exists(hedges_file):
            with open(hedges_file) as f:
                evorm_info['hyperedges'] = len(json.load(f))

    result = {
        'dataset': dataset_name,
        'model': model_name,
        'aligned_pairs': n_pairs,
        'time_seconds': round(elapsed, 1),
        'tokens': tokens,
        'tokens_per_pair': round(tokens / max(1, n_pairs), 1),
        'time_per_pair_s': round(elapsed / max(1, n_pairs), 2),
        'metrics': metrics,
        'evorm_state': evorm_info,
    }

    print(f"  Aligned: {n_pairs}, Time: {elapsed:.1f}s, Tokens: {tokens}")
    if metrics:
        print(f"  Hits@1: {metrics.get('Hits@1', 'N/A')}, Precision: {metrics.get('Precision', 'N/A')}")

    return result

# ===== RQ2: 跨模型分析 =====
def run_cross_model_experiment(data_dir, dataset_name):
    print(f"\n{'='*70}")
    print(f"RQ2: CROSS-MODEL ANALYSIS - {dataset_name}")
    print(f"{'='*70}")

    models = ['gpt-3.5-turbo', 'gpt-4o-mini']
    results = {}

    for model_name in models:
        result = run_main_experiment(data_dir, dataset_name, model_name)
        if result:
            results[model_name] = result

    return results

# ===== RQ3: 消融实验 =====
def run_ablation_experiment(data_dir, dataset_name, n_entities=200):
    print(f"\n{'='*70}")
    print(f"RQ3: ABLATION STUDY - {dataset_name} ({n_entities} entities)")
    print(f"{'='*70}")

    config = setup_api('gpt-3.5-turbo')
    ref_pairs_file = os.path.join(data_dir, "ref_pairs")
    aligned_file = os.path.join(data_dir, "message_pool", "aligned_entities.txt")
    important_file = os.path.join(data_dir, "message_pool", "important_entities.txt")
    retriever_file = os.path.join(data_dir, "message_pool", "retriever_outputs.txt")

    if not os.path.exists(retriever_file):
        print(f"  SKIP: no retriever_outputs.txt")
        return None

    # Load ref pairs
    ref_set = load_ref_pairs(ref_pairs_file)
    if not ref_set:
        print(f"  SKIP: no ref_pairs")
        return None

    # Get ordered KG1 entities
    kg1_order = []
    with open(retriever_file, 'r') as f:
        for line in f:
            parts = line.strip().split('\t')
            if len(parts) == 2:
                try:
                    e1 = int(parts[0])
                    if e1 not in kg1_order:
                        kg1_order.append(e1)
                except:
                    pass

    # Filter to only entities with reference pairs
    ref_kg1 = set(e1 for e1, e2 in ref_set)
    kg1_order_filtered = [e for e in kg1_order if e in ref_kg1]
    if len(kg1_order_filtered) < n_entities:
        n_entities = len(kg1_order_filtered)
        print(f"  WARNING: Only {n_entities} entities with ref pairs available")
    subset_kg1 = set(kg1_order_filtered[:n_entities])
    subset_ref = {e1: ref for e1, ref in ref_set if e1 in subset_kg1}
    ref_set_limited = set(subset_ref.items())

    # Filter retriever lines
    subset_lines = []
    with open(retriever_file, 'r') as f:
        for line in f:
            parts = line.strip().split('\t')
            if len(parts) == 2:
                try:
                    if int(parts[0]) in subset_kg1:
                        subset_lines.append(line)
                except:
                    pass

    # Backup
    bak = important_file + ".bak_evorm"
    if not os.path.exists(bak):
        shutil.copy(important_file, bak)

    sys.path.insert(0, '/root/autodl-tmp/AdaCoAgentEA')
    sys.path.insert(0, '/root/autodl-tmp/AdaCoAgentEA/Area2')

    configs = [
        ('Full System (EvoRM)', None),
        ('w/o Stage-1 (Symbolic)', 'no_stage1'),
        ('w/o Maintenance', 'no_maintenance'),
        ('w/o MLP Gate', 'no_mlp_gate'),
        ('w/o Hypergraph', 'no_hypergraph'),
        ('w/o Mlight', 'no_mlight'),
    ]

    results = {}
    for config_name, ablation_mode in configs:
        print(f"\n  --- {config_name} ---")

        with open(important_file, 'w') as f:
            f.writelines(subset_lines)

        clear_state(data_dir)

        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, ablation_mode=ablation_mode)
        elapsed = time.time() - start_time

        tokens = tokens_cal.global_tokens
        n_pairs = len(aligned_pairs)
        metrics = compute_metrics(aligned_file, ref_pairs_file, ref_subset=ref_set_limited)

        results[config_name] = {
            'aligned_pairs': n_pairs,
            'time_seconds': round(elapsed, 1),
            'tokens': tokens,
            'tokens_per_pair': round(tokens / max(1, n_pairs), 1),
            'metrics': metrics,
        }

        print(f"    Aligned: {n_pairs}, Hits@1: {metrics['Hits@1']}, Tokens: {tokens}")

    # Restore
    shutil.copy(bak, important_file)
    return results

# ===== RQ4: 效率分析 =====
def run_efficiency_experiment(data_dir, dataset_name):
    print(f"\n{'='*70}")
    print(f"RQ4: EFFICIENCY ANALYSIS - {dataset_name}")
    print(f"{'='*70}")

    config = setup_api('gpt-3.5-turbo')
    ref_pairs_file = os.path.join(data_dir, "ref_pairs")
    aligned_file = os.path.join(data_dir, "message_pool", "aligned_entities.txt")
    important_file = os.path.join(data_dir, "message_pool", "important_entities.txt")
    retriever_file = os.path.join(data_dir, "message_pool", "retriever_outputs.txt")

    if not os.path.exists(retriever_file):
        print(f"  SKIP: no retriever_outputs.txt")
        return None

    # Load all retriever lines
    with open(retriever_file, 'r') as f:
        all_lines = f.readlines()

    # Get ordered KG1 entities
    kg1_order = []
    for line in all_lines:
        parts = line.strip().split('\t')
        if len(parts) == 2:
            try:
                e1 = int(parts[0])
                if e1 not in kg1_order:
                    kg1_order.append(e1)
            except:
                pass

    # Backup
    bak = important_file + ".bak_evorm"
    if not os.path.exists(bak):
        shutil.copy(important_file, bak)

    sys.path.insert(0, '/root/autodl-tmp/AdaCoAgentEA')
    sys.path.insert(0, '/root/autodl-tmp/AdaCoAgentEA/Area2')

    # Filter to only entities with reference pairs
    ref_pairs = load_ref_pairs(ref_pairs_file)
    ref_kg1 = set(e1 for e1, e2 in ref_pairs)
    kg1_order_filt = [e for e in kg1_order if e in ref_kg1]

    # Test with different entity counts
    entity_counts = [100, 300, 500]
    results = {}

    for n_entities in entity_counts:
        print(f"\n  --- N_entities = {n_entities} ---")

        actual_n = min(n_entities, len(kg1_order_filt))
        subset_kg1 = set(kg1_order_filt[:actual_n])
        subset_lines = [l for l in all_lines 
                       if len(l.strip().split('\t')) == 2 
                       and int(l.strip().split('\t')[0]) in subset_kg1]

        with open(important_file, 'w') as f:
            f.writelines(subset_lines)

        clear_state(data_dir)

        import tokens_cal
        tokens_cal.global_tokens = 0

        from LLM1_label_selector import align_entities

        # Measure per-pair latency
        start_time = time.time()
        aligned_pairs = align_entities(data_dir, from_m3=False)
        total_time = time.time() - start_time

        tokens = tokens_cal.global_tokens
        n_pairs = len(aligned_pairs)

        results[f'n_{n_entities}'] = {
            'n_pairs': n_pairs,
            'total_time_s': round(total_time, 2),
            'time_per_pair_s': round(total_time / max(1, n_pairs), 3),
            'tokens': tokens,
            'tokens_per_pair': round(tokens / max(1, n_pairs), 1),
            'throughput_pairs_per_s': round(n_pairs / max(1, total_time), 2),
        }

        print(f"    Pairs: {n_pairs}, Time: {total_time:.1f}s, "
              f"Time/Pair: {total_time/max(1,n_pairs):.3f}s, "
              f"Tokens/Pair: {tokens/max(1,n_pairs):.1f}")

    # Restore
    shutil.copy(bak, important_file)
    return results

# ===== RQ5: 冷启动扩展 =====
def run_coldstart_experiment(data_dir, dataset_name):
    print(f"\n{'='*70}")
    print(f"RQ5: COLD-START SCALING - {dataset_name}")
    print(f"{'='*70}")

    config = setup_api('gpt-3.5-turbo')
    ref_pairs_file = os.path.join(data_dir, "ref_pairs")
    aligned_file = os.path.join(data_dir, "message_pool", "aligned_entities.txt")
    important_file = os.path.join(data_dir, "message_pool", "important_entities.txt")
    retriever_file = os.path.join(data_dir, "message_pool", "retriever_outputs.txt")

    if not os.path.exists(retriever_file):
        print(f"  SKIP: no retriever_outputs.txt")
        return None

    with open(retriever_file, 'r') as f:
        all_lines = f.readlines()

    kg1_order = []
    for line in all_lines:
        parts = line.strip().split('\t')
        if len(parts) == 2:
            try:
                e1 = int(parts[0])
                if e1 not in kg1_order:
                    kg1_order.append(e1)
            except:
                pass

    ref_pairs = load_ref_pairs(ref_pairs_file)
    ref_map = {e1: e2 for e1, e2 in ref_pairs}

    bak = important_file + ".bak_evorm"
    if not os.path.exists(bak):
        shutil.copy(important_file, bak)

    sys.path.insert(0, '/root/autodl-tmp/AdaCoAgentEA')
    sys.path.insert(0, '/root/autodl-tmp/AdaCoAgentEA/Area2')

    # Filter to only entities with reference pairs
    ref_kg1 = set(e1 for e1, e2 in ref_pairs)
    kg1_order_filt = [e for e in kg1_order if e in ref_kg1]

    warmup_sizes = [100, 200, 500, 1000]
    results = {}

    for n_warmup in warmup_sizes:
        print(f"\n  --- N_warmup = {n_warmup} ---")

        actual_n = min(n_warmup, len(kg1_order_filt))
        subset_kg1 = set(kg1_order_filt[:actual_n])
        subset_ref = {e1: ref_map[e1] for e1 in subset_kg1 if e1 in ref_map}
        ref_set_limited = set(subset_ref.items())

        subset_lines = [l for l in all_lines 
                       if len(l.strip().split('\t')) == 2 
                       and int(l.strip().split('\t')[0]) in subset_kg1]

        with open(important_file, 'w') as f:
            f.writelines(subset_lines)

        clear_state(data_dir)

        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
        n_pairs = len(aligned_pairs)
        metrics = compute_metrics(aligned_file, ref_pairs_file, ref_subset=ref_set_limited)

        evorm_info = {}
        evorm_state = os.path.join(data_dir, "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:
                    evorm_info['rules'] = len(json.load(f))

        results[f'n_{n_warmup}'] = {
            'aligned_pairs': n_pairs,
            'time_seconds': round(elapsed, 1),
            'tokens': tokens,
            'tokens_per_pair': round(tokens / max(1, n_pairs), 1),
            'metrics': metrics,
            'evorm_state': evorm_info,
        }

        print(f"    Aligned: {n_pairs}, Hits@1: {metrics['Hits@1']}, "
              f"Tokens/Pair: {tokens/max(1,n_pairs):.1f}, Rules: {evorm_info.get('rules', 0)}")

    shutil.copy(bak, important_file)
    return results

# ===== 主函数 =====
def main():
    parser = argparse.ArgumentParser(description='EvoRM Comprehensive Experiments')
    parser.add_argument('--dataset', type=str, default='icews_wiki',
                       choices=['icews_wiki', 'icews_yago', 'BETA'],
                       help='Dataset to run on')
    parser.add_argument('--all_datasets', action='store_true',
                       help='Run on all available datasets')
    parser.add_argument('--all', action='store_true', help='Run all experiments')
    parser.add_argument('--main', action='store_true', help='RQ1: Main experiment')
    parser.add_argument('--cross_model', action='store_true', help='RQ2: Cross-model')
    parser.add_argument('--ablation', action='store_true', help='RQ3: Ablation study')
    parser.add_argument('--efficiency', action='store_true', help='RQ4: Efficiency')
    parser.add_argument('--coldstart', action='store_true', help='RQ5: Cold-start')
    args = parser.parse_args()

    run_all = args.all or not any([args.main, args.cross_model, args.ablation, 
                                     args.efficiency, args.coldstart])

    datasets = [args.dataset]
    if args.all_datasets:
        datasets = ['icews_wiki', 'icews_yago', 'BETA']

    all_results = {}

    for ds_name in datasets:
        data_dir = os.path.join(DATA_BASE, ds_name)
        if not os.path.exists(data_dir):
            print(f"SKIP: {data_dir} not found")
            continue

        print(f"\n{'#'*70}")
        print(f"# DATASET: {ds_name}")
        print(f"{'#'*70}")

        ds_results = {}

        if run_all or args.main:
            ds_results['RQ1_main'] = run_main_experiment(data_dir, ds_name)

        if run_all or args.cross_model:
            ds_results['RQ2_cross_model'] = run_cross_model_experiment(data_dir, ds_name)

        if run_all or args.ablation:
            ds_results['RQ3_ablation'] = run_ablation_experiment(data_dir, ds_name)

        if run_all or args.efficiency:
            ds_results['RQ4_efficiency'] = run_efficiency_experiment(data_dir, ds_name)

        if run_all or args.coldstart:
            ds_results['RQ5_coldstart'] = run_coldstart_experiment(data_dir, ds_name)

        all_results[ds_name] = ds_results

    # 保存结果
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    result_file = os.path.join(RESULTS_DIR, f"comprehensive_{timestamp}.json")
    with open(result_file, 'w') as f:
        json.dump(all_results, f, indent=2)

    # 打印汇总
    print(f"\n{'='*80}")
    print("COMPREHENSIVE EXPERIMENT SUMMARY")
    print(f"{'='*80}")

    for ds_name, ds_results in all_results.items():
        print(f"\n--- {ds_name} ---")
        for exp_name, exp_result in ds_results.items():
            if exp_result and isinstance(exp_result, dict):
                if 'metrics' in exp_result and exp_result['metrics']:
                    print(f"  {exp_name}: Hits@1={exp_result['metrics'].get('Hits@1','N/A')}, "
                          f"Tokens={exp_result.get('tokens','N/A')}")

    print(f"\nResults saved to: {result_file}")

if __name__ == '__main__':
    main()