import pyomo
import pandas as pd
from pulp import *

class Scheduler:
    def __init__(self,
                 task_path='data/tasks.csv',
                 block_path='data/blocks.csv',
                 patient_path='data/patients.csv'):
        """
        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_tasks = pd.read_csv(task_path)
        except FileNotFoundError:
            print("Task data not found.")
        try:
            self.df_blocks = pd.read_csv(block_path)
        except FileNotFoundError:
            print("Session data not found")
        try:
            self.df_patients = pd.read_csv(patient_path)
        except FileNotFoundError:
            print("Patient data not found")

        self.create_lists()
        self.extend_data()

    def extend_data(self):

        # Date to integer all dates
        # tasks date
        newcol = [self.date2int(date) for date in self.df_tasks['date'].tolist()]
        self.df_tasks.insert(7, "dateint", newcol)

        # block date
        newcol2 = [self.date2int(date) for date in self.df_blocks['day'].tolist()]
        self.df_blocks.insert(3, "dateint", newcol2)
        return

    def create_lists(self):
        self.s = self.df_tasks['priority'].tolist()
        self.d = self.df_tasks['duration'].tolist()
        self.dates = self.df_tasks['date'].tolist()
        self.b = self.df_blocks['available'].tolist()
        self.block_dates= self.df_blocks['day'].tolist()

        self.patient_ids = self.df_patients['patient'].tolist()
        self.task_patient= self.df_tasks['patient'].tolist()
        # convert dates of tasks and blocks to (week,day)
        self.convert_dates()
        return

    def convert_dates(self):
        self.week_day = [self.date2week_day(date) for date in self.dates]
        self.block_week_day = [self.date2week_day(date) for date in self.block_dates]

    def date2int(self, date, dzero="25/11/2024"):
        return  self.date2week_day(date)[0]*7 + self.date2week_day(date)[1]

    def date2week_day(self, date, dzero="25/11/2024"):
        d0 = pd.to_datetime(dzero, 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


sc = Scheduler()
print(sc.df_tasks)
print(sc.df_blocks)
print(sc.df_patients)

# Declare variables for optimization
s = sc.s
print(f's = {s}')

d = sc.d
print(f'd = {d}')

b =  sc.b
print(f'b = {b}')

B = len(b)
n = len(s)
A = sum(b)

# Import PuLP

# Define the problem
prob = LpProblem("Schedule_Tasks", LpMaximize)

# Define y
y = LpVariable.dicts('Block', [(i, t) for i in range(n) for t in range(B)], cat = 'Binary')

# Definite objective function
prob += lpSum(s[i]*b[t]*y[(i,t)] for i in range(n) for t in range(B))

# CONSTRAINTS
# Constraint #1
prob += lpSum(y[(i,t)] for i in range(n) for t in range(B)) <= A

# Constraint #2
for i in range(n):
    prob += lpSum(y[(i,t)] for t in range(B)) <= d[i]

# Constraint #3
for t in range(B):
    prob += lpSum(y[(i,t)] for i in range(n)) <= 1

prob.solve()

# Visualize solution
tasks_blocks = pd.DataFrame(columns=['Task', 'Block'])

for i in range(n):
    for t in range(B):
        if y[(i,t)].varValue == 1:
            tasks_blocks = pd.concat([tasks_blocks,
                                     pd.DataFrame({'Task': [i], 'Block': [t], 'Patient': sc.task_patient[i] })],
                                     ignore_index=True)

print(tasks_blocks)

print(sc.df_tasks)
print(sc.df_patients)
print(sc.df_blocks)
print(type(sc.df_blocks))