import numpy as np
import torch

from sklearn.metrics import classification_report, confusion_matrix

import matplotlib.pyplot as plt
from pathlib import Path

from src.logger import get_logger
from src.data import load_dataset
from src.predict import CLASS_NAMES, load_trained_model

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

logger = get_logger(__name__)

def evaluate(
    model,
    test_loader
):
    """
    PyTorch evaluation pipeline.
    Uses DataLoader instead of raw numpy arrays.
    """

    logger.info("Evaluation started")

    # If loader not passed, create dataset
    if test_loader is None:
        _, test_loader = load_dataset()

    # Load PyTorch model
    model = load_trained_model()
    model.to(DEVICE)
    model.eval()

    all_preds = []
    all_labels = []

    with torch.no_grad():
        for images, labels in test_loader:

            images = images.to(DEVICE)
            labels = labels.to(DEVICE)

            outputs = model(images)

            probs = torch.softmax(outputs, dim=1)
            preds = torch.argmax(probs, dim=1)

            all_preds.append(preds.cpu().numpy())
            all_labels.append(labels.cpu().numpy())

    y_pred = np.concatenate(all_preds)
    y_true = np.concatenate(all_labels)


    # Classification Report

    logger.info("Classification report generated")

    report = classification_report(
        y_true,
        y_pred,
        target_names=CLASS_NAMES
    )

    print(report)

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

    with open("outputs/reports/classification_report.txt", "w") as f:
        f.write(report)


    # Confusion Matrix

    logger.info("Confusion matrix generated")

    cm = confusion_matrix(y_true, y_pred)
    print(cm)

    fig, ax = plt.subplots(figsize=(10, 8))

    im = ax.imshow(cm)

    ax.set_xticks(np.arange(len(CLASS_NAMES)))
    ax.set_yticks(np.arange(len(CLASS_NAMES)))

    ax.set_xticklabels(CLASS_NAMES, rotation=45)
    ax.set_yticklabels(CLASS_NAMES)

    plt.xlabel("Predicted")
    plt.ylabel("Actual")
    plt.title("Confusion Matrix")
    plt.colorbar(im)

    plt.tight_layout()

    plt.savefig("outputs/reports/confusion_matrix.png")
    plt.close()

    logger.info("Evaluation completed")

    return report, cm


if __name__ == "__main__":
    evaluate()