import matplotlib.pyplot as plt
import pandas as pd
import seaborn as sns
import numpy as np
from ipywidgets import interact, interactive, fixed, interact_manual, FloatSlider, Layout
import ipywidgets as widgets
sns.set(style='whitegrid', palette='deep', font_scale=1.1, rc={'figure.figsize': [8, 5]})
plt.rcParams['figure.figsize'] = [12, 8]


# load dataframe
filepath = "../datasets/AllErrors/U3_5.csv"
df = pd.read_csv(filepath)

# reorder columns
df = df[['theta', 'phi', 'lam', 'E', 'depol_prob', 't1', 't2', 'population', 'p0_0', 'p0_1', 'p1_0', 'p1_1']]

# filters
labels = ['theta', 'phi', 'lam', 'E']
depol_columns = ['depol_prob']
thermal_columns = ['t1', 't2', 'population']
readout_columns = ['p0_0', 'p0_1', 'p1_0', 'p1_1']

# filtered dataframes
ideal_only = df.query('depol_prob == 0 & t1 == inf & t2 == inf & p0_0 == 1 & p1_1 == 1')

df

# Explore Features

@interact
def show_gates_more_than(theta_range=widgets.FloatRangeSlider(value=[0, 6.3], min = 0.0, max=2*np.pi, layout=Layout(width='80%')),
                         lam_range=widgets.FloatRangeSlider(value=[0, 6.3], min = 0.0, max=2*np.pi, layout=Layout(width='80%')),
                        phi_values=widgets.SelectMultiple(options=df['phi'].unique(), value=tuple(df['phi'].unique()))):
    filtered = df.loc[(df['theta'].between(theta_range[0],theta_range[1])) & (df['lam'].between(lam_range[0],lam_range[1])) & (df['phi'].isin(phi_values))]
    sns.scatterplot(x='theta', y='E', hue="lam", data=filtered);

@interact
def show_gates_more_than(theta_range=widgets.FloatRangeSlider(value=[2.2, 2.6], min = 0.0, max=2*np.pi, layout=Layout(width='80%')),
                         lam_range=widgets.FloatRangeSlider(value=[0, 6.3], min = 0.0, max=2*np.pi, layout=Layout(width='80%')),
                        phi_values=widgets.SelectMultiple(options=df['phi'].unique(), value=tuple(df['phi'].unique()))):
    filtered = df.loc[(df['theta'].between(theta_range[0],theta_range[1])) & (df['lam'].between(lam_range[0],lam_range[1])) & (df['phi'].isin(phi_values))]
    sns.scatterplot(x='theta', y='E', hue="lam", data=filtered);

@interact
def show_gates_more_than(theta_range=widgets.FloatRangeSlider(value=[2.2, 2.6], min = 0.0, max=2*np.pi, layout=Layout(width='80%')),
                         lam_range=widgets.FloatRangeSlider(value=[0, 6.3], min = 0.0, max=2*np.pi, layout=Layout(width='80%')),
                        phi_values=widgets.SelectMultiple(options=df['phi'].unique(), value=tuple(df['phi'].unique()))):
    filtered = df.loc[(df['theta'].between(theta_range[0],theta_range[1])) & (df['lam'].between(lam_range[0],lam_range[1])) & (df['phi'].isin(phi_values))]
    sns.scatterplot(x='theta', y='E', hue="lam", data=filtered);

@interact
def show_gates_more_than(
theta_range=widgets.FloatRangeSlider(value=[0, 6.3], min = 0.0, max=2*np.pi, layout=Layout(width='80%'), description='theta_range'),
lam_range=widgets.FloatRangeSlider(value=[0, 6.3], min = 0.0, max=2*np.pi, layout=Layout(width='80%'), description='lam_range'),
phi_values=widgets.SelectMultiple(options=df['phi'].unique(), value=tuple(df['phi'].unique()), description='phi_values')):
    filtered = df.loc[(df['p1_1'] == 1.0) & (df['theta'].between(theta_range[0],theta_range[1])) & (df['lam'].between(lam_range[0],lam_range[1])) & (df['phi'].isin(phi_values))]
    display(theta_range, lam_range, phi_values);
    return sns.scatterplot(x='theta', y='E', hue='p0_0', data=filtered);
