import random

import numpy as np
import torch


SEED = 20260707


def set_seed(seed):
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)


def run_once():
    set_seed(SEED)

    model = torch.nn.Linear(4, 2)
    torch_values = torch.randn(4)
    numpy_values = np.random.rand(3)
    python_values = [random.random() for _ in range(3)]

    generator = torch.Generator().manual_seed(SEED)
    local_order = torch.randperm(8, generator=generator)

    return {
        "weight": model.weight.detach().flatten()[:4],
        "torch": torch_values,
        "numpy": numpy_values,
        "python": python_values,
        "order": local_order,
    }


first = run_once()
second = run_once()

torch_match = torch.equal(first["torch"], second["torch"])
weights_match = torch.equal(first["weight"], second["weight"])
numpy_match = np.array_equal(first["numpy"], second["numpy"])
python_match = first["python"] == second["python"]

print(f"seed: {SEED}")
print(f"torch values match: {torch_match}")
print(f"model weights match: {weights_match}")
print(f"numpy values match: {numpy_match}")
print(f"python values match: {python_match}")
print(f"local generator order: {first['order'].tolist()}")
print(f"runs match: {all([torch_match, weights_match, numpy_match, python_match])}")