# from PyQt5.QtWidgets.QWidget import sizeHint
from scipy.ndimage import rotate, map_coordinates
import numpy as np
from scipy.special.cython_special import inv_boxcox

'''
    тут собираются по различным траекториям несортированные данные в к-пространство
    
    на выходе хочется получать 5-мерный массив [x, y, z, contrast, coil], чтобы дальнейшие взаимодействие происходило одинаково
'''


def gather_data_along_trajectory(mat_data, nX, nY, nZ, nContr, nCoils=1):
    # out: [x, y, z, contrast, coil]

    gathered_data = np.zeros((nX, nY, nZ, nContr, nCoils), dtype=np.complex128)

    i = 0
    for coil in range(nCoils):
        for y in range(0, nY):
            for z in range(0, nZ):
                for contrast in range(0, nContr):
                    for x in range(0, nX):
                        gathered_data[x, y, z, contrast, coil] = mat_data[i]
                        i += 1
    return gathered_data


def gather_data_along_trajectory_nonlinear(mat_data, order, nX, nY, nZ, nContr, nCoils=1):
    # out: [x, y, z, contrast, coil]

    gathered_data = np.zeros((nX, nY, nZ, nContr, nCoils), dtype=np.complex128)

    N_center = (nY - 1) // 2

    if nY % 2 == 0:
        N_center += 1

    # print('N_center', N_center)
    step_y = order

    # print("STEP TSE", step_y)

    i = 0
    print(len(step_y))
    index_y = len(step_y)
    for coil in range(nCoils):
        for z in range(0, nZ):
            for ind in range(index_y):
                # for z in range(0, nZ):
                    for y in step_y[ind]:
                        # print(y, "step", step_y[ind])
                        if not isinstance(y, str):
                            # for z in range(0, nZ):
                                for x in range(0, nX):
                                    gathered_data[x, y, z, 0, coil] = mat_data[i]
                                    i += 1

                    # print(x, " ", y, " ", z)
    # print("TET")
    return gathered_data


def gather_data_along_trajectory_linear_epi(mat_data, nX, nY, nZ, nContr, nCoils=1):
    # out: [x, y, z, contrast, coil]

    gathered_data = np.zeros((nX, nY, nZ, nContr, nCoils), dtype=np.complex128)

    i = 0
    for coil in range(nCoils):
        for z in range(0, nZ):
            for y in range(0, nY):
                for contrast in range(0, nContr):
                    if y % 2 == 1:
                        for x in range(0, nX):
                            gathered_data[y, x, z, contrast, coil] = mat_data[i]
                            i += 1
                    else:
                        for x in range(nX - 1, -1, -1):
                            gathered_data[y, x, z, contrast, coil] = mat_data[i]
                            i += 1

    return gathered_data


def gather_data_along_trajectory_radial(mat_data, order, nX, nY, nZ, phi_wing, N_wings, nCoils=1):
    # out: [x, y, z, contrast, coil]
    print("traj_01")

    k_space_not_rotate = read_raw_data_radial(mat_data, order, nX, N_wings, nCoils)
    print("traj_02")
    k_space = rotate_radial(k_space_not_rotate, phi_wing, N_wings)
    print("traj_03")
    k_spaces_with_coils = np.stack(k_space, axis=0)
    print("traj_04")
    # k_space = interpolation_radial(arr_k_spaces)
    k_spaces_with_contrast = np.expand_dims(k_spaces_with_coils, axis=1)
    print("traj_05")

    k_spaces = np.transpose(k_spaces_with_contrast, (3, 4, 2, 1, 0))
    print(k_spaces.shape)

    return k_spaces


##########################################  RADIAL  ###############################################################
def rotate_radial(k_space_not_rotate_slices, phi_wing, N_wings):
    # k_space_not_rotate_slices = [coil, slice, n_wing, x, y]

    # print("ROTATE")
    k_spaces = []

    # result = []
    # result_slices = []
    # ind = 0
    # print(len(k_space_not_rotate_slices)) # 4

    for coil in range(k_space_not_rotate_slices.shape[0]):
        result = []
        result_slices = []
        for iSlice in range(k_space_not_rotate_slices.shape[1]):
            # print("iSlice: ", iSlice)
            # k_space_not_rotate = k_space_not_rotate_slices[iSlice]
            k_space_not_rotate = k_space_not_rotate_slices[coil, iSlice]
            ind = 0
            for i in range(0, phi_wing * N_wings, phi_wing):
                # print(ind)
                result.append(rotate(k_space_not_rotate[ind], angle=i, reshape=False, order=2))
                ind += 1
            result_slices.append(result.copy())
        k_space = interpolation_radial(result_slices)
        k_spaces.append(k_space.copy())
        print("len ",len(result))

    return k_spaces


def read_raw_data_radial(mat_data, order, nX, N_wings, nCoils):
    # out: [coil, slice, n_wing, x, y]

    # print(len(mat_data), " ", N_wings, " ", order)

    nSlices = int(len(mat_data) / N_wings / len(order[0]) / nX / nCoils)
    k_space_size = nX

    k_space = np.zeros((nCoils, nSlices, N_wings, k_space_size, k_space_size), dtype=np.complex128)

    center_k_space = k_space_size // 2  # положение по Х, центр меняться не будет, можно учитывать как Мх
    step_lines = order

    # k_spaces_not_rotate = []

    ind = 0
    for coil in range(nCoils):
        # for iSlice in range(nSlices):
        for blade_number in range(0, N_wings, 1):
            for iSlice in range(nSlices):

                for arr_line_number in step_lines:
                    # print(arr_line_number, "-arr_line_number")
                    for line_number in arr_line_number:
                        # print(line_number, "-line_number")
                        if not isinstance(line_number, str):
                            for index in range(nX):
                                # print(index, "-index")
                                radius = center_k_space + line_number - (len(arr_line_number) // 2)
                                k_space[coil, iSlice, blade_number, radius, index] = mat_data[ind]
                                ind += 1

                # k_spaces_not_rotate.append(k_space.copy())
            # k_space_with_slices[iSlice].append(k_spaces_not_rotate.copy())

    # print("k_space: ", len(k_space))
    return k_space


def interpolation_radial(arr_k_spaces_slices):
    # print("INTERPOL")
    result_slices = []

    for iSlice in range(len(arr_k_spaces_slices)):
        arr_k_spaces = arr_k_spaces_slices[iSlice]

        max_shape = tuple(np.max([arr.shape for arr in arr_k_spaces], axis=0))

        result = np.zeros(max_shape, dtype=np.complex128)

        for array in arr_k_spaces:
            scale_factors = [float(s1) / s2 for s1, s2 in zip(max_shape, array.shape)]
            grid = [np.linspace(0, dim - 1, num=int(dim * sf))[:max_shape[i]] for i, (dim, sf) in
                    enumerate(zip(array.shape, scale_factors))]
            coords = np.meshgrid(*grid, indexing="ij")

            interpolated_array = map_coordinates(array, coords, order=1, mode="nearest")
            result += np.nan_to_num(interpolated_array)

        result_slices.append(result.copy())

        # print(result)
    return result_slices