import numpy as np
from scipy.stats import ortho_group

import plotly
import plotly.plotly as py
import plotly.graph_objs as go

import matplotlib.pyplot as plt
import scipy.io as sio
from scipy import stats
%matplotlib inline

plotly.offline.init_notebook_mode(connected=True)
from IPython.core.display import display, HTML

# The polling here is to ensure that plotly.js has already been loaded before
# setting display alignment in order to avoid a race condition.
display(HTML(
    '<script>'
        'var waitForPlotly = setInterval( function() {'
            'if( typeof(window.Plotly) !== "undefined" ){'
                'MathJax.Hub.Config({ SVG: { font: "STIX-Web" }, displayAlign: "center" });'
                'MathJax.Hub.Queue(["setRenderer", MathJax.Hub, "SVG"]);'
                'clearInterval(waitForPlotly);'
            '}}, 250 );'
    '</script>'
))

n = 3 # order of matrices
number_tests = 1000

def operator_norm_random(f, nb_random_values=5000, min_coeff=-100, max_coeff=100):
    random_matrices = (max_coeff-min_coeff) * np.random.random_sample((nb_random_values, n, n)) + min_coeff
    return max(np.linalg.norm(f(x), ord=2)/np.linalg.norm(x, ord=2) for x in random_matrices)
    
def f(U, A):
    return U.dot(A.dot(U.T.conjugate()))
    
op_norms, spec_norms = [], []

for _ in range(number_tests):
    U, V = ortho_group.rvs(dim=n), ortho_group.rvs(dim=n)
    op_norms.append(operator_norm_random(lambda A: f(U, A)-f(V, A)))
    spec_norms.append(np.linalg.norm(U-V, ord=2))

results = (np.array(op_norms)<=np.array(spec_norms)).astype(int)
    
# Histogram of number of examples and counter-examples


trace = go.Bar(
    x = ['$\\textbf{Possibly (not sure): }\Vert f_U - f_V \Vert_{op} \overset{?}{\leq} \Vert U - V \Vert_2$', '$\Vert f_U - f_V \Vert_{op} > \Vert U - V \Vert_2$'],
    y = [np.sum(results), np.sum(1-results)],
    marker=dict(
        color='rgb(158,202,225)',
        line=dict(
            color='rgb(8,48,107)',
            width=1.5,
        )
    ),
    opacity=0.6
)


data = [trace]

layout = go.Layout(
    title='$\Vert f_U - f_V \Vert_{op} \overset{?}{\leq} \Vert U - V \Vert_2$',
    xaxis=dict(
        title='$\\text{Ordered pairs } (U, V) \\text{ such that ...}$'
    ),
    yaxis=dict(
        title='$\\text{Number of ordered pairs } (U, V)$'
    )
)
fig = go.Figure(data=data, layout=layout)

plotly.offline.iplot(fig)

trace2 = go.Scatter(
    x = list(range(number_tests)),
    y = np.sort(np.array(op_norms)-np.array(spec_norms)),
    mode = 'lines+markers',
    name = 'lines+markers'
)

data2 = [trace2]

layout2 = go.Layout(
    title='$\overbrace{\Vert f_U - f_V \Vert_{op}}^{\\text{estimated: lower bound}} - \Vert U - V \Vert_2$',
    xaxis=dict(
        title='$\\text{Ordered pairs } (U, V)$'
    ),
    yaxis=dict(
        title='$\\overbrace{\Vert f_U - f_V \Vert_{op}}^{\\text{estimated: lower bound}} - \Vert U - V \Vert_2$'
    )
)

fig2 = go.Figure(data=data2, layout=layout2)

plotly.offline.iplot(fig2)