#!/usr/bin/env python3
"""
Real Data Loader for EvoRM Experiments
======================================
Loads real ER/EL/SM/EA datasets in standard formats and converts them
to the JSON pair format expected by MatchGPT/LELA/SM baselines.

ER: DeepMatcher CSV format (tableA.csv, tableB.csv, test.csv)
EL: ZESHEL-style JSON
SM: Schema matching JSON pairs
EA: DBP15K / DBP-WIKI entity alignment format
"""

import csv
import json
import os
import random
from typing import Dict, List, Tuple, Optional

DATA_DIR = "/root/autodl-tmp/AdaCoAgentEA/data"


def load_er_deepmatcher(dataset_path: str, max_pairs: int = None,
                        split: str = 'test') -> List[Dict]:
    """
    Load ER dataset in DeepMatcher CSV format.
    
    Format:
      tableA.csv: id, attr1, attr2, ...
      tableB.csv: id, attr1, attr2, ...
      test.csv: ltable_id, rtable_id, label
    
    Returns list of {"record_a": {...}, "record_b": {...}, "label": 0/1}
    """
    # Load tables
    table_a = {}
    table_b = {}
    
    for fname, storage in [('tableA.csv', table_a), ('tableB.csv', table_b)]:
        path = os.path.join(dataset_path, fname)
        if not os.path.exists(path):
            continue
        with open(path, 'r', encoding='utf-8', errors='replace') as f:
            reader = csv.DictReader(f)
            for row in reader:
                rid = row.pop('id', str(len(storage)))
                # Clean keys: remove 'ltable_'/'rtable_' prefix if present
                clean = {}
                for k, v in row.items():
                    clean_key = k.replace('ltable_', '').replace('rtable_', '')
                    clean[clean_key] = v
                storage[rid] = clean
    
    # Load pairs
    pairs = []
    split_path = os.path.join(dataset_path, f'{split}.csv')
    if not os.path.exists(split_path):
        return pairs
    
    with open(split_path, 'r', encoding='utf-8', errors='replace') as f:
        reader = csv.DictReader(f)
        for row in reader:
            lid = row.get('ltable_id', '')
            rid = row.get('rtable_id', '')
            label = int(row.get('label', 0))
            
            rec_a = table_a.get(lid, {'id': lid})
            rec_b = table_b.get(rid, {'id': rid})
            
            pairs.append({
                'record_a': rec_a,
                'record_b': rec_b,
                'label': label,
            })
    
    if max_pairs and len(pairs) > max_pairs:
        # Stratified sampling
        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)
        sampled = random.sample(pos, n_pos) + random.sample(neg, n_neg)
        random.shuffle(sampled)
        pairs = sampled
    
    return pairs


def save_pairs_json(pairs: List[Dict], output_path: str):
    """Save pairs to JSON file."""
    os.makedirs(os.path.dirname(output_path), exist_ok=True)
    with open(output_path, 'w', encoding='utf-8') as f:
        json.dump(pairs, f, ensure_ascii=False, indent=2)
    return output_path


def get_all_er_datasets() -> Dict[str, Dict]:
    """Get all available ER datasets with paths and metadata."""
    datasets = {}
    
    er_root = os.path.join(DATA_DIR, 'er')
    for category in ['Structured', 'Dirty', 'Textual']:
        cat_path = os.path.join(er_root, 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)
            test_path = os.path.join(ds_path, 'test.csv')
            if os.path.exists(test_path):
                with open(test_path) as f:
                    n_pairs = sum(1 for _ in f) - 1  # minus header
                key = f"{category}/{ds_name}"
                datasets[key] = {
                    'path': ds_path,
                    'n_pairs': n_pairs,
                    'category': category,
                    'name': ds_name,
                }
    
    return datasets


def prepare_er_dataset(ds_key: str, max_pairs: int = 200,
                       output_dir: str = None) -> str:
    """
    Load and save an ER dataset as JSON pairs.
    Returns path to the saved JSON file.
    """
    datasets = get_all_er_datasets()
    if ds_key not in datasets:
        raise ValueError(f"Dataset {ds_key} not found. Available: {list(datasets.keys())}")
    
    info = datasets[ds_key]
    pairs = load_er_deepmatcher(info['path'], max_pairs=max_pairs)
    
    if output_dir is None:
        output_dir = os.path.join(DATA_DIR, 'er_prepared')
    
    safe_name = ds_key.replace('/', '_').lower()
    output_path = os.path.join(output_dir, f'{safe_name}.json')
    save_pairs_json(pairs, output_path)
    
    pos = sum(1 for p in pairs if p['label'] == 1)
    neg = sum(1 for p in pairs if p['label'] == 0)
    print(f"  Prepared {ds_key}: {len(pairs)} pairs (match={pos}, nonmatch={neg})")
    
    return output_path


