import torch
import torch.nn as nn
from .fourier_features import FourierFeatures

class RegionModel(nn.Module):
    def __init__(self):
        super().__init__()

        self.position_features = FourierFeatures(2, 256)
        self.position_encoder = nn.Linear(256, 2048)
        self.size_features = FourierFeatures(2, 256)
        self.size_encoder = nn.Linear(256, 2048)

        self.position_decoder = nn.Linear(2048, 2)
        self.size_decoder = nn.Linear(2048, 2)
        self.confidence_decoder = nn.Linear(2048, 1)

    def encode_position(self, position):
        return self.position_encoder(self.position_features(position))

    def encode_size(self, size):
        return self.size_encoder(self.size_features(size))

    def decode_position(self, x):
        return self.position_decoder(x)

    def decode_size(self, x):
        return self.size_decoder(x)

    def decode_confidence(self, x):
        return self.confidence_decoder(x)

    def encode(self, position, size):
        return torch.stack(
            [self.encode_position(position), self.encode_size(size)], dim=0
        )

    def decode(self, position_logits, size_logits):
        return (
            self.decode_position(position_logits),
            self.decode_size(size_logits),
            self.decode_confidence(size_logits),
        )