from entity import *
import random
import numpy as np
from sklearn.mixture import GMM
max_x = 30
max_y = 30
graph_xy = []
for x in range(max_x):
    for y in range(max_y):
        graph_xy.append((x,y))
random.shuffle(graph_xy)
# def test_orders():
#     order_list = [[] for _ in range(144)]
#     id = 0
#     for t in range(1,145):
#         step = 5
#         more_1 = random.randint(900,1200)
#         more_2 = random.randint(1300,1600)
#         less_1 = random.randint(900,1200)
#         for i in range(step):
#             step_list = []
#             if 69 <= t <= 78 or 105 <= t <= 114:
#                 for i in range(more_1//step):
#                     to_x = random.randint(0, 29)
#                     to_y = random.randint(0, 29)
#                     x = random.randint(0, 29)
#                     y = random.randint(0, 29)
#                     order = Order(id, x, y, to_x, to_y, t)
#                     step_list.append(order)
#                     id += 1
#                 for i in range(more_2//step):
#                     to_x = random.randint(0, 29)
#                     to_y = random.randint(0, 29)
#                     x = random.randint(18, 24)
#                     y = random.randint(18, 24)
#                     order = Order(id, x, y, to_x, to_y, t)
#                     step_list.append(order)
#                     id += 1
#             else:
#
#                 for i in range(less_1//step):
#                     to_x = random.randint(0, 29)
#                     to_y = random.randint(0, 29)
#                     x = random.randint(0, 29)
#                     y = random.randint(0, 29)
#                     order = Order(id , x, y, to_x, to_y, t)
#                     step_list.append(order)
#                     id += 1
#             order_list[t-1].append(step_list)
#     save(order_list,"test_orders.pkl")

def test_orders2():

    order_list = [[] for _ in range(144)]
    id = 0
    for t in range(1,145):
        more_1 = random.randint(800,900)
        # more_2 = random.randint(1300,1600)
        less_1 = random.randint(750,850)
        if 69 <= t <= 78 or 105 <= t <= 114:
            for i in range(more_1):
                to_xy = random.choice(graph_xy)
                to_x = to_xy[0]
                to_y = to_xy[1]
                xy = random.choice(graph_xy)
                x = xy[0]
                y = xy[1]
                order = Order(id, x, y, to_x, to_y, t)
                order_list[t-1].append(order)
                id += 1
            # for i in range(more_2):
            #     to_x = random.randint(0, 29)
            #     to_y = random.randint(0, 29)
            #     x = random.randint(18, 24)
            #     y = random.randint(18, 24)
            #     order = Order(id, x, y, to_x, to_y, t)
            #     order_list[t - 1].append(order)
            #     id += 1
        else:
            for i in range(less_1):
                to_xy = random.choice(graph_xy)
                to_x = to_xy[0]
                to_y = to_xy[1]
                xy = random.choice(graph_xy)
                x = xy[0]
                y = xy[1]
                order = Order(id , x, y, to_x, to_y, t)
                order_list[t-1].append(order)
                id += 1
    # 模拟早晚高峰
    gmm = GMM(n_components=2)
    gmm.weights_ = np.array([1 / 3, 2 / 3])
    gmm.means_ = np.array([[8, 8, 75],
                           [20, 20, 110]])
    gmm.covars_ = np.array([[4, 4, 5],
                            [4, 4, 5]])
    samp = gmm.sample(19100)
    # print(samp)
    for s in samp:
        x = int(round(s[0]))
        y = int(round(s[1]))
        t = int(round(s[2]))
        to_xy = random.choice(graph_xy)
        to_x = to_xy[0]
        to_y = to_xy[1]
        order = Order(id, x, y, to_x, to_y, t)
        order_list[t - 1].append(order)
        id+=1
    n = 0
    for o in order_list:
        print(n,len(o))
        n+=1
    new_order_list = [[] for _ in range(144)]
    for i in range(len(order_list)):
        ol = order_list[i]
        step = len(ol)//4
        random.shuffle(ol)
        start = 0
        for _ in range(4):
            step_list = ol[start:start+step]
            new_order_list[i].append(step_list)
            start += step
    save(new_order_list,"test_orders05.pkl")

def test_drivers():
    drivers_list = []
    for i in range(6000):
        d_xy = random.choice(graph_xy)
        d_x = d_xy[0]
        d_y = d_xy[1]
        driver = Driver(i, d_x, d_y,time=1)
        drivers_list.append(driver)
    save(drivers_list, "test_drivers.pkl")


if __name__ == '__main__':
    # test_orders()
    test_drivers()
    test_orders2()
    # test_orders2()
    # test_drivers2()
    pass