import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import Dataset, DataLoader
from torchvision import transforms
from PIL import Image
import requests
from io import BytesIO
import pandas as pd
from sklearn.model_selection import train_test_split
import numpy as np
import re
import Levenshtein
from symspellpy import SymSpell, Verbosity
import torch_xla
import torch_xla.core.xla_model as xm
import torch_xla.distributed.xla_multiprocessing as xmp
import torch_xla.distributed.parallel_loader as pl
import torch_xla.utils.utils as xu

# Define the character set
charset = '0123456789abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ .,%'
char_to_idx = {char: idx + 1 for idx, char in enumerate(charset)}
idx_to_char = {idx: char for char, idx in char_to_idx.items()}

# Check for TPU availability
device = xm.xla_device()

class CRNN(nn.Module):
    def __init__(self, num_classes, input_channels=1):
        super(CRNN, self).__init__()
        
        self.cnn = nn.Sequential(
            nn.Conv2d(input_channels, 64, kernel_size=3, padding=1),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(2, 2),
            nn.Dropout(0.2),
            
            nn.Conv2d(64, 128, kernel_size=3, padding=1),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(2, 2),
            
            nn.Conv2d(128, 256, kernel_size=3, padding=1),
            nn.BatchNorm2d(256),
            nn.ReLU(inplace=True),
            nn.MaxPool2d((2, 1)),
            
            nn.Conv2d(256, 512, kernel_size=3, padding=1),
            nn.BatchNorm2d(512),
            nn.ReLU(inplace=True),
            nn.MaxPool2d((2, 1)),
            
            nn.Conv2d(512, 512, kernel_size=2),
            nn.ReLU(inplace=True)
        )
        
        self.lstm = nn.GRU(
            input_size=512,
            hidden_size=256,
            num_layers=2,
            bidirectional=True,
            batch_first=True
        )
        
        self.fc = nn.Linear(512, num_classes)

    def forward(self, x):
        conv = self.cnn(x)
        b, c, h, w = conv.size()
        assert h == 1, f"Expected height == 1, got {h}"
        conv = conv.squeeze(2)
        conv = conv.permute(0, 2, 1)
        rnn_out, _ = self.lstm(conv)
        output = self.fc(rnn_out)
        return output.permute(1, 0, 2)

class OCRDataset(Dataset):
    def __init__(self, data_frame, transform=None):
        self.data_frame = data_frame.reset_index(drop=True)
        self.transform = transform

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

    def __getitem__(self, idx):
        img_url = self.data_frame.iloc[idx]['image_link']
        label = self.data_frame.iloc[idx]['entity_value']

        # Attempt to load the image
        try:
            response = requests.get(img_url, timeout=5)  # Set a timeout to avoid long waits
            img = Image.open(BytesIO(response.content)).convert('L')  # Convert to grayscale
        except Exception as e:
            # Skip the failed image and return None to indicate skipping
            print(f"Failed to load image at index {idx}, skipping: {e}")
            return None

        # Apply transformations if available
        if self.transform:
            img = self.transform(img)

        # Convert the label to a tensor of character indices
        label_idx = [char_to_idx.get(char, 0) for char in label]
        label_tensor = torch.tensor(label_idx, dtype=torch.long)

        return img, label_tensor, len(label_idx)

def collate_fn(batch):
    batch = [b for b in batch if b is not None]  # Filter out None entries
    images, labels, label_lengths = zip(*batch)
    images = torch.stack(images)
    labels = torch.cat(labels)
    label_lengths = torch.tensor(label_lengths, dtype=torch.long)
    return images, labels, label_lengths

# Data augmentation and preprocessing transformations
transform = transforms.Compose([
    transforms.Resize((32, 128)),
    transforms.RandomRotation(degrees=5),
    transforms.ColorJitter(brightness=0.5, contrast=0.5),
    transforms.ToTensor(),
    transforms.Normalize((0.5,), (0.5,)),
    transforms.RandomErasing(p=0.1)
])

# Spell correction setup
sym_spell = SymSpell(max_dictionary_edit_distance=2)
sym_spell.create_dictionary_entry('gram', 1)
sym_spell.create_dictionary_entry('kilogram', 1)

def correct_text(text):
    suggestions = sym_spell.lookup_compound(text, max_edit_distance=2)
    return suggestions[0].term if suggestions else text

def cer(prediction, ground_truth):
    prediction = prediction.replace(' ', '')
    ground_truth = ground_truth.replace(' ', '')
    distance = Levenshtein.distance(prediction, ground_truth)
    return distance / max(len(ground_truth), 1)

def wer(prediction, ground_truth):
    pred_words = prediction.split()
    gt_words = ground_truth.split()
    distance = Levenshtein.distance(' '.join(pred_words), ' '.join(gt_words))
    return distance / max(len(gt_words), 1)

def get_lr_scheduler(optimizer):
    return optim.lr_scheduler.ReduceLROnPlateau(optimizer, mode='min', factor=0.5, patience=2, verbose=True)

