#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Fri Jan  1 16:02:07 2021

@author: hasan
"""

import torch
import torch.nn as nn
import torch._utils
import torch.nn.functional as F


class BNReLU(nn.Sequential):
    def __init__(self,in_channels):
        super().__init__(nn.BatchNorm2d(in_channels),
                         nn.ReLU())

class SpatialGather_Module(nn.Module):
    """
        Aggregate the context features according to the initial 
        predicted probability distribution.
        Employ the soft-weighted method to aggregate the context.
    """
    def __init__(self, cls_num=0, scale=1):
        super(SpatialGather_Module, self).__init__()
        self.cls_num = cls_num
        #self.scale = scale

    def forward(self, feats, probs):
        batch_size, c, h, w = probs.size(0), probs.size(1), probs.size(2), probs.size(3)
        probs = probs.view(batch_size, c, -1)
        feats = feats.view(batch_size, feats.size(1), -1)
        feats = feats.permute(0, 2, 1) # batch x hw x c 
        probs = F.softmax(probs, dim=2) 

        ocr_context = torch.matmul(probs, feats).permute(0, 2, 1).unsqueeze(3)# batch x k x c
        return ocr_context


class ObjectAttentionBlock2D(nn.Module):
    '''
    The basic implementation for object context block
    Input:
        N X C X H X W
    Parameters:
        in_channels       : the dimension of the input feature map
        key_channels      : the dimension after the key/query transform
        scale             : choose the scale to downsample the input feature maps (save memory cost)
        bn_type           : specify the bn type
    Return:
        N X C X H X W
    '''
    def __init__(self, 
                 in_channels, 
                 key_channels,
                 cls_num = 1):
        super(ObjectAttentionBlock2D, self).__init__()
        self.in_channels = in_channels
        self.key_channels = key_channels
        self.cls_num = cls_num
        self.f_pixel = nn.Sequential(
            nn.Conv2d(in_channels=self.in_channels, out_channels=self.key_channels,
                kernel_size=1, stride=1, padding=0, bias=False),
            BNReLU(self.key_channels),
            nn.Conv2d(in_channels=self.key_channels, out_channels=self.key_channels,
                kernel_size=1, stride=1, padding=0, bias=False),
            BNReLU(self.key_channels),
        )
        
        self.f_object = nn.Sequential(
            nn.Conv2d(in_channels=self.in_channels, out_channels=self.key_channels,
                kernel_size=1, stride=1, padding=0, bias=False),
            BNReLU(self.key_channels),
            nn.Conv2d(in_channels=self.key_channels, out_channels=self.key_channels,
                kernel_size=1, stride=1, padding=0, bias=False),
            BNReLU(self.key_channels),
        )
        
        self.f_down = nn.Sequential(
            nn.Conv2d(in_channels=self.in_channels, out_channels=self.key_channels,
                kernel_size=1, stride=1, padding=0, bias=False),
            BNReLU(self.key_channels),
        )
        self.f_up = nn.Sequential(
            nn.Conv2d(in_channels=self.key_channels, out_channels=self.in_channels,
                kernel_size=1, stride=1, padding=0, bias=False),
            BNReLU(self.in_channels),
        )

    def forward(self, x, proxy):
        batch_size, h, w = x.size(0), x.size(2), x.size(3)

        query = self.f_pixel(x).view(batch_size, self.key_channels, -1)
        query = query.permute(0, 2, 1)
        key = self.f_object(proxy).view(batch_size, self.key_channels, -1)
        value = self.f_down(proxy).view(batch_size, self.key_channels, -1)
        value = value.permute(0, 2, 1)

        sim_map = torch.matmul(query, key)
        sim_map = (self.key_channels**-.5) * sim_map
        sim_map = F.softmax(sim_map, dim=-1) 
        if(self.cls_num ==1):
            sim_map = sim_map[...,0].unsqueeze(2)
            value = value[:,0,...].unsqueeze(1)
        print(sim_map.min())
        print(sim_map.max())
        print(sim_map.size())
        # add bg context ...
        context = torch.matmul(sim_map, value)
        context = context.permute(0, 2, 1).contiguous()
        context = context.view(batch_size, self.key_channels, *x.size()[2:])
        context = self.f_up(context)

        return context


class SpatialOCR_Module(nn.Module):
    """
    Implementation of the OCR module:
    We aggregate the global object representation to update the representation for each pixel.
    """
    def __init__(self, 
                 in_channels, 
                 key_channels, 
                 out_channels, 
                 dropout=0.1,
                 cls_num = 1):
        super(SpatialOCR_Module, self).__init__()
        self.object_context_block = ObjectAttentionBlock2D(in_channels, 
                                                           key_channels,
                                                           cls_num)
        _in_channels = 2 * in_channels

        self.conv_bn_dropout = nn.Sequential(
            nn.Conv2d(_in_channels, out_channels, kernel_size=1, padding=0, bias=False),
            BNReLU(out_channels),
            nn.Dropout2d(dropout)
        )

    def forward(self, feats, proxy_feats):
        context = self.object_context_block(feats, proxy_feats)
        output = self.conv_bn_dropout(torch.cat([context, feats], 1))

        return output

class OCR(nn.Module):
    def __init__(self,
                 in_channels,
                 key_channels=256,
                 out_channels=512,
                 dropout = 0 ,
                 num_cls = 1):
        super().__init__()
        
        
        self.aux_head = nn.Sequential(nn.Conv2d(in_channels = in_channels,
                                                out_channels = in_channels,
                                                kernel_size = 3,
                                                padding = 1),
                                      nn.BatchNorm2d(in_channels),
                                      nn.ReLU(inplace = True),
                                      nn.Conv2d(in_channels = in_channels,
                                                out_channels = num_cls if(num_cls!=1) else 2,
                                                kernel_size = 1),
                                      )
        self.conv3x3_ocr = nn.Sequential(nn.Conv2d(in_channels = in_channels,
                                                   out_channels = out_channels,
                                                   kernel_size = 3,
                                                   padding = 1),
                                         nn.BatchNorm2d(out_channels),
                                         nn.ReLU(inplace = True))
        self.ocr_gather_head = SpatialGather_Module(num_cls)
        
        self.ocr_distri_head = SpatialOCR_Module(out_channels,
                                                 key_channels,
                                                 out_channels,
                                                 dropout,
                                                 num_cls)
    def forward(self,feats):
        
        probs = self.aux_head(feats)
        feats = self.conv3x3_ocr(feats)
        context = self.ocr_gather_head(feats,probs)
        aug_feats = self.ocr_distri_head(feats,context)
        return aug_feats
        
'''
t = torch.rand((1,128,16,16)).cuda()
m = OCR(128)
m.eval()
m.cuda()
y = m(t)
print(y.size())
'''