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
import numpy as np
from tqdm import tqdm
import re
import Levenshtein
from symspellpy import SymSpell, Verbosity

# Define the character set
charset = '0123456789abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ .,%'
char_to_idx = {char: idx + 1 for idx, char in enumerate(charset)}  # Start indices from 1
idx_to_char = {idx: char for char, idx in char_to_idx.items()}
num_classes = len(charset) + 1  # +1 for the CTC blank label

# Check for CUDA availability
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# Define the CRNN model
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()
        conv = conv.squeeze(2)  # Remove the height dimension
        conv = conv.permute(0, 2, 1)  # [batch, width, channels]
        rnn_out, _ = self.lstm(conv)
        output = self.fc(rnn_out)
        return output.log_softmax(2).permute(1, 0, 2)  # CTC loss requires [time, batch, class]

# Define the OCRDataset class
# Define the OCRDataset class
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']
        # Download the image
        try:
            response = requests.get(img_url, timeout=10)
            response.raise_for_status()
            img = Image.open(BytesIO(response.content)).convert('L')  # Grayscale
            if self.transform:
                img = self.transform(img)
        except Exception as e:
            print(f"Failed to load image at index {idx}: {e}")
            return None  # Return None if the image cannot be loaded

        # Encode label
        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)


# Collate function for DataLoader
# Collate function for DataLoader
def collate_fn(batch):
    # Remove None images
    batch = [b for b in batch if b is not None] 
    if len(batch) == 0:
        return None, None, None
    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.ToTensor(),
    transforms.Normalize((0.5,), (0.5,)),
])

# Evaluation metrics
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)

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

# Early stopping
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

# Training function
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 batch in progress_bar:
            if batch[0] is None:
                continue  # Skip batches with no valid data
            images, labels, label_lengths = batch
            images = images.to(device)
            labels = labels.to(device)
            label_lengths = label_lengths.to(device)
            
            optimizer.zero_grad()
            outputs = model(images)
            
            batch_size = images.size(0)
            input_lengths = torch.full(size=(batch_size,), fill_value=outputs.size(0), dtype=torch.long).to(device)
            
            loss = criterion(outputs, labels, input_lengths, label_lengths)
            loss.backward()
            optimizer.step()
            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

# Validation function
def validate_model(model, test_loader, criterion):
    model.eval()
    total_loss = 0
    total_cer = 0
    total_wer = 0
    with torch.no_grad():
        for batch in tqdm(test_loader, desc="Validating", leave=False):
            if batch[0] is None:
                continue  # Skip batches with no valid data
            images, labels, label_lengths = batch
            images = images.to(device)
            labels = labels.to(device)
            label_lengths = label_lengths.to(device)
            
            outputs = model(images)
            
            batch_size = images.size(0)
            input_lengths = torch.full(size=(batch_size,), fill_value=outputs.size(0), dtype=torch.long).to(device)
            loss = criterion(outputs, labels, input_lengths, label_lengths)
            total_loss += loss.item()

            # Decode predictions and calculate CER/WER
            decoded_preds = decode_predictions(outputs)
            start = 0
            for i in range(batch_size):
                pred = decoded_preds[i]
                label_len = label_lengths[i]
                ground_truth = ''.join([idx_to_char.get(c.item(), '') for c in labels[start:start+label_len]])
                total_cer += cer(pred, ground_truth)
                total_wer += wer(pred, ground_truth)
                start += label_len
        
    avg_loss = total_loss / len(test_loader)
    total_samples = len(test_loader.dataset)
    avg_cer = total_cer / total_samples
    avg_wer = total_wer / total_samples
    return avg_loss, avg_cer, avg_wer

# Prediction decoding function
def decode_predictions(outputs):
    outputs = outputs.permute(1, 0, 2)  # [batch, time, classes]
    _, max_indices = outputs.max(2)
    predictions = []
    for indices in max_indices:
        pred_str = ''
        prev_idx = -1
        for idx in indices:
            idx = idx.item()
            if idx != prev_idx and idx != 0:  # Ignore duplicates and blanks
                pred_str += idx_to_char.get(idx, '')
            prev_idx = idx
        predictions.append(pred_str)
    return predictions

# Load the dataset
data_frame = pd.read_csv('/kaggle/input/dataset-new/shortened_dataset.csv')
data_frame = data_frame[['image_link', 'entity_name', 'entity_value']]
print(f"Total samples in dataset: {len(data_frame)}")

# For faster experimentation, you might consider using a smaller subset
# Uncomment the following line to use a subset:
# data_frame = data_frame.sample(n=10000, random_state=42).reset_index(drop=True)

# Split the dataset
from sklearn.model_selection import train_test_split
train_data, test_data = train_test_split(data_frame, test_size=0.1, random_state=42)

# Create dataset and dataloaders
train_dataset = OCRDataset(train_data, transform=transform)
test_dataset = OCRDataset(test_data, transform=transform)

# Adjust num_workers based on your system
num_workers = 4  # You can increase this number if your system allows
train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True, collate_fn=collate_fn, num_workers=num_workers)
test_loader = DataLoader(test_dataset, batch_size=32, shuffle=False, collate_fn=collate_fn, num_workers=num_workers)

# Initialize the model, criterion, and optimizer
model = CRNN(num_classes=num_classes).to(device)
criterion = nn.CTCLoss(blank=0, zero_infinity=True).to(device)
optimizer = optim.AdamW(model.parameters(), lr=0.0001)

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

what is wrong in this code? 
point out places for betterment 
don't tell to pre download the image as i have time constrain