!pip install aicrowd-cli > /dev/null

import os

AICROWD_DATASET_PATH = os.getenv("DATASET_PATH", os.getcwd()+"/13d77bb0-b325-4e95-a03b-833eb6694acd_a2_taxi_inputs.zip")
AICROWD_RESULTS_DIR = os.getenv("OUTPUTS_DIR", "results")

!unzip $AICROWD_DATASET_PATH

DATASET_DIR = 'inputs/'

import numpy as np

class TaxiEnv_HW2:
    def __init__(self, states, actions, probabilities, rewards, initial_policy):
        self.possible_states = states
        self._possible_actions = {st: ac for st, ac in zip(states, actions)}
        self._ride_probabilities = {st: pr for st, pr in zip(states, probabilities)}
        self._ride_rewards = {st: rw for st, rw in zip(states, rewards)}
        self.initial_policy = initial_policy
        self._verify()

    def _check_state(self, state):
        assert state in self.possible_states, "State %s is not a valid state" % state

    def _verify(self):
        """ 
        Verify that data conditions are met:
        Number of actions matches shape of next state and actions
        Every probability distribution adds up to 1 
        """
        ns = len(self.possible_states)
        for state in self.possible_states:
            ac = self._possible_actions[state]
            na = len(ac)

            rp = self._ride_probabilities[state]
            assert np.all(rp.shape == (na, ns)), "Probabilities shape mismatch"
        
            rr = self._ride_rewards[state]
            assert np.all(rr.shape == (na, ns)), "Rewards shape mismatch"

            assert np.allclose(rp.sum(axis=1), 1), "Probabilities don't add up to 1"

    def possible_actions(self, state):
        """ Return all possible actions from a given state """
        self._check_state(state)
        return self._possible_actions[state]

    def ride_probabilities(self, state, action):
        """ 
        Returns all possible ride probabilities from a state for a given action
        For every action a list with the returned with values in the same order as self.possible_states
        """
        actions = self.possible_actions(state)
        ac_idx = actions.index(action)
        return self._ride_probabilities[state][ac_idx]

    def ride_rewards(self, state, action):
        actions = self.possible_actions(state)
        ac_idx = actions.index(action)
        return self._ride_rewards[state][ac_idx]

def check_taxienv():
    # These are the values as used in the pdf, but they may be changed during submission, so do not hardcode anything

    states = ['A', 'B', 'C']

    actions = [['1','2','3'], ['1','2'], ['1','2','3']]

    probs = [np.array([[1/2,  1/4,  1/4],
                    [1/16, 3/4,  3/16],
                    [1/4,  1/8,  5/8]]),

            np.array([[1/2,   0,     1/2],
                    [1/16,  7/8,  1/16]]),

            np.array([[1/4,  1/4,  1/2],
                    [1/8,  3/4,  1/8],
                    [3/4,  1/16, 3/16]]),]

    rewards = [np.array([[10,  4,  8],
                        [ 8,  2,  4],
                        [ 4,  6,  4]]),

            np.array([[14,  0, 18],
                        [ 8, 16,  8]]),

            np.array([[10,  2,  8],
                        [6,   4,  2],
                        [4,   0,  8]]),]
    initial_policy = {'A': '1', 'B': '1', 'C': '1'}

    env = TaxiEnv_HW2(states, actions, probs, rewards, initial_policy)
    print("All possible states", env.possible_states)
    print("All possible actions from state B", env.possible_actions('B'))
    print("Ride probabilities from state A with action 2", env.ride_probabilities('A', '2'))
    print("Ride rewards from state C with action 3", env.ride_rewards('C', '3'))

    base_kwargs = {"states": states, "actions": actions, 
                "probabilities": probs, "rewards": rewards,
                "initial_policy": initial_policy}
    return base_kwargs

base_kwargs = check_taxienv()
env = TaxiEnv_HW2(**base_kwargs)

