import random, time
import numpy as np
import torch
import torch.nn as nn
import torch.optim as optim
from torch.distributions.categorical import Categorical
import gymnasium as gym

ENV_ID = "LunarLander-v2"
NUM_ENVS = 16
NUM_STEPS = 1024
TOTAL_TIMESTEPS = 1000000
LR = 2.5e-4
GAMMA = 0.999
GAE_LAMBDA = 0.98
NUM_MINIBATCHES = 32
UPDATE_EPOCHS = 4
CLIP_COEF = 0.2
ENT_COEF = 0.01
VF_COEF = 0.5
MAX_GRAD_NORM = 0.5
SEED = 1

BATCH_SIZE = NUM_ENVS * NUM_STEPS
MINIBATCH_SIZE = BATCH_SIZE // NUM_MINIBATCHES
NUM_UPDATES = TOTAL_TIMESTEPS // BATCH_SIZE
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print("device:", device, "updates:", NUM_UPDATES, flush=True)

random.seed(SEED); np.random.seed(SEED); torch.manual_seed(SEED)

def make_env(idx):
    def thunk():
        env = gym.make(ENV_ID)
        env = gym.wrappers.RecordEpisodeStatistics(env)
        env.reset(seed=SEED + idx)
        return env
    return thunk

envs = gym.vector.SyncVectorEnv([make_env(i) for i in range(NUM_ENVS)])
obs_dim = int(np.array(envs.single_observation_space.shape).prod())
act_dim = envs.single_action_space.n

def layer_init(layer, std=np.sqrt(2), bias=0.0):
    nn.init.orthogonal_(layer.weight, std)
    nn.init.constant_(layer.bias, bias)
    return layer

class Agent(nn.Module):
    def __init__(self):
        super().__init__()
        self.critic = nn.Sequential(
            layer_init(nn.Linear(obs_dim, 64)), nn.Tanh(),
            layer_init(nn.Linear(64, 64)), nn.Tanh(),
            layer_init(nn.Linear(64, 1), std=1.0))
        self.actor = nn.Sequential(
            layer_init(nn.Linear(obs_dim, 64)), nn.Tanh(),
            layer_init(nn.Linear(64, 64)), nn.Tanh(),
            layer_init(nn.Linear(64, act_dim), std=0.01))
    def get_value(self, x):
        return self.critic(x)
    def get_action_and_value(self, x, action=None):
        logits = self.actor(x)
        probs = Categorical(logits=logits)
        if action is None:
            action = probs.sample()
        return action, probs.log_prob(action), probs.entropy(), self.critic(x)

agent = Agent().to(device)
optimizer = optim.Adam(agent.parameters(), lr=LR, eps=1e-5)

obs = torch.zeros((NUM_STEPS, NUM_ENVS, obs_dim)).to(device)
actions = torch.zeros((NUM_STEPS, NUM_ENVS)).to(device)
logprobs = torch.zeros((NUM_STEPS, NUM_ENVS)).to(device)
rewards = torch.zeros((NUM_STEPS, NUM_ENVS)).to(device)
dones = torch.zeros((NUM_STEPS, NUM_ENVS)).to(device)
values = torch.zeros((NUM_STEPS, NUM_ENVS)).to(device)

global_step = 0
start_time = time.time()
next_obs, _ = envs.reset(seed=SEED)
next_obs = torch.Tensor(next_obs).to(device)
next_done = torch.zeros(NUM_ENVS).to(device)

for update in range(1, NUM_UPDATES + 1):
    frac = 1.0 - (update - 1.0) / NUM_UPDATES
    optimizer.param_groups[0]["lr"] = frac * LR
    ep_returns = []
    for step in range(NUM_STEPS):
        global_step += NUM_ENVS
        obs[step] = next_obs
        dones[step] = next_done
        with torch.no_grad():
            action, logprob, _, value = agent.get_action_and_value(next_obs)
            values[step] = value.flatten()
        actions[step] = action
        logprobs[step] = logprob
        next_obs_np, reward, term, trunc, info = envs.step(action.cpu().numpy())
        done = np.logical_or(term, trunc)
        rewards[step] = torch.tensor(reward).to(device).view(-1)
        next_obs = torch.Tensor(next_obs_np).to(device)
        next_done = torch.Tensor(done).to(device)
        if "final_info" in info:
            for item in info["final_info"]:
                if item and "episode" in item:
                    ep_returns.append(item["episode"]["r"])
    with torch.no_grad():
        next_value = agent.get_value(next_obs).reshape(1, -1)
        advantages = torch.zeros_like(rewards).to(device)
        lastgaelam = 0
        for t in reversed(range(NUM_STEPS)):
            if t == NUM_STEPS - 1:
                nextnonterminal = 1.0 - next_done
                nextvalues = next_value
            else:
                nextnonterminal = 1.0 - dones[t + 1]
                nextvalues = values[t + 1]
            delta = rewards[t] + GAMMA * nextvalues * nextnonterminal - values[t]
            advantages[t] = lastgaelam = delta + GAMMA * GAE_LAMBDA * nextnonterminal * lastgaelam
        returns = advantages + values
    b_obs = obs.reshape((-1, obs_dim))
    b_logprobs = logprobs.reshape(-1)
    b_actions = actions.reshape(-1)
    b_advantages = advantages.reshape(-1)
    b_returns = returns.reshape(-1)
    b_values = values.reshape(-1)
    b_inds = np.arange(BATCH_SIZE)
    for epoch in range(UPDATE_EPOCHS):
        np.random.shuffle(b_inds)
        for start in range(0, BATCH_SIZE, MINIBATCH_SIZE):
            mb_inds = b_inds[start:start + MINIBATCH_SIZE]
            _, newlogprob, entropy, newvalue = agent.get_action_and_value(b_obs[mb_inds], b_actions.long()[mb_inds])
            logratio = newlogprob - b_logprobs[mb_inds]
            ratio = logratio.exp()
            mb_adv = b_advantages[mb_inds]
            mb_adv = (mb_adv - mb_adv.mean()) / (mb_adv.std() + 1e-8)
            pg_loss1 = -mb_adv * ratio
            pg_loss2 = -mb_adv * torch.clamp(ratio, 1 - CLIP_COEF, 1 + CLIP_COEF)
            pg_loss = torch.max(pg_loss1, pg_loss2).mean()
            newvalue = newvalue.view(-1)
            v_loss = 0.5 * ((newvalue - b_returns[mb_inds]) ** 2).mean()
            entropy_loss = entropy.mean()
            loss = pg_loss - ENT_COEF * entropy_loss + v_loss * VF_COEF
            optimizer.zero_grad()
            loss.backward()
            nn.utils.clip_grad_norm_(agent.parameters(), MAX_GRAD_NORM)
            optimizer.step()
    mr = np.mean(ep_returns) if ep_returns else float("nan")
    sps = int(global_step / (time.time() - start_time))
    print(f"update {update}/{NUM_UPDATES} step {global_step} ep_rew_mean {mr:.1f} sps {sps}", flush=True)

envs.close()
torch.save(agent.state_dict(), "ppo_scratch_lunarlander.pt")
print("=== TRAINING DONE, weights saved ===", flush=True)