import numpy as np
import nibabel as nib
import os
import sys
import pdb
import glob

def load_data(path):
    data_dim, data, data_valid = load_ni(path)

    data = data.reshape(data_dim)
    data = np.moveaxis(data,1,2)
    data = np.moveaxis(data,2,0)
    data_dim = data.shape
    data = data.reshape((data_dim[0]*data_dim[1]*data_dim[2]))

    return data, data_dim

def load_ni(path):
    filename, file_extension = os.path.splitext(path)
    if file_extension=='.mnc':
        data_dim, data, data_valid = load_minc(path)
    else:
        data_dim, data, data_valid = load_nii(path)
    return data_dim, data, data_valid

def load_minc(path):
    ni_img = nib.load(path)
    ni_data = ni_img.get_data()
    ni_data_dim = ni_data.shape
    data_dim = ni_data_dim
    data = ni_data.reshape((ni_data_dim[0]*ni_data_dim[1]*ni_data_dim[2]))
    data_valid = True

    return data_dim, data, data_valid

def load_nii(path):
    ni_img = nib.load(path)
    ni_data = ni_img.get_data()
    ni_data_dim = ni_data.shape
    data_dim = ni_data_dim[0:3]
    if (len(ni_data_dim)==4):
       data = ni_data[:,:,:,0].reshape((ni_data_dim[0]*ni_data_dim[1]*ni_data_dim[2]))
    else:
       data = ni_data.reshape((ni_data_dim[0]*ni_data_dim[1]*ni_data_dim[2]))
    data_valid = True

    return data_dim, data, data_valid

files = glob.glob('/mnt/hippocampus/mri_data/brainweb/*.bc.nrm.nii')

for fid in files:
    in_path = fid
    in_filename, in_file_extension = os.path.splitext(in_path)
    out_path = in_filename + '.bin'

    data, data_dim = load_data(in_path)

    print data_dim
    with open(out_path,'wb') as f:
        f.write(np.asarray(data_dim,dtype=np.uint32).tobytes())
        f.write(np.asarray(data,dtype=np.float32).tobytes())

fid = '/mnt/hippocampus/mri_data/brainweb/phantom_1.0mm_normal_crisp.clean_ref.nii'
in_path = fid
in_filename, in_file_extension = os.path.splitext(in_path)
out_path = in_filename + '.bin'
data, data_dim = load_data(in_path)
print data_dim
with open(out_path,'wb') as f:
    f.write(np.asarray(data_dim,dtype=np.uint32).tobytes())
    f.write(np.asarray(data,dtype=np.float32).tobytes())
