#!/usr/bin/env python3
"""
EvoRM Complete Experiment Suite — Real Datasets
================================================
Runs ALL experiments from the TKDE EvoRM paper on REAL datasets.
"""

import os, sys, json, csv, time, hashlib, random, threading
from collections import defaultdict
from typing import Dict, List, Tuple, Optional

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_per_dataset': 100,
    'max_ea_pairs': 100,
    'data_dir': '/root/autodl-tmp/AdaCoAgentEA/data',
    'results_dir': '/root/autodl-tmp/AdaCoAgentEA/results',
    'evorm_persistence': '/tmp/evorm_experiments',
}

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

def load_deepmatcher_dataset(ds_path, max_pairs=None):
    """Load ER dataset in DeepMatcher format."""
    table_a, table_b = {}, {}
    for tbl_name, tbl_dict in [('tableA.csv', table_a), ('tableB.csv', table_b)]:
        tbl_path = os.path.join(ds_path, tbl_name)
        if not os.path.exists(tbl_path):
            tbl_path = tbl_path.replace('.csv', '_tok.csv')
        if os.path.exists(tbl_path):
            with open(tbl_path, 'r', encoding='utf-8', errors='replace') as f:
                reader = csv.DictReader(f)
                for row in reader:
                    tbl_dict[row.get('id', '')] = dict(row)
    
    pairs = []
    for split in ['test.csv', 'valid.csv']:
        split_path = os.path.join(ds_path, split)
        if not os.path.exists(split_path):
            continue
        with open(split_path, 'r', encoding='utf-8', errors='replace') as f:
            reader = csv.DictReader(f)
            for row in reader:
                left_id = row.get('ltable_id', '')
                right_id = row.get('rtable_id', '')
                pairs.append({
                    'record_a': table_a.get(left_id, {'id': left_id}),
                    'record_b': table_b.get(right_id, {'id': right_id}),
                    'label': int(row.get('label', 0)),
                })
        if pairs:
            break
    
    if max_pairs and 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)
        n_neg = min(len(neg), max_pairs - n_pos)
        pairs = random.sample(pos, n_pos) + random.sample(neg, n_neg)
        random.shuffle(pairs)
    
    return pairs

def load_dbp15k_dataset(data_dir, lang_pair, max_pairs=None):
    """Load DBP15K EA dataset."""
    pair_dir = os.path.join(data_dir, lang_pair)
    if not os.path.exists(pair_dir):
        return []
    
    def load_labels(path):
        labels = {}
        with open(path, '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')
    with open(ill_path, 'r', encoding='utf-8') as f:
        for line in f:
            parts = line.strip().split('\t')
            if len(parts) >= 2:
                s_uri, t_uri = parts[0], parts[1]
                pairs.append({
                    'source': {'name': s_labels.get(s_uri, s_uri.split('/')[-1]), 'uri': s_uri},
                    'target': {'name': t_labels.get(t_uri, t_uri.split('/')[-1]), 'uri': t_uri},
                    'label': 1,
                })
    
    # Add negative samples
    random.seed(42)
    s_uris = list(s_labels.keys())
    t_uris = list(t_labels.keys())
    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]))
    
    neg_pairs = []
    for _ in range(len(pairs)):
        s_uri = random.choice(s_uris)
        t_uri = random.choice(t_uris)
        if (s_uri, t_uri) not in ill_set:
            neg_pairs.append({
                'source': {'name': s_labels.get(s_uri, s_uri.split('/')[-1]), 'uri': s_uri},
                'target': {'name': t_labels.get(t_uri, t_uri.split('/')[-1]), 'uri': t_uri},
                'label': 0,
            })
        if len(neg_pairs) >= len(pairs):
            break
    
    all_pairs = pairs + neg_pairs[:len(pairs)]
    random.shuffle(all_pairs)
    
    if max_pairs and len(all_pairs) > max_pairs:
        all_pairs = all_pairs[:max_pairs]
    
    return all_pairs

def compute_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)}

