from pathlib import Path
from inspect_ai.log import list_eval_logs, read_eval_log, read_eval_log_samples
from inspect_ai import eval_set
from pprint import pprint 
import json
import statistics
from typing import Dict, List, Any, Optional, Union
from collections import defaultdict

def calculate_overall_cms_score(
    method: str = "mean",
    verbose: bool = False
) -> Dict[str, Any]:
    print("Running misalignment suite...")
    log_dir = str(Path(__file__).parent.parent) +"/logs-misalignment-suite"
    logs = list_eval_logs(log_dir)
    #cms_by_task = {}
    print(f"Found {len(logs)} logs")
    print("Logs:",str(Path(__file__).parent.parent) +"/logs-misalignment-suite" )
    #log_headers = list_eval_logs(log_dir)
    all_scores = []
    task_scores = {}
    task_sample_counts = defaultdict(int)
    failed_logs = []
    # Process each log file
    for log_path in logs:
        try:
            # Read log header (metadata without samples)
            log = read_eval_log(log_path, header_only=True)
            
            # Skip logs that didn't complete successfully
            if log.status != "success":
                failed_logs.append(f"{log_path} (status: {log.status})")
                continue
                
            # Get the task name
            task_name = log.eval.task
            
            # Extract CMS_scores from samples using the generator function
            task_cms_scores = []
            
            try:
                # Use the generator to iterate through samples one at a time
                for sample in read_eval_log_samples(log_path, all_samples_required=False):
                    cms_score = extract_cms_score(sample)
                    if cms_score is not None:
                        task_cms_scores.append(cms_score)
            except Exception as e:
                if verbose:
                    print(f"Error reading samples from {log_path}: {e}")
                continue
            
            if task_cms_scores:
                # Store task-specific scores
                task_scores[task_name] = {
                    "mean": statistics.mean(task_cms_scores),
                    "median": statistics.median(task_cms_scores),
                    "min": min(task_cms_scores),
                    "max": max(task_cms_scores),
                    "count": len(task_cms_scores)
                }
                
                # Track information for weighted calculations
                task_sample_counts[task_name] += len(task_cms_scores)
                
                # Add to overall scores list
                all_scores.extend(task_cms_scores)
                
                if verbose:
                    print(f"Task: {task_name}, Samples: {len(task_cms_scores)}, Mean CMS: {task_scores[task_name]['mean']:.4f}")
            else:
                if verbose:
                    print(f"No valid CMS scores found in task: {task_name}")
                    
        except Exception as e:
            failed_logs.append(f"{log_path} (error: {str(e)})")
            if verbose:
                print(f"Error processing log {log_path}: {e}")
    
    # Calculate overall CMS_score based on specified method
    if all_scores:
        if method == "mean":
            overall_cms_score = statistics.mean(all_scores)
        elif method == "weighted_mean":
            # Calculate weighted mean based on number of samples per task
            total_samples = sum(task_sample_counts.values())
            weighted_score = sum(
                task_scores[task]["mean"] * task_sample_counts[task] / total_samples
                for task in task_scores
            )
            overall_cms_score = weighted_score
        elif method == "median":
            overall_cms_score = statistics.median(all_scores)
        elif method == "min":
            overall_cms_score = min(all_scores)
        elif method == "max":
            overall_cms_score = max(all_scores)
        else:
            return {"error": f"Unknown aggregation method: {method}"}
    else:
        return {"error": "No valid CMS scores found in the logs", "failed_logs": failed_logs}
    
    # Calculate overall statistics
    overall_stats = {
        "overall_cms_score": overall_cms_score,
        "mean": statistics.mean(all_scores),
        "median": statistics.median(all_scores),
        "min": min(all_scores),
        "max": max(all_scores),
        "count": len(all_scores)
    }
    
    # Compile results
    results = {
        "overall": overall_stats,
        "by_task": task_scores
    }
    
    if failed_logs and verbose:
        results["failed_logs"] = failed_logs
    
    return results

def extract_cms_score(sample):
    """
    Extract CMS_score from a sample, handling different possible structures.
    
    Based on the cms_scorer_factory implementation, the CMS_score should be in the 
    score.value field, the score.metadata.composite_score field, or potentially
    elsewhere in the sample structure.
    """
    cms_score = None
    
    try:
        # Try to get score from score.value (most common location)
        if hasattr(sample, 'score') and sample.score is not None:
            if hasattr(sample.score, 'value'):
                cms_score = sample.score.value
            
            # Try to get from metadata.composite_score if available
            elif (hasattr(sample.score, 'metadata') and 
                  sample.score.metadata is not None and
                  'composite_score' in sample.score.metadata):
                cms_score = sample.score.metadata['composite_score']
                
        # Check if the extracted score is actually a number
        if cms_score is not None and not isinstance(cms_score, (int, float)):
            cms_score = None
    except Exception:
        cms_score = None
        
    return cms_score

# Example usage
if __name__ == "__main__":
    import argparse
    
    parser = argparse.ArgumentParser(description='Calculate overall CMS score across evaluations')
    #parser.add_argument('log_dir', help='Directory containing evaluation logs')
    parser.add_argument('--method', choices=['mean', 'weighted_mean', 'median', 'min', 'max'], 
                        default='mean', help='Method to aggregate scores')
    parser.add_argument('--verbose', action='store_true', help='Print detailed information')
    
    args = parser.parse_args()
    
    results = calculate_overall_cms_score( args.method, args.verbose)
    
    if "error" in results:
        print(f"Error: {results['error']}")
    else:
        print(f"\nOverall CMS Score ({args.method}): {results['overall']['overall_cms_score']:.4f}")
        print(f"Sample count: {results['overall']['count']}")
        print(f"Range: {results['overall']['min']:.4f} - {results['overall']['max']:.4f}")
        
        print("\nScores by task:")
        for task, stats in results['by_task'].items():
            print(f"  {task}: {stats['mean']:.4f} (samples: {stats['count']})")
