import random
from math import floor
import numpy as np
import pandas as pd
import seaborn as sns


def uniform(a: int, b: int) -> int:
    u = random.random()
    return a + floor((b - a + 1) * u)


if __name__ == "__main__":
    total_max_bias = 0
    avgs_max_bias = []
    total_double_sampling = 0
    avgs_double_sampling = []

    for e in range(1, 100):
        samples = np.array([uniform(-3, 3) for _ in range(5)])
        max_sample = max(samples)
        total_max_bias += max_sample
        avgs_max_bias.append(total_max_bias / e)

        samples2 = np.array([uniform(-3, 3) for _ in range(5)])
        total_double_sampling += samples2[np.argmax(samples)]
        avgs_double_sampling.append(total_double_sampling / e)

    df = pd.DataFrame({'Max of Samples': avgs_max_bias, 'Double Samples': avgs_double_sampling})
    import matplotlib.pyplot as plt
    sns.lineplot(data=df)
    plt.show()

from torch.utils.data import Dataset

class TSPDataset(Dataset):
	"each data item of form (input, input_len, output_in, output_out, output_len)"
	data: List[Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray, np.ndarray]]
	
	def __len__(self):
		return len(self.data)

	def __getitem__(self, index):
		input, input_len, output_in, output_out, output_len = self.data[index]
		return input, input_len, output_in, output_out, output_len

class RNNEncoder(nn.Module):
	rnn: Union[nn.LSTM, nn.GRU, nn.RNN]

	def __init__(self, rnn_type: str, bidirectional: bool, num_layers: int, input_size: int, hidden_size: int, dropout: float):
		super(RNNEncoder, self).__init__()
		if bidirectional:
			assert hidden_size % 2 == 0
			hidden_size = hidden_size // 2
		self.rnn = rnn_init(rnn_type, input_size=input_size, hidden_size=hidden_size, bidirectional=bidirectional,num_layers=num_layers, dropout=dropout)

	def forward(self, src: Tensor, src_lengths: Tensor, hidden: Tensor = None) -> Tuple[Tensor, Tensor]:
		lengths = src_lengths.view(-1).tolist()
		packed_src = pack_padded_sequence(src, lengths)
		memory_bank, hidden_final = self.rnn(packed_src, hidden)
		memory_bank = pad_packed_sequence(memory_bank)[0]
		return memory_bank, hidden_final

class Attention(nn.Module):
	linear_out: nn.Linear

	def __init__(self, dim: int):
		super(Attention, self).__init__()
		self.linear_out = nn.Linear(dim * 2, dim, bias=False)

	def score(self, src: Tensor, target: Tensor) -> Tensor:
		batch_size, src_len, dim = src.size()
		_, target_len, _ = target.size()
		target_ = target
		src_ = src.transpose(1, 2)
		return torch.bmm(target_, src_)

	def forward(self, src: Tensor, target: Tensor, src_lengths: Tensor) -> Tuple[Tensor, Tensor]:
		assert target.dim() == 3

		batch_size, src_len, dim = src.size()
		_, target_len, _ = target.size()

		align_score = self.score(src, target)

		mask = sequence_mask(src_lengths)
		# (batch_size, max_len) -> (batch_size, 1, max_len)
		mask = mask.unsqueeze(1)
		align_score.data.masked_fill_(~mask, -float('inf'))
		align_score = F.softmax(align_score, -1)

		c = torch.bmm(align_score, src)

		concat_c = torch.cat([c, target], -1)
		attn_h = self.linear_out(concat_c)

		return attn_h, align_score

State = Tuple[int, int, int, int]
Action = int

