#! /usr/bin/env python
# Copyright (c) 2023 Predibase, Inc., 2020 Uber Technologies, Inc.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================

from __future__ import annotations

from abc import ABC, abstractmethod
from collections.abc import Callable, Generator
from concurrent.futures import ThreadPoolExecutor
from contextlib import contextmanager
from dataclasses import dataclass
from typing import Any, TYPE_CHECKING

import numpy as np
import pandas as pd
import psutil
import torch
from tqdm import tqdm

from ludwig.api_annotations import DeveloperAPI
from ludwig.backend.utils.storage import StorageManager
from ludwig.constants import MODEL_LLM
from ludwig.data.cache.manager import PreprocessedDataCache
from ludwig.data.dataframe.base import DataFrameEngine
from ludwig.data.dataframe.pandas import PANDAS
from ludwig.data.dataset.base import DatasetManager
from ludwig.data.dataset.pandas import PandasDatasetManager
from ludwig.distributed import init_dist_strategy
from ludwig.models.base import BaseModel
from ludwig.schema.trainer import BaseTrainerConfig
from ludwig.types import HyperoptConfigDict
from ludwig.utils.audio_utils import read_audio_from_path
from ludwig.utils.batch_size_tuner import BatchSizeEvaluator
from ludwig.utils.dataframe_utils import from_batches, to_batches
from ludwig.utils.fs_utils import get_bytes_obj_from_path
from ludwig.utils.misc_utils import get_from_registry
from ludwig.utils.system_utils import Resources
from ludwig.utils.torch_utils import initialize_pytorch
from ludwig.utils.types import DataFrame, Series

if TYPE_CHECKING:
    from ludwig.trainers.base import BaseTrainer


@DeveloperAPI
@dataclass(frozen=True)
class BackendCapabilities:
    """Named feature flags that a :class:`Backend` can advertise.

    Use this instead of a raw ``dict[str, Any]`` so that callers get IDE
    completion and type-safe access::

        if backend.capabilities.distributed:
            ...

    All flags default to ``False``; subclasses set them in the class body::

        class RayBackend(Backend):
            capabilities = BackendCapabilities(distributed=True, hyperopt=True)
    """

    distributed: bool = False
    hyperopt: bool = False
    async_execution: bool = False
    cache_preprocessing: bool = True


@DeveloperAPI
class Backend(ABC):
    """Abstract base class for Ludwig execution backends.

    Required abstract methods (must be implemented by every subclass):
        initialize(), initialize_pytorch(), create_trainer(), sync_model(),
        broadcast_return(), is_coordinator(), df_engine,
        supports_multiprocessing, read_binary_files(), num_nodes,
        num_training_workers, get_available_resources(),
        max_concurrent_trials(), tune_batch_size(), batch_transform()

    Optional methods (have sensible defaults; override to add capability):
        supports_batch_size_tuning() — returns True; set to False for
        backends that cannot tune batch sizes (e.g. remote inference only).

    Set the ``capabilities`` class attribute to advertise named feature flags::

        class MyBackend(Backend):
            capabilities = BackendCapabilities(distributed=False)
    """

    capabilities: BackendCapabilities = BackendCapabilities()

    def __init__(
        self,
        dataset_manager: DatasetManager,
        cache_dir: str | None = None,
        credentials: dict[str, dict[str, Any]] | None = None,
    ):
        credentials = credentials or {}
        self._dataset_manager = dataset_manager
        self._storage_manager = StorageManager(**credentials)
        self._cache_manager = PreprocessedDataCache(self._dataset_manager, cache_dir)

    @property
    def storage(self) -> StorageManager:
        return self._storage_manager

    @property
    def cache(self) -> PreprocessedDataCache:
        return self._cache_manager

    @property
    def dataset_manager(self) -> DatasetManager:
        return self._dataset_manager

    @abstractmethod
    def initialize(self):
        raise NotImplementedError()

    @abstractmethod
    def initialize_pytorch(self, *args, **kwargs):
        raise NotImplementedError()

    @contextmanager
    @abstractmethod
    def create_trainer(self, config: BaseTrainerConfig, model: BaseModel, **kwargs) -> Generator:
        raise NotImplementedError()

    @abstractmethod
    def sync_model(self, model):
        raise NotImplementedError()

    @abstractmethod
    def broadcast_return(self, fn):
        raise NotImplementedError()

    @abstractmethod
    def is_coordinator(self):
        raise NotImplementedError()

    @property
    @abstractmethod
    def df_engine(self) -> DataFrameEngine:
        raise NotImplementedError()

    @property
    @abstractmethod
    def supports_multiprocessing(self):
        raise NotImplementedError()

    @abstractmethod
    def read_binary_files(self, column: Series, map_fn: Callable | None = None) -> Series:
        raise NotImplementedError()

    @property
    @abstractmethod
    def num_nodes(self) -> int:
        raise NotImplementedError()

    @property
    @abstractmethod
    def num_training_workers(self) -> int:
        raise NotImplementedError()

    @abstractmethod
    def get_available_resources(self) -> Resources:
        raise NotImplementedError()

    @abstractmethod
    def max_concurrent_trials(self, hyperopt_config: HyperoptConfigDict) -> int | None:
        raise NotImplementedError()

    @abstractmethod
    def tune_batch_size(self, evaluator_cls: type[BatchSizeEvaluator], dataset_len: int) -> int:
        """Returns best batch size (measured in samples / s) on the given evaluator.

        The evaluator class will need to be instantiated on each worker in the backend cluster, then call
        `evaluator.select_best_batch_size(dataset_len)`.
        """
        raise NotImplementedError()

    @abstractmethod
    def batch_transform(
        self, df: DataFrame, batch_size: int, transform_fn: Callable, name: str | None = None
    ) -> DataFrame:
        """Applies `transform_fn` to every `batch_size` length batch of `df` and returns the result."""
        raise NotImplementedError()

    def supports_batch_size_tuning(self) -> bool:
        return True


