from __future__ import annotations 

import os 
import time 
import json 
import numpy as np 
import torch 
import torch .nn as nn 
import torch .nn .functional as F 
from torch .distributions import Categorical ,Normal 
from typing import Optional ,Tuple ,List 

from server .rl .environment import (
DISCRETE_NVEC ,
N_CONTINUOUS ,
TOTAL_OBS_DIM ,
AlphaBypassEnv ,
)
from server .rl .reward import reward_to_label 






class PolicyNetwork (nn .Module ):

    def __init__ (
    self ,
    obs_dim :int =TOTAL_OBS_DIM ,
    hidden :int =512 ,
    discrete_nvec :List [int ]=DISCRETE_NVEC ,
    n_continuous :int =N_CONTINUOUS ,
    ):
        super ().__init__ ()
        self .discrete_nvec =discrete_nvec 
        self .n_continuous =n_continuous 


        self .trunk =nn .Sequential (
        nn .Linear (obs_dim ,hidden ),
        nn .LayerNorm (hidden ),
        nn .ReLU (),
        nn .Linear (hidden ,hidden ),
        nn .LayerNorm (hidden ),
        nn .ReLU (),
        nn .Linear (hidden ,hidden ),
        nn .LayerNorm (hidden ),
        nn .ReLU (),
        )


        self .discrete_heads =nn .ModuleList ([
        nn .Linear (hidden ,n )for n in discrete_nvec 
        ])


        self .cont_mu =nn .Linear (hidden ,n_continuous )
        self .cont_log_std =nn .Parameter (torch .zeros (n_continuous ))


        self .value_head =nn .Sequential (
        nn .Linear (hidden ,256 ),
        nn .ReLU (),
        nn .Linear (256 ,1 ),
        )

    def forward (self ,obs :torch .Tensor ):
        h =self .trunk (obs )
        logits =[head (h )for head in self .discrete_heads ]
        mu =torch .sigmoid (self .cont_mu (h ))
        log_std =self .cont_log_std .clamp (-4 ,0 )
        value =self .value_head (h ).squeeze (-1 )
        return logits ,mu ,log_std ,value 

    def get_action_and_log_prob (
    self ,
    obs :torch .Tensor ,
    action_masks :Optional [List [Optional [torch .Tensor ]]]=None ,
    )->Tuple [np .ndarray ,np .ndarray ,torch .Tensor ,torch .Tensor ]:

        logits ,mu ,log_std ,value =self .forward (obs )

        discrete_actions =[]
        log_probs_discrete =[]

        for i ,(lg ,n )in enumerate (zip (logits ,self .discrete_nvec )):
            if action_masks and action_masks [i ]is not None :

                mask =action_masks [i ].to (lg .device )
                lg =lg .masked_fill (~mask ,float ("-inf"))
            dist =Categorical (logits =lg )
            a =dist .sample ()
            discrete_actions .append (a .item ())
            log_probs_discrete .append (dist .log_prob (a ))

        log_prob_discrete =torch .stack (log_probs_discrete ).sum ()


        std =log_std .exp ()
        dist_cont =Normal (mu ,std )
        cont_sample =dist_cont .sample ()
        cont_action =cont_sample .clamp (0.0 ,1.0 )
        log_prob_cont =dist_cont .log_prob (cont_sample ).sum ()

        total_log_prob =log_prob_discrete +log_prob_cont 

        return (
        np .array (discrete_actions ,dtype =np .int32 ),
        cont_action .detach ().cpu ().numpy (),
        total_log_prob ,
        value ,
        )

    def evaluate_actions (
    self ,
    obs :torch .Tensor ,
    discrete_actions :torch .Tensor ,
    cont_actions :torch .Tensor ,
    )->Tuple [torch .Tensor ,torch .Tensor ,torch .Tensor ]:
        logits ,mu ,log_std ,value =self .forward (obs )

        log_prob_d =torch .zeros (obs .shape [0 ],device =obs .device )
        entropy_d =torch .zeros (obs .shape [0 ],device =obs .device )
        for i ,lg in enumerate (logits ):
            dist =Categorical (logits =lg )
            log_prob_d +=dist .log_prob (discrete_actions [:,i ])
            entropy_d +=dist .entropy ()

        std =log_std .exp ()
        dist_c =Normal (mu ,std )
        log_prob_c =dist_c .log_prob (cont_actions ).sum (-1 )
        entropy_c =dist_c .entropy ().sum (-1 )

        return log_prob_d +log_prob_c ,(entropy_d +entropy_c )/2 ,value 