class CartPoleAbstractAgent(metaclass=abc.ABCMeta):
    def __init__(self, buckets=(1, 2, 6, 12), discount=0.98, lr_min=0.1, epsilon_min=0.1):
        self.env = gym.make('CartPole-v0')

        env = self.env
        # [position, velocity, angle, angular velocity]
        self.dims_config = [(env.observation_space.low[0], env.observation_space.high[0], 1),
                            (-0.5, 0.5, 1),
                            (env.observation_space.low[2], env.observation_space.high[2], 6),
                            (-math.radians(50) / 1., math.radians(50) / 1., 12)]
        self.q = np.zeros(buckets + (self.env.action_space.n,))
        self.pi = np.zeros_like(self.q)
        self.pi[:] = 1.0 / env.action_space.n

    def to_bin_idx(self, val: float, lower: float, upper: float, bucket_num: int) -> int:
        percent = (val + abs(lower)) / (upper - lower)
        return min(bucket_num - 1, max(0, int(round((bucket_num - 1) * percent))))

    def discretize(self, obs: np.ndarray) -> State:
        discrete_states = tuple([self.to_bin_idx(obs[d], *self.dims_config[d]) for d in range(len(obs))])
        return discrete_states

def train(self, num_episodes=2000):
    for e in range(num_episodes):
        print(e)
        s: State = self.discretize(self.env.reset())

        self.adjust_learning_rate(e)
        self.adjust_epsilon(e)
        done = False

        while not done:
            action: Action = self.choose_action(s)
            obs, reward, done, _ = self.env.step(action)
            s_next: State = self.discretize(obs)
            a_next = self.choose_action(s_next)
            self.update_q(s, action, reward, s_next, a_next)
            s = s_next

def choose_action(self, state) -> Action:
    if np.random.random() < self.epsilon:
        return self.env.action_space.sample()
    else:
        return np.argmax(self.q[state])

class CartPoleAbstractAgent(metaclass=abc.ABCMeta):
    @abc.abstractmethod
    def update_q(self, s: State, a: Action, r, s_next: State, a_next: Action):
        pass

class SarsaAgent(CartPoleAbstractAgent):

    def update_q(self, s: State, a: Action, r, s_next: State, a_next: Action):
        self.q[s][a] += self.lr * (r + self.discount * (self.q[s_next][a_next]) - self.q[s][a])

class QLearningAgent(CartPoleAbstractAgent):

    def update_q(self, s: State, a: Action, r, s_next: State, a_next: Action):
        self.q[s][a] += self.lr * (r + self.discount * np.max(self.q[s_next]) - self.q[s][a])

class ExpectedSarsaAgent(CartPoleAbstractAgent):

    def update_q(self, s: State, a: Action, r, s_next: State, a_next: Action):
        self.q[s][a] = self.q[s][a] + self.lr * (r + self.discount * np.dot(self.pi[s_next], self.q[s_next]) - self.q[s][a])
        # update pi[s]
        best_a = np.random.choice(np.where(self.q[s] == max(self.q[s]))[0])
        n_actions = self.env.action_space.n
        self.pi[s][:] = self.epsilon / n_actions
        self.pi[s][best_a] = 1 - (n_actions - 1) * (self.epsilon / n_actions)

# AC
# Runtime: 32 ms, faster than 54.28% of Python3 online submissions for Pow(x, n).
# Memory Usage: 14.2 MB, less than 5.04% of Python3 online submissions for Pow(x, n).

class Solution:
    def myPow(self, x: float, n: int) -> float:
        ret = 1.0
        i = abs(n)
        while i != 0:
            if i & 1:
                ret *= x
            x *= x
            i = i >> 1
        return 1.0 / ret if n < 0 else ret

# AC
# Runtime: 256 ms, faster than 26.21% of Python3 online submissions for Fibonacci Number.
# Memory Usage: 29.4 MB, less than 5.25% of Python3 online submissions for Fibonacci Number.

class Solution:

    def fib(self, N: int) -> int:
        if N <= 1:
            return N

        import numpy as np
        F = np.array([[1, 1], [1, 0]])

        N -= 1
        powerDouble = F
        powers = np.array([[1, 0], [0, 1]])
        while N > 0:
            if N % 2 == 1:
                powers = np.matmul(powers, powerDouble)
            powerDouble = np.matmul(powerDouble, powerDouble)
            N = N // 2

        return powers[0][0]

# AC
# Runtime: 116 ms, faster than 26.25% of Python3 online submissions for Fibonacci Number.
# Memory Usage: 29.2 MB, less than 5.25% of Python3 online submissions for Fibonacci Number.

