import numpy as np
from scipy.optimize import least_squares
from brukerapi.dataset import Dataset
import matplotlib.pyplot as plt

# load data as dataset using brukerapi.dataset
dataset = Dataset(r"C:\Users\hanap\brukerapi-python-master\brukerapi\20220427_123935_hpbrainingel_rj_rlx02_1_2\5\pdata\1\2dseq")
snrmaskcensus = 10 # noise multiple to be exceeded
snrmaskquota  = 2  # ... by at least this many samples


# DATA MODEL FUNCTION
def fun12(x, tx, te, v):
    """Function for simultaneous model fitting by function y(t) = a*(1-exp(-b1*tx)*exp(-b2*te))+c
    inputs: x - starting parameters of fitted values, x = [a,b1, b2,c]
            tx - list of TRs, t = [TR1-TEK, ..., TRn-TEK], TEK - TE[max_index]
            te - list of TEs, t = [TE1, ..., TEn]
            v - measured values, v = [v1, ..., vn]
    output: difference between fitted and measured values"""
    return (x[0] * (1 - np.exp((-tx) * x[1])) * np.exp((-te) * x[2]) + x[3]) - v
#_enddef


#%% AUXILIARY FUNCTIONS
def cutoutliers(image,margin=0.05,eps=0.001):
    imagemax = image.max()
    imagecount = image.size
    v = imagemax/2
    dv = v/2
    while True:
        boutlier = (image>=v)
        noutlier = boutlier.sum()
        noutlierfract = noutlier/imagecount
        if noutlierfract<margin:
            v = v - dv
            dv = dv/2
        else:
            v = v + dv
            dv = dv/2
        #_endif
        if (abs(noutlierfract - margin) < eps) or (dv <= 0.001 * imagemax):
            break
        #_endif
    #_endwhile
    return v
#_enddef


def figoverview(T_s,npar,title,p0,p1=0,p2=0,p3=0):
    nT = T_s.size
    w = p0.shape[0]
    h = p0.shape[1]
    threshold = np.zeros((nT,npar))
    if npar >= 1:
        image0 = np.zeros((w, h))   # parameter 0: a (for one value of TR)
    if npar >= 2:
        image1 = np.zeros((w, h))   # parameter 1: b (for one value of TR)
    if npar >= 3:
        image2 = np.zeros((w, h))   # parameter 2: c (for one value of TR)
    if npar >= 4:
        image3 = np.zeros((w, h))   # parameter 2: c (for one value of TR)
    # image grid rows=Ts, columns=parameters
    fig = plt.figure()
    c = 1 # image index (starting at 1)
    for n in range(nT):
        for i in range(w):
            for j in range(h):
                if npar >= 1:
                    image0[i][j] = p0[i][j][n]
                if npar >= 2:
                    image1[i][j] = p1[i][j][n]
                if npar >= 3:
                    image2[i][j] = p2[i][j][n]
                if npar >= 4:
                    image3[i][j] = p3[i][j][n]
            #_endfor
        #_endfor
        if npar >= 1:
            th = cutoutliers(image0,0.005,0.00005)
            threshold[n,0] = th
            fig.add_subplot(nT, npar, c)
            plt.imshow(image0, cmap='gray', vmin=0, vmax=th)
            plt.axis('off')
            plt.colorbar()
            c += 1
        #_endif
        if npar >= 2:
            th = cutoutliers(image1,0.001,0.00005)
            threshold[n,1] = th
            fig.add_subplot(nT, npar, c)
            plt.imshow(image1, cmap='gray', vmin=0, vmax=th)
            plt.axis('off')
            plt.colorbar()
            c += 1
        #_endif
        if npar >= 3:
            th = cutoutliers(image2,0.001,0.00005)
            threshold[n,2] = th
            fig.add_subplot(nT, npar, c)
            plt.imshow(image2, cmap='gray', vmin=0, vmax=th)
            plt.axis('off')
            plt.colorbar()
            c += 1
        #_endif
        if npar >= 4:
            th = cutoutliers(image3,0.001,0.00005)
            threshold[n,3] = th
            fig.add_subplot(nT, npar, c)
            plt.imshow(image3, cmap='gray', vmin=0, vmax=th)
            plt.axis('off')
            plt.colorbar()
            c += 1
        #_endif

    #_endfor
    plt.suptitle(title, fontweight="bold")
    plt.show()
    return threshold
#_enddef


def fmaplin(val, binv=0):
    return val
