#!/usr/bin/env python3
"""
Unified Author Disambiguation Benchmark Script

A flexible benchmark runner that can work with any Excel file by mapping
column names to required fields.

Usage:
    python scripts/benchmark.py \\
        --input data/my_authors.xlsx \\
        --first_name "FirstName,GivenName" \\
        --last_name "LastName,FamilyName" \\
        --ground_truth "openalex_id" \\
        --affiliation "Institution" \\
        --context "Keywords,ResearchAreas" \\
        --n 100

Column Mapping:
    - Columns can be specified as comma-separated alternatives (first match wins)
    - If a column doesn't exist, it will be left empty
    - Context can combine multiple columns (e.g., keywords + papers)
"""

import argparse
import asyncio
import json
import os
import sys
from datetime import datetime
from pathlib import Path
from typing import Any, Dict, List, Optional

import pandas as pd
from tqdm.asyncio import tqdm_asyncio

# Add src directory to path
sys.path.insert(0, str(Path(__file__).parent.parent / "src"))

from production_agent import disambiguate_author


def find_column(df: pd.DataFrame, column_spec: str) -> Optional[str]:
    """
    Find the first matching column from a comma-separated list of alternatives.
    
    Args:
        df: DataFrame to search in
        column_spec: Comma-separated column names (e.g., "FirstName,GivenName,first_name")
    
    Returns:
        The first matching column name, or None if no match found
    """
    if not column_spec:
        return None
    
    candidates = [c.strip() for c in column_spec.split(',')]
    
    for candidate in candidates:
        if candidate in df.columns:
            return candidate
    
    return None


def get_column_value(row: pd.Series, df: pd.DataFrame, column_spec: str) -> str:
    """
    Get value from row using column specification.
    
    Args:
        row: The data row
        df: The full dataframe (for column lookup)
        column_spec: Comma-separated column names
    
    Returns:
        String value (empty string if not found or NaN)
    """
    col_name = find_column(df, column_spec)
    if not col_name:
        return ""
    
    value = row.get(col_name, "")
    if pd.isna(value):
        return ""
    
    return str(value).strip()


def combine_context(row: pd.Series, df: pd.DataFrame, context_spec: str) -> str:
    """
    Combine multiple columns into a single context string.
    
    Args:
        row: The data row
        df: The full dataframe
        context_spec: Comma-separated column names to combine
    
    Returns:
        Combined context string (comma-separated values)
    """
    if not context_spec:
        return ""
    
    values = []
    for col_spec in context_spec.split(','):
        col_spec = col_spec.strip()
        col_name = find_column(df, col_spec)
        if col_name:
            val = row.get(col_name, "")
            if not pd.isna(val) and str(val).strip():
                values.append(str(val).strip())
    
    return ", ".join(values)


def parse_ground_truth(value: Any) -> str:
    """
    Parse ground truth OpenAlex ID from various formats.
    
    Handles:
    - Full URLs: https://openalex.org/A5074091984
    - IDs only: A5074091984
    - Lists/comma-separated: A1, A2, A3
    
    Returns:
        Primary OpenAlex ID (first if multiple)
    """
    if pd.isna(value) or not value:
        return ""
    
    value_str = str(value).strip()
    
    # Extract from URL if present
    if "openalex.org/" in value_str:
        value_str = value_str.split("openalex.org/")[-1].split(",")[0].split(" ")[0]
    
    # Handle comma-separated
    if "," in value_str:
        value_str = value_str.split(",")[0].strip()
    
    return value_str


