import matplotlib.pyplot as plt
from PIL import Image
from matplotlib.image import imread
import scipy.io
import plotly
from plotly.subplots import make_subplots
import plotly.graph_objs as go
import numpy as np
from plotly import __version__
from plotly.offline import download_plotlyjs, init_notebook_mode, plot, iplot
config={'showLink': False, 'displayModeBar': False}

init_notebook_mode(connected=True)

import os
import markdown
import random
from scipy.integrate import quad

import warnings
warnings.filterwarnings('ignore')

def integrand(x, delta, t_2r):

    return (1/abs(3*x**2 - 1))*np.exp(-2*((2*np.pi*delta*t_2r)/(3*x**2 - 1))**2)

init_notebook_mode(connected=True)

delta = np.arange(-8000,8100,100) # in Hz
t_2r = np.arange(6e-6,24e-6,2e-6) # in s
G_gaussian = np.zeros((len(delta),len(t_2r)+1))
G_lor = np.zeros((len(delta),len(t_2r)+1))
G_superlor = np.zeros((len(delta),len(t_2r)+1))

G_gaussian[:,0] = delta
G_lor[:,0] = delta
G_superlor[:,0] = delta

for ii in range(len(t_2r)):
    for jj in range(len(delta)):
        G_gaussian[jj,ii+1] = (t_2r[ii]/np.sqrt(2*np.pi))*np.exp(-(2*np.pi*delta[jj]*t_2r[ii])**2/2)
        
        G_lor[jj,ii+1] = (t_2r[ii]/np.pi)*1/(1+((2*np.pi*delta[jj]*t_2r[ii])**2))
        
        integral = quad(integrand, 0, 1, args=(delta[jj],t_2r[ii]))
        G_superlor[jj,ii+1] = (t_2r[ii])*(np.sqrt(2/np.pi))*integral[0]


init_notebook_mode(connected=True)

lineshape1 = [dict(
        visible = False,
        x = G_gaussian[:,0],
        y = G_gaussian[:,ii+1],
        line = dict(color = "firebrick"),
        name = 'Gaussian',
        hovertemplate = 'Gaussian, G(\u0394,T<sub>2r</sub>) = %{y}<br>Frequency offset = %{x} Hz<extra></extra>') for ii in range(len(t_2r))]

lineshape1[4]['visible'] = True

lineshape2 = [dict(
        visible = False,
        x = G_lor[:,0],
        y = G_lor[:,ii+1],
        line = dict(color = "royalblue"),
        name = 'Lorentzian',
        hovertemplate = 'Lorentzian, G(\u0394,T<sub>2r</sub>) = %{y}<br>Frequency offset = %{x} Hz<extra></extra>') for ii in range(len(t_2r))]

lineshape2[4]['visible'] = True

lineshape3 = [dict(
        visible = False,
        x = G_superlor[:,0],
        y = G_superlor[:,ii+1],
        line = dict(color = "orange"),
        name = 'Super Lorentzian',
        hovertemplate = 'Super Lorentzian, G(\u0394,T<sub>2r</sub>) = %{y}<br>Frequency offset = %{x} Hz<extra></extra>') for ii in range(len(t_2r))]

lineshape3[4]['visible'] = True

data = lineshape1 + lineshape2 + lineshape3

steps = []
for i in range(len(t_2r)):
    step = dict(
        method = 'restyle',  
        args = ['visible', [False] * len(lineshape1)],
        label = str(round(t_2r[i], 7))
    )
    step['args'][1][i] = True # Toggle i'th trace to "visible"
    steps.append(step)

sliders = [dict(
    x = 0,
    y = -0.0,
    active = 3,
    currentvalue = {"prefix": "T<sub>2r</sub>: <b>"},
    pad = {"t": 50, "b": 10},
    steps = steps
)]

layout = go.Layout(
    plot_bgcolor='rgba(0,0,0,0)',
    width=580,
    height=450,
    margin=go.layout.Margin(
        l=120,
        r=80,
        b=60,
        t=10,
    ),
    annotations=[
        dict(
            x=0.5004254919715793,
            y=-0.2,
            showarrow=False,
            text='Frequency offset \u0394 (Hz)',
            font=dict(
                family='Times New Roman',
                size=22
            ),
            xref='paper',
            yref='paper'
        ),
        dict(
            x=-0.2,
            y=0.5,
            showarrow=False,
            text='Absorption lineshape G(\u0394,T<sub>2r</sub>)',
            font=dict(
                family='Times New Roman',
                size=22
            ),
            textangle=-90,
            xref='paper',
            yref='paper'
        ),
    ],
    xaxis=dict(
        autorange=False,
        range=[-8000, 8000],
        showgrid=False,
        linecolor='black',
        linewidth=2
    ),
    yaxis=dict(
        autorange=True,
        type="log",
        dtick=1,
        showgrid=False,
        linecolor='black',
        linewidth=2
    ),
    legend=dict(
        x=0.7,
        y=0.85,
        traceorder='normal',
        font=dict(
            family='Times New Roman',
            size=12,
            color='#000'
        ),
        bordercolor='#000000',
        borderwidth=2
    ), 
    sliders=sliders
)

fig = dict(data=data, layout=layout)

iplot(fig, filename = 'basic-line', config = config)