"""Reproduce the website's measured, local CPU training example."""
import json
import platform
from pathlib import Path

import torch
import flagquantum as fq

torch.manual_seed(7)

def circuit(parameters):
    return fq.Circuit(2).ry(0, parameters[0]).cx(0, 1)

model = fq.Module(circuit, n_parameters=1, init=torch.tensor([0.25]))
training = fq.train(
    model,
    optimizer=torch.optim.Adam(model.parameters(), lr=0.08),
    objective=lambda z: z.mean(),
    steps=60,
)
trained_circuit = circuit(next(model.parameters()).detach())
measurement = fq.expectation(fq.Z(0))
result = fq.run(trained_circuit, outputs=measurement)
record = {
    "scope": "Local CPU simulation; no remote hardware job submitted",
    "framework": fq.__version__,
    "torch": torch.__version__,
    "python": platform.python_version(),
    "qubits": 2,
    "parameters": 1,
    "optimizer": "Adam",
    "learning_rate": 0.08,
    "steps": training.completed_steps,
    "losses": list(training.losses),
    "expectation": float(result.expectation()),
    "theta": float(next(model.parameters()).detach()[0]),
}
print(json.dumps(record, indent=2))