"""

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 trainiqConfig

>>>

>>> config = trainiqConfig(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.3"

__author__ = "AutoML Team"


from typing import Any


import numpy as np


from .config import trainiqConfig

from .utils import get_logger, set_seed

from .research import WebResearchEngine, research


class trainiq:

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


    Parameters

    ----------

    config : trainiqConfig

        Configuration object. See :class:`trainiqConfig` for all options.


    Example

    -------

    >>> automl = AutoML(trainiqConfig(data_path="my_dataset.csv"))

    >>> results = automl.train()

    >>> automl.predict(new_data)
    """


    def __init__(self, config: trainiqConfig | None = None, **kwargs):

        if config is None:

            config = trainiqConfig(**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 = trainiqConfig(

                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()

