import scipy.io as sio
import matplotlib.pyplot as plt
import numpy as np
from scipy.optimize import curve_fit
import sys,os
from cfl import writecfl

os.environ['TOOLBOX_PATH'] = '/home/jpfitzer/bart-0.9.00'
os.environ['BART_TOOLBOX_PATH'] = '/home/jpfitzer/bart-0.9.00'
sys.path.append('/home/jpfitzer/bart-0.9.00/bart/python')

rawName = 'T2_CPMG.mat'
mat_data_0=sio.loadmat(rawName)

# Format: sampled, # kSpace [kRd, kPh, kSl, kSpace_echo_1, kSpace_echo_2, ..., kSpace_echo_nETL] (102400, 23)
# So the first three are the coordinates of the kspace, and the rest are the echoes
print(mat_data_0['kSpaces3D'].shape)

kSpaces3D = mat_data_0['kSpaces3D']
# self.mapVals['sampled'] = np.concatenate((kRD, kPH, kSL, dataAll_sampled), axis=1)

# nReadout, nPhase, nSlice
nPoints = (80, 80, 16)

echo_train_length = mat_data_0['kSpaces3D'].shape[1] - 3 # Because the first 3 are kRD, kPH, kSL -> should give 20
print(f"Echo train length: {echo_train_length}")

echo_spacing = mat_data_0['echoSpacing'][0][0]
print(f"Echo spacing: {echo_spacing}")

k_readout = kSpaces3D[:, 0]
k_phase = kSpaces3D[:, 1]
k_slice = kSpaces3D[:, 2]

# The rest of the data is the echoes
echos = kSpaces3D[:, 3:]

# Reshape the kspace data for bart
kSpace = echos.reshape(nPoints[2], nPoints[1], nPoints[0], echo_train_length)



print(kSpace.shape)

# cfl = writecfl('kSpace', kSpace)

# Create the image with bart fft -i 7 kSpace fft -> three dimensional
# Put the echos on the fifth dimension:
# bart transpose 3 5
# Put the slices on the correct dimension
# bart transpose 0 2

# traj = writecfl('traj', np.stack((k_readout, k_phase, k_slice), axis=1))

# Echo times with echo spacing
TE = np.linspace(echo_spacing, echo_spacing * echo_train_length, echo_train_length, endpoint=True)
print("TE: ", TE)

# Create the echotimes file:
# bart vec ... echo_times
# bart scale 0.001 echo_times echo_times_scaled
# Move the echo_times to the correct dimension:
# bart transpose 0 5 echo_times_scaled echo_times_final

# Fit the model:
# bart mobafit -T echo_times_final fft_transposed fit

# Now some values will be very large so we can apply a threshold to obtain a mask
# bart threshold -M 1000 reco/fit reco/mask

# Multiply the fit with the mask
# bart fmac reco/fit reco/fit reco/fit_mask

# Select slice
# bart slice 6 1 reco/fit_mask reco/R2_map

# Invert the data to get T2
# bart invert reco/R2_map reco/T2_map