#_enddef


def fmapsqrt(val, binv=0):
    if binv == 0:
        return np.sqrt(val)
    else:
        return val**2
    #_endif
#_enddef


def figonepar(T_s,fitpar,threshold,title,fmap):
    # size(fitpar)=[w,h], size(TR_s)=[nTR], threshold[:,fitparno]
    nT = T_s.size
    w = fitpar.shape[0]
    h = fitpar.shape[1]
    image = np.zeros((w, h))  # any parameter, one value of T
    fig = plt.figure()
    c = 1
    for n in range(nT):
        for i in range(w):
            for j in range(h):
                image[i][j] = fitpar[i][j][n]
            # _endfor
        # _endfor
        fig.add_subplot(1, 1, c)  # OK only for nT==1
        plt.imshow(fmap(image), cmap='gray', vmin=0, vmax=fmap(threshold[n]))
        plt.axis('off')
        cb = plt.colorbar()
        cbticks = cb.get_ticks()
        cblabels = [f'{u:.2f}' for u in fmap(cbticks,1)]
        cb.set_ticks(cbticks)
        #plt.title(f'for TR {c} = {T_s[c-1]:.3f} s')
        c += 1
    # _endfor
    plt.suptitle(title, fontweight="bold")
    plt.show()
#_enddef


#%% DO FITTING
# define parameters of dataset (width, height, TE, TR and their index in dataset)
# RAREVTR produces data with dimensions [r_read, r_phase, r_slice, TE, TR]
num_of_dim = dataset.dim
w = int(dataset.shape[0])   # image size in dimension 0 (readout)
h = int(dataset.shape[1])   # image size in dimension 1 (phase)
for num in range(2, num_of_dim):
    if dataset.dim_type[num] == '<FG_ECHO>':
        TE = dataset.TE     # vector of echo times TE
        TE_dim = num        # index of the TE dimension
    #_endif
    if dataset.dim_type[num] == '<FG_MOVIE>':
        TR = dataset.TR     # vector of repetition times TR
        TR_dim = num        # index of the TR dimension
    #_endif
#_endfor

# transformation of TE and TR parameters
nTE = len(TE)       # number of echo times TE
nTR = len(TR)       # number of repetition times TR
TE_s = TE/1000      # TE in seconds
TR_s = TR/1000      # TR in seconds

# image matrices of estimated parameters and their deviations, one for each TR
par_0 = np.zeros((w, h, 1))   # a
par_1 = np.zeros((w, h, 1))   # r1
par_2 = np.zeros((w, h, 1))   # r2
par_3 = np.zeros((w, h, 1))   # c
dev_0 = np.zeros((w, h, 1))   # stdev(a)
dev_1 = np.zeros((w, h, 1))   # stdev(r1)
dev_2 = np.zeros((w, h, 1))   # stdev(r2)
dev_3 = np.zeros((w, h, 1))   # stdev(c)
crb_0 = np.zeros((w, h, 1))   # crlb(a)
crb_1 = np.zeros((w, h, 1))   # crlb(r1)
crb_2 = np.zeros((w, h, 1))   # crlb(r2)
crb_3 = np.zeros((w, h, 1))   # crlb(c)

max_val = dataset.data.max()
# signal noise and standard deviation estimation
edge = np.reshape(dataset.data[1:3, :, 0, :, :], 2*h*nTE*nTR)                             # left edge values
edge = np.append(edge, np.reshape(dataset.data[(w-3):(w-1), :, 0, :, :], 2*h*nTE*nTR))    # right edge values
#edge = np.append(edge, np.reshape(dataset.data[:, 1:3, 0, :, :], 2*w*nTE*nTR))           # bottom edge values
#edge = np.append(edge, np.reshape(dataset.data[:, (h-3):(h-1), 0, :, :], 2*w*nTE*nTR))   # top edge values
bgndavg = np.mean(edge)             # mean of magnitudes at the edge for noise assessment, sigma*sqrt(pi/2)
bgndstd = np.std(edge)              # stdev of magnitudes at the edge could confirm Ricean distribution, sigma*sqrt(2-pi/2)
noise = bgndavg/np.sqrt(np.pi/2)    # estimate of single-channel (Re/Im) noise from zero-background noise
avg2std = bgndavg/bgndstd
avg2stdRice = 1/np.sqrt(4/np.pi-1)
print(f'max={max_val:.4f}, mean={bgndavg:.4f}, stdev={bgndstd:.4f}, sigma={noise:.4f}, mean/stdev={avg2std:.4f} (expected {avg2stdRice:.4f})')

