from __future__ import annotations

import math
import pickle
from functools import lru_cache
from pathlib import Path
from typing import Any

import numpy as np
import torch
import torch.nn as nn
from PIL import Image

IMAGE_HEIGHT = 64
IMAGE_WIDTH = 2048
HORIZONTAL_STRIDE = 4
RNN_HIDDEN = 256
RNN_LAYERS = 2
RNN_DROPOUT = 0.05
CHECKPOINT = Path(__file__).with_name("best-model.pt")
LANGUAGE_MODEL = Path(__file__).with_name("char-trigram-lm.pkl")
BEAM_CONFIG = {
    "beam_width": 10,
    "token_topk": 12,
    "lm_weight": 0.4,
    "token_bonus": 1.5,
}


class TemporalFeatureDropout(nn.Module):
    def __init__(self, probability: float = 0.0):
        super().__init__()
        self.p = float(probability)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        if not self.training or self.p == 0.0:
            return x
        keep = x.new_empty((x.shape[0], x.shape[1], 1)).bernoulli_(1.0 - self.p)
        return x * keep / (1.0 - self.p)


class KoshurCRNN(nn.Module):
    def __init__(self, n_classes: int, temporal_dropout_p: float = 0.0):
        super().__init__()
        self.cnn = nn.Sequential(
            nn.Conv2d(1, 48, 3, padding=1), nn.BatchNorm2d(48), nn.ReLU(), nn.MaxPool2d((2, 2)),
            nn.Conv2d(48, 96, 3, padding=1), nn.BatchNorm2d(96), nn.ReLU(), nn.MaxPool2d((2, 2)),
            nn.Conv2d(96, 160, 3, padding=1), nn.BatchNorm2d(160), nn.ReLU(), nn.MaxPool2d((2, 1)),
            nn.Conv2d(160, 192, 3, padding=1), nn.BatchNorm2d(192), nn.ReLU(),
        )
        self.temporal_dropout = TemporalFeatureDropout(temporal_dropout_p)
        self.rnn = nn.GRU(
            192 * (IMAGE_HEIGHT // 8), RNN_HIDDEN, num_layers=RNN_LAYERS,
            bidirectional=True, batch_first=True, dropout=RNN_DROPOUT,
        )
        self.head = nn.Linear(RNN_HIDDEN * 2, n_classes)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        z = self.cnn(x)
        z = z.permute(0, 3, 1, 2).contiguous().flatten(2)
        z = self.temporal_dropout(z)
        z, _ = self.rnn(z)
        return self.head(z).log_softmax(-1).permute(1, 0, 2)


@lru_cache(maxsize=1)
def load_model() -> tuple[KoshurCRNN, dict[int, str]]:
    checkpoint = torch.load(CHECKPOINT, map_location="cpu", weights_only=False)
    chars = list(checkpoint["chars"])
    architecture = checkpoint.get("architecture") or {}
    temporal_dropout_p = float(architecture.get("temporal_dropout_p", 0.0))
    model = KoshurCRNN(len(chars) + 1, temporal_dropout_p=temporal_dropout_p)
    model.load_state_dict(checkpoint["model_state_dict"], strict=True)
    model.eval()
    return model, {i + 1: char for i, char in enumerate(chars)}


@lru_cache(maxsize=1)
def load_language_model():
    # Trusted release artifact containing only fitted character n-gram counts.
    with LANGUAGE_MODEL.open("rb") as handle:
        return pickle.load(handle)


def preprocess(image: Image.Image) -> tuple[torch.Tensor, int, int]:
    if image is None:
        raise ValueError("Upload a cropped image containing one Kashmiri text line.")
    normal = image.convert("L")
    scale = IMAGE_HEIGHT / max(1, normal.height)
    resized_width = max(1, min(IMAGE_WIDTH, int(normal.width * scale)))
    normal = normal.resize((resized_width, IMAGE_HEIGHT), Image.Resampling.BICUBIC)
    canvas = Image.new("L", (IMAGE_WIDTH, IMAGE_HEIGHT), 255)
    canvas.paste(normal, (0, 0))
    pixels = 1.0 - np.asarray(canvas, dtype=np.float32) / 255.0
    input_length = max(
        1,
        min(IMAGE_WIDTH // HORIZONTAL_STRIDE, int(math.ceil(resized_width / HORIZONTAL_STRIDE))),
    )
    return torch.from_numpy(pixels).unsqueeze(0).unsqueeze(0), input_length, resized_width


def ctc_decode(ids: list[int], itos: dict[int, str]) -> str:
    output: list[str] = []
    previous = None
    for token in ids:
        if token != 0 and token != previous:
            output.append(itos.get(token, ""))
        previous = token
    return "".join(output)[::-1]


def recognize_line(image: Image.Image, *, decoder: str = "beam") -> tuple[str, dict[str, Any]]:
    if decoder not in {"beam", "greedy"}:
        raise ValueError("decoder must be 'beam' or 'greedy'")
    model, itos = load_model()
    tensor, input_length, resized_width = preprocess(image)
    with torch.inference_mode():
        logits = model(tensor)
    matrix = logits[:input_length, 0]
    diagnostics: dict[str, Any] = {
        "input_width": image.width,
        "input_height": image.height,
        "resized_width": resized_width,
        "ctc_frames": input_length,
        "decoder": decoder,
    }
    if decoder == "greedy":
        text = ctc_decode(matrix.argmax(-1).tolist(), itos)
    else:
        from ctc_prefix_beam import prefix_beam_search

        ids = prefix_beam_search(
            matrix.cpu().numpy(),
            lm=load_language_model(),
            **BEAM_CONFIG,
        )
        # Prefix-beam output is already CTC-collapsed in increasing-x order.
        text = "".join(itos[token] for token in ids)[::-1]
        diagnostics.update(BEAM_CONFIG)
    return text, diagnostics