"""
Input/output utilities for cognitive map evaluation.

This module provides functions to:
1. Load and parse evaluation files
2. Save evaluation results
3. Generate reports and statistics
"""

import os
import json
import glob
from datetime import datetime
import pandas as pd
from typing import Dict, List, Tuple, Any, Optional
from collections import defaultdict

def load_jsonl_data(jsonl_path: str) -> List[Dict]:
    """
    Load data from a JSONL file.
    
    Args:
        jsonl_path: Path to the JSONL file
        
    Returns:
        List of loaded JSON objects
    """
    data = []
    with open(jsonl_path, 'r', encoding='utf-8') as f:
        for line in f:
            data.append(json.loads(line.strip()))
    return data

def save_json_results(results: Dict, output_path: str) -> None:
    """
    Save evaluation results as JSON.
    
    Args:
        results: Evaluation results dictionary
        output_path: Path to save the results to
    """
    # Ensure directory exists
    os.makedirs(os.path.dirname(output_path), exist_ok=True)
    
    with open(output_path, 'w', encoding='utf-8') as f:
        json.dump(results, f, indent=2)
    
    print(f"Results saved to {output_path}")

def save_csv_results(results: List[Dict], output_path: str) -> None:
    """
    Save evaluation results as CSV.
    
    Args:
        results: List of result dictionaries
        output_path: Path to save the CSV to
    """
    # Ensure directory exists
    os.makedirs(os.path.dirname(output_path), exist_ok=True)
    
    # Create DataFrame and save
    df = pd.DataFrame(results)
    df.to_csv(output_path, index=False)
    
    print(f"CSV results saved to {output_path}")

def find_evaluation_files(eval_dir: str, pattern: str = "*.jsonl") -> List[str]:
    """
    Find all evaluation files matching a pattern in a directory.
    
    Args:
        eval_dir: Directory to search in
        pattern: Glob pattern to match
        
    Returns:
        List of matching file paths
    """
    return glob.glob(os.path.join(eval_dir, pattern))

def create_output_paths(base_name: str, output_dir: str) -> Dict[str, str]:
    """
    Create standardized output paths for results.
    
    Args:
        base_name: Base name for the result files
        output_dir: Directory to store the results
        
    Returns:
        Dictionary of output paths
    """
    # Ensure output directory exists
    os.makedirs(output_dir, exist_ok=True)
    
    # Create timestamp for unique filenames
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    
    # Create paths
    paths = {
        "json": os.path.join(output_dir, f"{base_name}_results_{timestamp}.json"),
        "csv": os.path.join(output_dir, f"{base_name}_results_{timestamp}.csv"),
        "settings": os.path.join(output_dir, f"{base_name}_settings_{timestamp}.json"),
    }
    
    return paths

def initialize_results_structure() -> Dict:
    """
    Initialize the results data structure.
    
    Returns:
        Empty results structure
    """
    # Basic settings to track - 确保顺序一致
    settings = ['around', 'rotation', 'translation', 'among', 'other']
    
    # Define which settings should be included in overall metrics
    # Translation is excluded from overall metrics but still tracked
    settings_to_include = {
        'around': True, 
        'rotation': True, 
        'translation': False,  # Exclude translation from overall metrics
        'among': True, 
        'other': True
    }
    
    # Initialize results
    results = {
        'total': 0,
        'gen_cogmap_correct': 0,
        'cogmap_similarity': {
            'total_valid': 0,
            'valid_percent': 0.0,
            'isomorphic_count': 0,
            'rotation_invariant_isomorphic_count': 0,
            'parsable_json_count': 0,
            'valid_format_count': 0,
            'avg_relative_position_accuracy': 0.0,
            'avg_facing_similarity': 0.0,
            'avg_directional_similarity': 0.0,
            'avg_overall_similarity': 0.0,
            'rotation_distribution': defaultdict(int)
        },
        'settings': {setting: {
            'total': 0, 
            'gen_cogmap_correct': 0,
            'include_in_overall': settings_to_include.get(setting, True),  # Flag to determine if setting should be included in overall metrics
            'cogmap_similarity': {
                'total_valid': 0,
                'valid_percent': 0.0,
                'isomorphic_count': 0,
                'rotation_invariant_isomorphic_count': 0,
                'parsable_json_count': 0,
                'valid_format_count': 0,
                'avg_relative_position_accuracy': 0.0,
                'avg_facing_similarity': 0.0,
                'avg_directional_similarity': 0.0,
                'avg_overall_similarity': 0.0
            }
        } for setting in settings}
    }
    
    return results

