import sys
import networkx as nx
from matrices import FileBackedObjectList, MatrixList, make_zero_matrix, add_with_matrix
from os import path
import json
import argparse
import time
from pyomo.environ import *
from pyomo.opt import SolverFactory as SolverFactory

# compute MWM based on --interval_length miliseconds


class TimeStat:
    def __init__(self, name="mwm"):
        self.name = name
        self.t0 = None
        self.sum = 0
        self.sum_sqr = 0
        self.max = None
        self.min = None
        self.count = 0
        
    def reset(self):
        self.t0 = None
        self.sum = 0
        self.sum_sqr = 0
        self.max = None
        self.min = None
        self.count = 0
        
    def start(self):
        self.t0 = time.time()
        
    def end(self):
        dif = time.time()-self.t0
        self.sum += dif
        self.sum_sqr += dif**2
        if self.count == 0:
            self.max = dif
            self.min = dif
        else:
            self.max = max([dif, self.max])
            self.min = min([dif, self.min])
        self.count += 1
        
    def get_avg(self):
        return self.sum * 1.0 / self.count
    
    def get_var(self):
        avg = self.get_avg()
        return self.sum_sqr * 1.0 / self.count - avg**2
    
    def to_dict(self):
        var = self.get_var()
        return dict(avg=self.get_avg(), var=var, std_err=var**0.5,
                    min=self.min, max=self.max, count=self.count,
                    sum=self.sum, sum_sqr=self.sum_sqr)
        
        
time_stat = TimeStat()


def calc_optimal_allocation(X, m, n, k=1):
    model = ConcreteModel()
    # n = the number of ToRs
    # m = the number of wavelengths
    rangeA = range(n)
    rangeB = range(n)
    model.y = Var(rangeA, rangeB, domain=Binary)
    
    def Y(i, j):
        # return model.y[i][j]
        return model.y[i, j]
    
    model.OBJ = Objective(expr=sum([X[i][j] * Y(i, j) for i in range(n) for j in range(n)]), sense=maximize)
    
    def diagonal_constraint_rule(model, i):
        # return the expression for the constraint for i
        return Y(i, i) == 0
    
    model.DiagonalConstraints = Constraint(range(n), rule=diagonal_constraint_rule)
    
    def symetric_constraint_rule(model, i, j):
        # return the expression for the constraint for i
        return Y(i, j) == Y(j, i)
    
    model.SymetricConstraints = Constraint(range(n), range(n), rule=symetric_constraint_rule)
    
    def ax_constraint_rule(model, i):
        # return the expression for the constraint for i
        # within the demand matrix - it is the same as going "stright" until j and "aside" from j
        # return sum([Y(i,j) for j in range(i+1, n)] + [Y(j,i) for j in range(i)]) <= 1
        return sum([Y(i, j) for j in range(n)]) <= k
    
    model.NodeConstraints = Constraint(range(n), rule=ax_constraint_rule)
    model.WaveConstraint = Constraint(expr=sum([Y(i, j) for i in range(n) for j in range(i + 1, n)]) <= m)
    opt = SolverFactory(
        'glpk')  # https://stackoverflow.com/questions/51802252/how-to-declare-two-dimensional-and-three-dimensional-vector-in-pyomo
    res = opt.solve(model)
    return model.y.get_values()


def compute_mwm_ex(traffic, n_tors, max_degree=1, n_waves=None, **kwargs):
    if n_waves is None:
        n_waves = n_tors * max_degree
    time_stat.start()
    D = calc_optimal_allocation(traffic, m=n_waves, n=n_tors, k=max_degree)
    time_stat.end()
    return list(list(x) for x in {tuple(sorted(k)) for k in D if D[k] == 1})


def compute_mwm_ex_approx(traffic, n_tors, max_degree=1, **kwargs):
    G = nx.Graph()
    G.add_nodes_from(range(n_tors))

    for i in range(n_tors):
        for j in range(n_tors):
            if i < j and (traffic[i][j] > 0 or traffic[j][i] > 0):
                G.add_edge(i, j, weight=traffic[i][j] + traffic[j][i])
                
    time_stat.start()
    total_mwm_res = []
    for i in range(max_degree):
        mwm_res = nx.max_weight_matching(G)
        for u, v in mwm_res:
            if u < v:
                G.remove_edge(u, v)
            else:
                G.remove_edge(v, u)
        total_mwm_res.extend(list(x) for x in mwm_res)
    time_stat.end()
    return total_mwm_res


def compute_mwm(traffic, n_tors):
    G = nx.Graph()
    G.add_nodes_from(range(n_tors))

    for i in range(n_tors):
        for j in range(n_tors):
            if i < j and (traffic[i][j] > 0 or traffic[j][i] > 0):
                G.add_edge(i, j, weight=traffic[i][j] + traffic[j][i])

    time_stat.start()
    mwm_res = nx.max_weight_matching(G)
    time_stat.end()
    mwm_list = [list(x) for x in mwm_res]
    return mwm_list