# visual inspection will help avoid noise misassessment; find a method that doesn't lock
fig = plt.figure()
fig.add_subplot(1, 1, 1)
plt.plot(edge, 'r.', markersize=1)
plt.axis('on')
plt.title(f'mean={bgndavg:.4f}, stdev={bgndstd:.4f}, sigma={noise:.4f}')
plt.suptitle("edge values for noise assessment", fontweight="bold")
plt.draw()
plt.show()

for z in range(w):
    print(f'{z} / {w}')
    for y in range(h):
            # data value for actual position (slice in TE-TR dimensions, reshaped to vector)
            v = dataset.data[z, y, 0, :, :].reshape(nTR*nTE)
            # condition for fitting only data with proven SNR
            if (v > snrmaskcensus*noise).sum() >= snrmaskquota:
                # recovery-time and echo-time matrix size modification
                TX_s = TR_s - TE_s[nTE - 1]
                mtx = np.multiply(np.ones(nTE)[:, None], TX_s).reshape(nTR * nTE)
                mte = np.multiply(TE_s[:, None], np.ones(nTR)).reshape(nTR * nTE)
                npar = 4
                x0 = [1.0, 1.0, 1.0, 0.0]  # starting parameters for fitting
                # fitting function y(t) = a*(1-exp(-b1*tx)*exp(-b2*te))+c
                res = least_squares(fun12, x0, args=(mtx, mte, v), bounds=([0, 0, 0, 0], [500, 500, 500, 50]))
                # ==========
                par_0[z, y] = res.x[0]
                par_1[z, y] = res.x[1]
                par_2[z, y] = res.x[2]
                par_3[z, y] = res.x[3]

                # ---------- stdev estimation - locally linearized model
                D = np.zeros((nTE * nTR, npar))  # partial derivatives (one pixel, one TR):
                                                 # D[s,p] = d(datasamplemodel[s])/d(parameter[p])
                # y = a*(1-exp(-r1*tx))*exp(-r2*te)+c
                # dy/da = (1-exp(-r1*tx))*exp(-r2*te)
                # dy/dr1 = a*tx*exp(-r1*tx)*exp(-r2*te) = a*tx*exp(-r1*tx-r2*te)
                # dy/dr2 = -a*te*(1-exp(-r1*tx))*exp(-r2*te)
                # dy/dc = 1
                D[:, 0] = (1 - np.exp(-mtx * res.x[1])) * np.exp(-mte * res.x[2])
                D[:, 1] = res.x[0] * mtx * np.exp(-res.x[1] * mtx) * np.exp(-res.x[2] * mte)
                D[:, 2] = -res.x[0] * mte * (1 - np.exp(-res.x[1] * mtx)) * np.exp(-res.x[2] * mte)
                D[:, 3] = 1


                U, s, Vh = np.linalg.svd(D, full_matrices=False)  # as if D = U*diag(s)*Vh
                smax = abs(s.max())
                bsnonzero = (abs(s) > smax * 1e-5)
                s_ = s[bsnonzero]
                U_ = U[:, bsnonzero]
                Vh_ = Vh[bsnonzero, :]
                S_ = np.diag(s_)
                B = np.matmul(np.matmul(np.transpose(Vh_), np.linalg.inv(S_)), np.transpose(U_))
                C = noise * np.sqrt(np.sum(B ** 2, axis=1))
                dev_0[z, y] = C[0]
                dev_1[z, y] = C[1]
                dev_2[z, y] = C[2]
                dev_3[z, y] = C[3]

                # ----------- stdev estimation - Cramer-Rao lower bounds
                Dh = np.transpose(D)
                F = np.matmul(Dh, D)  # Fisher*noise^2
                try:
                    Fi = np.linalg.inv(F)  # Fisher inverse has bounds on the diagonal
                    crb_0[z, y] = noise * np.sqrt(Fi[0][0])
                    crb_1[z, y] = noise * np.sqrt(Fi[1][1])
                    crb_2[z, y] = noise * np.sqrt(Fi[2][2])
                    crb_3[z, y] = noise * np.sqrt(Fi[3][3])
                except:  # negative ones signal that something is wrong
                    crb_0[z, y] = -1
                    crb_1[z, y] = -1
                    crb_2[z, y] = -1
                    crb_3[z, y] = -1
                try:
                    # Fisher inverse has bounds on the diagonal
                    Fi = 0                      # fixed background for debugging
                    Fi = np.linalg.inv(F)       # inv as default
                    if any(np.diag(Fi) < 0):
                        Fi = np.linalg.pinv(F, 1e-7) # pinv as substitute if diagonal negative (which is strange)
                    # _endif
                except:
                    Fi = np.linalg.pinv(F, 1e-7)     # pinv as substitute if inv failed
                #_endtry
                if Fi[0][0] < 0:
                    print(f'invFisher[{z},{y},0]={Fi[0][0]}<0')
                #_endif
                crb_0[z, y, 0] = noise * np.sqrt(abs(Fi[0][0]))
                if Fi[1][1] < 0:
                    print(f'invFisher[{z},{y},1]={Fi[1][1]}<0')
                #_endif
                crb_1[z, y, 0] = noise * np.sqrt(abs(Fi[1][1]))
                if Fi[2][2] < 0:
                    print(f'invFisher[{z},{y},2]={Fi[2][2]}<0')
                #_endif
                crb_2[z, y, 0] = noise * np.sqrt(abs(Fi[2][2]))
                if (Fi[3][3] < 0):
                    print(f'invFisher[{z},{y},3]={Fi[3][3]}<0')
                # _endif
                crb_3[z, y, 0] = noise * np.sqrt(abs(Fi[3][3]))
            #_endif
        #_endfor
    #_endfor
