# 项目文件：training_project/src/readingtorch/predict.py
import math

import numpy as np
import torch

from .checkpoint import load_checkpoint
from .models import build_model


class Predictor:
    def __init__(self, checkpoint, device="cpu"):
        self.saved = load_checkpoint(checkpoint)
        self.device = torch.device(device)
        self.model = build_model(self.saved["task"]).to(self.device)
        self.model.load_state_dict(self.saved["model"])
        self.model.eval()  # 关闭 Dropout 的随机丢弃，再关闭后续梯度记录。

    def pages(self, minutes, sessions):
        if self.saved["task"] != "regression":
            raise ValueError("这份检查点用于图形分类")
        if not all(math.isfinite(value) and value > 0 for value in [minutes, sessions]):
            raise ValueError("分钟数与次数应为正数")
        pre = self.saved["preprocess"]
        x = torch.tensor([[minutes, sessions]], dtype=torch.float32)
        x = ((x - pre["x_mean"]) / pre["x_scale"]).to(self.device)
        with torch.inference_mode():
            result = self.model(x).cpu() * pre["y_scale"] + pre["y_mean"]
        return result.item()