import pandas
import numpy as np
import os
import matplotlib.pyplot as plt
import math
import h5py

from ismrmrdtools import show, transform
import ismrmrd
import ismrmrd.xsd
# import ReadWrapper



def read_ocmr(filename):
# Input:  *.h5 file name
# Output: all_data    k-space data, orgnazide as {'kx'  'ky'  'kz'  'coil'  'phase'  'set'  'slice'  'rep'  'avg'}
#         param  some parameters of the scan
# 

# This is a function to read K-space from ISMRMD *.h5 data
# Modifid by Chong Chen (Chong.Chen@osumc.edu) based on the python script
# from https://github.com/ismrmrd/ismrmrd-python-tools/blob/master/recon_ismrmrd_dataset.py

    if not os.path.isfile(filename):
        print("%s is not a valid file" % filename)
        raise SystemExit
    dset = ismrmrd.Dataset(filename, 'dataset', create_if_needed=False)
    header = ismrmrd.xsd.CreateFromDocument(dset.read_xml_header())
    enc = header.encoding[0]

    # Matrix size
    eNx = enc.encodedSpace.matrixSize.x
    #eNy = enc.encodedSpace.matrixSize.y
    eNz = enc.encodedSpace.matrixSize.z
    eNy = (enc.encodingLimits.kspace_encoding_step_1.maximum + 1); #no zero padding along Ny direction

    # Field of View
    eFOVx = enc.encodedSpace.fieldOfView_mm.x
    eFOVy = enc.encodedSpace.fieldOfView_mm.y
    eFOVz = enc.encodedSpace.fieldOfView_mm.z
    
    # Save the parameters    
    param = dict();
    param['TRes'] =  str(header.sequenceParameters.TR)
    param['FOV'] = [eFOVx, eFOVy, eFOVz]
    param['TE'] = str(header.sequenceParameters.TE)
    param['TI'] = str(header.sequenceParameters.TI)
    param['echo_spacing'] = str(header.sequenceParameters.echo_spacing)
    param['flipAngle_deg'] = str(header.sequenceParameters.flipAngle_deg)
    param['sequence_type'] = header.sequenceParameters.sequence_type

    # Read number of Slices, Reps, Contrasts, etc.
    nCoils = header.acquisitionSystemInformation.receiverChannels
    try:
        nSlices = enc.encodingLimits.slice.maximum + 1
    except:
        nSlices = 1
        
    try:
        nReps = enc.encodingLimits.repetition.maximum + 1
    except:
        nReps = 1
               
    try:
        nPhases = enc.encodingLimits.phase.maximum + 1
    except:
        nPhases = 1;

    try:
        nSets = enc.encodingLimits.set.maximum + 1;
    except:
        nSets = 1;

    try:
        nAverage = enc.encodingLimits.average.maximum + 1;
    except:
        nAverage = 1;   
        
    # TODO loop through the acquisitions looking for noise scans
    firstacq=0
    for acqnum in range(dset.number_of_acquisitions()):
        acq = dset.read_acquisition(acqnum)

        # TODO: Currently ignoring noise scans
        if acq.isFlagSet(ismrmrd.ACQ_IS_NOISE_MEASUREMENT):
            #print("Found noise scan at acq ", acqnum)
            continue
        else:
            firstacq = acqnum
            print("Imaging acquisition starts acq ", acqnum)
            break

    # assymetry echo
    kx_prezp = 0;
    acq_first = dset.read_acquisition(firstacq)
    if  acq_first.center_sample*2 <  eNx:
        kx_prezp = eNx - acq_first.number_of_samples
         
    # Initialiaze a storage array
    param['kspace_dim'] = {'kx ky kz coil phase set slice rep avg'};
    all_data = np.zeros((eNx, eNy, eNz, nCoils, nPhases, nSets, nSlices, nReps, nAverage), dtype=np.complex64)

    # Loop through the rest of the acquisitions and stuff
    for acqnum in range(firstacq,dset.number_of_acquisitions()):
        acq = dset.read_acquisition(acqnum)

        # Stuff into the buffer
        y = acq.idx.kspace_encode_step_1
        z = acq.idx.kspace_encode_step_2
        phase =  acq.idx.phase;
        set =  acq.idx.set;
        slice =  acq.idx.slice;
        rep =  acq.idx.repetition;
        avg = acq.idx.average;        
        all_data[kx_prezp:, y, z, :,phase, set, slice, rep, avg ] = np.transpose(acq.data)
        
    return all_data,param


ocmr_data_attributes_location = '/home/tomerweiss/dor/OCMR/OCMR/ocmr_data_attributes.csv'
ocmr_data_location = '/home/tomerweiss/dor/OCMR/data/'
out_data_location = '/home/tomerweiss/dor/OCMR/data_processed/'
df = pandas.read_csv(ocmr_data_attributes_location)
df.dropna(how='all', axis=0, inplace=True)
df.dropna(how='all', axis=1, inplace=True)

try:
    os.mkdir(out_data)
except:
    pass

# get relevant files
files = df[df['smp'] == 'fs']['file name'].values

i = 0
for filename in files:
    i += 1
    print(i, '/', len(files))
    kData,param = read_ocmr(ocmr_data_location + filename) # [kx, ky, kz, coil, phase(time), set, slice, rep, avg],

    # Image reconstruction (SoS)
    '''
    im_sos and image are the reconstruced cardiac cine images, with and without readout oversampling, respectively.
    '''
    kData = np.mean(kData, axis = 8) # average the k-space if average > 1
   
    # reconstruction
    im_coil = transform.transform_kspace_to_image(kData, [0,1]) # IFFT (2D image)
    im_sos = np.sqrt(np.sum(np.abs(im_coil) ** 2, 3)) # Sum of Square
    RO = im_sos.shape[0]
    image = im_sos[math.floor(RO/4):math.floor(RO/4*3),:,:] # Remove RO oversampling


    # writing files
    kData = kData[:,:,0,:,:,0,:,0].transpose((4,3,2,0,1)) # [slice, frames, coils, h,w]
    target = im_sos[:,:,0,:,0,:,0].transpose((3,2,0,1)) # [slice, h,w]
    hf = h5py.File(out_data_location + filename, 'w')
    hf.create_dataset('kspace', data=kData)
    hf.create_dataset('reconstruction_rss', data=target)
    hf.close()

