from collections import defaultdict

import networkx as nx # Library for displaying graphs.

import matplotlib.pyplot as plt

import random

class DependencyScheduler$($object$)$:

def $\_\_$init$\_\_($self$)$:

self.tasks $=$ set$()$

# The successors of a task are the tasks that depend on it$,$ and can

# only be done once the task is completed.

self.successors $=$ defaultdict$($set$)$

# The predecessors of a task have to be done before the task.

self.predecessors $=$ defaultdict$($set$)$

self.completed$\_$tasks $=$ set$()$ # completed tasks

def add$\_$task$($self$,$ t$,$ dependencies$)$:

$"""$Adds a task t with given dependencies."""

# Makes sure we know about all tasks mentioned.

assert t not in self.tasks or len$($self$.$predecessors$[$t$])==0,$ "The task was already present."

self.tasks.add$($t$)$

self.tasks.update$($dependencies$)$

# The predecessors are the tasks that need to be done before.

self.predecessors$[$t$]=$ set$($dependencies$)$

# The new task is a successor of its dependencies.

for u in dependencies:

self.successors$[$u$].$add$($t$)$

def reset$($self$)$:

self.completed$\_$tasks $=$ set$()$

@property

def done$($self$)$:

return self.completed$\_$tasks $==$ self.tasks

def show$($self$)$:

$"""$We use the nx graph to display the graph."""

g $=$ nx$.$DiGraph$()$

g$.$add$\_$nodes$\_$from$($self$.$tasks$)$

g$.$add$\_$edges$\_$from$([($u$,$ v$)$ for u in self.tasks for v in self.successors$[$u$]])$

node$\_$colors $=[(\text{'}$green$\text{'}$ if v in self.completed$\_$tasks

else 'red' if v in self.available$\_$tasks

else 'gray'$)$

for v in self.tasks$]$

nx$.$draw$($g$,$ with$\_$labels$=$True, node$\_$color$=$node$\_$colors$)$

plt$.$show$()$

@property

def uncompleted$($self$)$:

$"""$Returns the tasks that have not been completed.

This is a property, so you can say scheduler.uncompleted rather than

scheduler.uncompleted$()"""$

return self.tasks $-$ self.completed$\_$tasks

@property

def available$\_$tasks$($self$)$:

$"""$Returns the available tasks.

This is just a placeholder; we will let you implement this."""

return set$()$

def mark$\_$completed$($self$,$ t$)$:

$"""$Marks that the task t is completed, and returns the set of tasks

that become available as a consequence.

This is just a placeholder, we will let you implement this."""

return set$()$

def $\_$check$($self$)$:

$"""$We check that if t is a successor of u$,$ then u is a predecessor

of t$."""$

for u in self.tasks:

for t in self.successors$[$u$]$:

assert u in self.predecessors$[$t$]$

#@title Implementation of $`$available$\_$tasks$`$ and $`$mark$\_$completed$`.$

def scheduler$\_$available$\_$tasks$($self$)$:

$"""$Returns the set of tasks that can be done in parallel.

A task can be done if all its predecessors have been completed.

And of course, we don't return any task that has already been

completed."""

### YOUR SOLUTION HERE

def scheduler$\_$mark$\_$completed$($self$,$ t$)$:

$"""$Marks the task t as completed, and returns the additional

set of tasks that can be done $($and that could not be

previously done$)$ once t is completed."""

### YOUR SOLUTION HERE

DependencyScheduler.available$\_$tasks $=$ property$($scheduler$\_$available$\_$tasks$)$

DependencyScheduler.mark$\_$completed $=$ scheduler$\_$mark$\_$completed

s $=$ DependencyScheduler$()$

s$.$add$\_$task$(\text{'}$a$\text{'},[\text{'}$b$\text{'},\text{'}$c$\text{'}])$

s$.$add$\_$task$(\text{'}$b$\text{'},[\text{'}$c$\text{'},\text{'}$e$\text{'}])$

s$.\_$check$()$

s$.$show$()$

# Tests $5$ points: Simple tests.

s $=$ DependencyScheduler$()$

s$.$add$\_$task$(\text{'}$a$\text{'},[])$

assert s$.$available$\_$tasks $==\{\text{'}$a$\text{'}\}$

# Tests $5$ points: Slightly more complicated.

s $=$ DependencyScheduler$()$

s$.$add$\_$task$(\text{'}$a$\text{'},[\text{'}$b$\text{'},\text{'}$c$\text{'}])$

s$.$add$\_$task$(\text{'}$b$\text{'},[\text{'}$c$\text{'},\text{'}$e$\text{'}])$

assert s$.$available$\_$tasks $==\{\text{'}$e$\text{'},\text{'}$c$\text{'}\}$

s $=$ DependencyScheduler$()$

s$.$add$\_$task$(\text{'}$a$\text{'},[\text{'}$b$\text{'}])$

s$.$add$\_$task$(\text{'}$b$\text{'},[\text{'}$a$\text{'}])$

assert s$.$available$\_$tasks $==$ set$()$

# Tests $5$ points: Now, let's test $`$mark$\_$completed$`.$ Simple tests first.

s $=$ DependencyScheduler$()$

s$.$add$\_$task$(\text{'}$a$\text{'},[])$

assert s$.$available$\_$tasks $==\{\text{'}$a$\text{'}\}$

r $=$ s$.$mark$\_$completed$(\text{'}$a$\text{'})$

assert r $==$ set$()$

s $=$ DependencyScheduler$()$

s$.$add$\_$task$(\text{'}$a$\text{'},[\text{'}$b$\text{'}])$

assert s$.$available$\_$tasks $==\{\text{'}$b$\text{'}\}$

r $=$ s$.$mark$\_$completed$(\text{'}$b$\text{'})$

assert r $==\{\text{'}$a$\text{'}\}$

# Tests $5$ points: Slightly more complicated.

s $=$ DependencyScheduler$()$

s$.$add$\_$task$(\text{'}$a$\text{'},[\text{'}$b$\text{'},\text{'}$c$\text{'}])$

assert s$.$available$\_$tasks $==\{\text{'}$b$\text{'},\text{'}$c$\text{'}\}$

r $=$ s$.$mark$\_$completed$(\text{'}$b$\text{'})$

assert r $==$ set$()$

assert s$.$available$\_$tasks $==\{\text{'}$c$\text{'}\}$

r $=$ s$.$mark$\_$completed$(\text{'}$c$\text{'})$

assert r $==\{\text{'}$a$\text{'}\}$

s $=$ DependencyScheduler$()$

s$.$add$\_$task$(\text{'}$a$\text{'},[\text{'}$b$\text{'},\text{'}$c$\text{'}])$

s$.$add$\_$task$(\text{'}$b$\text{'},[\text{'}$c$\text{'},\text{'}$e$\text{'}])$

s$.$add$\_$task$(\text{'}$c$\text{'},[])$

assert s$.$available$\_$tasks $==\{\text{'}$c$\text{'},\text{'}$e$\text{'}\}$

r $=$ s$.$mark$\_$completed$(\text{'}$e$\text{'})$

assert r $==$ set$()$

r $=$ s$.$mark$\_$completed$(\text{'}$c$\text{'})$

assert r $==\{\text{'}$b$\text{'}\}$

r $=$ s$.$mark$\_$completed$(\text{'}$b$\text{'})$

assert r $==\{\text{'}$a$\text{'}\}$

r $=$ s$.$mark$\_$completed$(\text{'}$a$\text{'})$

assert r $==$ set$()$

assert s$.$available$\_$tasks $==$ set$()$