class EarlyStopping:
    def __init__(self, patience=5, min_delta=0):
        self.patience = patience
        self.min_delta = min_delta
        self.counter = 0
        self.best_loss = None
        self.early_stop = False

    def __call__(self, val_loss):
        if self.best_loss is None:
            self.best_loss = val_loss
        elif val_loss > self.best_loss - self.min_delta:
            self.counter += 1
            if self.counter >= self.patience:
                self.early_stop = True
        else:
            self.best_loss = val_loss
            self.counter = 0

def train_model(model, train_loader, test_loader, criterion, optimizer, num_epochs=20):
    scheduler = get_lr_scheduler(optimizer)
    early_stopping = EarlyStopping(patience=5)

    for epoch in range(num_epochs):
        model.train()
        total_loss = 0
        progress_bar = tqdm(train_loader, desc=f"Epoch [{epoch+1}/{num_epochs}]", leave=False)

        for images, labels, label_lengths in progress_bar:
            if images is None:
                continue  # Skip this batch if it contains no valid images

            # Move data to the TPU device
            images = images.to(device)
            labels = labels.to(device)
            label_lengths = label_lengths.to(device)

            optimizer.zero_grad()
            outputs = model(images)

            # Apply log softmax to outputs before passing to CTCLoss
            outputs = torch.nn.functional.log_softmax(outputs, dim=-1)

            batch_size = images.size(0)
            seq_len = outputs.size(0)
            input_lengths = torch.full(size=(batch_size,), fill_value=seq_len, dtype=torch.long).to(device)

            if input_lengths.min() <= 0 or label_lengths.min() <= 0:
                print(f"Invalid lengths - input_lengths: {input_lengths}, label_lengths: {label_lengths}")
                continue

            loss = criterion(outputs, labels, input_lengths, label_lengths)
            if torch.isnan(loss) or torch.isinf(loss):
                print(f"Found invalid loss at epoch {epoch+1}, skipping this batch")
                continue  # Skip if nan or inf loss

            loss.backward()
            xm.optimizer_step(optimizer)  # Use xm.optimizer_step for TPU
            total_loss += loss.item()

            progress_bar.set_postfix(loss=loss.item())

        avg_loss = total_loss / len(train_loader)
        print(f"Epoch [{epoch+1}/{num_epochs}], Training Loss: {avg_loss:.4f}")

        val_loss, val_cer, val_wer = validate_model(model, test_loader, criterion)
        print(f"Validation Loss: {val_loss:.4f}, CER: {val_cer:.4f}, WER: {val_wer:.4f}")

        scheduler.step(val_loss)
        early_stopping(val_loss)
        if early_stopping.early_stop:
            print("Early stopping triggered")
            break

def validate_model(model, test_loader, criterion):
    model.eval()
    total_loss = 0
    total_cer = 0
    total_wer = 0
    with torch.no_grad():
        for images, labels, label_lengths in test_loader:
            images = images.to(device)
            labels = labels.to(device)
            label_lengths = label_lengths.to(device)

            outputs = model(images)

            batch_size = images.size(0)
            seq_len = outputs.size(0)
            input_lengths = torch.full(size=(batch_size,), fill_value=seq_len, dtype=torch.long).to(device)

            loss = criterion(outputs, labels, input_lengths, label_lengths)
            total_loss += loss.item()

            for i in range(batch_size):
                pred_str = decode_output(outputs[:, i, :])
                gt_str = decode_label(labels[i])
                total_cer += cer(pred_str, gt_str)
                total_wer += wer(pred_str, gt_str)

    avg_loss = total_loss / len(test_loader)
    avg_cer = total_cer / len(test_loader.dataset)
    avg_wer = total_wer / len(test_loader.dataset)
    return avg_loss, avg_cer, avg_wer

def decode_output(output):
    _, preds = torch.max(output, 2)
    preds = preds.contiguous().view(-1)
    pred_str = ''.join([idx_to_char.get(pred.item(), '') for pred in preds])
    return pred_str

def decode_label(label):
    return ''.join([idx_to_char.get(l.item(), '') for l in label])

# Load dataset (assuming a dataframe with 'image_link' and 'entity_value')
df = pd.read_csv('/path/to/your/dataset.csv')
train_df, test_df = train_test_split(df, test_size=0.2, random_state=42)

train_dataset = OCRDataset(train_df, transform=transform)
test_dataset = OCRDataset(test_df, transform=transform)

train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True, collate_fn=collate_fn)
test_loader = DataLoader(test_dataset, batch_size=32, shuffle=False, collate_fn=collate_fn)

# Initialize model, loss, and optimizer
model = CRNN(num_classes=len(charset)+1).to(device)
criterion = nn.CTCLoss(blank=0)
optimizer = optim.Adam(model.parameters(), lr=0.001)

# Train the model on TPU
train_model(model, train_loader, test_loader, criterion, optimizer, num_epochs=20)


would this run on kaggle tpu environment do did i made a mistake somewhere?