class Solution:

    def fib(self, N: int) -> int:
        if N <= 1:
            return N

        from numpy.linalg import matrix_power
        import numpy as np
        F = np.array([[1, 1], [1, 0]])
        F = matrix_power(F, N - 1)

        return F[0][0]

# AC
# Runtime: 36 ms, faster than 98.88% of Python3 online submissions for Number of Ways to Paint N × 3 Grid.
# Memory Usage: 13.9 MB, less than 58.66% of Python3 online submissions for Number of Ways to Paint N × 3 Grid.

class Solution:
    def numOfWays(self, n: int) -> int:
        MOD = 10 ** 9 + 7
        dp2, dp3 = 6, 6
        n -= 1
        while n > 0:
            dp2, dp3 = (dp2 * 3 + dp3 * 2) % MOD, (dp2 * 2 + dp3 * 2) % MOD
            n -= 1
        return (dp2 + dp3) % MOD

# AC
# Runtime: 88 ms, faster than 39.07% of Python3 online submissions for Number of Ways to Paint N × 3 Grid.
# Memory Usage: 30.2 MB, less than 11.59% of Python3 online submissions for Number of Ways to Paint N × 3 Grid.

class Solution:
    def numOfWays(self, n: int) -> int:
        import numpy as np

        MOD = int(1e9 + 7)
        ret = np.array([[6, 6]])
        m = np.array([[3, 2], [2, 2]])

        n -= 1
        while n > 0:
            if n % 2 == 1:
                ret = np.matmul(ret, m) % MOD
            m = np.matmul(m, m) % MOD
            n = n // 2
        return int((ret[0][0] + ret[0][1]) % MOD)

State = Tuple[int, int, bool]
Action = bool
Reward = float
Actions = np.ndarray
ActionValue = Dict[State, Actions]
Policy = Dict[State, Actions]

policy: Policy
A: ActionValue = policy[state]
action = np.random.choice([0, 1], p=A/sum(A))

def mc_control_exploring_starts(env: BlackjackEnv, num_episodes, discount_factor=1.0) \
        -> Tuple[ActionValue, Policy]:
    states = list(product(range(10, 22), range(1, 11), (True, False)))
    policy = {s: np.ones(env.action_space.n) * 1.0 / env.action_space.n for s in states}
    Q = defaultdict(lambda: np.zeros(env.action_space.n))
    returns_sum = defaultdict(float)
    returns_count = defaultdict(float)

    for episode_i in range(1, num_episodes + 1):
        s0 = random.choice(states)
        reset_env_with_s0(env, s0)
        episode_history = gen_custom_s0_stochastic_episode(policy, env, s0)

        G = 0
        for t in range(len(episode_history) - 1, -1, -1):
            s, a, r = episode_history[t]
            G = discount_factor * G + r
            if not any(s_a_r[0] == s and s_a_r[1] == a for s_a_r in episode_history[0: t]):
                returns_sum[s, a] += G
                returns_count[s, a] += 1.0
                Q[s][a] = returns_sum[s, a] / returns_count[s, a]
                best_a = np.argmax(Q[s])
                policy[s][best_a] = 1.0
                policy[s][1-best_a] = 0.0

    return Q, policy

def reset_env_with_s0(env: BlackjackEnv, s0: State) -> BlackjackEnv:
    env.reset()
    player_sum = s0[0]
    oppo_sum = s0[1]
    has_usable = s0[2]

    env.dealer[0] = oppo_sum
    if has_usable:
        env.player[0] = 1
        env.player[1] = player_sum - 11
    else:
        if player_sum > 11:
            env.player[0] = 10
            env.player[1] = player_sum - 10
        else:
            env.player[0] = 2
            env.player[1] = player_sum - 2
    return env