def print_results(results: Dict) -> None:
    """
    Print evaluation results in a readable format.
    
    Args:
        results: Evaluation results dictionary
    """
    # 获取总数和未过滤总数
    total = results['total']  # 过滤后的总数（不包括translation）
    unfiltered_total = results.get('unfiltered_total', total)  # 原始总数（全部setting）
    
    print(f"Total examples: {unfiltered_total} (Evaluated: {total}, excluding translation)")
    print(f"Gen CogMap accuracy: {results['gen_cogmap_accuracy']*100:.2f}% ({results['gen_cogmap_correct']}/{total})")
    
    # Print cognitive map similarity metrics
    print("\nCognitive Map Validation Metrics:")
    # Calculate validation rates for each stage
    parsable_json_count = results['cogmap_similarity'].get('parsable_json_count', 0)
    parsable_json_percent = 100 * parsable_json_count / total if total > 0 else 0
    
    valid_format_count = results['cogmap_similarity'].get('valid_format_count', 0)
    valid_format_percent = 100 * valid_format_count / total if total > 0 else 0
    
    valid_graph_count = results['cogmap_similarity']['total_valid']
    valid_graph_percent = results['cogmap_similarity']['valid_percent']
    
    # Print validation metrics
    print(f"Parsable JSON: {parsable_json_count}/{total} ({parsable_json_percent:.2f}%)")
    print(f"Valid graph structure: {valid_graph_count}/{total} ({valid_graph_percent:.2f}%)")
    
    # Only print similarity metrics if valid graph structures exist
    if valid_graph_count > 0:
        print("\nCognitive Map Similarity Metrics:")
        iso_count = results['cogmap_similarity']['isomorphic_count']
        iso_percent = 100 * iso_count / total if total > 0 else 0
        
        print(f"Isomorphic graphs: {iso_count}/{total} ({iso_percent:.2f}%)")
        print(f"Avg overall similarity: {results['cogmap_similarity'].get('avg_overall_similarity', 0)*100:.2f}%")
        print(f"Avg directional similarity: {results['cogmap_similarity'].get('avg_directional_similarity', 0)*100:.2f}%")
        print(f"Avg facing similarity: {results['cogmap_similarity']['avg_facing_similarity']*100:.2f}%")
        print(f"Best rotation distribution: {dict(results['cogmap_similarity']['rotation_distribution'])}")
    
    print("\nResults by setting:")
    # Print all settings including translation - 使用一致的顺序
    for setting in ['translation', 'rotation', 'among', 'around']:
        stats = results['settings'][setting]
        if stats['total'] > 0:
            setting_total = stats['total']
            included = "(included in overall)" if stats.get('include_in_overall', True) else "(reported separately)"
            print(f"  {setting} {included}:")
            print(f"    Total examples: {setting_total}")
            print(f"    Answer accuracy: {stats['gen_cogmap_accuracy']*100:.2f}% ({stats['gen_cogmap_correct']}/{setting_total})")
            
            # Calculate validation rates
            s_parsable_json_count = stats['cogmap_similarity'].get('parsable_json_count', 0)
            s_parsable_json_percent = 100 * s_parsable_json_count / setting_total if setting_total > 0 else 0
            
            s_valid_format_count = stats['cogmap_similarity'].get('valid_format_count', 0)
            s_valid_format_percent = 100 * s_valid_format_count / setting_total if setting_total > 0 else 0
            
            s_valid_graph_count = stats['cogmap_similarity']['total_valid']
            s_valid_graph_percent = stats['cogmap_similarity']['valid_percent']
            
            # Print concise validation metrics
            print(f"    Validation: JSON({s_parsable_json_percent:.2f}%) | Format({s_valid_format_percent:.2f}%) | Graph({s_valid_graph_percent:.2f}%)")
            
            # Only print similarity metrics if valid graph structures exist
            if s_valid_graph_count > 0:
                s_iso_count = stats['cogmap_similarity']['isomorphic_count']
                s_iso_percent = 100 * s_iso_count / setting_total
                
                # Print similarity metrics
                print(f"    Similarity: Isomorphic({s_iso_percent:.2f}%) | Overall({stats['cogmap_similarity'].get('avg_overall_similarity', 0)*100:.2f}%) | " +
                      f"Direction({stats['cogmap_similarity']['avg_relative_position_accuracy']*100:.2f}%) | " +
                      f"Facing({stats['cogmap_similarity']['avg_facing_similarity']*100:.2f}%)")

                # 如果有设置特定的旋转分布，显示它
                if 'rotation_distribution' in stats['cogmap_similarity'] and stats['cogmap_similarity']['rotation_distribution']:
                    print(f"    Rotation distribution: {dict(stats['cogmap_similarity']['rotation_distribution'])}")
                
    # 'other' category is not shown separately to simplify output

