import csv
import json
from pathlib import Path

import matplotlib.pyplot as plt
import torch
from torch.utils.data import DataLoader

from src.data import PathDataset, get_source_name


# ── Inference ──────────────────────────────────────────────────────────────

# Run model inference and return one prediction record per sample
# raw_dataset and indices supply path/source/label metadata, decoupling this
# function from the wrapper layout (Subset/TransformSubset) of `dataset`
def predict_rows(model, dataset, raw_dataset, indices, batch_size, device, *, num_workers=4):
    loader = DataLoader(dataset, batch_size=batch_size, num_workers=num_workers, shuffle=False)
    records = []

    model.eval().to(device)
    with torch.no_grad():
        offset = 0
        for images, labels in loader:
            logits = model(images.to(device)).squeeze(1).cpu()
            probs = torch.sigmoid(logits)
            preds = (probs >= 0.5).long()

            for i in range(len(labels)):
                sample_idx = indices[offset + i]
                path, label = raw_dataset.samples[sample_idx]
                records.append({
                    "path": str(path),
                    "basename": Path(path).name,
                    "source": get_source_name(path),
                    "label": int(label),
                    "pred": int(preds[i].item()),
                    "prob_fake": float(probs[i].item()),
                    "logit": float(logits[i].item()),
                })
            offset += len(labels)

    return records


# Re-run model scoring for an existing list of records
def rescore_rows(
    model, records, image_size, batch_size, device, preprocess=None, *, num_workers=4
):
    dataset = PathDataset(records, image_size=image_size, preprocess=preprocess)
    loader = DataLoader(dataset, batch_size=batch_size, num_workers=num_workers, shuffle=False)
    outputs = []

    model.eval().to(device)
    with torch.no_grad():
        for images, _, indices in loader:
            logits = model(images.to(device)).squeeze(1).cpu()
            probs = torch.sigmoid(logits)
            preds = (probs >= 0.5).long()

            for j in range(len(indices)):
                base = dict(records[int(indices[j].item())])
                base["prob_fake"] = float(probs[j].item())
                base["pred"] = int(preds[j].item())
                base["logit"] = float(logits[j].item())
                outputs.append(base)

    return outputs


# ── Export ─────────────────────────────────────────────────────────────────

# Writes all prediction records to CSV in a fixed column order
def save_preds(records, output_path):
    output_path.parent.mkdir(parents=True, exist_ok=True)
    fieldnames = ["path", "basename", "source", "label", "pred", "prob_fake", "logit"]
    with open(output_path, "w", newline="") as f:
        writer = csv.DictWriter(f, fieldnames=fieldnames)
        writer.writeheader()
        writer.writerows(records)


# Saves the top-k highest-confidence false positives and false negatives for error analysis
def save_errors(records, output_path, top_k=32):
    false_positives = sorted(
        (r for r in records if r["label"] == 0 and r["pred"] == 1),
        key=lambda r: r["prob_fake"],
        reverse=True,
    )[:top_k]
    false_negatives = sorted(
        (r for r in records if r["label"] == 1 and r["pred"] == 0),
        key=lambda r: r["prob_fake"],
    )[:top_k]
    output_path.parent.mkdir(parents=True, exist_ok=True)
    with open(output_path, "w") as f:
        json.dump({
            "top_false_positives": false_positives,
            "top_false_negatives": false_negatives,
        }, f, indent=2)


# Saves P(fake) histograms: one overall class comparison and one per source
def save_hists(records, output_dir):
    output_dir.mkdir(parents=True, exist_ok=True)

    fig, ax = plt.subplots(figsize=(8, 5))
    real_probs = [r["prob_fake"] for r in records if r["label"] == 0]
    fake_probs = [r["prob_fake"] for r in records if r["label"] == 1]
    ax.hist(real_probs, bins=30, alpha=0.6, label="real", density=True)
    ax.hist(fake_probs, bins=30, alpha=0.6, label="fake", density=True)
    ax.set_xlabel("Predicted P(fake)")
    ax.set_ylabel("Density")
    ax.set_title("Confidence by class")
    ax.legend()
    fig.tight_layout()
    fig.savefig(output_dir / "confidence_by_class.png", dpi=160)
    plt.close(fig)

    for source in sorted({r["source"] for r in records}):
        fig, ax = plt.subplots(figsize=(8, 5))
        source_probs = [r["prob_fake"] for r in records if r["source"] == source]
        ax.hist(source_probs, bins=30, alpha=0.8)
        ax.set_xlabel("Predicted P(fake)")
        ax.set_ylabel("Count")
        ax.set_title(f"Confidence distribution: {source}")
        fig.tight_layout()
        fig.savefig(output_dir / f"confidence_{source}.png", dpi=160)
        plt.close(fig)


# Saves the aggregated results dict as a formatted JSON file
def save_summary(summary, output_path):
    output_path.parent.mkdir(parents=True, exist_ok=True)
    with open(output_path, "w") as f:
        json.dump(summary, f, indent=2)