def mc_control_epsilon_greedy(env: BlackjackEnv, num_episodes, discount_factor=1.0, epsilon=0.1) \
        -> Tuple[ActionValue, Policy]:
    returns_sum = defaultdict(float)
    returns_count = defaultdict(float)

    states = list(product(range(10, 22), range(1, 11), (True, False)))
    policy = {s: np.ones(env.action_space.n) * 1.0 / env.action_space.n for s in states}
    Q = defaultdict(lambda: np.zeros(env.action_space.n))

    def update_epsilon_greedy_policy(policy: Policy, Q: ActionValue, s: State):
        policy[s] = np.ones(env.action_space.n, dtype=float) * epsilon / env.action_space.n
        best_action = np.argmax(Q[s])
        policy[s][best_action] += (1.0 - epsilon)

    for episode_i in range(1, num_episodes + 1):
        episode_history = gen_stochastic_episode(policy, env)

        G = 0
        for t in range(len(episode_history) - 1, -1, -1):
            s, a, r = episode_history[t]
            G = discount_factor * G + r
            if not any(s_a_r[0] == s and s_a_r[1] == a for s_a_r in episode_history[0: t]):
                returns_sum[s, a] += G
                returns_count[s, a] += 1.0
                Q[s][a] = returns_sum[s, a] / returns_count[s, a]
                update_epsilon_greedy_policy(policy, Q, s)

    return Q, policy

from gym.envs.toy_text import BlackjackEnv

State = Tuple[int, int, bool]
Action = bool
Reward = float
StateValue = Dict[State, float]
DeterministicPolicy = Callable[[State], Action]

class BlackjackEnv(gym.Env):

    def step(self, action):
        assert self.action_space.contains(action)
        if action:  # hit: add a card to players hand and return
            self.player.append(draw_card(self.np_random))
            if is_bust(self.player):
                done = True
                reward = -1.
            else:
                done = False
                reward = 0.
        else:  # stick: play out the dealers hand, and score
            done = True
            while sum_hand(self.dealer) < 17:
                self.dealer.append(draw_card(self.np_random))
            reward = cmp(score(self.player), score(self.dealer))
            if self.natural and is_natural(self.player) and reward == 1.:
                reward = 1.5
        return self._get_obs(), reward, done, {}

    def _get_obs(self):
        return (sum_hand(self.player), self.dealer[0], usable_ace(self.player))

def gen_episode_data(policy: DeterministicPolicy, env: BlackjackEnv) -> List[Tuple[State, Action, Reward]]:
    episode_history = []
    state = env.reset()
    done = False
    while not done:
        action = policy(state)
        next_state, reward, done, _ = env.step(action)
        episode_history.append((state, action, reward))
        state = next_state
    return episode_history

def fixed_policy(observation):
    """
    sticks if the player score is >= 20 and hits otherwise.
    """
    score, dealer_score, usable_ace = observation
    return 0 if score >= 20 else 1

def mc_prediction_first_visit(policy: DeterministicPolicy, env: BlackjackEnv,
                              num_episodes, discount_factor=1.0) -> StateValue:
    returns_sum = defaultdict(float)
    returns_count = defaultdict(float)

    for episode_i in range(1, num_episodes + 1):
        episode_history = gen_episode_data(policy, env)

        G = 0
        for t in range(len(episode_history) - 1, -1, -1):
            s, a, r = episode_history[t]
            G = discount_factor * G + r
            if not any(s_a_r[0] == s for s_a_r in episode_history[0: t]):
                returns_sum[s] += G
                returns_count[s] += 1.0

    V = defaultdict(float)
    V.update({s: returns_sum[s] / returns_count[s] for s in returns_sum.keys()})
    return V

def mc_prediction_every_visit(policy: DeterministicPolicy, env: BlackjackEnv,
                              num_episodes, discount_factor=1.0) -> StateValue:
    returns_sum = defaultdict(float)
    returns_count = defaultdict(float)

    for episode_i in range(1, num_episodes + 1):
        episode_history = gen_episode_data(policy, env)

        G = 0
        for t in range(len(episode_history) - 1, -1, -1):
            s, a, r = episode_history[t]
            G = discount_factor * G + r
            returns_sum[s] += G
            returns_count[s] += 1.0

    V = defaultdict(float)
    V.update({s: returns_sum[s] / returns_count[s] for s in returns_sum.keys()})
    return V

