import time
from typing import Union, Callable

from sampo.scheduler.base import SchedulerType, Scheduler
from sampo.scheduler.genetic.base import GeneticScheduler
from sampo.scheduler.heft.base import HEFTBetweenScheduler
from sampo.scheduler.heft.base import HEFTScheduler
from sampo.scheduler.resource.base import ResourceOptimizer
from sampo.scheduler.topological.base import TopologicalScheduler
from sampo.schemas.contractor import Contractor
from sampo.schemas.graph import WorkGraph
from sampo.schemas.landscape import LandscapeConfiguration
from sampo.schemas.schedule import Schedule
from sampo.schemas.time_estimator import WorkTimeEstimator
from sampo.utilities.resource_usage import resources_costs_sum


def get_scheduler_ctor(scheduling_algorithm_type: SchedulerType) \
        -> Callable[[SchedulerType, ResourceOptimizer, WorkTimeEstimator], Scheduler]:
    if scheduling_algorithm_type is SchedulerType.HEFTAddBetween:
        return HEFTBetweenScheduler
    if scheduling_algorithm_type is SchedulerType.HEFTAddEnd:
        return HEFTScheduler
    if scheduling_algorithm_type is SchedulerType.Topological:
        return TopologicalScheduler
    return GeneticScheduler


def generate_schedule(scheduling_algorithm_type: SchedulerType,
                      work_time_estimator: WorkTimeEstimator,
                      work_graph: WorkGraph,
                      contractors: Union[Contractor, list[Contractor]],
                      validate_schedule: bool,
                      landscape: LandscapeConfiguration = LandscapeConfiguration()) -> Schedule:
    scheduler = get_scheduler_ctor(scheduling_algorithm_type)(work_estimator=work_time_estimator)
    start_time = time.time()
    if isinstance(scheduler, GeneticScheduler):
        scheduler.number_of_generation = 2
        scheduler.set_use_multiprocessing(n_cpu=4)

    schedule = scheduler.schedule(work_graph,
                                  [contractors] if isinstance(contractors, Contractor) else contractors,
                                  validate=validate_schedule, landscape=landscape)

    print(f'Time: {(time.time() - start_time) * 1000} ms')
    print(f'Cost: {resources_costs_sum(schedule)}')
    return schedule
