import numpy as np
import pandas as pd
import json as json
from scipy import stats
from statsmodels.formula.api import ols

from o_plot import opl # a small local package dedicated to this project

import sys
print(sys.executable)

# loading the data
file_name = 'QAD_data_new.json'
f = open(file_name)
All_data = json.load(f)

#defining the pandas data frame for statistics
df_All= pd.DataFrame(All_data,columns=['shots', 'device', 'fidelity',
                                        'mitigation','model','id_gates',
                                        'QV', 'metric'])
# any shot number >= 488 indicates that the curve calculation was ended after reaching n = 500:
df_All.loc[df_All.shots>=488,"shots"]=np.nan
# add the variable neperian log of safe shot number:
df_All['log_shots'] = np.log(df_All['shots']) 

stat_model = ols("log_shots ~  metric + mitigation + model + fidelity + QV + id_gates + device",
                 df_All.query("device != 'ideal_device' ")).fit()
print(stat_model.summary())

# this for Jensen-Shannon metric 
s_metric = 'jensenshannon'
sm = np.array([96])
SAD=0
# ! will be unselected by running the next cell

# main option for metric: squared euclidean distance
# skip this cell if you don't want this option
s_metric = 'sqeuclidean'
sm = np.array([97])
SAD=2   

# this for measurement mitigation
mit = 'yes'
MIT=0
# ! will be unselected by running the next cell

# main option: no mitigation
# skip this cell if you don't want this option
mit = 'no'
MIT=-4

# select data according to the options
df_mod = df_All[df_All.mitigation == mit][df_All.metric == s_metric]

print("mitigation:",mit,"  metric:",s_metric )
df_mod.groupby('device')[['shots','log_shots']].describe(percentiles=[0.5])

ADD=0+SAD+MIT
opl.plot_curves(All_data, np.append(sm,ADD+np.array([4,5,12,13,20,21,28,29,36,37,44,45])),                                                      
            "No noise simulator vs empirical model - $\epsilon=0.001$ - no delay",
            ["metric","mitigation"],
            ["device","model"], right_xlimit = 70)

print("mitigation:",mit,"  metric:",s_metric )
stat_model = ols("log_shots ~  model + fidelity ", df_mod.query("id_gates == 0.0 ")).fit()
print(stat_model.summary())

ADD=48+SAD+MIT
opl.plot_curves(All_data, np.append(sm,ADD+np.array([4,5,12,13,20,21,28,29,36,37,44,45])),                                                      
            "No noise simulator vs empirical model - $\epsilon=0.001$ - with delay",
            ["metric","mitigation"],
            ["device","model"], right_xlimit = 70)

print("mitigation:",mit,"  metric:",s_metric )
stat_model = ols("log_shots ~  model + fidelity", df_mod.query("id_gates == 256.0")).fit() 
print(stat_model.summary())

print("mitigation:",mit,"  metric:",s_metric )
stat_model = ols("log_shots ~  model + fidelity + id_gates",
                 df_mod).fit() 
print(stat_model.summary())

# rem: in this figure, the ideal device calculated value is shown
# but not included in the regression model
df_short = df_All.query("model == 'empirical' & id_gates == 0.0")
df_short = df_short[df_short.metric == s_metric]
# here for adding the "ideal device" observation (11 shots for fidelity = 1.0)
df_ideal= df_All.query("device == 'ideal_device' ")
df_ideal = df_ideal[df_ideal.metric == s_metric]  
df_short = df_short[df_short['mitigation']==mit]
tit_head = "$\ln \; n_s$ for $\epsilon$ = 0.001 vs "
tit_tail =  "mitigation:"+mit+ "  metric:"+s_metric+"  model: empirical, no delay"
opl.plot_scatter(df_ideal.append(df_short),"fidelity",
     tit_head + "quantum readout error fidelity\n" +tit_tail,
     left_xlimit = 0.6, right_xlimit = 1.01,
     bottom_ylimit = 2, top_ylimit = 6)

print("mitigation:",mit,"  metric:",s_metric )
stat_model = ols("log_shots ~  fidelity",
                 df_mod.query("id_gates == 0.0 & model == 'empirical'")).fit() 
print(stat_model.summary())

# check if python linregress (figure) and OLS stat_model (table) agree
round(np.sqrt(stat_model.rsquared),3)

# calculate the expected value of n_s for a theoretical fidelity equal to 1.0
# to be compared to the oracle derived value for an ideal device: n_s = 11  
alpha, beta = stat_model.params
n_ideal = np.exp(alpha+beta)
n_ideal

ADD=0+SAD+MIT
opl.plot_curves(All_data, np.append(sm,ADD+np.array([4,12,20,28,36,44])),                                                      
            "Empirical model - $\epsilon=0.001$ - no delay",
            ["metric","mitigation"],
            ["device","model"], right_xlimit = 45, top_ylimit = 0.5)

opl.plot_scatter(df_short, "QV",
     tit_head + "quantum volume\n" + tit_tail,
     left_xlimit = 0, right_xlimit = 33,
     bottom_ylimit = 2, top_ylimit = 6)

print("mitigation:",mit,"  metric:",s_metric )
stat_model = ols("log_shots ~  QV",
                 df_mod.query("id_gates == 0.0 & model == 'empirical'")).fit() 
print(stat_model.summary())

# # check if python linregress (figure) and OLS stat_model (table) agree
round(np.sqrt(stat_model.rsquared),3)