def run_er_experiments():
    """Run ER experiments on DeepMatcher datasets."""
    from baselines.evorm_wrappers.matchgpt import run_matchgpt, run_anymatch
    
    er_dir = os.path.join(CONFIG['data_dir'], 'er')
    results = []
    
    datasets = []
    for category in ['Structured', 'Textual', 'Dirty']:
        cat_path = os.path.join(er_dir, category)
        if not os.path.exists(cat_path):
            continue
        for ds_name in sorted(os.listdir(cat_path)):
            ds_path = os.path.join(cat_path, ds_name)
            if os.path.isdir(ds_path):
                has_test = os.path.exists(os.path.join(ds_path, 'test.csv')) or os.path.exists(os.path.join(ds_path, 'valid.csv'))
                if has_test:
                    datasets.append((f"{ds_name}", ds_path, category))
    
    print(f"\n{'='*60}")
    print(f"ER Experiments: {len(datasets)} datasets")
    print(f"{'='*60}")
    
    for ds_name, ds_path, category in datasets:
        print(f"\n--- {ds_name} ({category}) ---")
        pairs = load_deepmatcher_dataset(ds_path, CONFIG['max_pairs_per_dataset'])
        if not pairs:
            print(f"  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"  Loaded {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_per_dataset'])
                    r['dataset'] = ds_name
                    r['category'] = category
                    r['task'] = 'ER'
                    results.append(r)
                    ev_tag = '+EvoRM' if use_evorm else ''
                    print(f"  MatchGPT[{backend}]{ev_tag}: 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_per_dataset'])
                r['dataset'] = ds_name
                r['category'] = category
                r['task'] = 'ER'
                results.append(r)
                ev_tag = '+EvoRM' if use_evorm else ''
                print(f"  AnyMatch{ev_tag}: 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_experiments():
    """Run EL experiments."""
    from baselines.evorm_wrappers.lela_el import run_lela
    
    el_dir = os.path.join(CONFIG["data_dir"], "el")
    if not os.path.exists(el_dir):
        el_dir = os.path.join(CONFIG['data_dir'], 'el')
    
    results = []
    el_datasets = []
    if os.path.exists(el_dir):
        for fname in sorted(os.listdir(el_dir)):
            if fname.endswith('.json'):
                ds_path = os.path.join(el_dir, fname)
                try:
                    with open(ds_path) as f:
                        data = json.load(f)
                    if isinstance(data, list) and len(data) > 0:
                        el_datasets.append((fname.replace('.json', ''), ds_path, len(data)))
                except:
                    pass
    
    if not el_datasets:
        print("No EL datasets found, skipping")
        return results
    
    print(f"\n{'='*60}")
    print(f"EL Experiments: {len(el_datasets)} datasets")
    print(f"{'='*60}")
    
    for ds_name, ds_path, n_items in el_datasets:
        print(f"\n--- {ds_name} ({n_items} items) ---")
        max_n = min(CONFIG['max_pairs_per_dataset'], n_items)
        
        for use_evorm in [False, True]:
            try:
                r = run_lela(ds_path, use_evorm=use_evorm, max_items=max_n)
                r['dataset'] = ds_name
                r['task'] = 'EL'
                results.append(r)
                ev_tag = '+EvoRM' if use_evorm else ''
                print(f"  LELA{ev_tag}: 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_experiments():
    """Run SM experiments."""
    from baselines.evorm_wrappers.lela_el import run_schema_matching
    
    sm_dir = os.path.join(CONFIG['data_dir'], 'sm')
    if not os.path.exists(sm_dir):
        print("No SM datasets found, skipping")
        return []
    
    results = []
    sm_datasets = []
    for fname in sorted(os.listdir(sm_dir)):
        if fname.endswith('.json'):
            ds_path = os.path.join(sm_dir, fname)
            try:
                with open(ds_path) as f:
                    data = json.load(f)
                if isinstance(data, list) and len(data) > 0:
                    sm_datasets.append((fname.replace('.json', ''), ds_path, len(data)))
            except:
                pass
    
    if not sm_datasets:
        print("No SM datasets found, skipping")
        return results
    
    print(f"\n{'='*60}")
    print(f"SM Experiments: {len(sm_datasets)} datasets")
    print(f"{'='*60}")
    
    for ds_name, ds_path, n_items in sm_datasets:
        print(f"\n--- {ds_name} ({n_items} items) ---")
        max_n = min(CONFIG['max_pairs_per_dataset'], n_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=max_n)
                    r['dataset'] = ds_name
                    r['task'] = 'SM'
                    results.append(r)
                    ev_tag = '+EvoRM' if use_evorm else ''
                    print(f"  {method}{ev_tag}: 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_experiments():
    """Run EA experiments on DBP15K."""
    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}")
    print(f"EA Experiments: {len(lang_pairs)} language pairs")
    print(f"{'='*60}")
    
    client = OpenAI(base_url=CONFIG['api_base'], api_key=CONFIG['api_key'])
    
    for lang_pair in lang_pairs:
        print(f"\n--- DBP15K {lang_pair} ---")
        pairs = load_dbp15k_dataset(ea_dir, lang_pair, CONFIG['max_ea_pairs'])
        if not pairs:
            print(f"  No data, skipping")
            continue
        
        pos_count = sum(1 for p in pairs if p['label'] == 1)
        print(f"  Loaded {len(pairs)} pairs (pos={pos_count})")
        
        for use_evorm in [False, True]:
            evorm = None
            if use_evorm:
                evorm = EvoRMPlugin(client=client, persistence_dir=f'{CONFIG["evorm_persistence"]}/ea_{lang_pair}')
            
            predictions, labels_list = [], []
            lock = threading.Lock()
            tokens_cal.global_tokens = 0
            t0 = time.time()
            
            def process_ea(item):
                s = item['source']
                t = item['target']
                label = 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', '')
                        if raw:
                            raw_upper = raw.strip().upper()
                            pred = 1 if 'MATCH' in raw_upper and 'NON' not in raw_upper else 0
                        else:
                            pred = 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
            metrics = compute_metrics(predictions, labels_list)
            metrics['dataset'] = f'DBP15K {lang_pair}'
            metrics['task'] = 'EA'
            metrics['method'] = f"EA{' + EvoRM' if use_evorm else ''}"
            metrics['time_seconds'] = round(elapsed, 1)
            metrics['tokens'] = tokens_cal.global_tokens
            if evorm:
                metrics['evorm_stats'] = evorm.get_stats()
            results.append(metrics)
            ev_tag = '+EvoRM' if use_evorm else ''
            print(f"  EA{ev_tag}: F1={metrics.get('F1', 'N/A')}, Prec={metrics.get('Precision', 'N/A')}, Rec={metrics.get('Recall', 'N/A')}")
    
    return results

def main():
    print("=" * 70)
    print("EvoRM Complete Experiment Suite — REAL Datasets")
    print("=" * 70)
    print(f"Model: {CONFIG['model']}")
    print(f"Max pairs per dataset: {CONFIG['max_pairs_per_dataset']}")
    
    all_results = []
    
    # 1. ER
    print("\n" + "=" * 70)
    print("PART 1: Entity Resolution (ER)")
    print("=" * 70)
    try:
        er_results = run_er_experiments()
        all_results.extend(er_results)
        print(f"\nER completed: {len(er_results)} experiments")
    except Exception as e:
        print(f"ER failed: {e}")
        import traceback; traceback.print_exc()
    
    # 2. EA
    print("\n" + "=" * 70)
    print("PART 2: Entity Alignment (EA)")
    print("=" * 70)
    try:
        ea_results = run_ea_experiments()
        all_results.extend(ea_results)
        print(f"\nEA completed: {len(ea_results)} experiments")
    except Exception as e:
        print(f"EA failed: {e}")
        import traceback; traceback.print_exc()
    
    # 3. EL
    print("\n" + "=" * 70)
    print("PART 3: Entity Linking (EL)")
    print("=" * 70)
    try:
        el_results = run_el_experiments()
        all_results.extend(el_results)
        print(f"\nEL completed: {len(el_results)} experiments")
    except Exception as e:
        print(f"EL failed: {e}")
        import traceback; traceback.print_exc()
    
    # 4. SM
    print("\n" + "=" * 70)
    print("PART 4: Schema Matching (SM)")
    print("=" * 70)
    try:
        sm_results = run_sm_experiments()
        all_results.extend(sm_results)
        print(f"\nSM completed: {len(sm_results)} experiments")
    except Exception as e:
        print(f"SM failed: {e}")
        import traceback; traceback.print_exc()
    
    # Save results
    timestamp = time.strftime('%Y%m%d_%H%M%S')
    results_path = os.path.join(CONFIG['results_dir'], f'real_data_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}")
    
    # Summary
    for task in ['ER', 'EA', 'EL', 'SM']:
        task_results = [r for r in all_results if r.get('task') == task]
        if not task_results:
            continue
        print(f"\n{task} ({len(task_results)} experiments):")
        for r in task_results:
            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()