# 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

  def test_cca_loss(self):
    dim = 4

    # Test with two completely shared dimensions
    x, y = self.create_test_data(frac=0.0)
    loss = cca.cca_loss(x, y, dim, 1e-4, 1e-2)
    print('New test_cca_loss is:', loss)
    self.assertAlmostEqual(loss, 2.05, delta=0.1)

    # Test with 1.5 shared dimensions (one full, one half dimension)
    x, y = self.create_test_data(frac=0.5)
    loss = cca.cca_loss(x, y, dim, 1e-4, 1e-2)
    print('New test_cca_loss is:', loss)
    self.assertAlmostEqual(loss, 1.72, delta=0.1)

    # Test with just one shared dimension
    x, y = self.create_test_data(frac=1.0)
    loss = cca.cca_loss(x, y, dim, 1e-4, 1e-2)
    print('New test_cca_loss is:', loss)
    self.assertAlmostEqual(loss, 1.0, delta=0.1)

  def test_cca_brain_model(self):
    """Test simple (no temporal shift) CCA."""
    print('\n\n**********test_cca_brain_model starting... *******')
    x1, x2 = self.create_test_data()
    cca_dims = 3
    batch_size = 1024
    frame_rate = 100.0  # Arbitrary


# 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


  def test_cca_loss(self):
    dim = 4

    # Test with two completely shared dimensions
    x, y = self.create_test_data(frac=0.0)
    loss = cca.cca_loss(x, y, dim, 1e-4, 1e-2)
    print('New test_cca_loss is:', loss)
    self.assertAlmostEqual(loss, 2.05, delta=0.1)

    # Test with 1.5 shared dimensions (one full, one half dimension)
    x, y = self.create_test_data(frac=0.5)
    loss = cca.cca_loss(x, y, dim, 1e-4, 1e-2)
    print('New test_cca_loss is:', loss)
    self.assertAlmostEqual(loss, 1.72, delta=0.1)

    # Test with just one shared dimension
    x, y = self.create_test_data(frac=1.0)
    loss = cca.cca_loss(x, y, dim, 1e-4, 1e-2)
    print('New test_cca_loss is:', loss)
    self.assertAlmostEqual(loss, 1.0, delta=0.1)

  def test_cca_brain_model(self):
    """Test simple (no temporal shift) CCA."""
    print('\n\n**********test_cca_brain_model starting... *******')
    x1, x2 = self.create_test_data()
    cca_dims = 3
    batch_size = 1024
    frame_rate = 100.0  # Arbitrary


    brain_dataset = brain_data.TestBrainData('input_1', 'output_1', frame_rate,