utils/logger.py

"""
utils/logger.py
---------------
Prediction logging utility.
Logs every API prediction to a CSV file for later analysis / research reporting.
"""

import csv
import os
import json
import datetime
from pathlib import Path


LOG_PATH = Path(__file__).parent.parent / "outputs" / "prediction_logs.csv"

FIELDNAMES = [
    "timestamp",
    "filename",
    "prediction",
    "confidence",
    "spectral_prob_fake",
    "spectral_confidence",
    "edge_prob_fake",
    "edge_confidence",
    "cnn_prob_fake",
    "cnn_confidence",
    "vit_prob_fake",
    "vit_confidence",
    "diffusion_prob_fake",
    "diffusion_confidence",
]


def log_prediction(filename: str, result: dict) -> None:
    """
    Append a prediction result dict to the CSV log.

    Args:
        filename : Original uploaded filename
        result   : Full prediction dict as returned by the fusion module
    """
    LOG_PATH.parent.mkdir(parents=True, exist_ok=True)

    write_header = not LOG_PATH.exists()

    row = {
        "timestamp": datetime.datetime.now().isoformat(),
        "filename": filename,
        "prediction": result.get("prediction", ""),
        "confidence": result.get("confidence", 0.0),
    }

    # Flatten branch scores
    branches = result.get("branches", {})
    for branch in ["spectral", "edge", "cnn", "vit", "diffusion"]:
        branch_data = branches.get(branch, {})
        row[f"{branch}_prob_fake"]   = round(branch_data.get("prob_fake", 0.5), 4)
        row[f"{branch}_confidence"]  = round(branch_data.get("confidence", 0.0), 4)

    with open(LOG_PATH, "a", newline="") as f:
        writer = csv.DictWriter(f, fieldnames=FIELDNAMES)
        if write_header:
            writer.writeheader()
        writer.writerow(row)


def get_log_summary() -> dict:
    """Return basic statistics from the prediction log."""
    if not LOG_PATH.exists():
        return {"total": 0, "real": 0, "ai_generated": 0}

    rows = []
    with open(LOG_PATH, "r") as f:
        reader = csv.DictReader(f)
        rows = list(reader)

    total = len(rows)
    ai_gen = sum(1 for r in rows if r["prediction"] == "AI-Generated")
    real   = total - ai_gen

    return {"total": total, "real": real, "ai_generated": ai_gen}