import os
from pathlib import Path

import numpy as np

import tensorflow as tf
import tensorflow.keras as ks
from tensorflow.keras.datasets import mnist

modelCheckpoints = Path('modelCheckpointsMnist')
os.makedirs(modelCheckpoints, exist_ok=True)

class ClipContraint(ks.constraints.Constraint):
    def __init__(self, minVal=-1, maxVal=1):
        self.minVal = minVal
        self.maxVal = maxVal
    
    def __call__(self, x):
        return tf.clip_by_value(x, self.minVal, self.maxVal)
    
    def get_config(self):
        config = {
            'minVal': self.minVal,
            'maxVal': self.maxVal
        }
        return config

class BinaryFunctions:
    
    @staticmethod
    def _round_through(x):
        rounded = tf.round(x)
        return x + tf.stop_gradient(rounded - x)
    
    @staticmethod
    def _hard_sigmoid(x):
        # same as x = (x+1)/2
        x = (0.5 * x) + 0.5
        return tf.clip_by_value(x, 0, 1)
    
    @staticmethod
    def _binary_tanh(x):
        return 2 * BinaryFunctions._round_through(BinaryFunctions._hard_sigmoid(x)) - 1
    
    @staticmethod
    def _binarize(x, H = 1.0):
        xb = H * BinaryFunctions._binary_tanh(x / H)
        return xb