!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
    flag = False
    while not flag:
    #policy evaluation step
      while True:
        delta =0
        values_old =values.copy()
        for s in states:
          a=policy[s]
          J= [ i*gamma for i in list(values_old.values())]
          values[s] = sum(taxienv.ride_probabilities(s,a)* ( taxienv.ride_rewards(s,a) + J ))    
          delta = max(delta, abs(values[s] - values_old[s]))
        if delta < 1e-8:
          break 
    #policy improvement step
      flag =True
      for s in states:
        actions = taxienv.possible_actions(s)
        rewards = {a:0 for a in actions}
        for a in actions:
          J= [ i*gamma for i in list(values.values()) ]
          rewards[a] = sum(taxienv.ride_probabilities(s, a)* ( taxienv.ride_rewards(s, a) + J ))

        action=max(rewards, key=rewards.get)
        
        if policy[s] != action:
          policy[s]=action
          flag= False
          



    # 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

    
    # 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

# 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
    flag = False
    while not flag:

    #policy evaluation step
      for i in range(m):
        values_old =values.copy()
        for s in states:
          a=policy[s]
          J= [ i*gamma for i in list(values_old.values())]
          values[s] = sum(taxienv.ride_probabilities(s,a)* ( taxienv.ride_rewards(s,a) + J ))    


    #policy improvement step
      flag =True
      for s in states:
        actions = taxienv.possible_actions(s)
        rewards = {a:0 for a in actions}
        for a in actions:
          J= [ i*gamma for i in list(values.values()) ]
          reward = sum(taxienv.ride_probabilities(s, a)* ( taxienv.ride_rewards(s, a) + J ))
          rewards[a]=reward
 

        action=max(rewards, key=rewards.get)
        if policy[s] != action:
          policy[s]=action
          flag= False
          

    # 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

    
    # 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)

# 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
    iter=0
    delta_values=[]
    while True:
      iter=iter+1
      delta=0
      values_old=values.copy()

      for s in states:
        actions=taxienv.possible_actions(s)
        rewards= {a:0 for a in actions}

        for a in actions:
                    J= [ i*gamma for i in list(values_old.values()) ]
                    rewards[a] = sum(taxienv.ride_probabilities(s, a)* ( taxienv.ride_rewards(s, a) + J ))
        action=max(rewards,key=rewards.get)
        values[s]=rewards[action]
        policy[s]=action

        delta=max(delta,abs(values[s]-values_old[s]))
      delta_values.append(delta)
      if delta < 1e-8:
        break



    # 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


    # Put your extra information needed for plots etc in this dictionary
    extra_info = {"iter":iter}

    ## 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

    iter=0
    delta_values=[]
    while True:
      iter=iter+1
      delta=0
      values_old=values.copy()
      for s in states:
        actions=taxienv.possible_actions(s)
        rewards= {a:0 for a in actions}
        for a in actions:
          J= [ i*gamma for i in list(values.values()) ]
          rewards[a] = sum(taxienv.ride_probabilities(s, a)* ( taxienv.ride_rewards(s, a) + J ))

        action=max(rewards,key=rewards.get)
        values[s]=rewards[action]
        policy[s]=action

        delta=max(delta,abs(values[s]-values_old[s]))
        delta_values.append(delta)

      if delta <1e-8:
        break
    
    
    # Put your extra information needed for plots etc in this dictionary
    extra_info = {"iter":iter}

    ## 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))

gamma = 0.9
results, extra_info = policy_iteration(env, gamma)
reward = results["Expected Reward"]
policy = results["Policy"]
print("Optimal Policy",policy)
print("Expected Rewards",reward)

results, extra_info = run_policy_iteration(env)
import matplotlib.pyplot as plt
plt.figure(figsize=(10, 5), dpi=100)
plt.title("Expected Reward Vs Gamma Values Plot")
plt.xlabel('Gamma values')
plt.ylabel('Expected rewards')
plt.xticks(list(results.keys()))

cities = ['A','B','C']
for city in cities:
  plt.plot(list(results.keys()),[i['Expected Reward'][city] for i in results.values()],marker='o',linewidth=0.5, markersize=3)

plt.legend(["City_A","City_B","City_C"])
plt.show()

##printing all rewards and policy wrt to gamma
print("Expected Reward")
for gamma in results.keys():
  rewards = results[gamma]['Expected Reward']
  print("Gamma",gamma,rewards)

print("\nPolicy")
for gamma in results.keys():
  policy = results[gamma]['Policy']
  print("Gamma",gamma,policy)