class TSPSolver:
    g: Graph

    def __init__(self, g: Graph):
        self.g = g

    def solve(self) -> int:
        """
        :param v:
        :param state:
        :return: -1 means INF
        """
        N = self.g.v_num
        dp = [[INT_INF for c in range(N)] for r in range(1 << N)]

        dp[(1 << N) - 1][0] = 0

        for state in range((1 << N) - 2, -1, -1):
            for v in range(N):
                for u in range(N):
                    if ((state >> u) & 1) == 0:
                        if dp[state | 1 << u][u] != INT_INF and self.g.edges[v][u] != INT_INF:
                            if dp[state][v] == INT_INF:
                                dp[state][v] = dp[state | 1 << u][u] + self.g.edges[v][u]
                            else:
                                dp[state][v] = min(dp[state][v], dp[state | 1 << u][u] + self.g.edges[v][u])
        return dp[0][0]

import matplotlib.pyplot as plt
points='0.607122 0.664447 0.953593 0.021519 0.757626 0.921024 0.586376 0.433565 0.786837 0.052959 0.016088 0.581436 0.496714 0.633571 0.227777 0.971433 0.665490 0.074331 0.383556 0.104392'
float_list = list(map(lambda x: float(x), points.split(' ')))

x,y = [],[]
for idx, p in enumerate(float_list):
  if idx % 2 == 0:
    x.append(p)
  else:
    y.append(p)

for i in range(0, len(x)):
  for j in range(0, len(x)):
    if i == j:
      continue
    plt.plot((x[i],x[j]),(y[i],y[j]))

plt.show()

tour_str = '1 3 8 6 10 9 5 2 4 7 1'
tour = list(map(lambda x: int(x), tour_str.split(' ')))

for i in range(0, len(tour)-1):
  p1 = tour[i] - 1
  p2 = tour[i + 1] - 1
  plt.plot((x[p1],x[p2]),(y[p1],y[p2]))
plt.show()

g: Graph = Graph(N)
for v in range(N):
	for u in range(N):
		diff_x = coordinates[v][0] - coordinates[u][0]
		diff_y = coordinates[v][1] - coordinates[u][1]
		dist: float = math.sqrt(diff_x * diff_x + diff_y * diff_y)
		g.setDist(u, v, dist)
		g.setDist(v, u, dist)

ret: float = FLOAT_INF
u_min: int = -1
	for u in range(self.g.v_num):
		if (state & (1 << u)) == 0:
			s: float = self._recurse(u, state | 1 << u)
				if s + edges[v][u] < ret:
					ret = s + edges[v][u]
					u_min = u
	dp[state][v] = ret
	self.parent[state][v] = u_min

def _form_tour(self):
	self.tour = [0]
	bit = 0
	v = 0
	for _ in range(self.g.v_num - 1):
		v = self.parent[bit][v]
		self.tour.append(v)
		bit = bit | (1 << v)
	self.tour.append(0)

tsp: TSPSolver = TSPSolver(g)
tsp.solve()

output_dist: float = 0.0
output_tour = list(map(lambda x: int(x) - 1, output.split(' ')))
for v in range(1, len(output_tour)):
	pre_v = output_tour[v-1]
	curr_v = output_tour[v]
	diff_x = coordinates[pre_v][0] - coordinates[curr_v][0]
	diff_y = coordinates[pre_v][1] - coordinates[curr_v][1]
	dist: float = math.sqrt(diff_x * diff_x + diff_y * diff_y)
	output_dist += dist

	passed = abs(tsp.dist - output_dist) < 10e-5
	if passed:
		print(f'passed dist={tsp.tour}')
	else:
		print(f'Min Tour Distance = {output_dist}, Computed Tour Distance = {tsp.dist}, Expected Tour = {output_tour}, Result = {tsp.tour}')

[[ 0. -1. -2. -3.]
 [-1. -2. -3. -2.]
 [-2. -3. -2. -1.]
 [-3. -2. -1.  0.]]

def action_value(env: GridWorldEnv, state: State, V: StateValue, gamma=1.0) -> ActionValue:
    q = np.zeros(env.nA)
    for a in range(env.nA):
        for prob, next_state, reward, done in env.P[state][a]:
            q[a] += prob * (reward + gamma * V[next_state])
    return q

