from pathlib import Path
import pandas as pd
import numpy as np
from _shared import load
ROOT = Path(__file__).resolve().parents[1]
OUT = ROOT / "outputs"
OUT.mkdir(exist_ok=True)
import json
import torch
from torch import nn
from torch.utils.data import TensorDataset, DataLoader
torch.set_num_threads(1)
torch.manual_seed(20261013)
torch.use_deterministic_algorithms(True)
raw=load("prediction")
assert raw["sample_id"].is_unique
cols=["x1","x2","x3"]
train=raw.loc[raw["split"].eq("train")].copy()
val=raw.loc[raw["split"].eq("validation")].copy()
test=raw.loc[raw["split"].eq("test")].copy()
assert all(len(d)>0 for d in [train,val,test])
assert not set(train["sample_id"]) & set(test["sample_id"])
fill=train[cols].mean()
assert fill.notna().all()
filled=train[cols].fillna(fill)
mean=filled.mean();scale=filled.std(ddof=0).replace(0,1)
def tensorize(frame):
    x=((frame[cols].fillna(fill)-mean)/scale).to_numpy(dtype=np.float32)
    y=frame["target"].to_numpy(dtype=np.float32).reshape(-1,1)
    assert np.isfinite(x).all() and np.isfinite(y).all()
    return torch.from_numpy(x),torch.from_numpy(y)
xtrain,ytrain=tensorize(train)
xval,yval=tensorize(val)
xtest,ytest=tensorize(test)
assert xtrain.shape==(len(train),3) and ytrain.shape==(len(train),1)
generator=torch.Generator().manual_seed(20261013)
loader=DataLoader(TensorDataset(xtrain,ytrain),batch_size=32,shuffle=True,generator=generator)
def build_model():
    return nn.Sequential(nn.Linear(3,16),nn.ReLU(),nn.Linear(16,1))
model=build_model().cpu()
optimizer=torch.optim.Adam(model.parameters(),lr=0.01)
loss_fn=nn.MSELoss()
best_loss=float("inf");best_state=None;best_epoch=0;history=[]
for epoch in range(80):
    model.train()
    total=0.0;n=0
    for xb,yb in loader:
        optimizer.zero_grad(set_to_none=True)
        prediction=model(xb)
        assert prediction.shape==yb.shape
        loss=loss_fn(prediction,yb)
        assert torch.isfinite(loss)
        loss.backward()
        assert all(p.grad is not None and torch.isfinite(p.grad).all() for p in model.parameters())
        optimizer.step()
        total+=loss.item()*len(xb);n+=len(xb)
    model.eval()
    with torch.no_grad():v=loss_fn(model(xval),yval).item()
    assert np.isfinite(v)
    history.append({"epoch":epoch+1,"train_mse":total/n,"validation_mse":v})
    if v<best_loss:
        best_loss=v;best_epoch=epoch+1
        best_state={k:value.detach().clone() for k,value in model.state_dict().items()}
assert best_state is not None
model.load_state_dict(best_state)
model.eval()
# 模型选择完成后只进行一次测试集评价。
with torch.no_grad():
    predicted=model(xtest)
    test_mse=loss_fn(predicted,ytest).item()
    baseline_mse=loss_fn(torch.full_like(ytest,ytrain.mean().item()),ytest).item()
checkpoint={"state_dict":best_state,"columns":cols,"fill":fill.tolist(),"mean":mean.tolist(),"scale":scale.tolist(),"best_epoch":best_epoch}
path=OUT/"project_model.pt"
torch.save(checkpoint,path)
saved=torch.load(path,map_location="cpu",weights_only=True)
restored=build_model()
restored.load_state_dict(saved["state_dict"])
restored.eval()
# 保存的预处理参数也参与重新推理，而不是仅检验网络权重。
f=pd.Series(saved["fill"],index=saved["columns"])
m=pd.Series(saved["mean"],index=saved["columns"])
s=pd.Series(saved["scale"],index=saved["columns"])
reloaded_x=torch.from_numpy(((test[saved["columns"]].fillna(f)-m)/s).to_numpy(dtype=np.float32))
with torch.no_grad():reloaded_prediction=restored(reloaded_x)
torch.testing.assert_close(predicted,reloaded_prediction,rtol=0,atol=0)
assert not restored.training and predicted.device.type=="cpu"
pd.DataFrame(history).to_csv(OUT/"project_model_history.csv",index=False)
pd.DataFrame({"sample_id":test["sample_id"].to_numpy(),"target":ytest.numpy().ravel(),"prediction":predicted.numpy().ravel()}).to_csv(OUT/"project_model_predictions.csv",index=False)
metrics={"torch_version":torch.__version__,"pandas_version":pd.__version__,"best_epoch":best_epoch,"validation_mse":best_loss,"test_mse":test_mse,"baseline_mse":baseline_mse,"roundtrip":"exact"}
(OUT/"project_model_metrics.json").write_text(json.dumps(metrics,indent=2),encoding="utf-8")
print(json.dumps(metrics,indent=2))
print("形状、梯度、CPU 推理及保存加载一致性已通过")