#!/usr/bin/env python3
# Copyright 2010-2021 Google LLC
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Scheduling cooking tasks in a bakery using no-wait jobshop scheduling.

We are scheduling a full day of baking:
  - Work starts at 4am.
  - The shop closes at 10PM.
  - Durations are in minutes. Time starts at midnight.
"""

import collections
from typing import Sequence
from absl import app
from absl import flags

from google.protobuf import text_format
from ortools.sat.python import cp_model

FLAGS = flags.FLAGS

flags.DEFINE_string('params', 'num_search_workers:16, max_time_in_seconds:30',
                    'Sat solver parameters.')
flags.DEFINE_string('proto_file', '',
                    'If not empty, output the proto to this file.')

# Recipes
CROISSANT = 'croissant'
APPLE_PIE = 'apple pie'
BRIOCHE = 'brioche'
CHOCOLATE_CAKE = 'chocolate cake'

# Skills
BAKING = 'baking'
PROOFING = 'proofing'
COOKING = 'cooking'
COOLING = 'cooling'
DECORATING = 'decorating'
DISPLAY = 'display'


class Task(object):
    """A unit baking task.

  - Simple baking tasks have a fixed duration. They are performed by workers.
  - Waiting/cooling/proofing tasks have a min and a max duration.
    They are performed by machine or they use space resources.
  """

    def __init__(self, name, min_duration, max_duration):
        self.name = name
        self.min_duration = min_duration
        self.max_duration = max_duration


class Skill(object):
    """The skill of a worker or the capability of a machine."""

    def __init__(self, name, efficiency):
        self.name = name
        # Efficiency is currently not used.
        self.efficiency = efficiency


class Recipe(object):
    """A recipe is a sequence of cooking tasks."""

    def __init__(self, name):
        self.name = name
        self.tasks = []

    def add_task(self, resource_name, min_duration, max_duration):
        self.tasks.append(Task(resource_name, min_duration, max_duration))
        return self


class Resource(object):
    """A resource is a worker, a machine, or just some space for cakes to rest.

  - Workers have a capacity of 1 and can have variable efficiency.
  - Machines and spaces have a capacity greater or equal to one, but the
    efficiency is fixed to 100.

    For a worker with efficiency k and a task of duration t, the resulting
    work will have a duration `ceil(t * k)`.
  """

    def __init__(self, name, capacity):
        self.name = name
        self.capacity = capacity
        self.skills = []

    def add_skill(self, skill_name, efficiency):
        self.skills.append(Skill(skill_name, efficiency))
        return self


class Order(object):
    """An order is a recipe that should be delivered at a given due date."""

    def __init__(self, unique_id, recipe_name, due_date, quantity):
        """Builds an order.

    Args:
      unique_id: A unique identifier for the order. Used to display the result.
      recipe_name: The name of the recipe. It must match one of the recipes.
      due_date: The due date in minutes since midnight.
      quantity: How many cakes to prepare.
    """
        self.unique_id = unique_id
        self.recipe_name = recipe_name
        self.due_date = due_date
        self.quantity = quantity


def set_up_data():
    """Set up the bakery problem data."""

    # Recipes.
    croissant_recipe = Recipe(CROISSANT)
    croissant_recipe.add_task(BAKING, 15, 15)
    croissant_recipe.add_task(PROOFING, 60, 90)
    croissant_recipe.add_task(COOKING, 20, 20)
    croissant_recipe.add_task(DISPLAY, 5, 5 * 60)

    apple_pie_recipe = Recipe(APPLE_PIE)
    apple_pie_recipe.add_task(BAKING, 25, 25)
    apple_pie_recipe.add_task(PROOFING, 15, 60)
    apple_pie_recipe.add_task(COOKING, 30, 30)
    apple_pie_recipe.add_task(DECORATING, 5, 5)
    apple_pie_recipe.add_task(DISPLAY, 5, 5 * 60)

    brioche_recipe = Recipe(BRIOCHE)
    brioche_recipe.add_task(BAKING, 20, 20)
    brioche_recipe.add_task(PROOFING, 60, 90)
    brioche_recipe.add_task(COOKING, 30, 30)
    brioche_recipe.add_task(DISPLAY, 5, 5 * 60)

    chocolate_cake_recipe = Recipe(CHOCOLATE_CAKE)
    chocolate_cake_recipe.add_task(BAKING, 15, 15)
    chocolate_cake_recipe.add_task(COOKING, 25, 25)
    chocolate_cake_recipe.add_task(DECORATING, 15, 15)
    chocolate_cake_recipe.add_task(DISPLAY, 5, 5 * 60)
    recipes = [
        croissant_recipe, apple_pie_recipe, brioche_recipe,
        chocolate_cake_recipe
    ]

    # Resources.
    baker1 = Resource('baker1', 1).add_skill(BAKING, 1.0)
    baker2 = Resource('baker2', 1).add_skill(BAKING, 1.0)
    decorator1 = Resource('decorator1', 1).add_skill(DECORATING, 1.0)
    waiting_space = Resource('waiting_space', 4).add_skill(PROOFING, 1.0)
    oven = Resource('oven', 4).add_skill(COOKING, 1.0)
    display_space = Resource('display_space', 12).add_skill(DISPLAY, 1.0)
    resources = [baker1, baker2, decorator1, waiting_space, oven, display_space]

    # Orders
    croissant_7am = Order('croissant_7am', CROISSANT, 7 * 60, 3)
    croissant_8am = Order('croissant_8am', CROISSANT, 8 * 60, 3)
    croissant_9am = Order('croissant_9am', CROISSANT, 9 * 60, 2)
    croissant_10am = Order('croissant_10am', CROISSANT, 10 * 60, 1)
    croissant_11am = Order('croissant_11am', CROISSANT, 11 * 60, 1)
    brioche_10am = Order('brioche_10am', BRIOCHE, 10 * 60, 8)
    brioche_12pm = Order('brioche_12pm', BRIOCHE, 12 * 60, 8)
    apple_pie_1pm = Order('apple_pie_1pm', APPLE_PIE, 13 * 60, 10)
    chocolate_4pm = Order('chocolate_4pm', CHOCOLATE_CAKE, 16 * 60, 10)
    orders = [
        croissant_7am, croissant_8am, croissant_9am, croissant_10am,
        croissant_11am, brioche_10am, brioche_12pm, apple_pie_1pm, chocolate_4pm
    ]

    return recipes, resources, orders


def solve_with_cp_sat(recipes, resources, orders):
    """Build the optimization model, and solve the problem."""

    model = cp_model.CpModel()
    horizon = 22 * 60  # 10PM.
    start_work = 4 * 60  # 4am.

    # Parse recipes.
    recipe_by_name = {}
    for recipe in recipes:
        recipe_by_name[recipe.name] = recipe

    # Parse resources.
    resource_by_name = {}
    resource_list_by_skill_name = collections.defaultdict(list)
    for resource in resources:
        resource_by_name[resource.name] = resource
        for skill in resource.skills:
            resource_list_by_skill_name[skill.name].append(resource)

    # Parse orders and create one optional copy per eligible resource and per
    # task.
    interval_list_by_resource_name = collections.defaultdict(list)
    orders_sequence_of_events = collections.defaultdict(list)
    sorted_orders = []
    tardiness_vars = []
    for order in orders:
        for batch in range(order.quantity):
            order_id = f'{order.unique_id}_{batch}'
            sorted_orders.append(order_id)
            previous_end = None
            due_date = order.due_date
            recipe = recipe_by_name[order.recipe_name]
            for task in recipe.tasks:
                skill_name = task.name
                suffix = f'_{order.unique_id}_batch{batch}_{skill_name}'

                start = None
                if previous_end is None:
                    start = model.NewIntVar(start_work, horizon,
                                            f'start{suffix}')
                    orders_sequence_of_events[order_id].append(
                        (start, f'start{suffix}'))
                else:
                    start = previous_end

                size = model.NewIntVar(task.min_duration, task.max_duration,
                                       f'size{suffix}')
                end = None
                if task == recipe.tasks[-1]:
                    # The order must end after the due_date. Ideally, exactly at the
                    # due_date.
                    tardiness = model.NewIntVar(0, horizon - due_date,
                                                f'end{suffix}')
                    end = tardiness + due_date

                    # Store the end_var for the objective.
                    tardiness_vars.append(tardiness)
                else:
                    end = model.NewIntVar(start_work, horizon, f'end{suffix}')
                orders_sequence_of_events[order_id].append(
                    (end, f'end{suffix}'))
                previous_end = end

                # Per resource copy.
                presence_literals = []
                for resource in resource_list_by_skill_name[skill_name]:
                    presence = model.NewBoolVar(
                        f'presence{suffix}_{resource.name}')
                    copy = model.NewOptionalIntervalVar(
                        start, size, end, presence,
                        f'interval{suffix}_{resource.name}')
                    interval_list_by_resource_name[resource.name].append(copy)
                    presence_literals.append(presence)

                # Only one copy will be performed.
                model.AddExactlyOne(presence_literals)

    # Create resource constraints.
    for resource in resources:
        intervals = interval_list_by_resource_name[resource.name]
        if resource.capacity == 1:
            model.AddNoOverlap(intervals)
        else:
            model.AddCumulative(intervals, [1] * len(intervals),
                                resource.capacity)

    # The objective is to minimize the sum of the tardiness values of each jobs.
    # The tardiness is difference between the end time of an order and its
    # due date.
    model.Minimize(sum(tardiness_vars))

    # Solve model.
    solver = cp_model.CpSolver()
    if FLAGS.params:
        text_format.Parse(FLAGS.params, solver.parameters)
    solver.parameters.log_search_progress = True
    status = solver.Solve(model)

    if status == cp_model.OPTIMAL or status == cp_model.FEASIBLE:
        for order_id in sorted_orders:
            print(f'{order_id}:')
            for time_expr, event_id in orders_sequence_of_events[order_id]:
                time = solver.Value(time_expr)
                print(f'  {event_id} at {time // 60}:{time % 60:02}')


def main(argv: Sequence[str]) -> None:
    if len(argv) > 1:
        raise app.UsageError('Too many command-line arguments.')

    recipes, resources, orders = set_up_data()
    solve_with_cp_sat(recipes, resources, orders)


if __name__ == '__main__':
    app.run(main)
