"""
automl_lib.model_selector
~~~~~~~~~~~~~~~~~~~~~~~~~
Logic for training multiple candidate models and selecting the best one.
"""

from __future__ import annotations

import copy
from typing import Any

import pandas as pd
from torch.utils.data import DataLoader

from .config import AutoMLConfig
from .model_zoo import list_models, build_model
from .trainer import Trainer, train_sklearn
from .utils import get_logger

logger = get_logger(__name__)

class ModelSelector:
    """Ranks and selects the best model from a set of candidates."""

    def __init__(self, config: AutoMLConfig):
        self.config = config
        self.leaderboard: list[dict[str, Any]] = []

    def select_best(
        self,
        train_loader: DataLoader,
        val_loader: DataLoader,
        data_type: str,
        task: str,
        num_classes: int,
        input_shape: tuple,
    ) -> tuple[Any, dict[str, Any]]:
        """Train all candidate models and return (best_model, results)."""
        candidates = list_models(data_type, task)
        logger.info(f"Comparing {len(candidates)} candidate models: {candidates}")

        for name in candidates:
            logger.info("-" * 40)
            logger.info(f"Evaluating model: {name}")
            
            # Create a temporary config for this model
            tmp_config = copy.deepcopy(self.config)
            tmp_config.model_name = name
            
            # Reduce epochs for selection to speed up comparison
            if self.config.epochs > 10:
                tmp_config.epochs = 10
            tmp_config.verbose = False 
            
            try:
                model = build_model(
                    name, 
                    input_shape=input_shape, 
                    num_classes=num_classes, 
                    config=tmp_config
                )
                
                is_sklearn = getattr(model, "_is_sklearn", False)
                if is_sklearn:
                    results = train_sklearn(model, train_loader, val_loader, task)
                else:
                    trainer = Trainer(model, tmp_config)
                    results = trainer.train(train_loader, val_loader, num_classes, task)
                
                metric_name = "val_acc" if task == "classification" else "val_loss"
                if is_sklearn:
                    score = results.get("val_acc" if task == "classification" else "val_loss", 0.0)
                else:
                    # For PyTorch models, get from history
                    if "history" in results and metric_name in results["history"]:
                        if task == "classification":
                            score = max(results["history"]["val_acc"]) if results["history"]["val_acc"] else 0.0
                        else:
                            score = min(results["history"]["val_loss"]) if results["history"]["val_loss"] else float('inf')
                    else:
                        score = results.get(f"best_{metric_name}", 0.0)
                
                # Normalize score for sorting (higher is better)
                sort_score = score if task == "classification" else -score
                
                self.leaderboard.append({
                    "model_name": name,
                    "score": score,
                    "sort_score": sort_score,
                    "results": results,
                    "model": model
                })
                logger.info(f"Model {name} finished with {metric_name}: {score:.4f}")

            except Exception as e:
                logger.error(f"Failed to evaluate {name}: {e}")
                continue

        if not self.leaderboard:
            raise RuntimeError("All candidate models failed during selection.")

        # Sort by score descending
        self.leaderboard.sort(key=lambda x: x["sort_score"], reverse=True)
        
        best_entry = self.leaderboard[0]
        logger.info("=" * 40)
        logger.info(f"Leaderboard (Best to Worst):")
        for i, entry in enumerate(self.leaderboard):
            logger.info(f"  {i+1}. {entry['model_name']}: {entry['score']:.4f}")
        logger.info(f"WINNER: {best_entry['model_name']}")
        logger.info("=" * 40)

        return best_entry["model"], best_entry["results"]

    def get_leaderboard_df(self) -> pd.DataFrame:
        """Return the leaderboard as a pandas DataFrame."""
        entries = []
        for e in self.leaderboard:
            entries.append({"model": e["model_name"], "score": e["score"]})
        return pd.DataFrame(entries)
