#!/usr/bin/env python3
"""Quick experiment run on real datasets — 50 pairs, key datasets only."""
import os, sys, json, csv, time, hashlib, random, threading, re
sys.path.insert(0, '/root/autodl-tmp/AdaCoAgentEA')

CONFIG = {
    'api_base': os.environ.get("OPENAI_BASE_URL", "https://api.openai.com/v1"),
    'api_key': "",
    'model': "gpt-3.5-turbo-1106",
    'max_workers': 5,
    'max_pairs': 50,
    'data_dir': '/root/autodl-tmp/AdaCoAgentEA/data',
    'results_dir': '/root/autodl-tmp/AdaCoAgentEA/results',
    'evorm_persistence': '/tmp/evorm_quick',
}

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

# ===== DATA LOADERS =====
def load_dm_dataset(ds_path, max_pairs=50):
    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):
                pairs.append({
                    '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', '')}),
                    'label': int(row.get('label', 0)),
                })
        if pairs: break
    if len(pairs) > max_pairs:
        random.seed(42)
        pos = [p for p in pairs if p['label'] == 1]
        neg = [p for p in pairs if p['label'] == 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_dbp15k(data_dir, lang_pair, max_pairs=50):
    pair_dir = os.path.join(data_dir, lang_pair)
    if not os.path.exists(pair_dir): return []
    def load_labels(p):
        labels = {}
        with open(p, 'r', encoding='utf-8') as f:
            for line in f:
                parts = line.strip().split('\t')
                if len(parts) >= 2: labels[parts[0]] = parts[1]
        return labels
    s_labels = load_labels(os.path.join(pair_dir, 's_labels'))
    t_labels = load_labels(os.path.join(pair_dir, 't_labels'))
    pairs = []
    ill_path = os.path.join(pair_dir, 'ent_ILLs')
    ill_set = set()
    with open(ill_path, 'r', encoding='utf-8') as f:
        for line in f:
            parts = line.strip().split('\t')
            if len(parts) >= 2:
                ill_set.add((parts[0], parts[1]))
                pairs.append({
                    'source': {'name': s_labels.get(parts[0], parts[0].split('/')[-1]), 'uri': parts[0]},
                    'target': {'name': t_labels.get(parts[1], parts[1].split('/')[-1]), 'uri': parts[1]},
                    'label': 1,
                })
    random.seed(42)
    s_uris, t_uris = list(s_labels.keys()), list(t_labels.keys())
    neg_pairs = []
    for _ in range(len(pairs)):
        su, tu = random.choice(s_uris), random.choice(t_uris)
        if (su, tu) not in ill_set:
            neg_pairs.append({
                'source': {'name': s_labels.get(su, su.split('/')[-1]), 'uri': su},
                'target': {'name': t_labels.get(tu, tu.split('/')[-1]), 'uri': tu},
                'label': 0,
            })
        if len(neg_pairs) >= len(pairs): break
    all_pairs = pairs + neg_pairs[:len(pairs)]
    random.shuffle(all_pairs)
    if len(all_pairs) > max_pairs: all_pairs = all_pairs[:max_pairs]
    return all_pairs

def load_json(path, max_items=50):
    if not os.path.exists(path): return []
    with open(path, 'r', encoding='utf-8') as f:
        try:
            data = json.load(f)
            if isinstance(data, list):
                return data[:max_items] if len(data) > max_items else data
        except: pass
    return []

def 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),
            'TP': tp, 'FP': fp, 'FN': fn, 'Total': len(predictions)}