def prepare_summary_record(results: Dict, filename: str, model: str, version: str, gen_cogmap: bool) -> Dict:
    """
    Prepare a summary record for a single evaluation file.
    
    Args:
        results: Evaluation results
        filename: Source filename
        model: Model name
        version: Model version
        gen_cogmap: Whether the model generated cognitive maps
        
    Returns:
        Summary record dictionary
    """
    return {
        'filename': filename,
        'model': model,
        'version': version,
        'gen_cogmap': gen_cogmap,
        'total_examples': results['total'],
        'gen_cogmap_accuracy': results['gen_cogmap_accuracy'],
        'valid_cogmaps': results['cogmap_similarity']['total_valid'],
        'valid_rate': results['cogmap_similarity']['valid_percent'] / 100.0,
        'isomorphic_count': results['cogmap_similarity']['isomorphic_count'],
        'isomorphic_rate': results['cogmap_similarity']['isomorphic_count'] / results['total'] if results['total'] > 0 else 0,
        'avg_relative_position_accuracy': results['cogmap_similarity']['avg_relative_position_accuracy'],
        'avg_facing_similarity': results['cogmap_similarity']['avg_facing_similarity'],
        'avg_directional_similarity': results['cogmap_similarity']['avg_directional_similarity'],
        'avg_overall_similarity': results['cogmap_similarity'].get('avg_overall_similarity', 0.0)
    }

def prepare_setting_record(setting_stats: Dict, filename: str, model: str, version: str, 
                        gen_cogmap: bool, setting: str) -> Dict:
    """
    Prepare a setting-specific record.
    
    Args:
        setting_stats: Setting-specific results
        filename: Source filename
        model: Model name
        version: Model version
        gen_cogmap: Whether the model generated cognitive maps
        setting: Setting name
        
    Returns:
        Setting-specific record dictionary
    """
    valid_count = setting_stats['cogmap_similarity']['total_valid']
    
    return {
        'filename': filename,
        'model': model,
        'version': version,
        'gen_cogmap': gen_cogmap,
        'setting': setting,
        'total': setting_stats['total'],
        'gen_cogmap_accuracy': setting_stats['gen_cogmap_accuracy'],
        'valid_cogmaps': valid_count,
        'valid_rate': setting_stats['cogmap_similarity']['valid_percent'] / 100.0,
        'isomorphic_count': setting_stats['cogmap_similarity']['isomorphic_count'],
        'isomorphic_rate': setting_stats['cogmap_similarity']['isomorphic_count'] / setting_stats['total'] if setting_stats['total'] > 0 else 0,
        'avg_relative_position_accuracy': setting_stats['cogmap_similarity']['avg_relative_position_accuracy'],
        'avg_facing_similarity': setting_stats['cogmap_similarity']['avg_facing_similarity'],
        'avg_directional_similarity': setting_stats['cogmap_similarity']['avg_directional_similarity'],
        'avg_overall_similarity': setting_stats['cogmap_similarity'].get('avg_overall_similarity', 0.0)
    } 