import os
import random
import numpy as np
from PIL import Image

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import Dataset, DataLoader
import torchvision.transforms as T

# -----------------------------
# Config
# -----------------------------
IMG_DIR = "images"
BATCH_SIZE = 64
EPOCHS = 100
LR = 3e-4
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"

IMG_HEIGHT = 30
IMG_WIDTH = 100

# -----------------------------
# Dataset
# -----------------------------
class OCRDataset(Dataset):
    def __init__(self, img_dir):
        self.img_dir = img_dir
        self.files = [f for f in os.listdir(img_dir) if f.endswith(".png")]
        random.shuffle(self.files)

        charset = set()
        for f in self.files:
            label = os.path.splitext(f)[0]
            charset.update(label)

        self.charset = sorted(list(charset))
        self.char2idx = {c: i + 1 for i, c in enumerate(self.charset)}
        self.idx2char = {i: c for c, i in self.char2idx.items()}

        self.transform = T.Compose([
            T.Grayscale(),
            T.Resize((IMG_HEIGHT, IMG_WIDTH)),

            T.RandomApply([
                T.RandomRotation(3),
                T.RandomAffine(degrees=0, shear=5, scale=(0.9, 1.1))
            ], p=0.7),

            T.ColorJitter(brightness=0.3, contrast=0.3),
            T.GaussianBlur(kernel_size=3, sigma=(0.1, 1.5)),

            T.ToTensor(),
            T.Normalize((0.5,), (0.5,))
        ])

    def __len__(self):
        return len(self.files)

    def __getitem__(self, idx):
        fname = self.files[idx]
        label = os.path.splitext(fname)[0]

        img_path = os.path.join(self.img_dir, fname)
        img = Image.open(img_path)
        img = self.transform(img)

        label_encoded = torch.tensor(
            [self.char2idx[c] for c in label], dtype=torch.long
        )

        return img, label_encoded, len(label_encoded)


def collate_fn(batch):
    imgs, labels, lengths = zip(*batch)

    imgs = torch.stack(imgs)
    labels_concat = torch.cat(labels)
    label_lengths = torch.tensor(lengths, dtype=torch.long)

    return imgs, labels_concat, label_lengths

# -----------------------------
# Model
# -----------------------------
class OCRModel(nn.Module):
    def __init__(self, num_classes):
        super().__init__()

        self.cnn = nn.Sequential(
            nn.Conv2d(1, 32, 3, padding=1),
            nn.BatchNorm2d(32),
            nn.ReLU(),
            nn.MaxPool2d((2, 2)),
            nn.Dropout2d(0.1),

            nn.Conv2d(32, 64, 3, padding=1),
            nn.BatchNorm2d(64),
            nn.ReLU(),
            nn.MaxPool2d((2, 1)),
            nn.Dropout2d(0.1),
        )

        with torch.no_grad():
            dummy = torch.zeros(1, 1, IMG_HEIGHT, IMG_WIDTH)
            out = self.cnn(dummy)
            _, c, h, w = out.shape
            self.rnn_input_size = c * h

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

        self.fc = nn.Linear(768, num_classes + 1)

    def forward(self, x):
        x = self.cnn(x)
        b, c, h, w = x.size()

        x = x.permute(0, 3, 1, 2)
        x = x.contiguous().view(b, w, c * h)

        x, _ = self.lstm(x)
        x = self.fc(x)
        x = x.log_softmax(2)

        x = x.permute(1, 0, 2)
        return x

# -----------------------------
# Decoder
# -----------------------------
def decode(output, idx2char):
    output = output.argmax(2)
    output = output.permute(1, 0)

    results = []
    for seq in output:
        prev = -1
        s = ""
        for i in seq:
            i = i.item()
            if i != prev and i != 0:
                s += idx2char.get(i, "")
            prev = i
        results.append(s)
    return results

# -----------------------------
# Training
# -----------------------------
def train():
    dataset = OCRDataset(IMG_DIR)

    dataloader = DataLoader(
        dataset,
        batch_size=BATCH_SIZE,
        shuffle=True,
        collate_fn=collate_fn
    )

    model = OCRModel(num_classes=len(dataset.charset)).to(DEVICE)

    criterion = nn.CTCLoss(blank=0, zero_infinity=True)

    optimizer = optim.Adam(
        model.parameters(),
        lr=LR,
        weight_decay=1e-5
    )

    scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(
        optimizer,
        mode='min',
        factor=0.5,
        patience=2
    )

    best_loss = float('inf')

    for epoch in range(EPOCHS):
        model.train()
        total_loss = 0

        for imgs, labels, label_lengths in dataloader:
            imgs = imgs.to(DEVICE)
            labels = labels.to(DEVICE)
            label_lengths = label_lengths.to(DEVICE)

            outputs = model(imgs)

            input_lengths = torch.full(
                size=(outputs.size(1),),
                fill_value=outputs.size(0),
                dtype=torch.long
            ).to(DEVICE)

            loss = criterion(outputs, labels, input_lengths, label_lengths)

            optimizer.zero_grad()
            loss.backward()
            optimizer.step()

            total_loss += loss.item()

        scheduler.step(total_loss)

        print(f"\nEpoch {epoch+1}/{EPOCHS}, Loss: {total_loss:.4f}")

        if total_loss < best_loss:
            best_loss = total_loss
            print(f"Saving model with {best_loss} loss")
            torch.save({
                "model": model.state_dict(),
                "charset": dataset.charset
            }, "best_ocr_model.pth")

        if epoch % 2 == 0:
            model.eval()
            with torch.no_grad():
                sample_imgs, _, _ = next(iter(dataloader))
                sample_imgs = sample_imgs.to(DEVICE)

                preds = model(sample_imgs)
                texts = decode(preds, dataset.idx2char)

                print("Sample predictions:")
                for t in texts[:5]:
                    print("  ", t)


if __name__ == "__main__":
    train()
