# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/12o_agents.dqn.categorical.ipynb.

# %% auto 0
__all__ = ['create_support', 'CategoricalDQN', 'distribute', 'final_distribute', 'categorical_update', 'show_q_distribution',
           'PartialCrossEntropy', 'CategoricalTargetQCalc', 'MultiModelRunner', 'show_q']

# %% ../nbs/12o_agents.dqn.categorical.ipynb 3
# Python native modules
import os
from collections import deque
# Third party libs
from fastcore.all import *
import torchdata.datapipes as dp
from torch.utils.data.dataloader_experimental import DataLoader2
from torch.utils.data.datapipes._typing import _DataPipeMeta, _IterDataPipeMeta
from torchdata.dataloader2.graph import find_dps,traverse,DataPipe,remove_dp,replace_dp
import torch
from torch.nn import *
import torch.nn.functional as F
from torch.optim import *
from fastai.torch_basics import *
from fastai.torch_core import *
# Local modules

from ...core import *
from ..core import *
from ...pipes.core import *
from ...fastai.data.block import *
from ...memory.experience_replay import *
from ..core import *
from ..discrete import *
from ...loggers.core import *
from ...loggers.jupyter_visualizers import *
from ...learner.core import *
from .basic import *
from .target import *

# %% ../nbs/12o_agents.dqn.categorical.ipynb 10
def create_support(v_min=-10,v_max=10,n_atoms=51)->Tuple[Tensor,float]:
    "Creates the support and returns the z_delta that was used."
    z_delta=(v_max-v_min)/(n_atoms-1)
    return (Tensor([i*z_delta for i in range(n_atoms)])+v_min,z_delta)

# %% ../nbs/12o_agents.dqn.categorical.ipynb 29
class CategoricalDQN(Module):
    def __init__(self,
                 state_sz:int,
                 action_sz:int,
                 n_atoms:int=51,
                 hidden=512,
                 v_min=-10,
                 v_max=10,
                 head_layer=Linear,
                 activation_fn=ReLU
                ):
        store_attr()
        self.layers=Sequential(
            Linear(state_sz,hidden),
            activation_fn(),
            head_layer(hidden,action_sz*n_atoms),
        )
        self.supports,self.z_delta=create_support(v_min=v_min,v_max=v_max,n_atoms=n_atoms)
        self.softmax=Softmax(dim=2)
    
    def to(self, *args, **kwargs):
        self = super().to(*args, **kwargs) 
        self.supports=self.supports.to(*args, **kwargs)
        return self
    
    def forward(self,x): 
        return self.layers(x).view(x.shape[0],self.action_sz,self.n_atoms)
    def policy(self,x): return (self.supports*self.p(x)).mean(dim=2)
    def q(self,x): return (self.supports*self.p(x)).sum(dim=2)
    def p(self,x): return self.softmax(self(x))

# %% ../nbs/12o_agents.dqn.categorical.ipynb 34
def distribute(projection,left,right,support_value,p_a,atom,done):
    "Does: m_l <- m_l + p_j(x_{t+1},a*)(u - b_j) operation for non-final states."
    diffs=torch.hstack((support_value-left,right-support_value))
    # If they are the same location, then just split the value in half, and add twice
    diffs[(left==right).reshape(-1,),:]=0.5 
    mask=~done.reshape(-1,)
    
    left_v=projection[mask].gather(1,left[mask])+(p_a[:,atom][mask]*diffs[:,0][mask]).reshape(-1,1)
    right_v=projection[mask].gather(1,right[mask])+(p_a[:,atom][mask]*diffs[:,1][mask]).reshape(-1,1)
    
    projection[mask]=projection[mask].scatter(dim=1,index=left[mask],src=left_v)
    projection[mask]=projection[mask].scatter(dim=1,index=right[mask],src=right_v)

def final_distribute(projection,left,right,support_value,p_a,atom,done):
    "Does: m_l <- m_l + p_j(x_{t+1},a*)(u - b_j) operation for final states."
    diffs=torch.hstack((support_value-left,right-support_value))
    # If they are the location, then just split the value in half, and add twice
    diffs[(left==right).reshape(-1,),:]=0.5 
    mask=done.reshape(-1,)
    
    left_v=diffs[:,0].reshape(-1,1)
    right_v=projection[mask].gather(1,right)+diffs[:,1].reshape(-1,1)

    projection[mask]=0.0
    projection[mask]=projection[mask].scatter(dim=1,index=left,src=left_v)
    projection[mask]=projection[mask].scatter(dim=1,index=right,src=right_v)

