import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.distributions import Categorical


class PolicyNetwork(nn.Module):
    def __init__(self, state_dim, n_actions, hidden=64):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(state_dim, hidden), nn.ReLU(),
            nn.Linear(hidden, n_actions),
        )

    def forward(self, x):
        return F.softmax(self.net(x), dim=-1)


class ValueNetwork(nn.Module):
    def __init__(self, state_dim, hidden=64):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(state_dim, hidden), nn.ReLU(),
            nn.Linear(hidden, 1),
        )

    def forward(self, x):
        return self.net(x).squeeze(-1)


def compute_returns(rewards, gamma):
    returns = []
    G = 0.0
    for r in reversed(rewards):
        G = r + gamma * G
        returns.insert(0, G)
    return torch.tensor(returns, dtype=torch.float32)


def reinforce_episode(env, policy, optimizer, gamma=0.99):
    state = env.reset()
    log_probs, rewards = [], []
    done = False
    while not done:
        probs = policy(torch.tensor(state, dtype=torch.float32))
        dist = Categorical(probs)
        action = dist.sample()
        log_probs.append(dist.log_prob(action))
        state, reward, done = env.step(int(action.item()))
        rewards.append(reward)

    returns = compute_returns(rewards, gamma)
    returns = (returns - returns.mean()) / (returns.std() + 1e-8)
    loss = -torch.sum(
        torch.stack(log_probs) * returns
    )
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()
    return sum(rewards)


def ppo_update(
    policy, value_fn, policy_opt, value_opt,
    states, actions, old_log_probs, returns, advantages,
    clip_eps=0.2, epochs=4,
):
    for _ in range(epochs):
        probs = policy(states)
        dist = Categorical(probs)
        new_log_probs = dist.log_prob(actions)

        ratio = torch.exp(new_log_probs - old_log_probs)
        surrogate1 = ratio * advantages
        surrogate2 = torch.clamp(
            ratio, 1 - clip_eps, 1 + clip_eps
        ) * advantages
        policy_loss = -torch.min(surrogate1, surrogate2).mean()

        policy_opt.zero_grad()
        policy_loss.backward()
        policy_opt.step()

        value_loss = F.mse_loss(value_fn(states), returns)
        value_opt.zero_grad()
        value_loss.backward()
        value_opt.step()
