"""
Classification metrics shared across all three courseworks.

Provides:
    - accuracy(): Top-1 accuracy
    - precision_recall_f1(): Per-class and macro-averaged metrics
    - confusion_matrix(): NxN confusion matrix
    - classification_report(): Formatted text report
"""

import numpy as np


def accuracy(y_pred, y_true):
    """
    Compute top-1 classification accuracy.

    Args:
        y_pred: (N,) predicted class indices
        y_true: (N,) ground truth class indices

    Returns:
        Scalar accuracy in [0, 1]
    """
    return np.mean(y_pred == y_true)


def precision_recall_f1(y_pred, y_true, num_classes=10):
    """
    Compute per-class precision, recall, and F1 score.

    Args:
        y_pred: (N,) predicted class indices
        y_true: (N,) ground truth class indices
        num_classes: Number of classes

    Returns:
        dict with keys:
            'precision': (num_classes,) per-class precision
            'recall': (num_classes,) per-class recall
            'f1': (num_classes,) per-class F1 score
            'macro_precision': scalar macro-averaged precision
            'macro_recall': scalar macro-averaged recall
            'macro_f1': scalar macro-averaged F1 score
    """
    precision = np.zeros(num_classes)
    recall = np.zeros(num_classes)
    f1 = np.zeros(num_classes)

    for c in range(num_classes):
        tp = np.sum((y_pred == c) & (y_true == c))
        fp = np.sum((y_pred == c) & (y_true != c))
        fn = np.sum((y_pred != c) & (y_true == c))

        precision[c] = tp / (tp + fp) if (tp + fp) > 0 else 0.0
        recall[c] = tp / (tp + fn) if (tp + fn) > 0 else 0.0
        f1[c] = (2 * precision[c] * recall[c] / (precision[c] + recall[c])
                 if (precision[c] + recall[c]) > 0 else 0.0)

    return {
        'precision': precision,
        'recall': recall,
        'f1': f1,
        'macro_precision': np.mean(precision),
        'macro_recall': np.mean(recall),
        'macro_f1': np.mean(f1),
    }


def confusion_matrix(y_pred, y_true, num_classes=10):
    """
    Compute confusion matrix.

    Args:
        y_pred: (N,) predicted class indices
        y_true: (N,) ground truth class indices
        num_classes: Number of classes

    Returns:
        (num_classes, num_classes) integer array where
        cm[i, j] = number of samples with true label i predicted as j
    """
    cm = np.zeros((num_classes, num_classes), dtype=np.int64)
    for t, p in zip(y_true, y_pred):
        cm[t, p] += 1
    return cm


def classification_report(y_pred, y_true, class_names, num_classes=10):
    """
    Generate a formatted classification report string.

    Args:
        y_pred: (N,) predicted class indices
        y_true: (N,) ground truth class indices
        class_names: List of class name strings
        num_classes: Number of classes

    Returns:
        Formatted string with per-class and overall metrics
    """
    metrics = precision_recall_f1(y_pred, y_true, num_classes)
    acc = accuracy(y_pred, y_true)

    lines = []
    lines.append(f"{'Class':<20s} {'Precision':>10s} {'Recall':>10s} {'F1':>10s} {'Support':>10s}")
    lines.append("-" * 60)

    for c in range(num_classes):
        support = np.sum(y_true == c)
        lines.append(
            f"{class_names[c]:<20s} "
            f"{metrics['precision'][c]:>10.4f} "
            f"{metrics['recall'][c]:>10.4f} "
            f"{metrics['f1'][c]:>10.4f} "
            f"{support:>10d}"
        )

    lines.append("-" * 60)
    lines.append(
        f"{'Macro Average':<20s} "
        f"{metrics['macro_precision']:>10.4f} "
        f"{metrics['macro_recall']:>10.4f} "
        f"{metrics['macro_f1']:>10.4f} "
        f"{len(y_true):>10d}"
    )
    lines.append(f"\nOverall Accuracy: {acc:.4f}")

    return "\n".join(lines)