async def run_single_test(
    test_id: int,
    row: pd.Series,
    df: pd.DataFrame,
    args: argparse.Namespace
) -> Dict[str, Any]:
    """Run disambiguation for a single author."""
    
    # Extract fields using column mapping
    first_name = get_column_value(row, df, args.first_name)
    last_name = get_column_value(row, df, args.last_name)
    affiliation = get_column_value(row, df, args.affiliation) if args.affiliation else ""
    context = combine_context(row, df, args.context) if args.context else ""
    
    # Get ground truth
    ground_truth_col = find_column(df, args.ground_truth)
    ground_truth = ""
    if ground_truth_col:
        ground_truth = parse_ground_truth(row.get(ground_truth_col))
    
    # Run disambiguation with concurrency control
    try:
        result = await disambiguate_author(
            first_name=first_name,
            last_name=last_name,
            institution=affiliation,
            context=context,
            find_email=False,
            max_concurrent=5  # Use agent's built-in concurrency control
        )
        
        # Normalize status
        status = result.get('status', 'error')
        if status == 'found':
            status = 'success'
        if status == 'error' and result.get('author_candidates'):
            status = 'success'
        
        # Extract candidates
        candidates = []
        
        # Try new format
        if result.get('author_candidates'):
            for candidate in result['author_candidates']:
                author = candidate.get('author', {})
                openalex_id = author.get('openalex_id', '')
                if 'openalex.org/' in openalex_id:
                    openalex_id = openalex_id.split('/')[-1]
                candidates.append({
                    'openalex_id': openalex_id,
                    'name': author.get('name', ''),
                    'confidence': candidate.get('confidence', 'unknown')
                })
        
        # Try old format
        elif result.get('author'):
            author = result.get('author', {})
            openalex_id = author.get('openalex_id', '')
            if 'openalex.org/' in openalex_id:
                openalex_id = openalex_id.split('/')[-1]
            if openalex_id:
                candidates.append({
                    'openalex_id': openalex_id,
                    'name': author.get('display_name', author.get('name', '')),
                    'confidence': result.get('confidence', 'unknown')
                })
        
        # Check if correct
        found_at_rank = None
        if ground_truth:
            for rank, candidate in enumerate(candidates, 1):
                if candidate['openalex_id'] == ground_truth:
                    found_at_rank = rank
                    break
        
        return {
            'test_id': test_id,
            'input': {
                'first_name': first_name,
                'last_name': last_name,
                'affiliation': affiliation,
                'context': context
            },
            'ground_truth': ground_truth,
            'result': {
                'status': status,
                'num_candidates': len(candidates),
                'found_at_rank': found_at_rank,
                'correct': found_at_rank == 1 if found_at_rank else False
            },
            'candidates_returned': candidates[:5]
        }
    
    except Exception as e:
        return {
            'test_id': test_id,
            'input': {
                'first_name': first_name,
                'last_name': last_name,
                'affiliation': affiliation,
                'context': context
            },
            'ground_truth': ground_truth,
            'error': str(e),
            'result': {
                'status': 'error',
                'num_candidates': 0,
                'correct': False
            }
        }


async def run_benchmark(args: argparse.Namespace):
    """Run the full benchmark."""
    
    print("=" * 70)
    print("AUTHOR DISAMBIGUATION BENCHMARK")
    print("=" * 70)
    print(f"Input file: {args.input}")
    print(f"Sample size: {args.n if args.n else 'all'}")
    print()
    
    # Load data
    df = pd.read_excel(args.input)
    print(f"✓ Loaded {len(df)} rows")
    
    # Verify columns exist
    print("\n📋 Column Mapping:")
    print(f"  First name: {find_column(df, args.first_name) or 'NOT FOUND'}")
    print(f"  Last name: {find_column(df, args.last_name) or 'NOT FOUND'}")
    if args.affiliation:
        print(f"  Affiliation: {find_column(df, args.affiliation) or 'NOT FOUND (optional)'}")
    if args.context:
        context_cols = [find_column(df, c.strip()) for c in args.context.split(',')]
        print(f"  Context: {', '.join([c for c in context_cols if c]) or 'NOT FOUND (optional)'}")
    print(f"  Ground truth: {find_column(df, args.ground_truth) or 'NOT FOUND'}")
    print()
    
    # Filter out rows without ground truth
    ground_truth_col = find_column(df, args.ground_truth)
    if ground_truth_col:
        df_original_len = len(df)
        df = df[~df[ground_truth_col].isna() & (df[ground_truth_col] != '')]
        filtered_count = df_original_len - len(df)
        if filtered_count > 0:
            print(f"✓ Filtered out {filtered_count} rows without ground truth")
    
    # Filter out rows without required context columns
    if args.context:
        context_cols = [c.strip() for c in args.context.split(',')]
        for context_col_spec in context_cols:
            context_col = find_column(df, context_col_spec)
            if context_col:
                df_original_len = len(df)
                df = df[~df[context_col].isna() & (df[context_col] != '')]
                filtered_count = df_original_len - len(df)
                if filtered_count > 0:
                    print(f"✓ Filtered out {filtered_count} rows without '{context_col}' data")
    
    # Filter out rows without affiliation (if specified)
    if args.affiliation:
        affiliation_col = find_column(df, args.affiliation)
        if affiliation_col:
            df_original_len = len(df)
            df = df[~df[affiliation_col].isna() & (df[affiliation_col] != '')]
            filtered_count = df_original_len - len(df)
            if filtered_count > 0:
                print(f"✓ Filtered out {filtered_count} rows without affiliation data")
    
    if len(df) == 0:
        print("\n❌ ERROR: No rows remaining after filtering!")
        return
    
    # Sample if needed
    if args.n and args.n < len(df):
        df = df.sample(n=args.n, random_state=args.seed)
        print(f"✓ Sampled {args.n} rows (random_state={args.seed})")
    
    # Concurrency control is handled by the agent (max_concurrent=5)
    print(f"✓ Concurrency control: 5 simultaneous SDK initializations (handled by agent)")
    
    # Run tests
    print(f"\n🚀 Running disambiguation on {len(df)} authors...\n")
    
    tasks = [
        run_single_test(i, row, df, args)
        for i, (_, row) in enumerate(df.iterrows())
    ]
    
    results = await tqdm_asyncio.gather(*tasks, desc="Testing")
    
    # Calculate metrics
    total = len(results)
    correct = sum(1 for r in results if r['result']['correct'])
    top_2 = sum(1 for r in results if r['result'].get('found_at_rank') in [1, 2])
    top_3 = sum(1 for r in results if r['result'].get('found_at_rank') in [1, 2, 3])
    errors = sum(1 for r in results if r['result']['status'] == 'error')
    
    # Print results
    print("\n" + "=" * 70)
    print("RESULTS")
    print("=" * 70)
    print(f"Total tests: {total}")
    print()
    print(f"Top-1 Accuracy: {correct}/{total} ({correct/total*100:.1f}%)")
    print(f"Top-2 Accuracy: {top_2}/{total} ({top_2/total*100:.1f}%)")
    print(f"Top-3 Accuracy: {top_3}/{total} ({top_3/total*100:.1f}%)")
    print()
    print(f"Errors: {errors}")
    print("=" * 70)
    
    # Save results
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    output_file = Path("output") / f"benchmark_{timestamp}.json"
    output_file.parent.mkdir(exist_ok=True)
    
    output_data = {
        'benchmark_info': {
            'input_file': args.input,
            'total_tests': total,
            'timestamp': timestamp
        },
        'column_mapping': {
            'first_name': args.first_name,
            'last_name': args.last_name,
            'affiliation': args.affiliation,
            'context': args.context,
            'ground_truth': args.ground_truth
        },
        'summary': {
            'total_tests': total,
            'top_1_accuracy': correct / total * 100,
            'top_2_accuracy': top_2 / total * 100,
            'top_3_accuracy': top_3 / total * 100,
            'errors': errors
        },
        'all_results': results
    }
    
    with open(output_file, 'w') as f:
        json.dump(output_data, f, indent=2)
    
    print(f"\n📄 Full report saved to: {output_file}")
    print()


