#!/usr/bin/env python
# coding: utf-8
#
# Author:   Kazuto Nakashima
# URL:      http://kazuto1011.github.io
# Created:  2018-03-26

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


class MSC(nn.Module):
    """
    Multi-scale inputs
    """

    def __init__(self, base, scales=None):
        super(MSC, self).__init__()
        self.base = base
        if scales:
            self.scales = scales
        else:
            self.scales = [0.5, 0.75]

    def forward(self, x):
        # Original
        logits = self.base(x)
        _, _, H, W = logits.shape
        interp = lambda l: F.interpolate(
            l, size=(H, W), mode="bilinear", align_corners=False
        )

        # Scaled
        logits_pyramid = []
        for scale in self.scales:
            h = F.interpolate(x, scale_factor=scale, mode="bilinear", align_corners=False)
            logits_pyramid.append(self.base(h))

        # Pixel-wise max
        logits_all = [logits] + [interp(l) for l in logits_pyramid]
        logits_max = torch.max(torch.stack(logits_all), dim=0)[0]
        logits_confident = torch.min(torch.stack(logits_all), dim=0)[0]
        shapes = [logits.shape] + [out.shape for out in logits_pyramid] + [logits_max.shape]
        outputs = [logits] + logits_pyramid + [logits_confident]
        out = torch.cat([torch.flatten(output,start_dim=2) for output in outputs],dim=2)
        return logits , out , shapes