def policy_improvement(env: GridWorldEnv, policy: Policy, V: StateValue, gamma=1.0) -> bool:
    policy_stable = True

    for s in range(env.nS):
        old_action = np.argmax(policy[s])
        Q_s = action_value(env, s, V)
        best_action = np.argmax(Q_s)
        policy[s] = np.eye(env.nA)[best_action]

        if old_action != best_action:
            policy_stable = False
    return policy_stable


def policy_iteration(env: GridWorldEnv, policy: Policy, gamma=1.0) -> Tuple[Policy, StateValue]:
    iter = 0
    while True:
        V = policy_evaluate(policy, env, gamma)
        policy_stable = policy_improvement(env, policy, V)
        iter += 1

        if policy_stable:
            return policy, V

def value_iteration(env:GridWorldEnv, gamma=1.0, theta=0.0001) -> Tuple[Policy, StateValue]:
    V = np.zeros(env.nS)
    while True:
        delta = 0
        for s in range(env.nS):
            action_values = action_value(env, s, V, gamma=gamma)
            best_action_value = np.max(action_values)
            delta = max(delta, np.abs(best_action_value - V[s]))
            V[s] = best_action_value
        if delta < theta:
            break

    policy = np.zeros([env.nS, env.nA])
    for s in range(env.nS):
        action_values = action_value(env, s, V, gamma=gamma)
        best_action = np.argmax(action_values)
        policy[s, best_action] = 1.0

    return policy, V

def build_reverse_mapping(env:GridWorldEnv) -> Dict[State, Set[State]]:
    MAX_R, MAX_C = env.shape[0], env.shape[1]
    mapping = {s: set() for s in range(0, MAX_R * MAX_C)}
    action_delta = {Action.UP: (-1, 0), Action.DOWN: (1, 0), Action.LEFT: (0, -1), Action.RIGHT: (0, 1)}
    for s in range(0, MAX_R * MAX_C):
        r = s // MAX_R
        c = s % MAX_R
        for a in list(Action):
            neighbor_r = min(MAX_R - 1, max(0, r + action_delta[a][0]))
            neighbor_c = min(MAX_C - 1, max(0, c + action_delta[a][1]))
            s_ = neighbor_r * MAX_R + neighbor_c
            mapping[s_].add(s)
    return mapping

def value_iteration_async(env:GridWorldEnv, gamma=1.0, theta=0.0001) -> Tuple[Policy, StateValue]:
    mapping = build_reverse_mapping(env)

    V = np.zeros(env.nS)
    changed_state_set = set(s for s in range(env.nS))

    iter = 0
    while len(changed_state_set) > 0:
        changed_state_set_ = set()
        for s in changed_state_set:
            action_values = action_value(env, s, V, gamma=gamma)
            best_action_value = np.max(action_values)
            v_diff = np.abs(best_action_value - V[s])
            if v_diff > theta:
                changed_state_set_.update(mapping[s])
                V[s] = best_action_value
        changed_state_set = changed_state_set_
        iter += 1

    policy = np.zeros([env.nS, env.nA])
    for s in range(env.nS):
        action_values = action_value(env, s, V, gamma=gamma)
        best_action = np.argmax(action_values)
        policy[s, best_action] = 1.0

    return policy, V

# AC
# Runtime: 36 ms, faster than 91.78% of Python3 online submissions for Permutations.
# Memory Usage: 13.9 MB, less than 66.52% of Python3 online submissions for Permutations.

from itertools import permutations
from typing import List


class Solution:
    def permute(self, nums: List[int]) -> List[List[int]]:
        return [p for p in permutations(nums)]

# AC
# Runtime: 84 ms, faster than 95.43% of Python3 online submissions for Combinations.
# Memory Usage: 15.2 MB, less than 68.98% of Python3 online submissions for Combinations.
from itertools import combinations
from typing import List

class Solution:
    def combine(self, n: int, k: int) -> List[List[int]]:
        return [c for c in combinations(list(range(1, n + 1)), k)]

[lst for lst in product([1, 2, 3], ['a', 'b'])]
[(i, s) for i in [1, 2, 3] for s in ['a', 'b']]

