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

time_to_run = 152000
allow = 2000

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

Rota = RotaBuilder(datetime.date(2021, 3, 8),
                   weeks_to_rota=weeks_to_rota,
                   balance_offset_modifier=1,
                   max_weekend_frequency=2,
                   max_night_frequency=2)

Rota.constraint_options["limit_to_1_st1_on_nights"] = True
Rota.constraint_options["ensure_1_st4_plus_on_nights"] = True
Rota.constraint_options["balance_nights"] = True
Rota.constraint_options["constrain_time_off_after_nights"] = True
Rota.constraint_options["balance_nights_across_sites"] = True
Rota.constraint_options["balance_blocks"] = True
Rota.constraint_options["balance_weekends"] = True
Rota.constraint_options["avoid_st1_first_month"] = False
Rota.constraint_options["max_weekends"] = 10
Rota.constraint_options["prevent_monday_and_tuesday_after_full_weekends"] = [
    "truro", "torquay", "exeter"
]
Rota.constraint_options["prevent_monday_after_full_weekends"] = ["plymouth"]
Rota.constraint_options["prevent_fridays_before_full_weekends"] = ["plymouth"]
Rota.constraint_options["prevent_thursdays_before_full_weekends"] = [
    "truro", "torquay", "exeter"
]

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

use_test_workers = False

load_leave = True
Rota.build_shifts_and_workers()

if load_leave:
    import leave

    leave, requests, site_prefs = leave.load_leave(Rota)

# 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(1, 9)
    ])
    Rota.add_workers([
        Worker(Rota, i, "Plym {}".format(i), "plymouth", 4, 50)
        for i in range(9, 17)
    ])
    Rota.add_workers([
        Worker(Rota,
               i,
               "Plym {}".format(i),
               "plymouth",
               4,
               80,
               end_date="2020/10/05") for i in range(17, 19)
    ])
    # Rota.add_workers([
    #     Worker(
    #         Rota,
    #         i,
    #         "Exe {}".format(i),
    #         "exeter",
    #         4,
    #     ) for i in range(20, 25)
    # ])
    # Rota.add_workers([
    #     Worker(
    #         Rota,
    #         i,
    #         "Torq {}".format(i),
    #         "torquay",
    #         2,
    #         50,
    #     ) for i in range(25, 29)
    # ])
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

            name, site, grade, fte, nwd, end_date, oop, twi, twi_we, night, night_we, extra = row

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

            nwds = nwd.split("/") if nwd else None
            end_date = end_date if end_date else None
            oop = oop.split("-") if oop else None

            previous_shifts = {}
            if twi:
                worked = twi.split(" ")[0]
                allocated = twi.split(" ")[1][1:-1]
                previous_shifts["{}_twilight".format(site.lower())] = (worked, allocated)
            if twi_we:
                worked = twi_we.split(" ")[0]
                allocated = twi_we.split(" ")[1][1:-1]
                previous_shifts["weekend_{}".format(site.lower())] = (worked, allocated)
            if night:
                worked = night.split(" ")[0]
                allocated = night.split(" ")[1][1:-1]
                previous_shifts["night_weekday".format(site.lower())] = (worked, allocated)
            if night_we:
                worked = night_we.split(" ")[0]
                allocated = night_we.split(" ")[1][1:-1]
                previous_shifts["night_weekend".format(site.lower())] = (worked, allocated)


            leave_requests = []
            work_requests = []
            site_pref = 1
            if load_leave:
                if name not in leave:
                    raise ValueError(
                        "{} not found it leave requests".format(name))

                leave_requests = [i[0] for i in leave[name] if i[0] != ""]
                work_requests = [i for i in requests[name] if i[0] != ""]
                site_pref = int(site_prefs[name][0])
            #print(nwds, end_date, oop)
            # Rota.add_worker(
            #     Worker(Rota, n, name, site.lower(), int(grade[2]), 100,
            #            ))
            # if name.startswith(("Alex", "Michal", "Ross")):
            #     Rota.add_worker(
            #         Worker(Rota,
            #                n,
            #                name,
            #                site.lower(),
            #                int(grade[2]),
            #                int(fte),
            #                nwds,
            #                end_date,
            #                oop,
            #                work_requests=(("2020/09/07", "night_weekday"), )))
            #     continue
            # if "Salma" in name:
            #     Rota.add_worker(
            #         Worker(Rota,
            #                n,
            #                name,
            #                site.lower(),
            #                int(grade[2]),
            #                int(fte),
            #                nwds,
            #                end_date,
            #                oop,
            #                work_requests=(("2020/12/25",
            #                                "exeter_twilight"), )))
            #     continue
            Rota.add_worker(
                Worker(Rota,
                       n,
                       name,
                       site.lower(),
                       int(grade[2]),
                       int(fte),
                       nwds,
                       end_date,
                       oop,
                       not_available_to_work=leave_requests,
                       work_requests=work_requests,
                       night_at_derriford=site_pref, 
                       previous_shifts=previous_shifts))

print("Building model")
Rota.build_shifts_and_workers()

Rota.build_model()

#print(Rota.get_worker_details())
print("Setting up solver")
opt = SolverFactory('cbc')  # choose a solver
#opt = SolverFactory('ipopt')  # choose a solver

print("Solving")
if use_neos:
    solver_manager = SolverManagerFactory('neos')  # Solve in neos server
    results = solver_manager.solve(Rota.model, opt=opt, logfile="test.log")
else:
    results = opt.solve(
        Rota.model,
        tee=True,
        options={
            "seconds": time_to_run,
            "allow": allow,
            "threads": 4,
        },
        logfile="test.log"
    )  # 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, show_unavailable=False)

#print(worker_timetable_brief)
#print(Rota.get_workers_total_fte())

ResultsHolder.get_weekend_details()
#ResultsHolder.get_night_details()

with open("test.html", "w") as f:
    f.write(ResultsHolder.get_worker_timetable_html(True))

ResultsHolder.export_rota_to_csv()

# 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