from pyomo.environ import *
from pyomo.opt import SolverFactory
from shifts import SingleShift, RotaBuilder, RotaResults, days
from workers import Worker

from copy import deepcopy

import math

import scipy.stats as stats

import statistics

import datetime
import itertools
import operator

weeks_to_rota = 16

sites = ("truro", "exeter", "torquay", "barnstaple", "plymouth")

Rota = RotaBuilder(datetime.date(2020, 9, 7),
                   weeks_to_rota=weeks_to_rota,
                   balance_offset_modifier=1)

Rota.add_shifts(
    SingleShift(("exeter", ), "exeter_twilight", 12.5, days[:5]),
    SingleShift(("truro", ),
                "truro_twilight",
                12.5,
                days[:5],
                assign_as_block=False),
    #SingleShift(("torquay", ), "torquay_twilight", 12.5, days[:5]),
    #SingleShift(("plymouth", ), "plymouth_twilight", 12.5, days[:5]),
    SingleShift(("exeter", ),
                "weekend_exeter",
                12.5,
                days[5:],
                force_as_block=True),
    SingleShift(("truro", ),
                "weekend_truro",
                12.5,
                days[5:],
                assign_as_block=True),
    # SingleShift(("torquay", ),
    #             "weekend_torquay",
    #             12.5,
    #             days[5:],
    #             assign_as_block=True),
    # SingleShift(("plymouth", ),
    #             "weekend_plymouth",
    #             8,
    #             days[5:],
    #             workers_required=2,
    #             assign_as_block=True),
    SingleShift((sites),
                "night",
                12.25,
                days,
                balance_offset=2,
                balance_weighting=1,
                workers_required=3,
                rota_on_nwds=True,
                constraints=["night"]),
)

use_test_workers = True

# Import trainee data
if use_test_workers:
    # Rota.add_workers([
    #     Worker(Rota,
    #            i,
    #            "Truro {}".format(i),
    #            "truro",
    #            3,
    #            nwd=['Fri', "Sat", "Sun"]) for i in range(1, 2)
    # ])
    Rota.add_workers([
        Worker(
            Rota,
            i,
            "Truro {}".format(i),
            "truro",
            4,
        ) for i in range(2, 14)
    ])
    # Rota.add_workers([
    #     Worker(Rota, i, "Plym {}".format(i), "plymouth", 4)
    #     for i in range(16, 30)
    # ])
    Rota.add_workers([
        Worker(
            Rota,
            i,
            "Exe {}".format(i),
            "exeter",
            4,
        ) for i in range(30, 40)
    ])
    Rota.add_workers([
        Worker(Rota,
               i,
               "Torq {}".format(i),
               "torquay",
               2,
               pref_not_to_work=(("2020/09/07"), )) for i in range(50, 60)
    ])
else:
    import csv
    with open('trainees.csv', newline='') as f:
        reader = csv.reader(f, delimiter=',', quotechar='|')
        next(reader)  # skip header

        n = 0
        for row in reader:
            n = n + 1
            #print(row)
            name, site, grade, fte, nwd, end_date, oop, extra = row

            # Ignore trainees if fte == 0 (what are they doing here anyway)
            if fte == 0:
                continue

            nwds = nwd.split(",").title() if nwd else None
            end_date = end_date if end_date else None
            oop = oop.split("-") if oop else None
            #print(nwds, end_date, oop)
            Rota.add_worker(
                Worker(Rota, n, name, site.lower(), int(grade[2]), int(fte),
                       nwds, end_date, oop))

print(0)
Rota.build_shifts_and_workers()

Rota.build_model()

print(1)
opt = SolverFactory('cbc')  # choose a solver
#opt = SolverFactory('ipopt')  # choose a solver

print(2)
results = opt.solve(
    Rota.model, tee=True, options={
        "seconds": 1200,
        "allow": 70000
    })  # solve the model with the, options="seconds=60" selected solver

print(results)
results.solver.status
if not results.solver.status:
    sys.exit(0)

ResultsHolder = RotaResults(Rota, results)

week_table = ResultsHolder.get_work_table()  # list with the required workers
worker_timetable = ResultsHolder.get_worker_timetable()
worker_timetable_brief = ResultsHolder.get_worker_timetable_brief(
    show_prefs=False)

print(worker_timetable_brief)

# for i in ResultsHolder.rota.model.blocks_worker_shift_count:
#     n = ResultsHolder.rota.model.blocks_worker_shift_count[i].value

#     if n > 0:
#         print(i, n)
#         #print(ResultsHolder.rota.model.blocks_assigned[i].value)
#     else:
#         pass
#         #print("SHIT")
# for i in ResultsHolder.rota.model.blocks_assigned:
#     n = ResultsHolder.rota.model.blocks_assigned[i].value
#     if n:
#         print(i, n)
#         #print(ResultsHolder.rota.model.blocks_assigned[i].value)
#     else:
#         pass

# ResultsHolder.export_rota_to_csv("rotas/"+rot)
# with open("cest", "a") as f:
#     f.write(rot)
#     f.write("\n\n")

#     for s in Rota.shifts:
#         f.write(s.get_shift_summary())
#         f.write("\n")
#     f.write("\n\n")

#     f.write(worker_timetable_brief)
#     f.write("\n\n\n")
#workers_no_pref = ResultsHolder.get_no_preference(Rota.model.no_pref)  # list with the non-satisfied workers (work on Sat but not on Sun)

# worker_night_block = ResultsHolder.get_night_blocks()
# worker_shift_summary = ResultsHolder.get_shift_summary()
# multi = ResultsHolder.get_multinight()

#break