#!/usr/bin/env python

from __future__ import print_function

import sys

import h5py
import numpy as np
from astropy.io import fits

arguments = sys.argv[1:]

if '--physics' in arguments:
    physics = True
    arguments.remove('--physics')
else:
    physics = False

if '--images' in arguments:
    images = True
    arguments.remove('--images')
else:
    images = False

if not physics and not images:
    print("Need to specify at least one of --images or --physics")
    sys.exit(0)

for filename in arguments:

    try:
        f = h5py.File(filename, 'r')
    except IOError:
        print("Reading %s failed" % filename)
        continue

    if images:
        if 'Peeled' in f or 'Binned' in f:
            if 'Peeled' in f:
                n_groups = len(f['Peeled'])
                for ig in range(n_groups):
                    group = f['Peeled/group_%05i' % (ig + 1)]
                    if 'images' in group:
                        image = np.array(group['images'])
                        image = image.swapaxes(-1, -2)
                        image = image.swapaxes(-2, -3)
                        fits.writeto(filename.replace('.rtout', '_%05i_images.fits' % (ig + 1)),
                                     image, clobber=True)
            if 'Binned' in f:
                group = f['Binned']
                if 'images' in group:
                    image = np.array(group['images'])
                    image = image.swapaxes(-1, -2)
                    image = image.swapaxes(-2, -3)
                    fits.writeto(filename.replace('.rtout', '_binned_images.fits'),
                                 image, clobber=True)
        else:
            print("No images in %s" % filename)

    if physics:
        for group_name in f:
            if 'iteration' in group_name:
                iteration = int(group_name.split('_')[1])
                for array_name in f[group_name]:
                    fits.writeto(filename.replace('.rtout', '_%s_%05i.fits' % (array_name, iteration)),
                                 np.array(f[group_name][array_name]), clobber=True)

    try:
        f.close()
    except:
        print("Error closing %s" % filename)
