"""Imports."""
import matplotlib.pyplot as plt
import numpy as np
import scipy as scp

from IPython.display import clear_output, display

import qiskit.pulse as pulse
import qiskit

"""Plotting style."""
%matplotlib inline
plt.rcParams.update({"font.size": 16, "font.weight": "bold"})

"""Load account: Only do this once."""
qiskit.IBMQ.load_account()

"""Get a provider and see available backends."""
provider = qiskit.IBMQ.get_provider(hub="ibm-q")
print("Available backends:", *provider.backends())

"""Select ibmq-poughkeepsie device to use Pulse."""
# backend = provider.get_backend("ibmq_poughkeepsie")
# config = backend.configuration()
# defaults = backend.defaults()

"""Get a pulse channel system."""
# system = pulse.PulseChannelSpec.from_backend(backend)

"""Video showing conceptual visualization of the above experiment on the Bloch sphere."""
from IPython.display import YouTubeVideo
YouTubeVideo("g10PIHQ65L4", width=640, height=360)

"""Load in the raw data."""
data = data = np.loadtxt("https://raw.githubusercontent.com/rmlarose/qcbq/master/data/t1_data.txt")
times = data[:, 1]  # microseconds
measurements = data[:, 0]  # probability of excited state

"""Plot the raw data."""
### Your code here!

"""Do the data fitting."""
def fit(time: float, T1: float) -> float:
    ### Your code here!

"""Use scipy.optimize.curve_fit to fit the data."""
### Your code here!

"""Plot your fit over the raw data here."""
### Your code here!
# Compute the fit

"""Video showing conceptual visualization of the above experiment on the Bloch sphere."""
from IPython.display import YouTubeVideo
YouTubeVideo("9Ekep8zgZHc", width=640, height=360)

"""Load in the raw data."""
t2data = np.loadtxt("https://raw.githubusercontent.com/rmlarose/qcbq/master/data/t2_data.txt")
times = t2data[:, 1]  # microseconds
measurements = t2data[:, 0]  # probability of ground state

"""Plot the raw data."""
### Your code here!

"""Define the fit function."""
def ramsey_fit(time: float, domega: float, T2: float) -> float:
    ### Your code here!

"""Do the fitting."""
### Your code here!

"""Plot the fit and data."""
### Your code here!

"""Video showing conceptual visualization of the above experiment on the Bloch sphere."""
from IPython.display import YouTubeVideo
YouTubeVideo("wkuqcCVhl04", width=640, height=360)

"""Load in the raw data."""
t2echo_data = np.loadtxt("https://raw.githubusercontent.com/rmlarose/qcbq/master/data/t2e_data.txt")
times = t2echo_data[:, 1]  # microseconds
measurements = t2echo_data[:, 0]  # probability of ground state

"""Plot the raw data."""
### Your code here!

"""Define a fit function."""
def hahn_fit(time: float, T2: float) -> float:
    ### Your code here!

"""Do the fitting."""
### Your code here!

"""Plot the fit function and data."""
### Your code here!

"""Function for the Bloch-Redfield density matrix."""
def rho(t: float, 
        alpha0: complex = np.sqrt(1 / 2), 
        beta0: complex = np.sqrt(1 / 2),
        T1: float = 40,
        T2: float = 70) -> np.ndarray:
    return np.array([
        [1 + (abs(alpha0)**2 - 1) * np.exp(- t / T1), alpha0 * np.conj(beta0) * np.exp(-t / T2)],
        [np.conj(alpha0) * beta0 * np.exp(-t / T2), abs(beta0)**2 * np.exp(-t / T1)]
    ])

"""Visualization of the Bloch-Redfield density matrix."""
import ipywidgets

@ipywidgets.interact
def vis(time=(0, 500, 1), T1=(20, 200, 1), T2=(20, 200, 1)):
    plt.title(r"$\rho_{BR}$")
    plt.imshow(np.real(rho(time, T1=T1, T2=T2)), cmap="Greens", vmin=0, vmax=1)
    plt.text(-0.2, 0.0, round(rho(time, T1=T1, T2=T2)[0, 0], 4))
    plt.text(-0.2, 1.0, round(rho(time, T1=T1, T2=T2)[1, 0], 4))
    plt.text(0.8, 0.0, round(rho(time, T1=T1, T2=T2)[0, 1], 4))
    plt.text(0.8, 1.0, round(rho(time, T1=T1, T2=T2)[1, 1], 4))
    plt.colorbar()
    plt.axis("off")
    plt.show()

"""Plot the matrix elements of rho over time."""
# Initial state
alpha = np.sqrt(0.3)
beta = np.sqrt(0.7)

# T1 and T2 times
T1 = 50
T2 = 70

# Time
time = np.linspace(0, 250, 100)

# Matrix elements
rho00 = 1 + (abs(alpha)**2 - 1) * np.exp(-time / T1)
rho01 = alpha * np.conj(beta) * np.exp(-time / T2)
rho10 = np.conj(rho01)
rho11 = 1 - rho00

# Plotting
plt.figure(figsize=(16, 5));
plt.plot(time, rho00, "--", label=r"$\rho_{00}$", linewidth=4);
plt.plot(time, np.real(rho01), "--", label=r"$\Re[\rho_{01}]$", linewidth=4);
plt.plot(time, np.real(rho10), "--", label=r"$\Re[\rho_{10}]$", linewidth=4);
plt.plot(time, rho11, "--", label=r"$\rho_{11}$", linewidth=4);
plt.grid();
plt.legend();
plt.xlabel("Time [microseconds]");