import torch
import torch.nn as nn

from model.encoder import Encoder
from model.decoder import Decoder

from model.utils import weight_init


class Trainer(nn.Module):
    def __init__(self, model_type='small'):
        super().__init__()
        if model_type == 'tiny':
            embed_dim = 192
        elif model_type == 'small':
            embed_dim = 384
        else:
            assert False, r'Trainer: check the vit model type'

        self.encoder = Encoder(model_type)

        self.decoder = Decoder(in_dim=[64, 128, 256, embed_dim])
        weight_init(self.decoder)
        
    def forward(self, x, y):
        fx, fy = self.encoder(x, y)
        pred = self.decoder(fx, fy)

        return pred