# 1.1 Policy Iteration
def policy_iteration(taxienv, gamma):
    # A list of all the states
    states = taxienv.possible_states
    # Initial values
    values = {s: 0 for s in states}

    # This is a dictionary of states to policies -> e.g {'A': '1', 'B': '2', 'C': '1'}
    policy = taxienv.initial_policy.copy()

    ## Begin code here

    # Hints - 
    # Do not hardcode anything
    # Only the final result is required for the results
    # Put any extra data in "extra_info" dictonary for any plots etc
    # Use the helper functions taxienv.ride_rewards, taxienv.ride_probabilities,  taxienv.possible_actions
    # For terminating condition use the condition exactly mentioned in the pdf

    while True :
        while True :
            delta = 0
            for state in taxienv.possible_states :
                J = values[state]
                action = policy[state]
                probabs = taxienv.ride_probabilities(state, action)
                rewards = taxienv.ride_rewards(state, action)
                sum = 0
                for k in range(len(probabs)) :
                    sum += probabs[k] * (rewards[k] + gamma * values[taxienv.possible_states[k]])
                values[state] = sum
                delta = max(delta, abs(J - values[state]))
            if delta < 1e-8 : 
                break
        done = 1
        for state in taxienv.possible_states :
            B = policy[state]
            temp_array = []
            for action in taxienv.possible_actions(state) :
                probabs = taxienv.ride_probabilities(state, action)
                rewards = taxienv.ride_rewards(state, action)
                sum = 0
                for k in range(len(probabs)) :
                    sum += probabs[k] * (rewards[k] + gamma * values[taxienv.possible_states[k]])
                temp_array.append(sum)
            policy[state] = taxienv.possible_actions(state)[np.argmax(temp_array)]
            if policy[state] != B :
                done = 0;
        if done == 1 : 
            break
    
    # Put your extra information needed for plots etc in this dictionary
    extra_info = {}

    """
    # Code for printing results 
    string = ""
    for key in results.keys() :
        string += "| $\gamma$ = " + str(key) + " | A | B | C |" + "\n"
        string += "|---|---|---|---|\n"
        for key1 in results[key]['Expected Reward'].keys() :
            value = results[key]['Expected Reward'][key1]
            results[key]['Expected Reward'][key1] = round(value, 3)
        string += "| Expected Reward |" + str(results[key]['Expected Reward']['A']) + "|" + str(results[key]['Expected Reward']['B']) + "|" + str(results[key]['Expected Reward']['C']) + "|\n"
        string += "| Policy |" + str(results[key]['Policy']['A']) + "|" + str(results[key]['Policy']['B']) + "|" + str(results[key]['Policy']['C']) + "|\n\n___\n\n"
    print(string)
             
    """
    ## Do not edit below this line

    # Final results
    return {"Expected Reward": values, "Policy": policy}, extra_info

# 1.2 Policy Iteration with different values of gamma
def run_policy_iteration(env):
    gamma_values = np.arange(5, 100, 5)/100
    results, extra_info = {}, {}
    for gamma in gamma_values:
        results[gamma], extra_info[gamma] = policy_iteration(env, gamma)
    return results, extra_info

results, extra_info = run_policy_iteration(env)

# 1.3 Modified Policy Iteration
def modified_policy_iteration(taxienv, gamma, m):
    # A list of all the states
    states = taxienv.possible_states
    # Initial values
    values = {s: 0 for s in states}

    # This is a dictionary of states to policies -> e.g {'A': '1', 'B': '2', 'C': '1'}
    policy = taxienv.initial_policy.copy()

    ## Begin code here

    # Hints - 
    # Do not hardcode anything
    # Only the final result is required for the results
    # Put any extra data in "extra_info" dictonary for any plots etc
    # Use the helper functions taxienv.ride_rewards, taxienv.ride_probabilities,  taxienv.possible_actions
    # For terminating condition use the condition exactly mentioned in the pdf

    H = {s: 0 for s in states}
    while True : 
        for k in range(m) :
            for state in taxienv.possible_states : 
                temp_array = []
                for action in taxienv.possible_actions(state) : 
                    rewards = taxienv.ride_rewards(state, action)
                    probabs = taxienv.ride_probabilities(state, action)
                    sum = 0
                    for i in range(len(probabs)) :
                        sum += probabs[i] * (rewards[i] + gamma * values[taxienv.possible_states[i]])
                    temp_array.append(sum)
                H[state] = max(temp_array)
            for state in taxienv.possible_states : 
                values[state] = H[state]
        done = 1
        for state in taxienv.possible_states : 
            B = policy[state]
            temp_array = []
            for action in taxienv.possible_actions(state) :
                rewards = taxienv.ride_rewards(state, action)
                probabs = taxienv.ride_probabilities(state, action)
                sum = 0
                for i in range(len(probabs)) :
                    sum += probabs[i] * (rewards[i] + gamma * values[taxienv.possible_states[i]])
                temp_array.append(sum)
            policy[state] = taxienv.possible_actions(state)[np.argmax(temp_array)]
            if policy[state] != B :
                done = 0
        if done == 1 : 
            break
    
    # Put your extra information needed for plots etc in this dictionary
    extra_info = {}

    ## Do not edit below this line


    # Final results
    return {"Expected Reward": values, "Policy": policy}, extra_info

