# Path: telluride_decoding/brain_data.py
# class BrainData(object):
# class TestBrainData(BrainData):
# class TFExampleData(BrainData):
#   def __init__(self, in_fields, out_field,
#                frame_rate,
#                pre_context=0,
#                post_context=0,
#                in2_fields=None,
#                in2_pre_context=0,
#                in2_post_context=0,
#                attended_field=None,
#                initial_batch_size=1000,
#                final_batch_size=1000,
#                repeat_count=1,
#                shuffle_buffer_size=1000,
#                data_dir=None,
#                data_pattern='',
#                train_file_pattern='',
#                validate_file_pattern='',
#                test_file_pattern=''):
#   def all_files(self, max_count=0):
#   def set_file_patterns(self, train, validate, test):
#   def create_dataset(self, mode='train', temporal_context=True):
#   def _get_data_file_names(self):
#   def filter_file_names(self, mode):
#   def final_shuffle_and_batch(self, mode, input_dataset, mixup_batch=False):
#       def mixup_batch_function(x, x2, y, a):
#   def add_temporal_context(self, dataset_without_context):
#     def window_one_stream_new(x, pre_context, post_context):
#     def window_data(x, x2, y, a, pre_context=0, post_context=0,
#                     in2_pre_context=0, in2_post_context=0):
#   def input_fields_width(self, input_number=1):
#   def output_field_width(self):
#   def create_dataset(self, mode='train', temporal_context=True,
#                      mixup_batch=False):
#   def preserve_test_data(self, input_data, output_data,
#                          input2_data=None, attention_data=None):
#   def _get_data_file_names(self):
#   def create_dataset(self, mode='train', temporal_context=True,
#                      mixup_batch=False):
#   def read_data_into_dataset(self, filenames, temporal_context=True):
#   def preprocess_list(self, name_params_list, frame_rate):
#   def parse_and_select_from_tfrecord(self, raw_proto):
#   def parse_and_select_from_tfrecord2(self, raw_proto):
# def discover_feature_shapes(tfrecord_file_name):
# def count_tfrecords(tfrecord_file_name):
# def create_brain_dataset(data_type, in_fields, out_field, frame_rate,
#                          pre_context=0,
#                          post_context=0,
#                          in2_fields=None,
#                          in2_pre_context=0,
#                          in2_post_context=0,
#                          attended_field=None,
#                          initial_batch_size=1000,
#                          final_batch_size=1000,
#                          repeat_count=1,
#                          shuffle_buffer_size=1000,
#                          data_dir=None,
#                          data_pattern='',
#                          train_file_pattern=None,
#                          validate_file_pattern=None,
#                          test_file_pattern=None):
# 
# Path: telluride_decoding/regression_data.py
# class DataLocation(object):
# class RegressionData(object):
# class RegressionDataTelluride4(RegressionData):
# class RegressionDataJensMemory(RegressionData):
# class RegressionDataJensImpaired(RegressionData):
# FLAGS = flags.FLAGS
# def loadmat(filename):
#   def _check_keys(key_dict):
#   def _todict(matobj):
# def download_from_gdrive(url, output, debug=False):
#   def parse_url(url):
#   def get_url_from_gdrive_confirmation(contents):
# def make_if_not_exists(directory):
#   def download_data(self, url, cache_dir, debug=False):
#   def is_data_local(self, cache_dir):
#   def download_data(self, url, cache_dir, debug=False):
#   def is_data_ingested(self, tf_dir, num_files=32):
#   def ingest_data(self, cache_dir, tf_dir, desired_frame_rate):
#   def name(self):
#   def is_data_local(self, cache_dir, num_subjects=22):
#   def download_data(self, url, cache_dir):
#   def is_data_ingested(self, tf_dir, num_subjects=22, num_trials=40):
#   def ingest_data(self, cache_dir, tf_dir, desired_frame_rate):
#   def name(self):
#   def is_data_local(self, cache_dir, num_subjects=44):
#   def download_data(self, url, cache_dir):
#   def is_data_ingested(self, tf_dir, num_subjects=44, num_trials=48):
#   def ingest_data(self, cache_dir, tf_dir, desired_frame_rate):
# def write_summary(cache_dir, tf_dir, frame_rate, all_ingested_files=None):
# def main(argv):


