# 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]

seed = 2
np.random.seed(seed)
n_points = 20
len_x = 500
len_y = 500
map = np.zeros([len_x, len_y]) - math.pi
x, y = np.meshgrid(range(len_x), range(len_y))

def gauss(a, std_x, std_y, center):
    x0 = int(center[0] * len_x)
    y0 = int(center[1] * len_y)
    return a * np.exp(-((((x-x0)**2)/(std_x**2)) + (((y-y0)**2)/(std_y**2))))


for i_point in range(n_points):
    point = (np.random.uniform(), np.random.uniform())
    amp = np.random.randint(4, 9)
    std_x = np.random.randint(50, 100)
    std_y = np.random.randint(50, 100)
    map += gauss(amp, std_x, std_y, point)

wrapped = np.angle(np.exp(1j*map))

fig = make_subplots(rows=1, cols=2, shared_xaxes=False, horizontal_spacing=0.2, subplot_titles=("Wrapped", "Unwrapped"), specs=[[{"type": "Heatmap"}, {"type": "Heatmap"}]])
fig.add_trace(go.Heatmap(z=wrapped, colorscale='gray', colorbar_x=0.47, colorbar=dict(title="Rad", titleside="top", tickmode="array", tickvals=[-math.pi, 0, math.pi], ticktext = [f'-{PI_UNICODE}', 0, f'{PI_UNICODE}'])), 1, 1)
fig.add_trace(go.Heatmap(z=map, colorscale='gray', colorbar=dict(title="Rad", titleside="top", tickmode="array",
                                                                 tickvals=[-2*math.pi, 0, 2*math.pi, 4*math.pi, 6*math.pi, 8*math.pi, 10*math.pi, 12*math.pi, 14*math.pi],
                                                                 ticktext = [f'-2{PI_UNICODE}', 0, f'2{PI_UNICODE}', f'4{PI_UNICODE}', f'6{PI_UNICODE}', f'8{PI_UNICODE}', f'10{PI_UNICODE}', f'12{PI_UNICODE}', f'14{PI_UNICODE}'])), 1, 2)
fig.update_layout({"height": 400, "width": 750})
fig.update_xaxes(showticklabels=False)
fig.update_yaxes(showticklabels=False)
fig.show()