%matplotlib inline

import numpy as np
from qiskit_experiments.visualization import (
    CurvePlotter,
    IQPlotter,
    MplDrawer,
    PlotStyle,
)


from generate_data import generate_data

data = generate_data()


data_keys = set()
for _, x in data.items():
    data_keys.update(x.keys())
print(data_keys)


# Create plotter and set options and style.
plotter = CurvePlotter(MplDrawer())
plotter.set_options(
    plot_sigma=[
        (1.0, 0.5)
    ],  # Controls confidence-intervals for `y_interp_err` data-keys.
)
plotter.set_figure_options(
    series_params={
        "A": {"symbol": "o", "color": "C0", "label": "Qubit 0"},
        "B": {"symbol": "X", "color": "C1", "label": "Qubit 1"},
        "C": {"symbol": "v", "color": "k", "label": "Ideal 0"},
        "D": {"symbol": "^", "color": "k", "label": "Ideal 1"},
    },
    xlabel="Parameter",
    ylabel="${\\langle{}Z\\rangle{}}$",
    figure_title="Expectation Values",
)


plotter.figure()


plotter.set_series_data("A", x=data["A"]["x"], y=data["A"]["y"])
plotter.figure()


plotter.set_series_data(
    "A",
    x=data["A"]["x"],
    y=data["A"]["y"],
    x_interp=data["A"]["x_interp"],
    y_interp=data["A"]["y_interp"],
    y_interp_err=data["A"]["y_interp_err"],
)
plotter.figure()


plotter.set_series_data(
    "B",
    x=data["B"]["x"],
    y=data["B"]["y"],
    x_interp=data["B"]["x_interp"],
    y_interp=data["B"]["y_interp"],
    y_interp_err=data["B"]["y_interp_err"],
)
plotter.figure()


plotter.set_series_data(
    "C",
    x=data["C"]["x"],
    y=data["C"]["y"],
    x_interp=data["C"]["x_interp"],
    y_interp=data["C"]["y_interp"],
)
plotter.set_series_data(
    "D",
    x=data["D"]["x"],
    y=data["D"]["y"],
    x_interp=data["D"]["x_interp"],
    y_interp=data["D"]["y_interp"],
)
plotter.set_options(
    style=PlotStyle(legend_loc="lower left"),
)
plotter.figure()


plotter.set_options(subplots=(2, 1))
plotter.set_figure_options(
    series_params={
        "A": {
            "symbol": "o",
            "color": "C0",
            "label": "Qubit 0",
            "canvas": 0,
        },  # Here we add "canvas"
        "B": {
            "symbol": "X",
            "color": "C1",
            "label": "Qubit 1",
            "canvas": 1,
        },  # entries to the
        "C": {
            "symbol": "v",
            "color": "k",
            "label": "Ideal 0",
            "canvas": 0,
        },  # `series_params` figure
        "D": {"symbol": "^", "color": "k", "label": "Ideal 1", "canvas": 1},  # option.
    },
)
plotter.figure()


from generate_data import generate_iq_data

iq_data, iq_discriminator = generate_iq_data()


data_keys = set()
for _, x in iq_data.items():
    data_keys.update(x.keys())
print(data_keys)


# Create plotter
iq_plotter = IQPlotter(MplDrawer())

# Set options.
iq_plotter.set_options(
    style=PlotStyle(
        figsize=(6, 4),
        legend_loc=None,
        textbox_rel_pos=(0.18, -0.30),
    ),
)
iq_plotter.set_figure_options(
    series_params={
        "0": {"label": "$|0>$"},
        "1": {"label": "$|1>$"},
        "2": {"label": "$|2>$"},
    },
    xlabel="In-Phase",
    ylabel="Quad.",
    xval_unit="Arb.",
    yval_unit="Arb.",
    figure_title="Single-Shots",
)

# Set data.
for series, series_data in iq_data.items():
    iq_plotter.set_series_data(series, **series_data)
# Generate figure.
iq_plotter.figure()


iq_plotter.set_supplementary_data(discriminator=iq_discriminator)
iq_plotter.figure()


import matplotlib.pyplot as plt

fig, axes = plt.subplots(1, 2, figsize=(10, 6))
plotter.set_options(axis=axes[1])
iq_plotter.set_options(axis=axes[0])

# Generate both figures, which adds graphics to each subplot.
plotter.figure()
iq_plotter.figure()

# Tighten layout
plt.tight_layout()