def main():
    parser = argparse.ArgumentParser(
        description='Run author disambiguation benchmark with flexible column mapping',
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
  # Simple: Use all defaults (name + affiliation + keywords)
  python scripts/benchmark.py -n 100
  
  # Name-only benchmark (no context)
  python scripts/benchmark.py -n 100 --context "" --affiliation ""
  
  # Keywords + subject areas as context
  python scripts/benchmark.py -n 100 --context "keywords,subject_areas"
  
  # With different random seed
  python scripts/benchmark.py -n 100 --seed 999
  
  # Custom file with explicit columns
  python scripts/benchmark.py \\
      --input data/authors.xlsx \\
      --first_name "FirstName" \\
      --last_name "LastName" \\
      --ground_truth "openalex_id" \\
      --affiliation "Institution" \\
      --context "Keywords,ResearchArea" \\
      -n 50
        """
    )
    
    parser.add_argument(
        '--input',
        default='data/embo_membership_candidates_with_author_ids.xlsx',
        help='Input Excel file path (default: data/embo_membership_candidates_with_author_ids.xlsx)'
    )
    
    parser.add_argument(
        '--first_name',
        default='first_name_plus_initials',
        help='Column name(s) for first name (default: first_name_plus_initials)'
    )
    
    parser.add_argument(
        '--last_name',
        default='last_name',
        help='Column name(s) for last name (default: last_name)'
    )
    
    parser.add_argument(
        '--ground_truth',
        default='matched_author_openalex_id',
        help='Column name(s) for ground truth OpenAlex ID (default: matched_author_openalex_id)'
    )
    
    parser.add_argument(
        '--affiliation',
        default='affiliation',
        help='Column name(s) for affiliation (default: affiliation)'
    )
    
    parser.add_argument(
        '--context',
        default='keywords',
        help='Column name(s) for context (default: keywords). Can specify multiple to combine (comma-separated)'
    )
    
    parser.add_argument(
        '-n',
        type=int,
        help='Number of samples to test (default: all rows)'
    )
    
    parser.add_argument(
        '--seed',
        type=int,
        default=42,
        help='Random seed for sampling (default: 42)'
    )
    
    args = parser.parse_args()
    
    # Run benchmark
    asyncio.run(run_benchmark(args))


if __name__ == "__main__":
    main()