class RolloutBuffer :
    def __init__ (self ):
        self .clear ()

    def clear (self ):
        self .obs :List [np .ndarray ]=[]
        self .discrete_actions :List [np .ndarray ]=[]
        self .cont_actions :List [np .ndarray ]=[]
        self .rewards :List [float ]=[]
        self .log_probs :List [torch .Tensor ]=[]
        self .values :List [torch .Tensor ]=[]
        self .dones :List [bool ]=[]

    def add (self ,obs ,d_action ,c_action ,reward ,log_prob ,value ,done ):
        self .obs .append (obs )
        self .discrete_actions .append (d_action )
        self .cont_actions .append (c_action )
        self .rewards .append (reward )
        self .log_probs .append (log_prob )
        self .values .append (value )
        self .dones .append (done )

    def compute_returns (self ,last_value :float ,gamma :float =0.99 ,gae_lambda :float =0.95 ):
        advantages =[]
        gae =0.0 
        values =[v .item ()for v in self .values ]+[last_value ]

        for t in reversed (range (len (self .rewards ))):
            delta =self .rewards [t ]+gamma *values [t +1 ]*(1 -self .dones [t ])-values [t ]
            gae =delta +gamma *gae_lambda *(1 -self .dones [t ])*gae 
            advantages .insert (0 ,gae )

        returns =[a +v .item ()for a ,v in zip (advantages ,self .values )]
        return advantages ,returns 

    def to_tensors (self ,device :torch .device ):
        obs =torch .FloatTensor (np .stack (self .obs )).to (device )
        d_act =torch .LongTensor (np .stack (self .discrete_actions )).to (device )
        c_act =torch .FloatTensor (np .stack (self .cont_actions )).to (device )
        return obs ,d_act ,c_act 






