# AUTOGENERATED! DO NOT EDIT! File to edit: nbs/17_cpc.ipynb (unless otherwise specified).

__all__ = ['CPCEncoder', 'CPCModel']

# Cell

import torch

import torch.nn.functional as F

import torch.nn as nn

from clinical_ts.basic_conv1d import _conv1d

import numpy as np

from .basic_conv1d import listify, bn_drop_lin

# Cell

class CPCEncoder(nn.Sequential):

    'CPC Encoder'

    def __init__(self, input_channels, strides=[5,4,2,2,2], kss=[10,8,4,4,4], features=[512,512,512,512],bn=False):

        assert(len(strides)==len(kss) and len(strides)==len(features))

        lst = []

        for i,(s,k,f) in enumerate(zip(strides,kss,features)):

            lst.append(_conv1d(input_channels if i==0 else features[i-1],f,kernel_size=k,stride=s,bn=bn))

        super().__init__(*lst)

        self.downsampling_factor = np.prod(strides)

        self.output_dim = features[-1]

        # output: bs, output_dim, seq//downsampling_factor

    def encode(self, input):

        #bs = input.size()[0]

        #ch = input.size()[1]

        #seq = input.size()[2]

        #segments = seq//self.downsampling_factor

        #input_encoded = self.forward(input[:,:,:segments*self.downsampling_factor]).transpose(1,2) #bs, seq//downsampling, encoder_output_dim (standard ordering for batch_first RNNs)

        input_encoded = self.forward(input).transpose(1,2)

        return input_encoded

# Cell

class CPCModel(nn.Module):

    "CPC model"

    def __init__(self, input_channels, strides=[5,4,2,2,2], kss=[10,8,4,4,4], features=[512,512,512,512],bn_encoder=False, n_hidden=512,n_layers=2,mlp=False,lstm=True,bias_proj=False, num_classes=None, concat_pooling=True, ps_head=0.5,lin_ftrs_head=[512],bn_head=True,skip_encoder=False):

        super().__init__()

        assert(skip_encoder is False or num_classes is not None)#pretraining only with encoder

        self.encoder = CPCEncoder(input_channels,strides=strides,kss=kss,features=features,bn=bn_encoder) if skip_encoder is False else None

        self.encoder_output_dim = self.encoder.output_dim if skip_encoder is False else None

        self.encoder_downsampling_factor = self.encoder.downsampling_factor if skip_encoder is False else None

        self.n_hidden = n_hidden

        self.n_layers = n_layers

        self.mlp = mlp

        self.num_classes = num_classes

        self.concat_pooling = concat_pooling

        self.rnn = nn.LSTM(self.encoder_output_dim if skip_encoder is False else input_channels,n_hidden,num_layers=n_layers,batch_first=True) if lstm is True else nn.GRU(self.encoder.output_dim,n_hidden,num_layers=n_layers,batch_first=True)

        if(num_classes is None): #pretraining

            if(mlp):# additional hidden layer as in simclr

                self.proj = nn.Sequential(nn.Linear(n_hidden, n_hidden),nn.ReLU(inplace=True),nn.Linear(n_hidden, self.encoder_output_dim,bias=bias_proj))

            else:

                self.proj = nn.Linear(n_hidden, self.encoder_output_dim,bias=bias_proj)

        else: #classifier

            #slightly adapted from RNN1d

            layers_head =[]

            if(self.concat_pooling):

                layers_head.append(AdaptiveConcatPoolRNN())

            #classifier

            nf = 3*n_hidden if concat_pooling else n_hidden

            lin_ftrs_head = [nf, num_classes] if lin_ftrs_head is None else [nf] + lin_ftrs_head + [num_classes]

            ps_head = listify(ps_head)

            if len(ps_head)==1:

                ps_head = [ps_head[0]/2] * (len(lin_ftrs_head)-2) + ps_head

            actns = [nn.ReLU(inplace=True)] * (len(lin_ftrs_head)-2) + [None]

            for ni,no,p,actn in zip(lin_ftrs_head[:-1],lin_ftrs_head[1:],ps_head,actns):

                layers_head+=bn_drop_lin(ni,no,bn_head,p,actn)

            self.head=nn.Sequential(*layers_head)

    def forward(self, input):

        # input shape bs,ch,seq

        if(self.encoder is not None):

            input_encoded = self.encoder.encode(input)

        else:

            input_encoded = input.transpose(1,2) #bs, seq, channels

        output_rnn, _ = self.rnn(input_encoded) #output_rnn: bs, seq, n_hidden

        if(self.num_classes is None):#pretraining

            return input_encoded, self.proj(output_rnn)

        else:#classifier

            output = output_rnn.transpose(1,2)#bs,n_hidden,seq (i.e. standard CNN channel ordering)

            if(self.concat_pooling is False):

                output = output[:,:,-1]

            return self.head(output)

    def get_layer_groups(self):

        return (self.encoder,self.rnn,self.head)

    def get_output_layer(self):

        return self.head[-1]

    def set_output_layer(self,x):

        self.head[-1] = x

    def cpc_loss(self,input, target=None, steps_predicted=5, n_false_negatives=9, negatives_from_same_seq_only=False, eval_acc=False):

        assert(self.num_classes is None)

        input_encoded, output = self.forward(input) #input_encoded: bs, seq, features; output: bs,seq,features

        input_encoded_flat = input_encoded.reshape(-1,input_encoded.size(2)) #for negatives below: -1, features

        bs = input_encoded.size()[0]

        seq = input_encoded.size()[1]

        loss = torch.tensor(0,dtype=torch.float32).to(input.device)

        tp_cnt = torch.tensor(0,dtype=torch.int64).to(input.device)

        for i in range(input_encoded.size()[1]-steps_predicted):

            positives = input_encoded[:,i+steps_predicted].unsqueeze(1) #bs,1,encoder_output_dim

            if(negatives_from_same_seq_only):

                idxs = torch.randint(0,(seq-1),(bs*n_false_negatives,)).to(input.device)

            else:#negative from everywhere

                idxs = torch.randint(0,bs*(seq-1),(bs*n_false_negatives,)).to(input.device)

            idxs_seq = torch.remainder(idxs,seq-1) #bs*false_neg

            idxs_seq2 = idxs_seq * (idxs_seq<(i+steps_predicted)).long() +(idxs_seq+1)*(idxs_seq>=(i+steps_predicted)).long()#bs*false_neg

            if(negatives_from_same_seq_only):

                idxs_batch = torch.arange(0,bs).repeat_interleave(n_false_negatives).to(input.device)

            else:

                idxs_batch = idxs//(seq-1)

            idxs2_flat = idxs_batch*seq+idxs_seq2 #for negatives from everywhere: this skips step i+steps_predicted from the other sequences as well for simplicity

            negatives = input_encoded_flat[idxs2_flat].view(bs,n_false_negatives,-1) #bs*false_neg, encoder_output_dim

            candidates = torch.cat([positives,negatives],dim=1)#bs,false_neg+1,encoder_output_dim

            preds=torch.sum(output[:,i].unsqueeze(1)*candidates,dim=-1) #bs,(false_neg+1)

            targs = torch.zeros(bs, dtype=torch.int64).to(input.device)

            if(eval_acc):

                preds_argmax = torch.argmax(preds,dim=-1)

                tp_cnt += torch.sum(preds_argmax == targs)

            loss += F.cross_entropy(preds,targs)

        if(eval_acc):

            return loss, tp_cnt.float()/bs/(input_encoded.size()[1]-steps_predicted)

        else:

            return loss

