#!/usr/bin/env python3
"""
EvoRM Ablation Study - All 6 modes
Uses ref_pairs order for entity selection (correct approach).
"""
import os, sys, json, time, shutil
from datetime import datetime

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

os.environ["OPENAI_API_BASE"] = "https://hk.xty.app/v1"
os.environ["OPENAI_API_KEY"] = "sk-7a7Ev4VcVyysPLT5hqtqIVD6PybzJ1ZlEIVZddIR3NtZvPgK"

DATA_DIR = "/root/autodl-tmp/AdaCoAgentEA/data/icews_wiki"
RESULTS_DIR = "/root/autodl-tmp/AdaCoAgentEA/results"
ALIGNED_FILE = os.path.join(DATA_DIR, "message_pool", "aligned_entities.txt")
REF_PAIRS_FILE = os.path.join(DATA_DIR, "ref_pairs")
IMPORTANT_FILE = os.path.join(DATA_DIR, "message_pool", "important_entities.txt")
RETRIEVER_FILE = os.path.join(DATA_DIR, "message_pool", "retriever_outputs.txt")
EVORM_STATE = os.path.join(DATA_DIR, "evorm_state")

os.makedirs(RESULTS_DIR, exist_ok=True)

def load_ref_pairs():
    pairs = {}
    with open(REF_PAIRS_FILE, 'r') as f:
        for line in f:
            parts = line.strip().split('\t')
            if len(parts) == 2:
                try:
                    pairs[int(parts[0])] = int(parts[1])
                except:
                    pass
    return pairs

def compute_metrics(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 = set(ref_subset.items())
    else:
        ref_set = set(load_ref_pairs().items())
    if not ref_set:
        return {'error': 'No ref pairs'}
    tp = len(aligned & ref_set)
    fp = len(aligned - ref_set)
    fn = len(ref_set - aligned)
    return {
        'Hits@1': round(tp / max(1, len(ref_set)), 4),
        'MRR': round(tp / max(1, len(ref_set)), 4),
        'Precision': round(tp / max(1, tp + fp), 4),
        'Recall': round(tp / max(1, tp + fn), 4),
        'F1': round(2*tp/max(1, 2*tp+fp+fn), 4),
        'TP': tp, 'FP': fp, 'FN': fn,
        'Total_Ref': len(ref_set),
        'Total_Aligned': len(aligned),
    }

def clear_state():
    if os.path.exists(EVORM_STATE):
        shutil.rmtree(EVORM_STATE)
    if os.path.exists(ALIGNED_FILE):
        os.remove(ALIGNED_FILE)

def create_subset_from_ref_order(n_entities, ref_map):
    """Create subset using ref_pairs order (not retriever order)."""
    # Get entities in ref_pairs order
    ref_order = []
    with open(REF_PAIRS_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 ref_order:
                        ref_order.append(e1)
                except:
                    pass
    
    actual_n = min(n_entities, len(ref_order))
    subset_kg1 = set(ref_order[:actual_n])
    
    # Filter retriever lines for these entities
    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
    
    with open(IMPORTANT_FILE, 'w') as f:
        f.writelines(subset_lines)
    
    subset_ref = {e1: ref_map[e1] for e1 in subset_kg1 if e1 in ref_map}
    print(f"  Created subset: {len(subset_kg1)} entities, {len(subset_ref)} ref pairs, {len(subset_lines)} retriever lines")
    return subset_ref

def main():
    n_entities = 200
    ref_map = load_ref_pairs()
    print(f"Loaded {len(ref_map)} reference pairs")
    
    import tokens_cal
    from LLM1_label_selector import align_entities
    
    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{'='*50}")
        print(f"  {config_name}")
        print(f"{'='*50}")
        
        subset_ref = create_subset_from_ref_order(n_entities, ref_map)
        clear_state()
        
        tokens_cal.global_tokens = 0
        start = time.time()
        aligned_pairs = align_entities(DATA_DIR, from_m3=False, ablation_mode=ablation_mode)
        elapsed = time.time() - start
        tokens = tokens_cal.global_tokens
        n_pairs = len(aligned_pairs)
        metrics = compute_metrics(ref_subset=subset_ref)
        
        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}, Time: {elapsed:.1f}s, Tokens: {tokens}")
        print(f"  Hits@1: {metrics['Hits@1']}, Precision: {metrics['Precision']}, Recall: {metrics['Recall']}")
    
    # Save
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    result_file = os.path.join(RESULTS_DIR, f"ablation_6modes_{timestamp}.json")
    with open(result_file, 'w') as f:
        json.dump({
            'experiment': 'EvoRM Ablation Study (6 modes)',
            'dataset': 'ICEWS-WIKI',
            'n_entities': n_entities,
            'results': results,
        }, f, indent=2)
    
    # Summary
    print(f"\n{'='*70}")
    print("ABLATION SUMMARY")
    print(f"{'='*70}")
    print(f"{'Configuration':<30} {'Hits@1':>8} {'Prec':>8} {'Recall':>8} {'F1':>8} {'Tok':>10}")
    print("-" * 70)
    for cname, _ in configs:
        if cname in results:
            m = results[cname]['metrics']
            print(f"{cname:<30} {m['Hits@1']:>8.4f} {m['Precision']:>8.4f} {m['Recall']:>8.4f} {m['F1']:>8.4f} {results[cname]['tokens']:>10}")
    
    print(f"\nResults saved to: {result_file}")

if __name__ == '__main__':
    main()