# Prepare Python environment
import numpy as np
import plotly.express as px
import os
import nibabel as nib
import numpy as np
from plotly.subplots import make_subplots
import plotly.graph_objects as go
import math
from sklearn.linear_model import LinearRegression
from dash import Dash, dcc, html
import pandas as pd
PI_UNICODE = "\U0001D70B"
GYRO_BAR_RATIO_H = 42.6e6  # [Hz/T]

# Taken from real brain data
phase1, phase2, phase3, phase4 = (2.9298516102709202, -0.21864564255753116, -2.2884910587688285, 2.392827225041896)
beg = 0
end = 0.015

t = np.array([0.00263, 0.00526, 0.009, 0.013])
n_echoes = len(t)
y = np.array([phase1, phase2, phase3, phase4])
y_unwrapped = np.array([phase1, phase2, phase3, phase4 - (2 * math.pi)])

reg = LinearRegression().fit(t.reshape(-1, 1), y_unwrapped.reshape(-1,1))
# Slope of linear regression reshaped into the shape of original 3D phase.
fieldmap_rad = reg.coef_[0]  # [rad / s]
fieldmap_intercept = reg.intercept_[0]  # [rad / s]

t_predict = np.array([beg, end])
y_predict = reg.predict(t_predict.reshape(-1,1))[:,0]


fig = go.Figure()
fig.add_trace(go.Scatter(x=t, y=y, mode='markers', marker=dict(color='blue'), name='Original Phase'))
fig.add_trace(go.Scatter(x=[t[3]], y=[phase4- (2*math.pi)], mode='markers', marker=dict(color='red'), name='Unwrapped Phase'))
fig.add_trace(go.Scatter(x=t_predict, y=y_predict, mode='lines', marker=dict(color='green'), name='Fit'))
fig.add_trace(go.Scatter(x=[beg, end], y=[math.pi, math.pi], mode='lines', line=dict(color='gray'), showlegend=False))
fig.add_trace(go.Scatter(x=[beg, end], y=[-math.pi, -math.pi], mode='lines', line=dict(color='gray'), showlegend=False))

fig.update_xaxes(title_text="Time (ms)", range=[beg, end])
fig.update_yaxes(title_text="Phase (rad)", tickmode = 'array', range=[-7,7],
                 tickvals = [-2*math.pi, -math.pi, 0, math.pi, 2*math.pi],
                 ticktext = [f'-2{PI_UNICODE}', f'-{PI_UNICODE}', '0', f'{PI_UNICODE}', f'2{PI_UNICODE}'])
fig.update_layout({"width": 800})

fig.show()