"""
captCHAD: Neural network for CAPTCHA optical character recognition.
Input: (B, 3, 64, 192) RGB images
Output: (T=48, B, 63) Logits for CTC loss (62 alphanumeric chars + 1 blank token)
Supported Length: 1 to 8 characters per image (optimized for 4–7 characters)
"""
import math
import string
from collections import defaultdict
import torch
import torch.nn as nn
import torch.nn.functional as F

# Complete 62-character alphanumeric set (digits + uppercase + lowercase)
CHARSET = string.digits + string.ascii_letters
BLANK_IDX = 0
NUM_CLASSES = len(CHARSET) + 1  # 63 classes (index 0 is CTC blank)

CHAR2IDX = {ch: i + 1 for i, ch in enumerate(CHARSET)}
IDX2CHAR = {i + 1: ch for i, ch in enumerate(CHARSET)}


def encode_string(text: str) -> list[int]:
    """Convert text string into list of token IDs (1-indexed)."""
    return [CHAR2IDX[c] for c in text if c in CHAR2IDX]


def decode_tokens(tokens: list[int]) -> str:
    """Decode token IDs into string, collapsing consecutive duplicates and stripping blanks."""
    res = []
    prev = 0
    for t in tokens:
        if t != BLANK_IDX:
            if t != prev:
                res.append(IDX2CHAR.get(t, ''))
        prev = t
    return "".join(res)


def decode_beam_search_single(log_probs: torch.Tensor, beam_width: int = 15) -> list[tuple[str, float]]:
    """
    Perform CTC Beam Search decoding on a single sequence of log probabilities.
    log_probs: (T, C) tensor
    Returns list of (decoded_text, log_score) sorted from highest to lowest score.
    """
    T, C = log_probs.shape
    beam = {(): (0.0, -float('inf'))}

    def logaddexp(a, b):
        if a == -float('inf'):
            return b
        if b == -float('inf'):
            return a
        m = max(a, b)
        return m + math.log(1.0 + math.exp(-abs(a - b)))

    for t in range(T):
        curr_beam = defaultdict(lambda: (-float('inf'), -float('inf')))
        lp = log_probs[t]
        topk_probs, topk_indices = torch.topk(lp, min(C, beam_width * 2))
        topk_probs = topk_probs.tolist()
        topk_indices = topk_indices.tolist()

        for prefix, (p_b, p_nb) in beam.items():
            p_tot = logaddexp(p_b, p_nb)
            for prob, c in zip(topk_probs, topk_indices):
                if c == BLANK_IDX:
                    nb_b, nb_nb = curr_beam[prefix]
                    curr_beam[prefix] = (logaddexp(nb_b, p_tot + prob), nb_nb)
                else:
                    new_prefix = prefix + (c,)
                    nb_b, nb_nb = curr_beam[new_prefix]
                    if prefix and prefix[-1] == c:
                        curr_beam[new_prefix] = (nb_b, logaddexp(nb_nb, p_b + prob))
                        old_b, old_nb = curr_beam[prefix]
                        curr_beam[prefix] = (old_b, logaddexp(old_nb, p_nb + prob))
                    else:
                        curr_beam[new_prefix] = (nb_b, logaddexp(nb_nb, p_tot + prob))

        sorted_prefixes = sorted(
            curr_beam.keys(),
            key=lambda p: logaddexp(curr_beam[p][0], curr_beam[p][1]),
            reverse=True
        )[:beam_width]
        beam = {p: curr_beam[p] for p in sorted_prefixes}

    results = []
    for p in beam:
        score = logaddexp(beam[p][0], beam[p][1])
        text = "".join([IDX2CHAR.get(tok, '') for tok in p])
        results.append((text, score))
    results.sort(key=lambda x: x[1], reverse=True)
    return results


