import simpy
import random
import numpy as np
from collections import deque, defaultdict

# Global counter for patient IDs
patient_counter = 0


# Define the Patient class, representing a patient in the system
class Patient:
    def __init__(self, env, urgency, department, service_time):
        global patient_counter
        self.id = patient_counter  # Assign a unique ID to each patient
        patient_counter += 1

        self.env = env
        self.urgency = urgency  # Urgency level of the patient (1 is most urgent, 3 is least)
        self.department = department  # Department the patient is assigned to
        self.service_time = service_time  # Time required to serve the patient
        self.arrival_time = env.now  # Time when the patient arrives
        self.wait_time = None  # Time the patient spends waiting before being served


# Define the Department class, representing a department in the hospital
class Department:
    def __init__(self, env, name, rl_agent):
        self.env = env
        self.name = name  # Name of the department
        self.queue = deque()  # Queue to hold patients waiting in the department
        self.rl_agent = rl_agent  # Reference to the RL agent for decision-making
        self.action = env.process(self.run())  # Start the department's main process
        self.served_patients = []  # List to store wait times of served patients
        self.patient_ids = []  # List to store IDs of patients waiting in the department

    def run(self):
        # Main loop for serving patients in the department
        while True:
            if self.queue:
                # Get the current state and select a patient based on urgency and arrival time (FCFS)
                state = self.rl_agent.get_state(self)
                patient = self.select_patient(state)

                if patient:
                    # Serve the selected patient and calculate wait time
                    yield self.env.timeout(patient.service_time)
                    patient.wait_time = self.env.now - patient.arrival_time
                    self.served_patients.append(patient.wait_time)

                    # Update the RL agent's Q-table with the observed reward
                    next_state = self.rl_agent.get_state(self)
                    reward = self.calculate_reward(patient)
                    self.rl_agent.update_q_table(state, self.queue.index(patient), reward, next_state)

                    # Log patient service information
                    print(
                        f"Patient {patient.id} with urgency {patient.urgency} served in {self.name} at time {self.env.now} after waiting {patient.wait_time}")

                    # Remove the patient from the queue and the patient ID list
                    self.queue.remove(patient)
                    self.patient_ids.remove(patient.id)
            else:
                # If the queue is empty, wait for a short time before checking again
                yield self.env.timeout(1)

    def add_patient(self, patient):
        # Add a new patient to the department's queue and track their ID
        self.queue.append(patient)
        self.patient_ids.append(patient.id)
        print(f"Patient {patient.id} with urgency {patient.urgency} arrives at {self.name} at time {self.env.now}")

    def calculate_reward(self, patient):
        # Calculate the reward based on the patient's urgency and wait time
        # The reward is higher for shorter wait times and more urgent patients
        base_reward = 100  # A base reward for serving the patient
        wait_time_penalty = -patient.wait_time  # Negative impact of waiting time
        urgency_multiplier = 4 - patient.urgency  # Urgency multiplier (higher for more urgent patients)

        # Final reward calculation
        reward = base_reward + (urgency_multiplier * wait_time_penalty)

        return reward

    def get_average_wait_time(self):
        # Calculate the average wait time for all served patients
        if self.served_patients:
            return np.mean(self.served_patients)
        return 0.0

    def select_patient(self, state):
        # Select the patient to be served next based on urgency and FCFS within the same urgency
        sorted_queue = sorted(self.queue, key=lambda p: (p.urgency, p.arrival_time))
        return sorted_queue[0] if sorted_queue else None