# AC
# Runtime: 24 ms, faster than 94.50% of Python3 online submissions for Letter Combinations of a Phone Number.
# Memory Usage: 13.7 MB, less than 83.64% of Python3 online submissions for Letter Combinations of a Phone Number.

from itertools import product
from typing import List


class Solution:
    def letterCombinations(self, digits: str) -> List[str]:
        if digits == "":
            return []
        mapping = {'2':'abc', '3':'def', '4':'ghi', '5':'jkl', '6':'mno', '7':'pqrs', '8':'tuv', '9':'wxyz'}
        iter_dims = [mapping[i] for i in digits]

        result = []
        for lst in product(*iter_dims):
            result.append(''.join(lst))

        return result

# AC
# Runtime: 28 ms, faster than 85.56% of Python3 online submissions for Fibonacci Number.
# Memory Usage: 13.8 MB, less than 58.41% of Python3 online submissions for Fibonacci Number.

class Solution:
    def fib(self, N: int) -> int:
        if N <= 1:
            return N
        i = 2
        for fib in self.fib_next():
            if i == N:
                return fib
            i += 1
            
    def fib_next(self):
        f_last2, f_last = 0, 1
        while True:
            f = f_last2 + f_last
            f_last2, f_last = f_last, f
            yield f

def outer_generator():
    for i in inner_generator():
        yield i

# AC
# Runtime: 48 ms, faster than 90.31% of Python3 online submissions for Kth Smallest Element in a BST.
# Memory Usage: 17.9 MB, less than 14.91% of Python3 online submissions for Kth Smallest Element in a BST.

class Solution:
    def kthSmallest(self, root: TreeNode, k: int) -> int:
        def ordered_iter(node):
            if node:
                for sub_node in ordered_iter(node.left):
                    yield sub_node
                yield node
                for sub_node in ordered_iter(node.right):
                    yield sub_node

        for node in ordered_iter(root):
            k -= 1
            if k == 0:
                return node.val

# AC
# Runtime: 56 ms, faster than 63.74% of Python3 online submissions for Kth Smallest Element in a BST.
# Memory Usage: 17.7 MB, less than 73.33% of Python3 online submissions for Kth Smallest Element in a BST.

class Solution:
    def kthSmallest(self, root: TreeNode, k: int) -> int:
        def ordered_iter(node):
            if node:
                yield from ordered_iter(node.left)
                yield node
                yield from ordered_iter(node.right)

        for node in ordered_iter(root):
            k -= 1
            if k == 0:
                return node.val

# AC
# Runtime: 112 ms, faster than 57.59% of Python3 online submissions for 24 Game.
# Memory Usage: 13.7 MB, less than 85.60% of Python3 online submissions for 24 Game.

import math
from itertools import permutations, product
from typing import List

class Solution:

    def iter_trees(self, op1, op2, op3, a, b, c, d):
        yield op1(op2(a, b), op3(c, d))
        yield op1(a, op2(op3(b, c), d))
        yield op1(a, op2(b, op3(c, d)))
        yield op1(op2(a, op3(b, c)), d)

    def judgePoint24(self, nums: List[int]) -> bool:
        mul = lambda x, y: x * y
        plus = lambda x, y: x + y
        div = lambda x, y: x / y if y != 0 else math.inf
        minus = lambda x, y: x - y

        op_lst = [plus, minus, mul, div]

        for ops in product(op_lst, repeat=3):
            for val in permutations(nums):
                for v in self.iter_trees(ops[0], ops[1], ops[2], val[0], val[1], val[2], val[3]):
                    if abs(v - 24) < 0.0001:
                        return True
        return False

# AC
# Runtime: 116 ms, faster than 56.23% of Python3 online submissions for 24 Game.
# Memory Usage: 13.9 MB, less than 44.89% of Python3 online submissions for 24 Game.

import math
from itertools import combinations, product, permutations
from typing import List

