import numpy as np
import pandas as pd
from scipy.optimize import minimize
import time
from collections import defaultdict

def maxentdefaultparams(algorithm_params):
    """
    Fill in default parameters for the MaxEnt algorithm.
    """
    # Create default parameters.
    default_params = {
        'seed': 0,
        'laplace_prior': 0,
        'all_features': 1,
        'true_features': 0
    }
    # Set parameters.
    for key, value in default_params.items():
        if key not in algorithm_params:
            algorithm_params[key] = value
    return algorithm_params

def maxentsoftmax(q):
    """
    Compute column-wise softmax of a matrix in a numerically stable way.
    """
    # Find maximum elements along with axis=1 (columns).
    maxx = np.max(q, axis=1)
    # Compute safe softmax
    v = maxx + np.log(np.sum(np.exp(q - maxx[:, np.newaxis]), axis=1))
    return v

def linearvalueiteration(mdp_data, r_sa, vinit=None):
    """
    Run value iteration to solve a linear MDP.
    """
    states = mdp_data['states']
    actions = mdp_data['actions']
    sa_p = mdp_data['sa_p']  # Shape: (states, actions, transitions)
    sa_s = mdp_data['sa_s']
    discount = mdp_data['discount']
    VITR_THRESH = 1e-4

    # Initialize value function
    if vinit is not None:
        v = vinit.copy()
    else:
        v = np.zeros(states)

    diff = 1.0
    while diff >= VITR_THRESH:
        vp = v.copy()
        # Compute q function
        vp_sa_s = vp[sa_s]  # Shape: (states, actions, transitions)
        expected_vp = np.sum(sa_p * vp_sa_s, axis=2)  # Shape: (states, actions)
        q = r_sa + discount * expected_vp

        # Compute softmax
        v = maxentsoftmax(q)

        diff = np.max(np.abs(v - vp))
    # Compute policy
    logp = q - v[:, np.newaxis]
    p = np.exp(logp)
    return v, q, p, logp

def linearmdpsolve(mdp_data, r_sa):
    """
    Solve a linear MDP and return value function, Q function, and policy.
    """
    v, q, p, _ = linearvalueiteration(mdp_data, r_sa)
    mdp_solution = {'v': v, 'q': q, 'p': p}
    return mdp_solution

def linearmdpfrequency(mdp_data, mdp_solution, initD, prevD=None):
    """
    Compute the occupancy measure of the linear MDP given a policy.
    """
    states = mdp_data['states']
    actions = mdp_data['actions']
    transitions = mdp_data['transitions']
    VITR_THRESH = 1e-4

    if prevD is not None:
        D = prevD.copy()
    else:
        D = np.zeros(states)

    if initD is None or len(initD) == 0:
        # Initialize uniform initial state distribution.
        initD = (1.0 / states) * np.ones(states)

    diff = 1.0

    # Build the state transition matrix under policy
    P_pi = np.zeros((states, states))
    sa_p = mdp_data['sa_p']  # Shape: (states, actions, transitions)
    sa_s = mdp_data['sa_s']
    p = mdp_solution['p']

    for s in range(states):
        for a in range(actions):
            for k in range(transitions):
                s_next = sa_s[s, a, k]
                if s_next >= 0:
                    prob = p[s, a] * sa_p[s, a, k]
                    P_pi[s_next, s] += prob

    while diff >= VITR_THRESH:
        Dp = D.copy()
        D = initD + mdp_data['discount'] * P_pi @ Dp
        diff = np.max(np.abs(D - Dp))
    return D

def maxentdiscounted(r, F, muE, mu_sa, mdp_data, initD, laplace_prior=0):
    """
    Compute MaxEnt objective and gradient. Discounted infinite-horizon version.
    """
    actions = mdp_data['actions']

    # Convert to full reward.
    wts = r.copy()
    r_full = F @ r  # r_full is of shape (states,)
    # Expand r_full to shape (states, actions)
    r_sa = np.tile(r_full[:, np.newaxis], (1, actions))  # shape (states, actions)

    # Run linear value iteration
    v, q, policy, logpolicy = linearvalueiteration(mdp_data, r_sa)

    # Compute value by adding up log example probabilities.
    val = np.sum(logpolicy * mu_sa)

    # Add laplace prior.
    if laplace_prior:
        val = val - laplace_prior * np.sum(np.abs(wts))

    # Invert for descent.
    val = -val

    # Compute gradient.
    # Compute state visitation count D.
    mdp_solution = {'p': policy}
    D = linearmdpfrequency(mdp_data, mdp_solution, initD)

    # Compute gradient.
    dr = muE - F.T @ D

    # Laplace prior.
    if laplace_prior:
        dr = dr - laplace_prior * np.sign(wts)

    # Invert for descent.
    dr = -dr
    return val, dr

