"""tfts Trainer"""

from collections.abc import Iterable
from contextlib import nullcontext
import logging
import os
from typing import Any, Callable, Dict, Generator, List, Optional, Tuple, Union

import numpy as np
import pandas as pd
import tensorflow as tf
from tensorflow.keras.layers import Input

from .constants import CONFIG_NAME, TF2_WEIGHTS_INDEX_NAME, TF2_WEIGHTS_NAME, TF_WEIGHTS_NAME, TFTS_HOME, TFTS_HUB_CACHE
from .models.base import BaseModel
from .training_args import TrainingArguments

__all__ = ["Trainer", "KerasTrainer", "Seq2seqKerasTrainer"]


logger = logging.getLogger(__name__)


class BaseTrainer(object):
    """Trainer for pipeline"""

    def __init__(
        self,
        model: Union[tf.keras.Model, "BaseModel"],
        args: Optional[TrainingArguments] = None,
        strategy: Optional[tf.distribute.Strategy] = None,
        **kwargs,
    ):
        self.model = model
        self.config = model.config if hasattr(model, "config") else None
        self.args = args or TrainingArguments(output_dir=TFTS_HUB_CACHE)
        self.strategy = strategy

        # with self.get_strategy_scope(strategy):
        #     self.model = self._setup_model(model)
        #     self.loss_fn = loss_fn
        #     self.metrics = metrics or []
        #     self.optimizer = optimizer or self._create_optimizer()
        #     self.lr_scheduler = lr_scheduler or self._create_lr_scheduler()
        #
        #     # Training state
        #     self.global_step = tf.Variable(0, trainable=False, dtype=tf.int32)
        #     if self.args.fp16:
        #         self._setup_mixed_precision()

    def evaluate(self):
        pass

    def get_train_dataloader(self):
        return

    def get_eval_dataloader(self):
        return

    def get_test_dataloader(self):
        return

    def get_learning_rates(self):
        return

    def create_accelerator_and_postprocess(self):
        return

    def get_strategy_scope(self):
        return self.strategy.scope() if self.strategy else nullcontext()

    def _create_optimizer(self) -> tf.keras.optimizers.Optimizer:
        """Create optimizer with specified parameters."""
        return tf.keras.optimizers.Adam(
            learning_rate=self.args.learning_rate,
            beta_1=self.args.adam_beta1,
            beta_2=self.args.adam_beta2,
            epsilon=self.args.adam_epsilon,
            weight_decay=self.args.weight_decay,
        )

    def _create_lr_scheduler(self) -> Optional[tf.keras.optimizers.schedules.LearningRateSchedule]:
        """Create learning rate scheduler based on arguments."""
        if self.args.lr_scheduler_type == "linear":
            return tf.keras.optimizers.schedules.PolynomialDecay(
                initial_learning_rate=self.args.learning_rate,
                decay_steps=self.args.max_steps if self.args.max_steps > 0 else self.args.num_train_epochs,
                end_learning_rate=0,
                power=1.0,
            )
        return None

    def _setup_mixed_precision(self) -> None:
        """Configure mixed precision training."""
        policy = tf.keras.mixed_precision.Policy("mixed_float16")
        tf.keras.mixed_precision.set_global_policy(policy)

    # def _setup_ema(self) -> None:
    #     """Configure Exponential Moving Average if enabled."""
    #     self.ema = None
    #     if self.config.use_ema:
    #         self.ema = tf.train.ExponentialMovingAverage(self.config.ema_decay)

    def get_inputs(self, train_dataset):
        if isinstance(train_dataset, tf.data.Dataset):
            # choose the first batch
            x = next(iter(train_dataset.take(1).as_numpy_iterator()))[0]
            inputs = self._prepare_inputs_for_model(x)

        elif isinstance(train_dataset, tf.keras.utils.Sequence):
            x, _ = train_dataset[0]
            inputs = self._prepare_inputs_for_model(x)

        elif isinstance(train_dataset, (list, tuple)):
            x = train_dataset[0]
            inputs = self._prepare_inputs_for_model(x)
        else:
            raise ValueError("Unsupported dataset type. Expected tf.data.Dataset, keras.utils.Sequence, or list/tuple.")
        return inputs

    def _prepare_inputs_for_model(
        self, x: Union[np.ndarray, pd.DataFrame]
    ) -> Union[Dict[str, tf.keras.layers.Input], List[tf.keras.layers.Input], tf.keras.layers.Input]:
        """
        Prepares the input layer(s) based on the shape of the provided data.

        Args:
            x: Input data (either a NumPy array or a Pandas DataFrame).

        Returns:
            The corresponding Keras Input layers.
        """
        if isinstance(x, dict):
            logger.debug("Preparing inputs from dict")
            return {key: Input(shape=item.shape[1:], name=key) for key, item in x.items()}
        elif isinstance(x, (list, tuple)):
            logger.debug("Preparing inputs from list or tuple")
            return [Input(shape=item.shape[1:], name=f"input_{i}") for i, item in enumerate(x)]
        else:
            logger.debug("Preparing single input")
            return Input(shape=x.shape[1:], name="input")

    def _save(self, output_dir: Optional[str] = None):
        output_dir = output_dir if output_dir is not None else TFTS_HOME
        os.makedirs(output_dir, exist_ok=True)
        logger.info(f"Saving model checkpoint to {output_dir}")
        # save_model = self.model.model if hasattr(self.model, "model") else self.model
        # self.model.save_pretrained(output_dir)

        # model save (due to after build_model, the model will be replaced to a tf.keras.model)
        save_directory = output_dir
        if os.path.isfile(save_directory):
            logger.error(f"Provided path ({save_directory}) should be a directory, not a file")
            return

        os.makedirs(save_directory, exist_ok=True)
        self.config.architectures = [self.model.__class__.__name__[2:]]
        self.config.save_pretrained(save_directory)

        weights_file = os.path.join(save_directory, TF2_WEIGHTS_NAME)  # Or the appropriate extension

        try:
            self.model.save_weights(weights_file)
            logging.info(f"Model weights successfully saved in {weights_file}")
        except Exception as e:
            logging.error(f"Failed to save model weights to {weights_file}: {e}")
            return


