"""Page orientation classifier: 0 / 90 / 180 / 270 degrees, voted over text-dense patches."""

from torch import nn

from .common import ConvBNAct

PATCH = 192
ANGLES = (0, 90, 180, 270)  # clockwise rotation the page has undergone


class OrientationNet(nn.Module):
    def __init__(self, widths=(16, 32, 64, 88, 120)):
        super().__init__()
        layers, cin = [], 1
        for i, w in enumerate(widths):
            layers += [ConvBNAct(cin, w, s=2), ConvBNAct(w, w, k=3, groups=w), ConvBNAct(w, w, k=1)]
            cin = w
        self.features = nn.Sequential(*layers)
        self.pool = nn.AdaptiveAvgPool2d(1)
        self.fc = nn.Sequential(nn.Dropout(0.2), nn.Linear(cin, len(ANGLES)))

    def forward(self, x):
        """x: (B, 1, 192, 192) in [0, 1]. Returns logits (B, 4)."""
        return self.fc(self.pool(self.features(x - 0.5)).flatten(1))