import pandas as pd
import pyomo.environ as pyo
import datetime
import matplotlib.cm as cm
import matplotlib.pyplot as plt

class Scheduler:
    def __init__(self, case_file_path, session_file_path, patient_file_path):
        """
        Read case and session data into Pandas DataFrames
        Args:
            case_file_path (str): path to case data in CSV format
            session_file_path (str): path to theatre session data in CSV format
        """
        try:
            self.df_cases = pd.read_csv(case_file_path)
        except FileNotFoundError:
            print(f"Case data not found. {case_file_path}")
        try:
            self.df_sessions = pd.read_csv(session_file_path)
        except FileNotFoundError:
            print("Session data not found")
        try:
            self.df_patients = pd.read_csv(patient_file_path)
        except FileNotFoundError:
            print("Patient data not found")
        self.solver = pyo.SolverFactory('glpk')
        self.model = self.create_model()
        self.build_model()

    def solve_model(self):
        self.solver_results = self.solver.solve(self.model, tee=True)

    def extract_results(self):
        results = [{"Case": case,
                    "Session": session,
                    #"Session Date": self.model.SESSION_DATES[session],
                    #"Case Deadline": self.model.CASE_DEADLINES[case],
                    #"Days before deadline": self.model.CASE_DEADLINES[case] - self.model.SESSION_DATES[session],
                    #"Start": self.model.CASE_START_TIME[case, session](),
                    "Assignment": self.model.SESSION_ASSIGNED[case, session]()}
                   for (case, session) in self.model.TASKS]
        self.df_times = pd.DataFrame(results)

    def create_model(self):
        return pyo.ConcreteModel()

    def set_constraint(self):
        self.model.c = pyo.ConstraintList()

    def solve(self):
        result = self.solver.solve(self.model)
        print(f'result was {result}')

    def build_model(self):
        self.add_cases()
        #self.set_options()
        self.add_sessions()
        self.add_tasks()
        self.set_decisions()
        self.set_obj()
        return

    def add_cases(self):
        # List of case IDs in surgical waiting list
        self.model.CASES = pyo.Set(initialize=self.df_cases["CaseID"].tolist())
        return

    def set_options(self):
        # Add solver parameters (time limit)
        options = {"seconds": 6}
        for key, value in options.items():
          self.solver.options[key] = value

    def add_sessions(self):
        # List of sessions IDs
        # TODO: Generate more sessions based on EKT Erhaltung
        self.model.SESSIONS = pyo.Set(initialize=self.df_sessions["SessionID"].tolist())
        return

    # List of job shop tasks
    # all possible combinations of cases and sessions
    def add_tasks(self):
        self.model.TASKS = pyo.Set(initialize=self.model.CASES * self.model.SESSIONS, dimen=2)
        return
    # Decision Variables

    ## Upper bound (minutes in a day)
    #ub = 1440
    ## Upper bound of session utilisation set to 85%
    #max_util = 0.85
    def set_decisions(self):

        # Binary flag, 1 if case is assigned to session, 0 otherwise
        self.model.SESSION_ASSIGNED = pyo.Var(self.model.TASKS, domain=pyo.Binary)
        # Start time of a case
        #self.model.CASE_START_TIME = pe.Var(self.model.TASKS, bounds=(0, ub), within=pe.PositiveReals)
        # Session utilisation
        num_cases = self.df_cases.shape[0]
        self.model.CASES_IN_SESSION = pyo.Var(self.model.SESSIONS, bounds=(0, num_cases), within=pyo.PositiveReals)
        self.model.UTILISATION = pyo.Var(self.model.SESSIONS, bounds=(0, 1), within=pyo.PositiveReals)

    def set_obj(self):

        # Objective
        def objective_function(model):
            return pyo.summation(model.CASES_IN_SESSION)
        self.model.OBJECTIVE = pyo.Objective(rule=objective_function, sense=pyo.maximize)

    def draw_gantt(self):

        df = self.df_times[self.df_times["Assignment"] == 1]
        cases = sorted(list(df['Case'].unique()))
        sessions = sorted(list(df['Session'].unique()))

        bar_style = {'alpha': 1.0, 'lw': 25, 'solid_capstyle': 'butt'}
        text_style = {'color': 'white', 'weight': 'bold', 'ha': 'center', 'va': 'center'}
        colors = cm.Dark2.colors

        df.sort_values(by=['Case', 'Session'])
        df.set_index(['Case', 'Session'], inplace=True)

        fig, ax = plt.subplots(1, 1)
        for c_ix, c in enumerate(cases, 1):
            for s_ix, s in enumerate(sessions, 1):
                if (c, s) in df.index:
                    xs = df.loc[(c, s), 'Start']
                    xf = df.loc[(c, s), 'Start'] + \
                         self.df_cases[self.df_cases["CaseID"] == c]["Median Duration"]
                    ax.plot([xs, xf], [s] * 2, c=colors[c_ix % 7], **bar_style)
                    ax.text((xs + xf) / 2, s, c, **text_style)

        ax.set_title('Assigning Ophthalmology Cases to Theatre Sessions')
        ax.set_xlabel('Time')
        ax.set_ylabel('Sessions')
        ax.grid(True)

        fig.tight_layout()
        plt.show()

def date2int(date):
    d0 = pd.to_datetime("25/11/2024", dayfirst=True)
    delta = (pd.to_datetime(date, dayfirst=True) - d0).days
    day = delta%7
    week = int((delta - day)/7)
    return week,day

def int2date(week, day):
    d0 = pd.to_datetime("25/11/2024", dayfirst=True)
    delta = datetime.timedelta(days = 7*week + day)
    date = d0 + delta
    return date

path = '/home/hmag/git/octopusx/ekt/data'
my = Scheduler(path+'/cases.csv', path+'/sessions.csv', path+'/patients.csv')

print(my.df_cases)
print(my.df_sessions)
#print(my.df_patients)
#for i in my.df_patients:
#    print(i)
print(my.df_cases['Date'][0])
date = my.df_cases['Date'][0]

print(date2int(date))
print(int2date(1,0))


#my.set_constraint()
#my.set_obj()
my.solve_model()

my.extract_results()
print(my.df_times)
my.draw_gantt()