class KerasTrainer(BaseTrainer):
    """Keras trainer from tf.keras"""

    def __init__(
        self,
        model: Union[tf.keras.Model, "BaseModel"],
        strategy: Optional[tf.distribute.Strategy] = None,
        args: Optional[TrainingArguments] = None,
        **kwargs: Dict[str, object],
    ) -> None:
        """
        Initializes the trainer with the model, loss function, optimizer, and other optional parameters.

        Args:
            model: A Keras Model or Sequential instance to train.
            strategy: Optional distribution strategy for multi-GPU or multi-node training.
            **kwargs: Additional arguments that are passed to the instance as attributes.
        """
        super().__init__(model, args, strategy, **kwargs)
        self.model = model
        self.config = model.config if hasattr(model, "config") else None

        for key, value in kwargs.items():
            setattr(self, key, value)

    def train(
        self,
        train_dataset: Union[tf.data.Dataset, List[tf.Tensor], Tuple[tf.Tensor, tf.Tensor]],
        valid_dataset: Optional[Union[tf.data.Dataset, List[tf.Tensor], Tuple[tf.Tensor, tf.Tensor]]] = None,
        loss_fn: Union[Callable, tf.keras.losses.Loss, str] = "mse",
        optimizer: Union[tf.keras.optimizers.Optimizer, str, Dict] = "adam",
        epochs: int = 10,
        batch_size: int = 64,
        steps_per_epoch: Optional[int] = None,
        metrics: Optional[Union[List[tf.keras.metrics.Metric], List[str]]] = None,
        callbacks: Optional[List[tf.keras.callbacks.Callback]] = None,
        run_eagerly: bool = True,
        verbose: int = 1,
        **kwargs: Dict[str, object],
    ) -> tf.keras.callbacks.History:
        """
        Trains the model on the provided dataset.

        Args:
            train_dataset: A tf.data.Dataset or list/tuple of tensors (x_train, y_train).
            valid_dataset: A tf.data.Dataset or list/tuple of tensors (x_valid, y_valid), optional.
            loss_fn: A callable or Keras loss function. Default is MeanSquaredError.
            optimizer: A Keras optimizer instance. Default is Adam with learning rate 0.003.
            epochs: Number of epochs to train for. Default is 10.
            batch_size: Number of samples per batch. Default is 64.
            steps_per_epoch: Number of steps per epoch. Optional.
            metrics:  List of metrics for monitoring during training. Optional.
            callbacks: List of keras callbacks during training. Optional.
            run_eagerly: Whether to run eagerly. Default is True.
            verbose: Verbosity level. Default is 1.
            **kwargs: Additional keyword arguments for callbacks.

        Returns:
            A History object containing training logs.
        """

        if not callbacks:
            callbacks: List[tf.keras.callbacks.Callback] = []

        with self.get_strategy_scope():
            # if lr_scheduler:  # just set Optimizer(learning_rate=lr_scheduler)
            #     callbacks.append(tf.keras.callbacks.LearningRateScheduler(lr_scheduler, verbose=True))

            if isinstance(optimizer, (str, dict)):
                optimizer = tf.keras.optimizers.get(optimizer)

            if not isinstance(self.model, tf.keras.Model):
                inputs = self.get_inputs(train_dataset)
                if "build_model" not in dir(self.model):
                    raise TypeError("Trainer model should either be `tf.keras.Model` or has `build_model()` method")
                self.model = self.model.build_model(inputs=inputs)

            self.model.compile(loss=loss_fn, optimizer=optimizer, metrics=metrics, run_eagerly=run_eagerly)

            trainable_params = np.sum([tf.keras.backend.count_params(w) for w in self.model.trainable_weights])
            tf.print(f"Trainable parameters: {trainable_params}")

            if isinstance(train_dataset, (list, tuple)):
                x_train, y_train = train_dataset

                history = self.model.fit(
                    x_train,
                    y_train,
                    validation_data=valid_dataset,
                    steps_per_epoch=steps_per_epoch,
                    epochs=epochs,
                    batch_size=batch_size,
                    verbose=verbose,
                    callbacks=callbacks,
                )
            else:
                history = self.model.fit(
                    train_dataset,
                    validation_data=valid_dataset,
                    steps_per_epoch=steps_per_epoch,
                    epochs=epochs,
                    batch_size=batch_size,
                    verbose=verbose,
                    callbacks=callbacks,
                )
        return history

    def fit(self, **params):
        return self.train(**params)

    def predict(self, x_test: tf.Tensor) -> tf.Tensor:
        return self.model(x_test)

    def get_model(self) -> tf.keras.Model:
        return self.model

    def save_model(self, output_dir: Optional[str] = None):
        # save the model, checkpoint_dir if you use Checkpoint callback to save your best weights
        output_dir = TFTS_HOME if output_dir is None else output_dir
        self._save(output_dir)

    def plot(self, history, true: np.ndarray, pred: np.ndarray):
        import matplotlib.pyplot as plt

        train_length = history.shape[1]
        pred_length = true.shape[1]
        example = np.random.choice(range(history.shape[0]))

        plt.plot(range(train_length), history[example, :, 0], label="History")
        plt.plot(range(train_length, train_length + pred_length), true[example, :, 0], label="True")
        plt.plot(range(train_length, train_length + pred_length), pred[example, :, 0], label="Predicted")
        plt.legend()