def run_modified_policy_iteration(env):
    m_values = np.arange(1, 15)
    gamma = 0.9
    results, extra_info = {}, {}
    for m in m_values:
        results[m], extra_info[m] = modified_policy_iteration(env, gamma, m)
    return results, extra_info

results, extra_info = run_modified_policy_iteration(env)

for key in results.keys() :
    print(key)
    print(results[key]['Expected Reward'])
    print(results[key]['Policy'])

# 1.4 Value Iteration
def value_iteration(taxienv, gamma):
    # A list of all the states
    states = taxienv.possible_states
    # Initial values
    values = {s: 0 for s in states}

    # This is a dictionary of states to policies -> e.g {'A': '1', 'B': '2', 'C': '1'}
    policy = taxienv.initial_policy.copy()

    ## Begin code here

    # Hints - 
    # Do not hardcode anything
    # Only the final result is required for the results
    # Put any extra data in "extra_info" dictonary for any plots etc
    # Use the helper functions taxienv.ride_rewards, taxienv.ride_probabilities,  taxienv.possible_actions
    # For terminating condition use the condition exactly mentioned in the pdf
    H = {s: 0 for s in states}
    values_dict = {}
    iteration = 0
    while True : 
        delta = 0
        for state in taxienv.possible_states : 
            temp_array = []
            for action in taxienv.possible_actions(state) : 
                rewards = taxienv.ride_rewards(state, action)
                probabs = taxienv.ride_probabilities(state, action)
                sum = 0
                for i in range(len(probabs)) :
                    sum += probabs[i] * (rewards[i] + gamma * values[taxienv.possible_states[i]])
                temp_array.append(sum)
            H[state] = max(temp_array)
            delta = max(delta, abs(H[state] - values[state]))
        for state in taxienv.possible_states : 
            values[state] = H[state]
            values_dict[iteration] = values.copy()
        if delta < 1e-8 :
            break
        iteration += 1

    for state in taxienv.possible_states : 
        temp_array = []
        for action in taxienv.possible_actions(state) :
            rewards = taxienv.ride_rewards(state, action)
            probabs = taxienv.ride_probabilities(state, action)
            sum = 0
            for i in range(len(probabs)) :
                sum += probabs[i] * (rewards[i] + gamma * values[taxienv.possible_states[i]])
            temp_array.append(sum)
        policy[state] = taxienv.possible_actions(state)[np.argmax(temp_array)]

    # Put your extra information needed for plots etc in this dictionary
    extra_info = {'Values_Dict': values_dict}

    ## Do not edit below this line

    # Final results
    return {"Expected Reward": values, "Policy": policy}, extra_info

def run_value_iteration(env):
    gamma_values = np.arange(5, 100, 5)/100
    results = {}
    results, extra_info = {}, {}
    for gamma in gamma_values:
        results[gamma], extra_info[gamma] = value_iteration(env, gamma)
    return results, extra_info
  
results, extra_info = run_value_iteration(env)

