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 = 26

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

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

rota_collections = {
    "current": deepcopy(Rota),
    "extended proc nights only": deepcopy(Rota),
    "no weekends": deepcopy(Rota),
    "no twighlights": deepcopy(Rota),
    "no nights": deepcopy(Rota),
    "nights + proc twighlights + normal weekends": deepcopy(Rota),
    "nights + proc twighlights + proc weekends": deepcopy(Rota),
}

rota_collections['current'].add_shifts(
    SingleShift(("exeter", ), "exeter_twilight", 12.5, days[:5]),
    SingleShift(("truro", ), "truro_twilight", 12.5, days[:5]),
    SingleShift(("torquay", ), "torquay_twilight", 12.5, days[:5]),
    SingleShift(("plymouth", ), "plymouth_twilight", 12.5, days[:5]),
    SingleShift(("exeter", ),
                "weekend_exeter",
                12.5,
                days[5:],
                assign_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.5, days, balance_weighting=1, workers_required=3, rota_on_nwds=True),
    SingleShift((sites),
                "night",
                12.25,
                days,
                balance_offset=2,
                balance_weighting=1,
                workers_required=3,
                rota_on_nwds=True),
)

rota_collections['extended proc nights only'].add_shifts(
    #SingleShift((sites), "night", 12.5, days, balance_weighting=1, workers_required=3, rota_on_nwds=True),
    SingleShift((sites),
                "night",
                13,
                days,
                balance_offset=2,
                balance_weighting=1,
                workers_required=3,
                rota_on_nwds=True), )

rota_collections['no weekends'].add_shifts(
    SingleShift(("exeter", ), "exeter_twilight", 12.5, days[:5]),
    SingleShift(("truro", ), "truro_twilight", 12.5, days[:5]),
    SingleShift(("torquay", ), "torquay_twilight", 12.5, days[:5]),
    SingleShift(("plymouth", ), "plymouth_twilight", 12.5, days[:5]),
    #SingleShift((sites), "night", 12.5, days, balance_weighting=1, workers_required=3, rota_on_nwds=True),
    SingleShift((sites),
                "night",
                13,
                days,
                balance_offset=2,
                balance_weighting=1,
                workers_required=3,
                rota_on_nwds=True),
)

rota_collections['no twighlights'].add_shifts(
    SingleShift(("exeter", ),
                "weekend_exeter",
                12.5,
                days[5:],
                assign_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.5, days, balance_weighting=1, workers_required=3, rota_on_nwds=True),
    SingleShift((sites),
                "night",
                13,
                days,
                balance_offset=2,
                balance_weighting=1,
                workers_required=3,
                rota_on_nwds=True),
)

rota_collections['no nights'].add_shifts(
    SingleShift(("exeter", ), "exeter_twilight", 12.5, days[:5]),
    SingleShift(("truro", ), "truro_twilight", 12.5, days[:5]),
    SingleShift(("torquay", ), "torquay_twilight", 12.5, days[:5]),
    SingleShift(("plymouth", ), "plymouth_twilight", 12.5, days[:5]),
    SingleShift(("exeter", ),
                "weekend_exeter",
                12.5,
                days[5:],
                assign_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.5, days, balance_weighting=1, workers_required=3, rota_on_nwds=True),
)

rota_collections['nights + proc twighlights + normal weekends'].add_shifts(
    SingleShift((sites), "proc_twilight", 12.5, days[:5], workers_required=4),
    SingleShift(("exeter", ),
                "weekend_exeter",
                12.5,
                days[5:],
                assign_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.5, days, balance_weighting=1, workers_required=3, rota_on_nwds=True),
    SingleShift((sites),
                "night",
                12.25,
                days,
                balance_offset=2,
                balance_weighting=1,
                workers_required=3,
                rota_on_nwds=True),
)

rota_collections['nights + proc twighlights + proc weekends'].add_shifts(
    SingleShift((sites),
                "proc_twilight_and_weekends",
                12.5,
                days,
                workers_required=3),
    #SingleShift((sites), "night", 12.5, days, balance_weighting=1, workers_required=3, rota_on_nwds=True),
    SingleShift((sites),
                "night",
                12.25,
                days,
                balance_offset=2,
                balance_weighting=1,
                workers_required=3,
                rota_on_nwds=True),
)

# Import trainee data
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)
        for rot in rota_collections:
            rota_collections[rot].add_worker(
                Worker(Rota, n, name, site.lower(), int(grade[2]), int(fte),
                       nwds, end_date, oop))

with open("rotas.html", "w") as f:
    f.write("<html><body>")

for rot in rota_collections:
    print("Running rota: ", rot)
    Rota = rota_collections[rot]

    Rota.build_shifts_and_workers()

    Rota.build_model()

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

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

    results.solver.status

    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)

    worker_shift_summary = ResultsHolder.get_shift_summary()

    ResultsHolder.export_rota_to_csv("rotas/" + rot)
    with open("rotas.html", "a") as f:
        f.write("----------------------------")
        f.write("<h2>{}</h2>".format(rot))
        f.write("----------------------------")
        f.write("\n\n")

        for s in Rota.shifts:
            f.write("<pre>{}</pre>".format(s.get_shift_summary()))
            f.write("\n")
        f.write("\n\n")

        f.write("<pre>{}</pre>".format(worker_timetable_brief))
        f.write("\n\n")
        f.write("<pre>{}</pre>".format(worker_shift_summary))
        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()
    multi = ResultsHolder.get_multinight()

    #break

with open("rotas.html", "a") as f:
    f.write("</body></html>")