
from .utils import format_snippet

def monte_carlo_code():
    code = '''
import numpy as np
import gymnasium as gym
import matplotlib.pyplot as plt

env = gym.make('FrozenLake-v1', is_slippery = 'False')

def monte_carlo(env, policy, episodes = 10000, df = 0.99):
  V = np.zeros(env.observation_space.n)
  returns = {s:[] for s in range(env.observation_space.n)}
  V_hist = []

  for ep in range(episodes):
    episode = []
    state, _ = env.reset()
    done = False

    while not done:
      action = policy[state]
      next_state, reward, terminated, truncated, _ = env.step(action)
      done = terminated or truncated
      episode.append((state, reward))
      state = next_state

    G = 0
    visited_states = set()
    for s, r in reversed(episode):
      G = df * G + r
      if s not in visited_states:
        returns[s].append(G)
        V[s] = np.mean(returns[s])
        visited_states.add(s)

    V_hist.append(V.copy())
  return V, V_hist

np.random.seed(42)

policy = {s: np.random.choice([0, 1, 2, 3]) for s in range(env.observation_space.n)}
V_MC, V_MC_hist = monte_carlo(env, policy)
print("Monte Carlo Value: ")
print(V_MC)

def convergence(V_track, title):
  plt.figure(figsize = (8, 5))
  for s in range(env.observation_space.n):
    values = [v[s] for v in V_track]
    plt.plot(values, label = f"State {s}")

  plt.title(title)
  plt.xlabel("Episodes")
  plt.ylabel("Value")
  plt.legend()
  plt.grid(True)
  plt.show()

convergence(V_MC_hist, "Monte Carlo Value Convergence")
'''
    return format_snippet(code)

def td_code():
    code = '''
import numpy as np
import gymnasium as gym
import matplotlib.pyplot as plt

env = gym.make('FrozenLake-v1', is_slippery = 'False')

def temp_diff(env, policy, episodes = 10000, lr = 0.05, df = 0.99):
  V = np.zeros(env.observation_space.n)
  V_hist = []

  for ep in range(episodes):
    episode = []
    state, _ = env.reset()
    done = False

    while not done:
      action  = policy[state]
      next_state, reward, terminated, truncated , _ = env.step(action)
      done = terminated or truncated
      V[state] = V[state] + lr * (reward + df * V[next_state] - V[state])
      state = next_state
    V_hist.append(V.copy())
  return V, V_hist

np.random.seed(42)

policy = {s: np.random.choice([0, 1, 2, 3]) for s in range(env.observation_space.n)}
V_TD, V_TD_hist = temp_diff(env, policy)

print("Temporal Difference Value: ")
print(V_TD)

def convergence(V_track, title):
  plt.figure(figsize = (8, 5))
  for s in range(env.observation_space.n):
    values = [v[s] for v in V_track]
    plt.plot(values, label = f"State {s}")

  plt.title(title)
  plt.xlabel("Episodes")
  plt.ylabel("Value")
  plt.legend()
  plt.grid(True)
  plt.show()

convergence(V_TD_hist, "Temporal Difference Value Convergence")
'''
    return format_snippet(code)

def complete_policy_iteration_code():
    code = '''
import gymnasium as gym
import numpy as np
import matplotlib.pyplot as plt

env = gym.make('FrozenLake-v1', is_slippery=True, render_mode='ansi')
P = env.unwrapped.P
env.reset()

nS = env.observation_space.n
nA = env.action_space.n

print(f'Action space: {nA}')
print(f'Observation space: {nS}')
print(f'Reward range: {env.unwrapped.reward_range}')

print(env.render())
print(P)

random_policy=np.ones([env.observation_space.n, env.action_space.n])/env.action_space.n
gamma=0.9

def policy_evaluation(env, policy, V, gamma=1.0, theta=1e-8):
    V_history = [V.copy()]
    while True:
        delta = 0
        for s in range(env.observation_space.n):
            v = 0
            for a, action_prob in enumerate(policy[s]):
                for prob, next_state, reward, done in P[s][a]:
                    v += action_prob * prob * (reward + gamma * V[next_state])
            delta = max(delta, abs(v - V[s]))
            V[s] = v
        V_history.append(V.copy())

        if delta < theta:
            break

    return V, V_history

def q_from_v(env, V, s, gamma=1):
  q = np.zeros(env.action_space.n)
  for a in range(env.action_space.n):
    for prob, next_state, reward, done in env.unwrapped.P[s][a]:
      q[a] += prob * (reward + gamma * V[next_state])
  return q

def plot_convergence(V_history):
  iterations = len(V_history) - 1
  max_diffs = []
  for i in range(iterations):
    diff = np.max(np.abs(V_history[i+1] - V_history[i]))
    max_diffs.append(diff)

  plt.figure()
  plt.plot(range(1, iterations + 1), max_diffs)
  plt.xlabel("Iteration")
  plt.ylabel("Max Change in V")
  plt.title("Policy Evaluation Convergence")
  plt.show()

def policy_iteration(env, gamma=1.0, theta=1e-8):
    nS = env.observation_space.n
    nA = env.action_space.n

    policy = np.ones([nS, nA]) / nA
    V = np.zeros(nS)

    iteration = 0
    policy_stable = False
    V_history = []

    while not policy_stable:
        iteration += 1

        V, current_V_history = policy_evaluation(env, policy, V, gamma, theta)
        if not V_history:
            V_history.extend(current_V_history)
        else:
             V_history.extend(current_V_history[1:])


        policy_stable = True
        for s in range(nS):
            old_action = np.argmax(policy[s])

            Q = q_from_v(env, V, s, gamma)
            best_action = np.argmax(Q)

            new_policy_s = np.eye(nA)[best_action]

            if not np.array_equal(policy[s], new_policy_s):
                policy_stable = False
                policy[s] = new_policy_s

    return policy, V, iteration, V_history

def plot(V, policy, discount_factor=1.0, draw_vals=True):
  nrow = env.unwrapped.nrow
  ncol = env.unwrapped.ncol
  nA = env.action_space.n
  arrow_symbols = {0: '←', 1: '↓', 2: '→', 3: '↑'}
  grid = np.reshape(V, (nrow, ncol))
  plt.figure(figsize=(6, 6))
  plt.imshow(grid, cmap='cool', interpolation='none')
  for s in range(nrow * ncol):
    row, col = divmod(s, ncol)
    best_action = np.argmax(policy[s])

    if draw_vals:
      plt.text(col, row, f'{V[s]:.2f}', ha='center', va='center', color='white', fontsize=10)
    else:
      plt.text(col, row, arrow_symbols[best_action], ha='center', va='center', color='white',
      fontsize=14)

  plt.title("Value Function" if draw_vals else "Optimal Policy") # Added title here
  plt.axis('off')
  plt.show()

V_random = np.random.rand(env.observation_space.n)
V_random.reshape(4, 4)

old_policy = random_policy
new_policy, latest_V, iterations, _ = policy_iteration(env, gamma=0.9)

print(f'Convergenced in {iterations} steps.')

plot(latest_V, new_policy, 1.0, draw_vals=True) #Value Function
plot(V_random, new_policy, 1.0, draw_vals=False) #Optimal Policy
'''
    return format_snippet(code)   

