"""Evolve a two-level system with a matrix exponential and inspect observables."""

from __future__ import annotations

import json

import numpy as np
from scipy.linalg import expm


def experiment() -> dict[str, object]:
    omega, duration = 2.0, 0.4
    x = np.array([[0, 1], [1, 0]], dtype=complex)
    z = np.diag([1, -1])
    hamiltonian = omega * x / 2
    state = expm(-1j * hamiltonian * duration) @ np.array([1, 0])
    return {
        "omega_rad_per_us": omega,
        "duration_us": duration,
        "probability_one": float(abs(state[1]) ** 2),
        "z_expectation": float(np.vdot(state, z @ state).real),
        "hamiltonian_expectation": float(np.vdot(state, hamiltonian @ state).real),
        "norm_squared": float(np.vdot(state, state).real),
    }


if __name__ == "__main__":
    print(json.dumps(experiment(), sort_keys=True))