# 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
def complex_difference(phase1, phase2):
    """ Calculates the complex difference between 2 phase arrays (phase2 - phase1)

    Args:
        phase1 (numpy.ndarray): Array containing phase data in radians
        phase2 (numpy.ndarray): Array containing phase data in radians. Must be the same shape as phase1.

    Returns:
        numpy.ndarray: The difference in phase between each voxels of phase2 and phase1 (phase2 - phase1)
    """

    # Calculate phasediff using complex difference
    comp_0 = np.exp(-1j * phase1)
    comp_1 = np.exp(1j * phase2)
    return np.angle(comp_0 * comp_1)

def umpire_3echoes(phases, times):
    """
    This function performs unwrapping using the UMPIRE algorithm with 3 echoes. UMPIRE requires echo times that are unevenly spaced.
    """
    
    # Complex difference
    dpTE2 = complex_difference(phases[1], phases[2])
    dpTE1 = complex_difference(phases[0], phases[1])
    dpd = complex_difference(dpTE1, dpTE2)
    # print("Diff in phase diff:" , dpd)
    dTEs = np.array([times[1]-times[0], times[2]-times[1]])
    dt_dpd = dTEs[1] - dTEs[0]
    
    # Slope
    slope = dpd / dt_dpd
    
    # n wraps in differences
    n_wraps_dp = np.round((dTEs - dTEs*slope) / (2*math.pi))
    
    # Remove wraps in differences
    dpTE1_prime = dpTE1 - (2*n_wraps_dp[0]*math.pi)
    dpTE2_prime = dpTE2 - (2*n_wraps_dp[1]*math.pi)
    
    # Calculate better slope
    slope_prime1 = dpTE1_prime / dTEs[0]
    slope_prime2 = dpTE2_prime / dTEs[1]
    slope_avg = (slope_prime1 + slope_prime2) / 2
    
    # Calculate wraps in original phase
    n_wraps = np.round((phases - t*slope_avg) / (2*math.pi))
    
    # Remove wraps
    unwrapped_with_phase_offset = phases - 2*math.pi*n_wraps
    
    # # Calculate receiver offset
    # r = (t[0] * unwrapped_with_phase_offset[1] - t[1] * unwrapped_with_phase_offset[0]) / dTEs[0]

    # # Remove receiver phase offset
    # phase_no_offset = complex_difference(r, unwrapped_with_phase_offset)
    # # Unwrap one last time
    # ns = np.round((phase_no_offset - t*slope_avg) / (2*math.pi))
    # unwrapped_umpire = phase_no_offset - 2*math.pi*ns
    
    return unwrapped_with_phase_offset

t = np.array([0.003, 0.011, 0.020])
y_unwrapped = np.array([1.0, 9.05, 17.75])
wrapped = copy.deepcopy(y_unwrapped)
wrapped[0] = np.angle(np.exp(1j*wrapped[0]))
wrapped[1] = np.angle(np.exp(1j*wrapped[1]))
wrapped[2] = np.angle(np.exp(1j*wrapped[2]))
beg = 0.0
end = 0.021

# Fit original data
reg1 = LinearRegression().fit(t.reshape(-1, 1), y_unwrapped.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]

# Unwrap with UMPIRE
unwrapped_umpire = umpire_3echoes(wrapped, t)

# Fit unwrapped data of UMPIRE
reg2 = LinearRegression().fit(t.reshape(-1, 1), unwrapped_umpire.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]

t = np.array([0.003, 0.011, 0.020])
y_unwrapped = np.array([1.0, 9.05, 17.75])
wrapped = copy.deepcopy(y_unwrapped)
wrapped[0] = np.angle(np.exp(1j*wrapped[0]))
wrapped[1] = np.angle(np.exp(1j*wrapped[1]))
wrapped[2] = np.angle(np.exp(1j*wrapped[2]))
beg = 0.0
end = 0.021

fig = go.Figure()
noises = np.arange(-0.5, 0.51, 0.01)
# Add traces, one for each slider step
for noise in noises:
    # Noisy unwrapped data
    unwrapped_noisy = copy.deepcopy(y_unwrapped)
    unwrapped_noisy[1] += noise
    fig.add_trace(
        go.Scatter(
            visible=False,
            mode='markers',
            marker=dict(color='red'),
            name="True Phase",
            x=t,
            y=unwrapped_noisy))

    # Fit of noisy unwrapped data
    reg1 = LinearRegression().fit(t.reshape(-1, 1), unwrapped_noisy.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]
    fig.add_trace(go.Scatter(visible=False, x=t_predict1, y=y_predict1, mode='lines', marker=dict(color='red'), name='True linear fit'))

    # Noisy wrapped data
    wrapped_noisy = copy.deepcopy(unwrapped_noisy)
    wrapped_noisy[0] = np.angle(np.exp(1j*wrapped_noisy[0]))
    wrapped_noisy[1] = np.angle(np.exp(1j*wrapped_noisy[1]))
    wrapped_noisy[2] = np.angle(np.exp(1j*wrapped_noisy[2]))
    fig.add_trace(
        go.Scatter(
            visible=False,
            mode='markers',
            marker=dict(color='blue'),
            name="Wrapped Phase",
            x=t,
            y=wrapped_noisy))

    # UMPIRE
    unwrapped_umpire = umpire_3echoes(wrapped_noisy, t)
    fig.add_trace(
        go.Scatter(
            visible=False,
            x=t, y=unwrapped_umpire, mode='markers', marker=dict(color='green'), name='Umpire'))

    # Fit unwrapped data of UMPIRE
    reg2 = LinearRegression().fit(t.reshape(-1, 1), unwrapped_umpire.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.add_trace(go.Scatter(visible=False, x=t_predict2, y=y_predict2, mode='lines', marker=dict(color='green'), name='Umpire fit'))
    
active = len(noises)//2
fig.data[active].visible = True
fig.data[active+1].visible = True
fig.data[active+2].visible = True
fig.data[active+3].visible = True
fig.data[active+4].visible = True

# Static plot
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=[-4,25],
                 tickvals = [-2*math.pi, 0, 2*math.pi, 4*math.pi, 6*math.pi],
                 ticktext = [f'-2{PI_UNICODE}', '0', f'2{PI_UNICODE}', f'4{PI_UNICODE}', f'6{PI_UNICODE}'])

# Create and add slider
phase_offsets = [f"{i:.2}" for i in noises]
steps = []
for i in range(len(noises)):
    step = dict(
        method="update",
        label=phase_offsets[i],
        args=[{"visible": [False] * 5*len(noises) + [True] * (len(fig.data) - 5*len(noises))}],  # layout attribute
    )
    step["args"][0]["visible"][5*i] = True
    step["args"][0]["visible"][5*i + 1] = True
    step["args"][0]["visible"][5*i + 2] = True
    step["args"][0]["visible"][5*i + 3] = True
    step["args"][0]["visible"][5*i + 4] = True
    steps.append(step)

sliders = [dict(
    active=active,
    currentvalue={"prefix": "2nd echo phase offset (rad): "},
    pad={"t": 50},
    steps=steps
)]

fig.update_layout({"width": 800},
                 sliders=sliders,
                 title_text="Effect of noise using UMPIRE phase unwrapping", title_x=0.5)

fig.show()