def complete_value_iteration_code():
    code = '''
import gymnasium as gym
import numpy as np
import matplotlib.pyplot as plt

env = gym.make('FrozenLake-v1', is_slippery=True, render_mode='ansi')
P = env.unwrapped.P
env.reset()

nS = env.observation_space.n
nA = env.action_space.n

print(f'Action space: {nA}')
print(f'Observation space: {nS}')
print(f'Reward range: {env.unwrapped.reward_range}')

print(env.render())
print(P)

random_policy=np.ones([env.observation_space.n, env.action_space.n])/env.action_space.n
gamma=0.9


def q_from_v(env, V, s, gamma=1):
  q = np.zeros(env.action_space.n)
  for a in range(env.action_space.n):
    for prob, next_state, reward, done in env.unwrapped.P[s][a]:
      q[a] += prob * (reward + gamma * V[next_state])
  return q

def plot_convergence(V_history):
  iterations = len(V_history) - 1
  max_diffs = []
  for i in range(iterations):
    diff = np.max(np.abs(V_history[i+1] - V_history[i]))
    max_diffs.append(diff)

  plt.figure()
  plt.plot(range(1, iterations + 1), max_diffs)
  plt.xlabel("Iteration")
  plt.ylabel("Max Change in V")
  plt.title("Policy Evaluation Convergence")
  plt.show()

def value_iteration(env, gamma=1.0, theta=1e-8):
    nS = env.observation_space.n
    nA = env.action_space.n
    
    V = np.zeros(nS)
    V_history = [V.copy()]  # track convergence
    
    while True:
        delta = 0
        
        for s in range(nS):
            # Compute Q(s,a) for each action
            Q = np.zeros(nA)
            for a in range(nA):
                for prob, next_state, reward, done in env.unwrapped.P[s][a]:
                    Q[a] += prob * (reward + gamma * V[next_state])
            
            # Best action value
            v_new = np.max(Q)
            delta = max(delta, abs(v_new - V[s]))
            V[s] = v_new
        
        V_history.append(V.copy())
        
        if delta < theta:
            break
    
    # Extract greedy policy
    policy = np.zeros([nS, nA])
    for s in range(nS):
        Q = q_from_v(env, V, s, gamma)
        best_a = np.argmax(Q)
        policy[s] = np.eye(nA)[best_a]
        
    return policy, V, V_history


def plot(V, policy, discount_factor=1.0, draw_vals=True):
  nrow = env.unwrapped.nrow
  ncol = env.unwrapped.ncol
  nA = env.action_space.n
  arrow_symbols = {0: '←', 1: '↓', 2: '→', 3: '↑'}
  grid = np.reshape(V, (nrow, ncol))
  plt.figure(figsize=(6, 6))
  plt.imshow(grid, cmap='cool', interpolation='none')
  for s in range(nrow * ncol):
    row, col = divmod(s, ncol)
    best_action = np.argmax(policy[s])

    if draw_vals:
      plt.text(col, row, f'{V[s]:.2f}', ha='center', va='center', color='white', fontsize=10)
    else:
      plt.text(col, row, arrow_symbols[best_action], ha='center', va='center', color='white',
      fontsize=14)

  plt.title("Value Function" if draw_vals else "Optimal Policy") # Added title here
  plt.axis('off')
  plt.show()

V_random = np.random.rand(env.observation_space.n)
V_random.reshape(4, 4)

old_policy = random_policy
new_policy, latest_V, iterations, _ = policy_iteration(env, gamma=0.9)

print(f'Convergenced in {iterations} steps.')

plot(latest_V, new_policy, 1.0, draw_vals=True) #Value Function
plot(V_random, new_policy, 1.0, draw_vals=False) #Optimal Policy
'''
    return format_snippet(code) 