from greedrl.feature import *
from cvrp import make_problem as make_cvrp_problem
from torch.utils.data import Dataset, IterableDataset, DataLoader


def make_problem(batch_count, batch_size, task_count):
    features = [continuous_feature('task_demand_x'),
                continuous_feature('distance_matrix')]

    problem_list = make_cvrp_problem(batch_count, batch_size, task_count)
    for problem in problem_list:
        problem.features = features

    return problem_list


class Dataset(IterableDataset):
    def __init__(self, batch_count, batch_size, task_count):
        self._batch_size = batch_size
        self._task_count = task_count
        self._batch_count = batch_count
        self._index = 0

    def __iter__(self):
        self._index = 0
        return self

    def __next__(self):
        if self._batch_count is not None \
                and self._index >= self._batch_count:
            raise StopIteration()

        p = make_problem(1, self._batch_size, self._task_count)[0]
        self._index += 1
        return p


def write_vrplib(filename, name, size, demand, capacity, location):
    with open(filename, 'w') as f:
        f.write('\n'.join([
            "{} : {}".format(k, v)
            for k, v in (
                ('NAME', name),
                ('TYPE', 'CVRP'),
                ('COMMENT', 'NONE'),
                ('DIMENSION', size + 1),
                ('EDGE_WEIGHT_TYPE', 'EUC_2D'),
                ('CAPACITY', capacity)
            )
        ]))

        f.write('\n')
        f.write('NODE_COORD_SECTION\n')

        f.write('\n'.join(['{}\t{}\t{}'.format(i + 1, x, y) for i, (x, y) in enumerate(location)]))

        f.write('\n')
        f.write('DEMAND_SECTION\n')
        f.write('\n'.join(['{}\t{}'.format(i + 1, d) for i, d in enumerate([0] + demand)]))

        f.write('\n')
        f.write('DEPOT_SECTION\n')
        f.write('1\n')
        f.write('-1\n')
        f.write('EOF\n')