from pathlib import Path

import numpy as np
import torch

from PIL import Image

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

from src.unknown import (
    detect_unknown,
    calculate_entropy
)

logger = get_logger(__name__)

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

CLASS_NAMES = [
    "airplane",
    "automobile",
    "bird",
    "cat",
    "deer",
    "dog",
    "frog",
    "horse",
    "ship",
    "truck",
    "unknown"
]


def load_trained_model(
        model_path="models/best_resnet_cifar10.pth"
):
    """
    Load best trained model.
    """

    logger.info(
        f"Loading model from {model_path}"
    )

    model = ResNet18(
        num_classes=11
    )

    model.load_state_dict(
        torch.load(
            model_path,
            map_location=DEVICE
        )
    )

    model.to(DEVICE)

    model.eval()

    logger.info(
        "Model loaded successfully"
    )

    return model


def load_external_image(
        image_path
):
    """
    Load image and preprocess
    for CIFAR prediction.
    """

    logger.info(
        f"Loading image: {image_path}"
    )

    image = Image.open(
        image_path
    ).convert("RGB")

    image = image.resize(
        (32, 32)
    )

    image = np.array(
        image
    ).astype(
        np.float32
    ) / 255.0

    logger.info(
        f"Image shape: {image.shape}"
    )

    return image


def preprocess_image(
        image
):
    """
    Convert image to PyTorch tensor.

    Input:
        (32,32,3)

    Output:
        (1,3,32,32)
    """

    image = torch.tensor(
        image,
        dtype=torch.float32
    )

    image = image.permute(
        2,
        0,
        1
    )

    image = image.unsqueeze(
        0
    )

    image = image.to(
        DEVICE
    )

    return image


def predict_single_image(
        model,
        image,
        confidence_threshold=0.60
):
    """
    Predict single image.
    """

    logger.info(
        "Running prediction"
    )

    image_tensor = preprocess_image(
        image
    )

    with torch.no_grad():

        outputs = model(
            image_tensor
        )

        probabilities = torch.softmax(
            outputs,
            dim=1
        )

    probabilities_np = (
        probabilities
        .cpu()
        .numpy()
    )

    predicted_idx = int(
        np.argmax(
            probabilities_np
        )
    )

    confidence = float(
        np.max(
            probabilities_np
        )
    )

    entropy = calculate_entropy(
        probabilities_np
    )

    logger.info(
        f"Entropy={entropy:.4f}"
    )

    if detect_unknown(
            probabilities_np,
            threshold=confidence_threshold
    ):

        predicted_class = "unknown"

    else:

        predicted_class = CLASS_NAMES[
            predicted_idx
        ]

    logger.info(
        f"Prediction={predicted_class}"
    )

    logger.info(
        f"Confidence={confidence:.4f}"
    )

    print("\nTop 3 Predictions\n")

    top3 = np.argsort(
        probabilities_np[0]
    )[-3:][::-1]

    for idx in top3:

        print(
            f"{CLASS_NAMES[idx]:12s}"
            f" : {probabilities_np[0][idx]:.4f}"
        )

    print(
        f"\nFinal Prediction : "
        f"{predicted_class}"
    )

    print(
        f"Confidence       : "
        f"{confidence:.4f}"
    )

    print(
        f"Entropy          : "
        f"{entropy:.4f}"
    )

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

    with open(
        "outputs/predictions/predictions.log",
        "a",
        encoding="utf-8"
    ) as f:

        f.write(
            f"Prediction={predicted_class}, "
            f"Confidence={confidence:.4f}, "
            f"Entropy={entropy:.4f}\n"
        )

    return (
        predicted_class,
        confidence,
        probabilities_np
    )


if __name__ == "__main__":

    model = load_trained_model()

    image = load_external_image(
        "inputs/sample_image.jpg"
    )

    prediction, confidence, _ = (
        predict_single_image(
            model,
            image
        )
    )

    print(
        f"\nResult: {prediction}"
    )

    print(
        f"Confidence: {confidence:.4f}"
    )