# AUTOGENERATED! DO NOT EDIT! File to edit: nbs/00_model.ipynb (unless otherwise specified).

__all__ = ['SpatialDownsampler', 'Upsampler', 'IrradianceModule', 'Eclipse']

# Cell
import torch
import torch.nn as nn
import torch.nn.functional as F

from .layers import *

# Cell
class SpatialDownsampler(nn.Module):

    def __init__(self, in_channels=3):
        super().__init__()
        self.conv1 = ConvBlock(in_channels, 64, kernel_size=7, stride=1)
        self.blocks = nn.Sequential(ResBlock(64, 64, kernel_size=3, stride=2),
                                    ResBlock(64, 128, kernel_size=3, stride=2),
                                    ResBlock(128,256, kernel_size=3, stride=2))

    def forward(self, x):
        return self.blocks(self.conv1(x))

# Cell
class Upsampler(nn.Module):
    def __init__(self, sizes=[128,128,64], n_out=3):
        super().__init__()
        zsizes = zip(sizes[:-1], sizes[1:])
        self.convs = nn.Sequential(*[Bottleneck(si, sf, upsample=True) for si,sf in zsizes],
                                   Bottleneck(sizes[-1], sizes[-1], upsample=True),
                                   ConvBlock(sizes[-1], n_out, kernel_size=1, activation='none'))

    def forward(self, x):
        return self.convs(x)

# Cell
class IrradianceModule(nn.Module):
    def __init__(self):
        super().__init__()
        self.convs = nn.Sequential(ConvBlock(128, 64),
                                   ConvBlock(64, 64),
                                   nn.AdaptiveMaxPool2d(1)
                                  )
        self.linear = nn.Sequential(nn.Flatten(),
                                    nn.BatchNorm1d(64),
                                    nn.Linear(64, 1)
                                   )
    def forward(self, x):
        return self.linear(self.convs(x))

# Cell
class Eclipse(nn.Module):

    def __init__(self, n_in=3, horizon=5):
        super().__init__()
        self.horizon = horizon
        self.spatial_downsampler = SpatialDownsampler(n_in)
        self.temporal_model = TemporalBlock(256, 128)
        self.future_prediction = FuturePrediction(128, 128, n_gru_blocks=4, n_res_layers=4)
        self.upsampler = Upsampler()
        self.irradiance = IrradianceModule()

    def zero_hidden(self, x, horizon):
        bs, ch, h, w = x.shape
        return x.new_zeros(bs, horizon, ch, h, w)

    def forward(self, imgs):
        x = torch.stack([self.spatial_downsampler(img) for img in imgs], dim=2)

        #encode temporal model
        states = self.temporal_model(x).permute(0, 2, 1, 3, 4).contiguous()
        print(f'{states.shape=}')

        #get hidden state
        present_state = states[:, :1].contiguous()
        print(f'{present_state.shape=}')


        # Prepare future prediction input
        hidden_state = present_state[:, 0]

        print(f'{hidden_state.shape=}')
        future_prediction_input = self.zero_hidden(hidden_state, self.horizon)

        # Recursively predict future states
        future_states = self.future_prediction(future_prediction_input, hidden_state)

        # Concatenate present state
        future_states = torch.cat([present_state, future_states], dim=1)
        print(f'{future_states.shape=}')

        #decode outputs
        preds = {'masks': [], 'irradiance': []}
        for state in future_states.unbind(dim=1):
            preds['masks'].append(self.upsampler(state))
            preds['irradiance'].append(self.irradiance(state))
        return preds
