import torch
import torch.nn as nn


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

        self.features = nn.Sequential(
            nn.Conv2d(1, 64, kernel_size=3, padding=1),
            nn.BatchNorm2d(64),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=2, stride=2),
            
            nn.Conv2d(64, 128, kernel_size=3, padding=1),
            nn.BatchNorm2d(128),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=2, stride=2),

            nn.Conv2d(128, 256, kernel_size=3, padding=1),
            nn.BatchNorm2d(256),
            nn.ReLU(inplace=True),

            nn.Conv2d(256, 256, kernel_size=3, padding=1),
            nn.BatchNorm2d(256),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=(2, 1)),

            nn.Conv2d(256, 512, kernel_size=3, padding=1),
            nn.BatchNorm2d(512),
            nn.ReLU(inplace=True),

            nn.Conv2d(512, 512, kernel_size=3, padding=1),
            nn.BatchNorm2d(512),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=(2, 1)),

            nn.Conv2d(512, 512, kernel_size=(2, 1), padding=0),
            nn.BatchNorm2d(512),
            nn.ReLU(inplace=True),
        )

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


class CRNN(nn.Module):
    def __init__(
        self,
        num_classes: int = 37,   # 36 символов + бланк
        hidden_size: int = 256,
        num_layers: int = 2,
    ):
        super().__init__()

        self.cnn = CNN()

        self.lstm = nn.LSTM(
            input_size=512,
            hidden_size=hidden_size,
            num_layers=num_layers,
            batch_first=True,
            bidirectional=True,
            dropout=0.3,
        )

        self.classifier = nn.Linear(hidden_size * 2, num_classes)

    def forward(self, x: torch.Tensor) -> torch.Tensor:

        features = self.cnn(x)
        features = features.squeeze(2)
        features = features.permute(0, 2, 1)
        lstm_out, _ = self.lstm(features)
        logits = self.classifier(lstm_out)
        logits = logits.permute(1, 0, 2)

        return logits


def get_model(num_classes: int = 37, device: str = "cpu") -> CRNN:
    model = CRNN(num_classes=num_classes)
    model = model.to(device)
    return model


def count_parameters(model: nn.Module) -> int:
    return sum(p.numel() for p in model.parameters() if p.requires_grad)


# ── Запуск для проверки ─────────────────────────────────────────────────────

if __name__ == "__main__":
    print("=== Тест model.py ===\n")

    device = "cuda" if torch.cuda.is_available() else "cpu"
    print(f"Устройство: {device}")

    # создаём модель
    model = get_model(num_classes=37, device=device)

    # считаем параметры
    params = count_parameters(model)
    print(f"Параметров модели: {params:,}")  # написало примерно 8 млн

    # тестовый проход
    batch_size = 4
    dummy_input = torch.randn(batch_size, 1, 32, 128).to(device)
    print(f"\nВход:  {list(dummy_input.shape)}")

    with torch.no_grad():
        output = model(dummy_input)

    print(f"Выход: {list(output.shape)}")  # [32, 4, 37]
    print(f"  T=32  — временных шагов (позиций символов)")
    print(f"  B=4   — размер батча")
    print(f"  C=37  — классов (символов)")

    # Проверка что ctc loss работает
    print("\nПроверка CTC Loss...")
    ctc_loss = nn.CTCLoss(blank=0, reduction="mean", zero_infinity=True)

    log_probs = output.log_softmax(2)
    input_lengths = torch.full((batch_size,), 32, dtype=torch.long)
    targets = torch.tensor([1, 2, 3, 4, 5, 6], dtype=torch.long)
    target_lengths = torch.tensor([3, 3], dtype=torch.long)

    # Только для 2 примеров из батча
    loss = ctc_loss(
        log_probs[:, :2, :],
        targets,
        input_lengths[:2],
        target_lengths,
    )
    print(f"CTC Loss: {loss.item():.4f}")

    print("\n✓ еее работает")