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

use_neos = False

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=4,
                   max_weekend_frequency=2,
                   max_night_frequency=2)

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_weekday",
                12.25,
                days[:4],
                balance_offset=5,
                balance_weighting=1,
                workers_required=3,
                force_as_block=False,
                rota_on_nwds=True,
                constraints=["night"]),
    SingleShift((sites),
                "night_weekend",
                12.25,
                days[4:],
                balance_offset=4,
                balance_weighting=1,
                workers_required=3,
                force_as_block=False,
                rota_on_nwds=True,
                constraints=["night"]),
)

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_weekday",
                13,
                days[:4],
                balance_offset=5,
                balance_weighting=1,
                workers_required=3,
                force_as_block=False,
                rota_on_nwds=True,
                constraints=["night"]),
    SingleShift((sites),
                "night_weekend",
                13,
                days[4:],
                balance_offset=4,
                balance_weighting=1,
                workers_required=3,
                force_as_block=False,
                rota_on_nwds=True,
                constraints=["night"]),
)

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_weekday",
                13,
                days[:4],
                balance_offset=5,
                balance_weighting=1,
                workers_required=3,
                force_as_block=False,
                rota_on_nwds=True,
                constraints=["night"]),
    SingleShift((sites),
                "night_weekend",
                13,
                days[4:],
                balance_offset=4,
                balance_weighting=1,
                workers_required=3,
                force_as_block=False,
                rota_on_nwds=True,
                constraints=["night"]),
)

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_weekday",
                13,
                days[:4],
                balance_offset=5,
                balance_weighting=1,
                workers_required=3,
                force_as_block=False,
                rota_on_nwds=True,
                constraints=["night"]),
    SingleShift((sites),
                "night_weekend",
                13,
                days[4:],
                balance_offset=4,
                balance_weighting=1,
                workers_required=3,
                force_as_block=False,
                rota_on_nwds=True,
                constraints=["night"]),
)

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_weekday",
                12.25,
                days[:4],
                balance_offset=5,
                balance_weighting=1,
                workers_required=3,
                force_as_block=False,
                rota_on_nwds=True,
                constraints=["night"]),
    SingleShift((sites),
                "night_weekend",
                12.25,
                days[4:],
                balance_offset=4,
                balance_weighting=1,
                workers_required=3,
                force_as_block=False,
                rota_on_nwds=True,
                constraints=["night"]),
)

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_weekday",
                12.25,
                days[:4],
                balance_offset=5,
                balance_weighting=1,
                workers_required=3,
                force_as_block=False,
                rota_on_nwds=True,
                constraints=["night"]),
    SingleShift((sites),
                "night_weekend",
                12.25,
                days[4:],
                balance_offset=4,
                balance_weighting=1,
                workers_required=3,
                force_as_block=False,
                rota_on_nwds=True,
                constraints=["night"]),
)

# 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 = rota_collections[rot]
            Rota.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>
        <head>
        <link rel="stylesheet" type="text/css" href="timetable.css">
        <script src="https://ajax.googleapis.com/ajax/libs/jquery/3.5.1/jquery.min.js"></script>
        <script src="s.js" defer></script>
        </head>
    <body>""")

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

    Rota.build_shifts_and_workers()

    print("-Building model")
    Rota.build_model()
    print("-Building model: complete")

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

    if use_neos:
        print("-Solving with neos")
        solver_manager = SolverManagerFactory('neos')  # Solve in neos server
        results = solver_manager.solve(Rota.model,
                                       opt=opt,
                                       logfile="{}.log".format(rot))
    else:
        print("-Solving with local cbc")
        results = opt.solve(
            Rota.model,
            tee=True,
            options={
                "seconds": 10200,
                "allow": 2000,
            },
            logfile="{}.log".format(rot)
        )  # 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(ResultsHolder.get_worker_timetable_html(table_name=rot))
        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>")