# AUTOGENERATED! DO NOT EDIT! File to edit: nbs/11a_agents.policy_gradient.ppo.ipynb (unless otherwise specified).

__all__ = ['Critic', 'ClampBlock', 'show_multisurface', 'ConstBlock', 'clamp_or_const', 'show_surface', 'SimpleGMM',
           'MultiCompGMM']

# Cell
# Python native modules
import os
# Third party libs
import torch
from torch.nn import *
from torch.distributions 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

# Cell
class Critic(Module):
    def __init__(self,state_sz:int,action_sz:int=0,hidden=512):
        "Init model that estimates a single q value based on `state_sz` optionally + `action_sz`"
        self.v=nn.Sequential(
            nn.Linear(state_sz+action_sz,hidden),
            nn.ReLU(),
            nn.Linear(hidden,hidden),
            nn.ReLU(),
            nn.Linear(hidden,1),
        )

    def forward(self, x): return self.v(x.float())

# Cell
class ClampBlock(Module):
    def __init__(self,in_sz:int,out_sz:int,activation_cls=nn.Softplus,clamp_min=0.3,clamp_max=10.0):
        "Linear Layer with clamped outputs."
        store_attr()
        self.fc=nn.Linear(self.in_sz,self.out_sz)
        self.activation=self.activation_cls()

    def forward(self,x): return torch.clamp(self.activation(self.fc(x)),
                                            self.clamp_min,self.clamp_max)

# Cell
def show_multisurface(surfaces,title=''):
    from IPython.display import HTML
    import plotly.graph_objects as go
    from plotly.subplots import make_subplots

    rows=(len(surfaces)+1)//2
    cols=(len(surfaces)+1)%2+1
    fig=make_subplots(rows=rows,cols=cols,
                      specs=[[{"type": "surface"} for _ in range(cols)]]*rows,
                      column_widths=[0.5,0.5])

    for i in range(0,len(surfaces)):
        fig.add_trace(go.Surface(z=to_detach(surfaces[i]),showscale=False),
                      row=(i+2)//2,col=i%2+1)
    fig.update_layout(title=title,autosize=False,
                      width=1000, height=500*rows,#template='plotly_dark',
                      margin=dict(l=50, r=50, b=80, t=90))
    return HTML(fig.to_html())

# Cell
class ConstBlock(Module):
    def __init__(self,out_sz:int,constant=1.0,**kwargs):
        "Module that produces a constant output."
        store_attr()

    def forward(self,x): return torch.full((x.shape[0],self.out_sz),self.constant)

# Cell
@delegates(ClampBlock)
def clamp_or_const(fix_variance:bool,**kwargs):
    return ConstBlock(**kwargs) if fix_variance else ClampBlock(**kwargs)

# Cell
def show_surface(surface,title=''):
    from IPython.display import HTML
    import plotly.graph_objects as go
    from plotly.subplots import make_subplots

    fig=go.Figure()
    fig.add_trace(go.Surface(z=to_detach(surface),showscale=False))
    fig.update_layout(title=title,autosize=False,
                      width=500, height=500,#template='plotly_dark',
                      margin=dict(l=50, r=50, b=80, t=90))
    return HTML(fig.to_html())

# Cell
class SimpleGMM(Module):
    def __init__(self,in_sz,out_sz,fix_variance:bool=False):
        "Single-component GMM parameterized by a fully connected layer with optional std layer."
        store_attr()
        self.mean_fc=nn.Linear(self.in_sz,self.out_sz)
        self.std_fc=clamp_or_const(fix_variance,in_sz=self.in_sz,out_sz=self.out_sz)

    def forward(self,x)->Independent:
        return Independent(Normal(self.mean_fc(x),self.std_fc(x)),1)

# Cell
class MultiCompGMM(Module):
    @delegates(ClampBlock)
    def __init__(self,in_sz,out_sz,n_components,fix_variance:bool=False,**kwargs):
        "Multi-component GMM parameterized by a fully connected layer with optional std layer."
        store_attr()
        self.logit_fc=nn.Linear(self.in_sz,self.n_components)
        self.mean_fcs=nn.ModuleList([nn.Linear(self.in_sz,self.out_sz)
                                     for _ in range(self.n_components)])
        self.std_fcs=nn.ModuleList([clamp_or_const(fix_variance,in_sz=self.in_sz,out_sz=self.out_sz)
                                    for _ in range(self.n_components)])
        self.means,self.logits,self.stds=[],[],[]
        self.log_softmax=LogSoftmax(dim=1)

    def forward(self,x)->MixtureSameFamily:
        self.means=torch.stack([o(x) for o in self.mean_fcs],dim=1)
        self.stds=torch.stack([o(x) for o in self.std_fcs],dim=1)
        self.logits=self.logit_fc(x)
        self.logits=self.log_softmax(self.logits)
        return MixtureSameFamily(
            mixture_distribution=Categorical(logits=self.logits),
            component_distribution=Independent(Normal(self.means,self.stds),1)
        )