##simply just running modified PI with given values
gamma = 0.9
m = 5
results, extra_info = modified_policy_iteration(env, gamma, m)
rewards = results["Expected Reward"]
policy = results["Policy"]
print("Optimal Policy  ",policy)
print("Expected Rewards",rewards)

gamma = 0.9
results, extra_info = value_iteration(env, gamma)
rewards = results["Expected Reward"]
policy = results["Policy"]
print("Optimal Policy",policy)
print("Expected Rewards",rewards)

gamma = 0.9
results, extra_info = gauss_seidel_value_iteration(env, gamma)
rewards = results["Expected Reward"]
policy = results["Policy"]
print("Optimal Policy",policy)
print("Expected Rewards",rewards)

results, extra_info = run_policy_iteration(env)
import matplotlib.pyplot as plt
plt.figure(figsize=(10, 5), dpi=100)
plt.title("Expected Reward Vs Gamma Values Plot")
plt.xlabel('Gamma values')
plt.ylabel('Expected rewards')
plt.xticks(list(results.keys()))

cities = ['A','B','C']
for city in cities:
  plt.plot(list(results.keys()),[i['Expected Reward'][city] for i in results.values()],marker='o',linewidth=0.5, markersize=3)

plt.legend(["City_A","City_B","City_C"])
plt.show()

""" our policy changes as follows

Gamma 0.05 {'A': '1', 'B': '1', 'C': '1'}
Gamma 0.1 {'A': '1', 'B': '1', 'C': '1'}
Gamma 0.15 {'A': '1', 'B': '2', 'C': '1'}
Gamma 0.2 {'A': '1', 'B': '2', 'C': '1'}
Gamma 0.25 {'A': '1', 'B': '2', 'C': '1'}
Gamma 0.3 {'A': '1', 'B': '2', 'C': '1'}
Gamma 0.35 {'A': '1', 'B': '2', 'C': '1'}
Gamma 0.4 {'A': '1', 'B': '2', 'C': '1'}
Gamma 0.45 {'A': '1', 'B': '2', 'C': '1'}
Gamma 0.5 {'A': '1', 'B': '2', 'C': '1'}
Gamma 0.55 {'A': '1', 'B': '2', 'C': '2'}
Gamma 0.6 {'A': '1', 'B': '2', 'C': '2'}
Gamma 0.65 {'A': '1', 'B': '2', 'C': '2'}
Gamma 0.7 {'A': '1', 'B': '2', 'C': '2'}
Gamma 0.75 {'A': '1', 'B': '2', 'C': '2'}
Gamma 0.8 {'A': '2', 'B': '2', 'C': '2'}
Gamma 0.85 {'A': '2', 'B': '2', 'C': '2'}
Gamma 0.9 {'A': '2', 'B': '2', 'C': '2'}
Gamma 0.95 {'A': '2', 'B': '2', 'C': '2'}

we can easily see how it changes from "1" ie cruise the streets to look for a passenger to "2" ie  go to nearest taxi line and wait

this can be reasoned as with low gamma driver doesn't have alot of trips he can make and therefore chooses "1" which is the short term 
maximizer , while "2" is the long term maximizer which he can choose when the car has low change of breaking down"""

gamma = 0.9
m = 5
results5, extra_info = modified_policy_iteration(env, gamma, m)
rewards5 = results5["Expected Reward"]
policy5 = results5["Policy"]
print("Optimal Policy" ,policy5)
print("Expected Rewards",rewards5)

gamma = 0.9
m = 10
results10, extra_info = modified_policy_iteration(env, gamma, m)
rewards10 = results10["Expected Reward"]
policy10 = results10["Policy"]

print("\nOptimal Policy",policy10)
print("Expected Rewards ",rewards10)

import matplotlib.pyplot as plt

resultsvi, extra_infovi = run_value_iteration(env)
resultsgs, extra_infogs = run_gauss_seidel_value_iteration(env)
plt.figure(figsize=(10, 5), dpi=100)
plt.xlabel('Gamma')
plt.ylabel('Iterations')
plt.xticks(list(extra_infovi.keys()))
VI_iter= [extra_infovi[gamma]['iter'] for gamma in extra_infovi.keys()]
GS_iter=[extra_infogs[gamma]['iter'] for gamma in extra_infogs.keys()]
plt.plot(list(extra_infovi.keys()), VI_iter, color='cyan', marker='o', linewidth=0.5, markersize=5, label = "VI")
plt.plot(list(extra_infogs.keys()), GS_iter, color='red', marker='o', linewidth=0.5, markersize=5, label = "GS")

plt.legend()
plt.show()

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