"""
Reinforcement Learning Utility Functions
This module provides commonly used functions for RL algorithms.
"""

import numpy as np


def epsilon_greedy_action(q_values, epsilon, num_actions):
    """
    Select an action using epsilon-greedy policy.
    
    Args:
        q_values (list or np.array): Q-values for each action in current state
        epsilon (float): Exploration rate (0 to 1)
        num_actions (int): Total number of possible actions
    
    Returns:
        int: Selected action index
    """
    if np.random.random() < epsilon:
        # Explore: choose random action
        return np.random.randint(0, num_actions)
    else:
        # Exploit: choose best action
        return np.argmax(q_values)


def q_learning_update(q_table, state, action, reward, next_state, alpha, gamma, num_actions):
    """
    Update Q-table using Q-learning algorithm.
    
    Args:
        q_table (dict or np.array): Q-table storing state-action values
        state: Current state (key if dict, index if array)
        action (int): Action taken
        reward (float): Reward received
        next_state: Next state after taking action
        alpha (float): Learning rate (0 to 1)
        gamma (float): Discount factor (0 to 1)
        num_actions (int): Number of possible actions
    
    Returns:
        Updated Q-table and new Q-value for the state-action pair
    """
    if isinstance(q_table, dict):
        # Dictionary-based Q-table
        current_q = q_table.get((state, action), 0.0)
        next_max_q = max([q_table.get((next_state, a), 0.0) for a in range(num_actions)])
        new_q = current_q + alpha * (reward + gamma * next_max_q - current_q)
        q_table[(state, action)] = new_q
    else:
        # Array-based Q-table
        current_q = q_table[state, action]
        next_max_q = np.max(q_table[next_state])
        new_q = current_q + alpha * (reward + gamma * next_max_q - current_q)
        q_table[state, action] = new_q
    
    return q_table, new_q


def softmax_action_selection(q_values, temperature=1.0):
    """
    Select an action using softmax (Boltzmann) distribution.
    
    Args:
        q_values (list or np.array): Q-values for each action
        temperature (float): Temperature parameter (higher = more exploration)
    
    Returns:
        int: Selected action index based on softmax probabilities
    """
    q_values = np.array(q_values)
    # Apply temperature scaling
    exp_values = np.exp((q_values - np.max(q_values)) / temperature)
    probabilities = exp_values / np.sum(exp_values)
    
    # Sample action based on probabilities
    action = np.random.choice(len(q_values), p=probabilities)
    return action


def initialize_q_table(num_states, num_actions, init_value=0.0):
    """
    Initialize a Q-table with specified dimensions.
    
    Args:
        num_states (int): Number of states
        num_actions (int): Number of actions
        init_value (float): Initial value for all Q-values
    
    Returns:
        np.array: Initialized Q-table of shape (num_states, num_actions)
    """
    return np.full((num_states, num_actions), init_value)

