import numpy as np
from scipy.optimize import minimize

def filldefaultparams(params, default_params):
    """Fill in missing parameters with default values."""
    for key, value in default_params.items():
        if key not in params:
            params[key] = value
    return params

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
    algorithm_params = filldefaultparams(algorithm_params, default_params)
    return algorithm_params

def maxentsoftmax(q):
    """Compute column-wise softmax of a matrix in a numerically stable way."""
    # q is a 2D array of shape (states, actions)
    maxx = np.max(q, axis=1)  # Max over actions for each state
    # Compute safe softmax
    v = maxx + np.log(np.sum(np.exp(q - maxx[:, np.newaxis]), axis=1))
    return v

def linearvalueiteration(mdp_data, r, vinit=None):
    """Run value iteration to solve a linear MDP."""
    states = mdp_data['states']
    actions = mdp_data['actions']
    VITR_THRESH = 1e-4  # Threshold for convergence
    if vinit is not None:
        v = vinit
    else:
        v = np.zeros(states)
    diff = 1.0
    sa_p = mdp_data['sa_p']
    sa_s = mdp_data['sa_s']
    discount = mdp_data['discount']
    transitions = sa_p.shape[2]

    while diff >= VITR_THRESH:
        vp = v.copy()
        expected_v = np.zeros((states, actions))
        for s in range(states):
            for a in range(actions):
                for k in range(transitions):
                    sp = sa_s[s, a, k]
                    prob = sa_p[s, a, k]
                    expected_v[s, a] += prob * vp[sp]
        q = r + discount * expected_v
        v = maxentsoftmax(q)
        diff = np.max(np.abs(v - vp))
    # Compute policy
    logp = q - v[:, np.newaxis]  # Subtract v from each row
    p = np.exp(logp)
    return v, q, p, logp

def linearmdpfrequency(mdp_data, mdp_solution, initD=None, prevD=None):
    """Compute the occupancy measure of the linear MDP given a policy."""
    sa_p = mdp_data['sa_p']
    sa_s = mdp_data['sa_s']
    discount = mdp_data['discount']
    states, actions, transitions = sa_p.shape
    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
    while diff >= VITR_THRESH:
        Dp = D.copy()
        Dpi = np.tile(mdp_solution['p'][:, :, np.newaxis], (1, 1, transitions)) * sa_p * \
              Dp[:, np.newaxis, np.newaxis] * discount
        D_new = initD.copy()
        for s in range(states):
            for a in range(actions):
                for k in range(transitions):
                    sp = sa_s[s, a, k]
                    increment = Dpi[s, a, k]
                    D_new[sp] += increment
        D = D_new
        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 for discounted infinite-horizon."""
    actions = mdp_data['actions']
    # Convert to full reward
    wts = r.copy()
    r = F.dot(r)
    # Replicate r across actions
    r_full = np.tile(r[:, np.newaxis], (1, actions))
    # Run linear value iteration
    v, q, policy, logpolicy = linearvalueiteration(mdp_data, r_full)
    # 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
    D = linearmdpfrequency(mdp_data, {'p': policy}, initD)
    dr = muE - F.T.dot(D)
    # Add Laplace prior to gradient
    if laplace_prior:
        dr = dr - laplace_prior * np.sign(wts)
    # Invert for descent
    dr = -dr
    return val, dr

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

def maxentrun(algorithm_params, mdp_data, mdp_model, feature_data, example_samples, true_features, verbosity=1):
    """Ziebart's Maximum Entropy IRL, with optional prior on reward values."""
    # Fill in default parameters
    algorithm_params = maxentdefaultparams(algorithm_params)
    # Set random seed
    np.random.seed(algorithm_params['seed'])
    # Initialize variables
    states = mdp_data['states']
    actions = mdp_data['actions']
    transitions = mdp_data['sa_p'].shape[2]
    N = len(example_samples)
    T = len(example_samples[0])  # Assuming example_samples is a list of trajectories
    mdp_solve = linearmdpsolve  # Since we're using a linear MDP solver

    # 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 traj in example_samples:
        for (s, a) in traj:
            mu_sa[s, a] += 1
            muE += F[s, :]

    # Generate initial state distribution for infinite horizon
    initD = np.zeros(states)
    for traj in example_samples:
        for (s, _) in traj:
            initD[s] += 1
    for traj in example_samples:
        for (s, a) in traj:
            for k in range(transitions):
                sp = mdp_data['sa_s'][s, a, k]
                initD[sp] -= mdp_data['discount'] * mdp_data['sa_p'][s, a, k]

    # Define the function to be minimized
    def func(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}
    # Initialize reward
    r0 = np.random.rand(features)
    # Run optimization
    res = minimize(func, r0, method='L-BFGS-B', jac=True, options=options)
    r = res.x
    # Convert to full tabulated reward
    wts = r.copy()
    r = F.dot(r)
    # Return corresponding reward function
    r_full = np.tile(r[:, np.newaxis], (1, actions))
    soln = mdp_solve(mdp_data, r_full)
    v = soln['v']
    q = soln['q']
    p = soln['p']
    # Construct returned structure
    irl_result = {
        'r': r_full,
        'v': v,
        'p': p,
        'q': q,
        'r_itr': [r_full],
        'model_itr': [wts],
        'model_r_itr': [r_full],
        'p_itr': [p],
        'model_p_itr': [p],
        'time': res.nit  # Number of iterations
    }
    return irl_result

Yukarıdaki kodu detaylı bir şekilde kontrol edebilir misin