class LocalDataProcessingMixin:
    @property
    def df_engine(self):
        return PANDAS

    @property
    def supports_multiprocessing(self):
        return True

    @staticmethod
    def read_binary_files(column: pd.Series, map_fn: Callable | None = None, file_size: int | None = None) -> pd.Series:
        column = column.fillna(np.nan).replace([np.nan], [None])  # normalize NaNs to None

        sample_fname = column.head(1).values[0]
        with ThreadPoolExecutor() as executor:  # number of threads is inferred
            if isinstance(sample_fname, str):
                if map_fn is read_audio_from_path:  # bypass torchaudio issue that no longer takes in file-like objects
                    result = executor.map(  # type: ignore[misc]
                        lambda path: map_fn(path) if path is not None else path, column.values
                    )
                else:
                    result = executor.map(
                        lambda path: get_bytes_obj_from_path(path) if path is not None else path, column.values
                    )
            else:
                # If the sample path is not a string, assume the paths has already been read in
                result = column.values

            if map_fn is not None and map_fn is not read_audio_from_path:
                result = executor.map(lambda x: map_fn(x) if x is not None else None, result)

        return pd.Series(result, index=column.index, name=column.name)

    @staticmethod
    def batch_transform(df: DataFrame, batch_size: int, transform_fn: Callable, name: str | None = None) -> DataFrame:
        name = name or "Batch Transform"
        batches = to_batches(df, batch_size)
        transform = transform_fn()
        out_batches = [transform(batch.reset_index(drop=True)) for batch in tqdm(batches, desc=name)]
        out_df = from_batches(out_batches).reset_index(drop=True)
        return out_df


class LocalTrainingMixin:
    @staticmethod
    def initialize():
        init_dist_strategy("local")

    @staticmethod
    def initialize_pytorch(*args, **kwargs):
        initialize_pytorch(*args, **kwargs)

    @staticmethod
    def create_predictor(model: BaseModel, **kwargs):
        from ludwig.models.predictor import get_predictor_cls

        return get_predictor_cls(model.type())(model, **kwargs)  # type: ignore[call-arg]

    def sync_model(self, model):
        pass

    @staticmethod
    def broadcast_return(fn):
        return fn()

    @staticmethod
    def is_coordinator() -> bool:
        return True

    @staticmethod
    def tune_batch_size(evaluator_cls: type[BatchSizeEvaluator], dataset_len: int) -> int:
        evaluator = evaluator_cls()
        return evaluator.select_best_batch_size(dataset_len)


class RemoteTrainingMixin:
    def sync_model(self, model):
        pass

    @staticmethod
    def broadcast_return(fn):
        return fn()

    @staticmethod
    def is_coordinator() -> bool:
        return True


@DeveloperAPI
class LocalBackend(LocalDataProcessingMixin, LocalTrainingMixin, Backend):
    BACKEND_TYPE = "local"

    _shared_instance: LocalBackend

    @classmethod
    def shared_instance(cls) -> LocalBackend:
        """Returns a shared singleton LocalBackend instance."""
        if not hasattr(cls, "_shared_instance"):
            cls._shared_instance = cls()
        return cls._shared_instance

    def __init__(self, **kwargs) -> None:
        super().__init__(dataset_manager=PandasDatasetManager(self), **kwargs)

    @property
    def num_nodes(self) -> int:
        return 1

    @property
    def num_training_workers(self) -> int:
        return 1

    def get_available_resources(self) -> Resources:
        return Resources(cpus=psutil.cpu_count(), gpus=torch.cuda.device_count())

    def max_concurrent_trials(self, hyperopt_config: HyperoptConfigDict) -> int | None:
        # Every trial will be run with Pandas and NO Ray Datasets. Allow Ray Tune to use all the
        # trial resources it wants, because there is no Ray Datasets process to compete with it for CPUs.
        return None

    def create_trainer(
        self,
        config: BaseTrainerConfig,
        model: BaseModel,
        **kwargs,
    ) -> BaseTrainer:  # type: ignore[override]
        from ludwig.trainers.registry import get_llm_trainers_registry, get_trainers_registry

        trainer_cls: type
        if model.type() == MODEL_LLM:
            trainer_cls = get_from_registry(config.type, get_llm_trainers_registry())
        else:
            trainer_cls = get_from_registry(model.type(), get_trainers_registry())

        return trainer_cls(config=config, model=model, **kwargs)