from gtrain import Model
import numpy as np
import tensorflow as tf


class NetForHypinv(Model):
    """
    Implementaion of the crutial function for the HypINV algorithm.
    Warning: Do not use this class but implement its subclass, for example see FCNetForHypinv
    """
    def __init__(self, weights):
        self.eval_session = None
        self.grad_session = None
        self.initial_x = None
        self.center = None
        self.weights = weights
        self.out_for_eval = None #(going to be filled in build_for_eval method)
        self.boundary_out_for_eval = None
        self.trained_x = None
        self.training_class_index = None
        self.x = None # tf variable for inversion (going to be filled in build method)
        self.x_for_eval = None
        self.out = None
        self.boundary_out = None # list of tf tensorf for each class of softmax class vs others output
        self.loss = None
        self.boundary_loss = None
        self.t = None #target
        self.boundary_t = None
        self.x1 = None # this attribute is used of purposes of modified loss function


    def __del__(self):
        # close arr sessions
        if self.eval_session:
            self.eval_session.close()
        if self.grad_session:
            self.grad_session.close()

    def set_initial_x(self, initial_x):
        # sets starting point for the search of the closest point
        self.initial_x = initial_x

    def set_center(self, center):
        # sets center point
        self.center = center / np.linalg.norm(center)

    def set_x1(self, x1):
        # sets x1 to which we want to found the cosest point x0
        self.x1 = x1

    def has_modified_loss(self):
        pass # if uses modified loss then it returns true

    def set_initial_x_in_session(self, x, session=None):
        # sets initial x in certain session
        if session is None:
            self.set_initial_x(x)
        else:
            pass # overide this method

    def eval(self, x):
        if len(x.shape) == 1:
            x = x.reshape((1,len(x)))
        if not self.eval_session:
            self.eval_session = tf.Session()
            with self.eval_session.as_default():
                self.build_for_eval()
            self.eval_session.run(tf.global_variables_initializer())
        return self.eval_session.run(self.out_for_eval, {self.x_for_eval: x})

    def boundary_eval(self, x, class_index):
        # evaluates binary classificaitons class_index and other classes
        if not self.eval_session:
            self.eval_session = tf.Session()
            with self.eval_session.as_default():
                self.build_for_eval()
            self.eval_session.run(tf.global_variables_initializer())
        return self.eval_session.run(self.boundary_out_for_eval[class_index], {self.x_for_eval: x})

    def get_boundary_gradient(self, x, class_index):
        # computes gradient of the boundary for specified class_index
        if not self.grad_session:
            self.grad_session = tf.Session()
            with self.grad_session.as_default():
                self.build_for_eval()
                self.grad = list()
                for i in range(len(self.weights[0][-1][0])):
                    self.grad.append(tf.gradients(self.boundary_out_for_eval[i], [self.x_for_eval])[0])
                self.grad_x = self.x_for_eval
        return self.grad_session.run(self.grad[class_index], {self.grad_x: x})

    def build_for_eval(self):
        # build model for evaluation
        pass #override this method (fill self.out_for_eval)

    def train_ended(self, session):
        self.trained_x = session.run(self.x)

    def build(self):
        # build model for training
        pass #override this method (fill self.x, self.out)

    def set_train_class(self, class_index):
        # sets class of the x1
        self.training_class_index = class_index

    # overided methods from gtrain.Model
    def get_loss(self):
        if self.training_class_index is None:
            return self.loss
        else:
            return self.boundary_loss[self.training_class_index]

    def get_hits(self):
        return self.get_loss()

    def get_count(self):
        return self.get_loss()

    def get_train_summaries(self):
        return []

    def get_dev_summaries(self):
        return []

    def get_placeholders(self):
        if self.training_class_index is None:
            return [self.t]
        else:
            return [self.boundary_t]



#________________________________________EXAMPLES_OF_NetForHypinv_CLASS_____________________________________________

class FCNetForHypinv(NetForHypinv):
    """
    Implementation of multi layer perceptron to by used in HypINV rule extraction algorithm
    """
    def __init__(self, weights, function=tf.sigmoid, use_modified_loss=False, mu = 0.01):
        """
        :param weights: saved as [list of weights for layers][0 weight, 1 bias]
        :param function: tf function for propagation. For example tf.nn.sigmoid, tf.atan
        :param use_modified_loss: weather the modified loss should be used
        :param mu: factor of the penalty terms that specified the distance between x0 and x1 and
            the distance x1 from the boundary
        """
        super(FCNetForHypinv, self).__init__(weights)
        self.function = function
        self.layer_sizes = [len(self.weights[0][0])]
        for bias in weights[1]:
            self.layer_sizes.append(len(bias))
        self.num_classes = self.layer_sizes[-1]
        self.initial_x = 
np.zeros([1, self.layer_sizes[0]])