"""Blue Marble model architecture (Apache-2.0)."""
import torch, torch.nn as nn, torch.nn.functional as F
S2_BANDS = ['B01', 'B02', 'B03', 'B04', 'B05', 'B06', 'B07', 'B08', 'B8A', 'B09', 'B10', 'B11', 'B12']

def conv_bn(i, o, k=3, s=1, g=1, d=1, act=True):
    layers = [nn.Conv2d(i, o, k, s, d * (k // 2), dilation=d, groups=g, bias=False), nn.BatchNorm2d(o)]
    if act:
        layers.append(nn.ReLU6(inplace=True))
    return nn.Sequential(*layers)

class SE(nn.Module):

    def __init__(self, c, r=4):
        super().__init__()
        h = max(8, c // r)
        self.f = nn.Sequential(nn.Conv2d(c, h, 1), nn.ReLU6(inplace=True), nn.Conv2d(h, c, 1), nn.Hardsigmoid())

    def forward(self, x):
        return x * self.f(x.mean((2, 3), keepdim=True))

class IRB(nn.Module):

    def __init__(self, i, o, s=1, e=3, k=5, se=False, d=1):
        super().__init__()
        h = int(round(i * e))
        self.res = s == 1 and i == o
        self.net = nn.Sequential(conv_bn(i, h, 1) if e != 1 else nn.Identity(), conv_bn(h, h, k, s, g=h, d=d), SE(h) if se else nn.Identity(), conv_bn(h, o, 1, act=False))

    def forward(self, x):
        y = self.net(x)
        return x + y if self.res else y

class SpectralStem(nn.Module):

    def __init__(self, c, in_bands=13):
        super().__init__()
        self.spectral = nn.Conv2d(in_bands, c, 1, bias=False)
        self.bn = nn.BatchNorm2d(c)
        self.act = nn.ReLU6(inplace=True)
        self.spatial = conv_bn(c, c, 3, g=c)
        self.fuse = conv_bn(c, c, 1)

    def forward(self, x):
        return self.fuse(self.spatial(self.act(self.bn(self.spectral(x)))))

    @torch.no_grad()
    def slice_bands(self, idx):
        new = nn.Conv2d(len(idx), self.spectral.out_channels, 1, bias=False).to(self.spectral.weight.device)
        new.weight.copy_(self.spectral.weight[:, idx])
        self.spectral = new

class Attn2d(nn.Module):

    def __init__(self, c, heads=4, mlp=2):
        super().__init__()
        self.h, self.d = (heads, c // heads)
        self.n1, self.q, self.k, self.v, self.o = (nn.BatchNorm2d(c), nn.Conv2d(c, c, 1), nn.Conv2d(c, c, 1), nn.Conv2d(c, c, 1), nn.Conv2d(c, c, 1))
        self.mlp = nn.Sequential(nn.BatchNorm2d(c), nn.Conv2d(c, c * mlp, 1), nn.ReLU6(inplace=True), nn.Conv2d(c * mlp, c, 1))

    def forward(self, x):
        b, c, h, w = x.shape
        n = self.n1(x)
        sp = lambda t: t.reshape(b, self.h, self.d, h * w).transpose(2, 3)
        q, k, v = (sp(self.q(n)), sp(self.k(n)), sp(self.v(n)))
        a = (q.float() @ k.float().transpose(2, 3) / self.d ** 0.5).softmax(-1)
        y = (a @ v.float()).to(x.dtype).transpose(2, 3).reshape(b, c, h, w)
        x = x + self.o(y)
        return x + self.mlp(x)
ARCHS = {'pico': dict(stem=8, w=(8, 16, 24, 40), d=(1, 2, 2, 2), e=2, k=3, rates=(1,), dec=16, full=False, se=False), 'nano': dict(stem=12, w=(16, 24, 48, 80), d=(1, 2, 3, 2), e=3, k=5, rates=(1, 2), dec=24, full=True, se=False), 'tiny': dict(stem=16, w=(24, 40, 80, 160), d=(2, 2, 4, 3), e=3, k=5, rates=(1, 2, 4), dec=40, full=True, se=True), 'small': dict(stem=24, w=(32, 56, 112, 208), d=(2, 3, 5, 3), e=3, k=7, rates=(1, 2, 4), dec=48, full=True, se=True), 'medium': dict(stem=32, w=(48, 96, 192, 320), d=(2, 3, 8, 4), e=4, k=7, rates=(1, 2, 4), dec=96, full=True, se=True, attn=2)}

class Encoder(nn.Module):

    def __init__(self, cfg, in_bands=13):
        super().__init__()
        self.stem = SpectralStem(cfg['stem'], in_bands)
        prev, self.stages = (cfg['stem'], nn.ModuleList())
        for si, (w, n) in enumerate(zip(cfg['w'], cfg['d'])):
            blocks = [IRB(prev, w, 2, cfg['e'], 3)]
            blocks += [IRB(w, w, 1, cfg['e'], cfg['k'], se=cfg['se'] and si >= 2) for _ in range(n - 1)]
            if si == len(cfg['w']) - 1:
                blocks += [Attn2d(w, heads=8) for _ in range(cfg.get('attn', 0))]
            self.stages.append(nn.Sequential(*blocks))
            prev = w
        self.neck = nn.ModuleList((IRB(prev, prev, 1, 2, 3, d=r) for r in cfg['rates']))
        self.channels = (cfg['stem'], *cfg['w'])

    def forward(self, x):
        out = [self.stem(x)]
        for s in self.stages:
            out.append(s(out[-1]))
        y = out[-1]
        for n in self.neck:
            y = n(y)
        out[-1] = y
        return out

class FPNDecoder(nn.Module):

    def __init__(self, chs, dec, full=True):
        super().__init__()
        self.full = full
        self.lat = nn.ModuleList((conv_bn(c, dec, 1) for c in chs[1:]))
        self.ref = nn.ModuleList((IRB(dec, dec, 1, 2, 3) for _ in chs[1:-1]))
        if full:
            self.lat0 = conv_bn(chs[0], dec, 1)
            self.ref0 = IRB(dec, dec, 1, 2, 3)
        self.out_channels = dec

    def forward(self, f):
        y = self.lat[-1](f[-1])
        for i in range(len(self.ref) - 1, -1, -1):
            y = F.interpolate(y, scale_factor=2.0, mode='bilinear', align_corners=False) + self.lat[i](f[i + 1])
            y = self.ref[i](y)
        if self.full:
            y = F.interpolate(y, scale_factor=2.0, mode='bilinear', align_corners=False) + self.lat0(f[0])
            y = self.ref0(y)
        return y

class BlueMarble(nn.Module):

    def __init__(self, arch='tiny', num_classes=2, in_bands=13):
        super().__init__()
        cfg = ARCHS[arch]
        self.arch, self.upsample = (arch, 1 if cfg['full'] else 2)
        self.encoder = Encoder(cfg, in_bands)
        self.decoder = FPNDecoder(self.encoder.channels, cfg['dec'], cfg['full'])
        self.head = nn.Sequential(nn.Dropout2d(0.1), nn.Conv2d(cfg['dec'], num_classes, 1))

    def forward(self, x):
        y = self.head(self.decoder(self.encoder(x)))
        if self.upsample > 1:
            y = F.interpolate(y, scale_factor=float(self.upsample), mode='bilinear', align_corners=False)
        return y

def fast_layout(model, in_shape, device='cuda'):
    import torch.nn.functional as F
    if not str(device).startswith('cuda'):
        return (model, False)
    try:
        m = model.to(memory_format=torch.channels_last)
        x = torch.randn(in_shape, device=device).to(memory_format=torch.channels_last)
        with torch.autocast('cuda', dtype=torch.bfloat16):
            out = m(x)
        out.float().mean().backward()
        m.zero_grad(set_to_none=True)
        return (m, True)
    except RuntimeError as e:
        print('channels_last unavailable, using NCHW:', str(e)[:120], flush=True)
        return (model.to(memory_format=torch.contiguous_format), False)