# 1.4 Gauss Seidel Value Iteration
def gauss_seidel_value_iteration(taxienv, gamma):
    # A list of all the states
    # For Gauss Seidel Value Iteration - iterate through the values in the same order
    states = taxienv.possible_states

    # Initial values
    values = {s: 0 for s in states}

    # This is a dictionary of states to policies -> e.g {'A': '1', 'B': '2', 'C': '1'}
    policy = taxienv.initial_policy.copy()

    # Hints - 
    # Do not hardcode anything
    # For Gauss Seidel Value Iteration - iterate through the values in the same order as taxienv.possible_states
    # Only the final result is required for the results
    # Put any extra data in "extra_info" dictonary for any plots etc
    # Use the helper functions taxienv.ride_rewards, taxienv.ride_probabilities,  taxienv.possible_actions
    # For terminating condition use the condition exactly mentioned in the pdf

    ## Begin code here
    iteration = 0
    values_dict = {}
    while True : 
        delta = 0 
        for state in taxienv.possible_states : 
            J = values[state]
            temp_array = []
            for action in taxienv.possible_actions(state) :
                rewards = taxienv.ride_rewards(state, action)
                probabs = taxienv.ride_probabilities(state, action)
                sum = 0
                for i in range(len(probabs)) :
                    sum += probabs[i] * (rewards[i] + gamma * values[taxienv.possible_states[i]])
                temp_array.append(sum)
            values[state] = max(temp_array)
            delta = max(delta, abs(J - values[state]))
        values_dict[iteration] = values.copy()
        if delta < 1e-8 :
            break
        iteration += 1
    
    for state in taxienv.possible_states :
        temp_array = []
        for action in taxienv.possible_actions(state) :
            rewards = taxienv.ride_rewards(state, action)
            probabs = taxienv.ride_probabilities(state, action)
            sum = 0
            for i in range(len(probabs)) :
                sum += probabs[i] * (rewards[i] + gamma * values[taxienv.possible_states[i]])
            temp_array.append(sum)
        policy[state] = taxienv.possible_actions(state)[np.argmax(temp_array)]
    # Put your extra information needed for plots etc in this dictionary


    extra_info = {'Values_Dict' : values_dict}

    ## Do not edit below this line

    # Final results
    return {"Expected Reward": values, "Policy": policy}, extra_info

def run_gauss_seidel_value_iteration(env):
    gamma_values = np.arange(5, 100, 5)/100
    results = {}
    results, extra_info = {}, {}
    for gamma in gamma_values:
        results[gamma], extra_info[gamma] = gauss_seidel_value_iteration(env, gamma)
    return results, extra_info

results, extra_info = run_gauss_seidel_value_iteration(env)

# Do not edit this cell
def get_results(kwargs):

    taxienv = TaxiEnv_HW2(**kwargs)

    policy_iteration_results = run_policy_iteration(taxienv)[0]
    modified_policy_iteration_results = run_modified_policy_iteration(taxienv)[0]
    value_iteration_results = run_value_iteration(taxienv)[0]
    gs_vi_results = run_gauss_seidel_value_iteration(taxienv)[0]

    final_results = {}
    final_results["policy_iteration"] = policy_iteration_results
    final_results["modifed_policy_iteration"] = modified_policy_iteration_results
    final_results["value_iteration"] = value_iteration_results
    final_results["gauss_seidel_iteration"] = gs_vi_results

    return final_results

# Do not edit this cell, generate results with it as is
if not os.path.exists(AICROWD_RESULTS_DIR):
    os.mkdir(AICROWD_RESULTS_DIR)

for params_file in os.listdir(DATASET_DIR):
  kwargs = np.load(os.path.join(DATASET_DIR, params_file), allow_pickle=True).item()
  results = get_results(kwargs)
  idx = params_file.split('_')[-1][:-4]
  np.save(os.path.join(AICROWD_RESULTS_DIR, 'results_' + idx), results)

# Check your score on the given test cases (There are more private test cases not provided)
target_folder = 'targets'
result_folder = AICROWD_RESULTS_DIR

def check_algo_match(results, targets):
    param_matches = []
    for k in results:
        param_results = results[k]
        param_targets = targets[k]
        policy_match = param_results['Policy'] == param_targets['Policy']
        rv = [v for k, v in param_results['Expected Reward'].items()]
        tv = [v for k, v in param_targets['Expected Reward'].items()]
        rewards_match = np.allclose(rv, tv, rtol=3)
        equal = rewards_match and policy_match
        param_matches.append(equal)
    return np.mean(param_matches)

def check_score(target_folder, result_folder):
    match = []
    for out_file in os.listdir(result_folder):
        res_file = os.path.join(result_folder, out_file)
        results = np.load(res_file, allow_pickle=True).item()
        idx = out_file.split('_')[-1][:-4]  # Extract the file number
        target_file = os.path.join(target_folder, f"targets_{idx}.npy")
        targets = np.load(target_file, allow_pickle=True).item()
        algo_match = []
        for k in targets:
            algo_results = results[k]
            algo_targets = targets[k]
            algo_match.append(check_algo_match(algo_results, algo_targets))
        match.append(np.mean(algo_match))
    return np.mean(match)

if os.path.exists(target_folder):
    print("Shared data Score (normalized to 1):", check_score(target_folder, result_folder))

## Visualize policy iteration with multiple values of 
import matplotlib.pyplot as plt
# 1.2 Policy Iteration with different values of gamma
def run_policy_iteration(env):
    gamma_values = np.arange(5, 100, 5)/100
    results, extra_info = {}, {}
    for gamma in gamma_values:
        results[gamma], extra_info[gamma] = policy_iteration(env, gamma)
    return results, extra_info

