"""
Makes training and test dataset for nowcasting model using SEVIR
"""

# -*- coding: utf-8 -*-
import argparse
import logging

import os
import h5py
os.environ["HDF5_USE_FILE_LOCKING"]='FALSE'

import sys
import numpy as np
import tensorflow as tf
from nowcast_generator import get_nowcast_test_generator

# parser = argparse.ArgumentParser(description='Make nowcast training & test datasets using SEVIR')
# parser.add_argument('--sevir_data', type=str, help='location of SEVIR dataset',default='../../data/sevir')
# parser.add_argument('--sevir_catalog', type=str, help='location of SEVIR dataset',default='../../data/CATALOG.csv')
# parser.add_argument('--output_location', type=str, help='location of SEVIR dataset',default='../../data/interim')
# parser.add_argument('--n_chunks', type=int, help='Number of chucks to use (increase if memory limited)',default=10)


#args = parser.parse_args()

def generate_data(sevir_location,sevir_catalog,output_location,n_chunks=10):
    """ 
    Runs data processing scripts to extract training set from SEVIR
    """
    logger = logging.getLogger(__name__)
    logger.info('making final data set from raw data')
    #trn_generator = get_nowcast_train_generator(sevir_catalog=args.sevir_catalog,sevir_location=args.sevir_data)
    tst_generator = get_nowcast_test_generator(sevir_catalog,sevir_location)
    
    #ogger.info('Reading/writing training data to %s' % ('%s/nowcast_training.h5' % args.output_location))
    #read_write_chunks('%s/nowcast_training.h5' % args.output_location,trn_generator,args.n_chunks)
    logger.info('Reading/writing testing data to ' + output_location+'/nowcast_testing.h5')
    read_write_chunks(output_location+'/nowcast_testing.h5',tst_generator,n_chunks)


def read_write_chunks( filename, generator, n_chunks ):
    logger = logging.getLogger(__name__)
    chunksize = len(generator)//n_chunks
    # get first chunk
    logger.info('Gathering chunk 0/%s:' % n_chunks)
    X,Y=generator.load_batches(n_batches=chunksize,offset=0,progress_bar=True)
    # Create datasets
    with h5py.File(filename, 'w') as hf:
      hf.create_dataset('IN', data=X[0],  maxshape=(None,X[0].shape[1],X[0].shape[2],X[0].shape[3]))
      hf.create_dataset('OUT', data=Y[0], maxshape=(None,Y[0].shape[1],Y[0].shape[2],Y[0].shape[3]))
    # Gather other chunks
    for c in range(1,n_chunks+1):
      offset = c*chunksize
      n_batches = min(chunksize,len(generator)-offset)
      if n_batches<0: # all done
        break
      logger.info('Gathering chunk %d/%s:' % (c,n_chunks))
      X,Y=generator.load_batches(n_batches=n_batches,offset=offset,progress_bar=True)
      with h5py.File(filename, 'a') as hf:
            hf['IN'].resize((hf['IN'].shape[0] + X[0].shape[0]), axis = 0)
            hf['OUT'].resize((hf['OUT'].shape[0] + Y[0].shape[0]), axis = 0)
            hf['IN'][-X[0].shape[0]:]  = X[0]
            hf['OUT'][-Y[0].shape[0]:] = Y[0]




[(31, 'logging.getLogger', 'logging.getLogger', False, 'import logging\n'), (34, 'nowcast_generator.get_nowcast_test_generator', 'get_nowcast_test_generator', False, 'from nowcast_generator import get_nowcast_test_generator\n'), (43, 'logging.getLogger', 'logging.getLogger', False, 'import logging\n'), (49, 'h5py.File', 'h5py.File', False, 'import h5py\n'), (60, 'h5py.File', 'h5py.File', False, 'import h5py\n')]