from pathlib import Path
import json
import pickle

import torch
import torch.nn as nn
from torch.optim import Adam
from torch.optim.lr_scheduler import ReduceLROnPlateau

from src.logger import get_logger
from src.model import build_model

logger = get_logger(__name__)


def train(
    train_loader,
    val_loader,
    target_epochs=50
):

    logger.info("Training pipeline started")


    # Directories

    Path("models").mkdir(exist_ok=True)
    Path("outputs").mkdir(exist_ok=True)
    Path("outputs/reports").mkdir(parents=True, exist_ok=True)

    model_path = Path("models/resnet_cifar10.pth")
    best_model_path = Path("models/best_resnet_cifar10.pth")
    history_file = Path("outputs/history.pkl")


    # Device

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    logger.info(f"Device: {device}")


    # Model

    model = build_model().to(device)

    criterion = nn.CrossEntropyLoss(
        label_smoothing=0.1
    )

    optimizer = Adam(
        model.parameters(),
        lr=1e-4,
        weight_decay=1e-4
    )

    scheduler = ReduceLROnPlateau(
        optimizer,
        mode="min",
        factor=0.5,
        patience=3
    )

    # Resume Training

    start_epoch = 0

    if model_path.exists():

        logger.info(
            "Loading checkpoint"
        )

        checkpoint = torch.load(
            model_path,
            map_location=device
        )

        model.load_state_dict(
            checkpoint["model"]
        )

        optimizer.load_state_dict(
            checkpoint["optimizer"]
        )

        scheduler.load_state_dict(
            checkpoint["scheduler"]
        )

        start_epoch = checkpoint[
            "epoch"
        ]

        logger.info(
            f"Resuming from epoch {start_epoch}"
        )


    # Loss + Optimizer

    criterion = nn.CrossEntropyLoss(label_smoothing=0.1)

    optimizer = Adam(
        model.parameters(),
        lr=1e-4,
        weight_decay=1e-4
    )

    scheduler = ReduceLROnPlateau(
        optimizer,
        mode="min",
        factor=0.5,
        patience=3
    )


    # History

    if history_file.exists():
        with open(history_file, "rb") as f:
            history = pickle.load(f)
    else:
        history = {
            "loss": [],
            "accuracy": [],
            "val_loss": [],
            "val_accuracy": []
        }


    # Early stopping

    best_val_loss = float("inf")
    patience = 5
    early_counter = 0


    # Training loop

    for epoch in range(start_epoch, target_epochs):

        logger.info(f"Epoch {epoch+1}/{target_epochs}")

        #  TRAIN 
        model.train()

        train_loss = 0.0
        correct = 0
        total = 0

        for images, labels in train_loader:
            images, labels = images.to(device), labels.to(device)

            optimizer.zero_grad()

            outputs = model(images)
            loss = criterion(outputs, labels)

            loss.backward()
            optimizer.step()

            train_loss += loss.item()

            preds = outputs.argmax(dim=1)
            correct += (preds == labels).sum().item()
            total += labels.size(0)

        train_loss /= len(train_loader)
        train_acc = correct / total

        #  VALIDATION 
        model.eval()

        val_loss = 0.0
        val_correct = 0
        val_total = 0

        with torch.no_grad():
            for images, labels in val_loader:
                images, labels = images.to(device), labels.to(device)

                outputs = model(images)
                loss = criterion(outputs, labels)

                val_loss += loss.item()

                preds = outputs.argmax(dim=1)
                val_correct += (preds == labels).sum().item()
                val_total += labels.size(0)

        val_loss /= len(val_loader)
        val_acc = val_correct / val_total

        scheduler.step(val_loss)

        logger.info(
            f"Train Loss: {train_loss:.4f} | Train Acc: {train_acc:.4f}"
        )
        logger.info(
            f"Val Loss: {val_loss:.4f} | Val Acc: {val_acc:.4f}"
        )

        #  HISTORY 
        history["loss"].append(train_loss)
        history["accuracy"].append(train_acc)
        history["val_loss"].append(val_loss)
        history["val_accuracy"].append(val_acc)

        #  SAVE LAST MODEL 
        torch.save(
        {
            "epoch": epoch + 1,

            "model":
            model.state_dict(),

            "optimizer":
            optimizer.state_dict(),

            "scheduler":
            scheduler.state_dict(),

            "best_val_loss":
            best_val_loss
        },
        model_path
    )

        #  SAVE BEST MODEL 
        if val_loss < best_val_loss:

            best_val_loss = val_loss

            early_counter = 0

            torch.save(
                model.state_dict(),
                best_model_path
            )

            logger.info(
                "New best model saved"
            )

        else:

            early_counter += 1

        #  SAVE STATE 
    

        with open(history_file, "wb") as f:
            pickle.dump(history, f)

        #  EARLY STOPPING 
        if early_counter >= patience:
            logger.info("Early stopping triggered")
            break

    logger.info("Training completed")
    return history