def maxentrun(algorithm_params, mdp_data, mdp_model_solve, feature_data,
              example_samples, true_features, verbosity=1):
    """
    Run the MaxEnt IRL algorithm.
    """
    algorithm_params = maxentdefaultparams(algorithm_params)

    # Set random seed.
    np.random.seed(algorithm_params['seed'])

    # Initialize variables.
    states = mdp_data['states']
    actions = mdp_data['actions']
    N = len(example_samples)
    T = len(example_samples[0])

    # Build feature membership matrix.
    if algorithm_params['all_features']:
        F = feature_data['splittable']
        # Add dummy feature.
        F = np.hstack([F, np.ones((states, 1))])
    elif algorithm_params['true_features']:
        F = true_features
    else:
        F = np.eye(states)

    # Count features.
    features = F.shape[1]

    # Compute feature expectations.
    muE = np.zeros(features)
    mu_sa = np.zeros((states, actions))
    for i in range(N):
        for t in range(T):
            s = example_samples[i][t][0]
            a = example_samples[i][t][1]
            mu_sa[s, a] += 1
            muE += F[s, :]

    # Generate initial state distribution for infinite horizon.
    initD = np.zeros(states)
    for i in range(N):
        s = example_samples[i][0][0]  # Initial state of each trajectory
        initD[s] += 1
    initD /= N  # Normalize to form a probability distribution

    # Define optimization function
    def fun(r):
        val, grad = maxentdiscounted(r, F, muE, mu_sa, mdp_data, initD, algorithm_params['laplace_prior'])
        return val, grad

    # Set up optimization options.
    options = {'disp': verbosity > 0, 'maxiter': 1000}

    # Initialize reward.
    r = np.random.rand(features)

    # Run unconstrained non-linear optimization.
    start_time = time.time()
    res = minimize(fun, r, method='L-BFGS-B', jac=True, options=options)
    r = res.x
    time_elapsed = time.time() - start_time
    if verbosity != 0:
        print(f'Optimization completed in {time_elapsed} seconds.')

    # Convert to full tabulated reward.
    wts = r.copy()
    r_full = F @ r
    r_sa = np.tile(r_full[:, np.newaxis], (1, actions))

    # Solving the MDP with the computed reward
    soln = mdp_model_solve(mdp_data, r_sa)
    v = soln['v']
    q = soln['q']
    p = soln['p']

    # Construct returned structure.
    irl_result = {
        'r': r_sa,
        'v': v,
        'p': p,
        'q': q,
        'r_itr': [r_sa],
        'model_itr': [wts],
        'model_r_itr': [r_sa],
        'p_itr': [p],
        'model_p_itr': [p],
        'time': time_elapsed
    }

    return irl_result

# Now, we use the functions above to read your data, construct the MDP, and run the MaxEnt IRL algorithm.

# --- Main code starts here ---


# Read the CSV data
data = pd.read_csv('/content/drive/MyDrive/BIST_10/MDP_IRL/MLB/AKBNK.csv')
states_array = data['State'].values
actions_array = data['Action'].values
next_states_array = data['Next_State'].values

# Step 1: Map states and actions to indices
# Create mappings for states and actions
unique_states = np.unique(np.concatenate([states_array, next_states_array]))
unique_actions = np.unique(actions_array)
state_to_idx = {state: idx for idx, state in enumerate(unique_states)}
action_to_idx = {action: idx for idx, action in enumerate(unique_actions)}
idx_to_state = {idx: state for idx, state in enumerate(unique_states)}
idx_to_action = {idx: action for idx, action in enumerate(unique_actions)}

num_states = len(unique_states)
num_actions = len(unique_actions)

# Convert states and actions to indices
s_indices = np.array([state_to_idx[s] for s in states_array])
a_indices = np.array([action_to_idx[a] for a in actions_array])
sp_indices = np.array([state_to_idx[sp] for sp in next_states_array])

# Step 2: Build transition probabilities and successor states
from collections import defaultdict

# Initialize transition counts and action counts
transition_counts = defaultdict(int)
action_counts = defaultdict(int)
sa_sp_dict = defaultdict(set)
for s_idx, a_idx, sp_idx in zip(s_indices, a_indices, sp_indices):
    transition_counts[(s_idx, a_idx, sp_idx)] += 1
    action_counts[(s_idx, a_idx)] += 1
    sa_sp_dict[(s_idx, a_idx)].add(sp_idx)

# Initialize sa_p with zeros
sa_p = np.zeros((num_states, num_actions, num_states))

