# -*- coding: utf-8 -*-
from numpy import argmax
from typing import Iterable, List, Optional

from ..models import BatchJobSchedule, BatchJob, FixedLengthJobPool, Schedule, TimeInterval
from . import AbstractScheduler


class BatchScheduler(AbstractScheduler):
    """
    Batch scheduling algorithm as it is described in "Optimal Batch Schedules for Parallel Machines" (Koehler and
    Khuller, 2013). The version represented in this class is for the case when jobs are not restricted to be agreeable
    and works for the number of machines m = 1, which corresponds to the Active Time Problem. The algorithm runs in
    O(n^3).
    """

    @staticmethod
    def _construct_initial_schedule(
            job_pool: FixedLengthJobPool,
            max_concurrency: int,
    ) -> Iterable[BatchJobSchedule]:
        jobs = sorted(job_pool.jobs)
        used = set()

        while len(used) != len(jobs):
            batch = None

            for job in reversed(jobs):
                if job in used:
                    continue

                if batch is None:
                    batch = BatchJobSchedule(set(), job.release_time, job.release_time + job_pool.duration - 1)

                if job.release_time <= batch.execution_start and batch.execution_end <= job.deadline:
                    batch.jobs.add(job)
                    used.add(job)

                if len(batch.jobs) == max_concurrency:
                    break

            yield batch

    @staticmethod
    def _push_forward(i: int, batches: List[BatchJobSchedule], batch_duration: int, number_of_machines: int) -> None:
        updated_execution_start = batches[i].execution_start

        for job in batches[i].jobs:
            updated_execution_start = max(updated_execution_start, job.release_time)

        if i >= 1:
            updated_execution_start = max(updated_execution_start, batches[i - 1].execution_start)
        if i >= number_of_machines:
            updated_execution_start = max(
                updated_execution_start,
                batches[i - number_of_machines].execution_start + batch_duration,
            )

        batches[i].execution_start = updated_execution_start
        batches[i].execution_end = updated_execution_start + batch_duration - 1

    @staticmethod
    def _move_back(i: int, batches: List[BatchJobSchedule], job: BatchJob, max_concurrency: int) -> Optional[int]:
        batches[i].jobs.remove(job)

        preceding_deadline_available_jobs = [
            (job, j) for j in range(i) for job in batches[j].jobs if batches[i].execution_end <= job.deadline
        ]

        if preceding_deadline_available_jobs:
            job_to_bring_forward, j = max(preceding_deadline_available_jobs)
            batches[j].jobs.remove(job_to_bring_forward)
            batches[i].jobs.add(job_to_bring_forward)

        for j in range(i - 1, -1, -1):
            if len(batches[j].jobs) != max_concurrency:
                batches[j].jobs.add(job)
                return j

            job_with_max_release_time = list(batches[j].jobs)[
                argmax([job.release_time for job in list(batches[j].jobs)])
            ]

            if job_with_max_release_time.release_time <= job.release_time:
                batches[j].jobs.remove(job_with_max_release_time)
                batches[j].jobs.add(job)
                job = job_with_max_release_time

        return None

    @staticmethod
    def _compute_for_multiple_machines(
            job_pool: FixedLengthJobPool,
            max_concurrency: int,
            number_of_machines: int,
    ) -> Optional[List[BatchJobSchedule]]:
        if min([job.deadline - job.release_time + 1 for job in job_pool.jobs]) < job_pool.duration:
            return None

        batches = sorted(BatchScheduler._construct_initial_schedule(job_pool, max_concurrency))

        i = 0

        while True:
            if i == len(batches):
                break

            BatchScheduler._push_forward(i, batches, job_pool.duration, number_of_machines)

            jobs_to_move_back = [
                job for job in batches[i].jobs if batches[i].execution_start + job_pool.duration - 1 > job.deadline
            ]

            if not jobs_to_move_back:
                i += 1
                continue

            next_i = BatchScheduler._move_back(i, batches, jobs_to_move_back[0], max_concurrency)

            if next_i is None:
                return None
            i = next_i

        return [batch for batch in batches if len(batch.jobs) != 0]

    def process(self, job_pool: FixedLengthJobPool, max_concurrency: int) -> Schedule:
        """
        Compute the optimal batch schedule for a given set of jobs and maximum concurrency.
        :param job_pool: Job pool of fixed length jobs.
        :param max_concurrency: Maximum number of jobs allowed to run concurrently.
        :return: Computed batch schedule.
        """
        if job_pool.size == 0:
            return Schedule(True, [], [])

        batches = self._compute_for_multiple_machines(job_pool, max_concurrency, 1)

        if batches is None:
            return Schedule(False, None, None)

        active_time_intervals = [TimeInterval(batch.execution_start, batch.execution_end) for batch in batches]

        return Schedule(
            True,
            TimeInterval.merge_time_intervals(active_time_intervals),
            batches,
        )
