import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import Dataset, DataLoader, IterableDataset
from PIL import Image
import numpy as np
import os
import matplotlib.pyplot as plt
from pathlib import Path
from tqdm import tqdm

try:
    from datasets import load_dataset
except ImportError:
    load_dataset = None

from .model import LightweightOCR, CharacterSet, save_checkpoint


# ========== DATASET ==========
class OCRDataset(Dataset):
    def __init__(self, labels_file, charset, img_height=32):
        self.img_height = img_height
        self.charset = charset
        self.samples = []

        labels_path = Path(labels_file)
        self.img_dir = labels_path.parent / "images"

        with open(labels_file, "r", encoding="utf-8") as f:
            for line in f:
                parts = line.strip().split("\t")
                if len(parts) == 2:
                    img_name, text = parts
                    charset.add_chars(text)
                    self.samples.append((img_name, text))

        print(f"  Loaded {len(self.samples)} samples")

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

    def __getitem__(self, idx):
        img_name, text = self.samples[idx]
        img_path = self.img_dir / img_name

        img = Image.open(img_path).convert("L")

        # Resize maintaining aspect ratio
        w, h = img.size
        new_h = self.img_height
        new_w = int(w * new_h / h)
        img = img.resize((new_w, new_h), Image.LANCZOS)

        # Normalize
        img = np.array(img).astype(np.float32) / 255.0
        img = torch.FloatTensor(img).unsqueeze(0)

        target = self.charset.encode(text)
        target = torch.LongTensor(target)

        return img, target, text


class HFIterableOCRDataset(IterableDataset):
    def __init__(
        self, dataset, charset, img_height=32, image_col="image", text_col="text"
    ):
        if load_dataset is None:
            raise ImportError("Please install 'datasets' library: pip install datasets")

        self.img_height = img_height
        self.charset = charset
        self.image_col = image_col
        self.text_col = text_col
        self.dataset = dataset

        # Iterate to build charset if needed
        if len(charset) < 10:
            print(f"  🔄 Scanning charset from '{text_col}' column (Streaming)...")
            count = 0
            # Scan first 5000 samples for safety
            max_scan = 5000
            for i, item in enumerate(dataset):
                if i >= max_scan:
                    break
                text = item.get(text_col, "")
                if text:
                    charset.add_chars(text)
                count += 1
                if count % 100 == 0:
                    print(f"    Scanned {count} samples...", end='\r')
            print(f"    Scanned {count} samples. Done.")
        else:
            print(f"  ℹ️  Skipping charset scan (charset already has {len(charset)} chars)")

    def __iter__(self):
        worker_info = torch.utils.data.get_worker_info()
        ds = self.dataset
        if worker_info is not None:
             if hasattr(ds, 'shard'):
                 try:
                    ds = ds.shard(num_shards=worker_info.num_workers, index=worker_info.id)
                 except:
                    pass
        
        for item in ds:
            try:
                img = item[self.image_col]
                text = item[self.text_col]

                if img.mode != "L":
                    img = img.convert("L")

                w, h = img.size
                new_h = self.img_height
                new_w = int(w * new_h / h)
                img = img.resize((new_w, new_h), Image.LANCZOS)

                img = np.array(img).astype(np.float32) / 255.0
                img = torch.FloatTensor(img).unsqueeze(0)

                target = self.charset.encode(text)
                target = torch.LongTensor(target)

                yield img, target, text
            except Exception:
                continue


class HFOCRDataset(Dataset):
    def __init__(
        self, dataset, charset, img_height=32, image_col="image", text_col="text"
    ):
        if load_dataset is None:
            raise ImportError("Please install 'datasets' library: pip install datasets")

        self.img_height = img_height
        self.charset = charset
        self.image_col = image_col
        self.text_col = text_col
        self.dataset = dataset

        # Iterate to build charset
        print(f"  🔄 Scanning charset from '{text_col}' column...")
        if text_col not in self.dataset.column_names:
            raise ValueError(
                f"Column '{text_col}' not found in dataset. Available: {self.dataset.column_names}"
            )

        # Optimize: select only text column to avoid decoding images during scan
        try:
            iter_ds = self.dataset.select_columns([text_col])
        except:
            iter_ds = self.dataset

        for item in tqdm(iter_ds, desc="Scanning charset"):
            text = item.get(text_col, "")
            if text:
                charset.add_chars(text)
        
        # Clear memory
        del iter_ds
        import gc
        gc.collect()

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

    def __getitem__(self, idx):
        item = self.dataset[idx]
        img = item[self.image_col]
        text = item[self.text_col]

        if img.mode != "L":
            img = img.convert("L")

        # Resize maintaining aspect ratio
        w, h = img.size
        new_h = self.img_height
        new_w = int(w * new_h / h)
        img = img.resize((new_w, new_h), Image.LANCZOS)

        # Normalize
        img = np.array(img).astype(np.float32) / 255.0
        img = torch.FloatTensor(img).unsqueeze(0)

        target = self.charset.encode(text)
        target = torch.LongTensor(target)

        return img, target, text


