#!/usr/bin/env python3
"""
EvoRM Complete Experiment Suite — All Paper Datasets + Ablation
===============================================================
Covers all available datasets across ER, EA, EL, SM tasks.
Includes ablation studies (no_stage1, no_maintenance, no_mlp_gate).
"""
import os, sys, json, time, threading

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

CONFIG = {
    'model': 'gpt-3.5-turbo-1106',
    'max_pairs': 50,
    'max_workers': 5,
    'data_dir': '/root/autodl-tmp/AdaCoAgentEA/data',
    'results_dir': '/root/autodl-tmp/AdaCoAgentEA/results',
    'evorm_persistence': '/tmp/evorm_complete',
    'mlp_n_warmup': 20,  # Lower for small-scale experiments
}

os.makedirs(CONFIG['results_dir'], exist_ok=True)

# ==============================================================================
# Utility functions
# ==============================================================================

def load_dm_dataset(ds_path, max_pairs):
    """Load DeepMatcher dataset. Joins test.csv with tableA/tableB."""
    import csv, random
    table_a, table_b = {}, {}
    for tn, td in [('tableA.csv', table_a), ('tableB.csv', table_b)]:
        tp = os.path.join(ds_path, tn)
        if not os.path.exists(tp):
            tp = tp.replace('.csv', '_tok.csv')
        if os.path.exists(tp):
            with open(tp, 'r', encoding='utf-8', errors='replace') as f:
                for row in csv.DictReader(f):
                    td[row.get('id', '')] = dict(row)
    pairs = []
    for split in ['test.csv', 'valid.csv']:
        sp = os.path.join(ds_path, split)
        if not os.path.exists(sp):
            continue
        with open(sp, 'r', encoding='utf-8', errors='replace') as f:
            for row in csv.DictReader(f):
                record_a = table_a.get(row.get('ltable_id', ''), {'id': row.get('ltable_id', '')})
                record_b = table_b.get(row.get('rtable_id', ''), {'id': row.get('rtable_id', '')})
                pairs.append((dict(record_a), dict(record_b), int(row.get('label', 0))))
        if pairs:
            break
    if len(pairs) > max_pairs:
        random.seed(42)
        pos = [p for p in pairs if p[2] == 1]
        neg = [p for p in pairs if p[2] == 0]
        n_pos = min(len(pos), max_pairs // 2)
        pairs = random.sample(pos, n_pos) + random.sample(neg, max_pairs - n_pos)
        random.shuffle(pairs)
    return pairs

def load_json(path, max_items):
    if not os.path.exists(path):
        return []
    with open(path, 'r') as f:
        try:
            data = json.load(f)
            return data[:max_items] if isinstance(data, list) else []
        except json.JSONDecodeError:
            items = [json.loads(l) for l in f.read().strip().split('\n') if l.strip()]
            return items[:max_items]

def compute_er_metrics(predictions, labels):
    tp = sum(1 for p, l in zip(predictions, labels) if p == 1 and l == 1)
    fp = sum(1 for p, l in zip(predictions, labels) if p == 1 and l == 0)
    fn = sum(1 for p, l in zip(predictions, labels) if p == 0 and l == 1)
    precision = tp / max(1, tp + fp)
    recall = tp / max(1, tp + fn)
    f1 = 2 * precision * recall / max(0.001, precision + recall)
    return {'F1': round(f1, 4), 'Precision': round(precision, 4), 'Recall': round(recall, 4)}

# ==============================================================================
# ER: Entity Resolution (MatchGPT + AnyMatch)
# ==============================================================================

def run_er():
    from baselines.evorm_wrappers.matchgpt import run_matchgpt, run_anymatch
    
    datasets = [
        ("Amazon-Google", "Structured"),
        ("Beer", "Structured"),
        ("DBLP-ACM", "Structured"),
        ("DBLP-GoogleScholar", "Structured"),
        ("Fodors-Zagats", "Structured"),
        ("Walmart-Amazon", "Structured"),
        ("iTunes-Amazon", "Structured"),
        ("Abt-Buy", "Textual"),
    ]
    
    results = []
    print(f"\n{'='*60}\nER Experiments: {len(datasets)} datasets\n{'='*60}")
    
    for ds_name, category in datasets:
        ds_path = os.path.join(CONFIG['data_dir'], 'er', category, ds_name)
        if not os.path.exists(ds_path):
            print(f"  {ds_name}: not found, skipping")
            continue
        
        pairs = load_dm_dataset(ds_path, CONFIG['max_pairs'])
        if not pairs:
            print(f"  {ds_name}: no data, skipping")
            continue
        
        pos_count = sum(1 for _, _, l in pairs if l == 1)
        print(f"\n--- {ds_name} ({category}): {len(pairs)} pairs (pos={pos_count}) ---")
        
        # Save as temp JSON for baseline functions
        tmp_path = f'/tmp/er_{ds_name.replace("-", "_")}.json'
        with open(tmp_path, 'w') as f:
            json.dump([{'record_a': p[0], 'record_b': p[1], 'label': p[2]} for p in pairs], f)
        
        for backend in ['default', 'mixtral', 'solar', 'beluga2']:
            for use_evorm in [False, True]:
                try:
                    r = run_matchgpt(tmp_path, backend=backend, use_evorm=use_evorm, max_pairs=CONFIG['max_pairs'])
                    r['dataset'] = ds_name; r['task'] = 'ER'
                    results.append(r)
                    ev = '+EvoRM' if use_evorm else ''
                    print(f"  MatchGPT[{backend}]{ev}: F1={r.get('F1', 'N/A')}")
                except Exception as e:
                    print(f"  MatchGPT[{backend}]{'+EvoRM' if use_evorm else ''}: ERROR - {e}")
        
        for use_evorm in [False, True]:
            try:
                r = run_anymatch(tmp_path, use_evorm=use_evorm, max_pairs=CONFIG['max_pairs'])
                r['dataset'] = ds_name; r['task'] = 'ER'
                results.append(r)
                ev = '+EvoRM' if use_evorm else ''
                print(f"  AnyMatch{ev}: F1={r.get('F1', 'N/A')}")
            except Exception as e:
                print(f"  AnyMatch{'+EvoRM' if use_evorm else ''}: ERROR - {e}")
    
    return results

# ==============================================================================
# EA: Entity Alignment (DBP15K)
# ==============================================================================

def run_ea():
    from openai import OpenAI
    from evorm_plugin import EvoRMPlugin
    from ThreadPoolExecutor import ThreadPoolExecutor
    from tqdm import tqdm
    import tokens_cal
    
    ea_dir = os.path.join(CONFIG['data_dir'], 'ea', 'dbp15k')
    results = []
    
    lang_pairs = [d for d in sorted(os.listdir(ea_dir)) if os.path.isdir(os.path.join(ea_dir, d))]
    lang_pairs = [lp for lp in lang_pairs if lp in ['fr_en', 'ja_en', 'zh_en']]
    
    if not lang_pairs:
        print("No DBP15K language pairs found")
        return results
    
    print(f"\n{'='*60}\nEA Experiments: {len(lang_pairs)} language pairs\n{'='*60}")
    
    client = OpenAI(base_url=os.environ.get("OPENAI_BASE_URL", "https://api.openai.com/v1"), api_key="")
    
    for lang_pair in lang_pairs:
        lp_dir = os.path.join(ea_dir, lang_pair)
        ref_path = os.path.join(lp_dir, 'ent_ILLs')
        if not os.path.exists(ref_path):
            continue
        
        with open(ref_path, 'r') as f:
            ref_lines = [l.strip().split('\t') for l in f if l.strip()][:CONFIG['max_pairs']]
        
        print(f"\n--- DBP15K {lang_pair}: {len(ref_lines)} pairs ---")
        
        for use_evorm in [False, True]:
            try:
                evorm = EvoRMPlugin(
                    client=client,
                    persistence_dir=f'{CONFIG["evorm_persistence"]}/ea_{lang_pair}',
                    mlp_n_warmup=CONFIG['mlp_n_warmup'],
                ) if use_evorm else None
                
                predictions, labels, lock = [], [], threading.Lock()
                t0 = time.time()
                tokens_cal.global_tokens = 0
                
                def process_pair(e1, e2):
                    es_ctx = {'entity_name': str(e1)}
                    et_ctx = {'entity_name': str(e2)}
                    prompt = f"Entity Alignment: Are entity {e1} and {e2} the same? Answer MATCH or NON-MATCH."
                    
                    if evorm:
                        eid = str(hash(f"{e1}_{e2}"))[:8]
                        decision, triggered, cands = evorm.stage1(eid, eid, es_ctx, et_ctx)
                        if decision is not None:
                            pred = decision
                        else:
                            result = evorm.stage2(eid, eid, es_ctx, et_ctx, triggered, cands, prompt)
                            pred = result.get('decision', 0) or 0
                            evorm.record_trajectory(eid, eid, es_ctx, et_ctx, decision=pred,
                                                    rationale=result.get('rationale', ''),
                                                    triggered_rules=triggered,
                                                    rule_feedback=result.get('rule_feedback', {}))
                    else:
                        resp = client.chat.completions.create(
                            model=CONFIG['model'],
                            messages=[{'role': 'user', 'content': prompt}],
                            temperature=0.1)
                        answer = resp.choices[0].message.content.strip().upper()
                        tokens_cal.update_add_var(resp.usage.total_tokens)
                        pred = 1 if 'MATCH' in answer and 'NON' not in answer else 0
                    
                    with lock:
                        predictions.append(pred)
                        labels.append(1)  # All ref pairs are matches
                
                executor = ThreadPoolExecutor(max_workers=CONFIG['max_workers'])
                for e1, e2 in tqdm(ref_lines, desc=f"EA {lang_pair}{' +EvoRM' if use_evorm else ''}"):
                    executor.submit(process_pair, e1, e2)
                executor.shutdown(wait=True)
                
                elapsed = time.time() - t0
                m = compute_er_metrics(predictions, labels)
                m['dataset'] = f'DBP15K {lang_pair}'
                m['task'] = 'EA'
                m['method'] = f"EA{' + EvoRM' if use_evorm else ''}"
                m['time_seconds'] = round(elapsed, 1)
                m['tokens'] = tokens_cal.global_tokens
                if evorm:
                    m['evorm_stats'] = evorm.get_stats()
                results.append(m)
                ev = '+EvoRM' if use_evorm else ''
                print(f"  EA{ev}: F1={m.get('F1', 'N/A')}, Prec={m.get('Precision', 'N/A')}, Rec={m.get('Recall', 'N/A')}")
            except Exception as e:
                print(f"  EA{'+EvoRM' if use_evorm else ''}: ERROR - {e}")
                import traceback; traceback.print_exc()
    
    return results

# ==============================================================================
# EL: Entity Linking (LELA)
# ==============================================================================

def run_el():
    from baselines.evorm_wrappers.lela_el import run_lela
    
    el_dir = os.path.join(CONFIG['data_dir'], 'el')
    datasets = []
    for fname in ['AIDA-CoNLL_real.json', 'WNED-CWEB_real.json']:
        fp = os.path.join(el_dir, fname)
        if os.path.exists(fp):
            datasets.append((fname.replace('_real.json', ''), fp))
    
    if not datasets:
        print("No EL datasets found, skipping")
        return []
    
    results = []
    print(f"\n{'='*60}\nEL Experiments: {len(datasets)} datasets\n{'='*60}")
    
    for ds_name, ds_path in datasets:
        items = load_json(ds_path, CONFIG['max_pairs'])
        print(f"\n--- {ds_name}: {len(items)} items ---")
        
        for use_evorm in [False, True]:
            try:
                r = run_lela(ds_path, use_evorm=use_evorm, max_items=CONFIG['max_pairs'])
                r['dataset'] = ds_name; r['task'] = 'EL'
                results.append(r)
                ev = '+EvoRM' if use_evorm else ''
                print(f"  LELA{ev}: Acc={r.get('Accuracy', 'N/A')}")
            except Exception as e:
                print(f"  LELA{'+EvoRM' if use_evorm else ''}: ERROR - {e}")
    
    return results

# ==============================================================================
# SM: Schema Matching
# ==============================================================================

def run_sm():
    from baselines.evorm_wrappers.lela_el import run_schema_matching
    
    sm_dir = os.path.join(CONFIG['data_dir'], 'sm')
    datasets = []
    for fname in ['T2Dv2_schema.json']:
        fp = os.path.join(sm_dir, fname)
        if os.path.exists(fp):
            datasets.append(('T2Dv2', fp))
    
    if not datasets:
        print("No SM datasets found, skipping")
        return []
    
    results = []
    print(f"\n{'='*60}\nSM Experiments: {len(datasets)} datasets\n{'='*60}")
    
    for ds_name, ds_path in datasets:
        items = load_json(ds_path, CONFIG['max_pairs'])
        print(f"\n--- {ds_name}: {len(items)} items ---")
        
        for method in ['llm_dp', 'rematch', 'matchmaker']:
            for use_evorm in [False, True]:
                try:
                    r = run_schema_matching(ds_path, method=method, use_evorm=use_evorm, max_items=CONFIG['max_pairs'])
                    r['dataset'] = ds_name; r['task'] = 'SM'
                    results.append(r)
                    ev = '+EvoRM' if use_evorm else ''
                    print(f"  {method}{ev}: F1={r.get('F1', 'N/A')}")
                except Exception as e:
                    print(f"  {method}{'+EvoRM' if use_evorm else ''}: ERROR - {e}")
    
    return results

# ==============================================================================
# Ablation: ER + SM
# ==============================================================================

def run_ablation():
    """Run ablation studies on ER and SM tasks."""
    from baselines.evorm_wrappers.matchgpt import run_matchgpt
    
    results = []
    print(f"\n{'='*60}\nAblation Experiments\n{'='*60}")
    
    # ER ablation on Beer (hardest dataset)
    ablation_modes = {
        'full': None,
        'no_stage1': 'no_stage1',
        'no_maintenance': 'no_maintenance',
        'no_mlp_gate': 'no_mlp_gate',
    }
    
    ds_path = os.path.join(CONFIG['data_dir'], 'er', 'Structured', 'Beer')
    if os.path.exists(ds_path):
        pairs = load_dm_dataset(ds_path, CONFIG['max_pairs'])
        tmp_path = '/tmp/er_beer_ablation.json'
        with open(tmp_path, 'w') as f:
            json.dump([{'record_a': p[0], 'record_b': p[1], 'label': p[2]} for p in pairs], f)
        
        print(f"\n--- ER Beer Ablation: {len(pairs)} pairs ---")
        for mode_name, ablation_mode in ablation_modes.items():
            try:
                r = run_matchgpt(tmp_path, backend='default', use_evorm=(mode_name != 'full' or True),
                                evorm_ablation_mode=ablation_mode, max_pairs=CONFIG['max_pairs'])
                r['dataset'] = 'Beer'; r['task'] = 'ER_Ablation'; r['ablation'] = mode_name
                results.append(r)
                print(f"  {mode_name}: F1={r.get('F1', 'N/A')}")
            except Exception as e:
                print(f"  {mode_name}: ERROR - {e}")
    
    # SM ablation on T2Dv2
    from baselines.evorm_wrappers.lela_el import run_schema_matching
    sm_path = os.path.join(CONFIG['data_dir'], 'sm', 'T2Dv2_schema.json')
    if os.path.exists(sm_path):
        print(f"\n--- SM T2Dv2 Ablation ---")
        for mode_name, ablation_mode in ablation_modes.items():
            try:
                use_evorm = (ablation_mode is None)  # 'full' mode uses EvoRM
                r = run_schema_matching(sm_path, method='llm_dp', use_evorm=True if mode_name == 'full' else True,
                                        evorm_ablation_mode=ablation_mode, max_items=CONFIG['max_pairs'])
                r['dataset'] = 'T2Dv2'; r['task'] = 'SM_Ablation'; r['ablation'] = mode_name
                results.append(r)
                print(f"  {mode_name}: F1={r.get('F1', 'N/A')}")
            except Exception as e:
                print(f"  {mode_name}: ERROR - {e}")
    
    return results

# ==============================================================================
# Main
# ==============================================================================

def main():
    print("=" * 70)
    print("EvoRM Complete Experiment Suite — ALL Paper Datasets + Ablation")
    print("=" * 70)
    print(f"Model: {CONFIG['model']}, Max pairs: {CONFIG['max_pairs']}")
    print(f"MLP Gate warmup: {CONFIG['mlp_n_warmup']} samples")
    
    all_results = []
    
    for name, func in [('ER', run_er), ('EA', run_ea), ('EL', run_el), ('SM', run_sm), ('Ablation', run_ablation)]:
        print(f"\n{'='*70}\nPART: {name}\n{'='*70}")
        try:
            r = func()
            all_results.extend(r)
            print(f"\n{name} completed: {len(r)} experiments")
        except Exception as e:
            print(f"{name} failed: {e}")
            import traceback; traceback.print_exc()
    
    timestamp = time.strftime('%Y%m%d_%H%M%S')
    results_path = os.path.join(CONFIG['results_dir'], f'complete_results_{timestamp}.json')
    with open(results_path, 'w') as f:
        json.dump(all_results, f, indent=2, ensure_ascii=False)
    
    print(f"\n{'='*70}")
    print(f"ALL DONE: {len(all_results)} experiments")
    print(f"Results: {results_path}")
    print(f"{'='*70}")
    
    # Print summary by task
    for task in ['ER', 'EA', 'EL', 'SM', 'ER_Ablation', 'SM_Ablation']:
        tr = [r for r in all_results if r.get('task') == task]
        if not tr:
            continue
        print(f"\n{task} ({len(tr)} experiments):")
        for r in tr:
            method = r.get('method', '?')
            ds = r.get('dataset', '?')
            abl = r.get('ablation', '')
            f1 = r.get('F1', r.get('Accuracy', 'N/A'))
            label = f"{ds} | {method}"
            if abl:
                label += f" [{abl}]"
            print(f"  {label}: {f1}")
    
    # Create done flag
    with open('/root/autodl-tmp/EXPERIMENT_COMPLETE', 'w') as f:
        f.write(f"Done at {timestamp}\n{len(all_results)} experiments\n")
    
    return all_results

if __name__ == '__main__':
    main()