# 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 = (-1, 2, 14)
phase_wrapped1 = 14-4*math.pi
phase_unwrapped1 = 14-2*math.pi
beg = 0
end = 0.016

t1 = np.array([0.00263, 0.00526])
t2 = np.array([0.00263, 0.00526, 0.015])
n_echoes1 = len(t1)
n_echoes2 = len(t2)
y1 = np.array([phase1, phase2])
y2 = np.array([phase1, phase2, phase3])

reg1 = LinearRegression().fit(t1.reshape(-1, 1), y1.reshape(-1,1))
# Slope of linear regression reshaped into the shape of original 3D phase.
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]

reg2 = LinearRegression().fit(t2.reshape(-1, 1), y2.reshape(-1,1))
# Slope of linear regression reshaped into the shape of original 3D phase.
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]

fig = go.Figure()
fig.add_trace(go.Scatter(x=t2, y=y2, mode='markers', marker=dict(color='green'), name='Unwrapped phase'))
fig.add_trace(go.Scatter(x=[t2[2]], y=[phase_wrapped1], mode='markers', marker=dict(color='blue'), name='Acquired wrapped phase'))
fig.add_trace(go.Scatter(x=[t2[2]], y=[phase_unwrapped1], mode='markers', marker=dict(color='orange'), name='Wrong unwrapped phase'))
fig.add_trace(go.Scatter(x=t_predict1, y=y_predict1, mode='lines', marker=dict(color='blue'), name='Fit with first 2 echoes'))
fig.add_trace(go.Scatter(x=t_predict2, y=y_predict2, mode='lines', marker=dict(color='green'), name='Fit with 3 echoes'))
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=[-5,15],
                 tickvals = [-2*math.pi, -math.pi, 0, math.pi, 2*math.pi, 3*math.pi, 4*math.pi, 5*math.pi],
                 ticktext = [f'-2{PI_UNICODE}', f'-{PI_UNICODE}', '0', f'{PI_UNICODE}', f'2{PI_UNICODE}', f'3{PI_UNICODE}', f'4{PI_UNICODE}', f'5{PI_UNICODE}'])
fig.update_layout({"width": 800}, title_text=f"Using fast {DELTA_UNICODE}TE to help temporally unwrapped 3rd echo", title_x=0.5)

fig.show()