# ===== EXPERIMENT RUNNERS =====
def run_er():
    """Run ER experiments on key datasets."""
    from baselines.evorm_wrappers.matchgpt import run_matchgpt, run_anymatch
    
    datasets = [
        ("Amazon-Google", "Structured"),
        ("Beer", "Structured"),
        ("DBLP-ACM", "Structured"),
        ("Fodors-Zagats", "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
        
        tmp_path = f'/tmp/er_{hashlib.md5(ds_name.encode()).hexdigest()[:8]}.json'
        with open(tmp_path, 'w') as f: json.dump(pairs, f)
        
        pos_count = sum(1 for p in pairs if p['label'] == 1)
        print(f"\n--- {ds_name} ({category}): {len(pairs)} pairs (pos={pos_count}) ---")
        
        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}")
        
        os.remove(tmp_path)
    
    return results

def run_el():
    """Run EL experiments."""
    from baselines.evorm_wrappers.lela_el import run_lela
    
    el_dir = os.path.join(CONFIG['data_dir'], 'el')
    results = []
    
    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

def run_sm():
    """Run SM experiments."""
    from baselines.evorm_wrappers.lela_el import run_schema_matching
    
    sm_dir = os.path.join(CONFIG['data_dir'], 'sm')
    results = []
    
    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

def run_ea():
    """Run EA experiments."""
    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 = []
    
    if not os.path.exists(ea_dir):
        print("No DBP15K data found")
        return 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=CONFIG['api_base'], api_key=CONFIG['api_key'])
    
    for lang_pair in lang_pairs:
        pairs = load_dbp15k(ea_dir, lang_pair, CONFIG['max_pairs'])
        if not pairs:
            print(f"  {lang_pair}: no data, skipping")
            continue
        
        pos_count = sum(1 for p in pairs if p['label'] == 1)
        print(f"\n--- DBP15K {lang_pair}: {len(pairs)} pairs (pos={pos_count}) ---")
        
        for use_evorm in [False, True]:
            evorm = EvoRMPlugin(client=client, persistence_dir=f'{CONFIG["evorm_persistence"]}/ea_{lang_pair}') if use_evorm else None
            predictions, labels_list = [], []
            lock = threading.Lock()
            tokens_cal.global_tokens = 0
            t0 = time.time()
            
            def process_ea(item):
                s, t, label = item['source'], item['target'], item['label']
                prompt = f"""Entity Alignment: Determine if these two entities refer to the same real-world entity.

Source Entity: {s.get('name', '')}
Target Entity: {t.get('name', '')}

Are these the same entity? Answer ONLY "MATCH" or "NON-MATCH".
Your answer:"""
                
                if evorm:
                    es_ctx = {'entity_name': s.get('name', ''), 'uri': s.get('uri', '')}
                    et_ctx = {'entity_name': t.get('name', ''), 'uri': t.get('uri', '')}
                    eid = s.get('name', '')[:20]; tid = t.get('name', '')[:20]
                    decision, triggered, cands = evorm.stage1(eid, tid, es_ctx, et_ctx)
                    if decision is not None:
                        pred = decision
                    else:
                        result = evorm.stage2(eid, tid, es_ctx, et_ctx, triggered, cands, prompt)
                        raw = result.get('raw_response', '')
                        pred = 1 if raw and 'MATCH' in raw.upper() and 'NON' not in raw.upper() else (result.get('decision', 0) or 0)
                        evorm.record_trajectory(eid, tid, es_ctx, et_ctx, decision=pred,
                                               rationale=result.get('rationale', ''),
                                               triggered_rules=triggered,
                                               rule_feedback=result.get('rule_feedback', {}))
                else:
                    response = client.chat.completions.create(
                        model=CONFIG['model'], messages=[{'role': 'user', 'content': prompt}], temperature=0.1)
                    answer = response.choices[0].message.content.strip().upper()
                    tokens_cal.update_add_var(response.usage.total_tokens)
                    pred = 1 if 'MATCH' in answer and 'NON' not in answer else 0
                
                with lock:
                    predictions.append(pred)
                    labels_list.append(label)
            
            executor = ThreadPoolExecutor(max_workers=CONFIG['max_workers'])
            for item in tqdm(pairs, desc=f"EA {lang_pair}{' +EvoRM' if use_evorm else ''}"):
                executor.submit(process_ea, item)
            executor.shutdown(wait=True)
            
            elapsed = time.time() - t0
            m = metrics(predictions, labels_list)
            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')}")
    
    return results

def main():
    print("=" * 70)
    print("EvoRM Quick Experiment Suite — REAL Datasets")
    print("=" * 70)
    print(f"Model: {CONFIG['model']}, Max pairs: {CONFIG['max_pairs']}")
    
    all_results = []
    
    for name, func in [('ER', run_er), ('EA', run_ea), ('EL', run_el), ('SM', run_sm)]:
        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'quick_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}")
    
    for task in ['ER', 'EA', 'EL', 'SM']:
        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', '?')
            f1 = r.get('F1', r.get('Accuracy', 'N/A'))
            print(f"  {ds} | {method}: {f1}")
    
    return all_results

if __name__ == '__main__':
    main()