def collate_fn(batch):
    images, targets, texts = zip(*batch)

    max_width = max([img.size(2) for img in images])
    batch_size = len(images)
    height = images[0].size(1)

    padded_images = torch.zeros(batch_size, 1, height, max_width)

    for i, img in enumerate(images):
        w = img.size(2)
        padded_images[i, :, :, :w] = img

    target_lengths = torch.LongTensor([len(t) for t in targets])
    targets_concat = torch.cat(targets)

    return padded_images, targets_concat, target_lengths, texts


def train_model(
    model,
    train_loader,
    val_loader,
    charset,
    num_epochs=200,
    device="cuda",
    save_dir="models",
    lr=0.001,
    weight_decay=0.0001,
):
    os.makedirs(save_dir, exist_ok=True)

    model = model.to(device)
    criterion = nn.CTCLoss(blank=0, zero_infinity=True)
    optimizer = optim.AdamW(model.parameters(), lr=lr, weight_decay=weight_decay)
    scheduler = optim.lr_scheduler.ReduceLROnPlateau(
        optimizer, mode="min", factor=0.5, patience=10
    )

    best_loss = float("inf")
    best_acc = 0

    print("\n" + "=" * 70)
    print("  🚀 Training Lightweight OCR")
    print("=" * 70)

    history = {"train_loss": [], "val_loss": [], "acc": []}

    for epoch in range(num_epochs):
        # Training
        model.train()
        train_loss = 0

        # Progress bar for training
        pbar = tqdm(train_loader, desc=f"Epoch {epoch+1}/{num_epochs}", unit="batch")

        for batch_idx, (images, targets, target_lengths, texts) in enumerate(pbar):
            images = images.to(device)
            targets = targets.to(device)
            target_lengths = target_lengths.to(device)

            optimizer.zero_grad()
            outputs = model(images)

            input_lengths = torch.full(
                size=(outputs.size(1),),
                fill_value=outputs.size(0),
                dtype=torch.long,
                device=device,
            )

            log_probs = nn.functional.log_softmax(outputs, dim=2)
            loss = criterion(log_probs, targets, input_lengths, target_lengths)

            loss.backward()
            torch.nn.utils.clip_grad_norm_(model.parameters(), 5.0)
            optimizer.step()

            train_loss += loss.item()

            # Update progress bar
            pbar.set_postfix({"loss": loss.item()})

        train_loss /= len(train_loader)

        # Validation
        model.eval()
        val_loss = 0
        correct = 0
        total = 0

        if len(val_loader) > 0:
            with torch.no_grad():
                for images, targets, target_lengths, texts in val_loader:
                    images = images.to(device)
                    targets = targets.to(device)
                    target_lengths = target_lengths.to(device)

                    outputs = model(images)

                    input_lengths = torch.full(
                        size=(outputs.size(1),),
                        fill_value=outputs.size(0),
                        dtype=torch.long,
                        device=device,
                    )

                    log_probs = nn.functional.log_softmax(outputs, dim=2)
                    loss = criterion(log_probs, targets, input_lengths, target_lengths)
                    val_loss += loss.item()

                    # Accuracy
                    _, preds = outputs.max(2)
                    preds = preds.transpose(0, 1)

                    for i, text in enumerate(texts):
                        pred_indices = preds[i].cpu().numpy()
                        pred_text = charset.decode(pred_indices)
                        if pred_text == text:
                            correct += 1
                        total += 1

            val_loss /= len(val_loader)
            accuracy = correct / total * 100 if total > 0 else 0
        else:
            val_loss = 0
            accuracy = 0

        scheduler.step(val_loss)

        print(
            f"Epoch [{epoch+1:3d}/{num_epochs}] "
            f"Train: {train_loss:.4f} | "
            f"Val: {val_loss:.4f} | "
            f"Acc: {accuracy:5.2f}%"
        )

        # Store history
        history["train_loss"].append(train_loss)
        history["val_loss"].append(val_loss)
        history["acc"].append(accuracy)

        # Save checkpoint every epoch
        save_checkpoint(
            model,
            charset,
            optimizer,
            epoch + 1,
            val_loss,
            accuracy,
            f"{save_dir}/checkpoint_epoch_{epoch+1}.kiri",
        )

        # Save best model
        if accuracy > best_acc or (accuracy == best_acc and val_loss < best_loss):
            best_loss = val_loss
            best_acc = accuracy

            save_checkpoint(
                model,
                charset,
                optimizer,
                epoch + 1,
                val_loss,
                accuracy,
                f"{save_dir}/model.kiri",
            )
            print(f"  ✓ Saved Best! Acc: {accuracy:.2f}%")
        
        if torch.cuda.is_available():
            torch.cuda.empty_cache()
            
    # Plot history
    try:
        plt.figure(figsize=(12, 5))

        plt.subplot(1, 2, 1)
        plt.plot(history["train_loss"], label="Train Loss")
        plt.plot(history["val_loss"], label="Val Loss")
        plt.xlabel("Epoch")
        plt.ylabel("Loss")
        plt.title("Training Loss")
        plt.legend()
        plt.grid(True)

        plt.subplot(1, 2, 2)
        plt.plot(history["acc"], label="Validation Accuracy", color="green")
        plt.xlabel("Epoch")
        plt.ylabel("Accuracy (%)")
        plt.title("Training Accuracy")
        plt.legend()
        plt.grid(True)

        plt.tight_layout()
        plt.savefig(f"{save_dir}/training_history.png")
        print(f"\n📊 Training plot saved to {save_dir}/training_history.png")
    except Exception as e:
        print(f"\n⚠️ Failed to save plot: {e}")

    print(f"\n🎉 Training complete! Best accuracy: {best_acc:.2f}%\n")


