"""

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 trainiqConfig

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: trainiqConfig):

        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)