# Compute sa_p[s_idx, a_idx, sp_idx]
for (s_idx, a_idx, sp_idx), count in transition_counts.items():
    total_count = action_counts[(s_idx, a_idx)]
    prob = count / total_count
    sa_p[s_idx, a_idx, sp_idx] = prob

# Now build sa_s and re-structure sa_p according to the number of transitions
# First, compute the number of transitions per (s_idx, a_idx)
num_transitions_per_sa = np.zeros((num_states, num_actions), dtype=int)
for (s_idx, a_idx), sp_set in sa_sp_dict.items():
    num_transitions_per_sa[s_idx, a_idx] = len(sp_set)

# Maximum number of transitions across all (s,a)
num_transitions = num_transitions_per_sa.max()

# Initialize sa_s and sa_p_new
sa_s = -np.ones((num_states, num_actions, num_transitions), dtype=int)
sa_p_new = np.zeros((num_states, num_actions, num_transitions))

# Fill in sa_s and sa_p_new
for s_idx in range(num_states):
    for a_idx in range(num_actions):
        sp_indices_sa = np.nonzero(sa_p[s_idx, a_idx, :])[0]
        probs = sa_p[s_idx, a_idx, sp_indices_sa]
        num_transitions_sa = len(sp_indices_sa)
        if num_transitions_sa > 0:
            sa_s[s_idx, a_idx, :num_transitions_sa] = sp_indices_sa
            sa_p_new[s_idx, a_idx, :num_transitions_sa] = probs
        else:
            # Handle cases where there is no data for this (s,a)
            # Assign uniform probabilities to all states
            sa_s[s_idx, a_idx, :] = np.arange(num_transitions)
            sa_p_new[s_idx, a_idx, :] = 1.0 / num_transitions

# Replace invalid entries in sa_s with 0
invalid_mask = (sa_s == -1)
sa_s[invalid_mask] = 0
sa_p_new[invalid_mask] = 0

# Step 3: Create mdp_data
mdp_data = {
    'states': num_states,
    'actions': num_actions,
    'transitions': num_transitions,
    'sa_p': sa_p_new,  # Shape: (states, actions, transitions)
    'sa_s': sa_s,      # Shape: (states, actions, transitions)
    'discount': 0.99
}

# Step 4: Prepare example samples
# Since we don't have trajectory identifiers, we will treat the entire sequence as a single trajectory
example_samples = []
trajectory = []
for s_idx, a_idx in zip(s_indices, a_indices):
    trajectory.append((s_idx, a_idx))
example_samples.append(trajectory)

# Step 5: Define feature data
# For simplicity, we'll use an identity matrix as features
feature_data = {
    'splittable': np.eye(num_states)
}

# Step 6: Define algorithm parameters
algorithm_params = {
    'seed': 42,
    'laplace_prior': 0.1,
    'all_features': 1,
    'true_features': 0
}

# Step 7: Define mdp_model_solve as linearmdpsolve
def mdp_model_solve(mdp_data, r_sa):
    return linearmdpsolve(mdp_data, r_sa)

# Step 8: Run MaxEnt IRL
irl_result = maxentrun(algorithm_params, mdp_data, mdp_model_solve, feature_data,
                       example_samples, true_features=None, verbosity=1)

# Step 9: Access the results
reward = irl_result['r']  # Shape: (states, actions)
value_function = irl_result['v']  # Shape: (states,)
policy = irl_result['p']  # Shape: (states, actions)

# Convert state and action indices back to actual states and actions
policy_dict = {}
for s_idx in range(num_states):
    state = idx_to_state[s_idx]
    action_probs = policy[s_idx]
    action_probs_dict = {idx_to_action[a_idx]: action_probs[a_idx] for a_idx in range(num_actions)}
    policy_dict[state] = action_probs_dict

# Step 10: Print the results
print("Reward Function (state-action values):")
for s_idx in range(num_states):
    state = idx_to_state[s_idx]
    print(f"State {state}:")
    for a_idx in range(num_actions):
        action = idx_to_action[a_idx]
        print(f"  Action {action}: Reward {reward[s_idx, a_idx]:.4f}")
print("\nValue Function:")
for s_idx in range(num_states):
    state = idx_to_state[s_idx]
    print(f"State {state}: Value {value_function[s_idx]:.4f}")
print("\nPolicy (probabilities for each action in each state):")
for state, actions in policy_dict.items():
    print(f"State {state}:")
    for action, prob in actions.items():
        print(f"  Action {action}: Probability {prob:.4f}")

Kodu kontrol edebilir misin? Değişmesi gereken her yeri yaz lütfen.