#_endfor

#%% MAKE GRAPHS
NA = np.zeros((1,))
#---- parameter values
# FIGURE 1: all fitted results, image grid rows=TRs, columns=parameters a, b, c
threshold = figoverview(NA,npar,"parameter values (A, R1, R2, C)",par_0,par_1,par_2,par_3)
if npar >= 1:
    # FIGURE 2: parameter a, all TRs
    figonepar(NA,par_0,threshold[:,0],"param A",fmaplin)
if npar >= 2:
    # FIGURE 3: parameter b, all TRs
    figonepar(NA,par_1,threshold[:,1],"param 1/T1",fmaplin)
if npar >= 3:
    # FIGURE 4: parameter b, all TRs
    figonepar(NA,par_2,threshold[:,2],"param 1/T2",fmaplin)
if npar >= 4:
    # FIGURE 5: parameter c, all TRs
    figonepar(NA,par_3,threshold[:,3],"param C",fmapsqrt)

#----- parameter stdevs
# FIGURE 6: all stdev estimates, image grid rows = TRs, columns = parameters a, b, c
threshold = figoverview(NA,npar,"parameter deviations (A, R1, R2, C)",dev_0,dev_1,dev_2,dev_3)
if npar >= 1:
    # FIGURE 7: stdev(a)
    figonepar(NA,dev_0,threshold[:,0],"deviation of param A",fmapsqrt)
if npar >= 2:
    # FIGURE 8: stdev(b)
    figonepar(NA,dev_1,threshold[:,1],"deviation of param 1/T1",fmapsqrt)
if npar >= 3:
    # FIGURE 9: stdev(b)
    figonepar(NA,dev_2,threshold[:,2],"deviation of param 1/T2",fmapsqrt)
if npar >= 4:
    # FIGURE 10: stdev(c)
    figonepar(NA,dev_3,threshold[:,3],"deviation of param C",fmapsqrt)

#----- parameter CRLBs
# FIGURE 11: all CRLB estimates, image grid rows = TRs, columns = parameters a, b, c
#threshold = figoverview(TR_s,npar,"parameter CRLBs",crb_0,crb_1,crb_2)
figoverview(NA,npar,"parameter CRLBs (A, R1, R2, C)",crb_0,crb_1,crb_2,crb_3)
if npar >= 1:
    # FIGURE 12: crlb(a)
    figonepar(NA,crb_0,threshold[:,0],"CRLB of param A",fmapsqrt)
if npar >= 2:
    # FIGURE 13: crlb(b)
    figonepar(NA,crb_1,threshold[:,1],"CRLB of param 1/T1",fmapsqrt)
if npar >= 3:
    # FIGURE 14: crlb(b)
    figonepar(NA,crb_2,threshold[:,2],"CRLB of param 1/T2",fmapsqrt)
if npar >= 4:
    # FIGURE 15: crlb(c)
    figonepar(NA,crb_3,threshold[:,3],"CRLB of param C",fmapsqrt)