#copied from RNN1d

class AdaptiveConcatPoolRNN(nn.Module):

    def __init__(self, bidirectional=False):

        super().__init__()

        self.bidirectional = bidirectional

    def forward(self,x):

        #input shape bs, ch, ts

        t1 = nn.AdaptiveAvgPool1d(1)(x)

        t2 = nn.AdaptiveMaxPool1d(1)(x)

        if(self.bidirectional is False):

            t3 = x[:,:,-1]

        else:

            channels = x.size()[1]

            t3 = torch.cat([x[:,:channels,-1],x[:,channels:,0]],1)

        out=torch.cat([t1.squeeze(-1),t2.squeeze(-1),t3],1) #output shape bs, 3*ch

        return out

#class CPCClassifier(nn.Module):

#    def __init__(self, cpcmodel, num_classes, concat_pooling=True, ps_head=0.5,lin_ftrs_head=None,bn_head=True):

#        super().__init__()

#        self.cpcmodel = cpcmodel

#        self.concat_pooling = concat_pooling

#        

#        #slightly adapted from RNN1d

#        layers_head =[]

#        if(self.concat_pooling):

#            layers_head.append(AdaptiveConcatPoolRNN())

#

#        #classifier

#        nf = 3*self.cpcmodel.encoder_output_dim if concat_pooling else self.cpcmodel.encoder_output_dim

#        lin_ftrs_head = [nf, num_classes] if lin_ftrs_head is None else [nf] + lin_ftrs_head + [num_classes]

#        ps_head = listify(ps_head)

#        if len(ps_head)==1:

#            ps_head = [ps_head[0]/2] * (len(lin_ftrs_head)-2) + ps_head

#        actns = [nn.ReLU(inplace=True)] * (len(lin_ftrs_head)-2) + [None]

#

#        for ni,no,p,actn in zip(lin_ftrs_head[:-1],lin_ftrs_head[1:],ps_head,actns):

#            layers_head+=bn_drop_lin(ni,no,bn_head,p,actn)

#        self.head=nn.Sequential(*layers_head)

#        

#    def forward(self,input):

#        output = self.cpcmodel(input)

#        if(self.concat_pooling is False):

#            output = output[:,:,-1]

#        return self.head(output)

#    

#    def get_layer_groups(self):

#        return (self.cpcmodel.encoder,self.cpcmodel.rnn,self.head)

#

#    def get_output_layer(self):

#        return self.head[-1]

#

#    def set_output_layer(self,x):

#        self.head[-1] = x