# %% ../nbs/12o_agents.dqn.categorical.ipynb 35
def categorical_update(support,delta_z,q,p,actions,rewards,dones,v_min=-10,
                       v_max=10,n_atoms=51,gamma=0.99,passes=None,nsteps=1,debug=False):
    if debug:
        print(f'support: {support.shape}, delta_z: {delta_z}, q: {q.shape}\n',
              f'\tp: {p.shape}, actions: {actions.shape}, rewards: {rewards.shape}\n',
              f'\tdones: {dones.shape}')
    
    bs=q.shape[0]
    passes=ifnone(passes,n_atoms)
    # Do this outside of the loop so we only have to do it once
    # Represents: p_j(x_{t+1},a*)
    p_a=p[torch.arange(bs),actions.reshape(-1,)]#.reshape(-1,1)
    # get a*
    next_actions=torch.argmax(q,dim=1)
    # m_i = 0 for i in [0,N-1]
    projection=torch.zeros((bs,n_atoms)).to(device=default_device())
    # j in [0, N - 1]
    for atom in range(0,passes):
        # Tz_j <- [r_t + gamma * z_j]_v_min^v_max
        target_z=rewards+(gamma**nsteps)*support[atom]
        target_z[dones.reshape(-1)]=rewards[dones.reshape(-1)].float()
        target_z=torch.clamp(target_z,v_min,v_max)
        # b_j <- (Tz_j - Vmin)/delta_z
        support_value=(target_z-v_min)/delta_z
        # l <- floor(b_j), u <- ceil(b_j)
        left,right=support_value.floor().long(),support_value.ceil().long()
        # m_l <- m_l + p_j(x_{t+1},a*)(u - b_j)
        distribute(projection,left,right,support_value,p_a,atom,dones)

    if dones.sum()>=1:
        target_z=rewards[dones.reshape(-1)]
        target_z=torch.clamp(target_z,v_min,v_max)
        support_value=(target_z-v_min)/delta_z
        left,right=support_value.floor().long(),support_value.ceil().long()
        final_distribute(projection,left,right,support_value,p_a,atom,dones)
    return projection

# %% ../nbs/12o_agents.dqn.categorical.ipynb 36
def show_q_distribution(cat_dist,title='Update Distributions'):
    "`cat_dist` being shape: (bs,n_atoms)"
    from IPython.display import HTML
    import plotly.graph_objects as go
    fig = go.Figure(data=[go.Surface(z=to_detach(cat_dist).numpy())])
    fig.update_layout(title=title,autosize=False,
                      width=500, height=500,#template='plotly_dark',
                      margin=dict(l=65, r=50, b=80, t=90))
    return HTML(fig.to_html())

# %% ../nbs/12o_agents.dqn.categorical.ipynb 39
def PartialCrossEntropy(p,q): return (-p*q).sum(dim=1).mean()

