import math
from tsp import solve_tsp
from map import Map as ClusterMap
try:
    import matplotlib.pyplot as plt
except ModuleNotFoundError:
    plt = None

def length(customer1, customer2):
    return math.sqrt((customer1.x - customer2.x)**2 + (customer1.y - customer2.y)**2)
distance=length


class Customer(object):
    def __init__(self, index, demand, x, y):
        self.id = index
        self.index = index
        self.demand = demand
        self.x = x
        self.y = y
        self.cluster_x = None
        self.cluster_y = None
        # A list tuples. The first field is the distance
        # and the second is another point (i.e. a neighbor).
        self.neighbors = []

    def add_neighbors(self, neighbors):
        neighbors = [(n, distance(self, n)) for n in neighbors]
        self.neighbors = sorted(neighbors, key=lambda t: t[1])

    def __str__(self):
        # m = "C({}, {})".format(self.cluster_x, self.cluster_y)
        m = "C({})".format(self.id)
        return m

    def __repr__(self):
        return self.__str__()


class Vrp(object):

    def __init__(self, input_data):
        lines = input_data.split('\n')
        parts = lines[0].split()
        self.customer_count = int(parts[0])
        self.vehicle_count = int(parts[1])
        self.vehicle_capacity = int(parts[2])
        self.vehicle_tours = []

        customers = []
        for i in range(1, self.customer_count + 1):
            line = lines[i]
            parts = line.split()
            demand = int(parts[0])
            x, y = map(float, parts[1:])
            c = Customer(i - 1, demand, x, y)
            customers.append(c)

        self.depot = customers[0]
        self.customers = customers[1:]

        self.map = ClusterMap(8)
        self.map.cluster(self.customers)
        self.plot_n = 0
        # print(f"{self.vehicle_count=} {self.vehicle_capacity=}")

    def order_customers_by_clusters(self):
        len_col = len(self.map.clusters)
        len_row = len(self.map.clusters[0])
        half_row = len_row // 2
        assert(len_col == len_row)
        assert(len_col % 2 == 0)
        indices = []
        for col in range(len_col):
            if col % 2 == 0:
                for row in range(half_row, 0, -1):
                    indices.append((col, row - 1))
            else:
                for row in range(half_row):
                    indices.append((col, row))

        for col in range(len_col, 0, -1):
            if col % 2 == 0:
                for row in range(half_row, len_row):
                    indices.append((col - 1, row))
            else:
                for row in range(len_row, half_row, -1):
                    indices.append((col - 1, row - 1))

        len_orig = len(self.customers)
        self.customers = []
        for col, row in indices:
            self.customers += self.map.clusters[row][col]
        assert(len_orig == len(self.customers))

    def plot(self):
        if plt is None:
            return

        plt.figure(dpi=300)

        for c in self.customers:
            plt.plot(c.x, c.y, 'rx')
            plt.text(c.x, c.y, f'  {c}')

        d = self.depot
        plt.plot(d.x, d.y, 'yo')
        self.map.plot_grid(plt)

        line_format = 'b:'
        line_width = 0.4
        for t in self.vehicle_tours:
            if not t:
                continue
            plt.plot([d.x, t[0].x], [d.y, t[0].y],
                     line_format, linewidth=line_width)
            for i in range(len(t) - 1):
                plt.plot([t[i].x, t[i + 1].x], [t[i].y, t[i + 1].y],
                         line_format, linewidth=line_width)
            plt.plot([t[-1].x, d.x], [t[-1].y, d.y],
                     line_format, linewidth=line_width)

        plt.axis('off')
        fig_file = "plots/step_{}.png".format(self.plot_n)
        plt.savefig(fig_file, bbox_inches='tight')
        self.plot_n += 1
        # plt.show()

    def solve_trivial(self):
        self.vehicle_tours = []

        remaining_customers = set(self.customers)

        for v in range(0, self.vehicle_count):
            self.vehicle_tours.append([])
            capacity_remaining = self.vehicle_capacity
            while sum([capacity_remaining >= customer.demand
                       for customer in remaining_customers]) > 0:
                used = set()
                order = sorted(remaining_customers,
                               key=lambda customer: -customer.demand)
                for customer in order:
                    if capacity_remaining >= customer.demand:
                        capacity_remaining -= customer.demand
                        self.vehicle_tours[v].append(customer)
                        used.add(customer)
                remaining_customers -= used

        assert sum([len(v) for v in self.vehicle_tours]) == len(self.customers)

    def optimize_tours(self):
        for t in self.vehicle_tours:
            t.append(self.depot)

        self.vehicle_tours = [solve_tsp(t) if len(t) > 1 else t
                              for t in self.vehicle_tours]
        new_tours = []
        for tour in self.vehicle_tours:
            for i, customer in enumerate(tour):
                if customer.id == 0:
                    depot_index = i
                    break
            new_tour = tour[i + 1:] + tour[:i]
            new_tours.append(new_tour)
            assert(len(tour) - 1== len(new_tour))
        self.vehicle_tours = new_tours

    def make_routes(self):
        self.order_customers_by_clusters()
        self.vehicle_tours = [[] for _ in range(self.vehicle_count)]
        self.vehicle_capacities = [self.vehicle_capacity
                                   for _ in range(self.vehicle_count)]
        for c in self.customers:
            for i, tour in enumerate(self.vehicle_tours):
                if c.demand < self.vehicle_capacities[i]:
                    self.vehicle_capacities[i] -= c.demand
                    tour.append(c)
                    break
            else:
                raise Exception("No tour for {}.".format(c))

    def to_output(self):
        obj = 0
        for v in range(0, self.vehicle_count):
            vehicle_tour = self.vehicle_tours[v]
            if len(vehicle_tour) > 0:
                obj += length(self.depot, vehicle_tour[0])
                for i in range(0, len(vehicle_tour) - 1):
                    obj += length(vehicle_tour[i],vehicle_tour[i + 1])
                obj += length(vehicle_tour[-1], self.depot)

        outputData = '%.2f' % obj + ' ' + str(0) + '\n'
        for v in range(0, self.vehicle_count):
            outputData += str(self.depot.id) + ' ' + ' '.join(
                [str(customer.id) for customer in self.vehicle_tours[v]]) \
                + ' ' + str(self.depot.id) + '\n'

        return outputData


def solve_it(input_data):
    vrp = Vrp(input_data)
    if len(vrp.customers) in [25, 199]:
        vrp.solve_trivial()
    else:
        vrp.make_routes()
    vrp.optimize_tours()
    vrp.plot()
    return vrp.to_output()


if __name__ == "__main__":
    file_location = "data/vrp_16_3_1"
    file_location = "data/vrp_51_5_1"
    with open(file_location, 'r') as f:
        print(solve_it(f.read()))
