import os

import glob

import pydicom

import numpy as np

import nibabel as nib

from dipy.align.reslice import reslice

from dipy.align.imaffine import AffineMap

def extract_cosines(ImageOrientationPatient):

    row_cosine    = np.array(ImageOrientationPatient[:3])

    column_cosine = np.array(ImageOrientationPatient[3:])

    slice_cosine  = np.cross(row_cosine, column_cosine)

    return row_cosine, column_cosine, slice_cosine

def read_RefDs(dicom_folder_path):

    cine_dicom_files = glob.glob(dicom_folder_path)

    SliceLocations   = [pydicom.read_file(cine_dicom_file).SliceLocation 

                        for cine_dicom_file in cine_dicom_files]

    SliceOriginID    = np.array(SliceLocations, dtype=float).argmin()

    RefDs = pydicom.read_file(cine_dicom_files[SliceOriginID])

    ImageOrientationPatient = RefDs.ImageOrientationPatient

    ImagePositionPatient    = np.array(list(RefDs.ImagePositionPatient), dtype=float)

    Zooms                   = np.array(list(RefDs.PixelSpacing) + [RefDs.SliceThickness], dtype=float)

    return ImageOrientationPatient, ImagePositionPatient, Zooms

def read_affine_info(dicom_folder_path, viewer='slicer'):

    ImageOrientationPatient, ImagePositionPatient, Zooms = read_RefDs(dicom_folder_path)

    affine_axial        = np.diag(list(Zooms)+[1])

    affine_axial[:3,3] += ImagePositionPatient

    row_cos, column_cos, slice_cos = extract_cosines(ImageOrientationPatient)

    ijk2ras = np.stack((row_cos, column_cos, slice_cos))

    if viewer == "slicer":

        ijk2ras = (ijk2ras*np.array([-1,-1,1])).T

        ImagePositionPatient = ImagePositionPatient*np.array([-1,-1,1])

    affine  = np.stack((ijk2ras[:,0]*Zooms[0],

                        ijk2ras[:,1]*Zooms[1],

                        ijk2ras[:,2]*Zooms[2],

                        ImagePositionPatient), axis=1)

    return ijk2ras, np.vstack((affine,[[0,0,0,1]])), affine_axial   

def reslice_3d(domain_nifti, codomain_nifti, mode='linear'):

    affine_map = AffineMap(affine=np.eye(4),

                  domain_grid_shape=domain_nifti.shape, domain_grid2world=domain_nifti.affine, 

                  codomain_grid_shape=codomain_nifti.shape, codomain_grid2world=codomain_nifti.affine)

    return nib.Nifti1Image(affine_map.transform(codomain_nifti.get_fdata(),mode),domain_nifti.affine)

def reslice_4d(domain_nifti, codomain_nifti, mode='linear'):

    reslice_arr = np.zeros(domain_nifti.shape[:3]+(codomain_nifti.shape[-1],))

    affine_map  = AffineMap(affine=np.eye(4),

                  domain_grid_shape=domain_nifti.shape[:3],     domain_grid2world=domain_nifti.affine, 

                  codomain_grid_shape=codomain_nifti.shape[:3], codomain_grid2world=codomain_nifti.affine)

    for frame_id in range(reslice_arr.shape[-1]):

        reslice_arr[:,:,:,frame_id] += affine_map.transform(codomain_nifti.get_fdata()[:,:,:,frame_id],mode)

    return nib.Nifti1Image(reslice_arr, domain_nifti.affine)

def resample_nifti(subject_nifti, resolution=(1.5, 1.5), Nz=16, order=1, mode='nearest'):

    """Resample a 3D or 4D (3D+time) cine-MRI nifti to a new in-plane `resolution` with `Nz` slices."""

    data   = subject_nifti.get_fdata()

    affine = subject_nifti.affine

    zooms  = subject_nifti.header.get_zooms()[:3]

    new_zooms = (resolution[0], resolution[1], (zooms[2] * data.shape[2]) / Nz)

    data_resampled, affine_resampled = reslice(data, affine, zooms, new_zooms, order=order, mode=mode)

    subject_nifti_resampled = nib.Nifti1Image(data_resampled, affine_resampled)

    x=subject_nifti_resampled.header.get_zooms()[:3]

    y=new_zooms

    if not np.allclose(x,y, rtol=1e-02):

        print(subject_nifti.affine,affine)

        print(zooms)

        print(x,y)

        warnings.warn('Output resolutions are different than expected!')

    return subject_nifti_resampled