from entity import *
import random
import numpy as np
import gc
from sklearn.externals import joblib

def get_data():
    data = load("train_data_02.pkl")
    new_data = [[] for _ in range(144)]
    for d in data:
        for match in d:
            # if match.order.arrive_time >= match.order.order_time or match.order.order_time >= 142:
                new_data[match.order.order_time-1].append(match)
    del data

    for t in range(len(new_data)):
        if 69<=t<=78 or 105<=t<=114:
            for n in range(random.randint(16000,20000)):
                d_x = random.randint(25, 35)
                d_y = random.randint(25, 35)
                driver = Driver(id, d_x, d_y)
                to_x = random.randint(0, 49)
                to_y = random.randint(0, 49)
                juli = 1000
                while juli > 6:
                    o_x = random.randint(d_x-6,d_x+6)
                    o_y = random.randint(d_y-6,d_y+6)
                    juli = abs(o_x - d_x) + abs(o_y - d_y)
                order = Order(0, o_x, o_y, to_x, to_y, t)
                match = Match(order, driver)

                new_data[t].append(match)
                print('\r' + str(t) + '  ' + str(n), end='')
    count = 0
    for i in new_data:
        count += 1
        print(str(count)+":",len(i))
    save(new_data,"train_data03.pkl")

if __name__ == '__main__':
    get_data()