# -*- coding: utf-8 -*-
# file: pipeline.py
# time: 18:38 12/04/2024
# author: YANG, HENG <hy345@exeter.ac.uk> (杨恒)
# github: https://github.com/yangheng95
# huggingface: https://huggingface.co/yangheng
# google scholar: https://scholar.google.com/citations?user=NPq5a_0AAAAJ&hl=en
# Copyright (C) 2019-2025. All Rights Reserved.
"""
This module provides the Pipeline class for creating and managing complete
machine learning workflows that combine models, tokenizers, datasets, and
trainers. Pipelines provide a unified interface for training, inference,
and model management.
"""

import json
import os

import autocuda
from transformers import AutoConfig, AutoTokenizer

from ..hub_utils import download_pipeline
from ..model_hub.model_hub import ModelHub
from omnigenbench.src.abc.abstract_model import OmniModel
from omnigenbench.src.misc.utils import env_meta_info, fprint
from omnigenbench.src.trainer.trainer import Trainer


class Pipeline:
    """
    The Pipeline class provides a unified interface for managing complete machine
    learning workflows. It handles model initialization, training, inference, and
    persistence. Pipelines can be loaded from pre-built configurations or created
    from scratch with custom components.

    Attributes:
        model (OmniModel): The underlying model for the pipeline.
        tokenizer: Tokenizer for preprocessing input sequences.
        dataset (dict): Dictionary containing train/validation/test datasets.
        metadata (dict): Environment and pipeline metadata.
        trainer (Trainer): Trainer instance for model training.
        device (str): Target device for model execution (CPU/GPU).
        name (str): Name identifier for the pipeline.

    Example:
        >>> from omnigenbench import Pipeline, OmniModelForSequenceClassification
        >>> # Create pipeline from model
        >>> model = OmniModelForSequenceClassification("model_path", tokenizer)
        >>> pipeline = Pipeline("my_pipeline", config_or_model=model)
        >>> # Use for inference
        >>> predictions = pipeline("ATCGATCG")
        >>> # Train the model
        >>> pipeline.train(datasets)
        >>> # Save pipeline
        >>> pipeline.save("./saved_pipeline")

    Note:
        - Pipelines automatically handle device placement and model optimization
        - Environment metadata is collected for reproducibility
        - Pipelines can be saved and loaded for easy deployment
        - Supports both local models and hub-based model loading
    """

    model: OmniModel = None
    tokenizer = None
    dataset: dict = None
    metadata: dict = None

    def __init__(
        self,
        name,
        *,
        config_or_model,
        tokenizer=None,
        datasets=None,
        trainer=None,
        **kwargs,
    ):
        """
        Initialize a Pipeline instance.

        Args:
            name (str): Name identifier for the pipeline.
            config_or_model (Union[str, OmniModel]): Model to use in the pipeline.
                Can be a string path/identifier or an OmniModel instance.
            tokenizer (optional): Tokenizer for preprocessing. If None, will be
                loaded from the model or model path. Defaults to None.
            datasets (dict, optional): Dictionary containing train/validation/test
                datasets. Keys should be 'train', 'valid', 'test'. Defaults to None.
            trainer (Trainer, optional): Trainer instance for model training.
                If None, a default trainer will be created. Defaults to None.
            **kwargs: Additional keyword arguments including:
                - device (str): Target device for model execution
                - trust_remote_code (bool): Whether to trust remote code in tokenizers
                - Other model-specific configuration parameters

        Raises:
            ValueError: If model initialization fails.
            ImportError: If required dependencies are not available.
            FileNotFoundError: If model path is invalid.

        Example:
            >>> # Create from model path
            >>> pipeline = Pipeline("rna_classification",
            ...                    config_or_model="yangheng/OmniGenome-186M")
            >>> # Create from model instance
            >>> model = OmniModelForSequenceClassification("model_path", tokenizer)
            >>> pipeline = Pipeline("custom_pipeline", config_or_model=model)

        Note:
            - The pipeline automatically handles model loading and device placement
            - Environment metadata is collected for tracking system information
            - If a model instance is provided, its tokenizer and metadata are used
        """
        self.metadata = env_meta_info()
        self.name = name
        self.tokenizer = tokenizer
        self.datasets = datasets
        self.trainer = trainer
        self.device = (
            autocuda.auto_cuda()
            if kwargs.get("device") is None
            else kwargs.get("device")
        )
        if not isinstance(config_or_model, str):
            self.model = config_or_model
            self.tokenizer = self.model.tokenizer
            self.metadata = self.model.metadata
        else:
            self.init_pipeline(
                config_or_model=config_or_model, tokenizer=tokenizer, **kwargs
            )

        self.model.to(self.device)

    def __call__(self, inputs, *args, **kwargs):
        """
        Call the pipeline for inference.

        This method provides a convenient interface for running inference
        through the pipeline. It delegates to the model's inference method.

        Args:
            inputs: Input data for inference (can be string, list, or tensor).
            *args: Additional positional arguments passed to model inference.
            **kwargs: Additional keyword arguments passed to model inference.

        Returns:
            dict: Inference results including predictions and confidence scores.

        Example:
            >>> pipeline = Pipeline("my_pipeline", config_or_model=model)
            >>> results = pipeline("ATCGATCG")
            >>> print(results['predictions'])
        """
        return self.model.inference(inputs, **kwargs)

    def to(self, device):
        """
        Move the pipeline to a specific device.

        Args:
            device (str): Target device ('cpu', 'cuda', 'cuda:0', etc.).

        Returns:
            Pipeline: Self for method chaining.

        Example:
            >>> pipeline = Pipeline("my_pipeline", config_or_model=model)
            >>> pipeline.to("cuda:0")  # Move to GPU
            >>> pipeline.to("cpu")     # Move to CPU
        """
        self.model.to(device)
        self.device = device
        return self

    def init_pipeline(self, *, config_or_model, tokenizer=None, **kwargs):
        """
        This method handles loading the model, tokenizer, and configuration
        from a model path or identifier. It tries to load from the ModelHub
        first, then falls back to HuggingFace transformers.

        Args:
            config_or_model (str): Path or identifier of the model to load.
            tokenizer (optional): Tokenizer instance. If None, will be loaded
                from the model path. Defaults to None.
            **kwargs: Additional keyword arguments for model loading including:

                - trust_remote_code (bool): Whether to trust remote code
                - device (str): Target device for the model
                - Other model-specific parameters

        Returns:
            Pipeline: Self for method chaining.

        Raises:
            ValueError: If model loading fails.
            ImportError: If required dependencies are not available.

        Example:
            >>> pipeline = Pipeline("my_pipeline")
            >>> pipeline.init_pipeline(config_or_model="yangheng/OmniGenome-186M")

        Note:
            - First attempts to load from OmniGenome ModelHub
            - Falls back to HuggingFace transformers if ModelHub fails
            - Automatically handles tokenizer loading and configuration
        """
        trust_remote_code = kwargs.get("trust_remote_code", True)
        try:  # for the models saved by OmniGenome and served by the model hub
            self.model = ModelHub.load(config_or_model, **kwargs)
            self.tokenizer = self.model.tokenizer
            self.metadata.update(self.model.metadata)
        except Exception as e:
            fprint(f"Fail to load the model from the model hub, the error is: {e}")

            config = AutoConfig.from_pretrained(
                config_or_model, trust_remote_code=trust_remote_code
            )
            if tokenizer is None:
                tokenizer = AutoTokenizer.from_pretrained(
                    config_or_model, trust_remote_code=trust_remote_code
                )
            self.model = OmniModel.from_pretrained(
                config_or_model,
                config=config,
                tokenizer=tokenizer,
                trust_remote_code=trust_remote_code,
                **kwargs,
            )
            self.tokenizer = self.model.tokenizer
            self.metadata.update(self.model.metadata)
        fprint(f"The pipeline has been initialized from {config_or_model}.")
        return self

    def train(self, datasets: dict = None, trainer=None, **kwargs):
        """
        This method initiates training of the model using the provided datasets
        and trainer configuration. If no trainer is provided, the pipeline's
        existing trainer will be used.

        Args:
            datasets (dict, optional): Dictionary containing train/validation/test
                datasets. If None, uses the pipeline's existing datasets.
                Keys should be 'train', 'valid', 'test'. Defaults to None.
            trainer (Trainer, optional): Trainer instance to use for training.
                If None, uses the pipeline's existing trainer. Defaults to None.
            **kwargs: Additional keyword arguments passed to the trainer.

        Raises:
            ValueError: If no trainer is available or datasets are invalid.
            RuntimeError: If training fails.

        Example:
            >>> pipeline = Pipeline("my_pipeline", config_or_model=model)
            >>> # Train with existing datasets
            >>> pipeline.train()
            >>> # Train with custom datasets
            >>> custom_datasets = {'train': train_data, 'valid': valid_data}
            >>> pipeline.train(datasets=custom_datasets)
            >>> # Train with custom trainer
            >>> from omnigenbench import Trainer
            >>> custom_trainer = Trainer(model, train_dataset=train_data)
            >>> pipeline.train(trainer=custom_trainer)

        Note:
            - Training uses the pipeline's current model and device
            - Progress and metrics are logged during training
            - The trained model is automatically saved in the pipeline
        """
        if trainer is not None:
            assert isinstance(trainer, Trainer)
            self.trainer = trainer

        self.trainer.train()

    def predict(self, inputs, **kwargs):
        """
        This method provides a high-level interface for generating predictions
        from the pipeline's model. It handles preprocessing and postprocessing
        automatically.

        Args:
            inputs: Input data for prediction. Can be:

                - str: Single sequence string
                - list: List of sequence strings
                - tensor: Preprocessed input tensors
            **kwargs: Additional keyword arguments passed to model prediction.

        Returns:
            dict: Prediction results including:
                - predictions: Predicted labels or values
                - confidence: Confidence scores (if available)
                - logits: Raw model outputs (if requested)

        Example:
            >>> pipeline = Pipeline("my_pipeline", config_or_model=model)
            >>> # Single prediction
            >>> result = pipeline.predict("ATCGATCG")
            >>> print(result['predictions'])
            >>> # Batch prediction
            >>> results = pipeline.predict(["ATCGATCG", "GCTAGCTA"])
            >>> print(results['predictions'])

        Note:
            - Input preprocessing is handled automatically
            - Results are formatted consistently across different model types
            - Confidence scores are included when available
        """
        return self.model.predict(inputs, **kwargs)

    def inference(self, inputs, **kwargs):
        """
        This method provides the complete inference pipeline including
        preprocessing, model forward pass, and postprocessing. It's the
        recommended method for production inference.

        Args:
            inputs: Input data for inference. Can be:
                - str: Single sequence string
                - list: List of sequence strings
                - tensor: Preprocessed input tensors
            **kwargs: Additional keyword arguments for inference including:
                - return_attention: Whether to return attention weights
                - return_hidden_states: Whether to return hidden states
                - temperature: Temperature for sampling (if applicable)

        Returns:
            dict: Complete inference results including:
                - predictions: Final predictions
                - confidence: Confidence scores
                - attention: Attention weights (if requested)
                - hidden_states: Hidden states (if requested)

        Example:
            >>> pipeline = Pipeline("my_pipeline", config_or_model=model)
            >>> # Basic inference
            >>> results = pipeline.inference("ATCGATCG")
            >>> print(results['predictions'])
            >>> # Inference with attention
            >>> results = pipeline.inference("ATCGATCG", return_attention=True)
            >>> print(results['attention'].shape)

        Note:
            - This is the most comprehensive inference method
            - Handles all preprocessing and postprocessing automatically
            - Returns rich information about the model's internal states
        """
        return self.model.inference(inputs, **kwargs)

    @staticmethod
    def load(pipeline_name_or_path, local_only=False, **kwargs):
        """
        This static method loads a complete pipeline including model, tokenizer,
        datasets, and trainer from a saved pipeline directory or hub identifier.

        Args:
            pipeline_name_or_path (str): Path to saved pipeline directory or
                hub identifier for downloading.
            local_only (bool, optional): If True, only load from local paths.
                If False, download from hub if not found locally. Defaults to False.
            **kwargs: Additional keyword arguments for pipeline initialization:
                - device: Target device for the model
                - name: Custom name for the pipeline
                - trust_remote_code: Whether to trust remote code

        Returns:
            Pipeline: Loaded pipeline instance ready for use.

        Raises:
            FileNotFoundError: If pipeline cannot be found locally and
                local_only is True.
            ValueError: If pipeline files are corrupted or invalid.
            ImportError: If required dependencies are not available.

        Example:
            >>> # Load from local path
            >>> pipeline = Pipeline.load("./saved_pipeline")
            >>> # Load from hub
            >>> pipeline = Pipeline.load("yangheng/OmniGenome-RNA-Classification")
            >>> # Use loaded pipeline
            >>> results = pipeline("ATCGATCG")

        Note:
            - Loads all pipeline components (model, tokenizer, datasets, trainer)
            - Automatically handles device placement
            - Preserves all training configurations and metadata
        """
        import dill

        if os.path.exists(pipeline_name_or_path):
            path = pipeline_name_or_path
        else:
            path = download_pipeline(
                pipeline_name_or_path, local_only=local_only, **kwargs
            )
        with open(f"{path}/datasets.pkl", "rb") as f:
            datasets = dill.load(f)
        with open(f"{path}/trainer.pkl", "rb") as f:
            trainer = dill.load(f)
        model = ModelHub.load(path, local_only=local_only, **kwargs)
        tokenizer = model.tokenizer
        pipeline = Pipeline(
            name=(
                pipeline_name_or_path
                if kwargs.get("name") is None
                else kwargs.get("name")
            ),
            config_or_model=model,
            tokenizer=tokenizer,
            datasets=datasets,
            trainer=trainer,
            **kwargs,
        )
        return pipeline

    def save(self, path, overwrite=False, **kwargs):
        """
        This method saves the complete pipeline including model, tokenizer,
        datasets, trainer, and metadata to a directory. The saved pipeline
        can be loaded later using Pipeline.load().

        Args:
            path (str): Directory path where to save the pipeline.
            overwrite (bool, optional): If True, overwrite existing directory.
                If False, raise error if directory exists. Defaults to False.
            **kwargs: Additional keyword arguments for model saving.

        Raises:
            FileExistsError: If path exists and overwrite is False.
            OSError: If there are issues creating the directory or writing files.
            RuntimeError: If saving fails due to model or data issues.

        Example:
            >>> pipeline = Pipeline("my_pipeline", config_or_model=model)
            >>> # Train the pipeline
            >>> pipeline.train(datasets)
            >>> # Save the trained pipeline
            >>> pipeline.save("./trained_pipeline", overwrite=True)
            >>> # Load the saved pipeline later
            >>> loaded_pipeline = Pipeline.load("./trained_pipeline")

        Note:
            - Saves all pipeline components (model, tokenizer, datasets, trainer)
            - Preserves training configurations and metadata
            - Model is temporarily moved to CPU during saving to avoid GPU memory issues
            - Creates a complete, self-contained pipeline directory
        """
        import dill

        if os.path.exists(path) and not overwrite:
            raise FileExistsError(
                f"The path {path} already exists, please set overwrite=True to overwrite it."
            )
        if not os.path.exists(path):
            os.makedirs(path)
        device = self.model.model.device
        self.model.model.to("cpu")
        with open(f"{path}/datasets.pkl", "wb") as f:
            dill.dump(self.datasets, f)
        with open(f"{path}/metadata.json", "w") as f:
            json.dump(self.metadata, f, indent=2)
        with open(f"{path}/tokenizer.pkl", "wb") as f:
            dill.dump(self.tokenizer, f)
        with open(f"{path}/trainer.pkl", "wb") as f:
            dill.dump(self.trainer, f)
        self.model.save(path, overwrite=overwrite, **kwargs)
        self.model.model.to(device)