class SqueezeExcitation(nn.Module):
    """Channel attention mechanism to adaptively suppress background grid/lines."""
    def __init__(self, channels: int, reduction: int = 4):
        super().__init__()
        mid = max(4, channels // reduction)
        self.fc = nn.Sequential(
            nn.AdaptiveAvgPool2d(1),
            nn.Conv2d(channels, mid, 1),
            nn.ReLU(inplace=True),
            nn.Conv2d(mid, channels, 1),
            nn.Hardsigmoid(inplace=True),
        )

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


class ConvBNAct(nn.Module):
    """Standard Convolution + BatchNorm + Hardswish block."""
    def __init__(self, in_c: int, out_c: int, k=3, s=1, p=1):
        super().__init__()
        self.block = nn.Sequential(
            nn.Conv2d(in_c, out_c, k, stride=s, padding=p, bias=False),
            nn.BatchNorm2d(out_c),
            nn.Hardswish(inplace=True)
        )

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


class DSBlock(nn.Module):
    """Depthwise-Separable block with Squeeze-and-Excitation."""
    def __init__(self, in_c: int, out_c: int, stride=(1, 1), se: bool = True):
        super().__init__()
        self.use_res = (stride == (1, 1) or stride == 1) and in_c == out_c
        self.conv = nn.Sequential(
            nn.Conv2d(in_c, in_c, 3, stride=stride, padding=1, groups=in_c, bias=False),
            nn.BatchNorm2d(in_c),
            nn.Hardswish(inplace=True),
            SqueezeExcitation(in_c) if se else nn.Identity(),
            nn.Conv2d(in_c, out_c, 1, bias=False),
            nn.BatchNorm2d(out_c),
            nn.Hardswish(inplace=True),
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        if self.use_res:
            return x + self.conv(x)
        return self.conv(x)


class captCHAD(nn.Module):
    """
    captCHAD OCR architecture with:
    1. Contrast-Invariant Preprocessing Stem (Normalized Luminance + Sobel X + Sobel Y)
    2. MobileNet Depthwise-Separable Spatial Backbone with Squeeze-and-Excitation
    3. 1-Layer Bidirectional GRU (Horizontal receptive field across sequence)
    4. CTC Linear Projection to 63 classes
    """
    def __init__(self, num_classes: int = NUM_CLASSES, hidden_dim: int = 52, use_contrast_stem: bool = True):
        super().__init__()
        self.use_contrast_stem = use_contrast_stem
        c1, c2, c3, c4, c5 = 24, 32, 48, 64, 88

        # Fixed analytical Sobel edge filters (0 trainable parameters)
        sobel_x = torch.tensor([[-1., 0., 1.], [-2., 0., 2.], [-1., 0., 1.]]).view(1, 1, 3, 3) / 4.0
        sobel_y = torch.tensor([[-1., -2., -1.], [0., 0., 0.], [1., 2., 1.]]).view(1, 1, 3, 3) / 4.0
        self.register_buffer('sobel_x', sobel_x)
        self.register_buffer('sobel_y', sobel_y)

        # Spatial Stem: (6, 64, 192) -> (c1, 32, 96)
        in_c = 6 if use_contrast_stem else 3
        self.stem = ConvBNAct(in_c, c1, k=3, s=(2, 2), p=1)
        self.b1 = DSBlock(c1, c1)

        # Stage 2: (c1, 32, 96) -> (c2, 16, 48)
        self.b2 = DSBlock(c1, c2, stride=(2, 2))
        self.b3 = DSBlock(c2, c2)

        # Stage 3: (c2, 16, 48) -> (c3, 8, 48) (horizontal sequence length T=48 preserved)
        self.b4 = DSBlock(c2, c3, stride=(2, 1))
        self.b5 = DSBlock(c3, c3)

        # Stage 4: (c3, 8, 48) -> (c4, 4, 48)
        self.b6 = DSBlock(c3, c4, stride=(2, 1))
        self.b7 = DSBlock(c4, c4)

        # Stage 5: (c4, 4, 48) -> (c5, 2, 48)
        self.b8 = DSBlock(c4, c5, stride=(2, 1))
        self.b9 = DSBlock(c5, c5)

        # Collapse height: (c5, 2, 48) -> (c5, 1, 48)
        self.pool = nn.AdaptiveAvgPool2d((1, None))

        # Bidirectional GRU: spans full sequence context horizontally
        self.gru = nn.GRU(c5, hidden_dim, num_layers=1, bidirectional=True, batch_first=True)

        # CTC classifier projection: 2*hidden_dim -> num_classes (63)
        self.fc = nn.Linear(hidden_dim * 2, num_classes)

    def extract_contrast_features(self, x: torch.Tensor) -> torch.Tensor:
        """
        Extract normalized luminance and horizontal/vertical Sobel gradients,
        concatenated with normalized RGB for complete chromatic + edge features.
        """
        r, g, b = x[:, 0:1], x[:, 1:2], x[:, 2:3]
        lum = 0.299 * r + 0.587 * g + 0.114 * b
        mean = lum.mean(dim=(-2, -1), keepdim=True)
        std = lum.std(dim=(-2, -1), keepdim=True) + 1e-5
        norm_lum = (lum - mean) / std

        grad_x = F.conv2d(norm_lum, self.sobel_x, padding=1)
        grad_y = F.conv2d(norm_lum, self.sobel_y, padding=1)
        return torch.cat([r, g, b, norm_lum, grad_x, grad_y], dim=1)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        Args:
            x: Input tensor of shape (B, 3, 64, 192)
        Returns:
            Logits of shape (T=48, B, NUM_CLASSES) formatted for PyTorch CTCLoss
        """
        if self.use_contrast_stem:
            x = self.extract_contrast_features(x)

        x = self.stem(x)
        x = self.b1(x)
        x = self.b2(x)
        x = self.b3(x)
        x = self.b4(x)
        x = self.b5(x)
        x = self.b6(x)
        x = self.b7(x)
        x = self.b8(x)
        x = self.b9(x)
        x = self.pool(x).squeeze(2)          # (B, c5, 48)
        x = x.permute(0, 2, 1)               # (B, 48, c5)
        gru_out, _ = self.gru(x)             # (B, 48, 2*hidden_dim)
        logits = self.fc(gru_out)            # (B, 48, num_classes)
        return logits.permute(1, 0, 2)       # (T=48, B, num_classes)

    def decode(self, logits: torch.Tensor) -> list[str]:
        """
        Greedy CTC decode (argmax per frame).
        Args:
            logits: (T, B, C) or (B, T, C)
        Returns:
            List of decoded strings of length B
        """
        if logits.dim() == 3 and logits.shape[1] != NUM_CLASSES and logits.shape[2] == NUM_CLASSES:
            tokens_batch = logits.argmax(dim=-1).permute(1, 0)
        elif logits.dim() == 3 and logits.shape[2] == NUM_CLASSES:
            tokens_batch = logits.argmax(dim=-1)
        else:
            raise ValueError(f"Unexpected logits shape: {logits.shape}")

        return [decode_tokens(tokens.tolist()) for tokens in tokens_batch]

    def decode_beam_search(self, logits: torch.Tensor, beam_width: int = 15) -> list[str]:
        """
        CTC Beam Search decoding over candidate sequence paths.
        Args:
            logits: (T, B, C)
        Returns:
            List of top-1 decoded strings of length B
        """
        if logits.shape[2] != NUM_CLASSES:
            logits = logits.permute(1, 0, 2)
        log_probs = logits.log_softmax(dim=-1)  # (T, B, C)
        T, B, C = log_probs.shape

        results = []
        for b in range(B):
            sample_lp = log_probs[:, b, :]  # (T, C)
            candidates = decode_beam_search_single(sample_lp, beam_width=beam_width)
            results.append(candidates[0][0] if candidates else "")
        return results


# Aliases for compatibility
CaptchaCRNN = captCHAD
CaptchaCTCNet = captCHAD


def get_model(num_classes: int = NUM_CLASSES, hidden_dim: int = 52) -> captCHAD:
    return captCHAD(num_classes=num_classes, hidden_dim=hidden_dim)


if __name__ == "__main__":
    model = get_model()
    n_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
    print(f"captCHAD initialized ({n_params:,} parameters).")

    dummy_input = torch.randn(4, 3, 64, 192)
    logits = model(dummy_input)
    print(f"Forward pass output shape: {logits.shape} (T, B, C)")
    decoded_greedy = model.decode(logits)
    print(f"Greedy decode: {decoded_greedy}")
    print("Self-test passed!")