# AUTOGENERATED! DO NOT EDIT! File to edit: nbs/10e_agents.dqn.categorical.ipynb (unless otherwise specified).

__all__ = ['create_support', 'CategoricalDQN', 'distribute', 'final_distribute', 'categorical_update',
           'show_q_distribution', 'PartialCrossEntropy', 'CategoricalDQNTrainer', 'CategoricalArgMaxFeed', 'show_q',
           'show_q_and_max_distribution']

# Cell
# Python native modules
from typing import *
# Third party libs
import torch
from torch.nn import *
from fastcore.all import *
from fastai.learner import *
from fastai.torch_basics import *
from fastai.torch_core import *
from fastai.callback.all import *
# Local modules
from ...data.block_simple import *
from ...data.gym import *
from ...agent import *
from ...core import *
from .core import *
from .targets import *

# Cell
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)

# Cell
class CategoricalDQN(Module):
    def __init__(self,state_sz:int,action_sz:int,n_atoms:int=51,hidden=512,
                 v_min=-10,v_max=10):
        store_attr()
        self.layers=Sequential(
            Linear(state_sz,hidden),
            ReLU(),
            Linear(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))

# Cell
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)

# Cell
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):
    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*support[atom]
        target_z[dones.reshape(-1)]=rewards[dones.reshape(-1)].float()
#         print(rewards,gamma*support[atom])
        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

# Cell
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())

# Cell
def PartialCrossEntropy(p,q): return (-p*q).sum(dim=1).mean()

# Cell
class CategoricalDQNTrainer(Callback):

    def __init__(self,n_batch=0,target_sync=300,discount=0.99,n_steps=1):
        store_attr()
        self._xb=None
        self.local_xb=None
        self.local_pred=None

    def before_fit(self):
        self.learn.target_model=deepcopy(self.learn.model.model)
        self.n_batch=0

    def after_pred(self):
        self.learn.yb=self.xb
        self.learn.xb=self.xb[0]

        self._xb=({k:v.clone() for k,v in self.xb.items()},)

        with torch.no_grad():
            distribution_m=categorical_update(self.learn.target_model.supports,
                                      self.learn.target_model.z_delta,
                                      self.learn.target_model.q(self.xb['next_state']),
                                      self.learn.target_model.p(self.xb['next_state']),
                                      self.xb['action'],self.xb['reward'],
                                      self.xb['done'])
        v=self.learn.model.model(self.xb['state'])
        self.local_v=v
        self.learn.pred=v[np.arange(v.shape[0]),self.xb['action'].reshape(-1,),:]
        self.learn.pred=F.log_softmax(self.learn.pred,dim=1)
        self.learn.xb=(distribution_m,)

        self.local_xb=self.learn.xb
        self.local_pred=self.learn.pred

    def before_backward(self): self.learn.xb=self._xb

    def after_batch(self):
        if self.n_batch%self.target_sync==0:
            self.target_model.load_state_dict(self.learn.model.state_dict())
        self.n_batch+=1

# Cell
class CategoricalArgMaxFeed(AgentCallback):
    def before_action(self):
        raw_action=self.agent.model.policy(self.experience['state'].to(default_device()))
        self.agent.raw_action_shape=raw_action.shape
        self.agent.action=torch.argmax(raw_action,dim=1).reshape(-1,1).to(default_device())

# Cell
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())

# Cell
def show_q_and_max_distribution(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())