import argparse
import logging
import os

import h5py

logger = logging.getLogger(__name__)


# Parse command line arguments
parser = argparse.ArgumentParser(description='Writes a HDF5 file of AustinMan or AustinWoman head only.', usage='python head_only_hdf5 filename')
parser.add_argument('filename', help='name and path to (HDF5) file containing AustinMan or AustinWoman model')
args = parser.parse_args()

# Read full body HDF5 file
f = h5py.File(args.filename, 'r')
dx_dy_dz = f.attrs['dx_dy_dz']
data = f['/data'][:, :, :]

# Define head as last 1/8 of total body height
nzhead = 7 * int(data.shape[2] / 8)

logger.info(f'Dimensions of head model: {data.shape[0]:g} x {data.shape[1]:g} x {data.shape[2] - nzhead:g} cells')

# Write HDF5 file
headfile = os.path.splitext(args.filename)[0] + '_head.h5'
f = h5py.File(headfile, 'w')
f.attrs['dx_dy_dz'] = dx_dy_dz
f['/data'] = data[:, :, nzhead:data.shape[2]]