"""Shared DPT reassemble + RefineNet cascade.

`DPTRefineNetStack` owns the `scratch` reassemble layers and the four
`FeatureFusionBlock` refinenets used by every DPT decoder in this repo (the depth
heads in ``dpt_decoder.py`` and the saliency decoder in
``dpt_segmentation_decoder.py``). Decoders subclass it so the parameter names stay
flat (``scratch.*`` / ``refinenet{1..4}.*``) and existing checkpoints keep loading;
each subclass provides its own input projection and output head.
"""

import torch.nn as nn

from .blocks import FeatureFusionBlock, _make_scratch


class DPTRefineNetStack(nn.Module):
    def __init__(self, features=256, use_bn=False, out_channels=(256, 512, 1024, 1024)):
        super().__init__()
        self.features = features
        self.scratch = _make_scratch(list(out_channels), features)
        self.refinenet4 = FeatureFusionBlock(features, nn.ReLU(inplace=False), bn=use_bn)
        self.refinenet3 = FeatureFusionBlock(features, nn.ReLU(inplace=False), bn=use_bn)
        self.refinenet2 = FeatureFusionBlock(features, nn.ReLU(inplace=False), bn=use_bn)
        self.refinenet1 = FeatureFusionBlock(features, nn.ReLU(inplace=False), bn=use_bn)

    def fuse(self, layers, keep_layer1_size=False):
        """Run the coarse-to-fine RefineNet cascade; returns the layer-1 feature map.

        ``keep_layer1_size=True`` stops the final block from doing its default 2x
        upsample (used by the high-res depth decoder, which upsamples in its head).
        """
        l1, l2, l3, l4 = layers
        l1 = self.scratch.layer1_rn(l1)
        l2 = self.scratch.layer2_rn(l2)
        l3 = self.scratch.layer3_rn(l3)
        l4 = self.scratch.layer4_rn(l4)
        path = self.refinenet4(l4, size=l3.shape[2:])
        path = self.refinenet3(path, l3, size=l2.shape[2:])
        path = self.refinenet2(path, l2, size=l1.shape[2:])
        if keep_layer1_size:
            path = self.refinenet1(path, l1, size=l1.shape[2:])
        else:
            path = self.refinenet1(path, l1)
        return path