import torch
import torch.nn as nn
import torchvision.transforms as T
from PIL import Image
import warnings

warnings.filterwarnings("ignore")

# -----------------------------
# Config
# -----------------------------
MODEL_PATH = "ocr_model.pth"
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"

IMG_HEIGHT = 30
IMG_WIDTH = 100


# -----------------------------
# Model (same as training)
# -----------------------------
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


# -----------------------------
# Load model
# -----------------------------
def load_model():
    checkpoint = torch.load(MODEL_PATH, map_location=DEVICE)

    charset = checkpoint["charset"]
    idx2char = {i + 1: c for i, c in enumerate(charset)}

    model = OCRModel(num_classes=len(charset))
    model.load_state_dict(checkpoint["model"])
    model.to(DEVICE)
    model.eval()

    return model, idx2char


# -----------------------------
# Image preprocessing
# -----------------------------
transform = T.Compose([
    T.Grayscale(),
    T.Resize((IMG_HEIGHT, IMG_WIDTH)),
    T.ToTensor(),
    T.Normalize((0.5,), (0.5,))
])


def predict(image_path):
    model, idx2char = load_model()

    img = Image.open(image_path)
    img = transform(img)
    img = img.unsqueeze(0).to(DEVICE)

    with torch.no_grad():
        output = model(img)
        text = decode(output, idx2char)[0]

    return text


# -----------------------------
# Run
# -----------------------------
if __name__ == "__main__":
    import sys

    if len(sys.argv) < 2:
        sys.exit("Usage: python3 infer.py image.png")

    img_path = sys.argv[1]
    result = predict(img_path)

    print(f"{result}")
