"""DPT-style depth decoders that consume spatial feature maps.

Both heads take a list of 4 spatial feature maps (FLUX + DINO + concepts) with
possibly different channel counts, run the shared DPT RefineNet cascade, and
upsample to a depth map:
  - DPTHeadSpatial : single 2x upsample in the head (bilinear resize to GT after).
  - DPTHeadHighRes : 4x learned 2x upsampling (16x) with optional image skip connections.
"""

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

from .dpt_backbone import DPTRefineNetStack


class DPTHeadSpatial(DPTRefineNetStack):
    def __init__(self, in_channels=[3840, 3840, 3840, 3840], features=256,
                 num_classes=1, use_bn=False, head_features=32):
        super().__init__(features=features, use_bn=use_bn)
        self.head_features = head_features
        self.num_classes = num_classes

        out_channels = [256, 512, 1024, 1024]
        self.projects = nn.ModuleList([nn.Conv2d(in_ch, oc, 1) for in_ch, oc in zip(in_channels, out_channels)])
        self.pre_fuse = nn.ModuleList([
            nn.Sequential(nn.GroupNorm(1, oc), nn.Conv2d(oc, oc, 1), nn.ReLU(inplace=True))
            for oc in out_channels])

        self.head = nn.Sequential(
            nn.Conv2d(features, features // 2, 3, padding=1),
            nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True),
            nn.Conv2d(features // 2, head_features, 3, padding=1),
            nn.ReLU(inplace=True),
            nn.Conv2d(head_features, num_classes, 1),
        )

    def forward(self, features):
        resized = [pf(proj(f)) for f, proj, pf in zip(features, self.projects, self.pre_fuse)]
        return self.head(self.fuse(resized))


class DPTHeadHighRes(DPTRefineNetStack):
    """High-resolution variant: progressive 4x (2x) learned upsampling = 16x total,
    with optional skip connections from the original image."""
    def __init__(self, in_channels=[3840, 3840, 3840, 3840], features=256,
                 num_classes=1, use_bn=False, head_features=64, use_skip_connections=True):
        super().__init__(features=features, use_bn=use_bn)
        self.head_features = head_features
        self.num_classes = num_classes
        self.use_skip_connections = use_skip_connections

        out_channels = [256, 512, 1024, 1024]
        self.projects = nn.ModuleList([nn.Conv2d(in_ch, oc, 1) for in_ch, oc in zip(in_channels, out_channels)])
        self.pre_fuse = nn.ModuleList([
            nn.Sequential(nn.GroupNorm(1, oc), nn.Conv2d(oc, oc, 1), nn.ReLU(inplace=True))
            for oc in out_channels])

        if use_skip_connections:
            def skip_enc(stride, out_ch):
                return nn.Sequential(
                    nn.Conv2d(3, out_ch, 3, stride=stride, padding=1), nn.ReLU(inplace=True),
                    nn.Conv2d(out_ch, out_ch, 3, padding=1), nn.ReLU(inplace=True))
            self.skip_enc_2x = skip_enc(2, 16)
            self.skip_enc_4x = skip_enc(4, 16)
            self.skip_enc_8x = skip_enc(8, 16)
            self.skip_enc_full = skip_enc(1, 8)

        def up_block(c_in, c_out):
            return nn.Sequential(
                nn.Conv2d(c_in, c_out, 3, padding=1), nn.ReLU(inplace=True),
                nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True),
                nn.Conv2d(c_out, c_out, 3, padding=1), nn.ReLU(inplace=True))

        skip = 16 if use_skip_connections else 0
        self.up1 = up_block(features, features // 2)
        self.fuse1 = nn.Conv2d(features // 2 + skip, features // 2, 1) if use_skip_connections else None
        self.up2 = up_block(features // 2, features // 4)
        self.fuse2 = nn.Conv2d(features // 4 + skip, features // 4, 1) if use_skip_connections else None
        self.up3 = up_block(features // 4, head_features)
        self.fuse3 = nn.Conv2d(head_features + skip, head_features, 1) if use_skip_connections else None
        self.up4 = up_block(head_features, head_features // 2)
        self.fuse4 = nn.Conv2d(head_features // 2 + (8 if use_skip_connections else 0), head_features // 2, 1) if use_skip_connections else None
        self.output = nn.Conv2d(head_features // 2, num_classes, 3, padding=1)

    def _apply_skip(self, x, fuse, skip):
        if not self.use_skip_connections or skip is None:
            return x
        if skip.shape[-2:] != x.shape[-2:]:
            skip = F.interpolate(skip, size=x.shape[-2:], mode='bilinear', align_corners=True)
        return fuse(torch.cat([x, skip], dim=1))

    def forward(self, features, image=None):
        if self.use_skip_connections and image is not None:
            skip_8x, skip_4x = self.skip_enc_8x(image), self.skip_enc_4x(image)
            skip_2x, skip_full = self.skip_enc_2x(image), self.skip_enc_full(image)
        else:
            skip_8x = skip_4x = skip_2x = skip_full = None

        resized = [pf(proj(f)) for f, proj, pf in zip(features, self.projects, self.pre_fuse)]
        path_1 = self.fuse(resized, keep_layer1_size=True)  # keep patch resolution

        x = self._apply_skip(self.up1(path_1), self.fuse1, skip_8x)
        x = self._apply_skip(self.up2(x), self.fuse2, skip_4x)
        x = self._apply_skip(self.up3(x), self.fuse3, skip_2x)
        x = self._apply_skip(self.up4(x), self.fuse4, skip_full)
        return self.output(x)