"""
automl_lib
~~~~~~~~~~
Universal AutoML Library – end-to-end ML/DL pipelines with a single call.

Quick Start
-----------
>>> from automl_lib import AutoML
>>> from automl_lib.config import AutoMLConfig
>>>
>>> config = AutoMLConfig(data_path="data.csv", target_column="label")
>>> automl = AutoML(config)
>>> results = automl.train()
>>> predictions = automl.predict(new_data)
>>> automl.export()
"""

from __future__ import annotations

__version__ = "0.1.0"
__author__ = "AutoML Team"

from typing import Any

import numpy as np

from .config import AutoMLConfig
from .utils import get_logger, set_seed


class AutoML:
    """High-level API that orchestrates the entire ML/DL pipeline.

    Parameters
    ----------
    config : AutoMLConfig
        Configuration object. See :class:`AutoMLConfig` for all options.

    Example
    -------
    >>> automl = AutoML(AutoMLConfig(data_path="my_dataset.csv"))
    >>> results = automl.train()
    >>> automl.predict(new_data)
    """

    def __init__(self, config: AutoMLConfig | None = None, **kwargs):
        if config is None:
            config = AutoMLConfig(**kwargs)
        self.config = config
        self.logger = get_logger("automl_lib")
        set_seed(config.seed)

        # Will be populated during train()
        self.model = None
        self.data_info: dict[str, Any] = {}
        self.data_type: str | None = None
        self.task: str | None = None
        self.input_shape: tuple | None = None
        self.num_classes: int | None = None
        self.results: dict[str, Any] = {}
        self._is_sklearn: bool = False

    # ══════════════════════════════════════════════════════════════════
    # TRAIN
    # ══════════════════════════════════════════════════════════════════
    def train(self) -> dict[str, Any]:
        """Run the full training pipeline: detect → load → build → train → evaluate.

        Returns a dict with training history and evaluation metrics.
        """
        from .detector import detect_data_type, detect_task
        from .data_loader import load_data
        from .model_zoo import list_models, build_model
        from .metrics import (
            compute_metrics, plot_training_history,
            plot_confusion_matrix,
        )

        self.logger.info("=" * 60)
        self.logger.info("AutoML Training Pipeline")
        self.logger.info("=" * 60)

        # ── Step 1: Detect data type ──────────────────────────────────
        self.data_type = self.config.data_type or detect_data_type(self.config.data_path)
        self.logger.info(f"Data type: {self.data_type}")

        # ── Step 2: Load & preprocess data ────────────────────────────
        data = load_data(self.config, self.data_type)
        train_loader = data["train_loader"]
        val_loader = data["val_loader"]
        self.num_classes = data["num_classes"]
        self.input_shape = data["input_shape"]
        self.data_info = data["dataset_info"]
        
        # Print dataset summary
        self.logger.info("=" * 60)
        self.logger.info("Dataset Summary")
        self.logger.info("=" * 60)
        for key, value in self.data_info.items():
            if key not in ("class_weights", "classes"):
                self.logger.info(f"  {key}: {value}")
        self.logger.info("=" * 60)

        # ── Step 3: Detect task ───────────────────────────────────────
        self.task = self.config.task or detect_task(
            data_type=self.data_type, num_classes=self.num_classes
        )
        self.config.task = self.task
        self.logger.info(f"Task: {self.task}")

        # ── Step 4: Select model ──────────────────────────────────────
        model_name = self.config.model_name
        if not model_name:
            from .model_selector import ModelSelector
            selector = ModelSelector(self.config)
            self.model, self.results = selector.select_best(
                train_loader, val_loader,
                data_type=self.data_type,
                task=self.task,
                num_classes=self.num_classes,
                input_shape=self.input_shape
            )
            model_name = self.config.model_name # selector might have updated it
        else:
            self.model = build_model(
                model_name,
                input_shape=self.input_shape,
                num_classes=self.num_classes,
                config=self.config,
            )

        self._is_sklearn = getattr(self.model, "_is_sklearn", False)

        # ── Step 5: Train (if not already trained by selector) ────────
        if not self.results:
            class_weights = self.data_info.get("class_weights")

            if self._is_sklearn:
                from .trainer import train_sklearn
                self.results = train_sklearn(
                    self.model, train_loader, val_loader, self.task
                )
            else:
                from .trainer import Trainer
                trainer = Trainer(self.model, self.config, class_weights=class_weights)
                self.results = trainer.train(
                    train_loader, val_loader,
                    num_classes=self.num_classes, task=self.task,
                )

        # ── Step 6: Evaluate ──────────────────────────────────────────
        if not self._is_sklearn:
            eval_metrics = compute_metrics(
                self.model, val_loader, self.task,
                num_classes=self.num_classes,
            )
            self.results["eval_metrics"] = eval_metrics

            # Plots
            if "history" in self.results and self.results["history"]:
                plot_training_history(self.results["history"], self.config.output_dir)
            if self.task == "classification" and eval_metrics.get("confusion_matrix") is not None:
                classes = self.data_info.get("classes")
                plot_confusion_matrix(
                    eval_metrics["confusion_matrix"],
                    class_names=classes,
                    output_dir=self.config.output_dir,
                )

        self.logger.info("=" * 60)
        self.logger.info("Training pipeline complete!")
        self.logger.info("=" * 60)
        return self.results

    # ══════════════════════════════════════════════════════════════════
    # PREDICT
    # ══════════════════════════════════════════════════════════════════
    def predict(self, data: Any) -> np.ndarray:
        """Run inference on new data using the trained model."""
        if self.model is None:
            raise RuntimeError("No trained model. Call train() first.")

        if self._is_sklearn:
            if isinstance(data, np.ndarray):
                return self.model.predict(data)
            import pandas as pd
            if isinstance(data, pd.DataFrame):
                return self.model.predict(data.values)

        from .inference import Predictor
        predictor = Predictor(self.model, self.config)
        return predictor.predict(data)

    # ══════════════════════════════════════════════════════════════════
    # EXPORT
    # ══════════════════════════════════════════════════════════════════
    def export(self, format: str | None = None) -> dict[str, str]:
        """Export the trained model to TorchScript / ONNX."""
        if self.model is None:
            raise RuntimeError("No trained model. Call train() first.")
        if self._is_sklearn:
            self.logger.warning("Sklearn models cannot be exported to TorchScript/ONNX. "
                                "Use joblib to serialize instead.")
            import joblib, os
            path = os.path.join(self.config.output_dir, "model_sklearn.joblib")
            joblib.dump(self.model.estimator, path)
            return {"joblib": path}

        if format:
            self.config.export_format = format

        from .exporter import export_model
        return export_model(self.model, self.input_shape, self.config)

    # ══════════════════════════════════════════════════════════════════
    # DEPLOY
    # ══════════════════════════════════════════════════════════════════
    def deploy(self, output_dir: str = "automl_output/api") -> str:
        """Generate a FastAPI deployment scaffold."""
        from .exporter import scaffold_api

        # Find exported model
        import os
        model_path = os.path.join(
            self.config.output_dir, "exported", "model_scripted.pt"
        )
        if not os.path.exists(model_path):
            self.logger.info("No exported model found; exporting first...")
            exported = self.export()
            model_path = exported.get("torchscript", exported.get("onnx", ""))

        return scaffold_api(model_path, output_dir)

    # ══════════════════════════════════════════════════════════════════
    # TUNE
    # ══════════════════════════════════════════════════════════════════
    def tune_and_train(self) -> dict[str, Any]:
        """Run hyperparameter tuning via Optuna, then train with best params."""
        from .tuner import tune_hyperparameters, suggest_params
        from .data_loader import load_data
        from .detector import detect_data_type, detect_task
        from .model_zoo import build_model, list_models
        from .trainer import Trainer

        self.data_type = self.config.data_type or detect_data_type(self.config.data_path)
        data = load_data(self.config, self.data_type)
        self.task = self.config.task or detect_task(data_type=self.data_type)
        self.num_classes = data["num_classes"]
        self.input_shape = data["input_shape"]

        def objective(trial):
            params = suggest_params(trial)
            cfg = AutoMLConfig(
                data_path=self.config.data_path,
                task=self.task,
                data_type=self.data_type,
                target_column=self.config.target_column,
                epochs=min(self.config.epochs, 10),  # quick eval
                batch_size=params["batch_size"],
                learning_rate=params["learning_rate"],
                optimizer=params["optimizer"],
                layers=[params["hidden_1"], params["hidden_2"]],
                dropout=params["dropout"],
                weight_decay=params["weight_decay"],
                seed=self.config.seed,
                output_dir=self.config.output_dir,
                early_stopping_patience=3,
                verbose=False,
            )
            candidates = list_models(self.data_type, self.task)
            model = build_model(
                candidates[0], input_shape=self.input_shape,
                num_classes=self.num_classes, config=cfg,
            )
            trainer = Trainer(model, cfg)
            results = trainer.train(
                data["train_loader"], data["val_loader"],
                num_classes=self.num_classes, task=self.task,
            )
            return max(results["history"]["val_acc"]) if self.task == "classification" else -min(results["history"]["val_loss"])

        direction = "maximize" if self.task == "classification" else "minimize"
        best_params = tune_hyperparameters(objective, self.config, direction)
        self.logger.info(f"Best HP: {best_params}")

        # Apply best params and do full training
        self.config.learning_rate = best_params.get("learning_rate", self.config.learning_rate)
        self.config.batch_size = best_params.get("batch_size", self.config.batch_size)
        self.config.optimizer = best_params.get("optimizer", self.config.optimizer)
        self.config.dropout = best_params.get("dropout", self.config.dropout)
        self.config.layers = [best_params.get("hidden_1", 256), best_params.get("hidden_2", 128)]

        return self.train()
