# ============================================================================
#  Copyright 2021. 
#
#
#  Author: <NAME> 
#  Contact: <EMAIL>, <EMAIL> 
#
#
# All Rights Reserved.
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ============================================================================



"""Tensorflow implementation of NVAE"""


import tensorflow as tf
import tensorflow.keras as tfk

import numpy as np

__all__ = [
           'DatasetSequence',
           ]

class DatasetSequence(tfk.utils.Sequence):
    def __init__(self, x_set, y_set, shape, batch_size, binary=False, flip_vertically=False, flip_horizontally=False):
        
        self.original_x, self.y = x_set, y_set      
        self.x=self.original_x 
                
        self.batch_size = batch_size
        
    
        self.size=x_set[0].shape[0]
        self.binary=binary
        self.shape=shape
        
        self.flip_vertically=flip_vertically
        
        
        self.flip_horizontally=flip_horizontally
    
    

    def __len__(self):
        return int(np.ceil(self.size / float(self.batch_size)))

    def __getitem__(self, idx):
        
        batch_x =tf.convert_to_tensor(self.x[0][idx * self.batch_size:min((idx + 1) * self.batch_size,self.size)])
        
        
    

        if self.flip_horizontally:
            batch_x=tf.image.random_flip_left_right(batch_x)
            
        if self.flip_vertically:
        
            batch_x=tf.image.random_flip_up_down(batch_x)
        
   
        batch_y = [tf.convert_to_tensor(y_i[idx * self.batch_size:(idx + 1) * self.batch_size]) for y_i in self.y]
        
        return batch_x, batch_y





# ============================================================================
#  Copyright 2021. 
#
#
#  Author: <NAME> 
#  Contact: <EMAIL>, <EMAIL> 
#
#
# All Rights Reserved.
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ============================================================================



"""Tensorflow implementation of NVAE"""


import tensorflow as tf
import tensorflow.keras as tfk

import numpy as np

__all__ = [
           'DatasetSequence',
           ]

class DatasetSequence(tfk.utils.Sequence):
    def __init__(self, x_set, y_set, shape, batch_size, binary=False, flip_vertically=False, flip_horizontally=False):
        
        self.original_x, self.y = x_set, y_set      
        self.x=self.original_x 
                
        self.batch_size = batch_size
        
    
        self.size=x_set[0].shape[0]
        self.binary=binary
        self.shape=shape
        
        self.flip_vertically=flip_vertically
        
        
        self.flip_horizontally=flip_horizontally
    
    

    def __len__(self):
        return int(np.ceil(self.size / float(self.batch_size)))

    def __getitem__(self, idx):
        
        batch_x =tf.convert_to_tensor(self.x[0][idx * self.batch_size:min((idx + 1) * self.batch_size,self.size)])
        
        
    

        if self.flip_horizontally:
            batch_x=tf.image.random_flip_left_right(batch_x)
            
        if self.flip_vertically:
        
            batch_x=tf.image.random_flip_up_down(batch_x)
        
   
        batch_y = [tf.convert_to_tensor(y_i[idx * self.batch_size:(idx + 1) * self.batch_size]) for y_i in self.y]
        
        return batch_x, batch_y





# ============================================================================ #  Copyright 2021. # # #  Author: <NAME> #  Contact: <EMAIL>, <EMAIL> # # # All Rights Reserved. # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # #     http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # ============================================================================ Tensorflow implementation of NVAE