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

ENV_ID="LunarLander-v2"; SEED=1
NUM_ENVS=8; NUM_STEPS=1024; TOTAL_TIMESTEPS=4_000_000
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
BATCH_SIZE=NUM_ENVS*NUM_STEPS; MINIBATCH_SIZE=BATCH_SIZE//NUM_MINIBATCHES
NUM_UPDATES=TOTAL_TIMESTEPS//BATCH_SIZE
random.seed(SEED); np.random.seed(SEED); torch.manual_seed(SEED)
device=torch.device("cuda" if torch.cuda.is_available() else "cpu")
print("device",device,"updates",NUM_UPDATES,flush=True)

def make_env():
    def thunk():
        e=gym.make(ENV_ID); e=gym.wrappers.RecordEpisodeStatistics(e); return e
    return thunk
envs=gym.vector.SyncVectorEnv([make_env() for _ 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()
next_obs,_=envs.reset(seed=SEED); next_obs=torch.Tensor(next_obs).to(device)
next_done=torch.zeros(NUM_ENVS).to(device)
recent=[]
for update in range(1,NUM_UPDATES+1):
    frac=1.0-(update-1.0)/NUM_UPDATES; optimizer.param_groups[0]["lr"]=frac*LR
    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
        no,r,term,trunc,info=envs.step(action.cpu().numpy())
        done=np.logical_or(term,trunc); rewards[step]=torch.tensor(r).to(device).view(-1)
        next_obs=torch.Tensor(no).to(device); next_done=torch.Tensor(done).to(device)
        if "final_info" in info:
            for it in info["final_info"]:
                if it and "episode" in it: recent.append(float(it["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 st in range(0,BATCH_SIZE,MINIBATCH_SIZE):
            mb=b_inds[st:st+MINIBATCH_SIZE]
            _,newlogprob,entropy,newvalue=agent.get_action_and_value(b_obs[mb],b_actions.long()[mb])
            logratio=newlogprob-b_logprobs[mb]; ratio=logratio.exp()
            mb_adv=b_advantages[mb]; mb_adv=(mb_adv-mb_adv.mean())/(mb_adv.std()+1e-8)
            pg1=-mb_adv*ratio; pg2=-mb_adv*torch.clamp(ratio,1-CLIP_COEF,1+CLIP_COEF)
            pg_loss=torch.max(pg1,pg2).mean()
            newvalue=newvalue.view(-1); v_loss=0.5*((newvalue-b_returns[mb])**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()
    if update%5==0 or update==NUM_UPDATES:
        m=np.mean(recent[-100:]) if recent else float('nan')
        print(f"update {update}/{NUM_UPDATES} step {global_step} ep_rew_mean {m:.1f} sps {int(global_step/(time.time()-start))}",flush=True)
torch.save(agent.state_dict(),"ppo_scratch_lunarlander.pt")
print("=== TRAINING DONE, weights saved ===",flush=True)