from qiskit_experiments.library.randomized_benchmarking import RBUtils

# Run old function with num_qubits = 2
num_qubits = 2
t1s = [100 for _ in range(num_qubits)]
t2s = [100 for _ in range(num_qubits)]
gate_length = 5

RBUtils.coherence_limit(
    nQ=num_qubits,
    T1_list=t1s,
    T2_list=t2s,
    gatelen=gate_length
)

# Run old function with num_qubits = 3
num_qubits = 3
t1s = [100 for _ in range(num_qubits)]
t2s = [100 for _ in range(num_qubits)]
gate_length = 5

RBUtils.coherence_limit(
    nQ=num_qubits,
    T1_list=t1s,
    T2_list=t2s,
    gatelen=gate_length
)

# Run new function with num_qubits = 2
num_qubits = 2
t1s = [100 for _ in range(num_qubits)]
t2s = [100 for _ in range(num_qubits)]
gate_length = 5

RBUtils.coherence_limit_error(
    num_qubits=num_qubits,
    gate_length=gate_length,
    t1s=t1s,
    t2s=t2s,
)

import numpy as np

np.isclose(
    RBUtils.coherence_limit(
        nQ=num_qubits,
        T1_list=t1s,
        T2_list=t2s,
        gatelen=gate_length
    ),
    RBUtils.coherence_limit_error(
        num_qubits=num_qubits,
        gate_length=gate_length,
        t1s=t1s,
        t2s=t2s,
    ),
    15
)

# Run new function with num_qubits = 3
num_qubits = 3
t1s = [100 for _ in range(num_qubits)]
t2s = [100 for _ in range(num_qubits)]
gate_length = 5

RBUtils.coherence_limit_error(
    num_qubits=num_qubits,
    gate_length=gate_length,
    t1s=t1s,
    t2s=t2s,
)

# Run new function with num_qubits = 9
num_qubits = 9
t1s = [100 for _ in range(num_qubits)]
t2s = [100 for _ in range(num_qubits)]
gate_length = 5

RBUtils.coherence_limit_error(
    num_qubits=num_qubits,
    gate_length=gate_length,
    t1s=t1s,
    t2s=t2s,
)

def coherence_limit(nQ=2, T1_list=None, T2_list=None, gatelen=0.1):
    T1 = np.array(T1_list)

    if T2_list is None:
        T2 = 2 * T1
    else:
        T2 = np.array(T2_list)

    if len(T1) != nQ or len(T2) != nQ:
        raise ValueError("T1 and/or T2 not the right length")

    coherence_limit_err = 0

    if nQ == 1:

        coherence_limit_err = 0.5 * (
            1.0 - 2.0 / 3.0 * np.exp(-gatelen / T2[0]) - 1.0 / 3.0 * np.exp(-gatelen / T1[0])
        )

    elif nQ == 2:

        T1factor = 0
        T2factor = 0

        for i in range(2):
            T1factor += 1.0 / 15.0 * np.exp(-gatelen / T1[i])
            T2factor += (
                2.0
                / 15.0
                * (
                    np.exp(-gatelen / T2[i])
                    + np.exp(-gatelen * (1.0 / T2[i] + 1.0 / T1[1 - i]))
                )
            )

        T1factor += 1.0 / 15.0 * np.exp(-gatelen * np.sum(1 / T1))
        T2factor += 4.0 / 15.0 * np.exp(-gatelen * np.sum(1 / T2))

        coherence_limit_err = 0.75 * (1.0 - T1factor - T2factor)

    else:
        raise ValueError("Not a valid number of qubits")

    return coherence_limit_err

def coherence_limit_error(
    num_qubits: int, gate_length: float, t1s: Sequence, t2s: Optional[Sequence] = None
):
     t1s = np.array(t1s)
    if t2s is None:
        t2s = 2 * t1s
    else:
        t2s = np.array([min(t2, 2 * t1) for t1, t2 in zip(t1s, t2s)])

    if len(t1s) != num_qubits or len(t2s) != num_qubits:
        raise ValueError("Length of t1s/t2s must equal num_qubits")

    def thermal_relaxation_choi(t1, t2, time):  # without excitation
        return qi.Choi(
            np.array(
                [
                    [1, 0, 0, np.exp(-time / t2)],
                    [0, 0, 0, 0],
                    [0, 0, 1 - np.exp(-time / t1), 0],
                    [np.exp(-time / t2), 0, 0, np.exp(-time / t1)],
                ]
            )
        )

    chois = [thermal_relaxation_choi(t1, t2, gate_length) for t1, t2 in zip(t1s, t2s)]
    traces = [np.real(np.trace(np.array(qi.SuperOp(choi)))) for choi in chois]
    d = 2**num_qubits
    return d / (d + 1) * (1 - functools.reduce(operator.mul, traces) / (d * d))