#!/usr/bin/env python3
"""
plot of the theory and experimental values of S.
The exerimental data (already processed) is contained in ASCII .dat files
The theoretical data is in matlab binary files .mat

2016.11.22 Alessandro Cere
"""

import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
from scipy.io import loadmat

sns.set_context("poster")
sns.set_style("white")
sns.set_style("ticks", {"xtick.direction": "in",
                        "ytick.direction": "in"})

# Experimental data readout
dat_file = 'par_chsh.dat'
err_file = 'par_chsh_err.dat'
N_max = 15

with open(dat_file) as f:
    a = f.readline().strip().split()
    angles = np.array(a[1:]).astype('float')

    b = np.array([line.strip().split()
                  for line
                  in f])
    N = b[:, 0].astype('int')
    bells = b[:, 1:].astype('float')


with open(err_file) as f:
    a = f.readline().strip().split()
    angles = np.array(a[1:]).astype('float')

    b = np.array([line.strip().split()
                  for line
                  in f])
    # N = b[:, 0].astype('int')
    bells_err = b[:, 1:].astype('float')

# Loading the numerical simulations
mat = loadmat('parityCHSHalpha.mat')
N_th = mat['nRange'].astype('int')[0]
viol_th = mat['parityViol'].astype('float')
# viol_th = np.c_[np.zeros(21) + 2, viol_th.T]
angles_th = mat['alphaRange'].astype('float')[0]
# angles_th = np.insert(angles_th, 0, 0)

""" plotting"""
f, ax = plt.subplots()
f.set_size_inches(14, 9)

[plt.errorbar(angles,
              np.abs(bell),
              yerr=err,
              fmt='o:',
              label='N={}'.format(n))
 for bell, err, n
 in zip(bells[:N_max], bells_err, N)]

plt.gca().set_prop_cycle(None)

[plt.plot(angles_th, np.abs(bell), '-')
 for bell, n
 in zip(viol_th, N_th[:N_max])]

plt.xlabel('angle (degrees)')
plt.ylabel('Bell parameter S')


plt.legend(bbox_to_anchor=(.1, 0.6), loc=2, borderaxespad=0.)
# plt.legend()
plt.grid(1)
# plt.tight_layout()
plt.savefig('par_chsh.pdf')

plt.show()