class PPOTrainer :
    def __init__ (
    self ,
    env :AlphaBypassEnv ,
    device_str :str ="cuda",
    lr :float =3e-4 ,
    gamma :float =0.99 ,
    gae_lambda :float =0.95 ,
    clip_eps :float =0.2 ,
    entropy_coef :float =0.01 ,
    vf_coef :float =0.5 ,
    max_grad_norm :float =0.5 ,
    update_epochs :int =4 ,
    steps_per_update :int =8 ,
    checkpoint_dir :str ="checkpoints",
    checkpoint_every :int =100 ,
    ):
        self .env =env 
        self .device =torch .device (device_str if torch .cuda .is_available ()else "cpu")
        print (f"[PPO] device: {self .device }")

        self .policy =PolicyNetwork ().to (self .device )
        self .optimizer =torch .optim .Adam (self .policy .parameters (),lr =lr )
        self .scheduler =torch .optim .lr_scheduler .ExponentialLR (self .optimizer ,gamma =0.999 )

        self .gamma =gamma 
        self .gae_lambda =gae_lambda 
        self .clip_eps =clip_eps 
        self .entropy_coef =entropy_coef 
        self .vf_coef =vf_coef 
        self .max_grad_norm =max_grad_norm 
        self .update_epochs =update_epochs 
        self .steps_per_update =steps_per_update 
        self .checkpoint_dir =checkpoint_dir 
        self .checkpoint_every =checkpoint_every 

        os .makedirs (checkpoint_dir ,exist_ok =True )

        self .total_episodes =0 
        self .best_reward =-float ("inf")
        self .reward_history :List [float ]=[]

    def _build_action_masks (self ,obs_tensor :torch .Tensor )->List [Optional [torch .Tensor ]]:

        return [None ]*len (DISCRETE_NVEC )

    def collect_rollout (self )->RolloutBuffer :
        buffer =RolloutBuffer ()
        obs =self .env ._build_obs ()

        for _ in range (self .steps_per_update ):
            obs_t =torch .FloatTensor (obs ).unsqueeze (0 ).to (self .device )

            with torch .no_grad ():
                masks =self ._build_action_masks (obs_t )
                d_action ,c_action ,log_prob ,value =self .policy .get_action_and_log_prob (
                obs_t .squeeze (0 ),masks 
                )

            next_obs ,reward ,done ,info =self .env .step (d_action ,c_action )

            self .total_episodes +=1 
            self .reward_history .append (reward )

            print (
            f"[Ep {self .total_episodes :04d}] "
            f"reward={reward :+.4f} {reward_to_label (reward )} | "
            f"transport={info ['transport']:5s} dest={info ['dest']:30s} | "
            f"stable={info ['stability']:.2f} "
            f"speed={info ['throughput_mbps']:.2f}Mbps"
            )

            buffer .add (obs ,d_action ,c_action ,reward ,log_prob ,value ,done )
            obs =next_obs 

            if done :
                obs =self .env .reset ()


            if self .total_episodes %self .checkpoint_every ==0 :
                self .save_checkpoint ()

        return buffer 

    def update (self ,buffer :RolloutBuffer ):
        print (f"\n[PPO] ── Update #{self .total_episodes //self .steps_per_update } ──────────────────────────")
        print (f"[PPO] Buffer: {len (buffer .rewards )} episodes | "
        f"mean_reward={sum (buffer .rewards )/len (buffer .rewards ):+.4f} | "
        f"positive={sum (1 for r in buffer .rewards if r >0 )}/{len (buffer .rewards )}")
        obs_t ,d_act_t ,c_act_t =buffer .to_tensors (self .device )


        with torch .no_grad ():
            last_obs =torch .FloatTensor (self .env ._build_obs ()).to (self .device )
            _ ,_ ,_ ,last_val =self .policy .forward (last_obs .unsqueeze (0 ))
            last_value =last_val .item ()

        advantages ,returns =buffer .compute_returns (last_value ,self .gamma ,self .gae_lambda )
        adv_t =torch .FloatTensor (advantages ).to (self .device )
        ret_t =torch .FloatTensor (returns ).to (self .device )


        adv_t =(adv_t -adv_t .mean ())/(adv_t .std ()+1e-8 )

        old_log_probs =torch .stack (buffer .log_probs ).to (self .device ).detach ()

        for _ in range (self .update_epochs ):
            log_probs ,entropy ,values =self .policy .evaluate_actions (obs_t ,d_act_t ,c_act_t )

            ratio =(log_probs -old_log_probs ).exp ()
            surr1 =ratio *adv_t 
            surr2 =ratio .clamp (1 -self .clip_eps ,1 +self .clip_eps )*adv_t 

            policy_loss =-torch .min (surr1 ,surr2 ).mean ()
            value_loss =F .mse_loss (values ,ret_t )
            entropy_loss =-entropy .mean ()

            loss =policy_loss +self .vf_coef *value_loss +self .entropy_coef *entropy_loss 

            self .optimizer .zero_grad ()
            loss .backward ()
            nn .utils .clip_grad_norm_ (self .policy .parameters (),self .max_grad_norm )
            self .optimizer .step ()

        self .scheduler .step ()

    def train (self ,total_episodes :int =10000 ):

        print (f"\n{'='*60 }")
        print (f"  AlphaBypass — PPO Training")
        print (f"  Target: {total_episodes } episodes")
        print (f"  Device: {self .device }")
        print (f"{'='*60 }\n")

        obs =self .env .reset ()

        while self .total_episodes <total_episodes :
            buffer =self .collect_rollout ()
            self .update (buffer )


            if len (self .reward_history )>=20 :
                recent =self .reward_history [-20 :]
                print (
                f"\n[Stats] last 20 episodes: "
                f"mean={np .mean (recent ):+.4f} "
                f"max={np .max (recent ):+.4f} "
                f"min={np .min (recent ):+.4f}\n"
                )

    def save_checkpoint (self ,tag :str =""):
        path =os .path .join (
        self .checkpoint_dir ,
        f"checkpoint_ep{self .total_episodes :05d}{tag }.pt"
        )
        torch .save ({
        "episode":self .total_episodes ,
        "policy_state":self .policy .state_dict (),
        "optimizer_state":self .optimizer .state_dict (),
        "reward_history":self .reward_history ,
        "best_reward":self .best_reward ,
        },path )
        print (f"[Checkpoint] saved → {path }")


        r =np .mean (self .reward_history [-10 :])if len (self .reward_history )>=10 else -999 
        if r >self .best_reward :
            self .best_reward =r 
            best_path =os .path .join (self .checkpoint_dir ,"best.pt")
            torch .save (torch .load (path ),best_path )
            print (f"[Checkpoint] 🏆 new best ({r :+.4f}) → {best_path }")

    def load_checkpoint (self ,path :str ):
        ck =torch .load (path ,map_location =self .device )
        self .policy .load_state_dict (ck ["policy_state"])
        self .optimizer .load_state_dict (ck ["optimizer_state"])
        self .total_episodes =ck ["episode"]
        self .reward_history =ck .get ("reward_history",[])
        self .best_reward =ck .get ("best_reward",-float ("inf"))
        print (f"[Checkpoint] loaded from ep {self .total_episodes }")