import numpy as np
import matplotlib.pyplot as plt
import torch

from pathlib import Path

from src.logger import get_logger
from src.model import ResNet18
from src.predict import CLASS_NAMES

logger = get_logger(__name__)

DEVICE = torch.device(
    "cuda"
    if torch.cuda.is_available()
    else "cpu"
)


def load_model_for_evaluation():
    """
    Load trained ResNet checkpoint.
    """

    logger.info(
        "Loading model for evaluation"
    )

    model = ResNet18(
        num_classes=11
    )

    checkpoint = torch.load(
        "models/resnet_cifar10.pth",
        map_location=DEVICE
    )

    model.load_state_dict(
        checkpoint["model"]
    )

    model.to(DEVICE)

    model.eval()

    logger.info(
        "Model loaded successfully"
    )

    return model


def show_errors(
        test_loader,
        max_images=9
):
    """
    Display misclassified images.

    Parameters
    ----------
    test_loader : DataLoader
        PyTorch test dataloader

    max_images : int
        Number of errors to visualize
    """

    logger.info(
        "Starting error analysis"
    )

    model = load_model_for_evaluation()

    all_images = []
    all_labels = []
    all_predictions = []

    logger.info(
        "Running batched inference"
    )

    with torch.no_grad():

        for images, labels in test_loader:

            images = images.to(
                DEVICE
            )

            outputs = model(
                images
            )

            predictions = torch.argmax(
                outputs,
                dim=1
            )

            all_predictions.extend(
                predictions.cpu().numpy()
            )

            all_labels.extend(
                labels.cpu().numpy()
            )

            all_images.extend(
                images.cpu()
            )

    y_pred = np.array(
        all_predictions
    )

    y_true = np.array(
        all_labels
    )

    errors = np.where(
        y_pred != y_true
    )[0]

    logger.info(
        f"Misclassified samples: {len(errors)}"
    )

    if len(errors) == 0:

        logger.info(
            "No misclassifications found"
        )

        return

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

    plt.figure(
        figsize=(12, 10)
    )

    num_images = min(
        max_images,
        len(errors)
    )

    for i in range(num_images):

        idx = errors[i]

        image = all_images[idx]

        if isinstance(
            image,
            torch.Tensor
        ):
            image = image.numpy()

        # NCHW -> NHWC

        image = np.transpose(
            image,
            (1, 2, 0)
        )

        image = np.clip(
            image,
            0,
            1
        )

        actual = int(
            y_true[idx]
        )

        predicted = int(
            y_pred[idx]
        )

        plt.subplot(
            3,
            3,
            i + 1
        )

        plt.imshow(
            image
        )

        plt.title(
            f"Actual: {CLASS_NAMES[actual]}\n"
            f"Pred: {CLASS_NAMES[predicted]}"
        )

        plt.axis(
            "off"
        )

    plt.tight_layout()

    save_path = (
        "outputs/plots/error_analysis.png"
    )

    plt.savefig(
        save_path,
        dpi=300,
        bbox_inches="tight"
    )

    plt.show()

    plt.close()

    logger.info(
        f"Error analysis saved to {save_path}"
    )


if __name__ == "__main__":

    from src.data import load_dataset

    train_loader, test_loader = (
        load_dataset()
    )

    show_errors(
        test_loader
    )