"""High-level grading functionality"""
import json
from datetime import datetime
from pathlib import Path

import pandas as pd
from tqdm import tqdm

from mlebench.data import get_leaderboard, is_dataset_prepared
from mlebench.grade_helpers import CompetitionReport
from mlebench.registry import Competition, Registry
from mlebench.registry import registry as DEFAULT_REGISTRY
from mlebench.utils import get_logger, get_timestamp, load_answers, purple, read_csv, read_jsonl

logger = get_logger(__name__)


def grade_jsonl(
    path_to_submissions: Path,
    output_dir: Path,
    registry: Registry = DEFAULT_REGISTRY,
):
    """
    Grades multiple submissions stored in a JSONL file.
    Saves the aggregated report as a JSON file.
    """

    submissions = read_jsonl(path_to_submissions, skip_commented_out_lines=True)
    competitions_reports = []

    for submission in tqdm(submissions, desc="Grading submissions", unit="submission"):
        submission_path = Path(str(submission["submission_path"]))
        competition_id = submission["competition_id"]
        competition = registry.get_competition(competition_id)
        single_report = grade_csv(submission_path, competition)
        competitions_reports.append(single_report)

    aggregated_report = aggregate_reports(competitions_reports)
    timestamp = get_timestamp()
    save_path = output_dir / f"{timestamp}_grading_report.json"
    logger.info(
        json.dumps(
            {k: v for k, v in aggregated_report.items() if k != "competition_reports"}, indent=4
        )
    )

    output_dir.mkdir(exist_ok=True)
    with open(save_path, "w") as f:
        json.dump(aggregated_report, f, indent=2)
    logger.info(purple(f"Saved summary report to {save_path}"))


def grade_csv(path_to_submission: Path, competition: Competition) -> CompetitionReport:
    """Grades a submission CSV for the given competition."""

    if not is_dataset_prepared(competition, grading_only=True):
        raise ValueError(
            f"Dataset for competition `{competition.id}` is not prepared! "
            f"Please run `mlebench prepare -c {competition.id}` to prepare the dataset."
        )

    score = None
    submission_exists = path_to_submission.is_file() and path_to_submission.suffix.lower() == ".csv"

    if submission_exists:
        submission_df = read_csv(path_to_submission)
        logger.info(purple(f"Load answers from {competition.answers}"))
        answers = load_answers(competition.answers)
        score = competition.grader(submission_df, answers)
    else:
        logger.warning(
            f"Invalid submission file: {path_to_submission}. Please check that the file exists and it is a CSV."
        )

    valid_submission = score is not None

    # Helper function to determine is_lower_better from metric name
    def infer_is_lower_better(metric_name: str) -> bool:
        """
        Infer whether lower is better based on metric name.
        Returns True if lower is better (error metrics), False if higher is better (accuracy metrics).
        """
        metric_name_lower = metric_name.lower()

        # Error/loss metrics (lower is better)
        error_metrics = [
            'rmse', 'rmsle', 'mae', 'mse', 'mape',
            'log_loss', 'crossentropy', 'kld', 'error',
            'loss', 'distance', 'deviation'
        ]

        # Accuracy/score metrics (higher is better)
        accuracy_metrics = [
            'accuracy', 'precision', 'recall', 'f1',
            'auc', 'roc', 'score', 'r2', 'correlation',
            'iou', 'dice', 'map', 'ndcg'
        ]

        for error_metric in error_metrics:
            if error_metric in metric_name_lower:
                return True

        for accuracy_metric in accuracy_metrics:
            if accuracy_metric in metric_name_lower:
                return False

        # Default: assume lower is better for safety
        return True

    # Try to get leaderboard, but make it optional
    try:
        competition_leaderboard = get_leaderboard(competition)
        rank_info = competition.grader.rank_score(score, competition_leaderboard)
        is_lower_better = competition.grader.is_lower_better(competition_leaderboard)
    except (AssertionError, FileNotFoundError):
        # Leaderboard not found, infer is_lower_better from grader name
        logger.warning(f"Leaderboard not found for competition {competition.id}, skipping rank calculation")
        competition_leaderboard = None
        rank_info = {
            "gold_threshold": None,
            "silver_threshold": None,
            "bronze_threshold": None,
            "median_threshold": None,
            "gold_medal": False,
            "silver_medal": False,
            "bronze_medal": False,
            "above_median": False,
        }
        # Infer is_lower_better from grader name
        is_lower_better = infer_is_lower_better(competition.grader.name)

    return CompetitionReport(
        competition_id=competition.id,
        score=score,
        gold_threshold=rank_info["gold_threshold"],
        silver_threshold=rank_info["silver_threshold"],
        bronze_threshold=rank_info["bronze_threshold"],
        median_threshold=rank_info["median_threshold"],
        any_medal=rank_info["gold_medal"] or rank_info["silver_medal"] or rank_info["bronze_medal"],
        gold_medal=rank_info["gold_medal"],
        silver_medal=rank_info["silver_medal"],
        bronze_medal=rank_info["bronze_medal"],
        above_median=rank_info["above_median"],
        submission_exists=submission_exists,
        valid_submission=valid_submission,
        is_lower_better=is_lower_better,
        created_at=datetime.now(),
        submission_path=str(path_to_submission),
    )


