from __future__ import division

%pylab inline

try:
    import mpltools.style
    mpltools.style.use('ggplot')
except ImportError:
    print (
        "Not able to load ggplot style--- this won't matter too much, "
        "but your figures will look a little different."
    )

rcParams['axes.labelsize'] = 20
rcParams['axes.titlesize'] = 22
rcParams['legend.fontsize'] = 18
rcParams['xtick.labelsize'] = 14
rcParams['ytick.labelsize'] = 14
rcParams['lines.linewidth'] = 1.5

from scipy.linalg import inv, eig, sqrtm
from scipy.stats import chisqprob
from scipy.optimize import curve_fit, minimize
from scipy.io import loadmat

from IPython.display import display, Latex, HTML, Javascript, clear_output

import os
import glob
import re

from qinfer.rb import RandomizedBenchmarkingModel, p, F
from qinfer.smc import SMCUpdater, SMCUpdaterBCRB
from qinfer.resamplers import LiuWestResampler
from qinfer.derived_models import BinomialModel, DifferentiableBinomialModel
from qinfer.distributions import (
    ConstantDistribution, ProductDistribution, UniformDistribution, PostselectedDistribution, MultivariateNormalDistribution
)

def set_fields(arr, ix, fields):
    for key, val in fields.iteritems():
        if isinstance(val, np.ndarray) and len(val.shape) > 0:
            arr[ix][key][:] = val[:]
        else:
            arr[ix][key] = val

FIGURES_DIR = os.path.abspath("../figures/")
if not os.path.exists(FIGURES_DIR):
    os.mkdir(FIGURES_DIR)
    
print "Figures will be stored to {}.".format(FIGURES_DIR)

_savefig = savefig
def savefig(name):
    for fmt in ('pdf', 'svg', 'png'):
        _savefig(os.path.join(FIGURES_DIR, "{}.{}".format(name, fmt)), format=fmt, bbox_inches='tight')

import uuid

class ProgressBar(object):
    def __init__(self):
        self._divid = "div-{}".format(uuid.uuid4())
        self._value = 0

    def __del__(self):
        self.hide()
        
    def hide(self):
        display(Javascript('$("#{}").hide()'.format(self._divid)))
        
    @property
    def html(self):
        return """
            <div id="{}"></div>
        """.format(self._divid)
        
    @property
    def value(self):
        return self._value
    @value.setter
    def value(self, newval):
        self._value = newval
        self.refresh()
        
    def refresh(self):
        display(Javascript('$("#{}").progressbar({{value: {}}})'.format(self._divid, self._value)))
        
    def show(self):
        display(HTML(self.html))
        self.refresh()

import warnings
warnings.simplefilter('ignore', DeprecationWarning)

import benchmarking_models as bm

model = DifferentiableBinomialModel(RandomizedBenchmarkingModel(order=0))
il_model = DifferentiableBinomialModel(RandomizedBenchmarkingModel(interleaved=True, order=0))

ms = np.arange(0, 400, 10)

expparams = np.empty(ms.shape, dtype=model.expparams_dtype)
expparams['m'] = ms
expparams['n_meas'] = 10000

true_model = np.array([[0.9994, 0.5, 0.5], [0.9998, 0.5, 0.5]])

example_signal = model.simulate_experiment(true_model, expparams) / expparams['n_meas']

plot(ms, example_signal[0, 0], 'k:', label=r'$p = {}$'.format(true_model[0, 0]))
plot(ms, example_signal[0, 1], 'k-', label=r'$p = {}$'.format(true_model[1, 0]))
xlabel('$m$')
ylabel(r'$\hat{F}_g(m)$')
legend(loc='best')

savefig('ex-signal-noninterleaved')

ms = np.arange(0, 400, 10)
expparams = np.empty(ms.shape, dtype=il_model.expparams_dtype)
expparams['m'] = ms
expparams['n_meas'] = 10000
expparams = np.repeat(expparams, 2)

