# 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
import copy
PI_UNICODE = "\U0001D70B"
DELTA_UNICODE = "\u0394"
GYRO_UNICODE = "\U0001D6FE"
GREEK_DELTA_UNICODE = "\u03B4"

phase1, phase2, phase3 = (-4, -2, 0)
beg = 0
end = 0.01
t = np.array([0.00263, 0.00526, 0.009])
n_echoes = len(t)
y1 = np.array([phase1, phase2, phase3])
y2 = y1 + 2 * math.pi

# Linear fit of the first line
reg1 = LinearRegression().fit(t.reshape(-1, 1), y1.reshape(-1,1))
fieldmap_rad1 = reg1.coef_[0]  # [rad / s]
fieldmap_intercept1 = reg1.intercept_[0]  # [rad / s]
t_predict1 = np.array([beg, end])
y_predict1 = reg1.predict(t_predict1.reshape(-1,1))[:,0]

# Linear fit of the second line
reg2 = LinearRegression().fit(t.reshape(-1, 1), y2.reshape(-1,1))
fieldmap_rad2 = reg2.coef_[0]  # [rad / s]
fieldmap_intercept2 = reg2.intercept_[0]  # [rad / s]
t_predict2 = np.array([beg, end])
y_predict2 = reg2.predict(t_predict2.reshape(-1,1))[:,0]

# Plot
fig = go.Figure()
fig.add_trace(go.Scatter(x=t, y=y1, mode='markers', marker=dict(color='blue'), name='Unwrapped solution 1'))
fig.add_trace(go.Scatter(x=t_predict1, y=y_predict1, mode='lines', marker=dict(color='blue'), name='Fit'))
fig.add_trace(go.Scatter(x=t, y=y2, mode='markers', marker=dict(color='red'), name='Unwrapped solution 2'))
fig.add_trace(go.Scatter(x=t_predict2, y=y_predict2, mode='lines', marker=dict(color='red'), name='Fit'))
fig.add_annotation(x=0.004, y=-3, text=f"Slope: {GYRO_UNICODE}{DELTA_UNICODE}B*t")
fig.add_annotation(x=0.004, y=3.3, text=f"Slope: {GYRO_UNICODE}{DELTA_UNICODE}B*t")
fig.update_xaxes(title_text="Time (ms)", range=[beg, end])
fig.update_yaxes(title_text="Phase (rad)", tickmode = 'array', range=[-8,8],
                 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()