"""
Author: Weisen Pan
Date: 2023-10-24
"""
# multiprocessing_run in tools is used for multi-process mode training; 
# average_slowdown and average_completion are used to extract calculation
# statistics from an Episode class object

import time
import numpy as np
import tensorflow as tf


def average_metric(exp, metric_func):
    """
    Calculate the average of a given metric for all tasks.

    Args:
    - exp: The simulation experiment.
    - metric_func: Function to compute the desired metric.

    Returns:
    - float: Average of the metric across tasks.
    """
    total = sum(metric_func(task) for job in exp.simulation.cluster.jobs for task in job.tasks)
    number_of_tasks = sum(1 for job in exp.simulation.cluster.jobs for _ in job.tasks)
    return total / number_of_tasks


def completion_metric(task):
    """
    Compute the completion metric for a task.

    Args:
    - task: The task object.

    Returns:
    - float: Completion metric for the task.
    """
    return task.finished_timestamp - task.started_timestamp


def slowdown_metric(task):
    """
    Compute the slowdown metric for a task.

    Args:
    - task: The task object.

    Returns:
    - float: Slowdown metric for the task.
    """
    return (task.finished_timestamp - task.started_timestamp) / task.task_config.duration


def average_completion(exp):
    """
    Compute the average completion time for all tasks.

    Args:
    - exp: The simulation experiment.

    Returns:
    - float: Average completion time.
    """
    return average_metric(exp, completion_metric)


def average_slowdown(exp):
    """
    Compute the average slowdown for all tasks.

    Args:
    - exp: The simulation experiment.

    Returns:
    - float: Average slowdown.
    """
    return average_metric(exp, slowdown_metric)


def multiprocessing_run(episode, trajectories, makespans, avg_completions, avg_slowdowns):
    """
    Run an episode in a multiprocessing environment and gather results.

    Args:
    - episode: The simulation episode.
    - trajectories: List to collect trajectories.
    - makespans: List to collect makespans.
    - avg_completions: List to collect average completions.
    - avg_slowdowns: List to collect average slowdowns.
    """
    # Set random seeds
    np.random.seed(int(time.time()))
    tf.random.set_random_seed(time.time())
    
    # Execute the episode
    episode.run()
    
    # Gather and append results to respective lists
    trajectories.append(episode.simulation.scheduler.algorithm.current_trajectory)
    makespans.append(episode.simulation.env.now)
    avg_completions.append(average_completion(episode))
    avg_slowdowns.append(average_slowdown(episode))