# Make the first half of the data the reference.
expparams['reference'][:len(expparams)//2] = True

# For the interleaved signals, m is doubled.
expparams['reference'][len(expparams)//2:] = False
expparams['m'][len(expparams)//2:] *= 2

true_model = np.array([[0.99994, 0.99999, 0.5, 0.5], [0.9999, 0.99999, 0.5, 0.5]])
example_il_signal = il_model.simulate_experiment(true_model, expparams) / expparams['n_meas']

ylim((0.965, 1.001))
plot(ms, example_il_signal[0, 0, :len(expparams)/2], 'k:', label='Reference')
plot(ms, example_il_signal[0, 0, len(expparams)/2:], 'k-', label='Interleaved')
legend(loc='lower left')

xlabel('$m$')
ylabel(r'$\hat{F}_g(m)$')

savefig('ex-signal-interleaved')

def achievable_non_interleaved(m_max, K, p, A, B):
    ms = np.arange(m_max)
    model = DifferentiableBinomialModel(RandomizedBenchmarkingModel())
    
    expparams = np.empty(ms.shape, dtype=model.expparams_dtype)
    expparams['m'] = ms
    expparams['n_meas'] = K
    
    true_model = np.array([[p, A, B]])
    
    fi = np.sum(model.fisher_information(true_model, expparams), axis=-1)[:, :, 0]
    
    return sqrtm(inv(fi))[0, 0]

def achievable_interleaved(m_max_ref, m_max_C, K, p_tilde, p_ref, A, B, m_step_ref=1, m_step_C=1):
    ms_ref = np.arange(0, m_max_ref, m_step_ref)
    ms_C   = np.arange(0, m_max_C,   m_step_C)
    
    model = DifferentiableBinomialModel(RandomizedBenchmarkingModel(interleaved=True))
    
    true_model = np.array([[p_tilde, p_ref, A, B]])
    
    # Calculate the FI for the reference and p_C measurements separately,
    # then add them.
    fi = np.zeros((4, 4))
    
    for ms, reference in [(ms_ref, True), (ms_C, False)]:
        expparams = np.empty(ms.shape, dtype=model.expparams_dtype)
        expparams['m'] = ms
        expparams['n_meas'] = K       
        expparams['reference'] = reference
        
        fi += np.sum(model.fisher_information(true_model, expparams), axis=-1)[:, :, 0]
    
    return sqrtm(inv(fi))[0, 0]

achievable_non_interleaved(40, 40, p(0.9994), 1.5, -.6)

def best_m(p_tilde, p_ref, A, B):
    objective = lambda m: (A**2 *m**2 *p_tilde**(-2 + 2*m)*p_ref**(2*m))/((-1 + B + A* p_tilde**m *p_ref**m)* (B + A*p_tilde**m *p_ref**m))
    result = minimize(objective,
                      100, method='nelder-mead'
                      )
    return result.x   

figure(figsize=(4,3))
As = linspace(0, 0.5, 100)

best_ms = []
for A in As:
    best_ms.append(np.ceil(best_m(p(0.9994), p(1 - 0.0011), A, 0.5)))

xlabel('$A$')
ylabel(r'$m_{\mathrm{opt}}(A)$')
plot(As, best_ms, 'k-')

savefig('best-m-vs-A')

figure(figsize=(4, 3))

Bs = linspace(.25, .75, 100)
best_ms = [] 

for B in Bs:
    best_ms.append(np.ceil(best_m(p(0.9994), p(1 - 0.0011), 0.25, B)))

xlabel('$B$')
ylabel(r'$m_{\mathrm{opt}}(B)$')
plot(Bs, best_ms, 'k-')

savefig('best-m-vs-B')

ps = 1 - np.logspace(-4, -0.5, 100)
m_opts = []

for p_val in ps:
    m_opts.append(best_m(p_val, p_val, 1, 0))
    
loglog((1 - ps), np.ceil(m_opts), 'k-')
gca().invert_xaxis()
xlabel('$1 - F$')
ylabel(r'$m_{\mathrm{opt}}(F)$')

savefig('limit-best-m')

risk_model = DifferentiableBinomialModel(RandomizedBenchmarkingModel(interleaved=True))
risk_prior = MultivariateNormalDistribution(
        mean=np.array([0.95, 0.95, 0.3, 0.5]),
        cov=np.diag([0.01, 0.01, 0.01, 0.01])**2
    )

def risk(refms, Cms, K, n_trials=100, n_particles = 1000):
    """
    Evaluates the risk incurred by using sequential Monte Carlo (SMC)
    to estimate benchmarking models.
    
    Parameters
    ----------
    refms : list
        List of values of $m$ to sample the reference curve at.
    Cms : list
        List of values of $m$ to sample the interleaved curve at.
    K : int
        Number of shots to take per sequence length.
    n_trials : int
        Number of times to run the SMC estimation procedure in
        estimating the risk.
    n_particles : int
        Number of SMC particles to use in estimating the benchmarking model.
    """
    
    # Allocate arrays to hold results.
    loss     = np.empty((n_trials,))
    postvar  = np.empty((n_trials,))
    bcrb_pps = np.empty((n_trials,))
    
    # TODO: fill these arrays!
    loss_trs = np.empty((n_trials,))
    bcrb_trs = np.empty((n_trials,))
    postvar_trs  = np.empty((n_trials,))
    
    # Make and show a progress bar.
    prog = ProgressBar()
    prog.show()
    
    for idx_trial in xrange(n_trials):
        true_model = risk_prior.sample()
        updater = SMCUpdaterBCRB(risk_model, n_particles, risk_prior, resampler=LiuWestResampler(0.95),
                                 initial_bim=None, adaptive=False)
         
        eps = np.array((
            [(m, True, K) for m in refms] +
            [(m, False, K) for m in Cms]
        ), dtype=risk_model.expparams_dtype)
        # Randomizing the order we perform experiments in may help?
        # np.random.shuffle(eps)
        data = risk_model.simulate_experiment(true_model, eps)[0, 0, :]
        
        for epvec, datum in zip(eps, data):
            updater.update(datum, epvec.reshape((1,)))
            
        mu = updater.est_mean()
        loss[idx_trial] = np.real(mu[0] - true_model[0, 0])**2
        loss_trs[idx_trial] = np.real(np.sum((mu - true_model[0])**2))
        
        prog.value = 100 * ((idx_trial + 1) / n_trials)
        
        inv_bim = inv(updater.current_bim)
        bcrb_pps[idx_trial] = inv_bim[0, 0]
        bcrb_trs[idx_trial] = np.trace(inv_bim)
        cov = updater.est_covariance_mtx()
        postvar[idx_trial] = cov[0, 0]
        postvar_trs[idx_trial] = np.trace(cov)
        
    return {
        'smc_risk_p': np.array([np.mean(loss), np.std(loss)]),
        'smc_risk_tr': np.array([np.mean(loss_trs), np.std(loss_trs)]),
        'smc_median_tr': np.median(loss_trs),
        'bcrb_p': np.mean(bcrb_pps),
        'postvar_p': np.mean(postvar),
        'bcrb_tr': np.mean(bcrb_trs),
        'postvar_tr': np.mean(postvar_trs)
    }

def exp_trial_fn(m, p, A, B):
    F_g = A * p**m + B
    return np.where(
        np.all([
            p >= 0, p <= 1,
            A >= -1, A <= 1,
            B >= 0, B <= 1
        ], axis=0),
        F_g,
        1e10 # Something far from actual data.
    )

def risk_lsf(refms, Cms, K, n_trials=100):
    """
    Evaluates the risk incurred by using least squares fitting (LSF)
    to estimate benchmarking models.
    
    Parameters
    ----------
    refms : list
        List of values of $m$ to sample the reference curve at.
    Cms : list
        List of values of $m$ to sample the interleaved curve at.
    K : int
        Number of shots to take per sequence length.
    n_trials : int
        Number of times to run the SMC estimation procedure in
        estimating the risk.
    """
    
    loss = np.empty((n_trials,))
    loss_tr = np.empty((n_trials,))
    
    prog = ProgressBar()
    prog.show()
    
    for idx_trial in xrange(n_trials):
        true_model = risk_prior.sample()
        
        # Generate simulated data, using the same model as was used in the SMC risk
        # evaluation.
        ref_eps = np.array([
            (m, True, K) for m in refms
        ], dtype=risk_model.expparams_dtype)
        C_eps = np.array([
            (m, False, K) for m in Cms
        ], dtype=risk_model.expparams_dtype)
        
        ref_data = risk_model.simulate_experiment(true_model, ref_eps)[0, 0, :] / K
        C_data = risk_model.simulate_experiment(true_model, C_eps)[0, 0, :] / K
        
        # Obtain least squares fits to each of the two curves.
        p0 = np.real(risk_prior.sample()[0, 1:]) # Another hack...
        res = curve_fit(exp_trial_fn, xdata=refms, ydata=ref_data, p0=p0, maxfev=10000)
        p_ref, A_ref, B_ref = res[0]
        
        res = curve_fit(exp_trial_fn, xdata=Cms, ydata=C_data, p0=p0, maxfev=10000)
        p_C, A_C, B_C = res[0]
        
        mu = p_C / p_ref
        
        prog.value = 100 * ((idx_trial + 1) / n_trials)
        
        loss[idx_trial] = (mu - true_model[0, 0])**2
        loss_tr[idx_trial] = np.sum((np.array([mu, p_ref, (A_ref + A_C) / 2, (B_ref + B_C) / 2]) - true_model[0])**2)
        
    return {
        'lsf_risk_p': np.array([np.mean(loss), np.std(loss)]),
        'lsf_risk_tr': np.array([np.mean(loss_tr), np.std(loss_tr)])
    }

refms = np.arange(1,100,1).astype(int)
Cms = np.arange(1,50,1).astype(int)

n_trials = 100

Ks = np.logspace(0, 5, 4).astype(int)

performance = np.zeros(
    (Ks.shape[0],),
    dtype=[
        ('smc_risk_p', '2float'),
        ('smc_risk_tr', '2float'),
        ('lsf_risk_p', '2float'),
        ('lsf_risk_tr', '2float'),
        ('smc_median_tr', float),
        
        ('bcrb_p', float),
        ('postvar_p', float),
        ('bcrb_tr', float),
        ('postvar_tr', float),
    ]
).view(np.recarray)

with warnings.catch_warnings():
    warnings.simplefilter('ignore')
    for idx_K, K in enumerate(Ks):
        print "K = {}, SMC:".format(K)
        set_fields(performance, idx_K, risk(refms, Cms, K, n_trials=n_trials, n_particles=10000))
        clear_output()
        
        print "K = {}, LSF:".format(K)
        set_fields(performance, idx_K, risk_lsf(refms, Cms, K, n_trials=n_trials))
        clear_output()
        

figure(figsize=(9, 5))

loglog(Ks, performance.bcrb_p, 'k-', linewidth=2, label='BCRB')


plot(Ks, performance.smc_risk_p[:, 0], 'bx-', label='SMC Risk')
plot(Ks, performance.lsf_risk_p[:, 0], 'rx--', label='LSF Risk')
loglog(Ks, performance.postvar_p, 'mx:', label='Posterior Variance')

ylabel('Mean squared error')
xlabel('Number of Shots per Sequence')

# We'll make the legend in the next cell.
ax = gca()
ax.yaxis.grid(color='gray', linestyle='dotted')

savefig('risk-comparison')

figure(figsize=(9, 5))

loglog(Ks, performance.bcrb_tr, 'k-', linewidth=2, label='BCRB')

plot(Ks, performance.smc_risk_tr[:, 0], 'bx-', label='SMC Risk')
plot(Ks, performance.lsf_risk_tr[:, 0], 'rx--', label='LSF Risk')
loglog(Ks, performance.postvar_tr, 'mx:', linewidth= 5, label='Posterior Variance')

ylabel('Risk')
xlabel('Number of Shots per Sequence')

#legend()

savefig('risk-tr-comparison')

ms = np.logspace(1.5, 3, 4).astype(int)
K = 1000
n_trials = 100

performance_v_m = np.zeros(
    (ms.shape[0],),
    dtype=[
        ('smc_risk_p', '2float'),
        ('smc_risk_tr', '2float'),
        ('lsf_risk_p', '2float'),
        ('lsf_risk_tr', '2float'),
        ('smc_median_tr', float),
        
        ('bcrb_p', float),
        ('postvar_p', float),
        ('bcrb_tr', float),
        ('postvar_tr', float),
    ]
).view(np.recarray)

with warnings.catch_warnings():
    warnings.simplefilter('ignore')
    for idx_m, m_max in enumerate(ms):
        marr = np.arange(1, m_max, 10).astype(int)
        
        print "m_max = {}, SMC:".format(m_max)
        set_fields(performance_v_m, idx_m, risk(marr, marr, K, n_trials=n_trials,n_particles=10000))
        clear_output()
        
        print "m_max = {}, SMC:".format(m_max)
        set_fields(performance_v_m, idx_m, risk_lsf(marr, marr, K, n_trials=n_trials))
        clear_output()

figure(figsize=(9, 5))
xlim(10.0**1.5, 10**3)

#loglog(Ks, crbs, label='Cramer-Rao Bound')
loglog(ms, performance_v_m.bcrb_tr, 'k-', linewidth=2, label='BCRB')

#errorbar(Ks, risks_smc[:, 0], yerr=risks_smc[:, 1], label='SMC Risk')
#errorbar(Ks, risks_lsf[:, 0], yerr=risks_lsf[:, 1], label='LSF Risk')

# Error bars don't seem to be working, so switching to ordinary plots for now.
#plot(ms, performance_v_m.smc_median_tr[:], 'x-.', color='cyan', label='SMC Median Loss')
plot(ms, performance_v_m.smc_risk_tr[:, 0], 'bx-', label='SMC Risk')
plot(ms, performance_v_m.lsf_risk_tr[:, 0], 'rx--', label='LSF Risk')
loglog(ms, performance_v_m.postvar_tr, 'mx:', linewidth = 5, label='Posterior Variance')


ylabel('Risk')
xlabel('Maximum Sequence Length Considered')

legend()

# We use the same legend as above, so no additional legend here.

savefig('risk-tr-comparison-vs-m-max')

n_sigs = 3
sigs = 1/logspace(2,n_sigs+1,n_sigs)

ms_ref = xrange(1, 101,10) #np.logspace(0, 10, num=20, base=2).astype(int)
ms_C   = xrange(1, 202,10) #np.logspace(0, 10, num=20, base=2).astype(int)
K = 10000

phys_model = BinomialModel(RandomizedBenchmarkingModel(interleaved=True))
mean=np.array([0.99, 0.99, 0.32, 0.5])

err = np.zeros(n_sigs)

for idx_sig in xrange(n_sigs):
    
    cov=np.diag(sigs[idx_sig]*ones((4,)))**2

    phys_prior =  PostselectedDistribution(MultivariateNormalDistribution(
            mean=mean,
            cov=cov
        ), phys_model)
    
    updater = SMCUpdater(
        phys_model, 10000, phys_prior,
        resampler=LiuWestResampler(0.9)
    )
     
    eps = np.array((
        [(m, True, K) for m in ms_ref] +
        [(m, False, K) for m in ms_C]
    ), dtype=phys_model.expparams_dtype)
    
    data = phys_model.simulate_experiment(mean[np.newaxis], eps)
        
    for epvec, datum in zip(eps, data):
        updater.update(datum, epvec.reshape((1,)))
            
    mu = updater.est_mean()
    err[idx_sig] = np.real(mu[1] - mean[1])**2
    

print err

rho_psi = np.diag(np.array([0.9, 0.1], dtype=complex)).flatten(order='F')
E_psi = rho_psi.conj().transpose()

gateset_dir = os.path.abspath('../gates/Sim_03_12_1742/')
gateset_files = glob.glob(os.path.join(gateset_dir, '*.mat'))

# QuaEC uses slightly different names from QuantumUtils`, so we define a mapping here.
gateset_namemap = {
    'Had': 'H',
    'Z90': 'R_pi4',
    'Z45': 'T',
    'X180': 'X',
    'Y180': 'Y',
    'Z180': 'Z'
}

gateset = {
    # Extract the name of the gate from strings like "pulseHad.mat".
    gateset_namemap[re.match('pulse(.*)\.mat', os.path.split(gateset_file)[1]).groups()[0]]:
    # Load the actual file, and grab the supermatrix representation of the actual gate.
    loadmat(gateset_file)['S']
    for gateset_file in gateset_files
}

gateset['I'] = np.eye(4, dtype=complex)

ideal_simulator = bm.RBPhysicalSimulator(paranoid=True)
simulator = bm.RBPhysicalSimulator(gateset, E_psi, rho_psi)

np.set_printoptions(precision=4, linewidth=100)

gs_true = simulator.interleaved_model_parameters((0, 1))
print gs_true
best_m(*gs_true)

log2(0.5) / log2(gs_true[0] * gs_true[1])

ms = np.arange(2, 100, 20)
target = (0, 1) # (identity, X)
n_seq_per_m = 20
ref_ps = np.array([
    np.mean([simulator.p_survival(simulator.random_sequence(m)) for idx in xrange(n_seq_per_m)])
    for m in ms
])
interleaved_ps = np.array([
    np.mean([simulator.p_survival(simulator.random_interleaved_sequence(m, target)) for idx in xrange(n_seq_per_m)])
    for m in ms
])

figure(figsize=(11,7))

p_til, p_ref, A, B = gs_true
p_C = p_til * p_ref

plot(ms, ref_ps, 'xk', label='Reference Signal', markersize=10)
plot(ms, A * p_ref**ms + B, '--k', label='Reference Model')
plot(ms, interleaved_ps, '.k', label='Interleaved Signal for $U = X$',  markersize=10)
plot(ms, A * p_C**(ms/2) + B, '-k', label='Interleaved Model')
legend()
xlabel('$m$')
ylabel(r'$\hat{p}(m)$')

savefig('example-phys-signal')

@np.vectorize
def phys_simulate_experiment(expparam_vec):
    print '.', 
    m, ref, n_meas = expparam_vec
    if ref:
        return np.sum([
            simulator.sample_interleaved_sequence(int(m), (0, 1))
            for idx in xrange(n_meas)
        ])
    
    else:        
        return np.sum([
            # Note that the sequence length is defined differently
            # for the interleaved and reference cases by a factor of 2.
            simulator.sample_random_sequence(2 * int(m))
            for idx in xrange(n_meas)
        ])
    

ms_ref = xrange(1, 201,10) #np.logspace(0, 10, num=20, base=2).astype(int)
ms_C   = xrange(2, 202,10) #np.logspace(0, 10, num=20, base=2).astype(int)
K = 1000

mean=np.array([0.95, 0.95, 0.3, 0.5])
cov=np.diag([0.01, 0.01, 0.01, 0.01])**2

phys_model = BinomialModel(RandomizedBenchmarkingModel(interleaved=True))
phys_prior =  PostselectedDistribution(MultivariateNormalDistribution(
        mean=mean,
        cov=cov
    ), phys_model)
print "The true state is ", np.sqrt(np.dot(mean - gs_true, np.dot(inv(cov),mean - gs_true))),"sigma from the prior mean."

print "The true state is outside of a", 100*(1 - chisqprob(np.dot(mean - gs_true, np.dot(inv(cov),mean - gs_true)),4)),"% credible region."


updater = SMCUpdater(
    phys_model, 20000, phys_prior,
    resampler=LiuWestResampler(0.9)
)
 
eps = np.array((
    [(m, True, K) for m in ms_ref] +
    [(m, False, K) for m in ms_C]
), dtype=phys_model.expparams_dtype)

print 'Simulating data...\r',
data = phys_simulate_experiment(eps)

print 'Performing SMC updates...\r',
for epvec, datum in zip(eps, data):
    updater.update(datum, epvec.reshape((1,)))


print 'Performing LSF...\t',
ref_data = data[:len(ms_ref)] / K
C_data = data[len(ms_ref):] / K

p0 = np.real(phys_prior.sample()[0, 1:]) # Another hack...
res = curve_fit(exp_trial_fn, xdata=ms_ref, ydata=ref_data, p0=p0, maxfev=10000)
p_ref, A_ref, B_ref = res[0]

res = curve_fit(exp_trial_fn, xdata=ms_C, ydata=C_data, p0=p0, maxfev=10000)
p_C, A_C, B_C = res[0]

lsf_estimate = np.array([p_C / p_ref, p_ref, (A_ref + A_C) / 2.0, (B_ref + B_C) / 2.0])

phys_table = r"""
    \begin{{array}}{{l|cccc}}
        & \tilde{{p}} & p_{{\text{{ref}}}} & A_0 & B_0 \\
        \hline
        \text{{True}} & {true[0]:.4f} & {true[1]:.4f} & {true[2]:.4f} & {true[3]:.4f} \\
        \text{{SMC Estimate}} & {smc_mean[0]:.4f} & {smc_mean[1]:.4f} & {smc_mean[2]:.4f} & {smc_mean[3]:.4f} \\
        \text{{LSF Estimate}} & {lsf_est[0]:.4f} & {lsf_est[1]:.4f} & {lsf_est[2]:.4f} & {lsf_est[3]:.4f} \\
        \hline
        \text{{SMC Error}}    & {smc_error[0]:.4f} & {smc_error[1]:.4f} & {smc_error[2]:.4f} & {smc_error[3]:.4f} \\
        \text{{LSF Error}} & {lsf_error[0]:.4f} & {lsf_error[1]:.4f} & {lsf_error[2]:.4f} & {lsf_error[3]:.4f}
    \end{{array}}
""".format(
    true=gs_true,
    smc_mean=updater.est_mean(),
    smc_stddev=np.real(np.diag(sqrtm(updater.est_covariance_mtx()))),
    smc_error=np.abs(updater.est_mean() - gs_true),
    lsf_est=lsf_estimate,
    lsf_error=np.abs(lsf_estimate - gs_true)
)

print "The true state is inside of a", 100*(1- chisqprob(np.dot(updater.est_mean() - gs_true, np.dot(inv(updater.est_covariance_mtx()),updater.est_mean() - gs_true)),4)),"% credible region."

display(Latex(phys_table))

figure(figsize=(9, 6))
plot(ms_ref, K * data[:len(ms_ref)] / K, 'k--', label='Reference')
plot(ms_C, K * data[len(ms_ref):] / K, label='Interleaved')
legend(loc='lower left')

xlabel('$m$')
ylabel('Number of survivals')
savefig('phys-gate-final-data-bad')

print phys_table

# Make another updater and use it for plotting below.
prior_updater = SMCUpdater(
    phys_model, 10000, phys_prior,
    resampler=LiuWestResampler(0.9)
)

figure(figsize=(12,6))
updater.plot_posterior_marginal(0, res=200)
prior_updater.plot_posterior_marginal(0, res=200)

ymin, ymax = ylim()
vlines(lsf_estimate[0], ymin, ymax * 2, linestyles=':')
vlines(gs_true[0], ymin, ymax * 2)
ylim(ymin, ymax)

xlabel(r'$\tilde{p}$')
ylabel(r'$\Pr(\tilde{p})$')
legend(['SMC Posterior', 'Bad Prior', 'LSF Estimate', 'True'], loc='upper left')

savefig('bad-prior-distns')

ms_ref = xrange(1, 101,10) #np.logspace(0, 10, num=20, base=2).astype(int)
ms_C   = xrange(2, 202,10) #np.logspace(0, 10, num=20, base=2).astype(int)
K = 100

mean=np.array([0.99, 0.99, 0.3, 0.5])
cov=np.diag([0.01, 0.01, 0.01, 0.01])**2

phys_model = BinomialModel(RandomizedBenchmarkingModel(interleaved=True))
phys_prior =  PostselectedDistribution(MultivariateNormalDistribution(
        mean=mean,
        cov=cov
    ), phys_model)
print "The true state is ", np.sqrt(np.dot(mean - gs_true, np.dot(inv(cov),mean - gs_true))),"sigma from the prior mean."
print "The true state is outside of a", 100*(1 - chisqprob(np.dot(mean - gs_true, np.dot(inv(cov),mean - gs_true)),4)),"% credible region."


updater = SMCUpdater(
    phys_model, 10000, phys_prior,
    resampler=LiuWestResampler(0.9)
)
 
eps = np.array((
    [(m, True, K) for m in ms_ref] +
    [(m, False, K) for m in ms_C]
), dtype=phys_model.expparams_dtype)

print 'Simulating data...\r',
data = phys_simulate_experiment(eps)

print 'Performing SMC updates...\r',
for epvec, datum in zip(eps, data):
    updater.update(datum, epvec.reshape((1,)))

print 'Performing LSF...\t',
ref_data = data[:len(ms_ref)] / K
C_data = data[len(ms_ref):] / K

p0 = np.real(phys_prior.sample()[0, 1:]) # Another hack...
res = curve_fit(exp_trial_fn, xdata=ms_ref, ydata=ref_data, p0=p0, maxfev=10000)
p_ref, A_ref, B_ref = res[0]

res = curve_fit(exp_trial_fn, xdata=ms_C, ydata=C_data, p0=p0, maxfev=10000)
p_C, A_C, B_C = res[0]

lsf_estimate = np.array([p_C / p_ref, p_ref, (A_ref + A_C) / 2.0, (B_ref + B_C) / 2.0])

phys_table = r"""
    \begin{{array}}{{l|cccc}}
        & \tilde{{p}} & p_{{\text{{ref}}}} & A_0 & B_0 \\
        \hline
        \text{{True}} & {true[0]:.4f} & {true[1]:.4f} & {true[2]:.4f} & {true[3]:.4f} \\
        \text{{SMC Estimate}} & {smc_mean[0]:.4f} & {smc_mean[1]:.4f} & {smc_mean[2]:.4f} & {smc_mean[3]:.4f} \\
        \text{{LSF Estimate}} & {lsf_est[0]:.4f} & {lsf_est[1]:.4f} & {lsf_est[2]:.4f} & {lsf_est[3]:.4f} \\
        \hline
        \text{{SMC Error}}    & {smc_error[0]:.4f} & {smc_error[1]:.4f} & {smc_error[2]:.4f} & {smc_error[3]:.4f} \\
        \text{{LSF Error}} & {lsf_error[0]:.4f} & {lsf_error[1]:.4f} & {lsf_error[2]:.4f} & {lsf_error[3]:.4f}
    \end{{array}}
""".format(
    true=gs_true,
    smc_mean=updater.est_mean(),
    smc_stddev=np.real(np.diag(sqrtm(updater.est_covariance_mtx()))),
    smc_error=np.abs(updater.est_mean() - gs_true),
    lsf_est=lsf_estimate,
    lsf_error=np.abs(lsf_estimate - gs_true)
)

print "The true state is inside of a", 100*(1- chisqprob(np.dot(updater.est_mean() - gs_true, np.dot(inv(updater.est_covariance_mtx()),updater.est_mean() - gs_true)),4)),"% credible region."

display(Latex(phys_table))

figure(figsize=(9, 6))
plot(ms_ref, data[:len(ms_ref)] / K, 'k--', label='Reference')
plot(ms_C, data[len(ms_ref):] / K, label='Interleaved')
legend(loc='lower left')

xlabel('$m$')
ylabel('Number of survivals')
savefig('phys-gate-final-data-good')

print phys_table

# Make another updater and use it for plotting below.
prior_updater = SMCUpdater(
    phys_model, 40000, phys_prior,
    resampler=LiuWestResampler(0.9)
)

figure(figsize=(12,6))
updater.plot_posterior_marginal(0, res=200)
prior_updater.plot_posterior_marginal(0, res=200)

ymin, ymax = ylim()
vlines(lsf_estimate[0], ymin, ymax * 2, linestyles=':')
vlines(gs_true[0], ymin, ymax * 2)
ylim(ymin, ymax)

xlabel(r'$\tilde{p}$')
ylabel(r'$\Pr(\tilde{p})$')
legend(['SMC Posterior', 'Prior', 'LSF Estimate', 'True'], loc='upper left')

savefig('good-prior-distns')