#!/usr/bin/env python

"""
Try out different partition assignment algorithms on a simulation trace, and
evaluate their performance.
"""

import argparse
import cvxpy
import math
import matplotlib.pyplot as plt
import numpy as np
import os
import pickle

from assign_partitions import AssignmentAlgorithm

def ParseArguments():
    parser = argparse.ArgumentParser(description="Plot times")
    parser.add_argument('--input', type=str, required=True,
                        help="Directory for input files")
    parser.add_argument('--output', type=str, required=True,
                        help="Directory for output files")
    parser.add_argument('--workers', type=int, required=True,
                        help="Number of workers")
    parser.add_argument('--cores', type=int, required=True,
                        help="Number of cores per worker")
    parser.add_argument('--window', type=int, default=10,
                        help="Window (horizon) for analysis")
    return parser.parse_args()


"""
Compute lower bounds by computing solution to convex opt problem relaxation.
"""
class Relaxation(AssignmentAlgorithm):
    def __init__(self, num_workers, num_cores_per_worker, partition_time,
                 assignment_window):
        super(Relaxation, self).__init__(
          num_workers, num_cores_per_worker, partition_time, -1)
        self.assignment_window = assignment_window
        self.partition_time = partition_time

    """
    Formulate the convex optimization problem and solve it.
    """
    def ComputeFractionalAssignment(self, start_iter, end_iter):
        print("Defining problem ...")
        num_workers = self.num_workers
        num_partitions = self.num_partitions
        partition_time = self.partition_time[:,start_iter:end_iter]
        partition_mean = partition_time.mean(axis=1).flatten()
        order = np.flip(np.argsort(partition_mean), axis=0)
        subset1 = order[0:num_workers]
        subset2 = order[num_partitions/2:-1]
        #print(partition_mean[subset1])
        #print(partition_mean[subset2])
        cost_list = list()
        constraint_list = list()
        A = cvxpy.Variable(num_workers, num_partitions)
        worker_matrix = A*partition_time
        for t in range(end_iter-start_iter):
            cost = cvxpy.max_entries(worker_matrix[:,t])
            cost_list.append(cost)
        for p in range(num_partitions):
            constraint_list.append(cvxpy.sum_entries(A[:,p]) == 1)
        for i,p in enumerate(subset1):
            w = i%num_workers
            constraint_list.append(A[w,p] == 1)
        for i,p in enumerate(subset1):
            w = i%num_workers
            constraint_list.append(A[w,p] == 1)
        constraint_list.append(A >= 0)
        constraint_list.append(A <= 1)
        objective = cvxpy.Minimize(sum(cost_list))
        prob = cvxpy.Problem(objective, constraint_list)
        print("Solving problem ...")
        prob.solve()
        print("Problem solved ...")
        return(A.value, prob.value)

    """
    Compute fractional assignments for each window, for each partition.
    """
    def AssignPartitions(self):
        print("Applying assignments using relaxation approximation ...")
        num_windows = int(math.ceil(float(self.num_iterations)/self.assignment_window))
        best_case_cost = self.partition_time.sum()/self.num_workers
        cost = 0
        for i in range(num_windows):
            start_iter = i*self.assignment_window
            end_iter = min(start_iter + self.assignment_window, self.num_iterations)
            A, c = self.ComputeFractionalAssignment(start_iter, end_iter)
            cost += c
        A = A.T
        self.partition_to_worker[:,start_iter:end_iter] = np.argmax(A, axis=1)
        print("Cost ratio = %f " % (cost/best_case_cost))


if __name__ == '__main__':
    args = ParseArguments()
    # Load data.
    partition_time = np.load(os.path.join(args.input, 'partition_time.npy'))
    # Solve relaxed problem.
    relaxation = Relaxation(
      args.workers, args.cores, partition_time[:,1:], args.window)
    relaxation.RunEvaluation(
      os.path.join(args.output, 'relaxation_approx'))
