import numpy as np
import os
import matplotlib
matplotlib.use('Agg')  # headless backend — must be set before pyplot import
import matplotlib.pyplot as plt
import sys


# os.chdir('C:/')
# sys.path.append("C:/Users/iuliia/recoUI/serv")
def reconstruction(matrix,k_space_min, k_space_max, nCoil, nContrast, nSlice, phase_shift=0, return_image=True):
    
    for i in range(matrix.shape[1]):
        phase_factors = np.exp(-1j * phase_shift) ** i
        matrix[:, i] = matrix[:, i] * phase_factors

    fft_data = np.fft.ifft2(matrix)
    fft_shifted_data = np.fft.ifftshift(fft_data)
    amplitude = np.abs(fft_shifted_data)

    plt.imshow(amplitude, cmap='gray', vmin=amplitude.min(), vmax=amplitude.max())
    plt.colorbar()
    plt.savefig(
        os.path.join('data', f'im_contrast_{nCoil}_{nContrast}_k_{nSlice}.png'))
    plt.close()

    log_matrix = np.log(np.abs(matrix))
    print(k_space_min, " ", k_space_max)
    plt.imshow(log_matrix, cmap='gray', vmin=np.log(np.abs(k_space_min)), vmax=np.log(np.abs(k_space_max)))
    plt.colorbar()
    plt.savefig(
        os.path.join('data', f'k_contrast_{nCoil}_{nContrast}_k_{nSlice}.png'))
    plt.close()

    if return_image:
        return amplitude


def reco_3d(matrix, maximum_k_space, minimum_k_space, nCoil, nContrast, phase_shift=0, return_image=True):
    '''
        :param matrix: [x,y,z]
        :param nContrast: количество контрастов
        :param phase_shift: угол
        :return: 
    '''

    for i in range(matrix.shape[1]):
        phase_factors = np.exp(-1j * phase_shift) ** i
        matrix[:, i, :] = matrix[:, i, :] * phase_factors

    fft_data = np.fft.ifftn(matrix)
    fft_shifted_data = np.fft.ifftshift(fft_data)
    amplitude = np.abs(fft_shifted_data)
    # print("TTTT")

    for slice in range(amplitude.shape[2]):
        plt.imshow(amplitude[:, :, slice], cmap='gray', vmin=amplitude.min(),
                   vmax=amplitude.max())
        plt.colorbar()
        plt.savefig(
            os.path.join('data', f'im_contrast_{nCoil}_{nContrast}_k_{slice}.png'))
        plt.close()

        # print("slice")

        log_matrix = np.log(np.abs(matrix))
        plt.imshow(log_matrix[:, :, slice], cmap='gray', vmin=np.log(np.abs(minimum_k_space)),
                   vmax=np.log(np.abs(maximum_k_space)))
        plt.savefig(
            os.path.join('data', f'k_contrast_{nCoil}_{nContrast}_k_{slice}.png'))
        plt.close()

    if return_image:
        return amplitude


def save_ffft(data, spoil=0):
    '''
        :param data: [x, y, z, contrast, coil]
        :param spoil: градус сдвига
    '''
    print(data.shape)

    k_space_max = np.max(data)
    k_space_min = np.min(data)

    num_slice = data.shape[2]  # Размер третьего измерения slice
    num_contrast = data.shape[3]  # Размер четвертого измерения contrast
    num_coils = data.shape[4]  # Размер по оси coil

    all_img = np.zeros((num_contrast, num_slice, data.shape[0], data.shape[1]), dtype=np.float32)

    for number_matrix_k_space in range(num_slice):
        for contrast_nb in range(num_contrast):
            img_sum = None

            for coil_nb in range(num_coils):
                coil_data = data[:, :, number_matrix_k_space, contrast_nb, coil_nb]
                img_coil = reconstruction(coil_data, k_space_min, k_space_max,
                                          coil_nb, contrast_nb, number_matrix_k_space,
                                          spoil,
                                          return_image=True )

                img_abs_two = img_coil ** 2

                if img_sum is None:
                    img_sum = img_abs_two
                else:
                    img_sum += img_abs_two

            all_img[contrast_nb, number_matrix_k_space] \
                = img_sum.copy()  # тут сохраняется ИЗОБРАЖЕНИЕ, квадрат модуля для ВСЕХ катушек данного констраста и слоя

    global_max = np.max(all_img)
    global_min = np.min(all_img)

    save_img(all_img, global_max, global_min)


def save_ffft_3d(data, phase=0):
    '''
        :param data:  [x, y, z, contrast, coil]
        :param phase: spoil
    '''
    num_contrast = data.shape[3]  # Размер четвертого измерения
    num_coils = data.shape[4]
    num_slice = data.shape[2]

    k_space_max = np.max(data)
    k_space_min = np.min(data)

    all_img = np.zeros((num_contrast, num_slice, data.shape[0], data.shape[1]), dtype=np.float32)

    # print("phase: ", phase)
    # for coil_nb in range(num_coils):
    for contrast_nb in range(num_contrast):
        img_sum = None

        for coil_nb in range(num_coils):
            coil_data = data[:, :, :, contrast_nb, coil_nb]
            img_coil = reco_3d(coil_data, k_space_max, k_space_min,
                               coil_nb, contrast_nb,
                               phase,
                               return_image=True)
            
            print(img_coil.shape)

            img_abs_two = img_coil ** 2

            if img_sum is None:
                img_sum = img_abs_two
            else:
                img_sum += img_abs_two


        img_sum = np.transpose(img_sum, (2, 0, 1))

        all_img[contrast_nb] \
            = img_sum.copy()  # тут сохраняется ИЗОБРАЖЕНИЕ, квадрат модуля для ВСЕХ катушек данного констраста
        
        print(contrast_nb)
        print(all_img.shape)

    # all_img = np.transpose(all_img, (1, 2, 3, 0))

    global_max = np.max(all_img)
    global_min = np.min(all_img)
    
    print(global_max, global_min)
    print(all_img)

    save_img(all_img, global_max, global_min)


def save_img(data, maximum, minimum):
    '''
        :param data: all_img[contrast_nb, number_matrix_k_space] = img_sum.[x,y] = [contrast, slice, x, y] 
        :return: 
    '''
    
    contrasts = data.shape[0]
    slices = data.shape[1]

    for contrast in range(contrasts):
        for slice in range(slices):
            plt.imshow(data[contrast, slice], cmap='gray', vmin=minimum, vmax=maximum)
            plt.colorbar()
            plt.savefig(
                os.path.join('data', f'im_{contrast}_k_{slice}.png'))
            plt.close()