results, extra_info = run_policy_iteration(env)

rewards = {}
policies = {}
for key in results.keys() :
    rewards[key] = results[key]['Expected Reward']
    policies[key] = results[key]['Policy']

rewards_A = []
rewards_B = []
rewards_C = []

gamma_list = list(results.keys())
for gamma in gamma_list :
    rewards_A.append(rewards[gamma]['A'])
    rewards_B.append(rewards[gamma]['B'])
    rewards_C.append(rewards[gamma]['C'])

plt.title('Rewards vs $\gamma$', fontsize = 20)
plt.plot(gamma_list, rewards_A, marker = 'o', linestyle = "--", color = 'red', label = "State A")
plt.plot(gamma_list, rewards_B, marker = '^', linestyle = "--", color = 'blue', label = "State B")
plt.plot(gamma_list, rewards_C, marker = '>', linestyle = "--", color = 'green', label = "State C")
plt.legend()
plt.xlabel('$\gamma$', fontsize = 15)
plt.ylabel('Rewards', fontsize = 15)
fig = plt.gcf()
fig.set_size_inches(10, 8)
plt.show()

def plotter(gamma) :
    gs_results, gs_extra_info = gauss_seidel_value_iteration(env, gamma)
    results, extra_info = value_iteration(env, gamma)

    gs_value_dict = gs_extra_info['Values_Dict']
    value_dict = extra_info['Values_Dict']

    iteration_value_dict = list(value_dict.keys())
    iteration_gs_value_dict = list(gs_value_dict.keys())

    value = {s: [] for s in ['A', 'B', 'C']}
    gs_value = {s: [] for s in ['A', 'B', 'C']}

    for key in iteration_value_dict :
        for alpha in ['A', 'B', 'C'] :
            value[alpha].append(value_dict[key][alpha])

    for key in iteration_gs_value_dict :
        for alpha in ['A', 'B', 'C'] :
            gs_value[alpha].append(gs_value_dict[key][alpha])

    fig, axs = plt.subplots(1, 3)
    axs[0].plot(iteration_value_dict, value['A'], color = 'red', linestyle = '--', marker = 'o', label = 'VI')
    axs[0].plot(iteration_gs_value_dict, gs_value['A'], color = 'red', linestyle = 'solid', marker = 'o', label = 'Gauss-Seidel VI')
    axs[0].set_xlabel("Iteration", fontsize = 15)
    axs[0].set_ylabel("Values for State A", fontsize = 15)
    axs[0].legend(loc = 'lower right', fontsize = 15)
    axs[0].set_title("Values vs Iteration for State A for $\gamma$ = " + str(gamma), fontsize = 20)

    axs[1].plot(iteration_value_dict, value['B'], color = 'blue', linestyle = '--', marker = '^', label = 'VI')
    axs[1].plot(iteration_gs_value_dict, gs_value['B'], color = 'blue', linestyle = 'solid', marker = '^', label = 'Gauss-Seidel VI')
    axs[1].set_xlabel("Iteration", fontsize = 15)
    axs[1].set_ylabel("Values for State B", fontsize = 15)
    axs[1].legend(loc = 'lower right', fontsize = 15)
    axs[1].set_title("Values vs Iteration for State B for $\gamma$ = " + str(gamma), fontsize = 20)

    axs[2].plot(iteration_value_dict, value['C'], color = 'green', linestyle = 'dashed', marker = '>', label = 'VI')
    axs[2].plot(iteration_gs_value_dict, gs_value['C'], color = 'green', linestyle = 'solid', marker = '>', label = 'Gauss-Seidel VI')
    axs[2].set_xlabel("Iteration", fontsize = 15)
    axs[2].set_ylabel("Values for State C", fontsize = 15)
    axs[2].legend(loc = 'lower right', fontsize = 15)
    axs[2].set_title("Values vs Iteration for State C for $\gamma$ = " + str(gamma), fontsize = 20)

    fig = plt.gcf()
    fig.set_size_inches(40, 10)

plotter(0.1)
plotter(0.4)
plotter(0.5)
plotter(0.7)
plotter(0.9)

!DATASET_PATH=$AICROWD_DATASET_PATH aicrowd notebook submit -c iit-m-rl-assignment-2-taxi -a assets