class Seq2seqKerasTrainer(KerasTrainer):
    """As the transformers forum mentioned: https://discuss.huggingface.co/t/trainer-vs-seq2seqtrainer/3145/2
    Seq2SeqTrainer is mostly about predict_with_generate."""

    def __init__(self, *args, **kwargs):
        super(Seq2seqKerasTrainer, self).__init__(*args, **kwargs)


class Trainer(object):
    """Custom trainer for tensorflow with support for CPU, GPU, and multi-GPU."""

    def __init__(
        self,
        model,
        strategy: Optional[tf.distribute.Strategy] = None,
        **kwargs: Dict[str, Any],
    ) -> None:
        self.model = model
        self.strategy = strategy

        for key, value in kwargs.items():
            setattr(self, key, value)

    def train(
        self,
        train_loader: Union[tf.data.Dataset, Generator],
        valid_loader: Union[tf.data.Dataset, Generator, None] = None,
        loss_fn: Union[Callable] = tf.keras.losses.MeanSquaredError(),
        optimizer: tf.keras.optimizers.Optimizer = tf.keras.optimizers.Adam(0.003),
        lr_scheduler: Optional[tf.keras.optimizers.schedules.LearningRateSchedule] = None,
        epochs: int = 10,
        learning_rate: float = 3e-4,
        verbose: int = 1,
        eval_metric: Union[Callable, List[Callable], None] = None,
        model_dir: Optional[str] = None,
        use_ema: bool = False,
        stop_no_improve_epochs: Optional[int] = None,
        max_grad_norm: float = 5.0,
        transform: Optional[Callable] = None,
    ) -> None:
        """
        Trains the model using the provided data loaders.

        Parameters
        ----------
        train_loader : Union[tf.data.Dataset, Generator]
            The training data loader, which can be a `tf.data.Dataset` or a Python generator.
        valid_loader : Union[tf.data.Dataset, Generator, None], optional
            The validation data loader, by default None.
        epochs : int, optional
            The number of epochs to train the model, by default 10.
        learning_rate : float, optional
            The initial learning rate for the optimizer, by default 3e-4.
        verbose : int, optional
            The verbosity level (0 = silent, 1 = progress bar, 2 = one line per epoch), by default 1.
        eval_metric : Union[Callable, List[Callable], None], optional
            The evaluation metric(s) to use for validation, by default None.
        model_dir : Optional[str], optional
            The directory to save the model weights, by default "../weights".
        use_ema : bool, optional
            Whether to use exponential moving average (EMA) for the model weights, by default False.
        stop_no_improve_epochs : Optional[int], optional
            If provided, training will stop if the validation metric does not improve for the specified
            number of epochs, by default None.
        max_grad_norm : float, optional
            the max gradient while backprop.
        transform : Optional[Callable], optional
            A function to transform the data before feeding it to the model, by default None.
        """
        self.loss_fn = loss_fn
        self.optimizer = optimizer
        self.lr_scheduler = lr_scheduler
        self.learning_rate = learning_rate
        self.eval_metric = eval_metric if isinstance(eval_metric, Iterable) else [eval_metric]
        self.use_ema = use_ema
        self.transform = transform
        self.max_grad_norm = max_grad_norm
        self.global_step = tf.Variable(0, trainable=False, dtype=tf.int32)

        if use_ema:
            self.ema = tf.train.ExponentialMovingAverage(0.9).apply(self.model.trainable_variables)

        if model_dir is None:
            model_dir = TFTS_HUB_CACHE

        if stop_no_improve_epochs is not None:
            no_improve_epochs: int = 0
            best_metric: float = float("inf")

        if not isinstance(self.model, tf.keras.Model):
            if "build_model" not in dir(self.model):
                raise TypeError("Trainer model should either be tf.keras.Model or has the build_model method")
            x = list(train_loader.take(1).as_numpy_iterator())[0][0]
            if isinstance(x, dict):
                inputs = {key: Input(item.shape[1:]) for key, item in x.items()}
            else:
                inputs = Input(x.shape[1:])
            self.model = self.model.build_model(inputs=inputs)

        for epoch in range(epochs):
            train_loss, train_scores = self.train_loop(train_loader)
            log_str = f"Epoch: {epoch + 1}, Train Loss: {train_loss:.4f}"  # noqa

            if valid_loader is not None:
                valid_loss, valid_scores = self.valid_loop(valid_loader)
                log_str += f", Valid Loss: {valid_loss:.4f}"  # noqa
                log_str + ",".join([" Valid Metrics{}: {:.4f}".format(i, me) for i, me in enumerate(valid_scores)])

                if (stop_no_improve_epochs is not None) and (eval_metric is not None):
                    if valid_scores[0] >= best_metric:
                        best_metric = valid_scores[0]
                        no_improve_epochs = 0
                    else:
                        no_improve_epochs += 1
                    if no_improve_epochs >= stop_no_improve_epochs:
                        logger.info("Tried the best, no improved and stop training")
                        break

            logger.info(log_str)

        # self.export_model(model_dir, only_pb=True)  # save the model

    def fit(self, **params):
        return self.train(**params)

    def train_loop(self, train_loader):
        train_loss: float = 0.0
        y_trues, y_preds = [], []

        for step, (x_train, y_train) in enumerate(train_loader):
            y_pred, step_loss = self.train_step(x_train, y_train)
            train_loss += step_loss
            y_preds.append(y_pred)
            y_trues.append(y_train)

        scores = []
        if self.eval_metric is not None:
            y_preds = tf.concat(y_preds, axis=0)
            y_trues = tf.concat(y_trues, axis=0)

            for metric in self.eval_metric:
                scores.append(metric(y_trues, y_preds))
        return train_loss / (step + 1), scores

    def train_step(self, x_train, y_train):
        with tf.GradientTape() as tape:
            y_pred = self.model(x_train, training=True)
            loss = self.loss_fn(y_train, y_pred)

        gradients = tape.gradient(loss, self.model.trainable_variables)
        gradients = [(tf.clip_by_value(grad, -self.max_grad_norm, self.max_grad_norm)) for grad in gradients]
        _ = self.optimizer.apply_gradients(zip(gradients, self.model.trainable_variables))

        if self.lr_scheduler is not None:
            lr = self.lr_scheduler(self.global_step)
            self.optimizer.learning_rate.assign(lr)
        else:
            lr = self.learning_rate
        self.optimizer.learning_rate.assign(lr)
        self.global_step.assign_add(1)
        # logger.info(f'Step: {self.global_step.numpy()}, Loss: {loss}'
        return y_pred, loss

    def valid_loop(self, valid_loader):
        valid_loss: float = 0.0
        y_valid_trues, y_valid_preds = [], []

        for valid_step, (x_valid, y_valid) in enumerate(valid_loader):
            y_valid_pred, valid_step_loss = self.valid_step(x_valid, y_valid)
            valid_loss += valid_step_loss
            y_valid_trues.append(y_valid)
            y_valid_preds.append(y_valid_pred)

        valid_scores = []
        if self.eval_metric:
            y_valid_preds = tf.concat(y_valid_preds, axis=0)
            y_valid_trues = tf.concat(y_valid_trues, axis=0)

            for metric in self.eval_metric:
                valid_scores.append(metric(y_valid_trues, y_valid_preds))
        return valid_loss / (valid_step + 1), valid_scores

    def valid_step(self, x_valid, y_valid):

        y_valid_pred = self.model(x_valid, training=False)
        valid_loss = self.loss_fn(y_valid, y_valid_pred)
        return y_valid_pred, valid_loss

    def predict(self, test_loader):
        y_test_trues, y_test_preds = [], []
        for x_test, y_test in test_loader:
            y_test_pred = self.model(x_test, training=False)
            y_test_preds.append(y_test_pred)
            y_test_trues.append(y_test)

        y_test_trues = tf.concat(y_test_trues, axis=0)
        y_test_preds = tf.concat(y_test_preds, axis=0)
        return tf.squeeze(y_test_trues, axis=-1), y_test_preds

    def save_model(self, model_dir, only_pb=True):
        # save the model
        if not model_dir.endswith(".keras"):
            model_dir = f"{model_dir}.keras"
        self.model.save(model_dir)
        logger.info(f"Model successfully saved in {model_dir}")

        if not only_pb:
            self.model.save_weights(f"{model_dir}.ckpt")
            logger.info(f"Model weights successfully saved in {model_dir}.ckpt")