def validate_submission(submission: Path, competition: Competition) -> tuple[bool, str]:
    """
    Validates a submission for the given competition by actually running the competition grader.
    This is designed for end users, not developers (we assume that the competition grader is
    correctly implemented and use that for validating the submission, not the other way around).
    """
    if not submission.is_file():
        return False, f"Submission invalid! Submission file {submission} does not exist."

    if not submission.suffix.lower() == ".csv":
        return False, "Submission invalid! Submission file must be a CSV file."

    if not is_dataset_prepared(competition, grading_only=True):
        raise ValueError(
            f"Dataset for competition `{competition.id}` is not prepared! "
            f"Please run `mlebench prepare -c {competition.id}` to prepare the dataset."
        )

    try:
        competition.grader.grade_fn(read_csv(submission), read_csv(competition.answers))
    except Exception as e:
        return (
            False,
            f"Submission invalid! The attempt to grade the submission has resulted in the following error message:\n{e}",
        )

    return True, "Submission is valid."


def aggregate_reports(competition_reports: list[CompetitionReport]) -> dict:
    """
    Builds the summary report from a list of competition reports.
    If pass_at_n is True, then aggregate performence of competitions by selecting the best performance per
    competition, otherwise sum metrics
    """

    total_gold_medals = sum(report.gold_medal for report in competition_reports)
    total_silver_medals = sum(report.silver_medal for report in competition_reports)
    total_bronze_medals = sum(report.bronze_medal for report in competition_reports)
    total_above_median = sum(report.above_median for report in competition_reports)
    total_submissions = sum(report.submission_exists for report in competition_reports)
    total_valid_submissions = sum(report.valid_submission for report in competition_reports)

    summary_report = {
        "total_runs": int(len(competition_reports)),
        "total_runs_with_submissions": int(total_submissions),
        "total_valid_submissions": int(total_valid_submissions),
        "total_medals": int(total_gold_medals + total_silver_medals + total_bronze_medals),
        "total_gold_medals": int(total_gold_medals),
        "total_silver_medals": int(total_silver_medals),
        "total_bronze_medals": int(total_bronze_medals),
        "total_above_median": int(total_above_median),
        "competition_reports": [cr.to_dict() for cr in competition_reports],
    }

    return summary_report

def simple_accuracy_grader(submission: pd.DataFrame, answers: pd.DataFrame) -> float:
    """
    Computes simple accuracy between submission and answers.
    Assumes matching order or joins on an ID column.
    """
    try:
        # 1. Try to find ID column to merge on
        id_col = None
        common_cols = set(submission.columns).intersection(set(answers.columns))
        for col in common_cols:
            if col.lower() in ['id', 'key', 'index', 'sample_id', 'image_id', 'patient_id']:
                id_col = col
                break
                
        if id_col:
            merged = submission.merge(answers, on=id_col, suffixes=('_pred', '_true'))
            # Compare all other columns
            pred_cols = [c for c in submission.columns if c != id_col]
            true_cols = [c for c in answers.columns if c != id_col]
            
            # If there's only one other column, compare it
            if len(pred_cols) == 1 and len(true_cols) == 1:
                return float((merged[pred_cols[0]] == merged[true_cols[0]]).mean())
            
            # Iterate over shared columns that are NOT id
            shared_content_cols = [c for c in common_cols if c != id_col]
            if shared_content_cols:
                matches = []
                for c in shared_content_cols:
                    # In merged df, they are c_pred and c_true
                    matches.append((merged[f"{c}_pred"] == merged[f"{c}_true"]).mean())
                if matches:
                    return float(sum(matches) / len(matches))

        # 2. Fallback: Position-based comparison (assuming sorted)
        if len(submission) == len(answers):
            # Compare last column if dimensions match
            if submission.shape[1] == answers.shape[1]:
                return float((submission.iloc[:, -1] == answers.iloc[:, -1]).mean())
            
            # Or compare specific common columns
            common_cols = set(submission.columns).intersection(set(answers.columns))
            if common_cols:
                matches = []
                for c in common_cols:
                    matches.append((submission[c] == answers[c]).mean())
                if matches:
                    return float(sum(matches) / len(matches))
    except Exception as e:
        logger.error(f"simple_accuracy_grader failed: {e}")
        return 0.0

    return 0.0