"""Example custom case for tools/run_tn_custom.py."""

from __future__ import annotations

import math

import numpy as np
from qibo import Circuit, gates


def build_circuit(nqubits, nlayers, seed):
    rng = np.random.default_rng(seed)
    circuit = Circuit(nqubits)
    for layer in range(nlayers):
        for qubit in range(nqubits):
            circuit.add(gates.RY(qubit, theta=rng.uniform(-math.pi, math.pi)))
            circuit.add(gates.RZ(qubit, theta=rng.uniform(-math.pi, math.pi)))
        for qubit in range(layer % 2, nqubits - 1, 2):
            circuit.add(gates.RXX(qubit, qubit + 1, theta=rng.uniform(-0.7, 0.7)))
            circuit.add(gates.RZZ(qubit, qubit + 1, theta=rng.uniform(-0.7, 0.7)))
    return circuit


def build_observable(nqubits, seed):
    return {
        "terms": [
            {
                "coefficient": 1.0 / max(1, nqubits - 1),
                "operators": [("Z", site), ("Z", site + 1)],
            }
            for site in range(nqubits - 1)
        ]
    }