import os
import time
import torch
import torch.nn as nn
from torch.optim import Adam
from torch.optim.lr_scheduler import ReduceLROnPlateau
from tqdm import tqdm

from model import get_model, count_parameters
from dataset import get_dataloaders, decode_prediction, ALPHABET


def greedy_decode(logits: torch.Tensor) -> list[str]:
    preds = logits.argmax(dim=2)       
    preds = preds.permute(1, 0)        
    results = []
    for seq in preds:
        results.append(decode_prediction(seq.tolist()))
    return results


def compute_accuracy(predictions: list[str], targets: list[str]) -> float:
    correct = sum(p == t for p, t in zip(predictions, targets))
    return correct / len(targets) if targets else 0.0


def train_one_epoch(model, loader, optimizer, criterion, device):
    model.train()
    total_loss = 0.0

    progress = tqdm(loader, desc="Обучение")
    for batch in progress:
        images = batch["images"].to(device)
        labels = batch["labels"]  # Оставляем на CPU для быстрой склейки
        label_lengths = batch["label_lengths"]

        optimizer.zero_grad()

        logits = model(images)                       # [T, B, C]
        log_probs = logits.log_softmax(2)

        T, B, _ = logits.shape
        input_lengths = torch.full((B,), T, dtype=torch.long, device=device)

        # Быстрая склейка таргетов на CPU, затем один перенос на GPU
        targets_flat = torch.cat([
            labels[i, :label_lengths[i]] for i in range(B)
        ]).to(device)
        
        label_lengths = label_lengths.to(device)

        loss = criterion(log_probs, targets_flat, input_lengths, label_lengths)

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

        total_loss += loss.item()
        progress.set_postfix(loss=f"{loss.item():.4f}")

    return total_loss / len(loader)


@torch.no_grad()
def validate(model, loader, criterion, device):
    model.eval()
    total_loss = 0.0
    all_preds, all_targets = [], []

    for batch in tqdm(loader, desc="Валидация"):
        images = batch["images"].to(device)
        labels = batch["labels"]
        label_lengths = batch["label_lengths"]
        label_strings = batch["label_strings"]

        logits = model(images)
        log_probs = logits.log_softmax(2)

        T, B, _ = logits.shape
        input_lengths = torch.full((B,), T, dtype=torch.long, device=device)

        targets_flat = torch.cat([
            labels[i, :label_lengths[i]] for i in range(B)
        ]).to(device)
        
        label_lengths = label_lengths.to(device)

        loss = criterion(log_probs, targets_flat, input_lengths, label_lengths)
        total_loss += loss.item()

        preds = greedy_decode(logits)
        all_preds.extend(preds)
        all_targets.extend(label_strings)

    avg_loss = total_loss / len(loader)
    accuracy = compute_accuracy(all_preds, all_targets)
    return avg_loss, accuracy, list(zip(all_preds[:5], all_targets[:5]))


def train(
    data_dir: str = "/mnt/d/datasetsss/mjsynth",
    epochs: int = 20,
    batch_size: int = 128,
    learning_rate: float = 1e-3,
    max_train_samples: int = None,
    save_dir: str = "../models",
    resume_from: str = None,
):
    device = "cuda" if torch.cuda.is_available() else "cpu"
    print(f"Устройство: {device}")

    os.makedirs(save_dir, exist_ok=True)

    print("\nЗагрузка данных...")
    train_loader, val_loader = get_dataloaders(
        data_dir=data_dir,
        batch_size=batch_size,
        num_workers=4,
        max_train_samples=max_train_samples,
    )

    num_classes = len(ALPHABET) + 1   # +1 для BLANK
    model = get_model(num_classes=num_classes, device=device)
    print(f"Параметров в модели: {count_parameters(model):,}")

    criterion = nn.CTCLoss(blank=0, reduction="mean", zero_infinity=True)
    optimizer = Adam(model.parameters(), lr=learning_rate)
    scheduler = ReduceLROnPlateau(optimizer, mode="max", factor=0.5, patience=2)

    start_epoch = 0
    best_accuracy = 0.0

    if resume_from and os.path.exists(resume_from):
        print(f"\nЗагружаем чекпоинт: {resume_from}")
        checkpoint = torch.load(resume_from, map_location=device)
        model.load_state_dict(checkpoint["model_state_dict"])
        optimizer.load_state_dict(checkpoint["optimizer_state_dict"])
        if "scheduler_state_dict" in checkpoint:
            scheduler.load_state_dict(checkpoint["scheduler_state_dict"])
        start_epoch = checkpoint.get("epoch", 0)
        best_accuracy = checkpoint.get("accuracy", 0.0)
        print(f"Продолжаем с эпохи {start_epoch}, лучшая точность: {best_accuracy:.4f}")

    print(f"\nНачинаем обучение: {epochs} эпох, batch_size={batch_size}\n")

    for epoch in range(start_epoch, epochs):
        start_time = time.time()
        print(f"\n=== Эпоха {epoch + 1}/{epochs} ===")

        train_loss = train_one_epoch(model, train_loader, optimizer, criterion, device)
        val_loss, val_accuracy, examples = validate(model, val_loader, criterion, device)

        scheduler.step(val_accuracy)
        elapsed = time.time() - start_time

        print(f"\nTrain Loss: {train_loss:.4f} | Val Loss: {val_loss:.4f} | Val Accuracy: {val_accuracy:.4f}")
        print(f"Время эпохи: {elapsed:.1f} сек")
        print("Примеры предсказаний (предсказано → правильно):")
        for pred, target in examples:
            mark = "✓" if pred == target else "✗"
            print(f"  {mark} '{pred}' → '{target}'")

        checkpoint = {
            "epoch": epoch + 1,
            "model_state_dict": model.state_dict(),
            "optimizer_state_dict": optimizer.state_dict(),
            "scheduler_state_dict": scheduler.state_dict(),
            "accuracy": val_accuracy,
        }
        torch.save(checkpoint, os.path.join(save_dir, "last_checkpoint.pth"))

        if val_accuracy > best_accuracy:
            best_accuracy = val_accuracy
            torch.save(checkpoint, os.path.join(save_dir, "best_model.pth"))
            print(f"Новая лучшая модель сохранена! Точность: {best_accuracy:.4f}")

    print(f"\nОбучение завершено. Лучшая точность: {best_accuracy:.4f}")
    return model


if __name__ == "__main__":
    train(
        data_dir="/mnt/d/datasatessss/mjsynth/mnt/ramdisk/max/90kDICT32px",
        epochs=20,
        batch_size=128,
        learning_rate=1e-3,
        max_train_samples=200_000,
        save_dir="../models",
    )