class Solution:

    def judgePoint24(self, nums: List[int]) -> bool:
        mul = lambda x, y: x * y
        plus = lambda x, y: x + y
        div = lambda x, y: x / y if y != 0 else math.inf
        minus = lambda x, y: x - y

        op_lst = [plus, minus, mul, div]

        def recurse(lst: List[int]):
            if len(lst) == 2:
                for op, values in product(op_lst, permutations(lst)):
                    yield op(values[0], values[1])
            else:
                # choose 2 indices from lst of length n
                for choosen_idx_lst in combinations(list(range(len(lst))), 2):
                    # remaining indices not choosen (of length n-2)
                    idx_remaining_set = set(list(range(len(lst)))) - set(choosen_idx_lst)

                    # remaining values not choosen (of length n-2)
                    value_remaining_lst = list(map(lambda x: lst[x], idx_remaining_set))
                    for op, idx_lst in product(op_lst, permutations(choosen_idx_lst)):
                        # 2 choosen values are lst[idx_lst[0]], lst[idx_lst[1]
                        value_remaining_lst.append(op(lst[idx_lst[0]], lst[idx_lst[1]]))
                        yield from recurse(value_remaining_lst)
                        value_remaining_lst = value_remaining_lst[:-1]

        for v in recurse(nums):
            if abs(v - 24) < 0.0001:
                return True

MDP = Dict[State, Dict[Action, List[Tuple[Prob, State, Reward, bool]]]]

# P[state][action] = [
#    (prob1, next_state1, reward1, is_done),
#    (prob2, next_state2, reward2, is_done), ...]

class Action(Enum):
    UP = 0
    DOWN = 1
    LEFT = 2
    RIGHT = 3

State = int
Reward = float
Prob = float
Policy = Dict[State, Dict[Action, Prob]]
Value = List[float]
StateSet = Set[int]
NonTerminalStateSet = Set[int]
MDP = Dict[State, Dict[Action, List[Tuple[Prob, State, Reward, bool]]]]
# P[s][a] = [(prob, next_state, reward, is_done), ...]

class GridWorldEnv(discrete.DiscreteEnv):
    """
    Grid World environment described in Sutton and Barto Reinforcement Learning 2nd, chapter 4.
    """

    def __init__(self, shape=[4,4]):
        self.shape = shape
        nS = np.prod(shape)
        nA = len(list(Action))
        MAX_R = shape[0]
        MAX_C = shape[1]
        self.grid = np.arange(nS).reshape(shape)
        isd = np.ones(nS) / nS

        # P[s][a] = [(prob, next_state, reward, is_done), ...]
        P: MDP = {}
        action_delta = {Action.UP: (-1, 0), Action.DOWN: (1, 0), Action.LEFT: (0, -1), Action.RIGHT: (0, 1)}
        for s in range(0, MAX_R * MAX_C):
            P[s] = {a.value : [] for a in list(Action)}
            is_terminal = self.is_terminal(s)
            if is_terminal:
                for a in list(Action):
                    P[s][a.value] = [(1.0, s, 0, True)]
            else:
                r = s // MAX_R
                c = s % MAX_R
                for a in list(Action):
                    neighbor_r = min(MAX_R-1, max(0, r + action_delta[a][0]))
                    neighbor_c = min(MAX_C-1, max(0, c + action_delta[a][1]))
                    s_ = neighbor_r * MAX_R + neighbor_c
                    P[s][a.value] = [(1.0, s_, -1, False)]

        super(GridWorldEnv, self).__init__(nS, nA, P, isd)

def policy_evaluate(policy: Policy, env: GridWorldEnv, gamma=1.0, theta=0.0001):
    V = np.zeros(env.nS)
    while True:
        delta = 0
        for s in range(env.nS):
            v = 0
            for a, action_prob in enumerate(policy[s]):
                for prob, next_state, reward, done in env.P[s][a]:
                    v += action_prob * prob * (reward + gamma * V[next_state])
            delta = max(delta, np.abs(v - V[s]))
            V[s] = v
        if delta < theta:
            break
    return np.array(V)

[[  0.         -13.99931242 -19.99901152 -21.99891199]
 [-13.99931242 -17.99915625 -19.99908389 -19.99909436]
 [-19.99901152 -19.99908389 -17.99922697 -13.99942284]
 [-21.99891199 -19.99909436 -13.99942284   0.        ]]