from pprint import pprint
import typing
import enum
import cvxpy as cvx
from cvxpy import *
#from cvxpy import Minimize
import operator
from functools import reduce
import numpy as np
from texttable import Texttable
from collections import defaultdict

# Create two scalar optimization variables.
constraints = []

tasks = 2

@enum.unique
class WorkloadClass(enum.Enum):
	LowLatency = enum.auto()
	Opportunistic = enum.auto()


class CPU(typing.NamedTuple):
	core_id: int
	package_id: int
	thread_id: int

class Workload(typing.NamedTuple):
	core_count: int
	workload_class: WorkloadClass = WorkloadClass.LowLatency

class Objective(typing.NamedTuple):
	objective: typing.Union[Minimize, Maximize]
	priority: int
	target: int


cores = [
	CPU(0, 0, 0),
	CPU(1, 0, 0),
	CPU(2, 0, 1),
	CPU(3, 0, 1),
	CPU(4, 0, 2),
	CPU(5, 0, 2),
	CPU(6, 0, 3),
	CPU(7, 0, 3),
	CPU(8, 1, 4),
	CPU(9, 1, 4),
	CPU(10, 1, 5),
	CPU(11, 1, 5),
	CPU(12, 1, 6),
	CPU(13, 1, 6),
	CPU(14, 1, 7),
	CPU(15, 1, 7)
]

cores_by_thread = defaultdict(list)
cores_by_package = defaultdict(list)
for i in cores:
	cores_by_thread[i.thread_id].append(i)

for i in cores:
	cores_by_package[i.package_id].append(i)

print(cores_by_thread)
print(cores_by_package)

workloads = [
	Workload(2),
	Workload(3),
	Workload(6)
]


# 2 tasks, up to 8 cores per task
assignments = cvx.Variable((len(workloads), len(cores)), boolean=True)
constraints = []
objectives = []

for core_assignments in cvx.atoms.affine.transpose.transpose(assignments):
	# It's priority 10000 to try to avoid scheduling multiple workloads on a single core
	objectives.append(Objective(Minimize(cvx.sum(core_assignments)), 10, 1))

for idx, workload in enumerate(workloads):
	constraints.append(cvx.sum(assignments[idx]) == workload.core_count)

# Now we try to make sure tha a given workload tries to stay on one package, then hyperthead
#for assignment in assignments:
for assignment in assignments:
	packages = []
	for _, package_cores in cores_by_package.items():
		packages.append(cvx.maximum(*[assignment[core.core_id] for core in package_cores]))
	objectives.append(Objective(Minimize(reduce(operator.add, packages)), 9, 1))


for assignment in assignments:
	threads = []
	for _, package_threads in cores_by_thread.items():
		threads.append(cvx.maximum(*[assignment[core.thread_id] for core in package_threads]))
	objectives.append(Objective(Minimize(reduce(operator.add, threads)), 8, 1))


#objective = cvx.transforms.scalarize.targets_and_priorities(*zip(*[(i.objective, i.priority, i.target) for i in objectives]))
objective = cvx.transforms.scalarize.targets_and_priorities(*zip(*[(i.objective, i.priority, i.target) for i in objectives]))
prob = cvx.Problem(objective, constraints)
print(prob.solve())
print(prob.status)
print(prob.value)
print(assignments.value)
#print(np.clip(assignments.value, 0, 1))
result = np.round(assignments.value).astype(int)

table = Texttable()
labels = ["Cores", "Package", "Thread"] + ["Task {}".format(task_num) for task_num in range(result.shape[0])]
table.add_row(labels)

print(cores)
for core_num, i in enumerate(result.transpose()):
	print(core_num)
	c = cores[core_num]
	table.add_row(['core {}'.format(c.core_id), 'package_id {}'.format(c.package_id), 'thread_id {}'.format(c.thread_id)] + list(i))
print(table.draw())