# 项目文件：training_project/src/readingtorch/engine.py
import json
import math
from pathlib import Path

import torch
from torch import nn

from .checkpoint import load_checkpoint, save_checkpoint, snapshot
from .data import load_bundle, make_loader
from .models import build_model


def one_epoch(model, loader, loss_fn, device, task, y_scale=1.0, optimizer=None):
    training = optimizer is not None
    model.train(training)
    total, count, correct = 0.0, 0, 0
    for x, y in loader:
        x, y = x.to(device), y.to(device)
        # 训练建立计算图，验证直接计算；模式切换和梯度开关各自设置。
        with torch.set_grad_enabled(training):
            if training:
                optimizer.zero_grad(set_to_none=True)
            output = model(x)
            loss = loss_fn(output, y)
            if not torch.isfinite(loss):
                raise ValueError("损失出现了 NaN 或无穷值，请检查输入与学习率")
            if training:
                loss.backward()
                optimizer.step()
        total += loss.item() * len(x)  # 按样本数加权，末批较小时仍统计正确。
        count += len(x)
        if task == "images":
            correct += (output.argmax(1) == y).sum().item()
    result = {"loss": total / count}
    if task == "images":
        result["accuracy"] = correct / count
    else:
        result["rmse"] = math.sqrt(result["loss"]) * y_scale
    return result