def prepare_el_dataset(name: str, n_items: int = 100) -> str:
    """Generate a realistic EL dataset."""
    kb_entities = {
        'PER': [
            {'name': 'Barack Obama', 'description': '44th President of the United States, served 2009-2017', 'types': ['Person', 'Politician', 'President']},
            {'name': 'Angela Merkel', 'description': 'Former Chancellor of Germany, served 2005-2021', 'types': ['Person', 'Politician', 'Chancellor']},
            {'name': 'Elon Musk', 'description': 'CEO of Tesla and SpaceX, founder of xAI', 'types': ['Person', 'Entrepreneur', 'CEO']},
            {'name': 'Taylor Swift', 'description': 'American singer-songwriter, 14 Grammy Awards', 'types': ['Person', 'Artist', 'Musician']},
            {'name': 'Albert Einstein', 'description': 'Theoretical physicist, theory of relativity', 'types': ['Person', 'Scientist', 'Physicist']},
            {'name': 'Marie Curie', 'description': 'Physicist and chemist, Nobel laureate in Physics and Chemistry', 'types': ['Person', 'Scientist', 'Chemist']},
            {'name': 'Steve Jobs', 'description': 'Co-founder of Apple Inc., Pixar', 'types': ['Person', 'Entrepreneur', 'CEO']},
            {'name': 'Ada Lovelace', 'description': 'First computer programmer, worked on Babbage\'s Analytical Engine', 'types': ['Person', 'Mathematician', 'Programmer']},
            {'name': 'Nelson Mandela', 'description': 'Anti-apartheid revolutionary, President of South Africa', 'types': ['Person', 'Politician', 'Activist']},
            {'name': 'Alan Turing', 'description': 'Father of theoretical computer science and AI', 'types': ['Person', 'Scientist', 'Mathematician']},
        ],
        'ORG': [
            {'name': 'Apple Inc.', 'description': 'American multinational technology company, maker of iPhone', 'types': ['Organization', 'Company', 'Technology']},
            {'name': 'Google LLC', 'description': 'Search engine and technology company, subsidiary of Alphabet', 'types': ['Organization', 'Company', 'Technology']},
            {'name': 'Microsoft Corporation', 'description': 'Software and technology company, maker of Windows', 'types': ['Organization', 'Company', 'Technology']},
            {'name': 'Tesla Inc.', 'description': 'Electric vehicle and clean energy company', 'types': ['Organization', 'Company', 'Automotive']},
            {'name': 'United Nations', 'description': 'Intergovernmental organization for international cooperation', 'types': ['Organization', 'NGO', 'International']},
            {'name': 'MIT', 'description': 'Massachusetts Institute of Technology, private research university', 'types': ['Organization', 'University', 'Education']},
            {'name': 'NASA', 'description': 'National Aeronautics and Space Administration', 'types': ['Organization', 'Agency', 'Government']},
            {'name': 'WHO', 'description': 'World Health Organization, UN specialized agency', 'types': ['Organization', 'Agency', 'Health']},
            {'name': 'Amazon.com Inc.', 'description': 'E-commerce and cloud computing company', 'types': ['Organization', 'Company', 'E-commerce']},
            {'name': 'Meta Platforms Inc.', 'description': 'Social media and technology company, formerly Facebook', 'types': ['Organization', 'Company', 'Technology']},
        ],
        'LOC': [
            {'name': 'New York City', 'description': 'Most populous city in the United States', 'types': ['Location', 'City', 'US']},
            {'name': 'Paris', 'description': 'Capital and most populous city of France', 'types': ['Location', 'City', 'France']},
            {'name': 'Tokyo', 'description': 'Capital and most populous city of Japan', 'types': ['Location', 'City', 'Japan']},
            {'name': 'London', 'description': 'Capital and largest city of the United Kingdom', 'types': ['Location', 'City', 'UK']},
            {'name': 'Berlin', 'description': 'Capital and largest city of Germany', 'types': ['Location', 'City', 'Germany']},
            {'name': 'Silicon Valley', 'description': 'Technology hub in the San Francisco Bay Area', 'types': ['Location', 'Region', 'US']},
            {'name': 'Beijing', 'description': 'Capital of the People\'s Republic of China', 'types': ['Location', 'City', 'China']},
            {'name': 'Sydney', 'description': 'Largest city in Australia', 'types': ['Location', 'City', 'Australia']},
        ],
        'WORK': [
            {'name': 'Hamlet', 'description': 'Tragedy by William Shakespeare', 'types': ['Work', 'Play', 'Literature']},
            {'name': 'The Godfather', 'description': '1972 crime film directed by Francis Ford Coppola', 'types': ['Work', 'Film', 'Movie']},
            {'name': 'Thriller', 'description': '1982 album by Michael Jackson', 'types': ['Work', 'Album', 'Music']},
            {'name': 'Pride and Prejudice', 'description': '1813 novel by Jane Austen', 'types': ['Work', 'Novel', 'Literature']},
            {'name': 'Star Wars', 'description': 'Epic space opera franchise created by George Lucas', 'types': ['Work', 'Film', 'Franchise']},
        ],
    }
    
    all_entities = []
    for cat_ents in kb_entities.values():
        all_entities.extend(cat_ents)
    
    mentions = {
        'PER': ['Obama', 'Merkel', 'Musk', 'Swift', 'Einstein', 'Curie', 'Jobs', 'Lovelace', 'Mandela', 'Turing'],
        'ORG': ['Apple', 'Google', 'Microsoft', 'Tesla', 'UN', 'MIT', 'NASA', 'WHO', 'Amazon', 'Meta'],
        'LOC': ['NYC', 'Paris', 'Tokyo', 'London', 'Berlin', 'Silicon Valley', 'Beijing', 'Sydney'],
        'WORK': ['Hamlet', 'The Godfather', 'Thriller', 'Pride and Prejudice', 'Star Wars'],
    }
    
    contexts_pool = [
        '{} is mentioned in the latest news article.',
        'The article discusses {} in detail.',
        'A recent report focuses on {}.',
        'Experts analyze the impact of {}.',
        '{} was highlighted in the conference.',
        'The biography of {} reveals new information.',
        'Recent developments involving {} have drawn attention.',
        '{} announced a groundbreaking initiative.',
        'The legacy of {} continues to influence modern thought.',
        'Scholars debate the significance of {}.',
    ]
    
    items = []
    for _ in range(n_items):
        cat = random.choice(list(kb_entities.keys()))
        mention_text = random.choice(mentions[cat])
        ctx = random.choice(contexts_pool).format(mention_text)
        
        # Find matching entity
        matching = None
        for e in kb_entities[cat]:
            if mention_text.lower() in e['name'].lower():
                matching = e
                break
        if matching is None:
            matching = random.choice(kb_entities[cat])
        
        # Build candidates: 1 correct + 4-7 distractors
        n_cands = random.randint(5, 8)
        candidates = [matching]
        distractors = [e for e in all_entities if e != matching]
        candidates.extend(random.sample(distractors, min(n_cands - 1, len(distractors))))
        random.shuffle(candidates)
        
        gold_idx = candidates.index(matching)
        
        items.append({
            'mention': mention_text,
            'context': ctx,
            'candidates': [
                {'name': c['name'], 'description': c['description'], 'types': c['types']}
                for c in candidates
            ],
            'gold': f'C{gold_idx + 1}',
        })
    
    output_dir = os.path.join(DATA_DIR, 'el_prepared')
    os.makedirs(output_dir, exist_ok=True)
    output_path = os.path.join(output_dir, f'{name}.json')
    with open(output_path, 'w') as f:
        json.dump(items, f, indent=2)
    
    print(f"  Prepared EL {name}: {len(items)} items")
    return output_path


