#! /usr/bin/env python

from itertools import product
import numpy as np
import traceback
import os

import sys
sys.path.append('..')

from lib.problem import Problem
from lib.partitioning import FMPartitioning, SpectralClustering
from lib.partitioning.utils import all_partitions_contiguous
from lib.algorithms import NcfEpi
from benchmarks.benchmark_consts import HOLDOUT_PROBLEMS

OUTPUT_CSV = 'grid-search.csv'
LOG_DIR = 'grid-search-logs'


def print_(*args, file=None):
    if file is None:
        file = sys.stdout
    print(*args, file=file)
    file.flush()


def grid_search(
        problem_name,
        topo_fname,
        tm_fname,
        num_paths_to_sweep=[4],
        edge_disjoint_to_sweep=[True, False],
        dist_metrics_to_sweep=['inv-cap'],
        partition_algos_to_sweep=['fm_partitioning', 'spectral_clustering'],
        num_parts_scale_factors_to_sweep=[1, 2, 3, 4]):

    problem = Problem.from_file(topo_fname, tm_fname)
    assert problem_name == problem.name
    print_(problem.name, tm_fname)
    traffic_seed = problem.traffic_matrix.seed
    total_demand = np.sum(problem.traffic_matrix.tm)
    print_('traffic seed: {}'.format(traffic_seed))
    print_('traffic matrix model: {}'.format(problem.traffic_matrix.model))
    print_('traffic scale factor: {}'.format(
        problem.traffic_matrix.scale_factor))
    print_('total demand: {}'.format(total_demand))

    num_parts_to_sweep = [
        sf * int(np.sqrt(len(problem.G.nodes)))
        for sf in num_parts_scale_factors_to_sweep
    ]

    for partition_algo, num_partitions_to_set, num_paths, edge_disjoint, dist_metric in product(
            partition_algos_to_sweep, num_parts_to_sweep, num_paths_to_sweep,
            edge_disjoint_to_sweep, dist_metrics_to_sweep):
        if partition_algo == 'fm_partitioning':
            partitioner = FMPartitioning(num_partitions_to_set)
        elif partition_algo == 'spectral_clustering':
            partitioner = SpectralClustering(num_partitions_to_set)

        print_(
            '\nNCFlow, {} partitioner, {} partitions, {} paths, edge disjoint {}, dist metric {}'
            .format(partition_algo, num_partitions_to_set, num_paths,
                    edge_disjoint, dist_metric))
        run_nc_dir = os.path.join(
            LOG_DIR, 'ncflow', partition_algo,
            '{}-partitions'.format(num_partitions_to_set),
            '{}-paths'.format(num_paths),
            'edge-disjoint-{}'.format(edge_disjoint),
            '{}-dist-metric'.format(dist_metric))
        if not os.path.exists(run_nc_dir):
            os.makedirs(run_nc_dir)

        with open(
                os.path.join(
                    run_nc_dir,
                    '{}-ncflow-{}_partitioner-{}_partitions-{}_paths-{}_edge_disjoint-{}_dist_metric.txt'
                    .format(problem.name, partition_algo,
                            num_partitions_to_set, num_paths, edge_disjoint,
                            dist_metric)), 'w') as log:
            partition_vector = partitioner.partition(problem)
            if not all_partitions_contiguous(problem, partition_vector):
                print_(
                    'Topology {}, partitioner {}, num_partitions_to_set {} did not find a valid partition'
                    .format(topo_fname, partition_algo, num_partitions_to_set))
                continue

            try:
                ncflow = NcfEpi.new_max_flow(num_paths,
                                             edge_disjoint=edge_disjoint,
                                             dist_metric=dist_metric,
                                             out=log)
                ncflow.solve(problem, partitioner)

                num_partitions = len(np.unique(ncflow._partition_vector))
                size_of_largest_partition = partitioner.size_of_largest_partition
                runtime = ncflow.runtime_est(14)
                total_flow = ncflow.obj_val
                with open(OUTPUT_CSV, 'a') as w:
                    print_('{},{},{},{},{},{},{},{},{},{}'.format(
                        problem.name, os.path.basename(tm_fname),
                        partition_algo, num_partitions,
                        size_of_largest_partition, num_paths, edge_disjoint,
                        dist_metric, total_flow, runtime),
                           file=w)
            except:
                print_(
                    'TM {}, {} partitioner, {} partitions, {} paths, edge disjoint {}, dist metric {} failed'
                    .format(tm_fname, partition_algo, num_partitions_to_set,
                            num_paths, edge_disjoint, dist_metric))
                traceback.print_exc(file=sys.stdout)


if __name__ == '__main__':
    with open(OUTPUT_CSV, 'a') as w:
        print_(
            'problem,tm_fname,partition_algo,num_partitions,size_of_largest_partition,num_paths,edge_disjoint,dist_metric,total_flow,runtime',
            file=w)

    for problem_name, topo_fname, tm_fname in HOLDOUT_PROBLEMS:
        grid_search(problem_name, topo_fname, tm_fname)