def train_command(args):
    IMAGE_HEIGHT = args.height
    BATCH_SIZE = args.batch_size
    NUM_EPOCHS = args.epochs
    HIDDEN_SIZE = args.hidden_size

    device = args.device
    if device == "cuda" and not torch.cuda.is_available():
        print("⚠️ CUDA not available, using CPU")
        device = "cpu"

    print(f"\n🖥️  Device: {device}\n")

    # Load charset from pretrained model if available
    charset = CharacterSet()
    if (
        hasattr(args, "from_model")
        and args.from_model
        and os.path.exists(args.from_model)
    ):
        try:
            checkpoint = torch.load(args.from_model, map_location="cpu")
            if "charset" in checkpoint:
                print(f"🔄 Loading charset from {args.from_model}")
                charset = CharacterSet.from_checkpoint(checkpoint)
        except Exception:
            pass  # Will be handled later or just start fresh

    print("📂 Loading datasets...")

    train_dataset = None
    val_dataset = None

    if hasattr(args, "hf_dataset") and args.hf_dataset:
        print(f"  ⬇️ Loading HF dataset: {args.hf_dataset}")
        subset = getattr(args, "hf_subset", None)
        streaming = getattr(args, "hf_streaming", False)

        if streaming:
            print("  🌊 Streaming mode enabled")

        # Load train split
        try:
            ds = load_dataset(args.hf_dataset, subset, split=args.hf_train_split, streaming=streaming)
        except Exception as e:
            print(f"❌ Error loading dataset: {e}")
            return

        val_ds = None
        # Try finding val split
        val_splits = (
            [args.hf_val_split] if args.hf_val_split else ["val", "test", "validation"]
        )

        for split in val_splits:
            if not split:
                continue
            try:
                # We check if we can load it
                candidate = load_dataset(args.hf_dataset, subset, split=split, streaming=streaming)
                val_ds = candidate
                print(f"  ✓ Found validation split: '{split}'")
                break
            except:
                pass

        if val_ds is None:
            if not streaming:
                print(
                    f"  ⚠️ No validation split found. Splitting {args.hf_val_percent*100}% from train..."
                )
                try:
                    split_ds = ds.train_test_split(test_size=args.hf_val_percent)
                    ds = split_ds["train"]
                    val_ds = split_ds["test"]
                except Exception as e:
                    print(f"  ⚠️ Could not split dataset: {e}. Using train for validation.")
                    val_ds = ds
            else:
                 print(f"  ⚠️ No validation split found and streaming is on. Validation will be skipped.")

        if streaming:
            train_dataset = HFIterableOCRDataset(
                ds,
                charset,
                img_height=IMAGE_HEIGHT,
                image_col=args.hf_image_col,
                text_col=args.hf_text_col,
            )
            val_dataset = HFIterableOCRDataset(
                val_ds,
                charset,
                img_height=IMAGE_HEIGHT,
                image_col=args.hf_image_col,
                text_col=args.hf_text_col,
            ) if val_ds else None
        else:
            train_dataset = HFOCRDataset(
                ds,
                charset,
                img_height=IMAGE_HEIGHT,
                image_col=args.hf_image_col,
                text_col=args.hf_text_col,
            )
            val_dataset = HFOCRDataset(
                val_ds,
                charset,
                img_height=IMAGE_HEIGHT,
                image_col=args.hf_image_col,
                text_col=args.hf_text_col,
            )
    else:
        if not os.path.exists(args.train_labels):
            print(f"❌ Training labels not found: {args.train_labels}")
            return

        train_dataset = OCRDataset(args.train_labels, charset, img_height=IMAGE_HEIGHT)

        if args.val_labels and os.path.exists(args.val_labels):
            val_dataset = OCRDataset(args.val_labels, charset, img_height=IMAGE_HEIGHT)
        else:
            print(
                f"⚠️ Validation labels not found. Using training set for validation (not recommended)."
            )
            val_dataset = train_dataset

    train_len = len(train_dataset) if hasattr(train_dataset, "__len__") else "Unknown (Streaming)"
    val_len = len(val_dataset) if val_dataset and hasattr(val_dataset, "__len__") else ("Unknown (Streaming)" if val_dataset else 0)
    print(f"\n📊 Dataset: {train_len} train, {val_len} val")
    print(f"📝 Characters: {len(charset)}\n")
    
    # Clear memory before training
    import gc
    gc.collect()
    if torch.cuda.is_available():
        torch.cuda.empty_cache()

    os.makedirs(args.output_dir, exist_ok=True)
    # charset.save(f'{args.output_dir}/charset_lite.txt') # Not needed if using .kiri

    # shuffle=True is not supported for IterableDataset
    is_iterable = isinstance(train_dataset, IterableDataset)
    train_loader = DataLoader(
        train_dataset,
        batch_size=BATCH_SIZE,
        shuffle=not is_iterable,
        num_workers=4 if device == "cuda" else 0,
        collate_fn=collate_fn,
    )
    if val_dataset:
        val_loader = DataLoader(
            val_dataset,
            batch_size=BATCH_SIZE,
            shuffle=False,
            num_workers=4 if device == "cuda" else 0,
            collate_fn=collate_fn,
        )
    else:
        val_loader = []

    model = LightweightOCR(num_chars=len(charset), hidden_size=HIDDEN_SIZE)

    # Load pretrained model if specified
    if hasattr(args, "from_model") and args.from_model:
        print(f"🔄 Loading pretrained model from {args.from_model}")
        if os.path.exists(args.from_model):
            try:
                checkpoint = torch.load(args.from_model, map_location=device)
                if "model_state_dict" in checkpoint:
                    state_dict = checkpoint["model_state_dict"]
                    model_dict = model.state_dict()

                    # Filter out unnecessary keys or shape mismatches
                    pretrained_dict = {
                        k: v
                        for k, v in state_dict.items()
                        if k in model_dict and v.shape == model_dict[k].shape
                    }

                    # Log what was skipped
                    skipped_keys = [k for k in state_dict if k not in pretrained_dict]
                    if skipped_keys:
                        print(
                            f"  ⚠️ Skipped mismatched layers (fine-tuning): {skipped_keys[:5]} ..."
                        )

                    # Overwrite entries in the existing state dict
                    model_dict.update(pretrained_dict)
                    model.load_state_dict(model_dict)
                    print("  ✓ Weights loaded successfully")
                else:
                    print("  ⚠️ Invalid checkpoint format (missing model_state_dict)")
            except Exception as e:
                print(f"  ❌ Error loading model: {e}")
        else:
            print(f"  ❌ Pretrained model not found: {args.from_model}")

    total_params = sum(p.numel() for p in model.parameters())
    model_size_mb = total_params * 4 / 1024 / 1024

    print(f"🏗️  Model: {total_params:,} params ({model_size_mb:.2f} MB)\n")

    train_model(
        model,
        train_loader,
        val_loader,
        charset,
        num_epochs=NUM_EPOCHS,
        device=device,
        save_dir=args.output_dir,
        lr=args.lr,
        weight_decay=args.weight_decay,
    )
