import csv

decision = [0, 1]  # 0 - go to bin, 1 - pick up
levels = [1, 2, 3, 4, 5]
# 1 - 0;20  2 - 20;40  3 - 40;60  4 - 60;80  5 - 80;100
# 1 - 0;40  2 - 40;80  3 - 80;120  4 - 120;160  5 - 160+


def enough_free_space(available_space, trash_size, available_mass, mass_trash):
    if available_space + trash_size <= 5 and available_mass + mass_trash <= 5:
        return True
    return False


def where_is_closer(bin_distance, trash_distance):
    if bin_distance <= trash_distance:
        return 0
    return 1


with open('tree_dataset.csv', 'w', newline='') as csv_file:
    file_writer = csv.writer(csv_file)
    file_writer.writerow(["dis_dump", "dis_trash", "mass", "space", "trash_mass", "trash_space", "decision"])

    counter = 0

    for dis_dump in levels:
        for dis_trash in levels:
            for mass in levels:
                for space in levels:
                    for trash_mass in levels:
                        for trash_space in levels:
                            if counter % 10 == 0:
                                if dis_dump == 1 and space >= 1 and mass >= 1:
                                    file_writer.writerow([dis_dump, dis_trash, mass, space, trash_mass, trash_space, 0])
                                elif dis_trash == 1 and enough_free_space(space, trash_space, mass, trash_mass):
                                    file_writer.writerow([dis_dump, dis_trash, mass, space, trash_mass, trash_space, 1])
                                elif mass == 4 or space == 4 and not enough_free_space(space, trash_space, mass, trash_mass):
                                    file_writer.writerow([dis_dump, dis_trash, mass, space, trash_mass, trash_space, 0])
                                elif mass == 5 or space == 5:
                                    file_writer.writerow([dis_dump, dis_trash, mass, space, trash_mass, trash_space, 0])
                                elif mass <= 3 and space <= 3 and enough_free_space(space, trash_space, mass, trash_mass):
                                    file_writer.writerow([dis_dump, dis_trash, mass, space, trash_mass, trash_space, 1])
                                elif mass == 4 or space == 4 and enough_free_space(space, trash_space, mass, trash_mass):
                                    file_writer.writerow([dis_dump, dis_trash, mass, space, trash_mass, trash_space, where_is_closer(dis_dump, dis_trash)])
                                elif not enough_free_space(space, trash_space, mass, trash_mass):
                                    file_writer.writerow([dis_dump, dis_trash, mass, space, trash_mass, trash_space, 0])
                                else:
                                    file_writer.writerow([dis_dump, dis_trash, mass, space, trash_mass, trash_space, None])

                            counter += 1