import os
import tensorflow.compat.v2 as tf
from absl import flags
from absl.testing import absltest
from telluride_decoding import brain_data
from telluride_decoding import regression_data

    self.assertEqual(brain_data.count_tfrecords(test_file), (8297, False))


class JensMemoryDataTest(absltest.TestCase):

  def setUp(self):
    super(JensMemoryDataTest, self).setUp()
    self._test_data_dir = os.path.join(
        flags.FLAGS.test_srcdir, '__main__',
        'test_data/')

  def test_data_ingestion(self):
    cache_dir = os.path.join(self._test_data_dir, 'jens_memory')
    tmp_dir = self.create_tempdir().full_path
    tf_dir = os.path.join(tmp_dir, 'jens_memory')
    num_subjects = 1   # Only 1 of 22 subjects loaded for test.
    num_trials = 5   # That one subject has been shortened to 5/40 trials.

    # Create the data object and make sure we have the downloaded archive file.
    rd = regression_data.RegressionDataJensMemory()
    self.assertTrue(rd.is_data_local(cache_dir, num_subjects))

    # Now ingest the data, making sure it's not present at start, then present.
    self.assertFalse(rd.is_data_ingested(tmp_dir, num_subjects))
    rd.ingest_data(cache_dir, tf_dir, 128)
    self.assertTrue(rd.is_data_ingested(tf_dir, num_subjects, num_trials))

    # Check the data files.
    test_file = os.path.join(tf_dir, 'subject_01', 'trial_01.tfrecords')