def compute_optimal_matching(traffic, n_tors, max_degree=1, mwm_approx=False, **kwargs):
    if max_degree == 1:
        return compute_mwm(traffic, n_tors)
    if mwm_approx:
        return compute_mwm_ex_approx(traffic, n_tors, max_degree=max_degree, **kwargs)
    return compute_mwm_ex(traffic, n_tors, max_degree=max_degree, **kwargs)


def compute_top(traffic, n_tors, top):
    for i in range(n_tors):
        for j in sorted(range(n_tors), key=lambda k: traffic[i][k])[:-top]:
            traffic[i][j] = 0


class DictList(FileBackedObjectList):
    def __init__(self, path_pattern, separator=" "):
        super(DictList, self).__init__(path_pattern)
        self.separator = separator
        self.path = self.path_pattern.replace("%", "")
        if path.isfile(self.path):
            self.data = json.load(open(self.path))
        else:
            self.data = {}

    def __getitem__(self, key):
        return self.data[str(key)]

    def __setitem__(self, key, value):
        self.data[str(key)] = value

    def __contains__(self, key):
        return str(key) in self.data
    
    def __del__(self):
        json.dump(self.data, open(self.path, "w"))


def compute_per_mili_mwm(check_existing=True, top=None, n_milis=5000, n_tors=80, output_dir=".", **kwargs):
    per_mili_pattern = path.join(output_dir, "matrix_mili_%d")
    per_mili_matrix_list = MatrixList(per_mili_pattern)
    
    per_mili_match_pattern = path.join(output_dir, "mwm_mili_%d"+("_deg_%d" % kwargs.get("max_degree", 1)))
    per_mili_match_list = DictList(per_mili_match_pattern)
    
    if top:
        per_mili_top_match_pattern = path.join(output_dir, "mwm_mili_%d_top_"+ str(top) + ("_deg_%d" % kwargs.get("max_degree", 1)))
        per_mili_top_match_list = DictList(per_mili_top_match_pattern)
    
    traffic = None
    for z in range(n_milis):
        print "\r", z, "/", n_milis,
        if not check_existing or not z in per_mili_match_list:
            traffic = list(per_mili_matrix_list[z])
            MWM = compute_optimal_matching(traffic, n_tors, **kwargs)
            per_mili_match_list[z] = MWM
        if top and (not check_existing or not (z, top) in per_mili_top_match_list):
            if traffic is None:
                traffic = list(per_mili_matrix_list[z])
            compute_top(traffic, n_tors, top)
            MWM = compute_optimal_matching(traffic, n_tors, **kwargs)
            per_mili_top_match_list[z] = MWM
        traffic = None
    

def compute_intervals_mwm(check_existing=True, interval_length=1, top=None, n_milis=5000, n_tors=80, output_dir=".", **kwargs):
    if interval_length == 1:
        return compute_per_mili_mwm(check_existing=check_existing, top=top, n_milis=n_milis, n_tors=n_tors, output_dir=output_dir, **kwargs)
    per_mili_pattern = path.join(output_dir, "matrix_mili_%d")
    per_mili_matrix_list = MatrixList(per_mili_pattern)
    per_interval_match_pattern = path.join(output_dir, "mwm_agg_%d_%d-%d"+("_deg_%d" % kwargs.get("max_degree", 1)))
    per_interval_match_list = DictList(per_interval_match_pattern)

    if top:
        per_interval_top_match_pattern = path.join(output_dir, "mwm_agg_%d_%d-%d_top_"+str(top)+("_deg_%d" % kwargs.get("max_degree", 1)))
        per_interval_top_match_list = DictList(per_interval_top_match_pattern)
        
    def get_traffic(z):
        traffic = make_zero_matrix(n_tors, n_tors)
        for i in range(interval_length):
            matrix = per_mili_matrix_list[z * interval_length + i]
            add_with_matrix(traffic, matrix)
        return traffic

    traffic = None
    for z in range(n_milis / interval_length):
        print "\r", z, "/", n_milis / interval_length,
        first = z*interval_length
        last = (z+1)*interval_length - 1
        if not check_existing or (interval_length, first, last) not in per_interval_match_list:
            traffic = get_traffic(z)
            MWM = compute_optimal_matching(traffic, n_tors, **kwargs)
            per_interval_match_list[(interval_length, first, last)] = MWM
        if top and (not check_existing or (interval_length, first, last, top) not in per_interval_top_match_list):
            if traffic is None:
                traffic = get_traffic(z)
            compute_top(traffic, n_tors, top)
            MWM = compute_optimal_matching(traffic, n_tors, **kwargs)
            per_interval_top_match_list[(interval_length, first, last)] = MWM
        traffic = None


def main():
    parser = argparse.ArgumentParser(description='Compute max pair matching per interval.')
    parser.add_argument('--conf', default="conf.json", type=open,
                        help='configuration file (default: conf.json)')
    parser.add_argument('--interval_length', default=1, type=int,
                        help='each interval length (default: 1)')
    
    args = parser.parse_args()
    conf = json.load(args.conf)
    conf["interval_length"] = args.interval_length
    compute_intervals_mwm(**conf)


if __name__ == "__main__":
    main()