import torch
from torch import nn


torch.manual_seed(11)
model_path = "weather-score-state-dict.pth"


class WeatherScoreModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(3, 5),
            nn.ReLU(),
            nn.Linear(5, 1),
        )

    def forward(self, inputs):
        return self.net(inputs)


def build_model():
    model = WeatherScoreModel()
    model.eval()
    return model


sample = torch.tensor([[0.6, 0.2, 0.9]], dtype=torch.float32)

trained_model = build_model()
with torch.inference_mode():
    expected_output = trained_model(sample)

torch.save(trained_model.state_dict(), model_path)
print(f"saved={model_path}")

loaded_model = build_model()
state_dict = torch.load(model_path, map_location="cpu", weights_only=True)
load_result = loaded_model.load_state_dict(state_dict)
loaded_model.eval()

with torch.inference_mode():
    loaded_output = loaded_model(sample)

print(f"keys={len(state_dict)}")
print(f"load_result={load_result}")
print(f"outputs_match={torch.allclose(expected_output, loaded_output)}")
print(f"output_shape={tuple(loaded_output.shape)}")
print(f"prediction={loaded_output.item():.4f}")