# Path: telluride_decoding/brain_data.py
# class BrainData(object):
# class TestBrainData(BrainData):
# class TFExampleData(BrainData):
#   def __init__(self, in_fields, out_field,
#                frame_rate,
#                pre_context=0,
#                post_context=0,
#                in2_fields=None,
#                in2_pre_context=0,
#                in2_post_context=0,
#                attended_field=None,
#                initial_batch_size=1000,
#                final_batch_size=1000,
#                repeat_count=1,
#                shuffle_buffer_size=1000,
#                data_dir=None,
#                data_pattern='',
#                train_file_pattern='',
#                validate_file_pattern='',
#                test_file_pattern=''):
#   def all_files(self, max_count=0):
#   def set_file_patterns(self, train, validate, test):
#   def create_dataset(self, mode='train', temporal_context=True):
#   def _get_data_file_names(self):
#   def filter_file_names(self, mode):
#   def final_shuffle_and_batch(self, mode, input_dataset, mixup_batch=False):
#       def mixup_batch_function(x, x2, y, a):
#   def add_temporal_context(self, dataset_without_context):
#     def window_one_stream_new(x, pre_context, post_context):
#     def window_data(x, x2, y, a, pre_context=0, post_context=0,
#                     in2_pre_context=0, in2_post_context=0):
#   def input_fields_width(self, input_number=1):
#   def output_field_width(self):
#   def create_dataset(self, mode='train', temporal_context=True,
#                      mixup_batch=False):
#   def preserve_test_data(self, input_data, output_data,
#                          input2_data=None, attention_data=None):
#   def _get_data_file_names(self):
#   def create_dataset(self, mode='train', temporal_context=True,
#                      mixup_batch=False):
#   def read_data_into_dataset(self, filenames, temporal_context=True):
#   def preprocess_list(self, name_params_list, frame_rate):
#   def parse_and_select_from_tfrecord(self, raw_proto):
#   def parse_and_select_from_tfrecord2(self, raw_proto):
# def discover_feature_shapes(tfrecord_file_name):
# def count_tfrecords(tfrecord_file_name):
# def create_brain_dataset(data_type, in_fields, out_field, frame_rate,
#                          pre_context=0,
#                          post_context=0,
#                          in2_fields=None,
#                          in2_pre_context=0,
#                          in2_post_context=0,
#                          attended_field=None,
#                          initial_batch_size=1000,
#                          final_batch_size=1000,
#                          repeat_count=1,
#                          shuffle_buffer_size=1000,
#                          data_dir=None,
#                          data_pattern='',
#                          train_file_pattern=None,
#                          validate_file_pattern=None,
#                          test_file_pattern=None):
# 
# Path: telluride_decoding/regression_data.py
# class DataLocation(object):
# class RegressionData(object):
# class RegressionDataTelluride4(RegressionData):
# class RegressionDataJensMemory(RegressionData):
# class RegressionDataJensImpaired(RegressionData):
# FLAGS = flags.FLAGS
# def loadmat(filename):
#   def _check_keys(key_dict):
#   def _todict(matobj):
# def download_from_gdrive(url, output, debug=False):
#   def parse_url(url):
#   def get_url_from_gdrive_confirmation(contents):
# def make_if_not_exists(directory):
#   def download_data(self, url, cache_dir, debug=False):
#   def is_data_local(self, cache_dir):
#   def download_data(self, url, cache_dir, debug=False):
#   def is_data_ingested(self, tf_dir, num_files=32):
#   def ingest_data(self, cache_dir, tf_dir, desired_frame_rate):
#   def name(self):
#   def is_data_local(self, cache_dir, num_subjects=22):
#   def download_data(self, url, cache_dir):
#   def is_data_ingested(self, tf_dir, num_subjects=22, num_trials=40):
#   def ingest_data(self, cache_dir, tf_dir, desired_frame_rate):
#   def name(self):
#   def is_data_local(self, cache_dir, num_subjects=44):
#   def download_data(self, url, cache_dir):
#   def is_data_ingested(self, tf_dir, num_subjects=44, num_trials=48):
#   def ingest_data(self, cache_dir, tf_dir, desired_frame_rate):
# def write_summary(cache_dir, tf_dir, frame_rate, all_ingested_files=None):
# def main(argv):


# Path: test/regression_data_test.py
import os
import tensorflow.compat.v2 as tf
from absl import flags
from absl.testing import absltest
from telluride_decoding import brain_data
from telluride_decoding import regression_data


    self.assertEqual(brain_data.count_tfrecords(test_file), (8297, False))


class JensMemoryDataTest(absltest.TestCase):

  def setUp(self):
    super(JensMemoryDataTest, self).setUp()
    self._test_data_dir = os.path.join(
        flags.FLAGS.test_srcdir, '__main__',
        'test_data/')

  def test_data_ingestion(self):
    cache_dir = os.path.join(self._test_data_dir, 'jens_memory')
    tmp_dir = self.create_tempdir().full_path
    tf_dir = os.path.join(tmp_dir, 'jens_memory')
    num_subjects = 1   # Only 1 of 22 subjects loaded for test.
    num_trials = 5   # That one subject has been shortened to 5/40 trials.

    # Create the data object and make sure we have the downloaded archive file.
    rd = regression_data.RegressionDataJensMemory()
    self.assertTrue(rd.is_data_local(cache_dir, num_subjects))

    # Now ingest the data, making sure it's not present at start, then present.
    self.assertFalse(rd.is_data_ingested(tmp_dir, num_subjects))
    rd.ingest_data(cache_dir, tf_dir, 128)
    self.assertTrue(rd.is_data_ingested(tf_dir, num_subjects, num_trials))

    # Check the data files.
    test_file = os.path.join(tf_dir, 'subject_01', 'trial_01.tfrecords')

    features = brain_data.discover_feature_shapes(test_file)