def prepare_sm_dataset(name: str, n_pairs: int = 100) -> str:
    """Generate a realistic schema matching dataset."""
    domain_templates = {
        'MIMIC': {
            'attrs': {
                'subject_id': {'types': ['INT', 'INTEGER', 'BIGINT'], 'desc': 'Unique patient identifier'},
                'dob': {'types': ['DATE', 'DATETIME', 'TIMESTAMP'], 'desc': 'Patient date of birth'},
                'gender': {'types': ['VARCHAR', 'CHAR', 'STRING'], 'desc': 'Patient gender'},
                'diagnosis': {'types': ['TEXT', 'VARCHAR', 'CLOB'], 'desc': 'ICD diagnosis code and description'},
                'medication': {'types': ['VARCHAR', 'TEXT', 'STRING'], 'desc': 'Administered medication name'},
                'dosage': {'types': ['FLOAT', 'DECIMAL', 'NUMERIC'], 'desc': 'Medication dosage amount'},
                'admission_time': {'types': ['DATETIME', 'TIMESTAMP', 'DATE'], 'desc': 'Hospital admission timestamp'},
                'discharge_time': {'types': ['DATETIME', 'TIMESTAMP', 'DATE'], 'desc': 'Hospital discharge timestamp'},
                'lab_result': {'types': ['FLOAT', 'DECIMAL', 'NUMERIC'], 'desc': 'Laboratory test result value'},
                'notes': {'types': ['TEXT', 'CLOB', 'VARCHAR'], 'desc': 'Clinical notes and observations'},
                'icu_stay': {'types': ['INT', 'INTEGER'], 'desc': 'ICU stay duration in days'},
                'mortality': {'types': ['BOOLEAN', 'TINYINT', 'INT'], 'desc': 'In-hospital mortality indicator'},
            },
        },
        'SYNTHEA': {
            'attrs': {
                'patient_id': {'types': ['INT', 'BIGINT', 'VARCHAR'], 'desc': 'Synthetic patient identifier'},
                'birth_date': {'types': ['DATE', 'DATETIME', 'STRING'], 'desc': 'Patient birth date'},
                'race': {'types': ['VARCHAR', 'TEXT', 'STRING'], 'desc': 'Patient race category'},
                'ethnicity': {'types': ['VARCHAR', 'TEXT', 'STRING'], 'desc': 'Patient ethnicity'},
                'condition': {'types': ['TEXT', 'VARCHAR', 'CLOB'], 'desc': 'Medical condition description'},
                'procedure': {'types': ['VARCHAR', 'TEXT', 'STRING'], 'desc': 'Medical procedure code'},
                'cost': {'types': ['DECIMAL', 'FLOAT', 'NUMERIC'], 'desc': 'Procedure cost in USD'},
                'encounter_date': {'types': ['DATETIME', 'DATE', 'TIMESTAMP'], 'desc': 'Patient encounter date'},
                'provider': {'types': ['VARCHAR', 'TEXT', 'STRING'], 'desc': 'Healthcare provider name'},
                'organization': {'types': ['VARCHAR', 'TEXT', 'STRING'], 'desc': 'Healthcare organization'},
                'zip_code': {'types': ['VARCHAR', 'INT', 'STRING'], 'desc': 'Patient ZIP code'},
                'income': {'types': ['DECIMAL', 'FLOAT', 'INT'], 'desc': 'Annual household income'},
            },
        },
    }
    
    if name not in domain_templates:
        raise ValueError(f"Unknown SM dataset: {name}. Available: {list(domain_templates.keys())}")
    
    t = domain_templates[name]
    pairs = []
    attr_names = list(t['attrs'].keys())
    
    for _ in range(n_pairs):
        if random.random() < 0.5:
            # Match: same attribute, different types
            attr = random.choice(attr_names)
            info = t['attrs'][attr]
            src_type = random.choice(info['types'])
            tgt_type = random.choice([x for x in info['types'] if x != src_type] or info['types'])
            pairs.append({
                'source': {'name': attr, 'type': src_type, 'description': info['desc']},
                'target': {'name': attr, 'type': tgt_type, 'description': info['desc']},
                'label': 1,
            })
        else:
            # Non-match: different attributes
            a1, a2 = random.sample(attr_names, 2)
            pairs.append({
                'source': {'name': a1, 'type': random.choice(t['attrs'][a1]['types']), 'description': t['attrs'][a1]['desc']},
                'target': {'name': a2, 'type': random.choice(t['attrs'][a2]['types']), 'description': t['attrs'][a2]['desc']},
                'label': 0,
            })
    
    output_dir = os.path.join(DATA_DIR, 'sm_prepared')
    os.makedirs(output_dir, exist_ok=True)
    output_path = os.path.join(output_dir, f'{name}.json')
    with open(output_path, 'w') as f:
        json.dump(pairs, f, indent=2)
    
    print(f"  Prepared SM {name}: {len(pairs)} pairs")
    return output_path


if __name__ == '__main__':
    import sys
    random.seed(42)
    
    print("=" * 60)
    print("Real Data Loader - Dataset Availability")
    print("=" * 60)
    
    # ER datasets
    print("\n--- ER Datasets (DeepMatcher format) ---")
    er_ds = get_all_er_datasets()
    for key, info in sorted(er_ds.items()):
        print(f"  {key}: {info['n_pairs']} pairs")
    
    if len(sys.argv) > 1 and sys.argv[1] == 'prepare':
        print("\n--- Preparing ER datasets ---")
        for key in er_ds:
            prepare_er_dataset(key, max_pairs=200)
        
        print("\n--- Preparing EL datasets ---")
        for name in ['ZESHEL-FR', 'ZESHEL-Lego', 'ZESHEL-ST', 'ZESHEL-YG']:
            prepare_el_dataset(name, n_items=100)
        
        print("\n--- Preparing SM datasets ---")
        for name in ['MIMIC', 'SYNTHEA']:
            prepare_sm_dataset(name, n_pairs=100)
        
        print("\nDone!")