# %% ../nbs/12o_agents.dqn.categorical.ipynb 40
class CategoricalTargetQCalc(dp.iter.IterDataPipe):
    debug = False
    
    def __init__(self,source_datapipe,discount=0.99,nsteps=1,device=None,double_dqn_strategy=False):
        self.source_datapipe = source_datapipe
        self.discount = discount
        self.nsteps = nsteps
        self.device = device
        self.double_dqn_strategy = double_dqn_strategy
        if source_datapipe is not None: self.learner = find_dp(traverse(self),LearnerBase)
        
    def reset(self):
        self.learner = find_dp(traverse(self),LearnerBase)
        
    def __iter__(self):
        for batch in self.source_datapipe:            
            if self.device is not None: batch = batch.device(self.device)
            actions = batch.action.long()
            self.learner.done_mask = batch.terminated.reshape(-1,)
            with torch.no_grad():
                target_actions = actions
                if self.double_dqn_strategy:
                    target_actions = self.learner.model.policy(batch.next_state).argmax(dim=1).reshape(-1,)
                    target_actions = target_actions.long()
                
                distribution_m = categorical_update(self.learner.target_model.supports,
                                          self.learner.target_model.z_delta,
                                          self.learner.target_model.q(batch.next_state),
                                          self.learner.target_model.p(batch.next_state),
                                          target_actions,batch.reward,
                                          self.learner.done_mask,nsteps=self.nsteps,
                                          debug=self.debug)
            self.learner.target_qs = distribution_m # (distribution_m,)
            v = self.learner.model(batch.state)
            self.learner.pred_raw = v[np.arange(v.shape[0]),actions.reshape(-1,),:]
            self.learner.pred = F.log_softmax(self.learner.pred_raw,dim=1)
            yield batch
    
    @classmethod
    def replace_remove_dp(cls,
                          old_dp=(QCalc,TargetModelQCalc),
                          rm_dp=(TargetCalc,),
                          device=None,
                          discount=0.99,
                          nsteps=1,
                          double_dqn_strategy=False
                         ) -> Callable[[DataPipe],DataPipe]:
        def _replace_remove_dp(pipe):
            found = False
            for _old_dp in old_dp:
                try:
                    old_dp_instance = find_dp(traverse(pipe),_old_dp)
                    v = replace_dp(
                        traverse(pipe),
                        old_dp_instance,
                        cls(old_dp_instance.source_datapipe,device=device,discount=discount,nsteps=nsteps,
                            double_dqn_strategy=double_dqn_strategy)
                    )
                    pipe = list(v.values())[0][0]
                    found = True
                except LookupError: pass
            if not found: 
                warn(f'Unable to find any: {old_dp} to replace with {cls} given {traverse(pipe)}')
            found = False
            for _rm_dp in rm_dp:
                try:
                    old_dp_instance = find_dp(traverse(pipe),_rm_dp)
                    v = remove_dp(
                        traverse(pipe),
                        old_dp_instance
                    )
                    pipe = list(v.values())[0][0]
                    found = True
                except LookupError: pass
                except RuntimeError:
                    print(traverse(pipe))
                    raise
            if not found: 
                warn(f'Unable to find any: {rm_dp} to replace with {cls} given {traverse(pipe)}')
            return pipe
        return _replace_remove_dp

# %% ../nbs/12o_agents.dqn.categorical.ipynb 41
class MultiModelRunner(dp.iter.IterDataPipe):
    "If a model contains multiple models, then we support selecting a sub model."
    def __init__(self,
                 source_datapipe,
                 device:Optional[str]=None,
                 model_name:str='policy'
                ): 
        self.source_datapipe = source_datapipe
        self.agent_base = find_dp(traverse(self),AgentBase)
        self.model = getattr(self.agent_base.model,model_name)
        self.device = device
    
    def __iter__(self):
        for x in self.source_datapipe:
            if self.device is not None: 
                x = x.to(torch.device(self.device))
            if len(x.shape)==1: x = x.unsqueeze(0)
            yield self.model(x)
            
    @classmethod
    def _replace_dp(cls,pipe,old_dp=SimpleModelRunner,device=None,model_name:str='policy'):
        old_dp_instance = find_dp(traverse(pipe),old_dp)
        v = replace_dp(
            traverse(pipe),
            old_dp_instance,
            cls(old_dp_instance.source_datapipe,device=device,model_name=model_name)
        )
        return list(v.values())[0][0]
            
    @classmethod
    def replace_dp(cls,old_dp=SimpleModelRunner,device=None,model_name:str='policy') -> Callable[[DataPipe],DataPipe]:
        return partial(cls._replace_dp,old_dp=old_dp,device=device,model_name=model_name)

# %% ../nbs/12o_agents.dqn.categorical.ipynb 45
def show_q(cat_dist,title='Update Distributions'):
    "`cat_dist` being shape: (bs,n_atoms)"
    from IPython.display import HTML
    import plotly.graph_objects as go
    from plotly.subplots import make_subplots
    
    distributions=to_detach(cat_dist).numpy()
    actions=np.argmax(distributions,axis=1).reshape(-1,)
    fig = make_subplots(rows=1, cols=2,
                        specs=[[{"type": "surface"},{"type": "xy"}]])
    fig.add_trace(go.Surface(z=distributions, showscale=False),row=1, col=1)
    fig.add_trace(go.Scatter(x=np.arange(len(actions)),y=actions),row=1, col=2)
    fig.update_layout(title=title,autosize=False,
                      width=1000, height=500,#template='plotly_dark',
                      margin=dict(l=65, r=50, b=80, t=90))
    return HTML(fig.to_html())
