import torch
import numpy as np
import pickle

def save_random_problems(batch_size, problem_size, alpha=0.2, n=40):
    depot_xy = torch.rand(size=(batch_size, 1, 2))
    # shape: (batch, 1, 2)

    node_xy = torch.rand(size=(batch_size, problem_size, 2))
    # shape: (batch, problem, 2)

    if problem_size == 20:
        demand_scaler = 30
    elif problem_size == 50:
        demand_scaler = 40
    elif problem_size == 100:
        demand_scaler = 50
    elif problem_size == 200:
        demand_scaler = 100
    elif problem_size == 500:
        demand_scaler = 150
    elif problem_size == 1000:
        demand_scaler = 200
    else:
        raise NotImplementedError

    node_demand = torch.randint(1, 10, size=(batch_size, problem_size)) / float(demand_scaler)
    # shape: (batch, problem)
    coords = torch.cat((depot_xy, node_xy), dim=1)
    coord_diff = coords.unsqueeze(2) - coords.unsqueeze(1)  # (B, N, N, 2)
    dists = torch.norm(coord_diff, dim=-1)  # (B, N, N)
    idx = torch.randperm(problem_size+1)[:2 * n]
    group_a = idx[:n]
    group_b = idx[n:]

    # Create mask for (i in group_a, j in group_b) and vice versa
    mask = torch.zeros(batch_size, problem_size+1, problem_size+1, dtype=torch.bool)
    for i in group_a:
        for j in group_b:
            mask[:, i, j] = True
            mask[:, j, i] = True  # if symmetry is needed

    # Step 3: Generate random increases (between 0 and alpha × distance)
    increase = torch.rand_like(dists) * alpha * dists
    perturbed_dists = dists + increase * mask  #

    my_dict = {
        'depot_xy': depot_xy,
        'node_xy': node_xy,
        'node_demand': node_demand,
        'distances': perturbed_dists
    }
    torch.save(my_dict, f'AVRP_{batch_size}_{problem_size}.pt')
    with open(f'AVRP_{batch_size}_{problem_size}.pkl', 'wb') as f:
        pickle.dump(my_dict, f)
    return depot_xy, node_xy, node_demand


def get_random_problems(batch_size, problem_size, alpha=0.2, n=40):
    depot_xy = torch.rand(size=(batch_size, 1, 2))
    # shape: (batch, 1, 2)

    node_xy = torch.rand(size=(batch_size, problem_size, 2))
    # shape: (batch, problem, 2)

    if problem_size == 20:
        demand_scaler = 30
    elif problem_size == 50:
        demand_scaler = 40
    elif problem_size == 100:
        demand_scaler = 50
    elif problem_size == 200:
        demand_scaler = 100
    elif problem_size == 400:
        demand_scaler = 150
    elif problem_size == 1000:
        demand_scaler = 200
    else:
        raise NotImplementedError

    node_demand = torch.randint(1, 10, size=(batch_size, problem_size)) / float(demand_scaler)
    # shape: (batch, problem)
    coords = torch.cat((depot_xy, node_xy), dim=1)
    coord_diff = coords.unsqueeze(2) - coords.unsqueeze(1)  # (B, N, N, 2)
    dists = torch.norm(coord_diff, dim=-1)  # (B, N, N)
    idx = torch.randperm(problem_size+1)[:2 * n]
    group_a = idx[:n]
    group_b = idx[n:]

    # Create mask for (i in group_a, j in group_b) and vice versa
    mask = torch.zeros(batch_size, problem_size+1, problem_size+1, dtype=torch.bool)
    for i in group_a:
        for j in group_b:
            mask[:, i, j] = True
            mask[:, j, i] = True  # if symmetry is needed

    # Step 3: Generate random increases (between 0 and alpha × distance)
    increase = torch.rand_like(dists) * alpha * dists
    perturbed_dists = dists + increase * mask  #

    return depot_xy, node_xy, node_demand, perturbed_dists


def augment_xy_data_by_8_fold(xy_data):
    # xy_data.shape: (batch, N, 2)

    x = xy_data[:, :, [0]]
    y = xy_data[:, :, [1]]
    # x,y shape: (batch, N, 1)

    dat1 = torch.cat((x, y), dim=2)
    dat2 = torch.cat((1 - x, y), dim=2)
    dat3 = torch.cat((x, 1 - y), dim=2)
    dat4 = torch.cat((1 - x, 1 - y), dim=2)
    dat5 = torch.cat((y, x), dim=2)
    dat6 = torch.cat((1 - y, x), dim=2)
    dat7 = torch.cat((y, 1 - x), dim=2)
    dat8 = torch.cat((1 - y, 1 - x), dim=2)

    aug_xy_data = torch.cat((dat1, dat2, dat3, dat4, dat5, dat6, dat7, dat8), dim=0)
    # shape: (8*batch, N, 2)

    return aug_xy_data


if __name__ == '__main__':
    save_random_problems(100, 500, 0.2, 200)