# Define the RL Agent class, responsible for decision-making and learning
class RLAgent:
    def __init__(self, departments):
        self.q_table = {}  # Q-table for storing state-action values
        self.departments = departments  # List of departments in the system

    def get_state(self, department):
        # Get the current state, represented by the length of the queue in each department
        return tuple(len(d.queue) for d in self.departments)

    def select_action(self, state, queue_length):
        # Select an action (which patient to serve) using an epsilon-greedy strategy
        if random.random() < 0.1:
            # Exploration: Choose a random action
            return random.choice(range(queue_length))
        # Exploitation: Choose the best action based on the Q-values
        q_values = [self.q_table.get((state, a), 0) for a in range(queue_length)]
        return np.argmax(q_values)

    def update_q_table(self, state, action, reward, next_state):
        # Update the Q-value for the given state-action pair using the reward and future rewards
        old_value = self.q_table.get((state, action), 0)
        future_rewards = max([self.q_table.get((next_state, a), 0) for a in range(len(self.departments))], default=0)
        self.q_table[(state, action)] = old_value + 0.1 * (reward + 0.9 * future_rewards - old_value)


# Function to generate patients and assign them to departments
def patient_generator(env, departments, arrival_rate):
    while True:
        # Stop generating new patients after time 500
        if env.now > 500:
            break

        # Generate a new patient with weighted urgency levels
        urgency_distribution = [0.1, 0.3, 0.6]  # Probability distribution for urgency levels [1, 2, 3]
        urgency = random.choices([1, 2, 3], weights=urgency_distribution, k=1)[0]

        department = random.choice(departments)

        # Determine service time based on the department
        if department.name == "Department 1":
            service_time = max(1, np.random.normal(10, 3))  # Mean 10, SD 3
        elif department.name == "Department 2":
            service_time = max(1, np.random.normal(8, 3))  # Mean 8, SD 3
        elif department.name == "Department 3":
            service_time = 11  # Fixed service time of 11 minutes

        patient = Patient(env, urgency, department, service_time)
        department.add_patient(patient)

        # Wait for the next patient to arrive based on the arrival rate
        yield env.timeout(random.expovariate(arrival_rate))


# Function to reset the departments' state between simulation episodes
def reset_departments(departments):
    for department in departments:
        department.queue.clear()
        department.patient_ids.clear()
        department.served_patients.clear()


# Function to run one episode of the simulation and return the average wait time
def run_episode(env, rl_agent, departments, simulation_time):
    # Start the patient generator process
    env.process(patient_generator(env, departments, arrival_rate=0.3))  # Increase the patient arrival rate

    # Run the simulation until time 540 or until all patients are served
    env.run(until=simulation_time)

    # After 500, ensure the simulation continues until all patients are served
    while any(len(dept.queue) > 0 for dept in departments):
        env.run(until=env.now + 1)  # Continue running until all queues are empty

    # Calculate the average wait time for all departments
    average_wait_times = [dept.get_average_wait_time() for dept in departments]
    return np.mean(average_wait_times)


# Function to set up the simulation environment and run multiple episodes
def run_simulation(episodes=100, simulation_time=540):
    first_round_avg = 0
    last_round_avg = 0

    for episode in range(episodes):
        print(f"\nStarting episode {episode + 1}")

        # Reset the patient counter at the start of each episode
        global patient_counter
        patient_counter = 0

        env = simpy.Environment()
        rl_agent = RLAgent(departments=[])
        departments = [Department(env, f"Department {i + 1}", rl_agent) for i in range(3)]
        rl_agent.departments = departments

        # Run one episode and record the average wait time
        avg_wait_time = run_episode(env, rl_agent, departments, simulation_time)
        print(f"Average wait time for episode {episode + 1}: {avg_wait_time}")

        # Store the average wait time for the first and last episodes
        if episode == 0:
            first_round_avg = avg_wait_time
        if episode == episodes - 1:
            last_round_avg = avg_wait_time

        # Reset the departments for the next episode
        reset_departments(departments)

    # Print the average wait times for the first and last episodes
    print(f"\nAverage wait time in the first episode: {first_round_avg}")
    print(f"Average wait time in the last episode: {last_round_avg}")


# Run the simulation
if __name__ == "__main__":
    run_simulation(episodes=100, simulation_time=540)