import glob
import random

import tqdm


PROGRESS_BAR_FORMAT = "{desc} - {percentage:3.0f}% |{bar:40}| {n_fmt}/{total_fmt} \
    [{elapsed}<{remaining},{rate_fmt}{postfix}]"


def format_num(num, digit=4):
    """Formatting the float values as string."""
    rounding = f".{digit}g"
    num_str = f"{num:{rounding}}".replace("+0", "+").replace("-0", "-")
    num = str(num)
    return num_str if len(num_str) < len(num) else num


def gen_bar_postfix(history, targets=["loss", "accuracy"], name_transform=["loss", "acc"]):  # pylint: disable=W0102
    """Generate string of metrics status to be appended to the end of the progress bar.

    Parameters
    ----------
    history: dict
        History records generated by ``train_steps``.
    targets: list[str]
        List of metric's names to be extracted as the postfix.
    name_transform: list[str]
        The alias metric name that will be showed on the bar.
        Should be the same length, same order as ``targets``.

    Returns
    -------
    postfix: str
        The extracted metrics information.
    """
    info = []
    for target, name in zip(targets, name_transform):
        if target not in history:
            continue
        val = history[target]
        val_str = format_num(val)
        info_str = f"{name}: {val_str}"
        info.append(info_str)
    return ", ".join(info)


def train_steps(model, dataset, steps=None, bar_title=None, validate=False):
    """A single training epoch with multiple steps.

    Customized training epoch compared to the built-in ``.fit(...)`` function
    of tensorflow keras model. The major difference is that the ``.fit()``
    requires the dataset to yield either (feature, target) or
    (feature, target, weight) pairs, which losses the flexibility of yielding
    different numbers of elements for each iteration. And thus we'd decide to
    implement our own training logic and relevant utilities same as provided
    in tensorflow like `callbacks`.

    Parameters
    ----------
    model:
        Compiled tf.keras model.
    dataset:
        The loaded tf.data.Dataset object that yields (feature, target) pairs
        at the first two elements, indicating that you can yields more than
        two elements for each iteration, but only the first two will be used
        for training.
    steps: int
        Total number of steps that the dataset object will yield. This is used
        for visualizing the training progress.
    bar_title: str
        Additional title to be printed at the start of the progress bar.
    validate: bool
        Indicating whether it is now in validation stage or it is within
        training loop that should update the weights of the model.

    Returns
    -------
    history: dict
        The history of scores for each metric during each epoch.
    """
    iter_bar = tqdm.tqdm(dataset, total=steps, desc=bar_title, bar_format=PROGRESS_BAR_FORMAT)

    for iters, data in enumerate(iter_bar):
        feat, label = data[:2]  # Assumed the first two elements are feature and label, respectively.
        if validate:
            step_result = model.test_on_batch(feat, label, return_dict=True)
        else:
            step_result = model.train_on_batch(feat, label, return_dict=True)

        if iters == 0:
            # model.metrics_names is only available after the first train_on_batch
            metrics = model.metrics_names
            history = {metric: 0 for metric in metrics}
            history.update({f"{metric}_sum": 0 for metric in metrics})

        for metric in metrics:
            history[f"{metric}_sum"] += step_result[metric]
            history[metric] = history[f"{metric}_sum"] / (iters + 1)
        iter_bar.set_postfix_str(gen_bar_postfix(history))

    # Remove metric_sum columns in the history
    history = {metric: history[metric] for metric in metrics}
    return history


def execute_callbacks(callbacks, func_name, **kwargs):
    """Execute callbacks at different training stage."""
    if callbacks is not None:
        for callback in callbacks:
            getattr(callback, func_name)(**kwargs)


def train_epochs(
    model,
    train_dataset,
    validate_dataset=None,
    epochs=10,
    steps=100,
    val_steps=100,
    callbacks=None,
    **kwargs
):
    """Logic of training loop.

    The main loop of the training, with events-based life-cycle management
    that triggers different events for all callbacks. Event types are the
    same as the original tensorflow implementation.

    Event types and their order:

    .. code-block:: none

        <start training>
         |
         |-on_train_begin
        T|  |-on_epoch_begin
        R|  |
        A| L|-on_train_batch_begin
        I| O|-on_train_batch_end
        N| O|
        I| P|-on_test_batch_begin
        N|  |-on_test_batch_end
        G|  |
         |  |-on_epoch_end
         |-on_train_end
         |
        <finished training>

    Parameters
    ----------
    model:
        Compiled tensorflow keras model.
    train_dataset:
        The tf.data.Dataset instance for training.
    validate_dataset:
        The tf.data.Dataset instance for validation. If not given, validation
        stage will be skipped.
    epochs: int
        Number of maximum training epochs.
    steps: int
        Number of training steps for each epoch. Should be the same as
        when initiating the dataset instance.
    val_steps: int
        Number of validation steps for each epoch.Should be the same as
        when initiating the dataset instance.
    callbacks:
        List of callback instances.

    Returns
    -------
    history: dict
        Score history of each metrics during each epoch of both training
        and validation.

    See Also
    --------
    omnizart.callbacks:
        Implementation and available callbacks for training.
    """
    history = {"train": [], "validate": []}
    execute_callbacks(callbacks, "_set_model", model=model)
    execute_callbacks(callbacks, "on_train_begin")
    for epoch_idx in range(epochs):
        # Epoch begin
        execute_callbacks(callbacks, "on_epoch_begin", epoch=epoch_idx+1)  # noqa: E226
        if model.stop_training:
            break

        print(f"Epoch: {epoch_idx+1}/{epochs}")

        # Train batch begin
        execute_callbacks(callbacks, "on_train_batch_begin")
        results = train_steps(model, dataset=train_dataset, steps=steps, bar_title="Train   ", **kwargs)

        # Train batch end
        execute_callbacks(callbacks, "on_train_batch_end")
        history["train"].append(results)

        # Test batch begin
        execute_callbacks(callbacks, "on_test_batch_begin")
        val_results = {}
        if validate_dataset is not None:
            val_results = train_steps(
                model, dataset=validate_dataset, steps=val_steps, validate=True, bar_title="Validate", **kwargs
            )

            # Test batch end
            execute_callbacks(callbacks, "on_test_batch_end")
            history["validate"].append(val_results)

        # Epoch end
        execute_callbacks(callbacks, "on_epoch_end", epoch=epoch_idx+1, history=history)  # noqa: E226

    execute_callbacks(callbacks, "on_train_end")
    return history


def get_train_val_feat_file_list(feature_folder, split=0.9):
    feat_files = glob.glob(f"{feature_folder}/*.hdf")
    sidx = round(len(feat_files) * split)
    random.shuffle(feat_files)
    train_files = feat_files[:sidx]
    val_files = feat_files[sidx:]
    return train_files, val_files
