# 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/cca.py
# def rmss(x):
# def cca_pearson_correlation(x, y):
# def cca_pearson_correlation_first(x, y):
# def cca_pearson_correlation_second(x, y):
#   def __init__(self, requested_cca_dims, **kwargs):
#   def set_initial_weights(self, mean1, mean2, rot1, rot2):
#   def build(self, input_shapes):
#   def call(self, inputs):
#   def get_config(self):
#   def __init__(self, input_dataset, cca_dims=5, regularization_lambda=0.0,
#                **kwargs):
#   def call(self, input_data):
#   def compile(self, optimizer=tf.keras.optimizers.RMSprop,
#               loss=cca_pearson_correlation_first,
#               metrics=cca_pearson_correlation_first,
#               learning_rate=1e-3, **kwargs):
#   def fit(self, dataset, epochs=1):
# def calculate_cca_parameters_from_dataset(dataset, dim, regularization=0.1,
#                                           mini_batch_count=1000, eps_eig=1e-12):
# def cca_loss(x, y, dim, rcov1, rcov2, eps_eig=1e-12):
# class BrainCcaLayer(tf.keras.layers.Layer):
# class BrainModelCCA(brain_model.BrainModel):


from absl import flags
from absl.testing import absltest
from absl.testing import parameterized
from telluride_decoding import brain_data
from telluride_decoding import cca
import numpy as np
import tensorflow.compat.v2 as tf
    bm_cca = cca.BrainModelCCA(test_dataset, cca_dims=cca_dims)
    bm_cca.compile(optimizer=tf.keras.optimizers.RMSprop(learning_rate=1e-3),
                   metrics=[cca.cca_pearson_correlation_first,],
                   loss=['mse'])

    for x, _ in test_dataset.take(1):
      cca_input = x

    result = bm_cca(cca_input)   # Be sure call it once so network gets built
    self.assertEqual(result.shape[0], batch_size)
    self.assertEqual(result.shape[1], 2*cca_dims)

    print('*** Now train the model ***')
    bm_cca.fit(test_dataset)
    bm_cca.summary()

    print('*** Save some predictions for later ***')
    for x, _ in test_dataset.take(1):
      true_x = x
      predict_y = bm_cca.predict(true_x)

    print('*** Now save the model ***')
    saved_model_file = '/tmp/tf_cca_saved_model.tf'
    bm_cca.save(saved_model_file, save_format='tf')
    self.assertTrue(tf.io.gfile.exists(saved_model_file))

    print('*** Now load and test the model ***')
    # Custom_objects parameter in the load_model call doesn't seem to work, but
    # with ....CustomObjectScope does.
    with tf.keras.utils.CustomObjectScope(

# 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/cca.py
# def rmss(x):
# def cca_pearson_correlation(x, y):
# def cca_pearson_correlation_first(x, y):
# def cca_pearson_correlation_second(x, y):
#   def __init__(self, requested_cca_dims, **kwargs):
#   def set_initial_weights(self, mean1, mean2, rot1, rot2):
#   def build(self, input_shapes):
#   def call(self, inputs):
#   def get_config(self):
#   def __init__(self, input_dataset, cca_dims=5, regularization_lambda=0.0,
#                **kwargs):
#   def call(self, input_data):
#   def compile(self, optimizer=tf.keras.optimizers.RMSprop,
#               loss=cca_pearson_correlation_first,
#               metrics=cca_pearson_correlation_first,
#               learning_rate=1e-3, **kwargs):
#   def fit(self, dataset, epochs=1):
# def calculate_cca_parameters_from_dataset(dataset, dim, regularization=0.1,
#                                           mini_batch_count=1000, eps_eig=1e-12):
# def cca_loss(x, y, dim, rcov1, rcov2, eps_eig=1e-12):
# class BrainCcaLayer(tf.keras.layers.Layer):
# class BrainModelCCA(brain_model.BrainModel):


# Path: test/cca_test.py
from absl import flags
from absl.testing import absltest
from absl.testing import parameterized
from telluride_decoding import brain_data
from telluride_decoding import cca
import numpy as np
import tensorflow.compat.v2 as tf

    bm_cca = cca.BrainModelCCA(test_dataset, cca_dims=cca_dims)
    bm_cca.compile(optimizer=tf.keras.optimizers.RMSprop(learning_rate=1e-3),
                   metrics=[cca.cca_pearson_correlation_first,],
                   loss=['mse'])

    for x, _ in test_dataset.take(1):
      cca_input = x

    result = bm_cca(cca_input)   # Be sure call it once so network gets built
    self.assertEqual(result.shape[0], batch_size)
    self.assertEqual(result.shape[1], 2*cca_dims)

    print('*** Now train the model ***')
    bm_cca.fit(test_dataset)
    bm_cca.summary()

    print('*** Save some predictions for later ***')
    for x, _ in test_dataset.take(1):
      true_x = x
      predict_y = bm_cca.predict(true_x)

    print('*** Now save the model ***')
    saved_model_file = '/tmp/tf_cca_saved_model.tf'
    bm_cca.save(saved_model_file, save_format='tf')
    self.assertTrue(tf.io.gfile.exists(saved_model_file))

    print('*** Now load and test the model ***')
    # Custom_objects parameter in the load_model call doesn't seem to work, but
    # with ....CustomObjectScope does.
    with tf.keras.utils.CustomObjectScope(

        {'BrainCcaLayer': cca.BrainCcaLayer,