import numpy as np
import data
import matplotlib.pyplot as plt
import tensorflow as tf


class TFDeep:
  def __init__(self, Dim, param_delta=0.1,param_lambda=0.0001):
    """Arguments:
       - D: dimensions of each datapoint
       - C: number of classes
       - param_delta: training step
    """

    # definicija podataka i parametara:
    D = Dim[0]
    n = len(Dim)
    C = Dim[n-1]
    S = n-2 #slojevi

    self.X = tf.placeholder(tf.float32, [None,D])
    self.Yoh_ = tf.placeholder(tf.float32, [None, C])
    self.W = []
    self.b = []
    self.h = []
    self.Dim = Dim

    #definicija W-ova i b-ova
    for i in range (1, n):
        self.W.append(tf.Variable(tf.random_normal([Dim[i-1], Dim[i]]), name = "W" + str(i), trainable=True))
        self.b.append(tf.Variable(tf.zeros([Dim[i]]), name = 'b' + str(i), trainable=True))

    #hidden layers
    for i in range(0,S):
        if i==0:
            self.h.append(tf.nn.relu(tf.matmul(self.X,self.W[i]) + self.b[i]))
        else:
            self.h.append(tf.nn.relu(tf.matmul(self.h[i-1],self.W[i]) + self.b[i]))

    # formulacija modela: izračunati self.probs
    #   koristiti: tf.matmul, tf.nn.softmax
    if (S==0):
        self.probs = tf.nn.softmax(tf.matmul(self.X,self.W[0]) + self.b[0])
    else:
        self.probs = tf.nn.softmax(tf.matmul(self.h[S-1],self.W[S]) + self.b[S])

    # formulacija gubitka: self.loss
    #   koristiti: tf.log, tf.reduce_sum, tf.reduce_mean
    for i in range(0,len(self.W)):
        w_squared = tf.reduce_sum(tf.nn.l2_loss(self.W[i]))
    self.loss = tf.reduce_mean(-tf.reduce_sum(self.Yoh_*tf.log(self.probs),reduction_indices=1)- param_lambda*w_squared)
    #self.loss = tf.reduce_mean(-tf.reduce_sum(self.Yoh_*tf.log(self.probs),reduction_indices=1))

    # formulacija operacije učenja: self.train_step
    #   koristiti: tf.train.GradientDescentOptimizer,
    #              tf.train.GradientDescentOptimizer.minimize
    self.trainer = tf.train.GradientDescentOptimizer(param_delta)
    self.train_step = self.trainer.minimize(self.loss)

    # instanciranje izvedbenog konteksta: self.session
    #   koristiti: tf.Session
    self.session = tf.Session()

  def train(self, X, Yoh_, param_niter, batch = False, nmb_batch = 0 ):
    """Arguments:
       - X: actual datapoints [NxD]
       - Yoh_: one-hot encoded labels [NxC]
       - param_niter: number of iterations
    """
    # incijalizacija parametara
    #   koristiti: tf.initializers.global_variables
    self.session.run(tf.global_variables_initializer())

    # optimizacijska petlja
    #   koristiti: tf.Session.run
    if (batch):
        size = int(np.shape(X)[0]/nmb_batch)
        for i in range(0,nmb_batch):
            start = i*size
            end = start+size
            X_b = X[start:end]
            Y_b = Yoh_[start:end]

            for j in range(0,param_niter):

                val_loss, _, val_W,val_B = self.session.run([self.loss, self.train_step, self.W, self.b], feed_dict={self.X: X_b, self.Yoh_: Y_b})
                if j==0 or j==450 or j==900:
                    print(i,j,val_loss)

    else:
        for i in range(0,param_niter):
            val_loss, _, val_W,val_B = self.session.run([self.loss, self.train_step, self.W, self.b], feed_dict={self.X: X, self.Yoh_: Yoh_})
            print(i,val_loss)#, val_W, val_B)
    return

  def eval(self, X):
    """Arguments:
       - X: actual datapoints [NxD]
       Returns: predicted class probabilites [NxC]
    """
    #   koristiti: tf.Session.run
    val_probs = self.session.run([self.probs], feed_dict={self.X: X})
    return val_probs

  def count_params(self):
      ##????LOOOOL
      var = tf.trainable_variables()
      nb_par = 0
      for i in range(1,len(self.Dim)):
          nb_par += self.Dim[i-1]*self.Dim[i]+self.Dim[i]
      return var, nb_par

if __name__ == "__main__":
  # inicijaliziraj generatore slučajnih brojeva
  np.random.seed(100)
  tf.set_random_seed(100)

  # instanciraj podatke X i labele Yoh_
  #X,Y_ = data.sample_gauss(2, 100)
  X,Y_ = data.sample_gmm_2d(6,2,10)
  Yoh_=data.class_to_onehot(Y_)
  N=np.shape(X)[0]
  # izgradi graf:
  tflr = TFDeep([2,2])

  # nauči parametre:
  tflr.train(X, Yoh_, 20000)

  # dohvati vjerojatnosti na skupu za učenje
  probs = np.transpose(tflr.eval(X))
  Y= np.argmax(probs,axis=0)
  Y = Y.reshape(N,)
  #print('Y-',Y_)

  # ispiši performansu (preciznost i odziv po razredima)
  ac,pr,M = data.eval_perf_multi(Y, Y_)
  print('Prec and Rec',pr)
  var, cnt = tflr.count_params()
  print('Trainable var:', var)
  print('Broj parametara:', cnt)

  # iscrtaj rezultate, decizijsku plohu
  decfun = lambda X: np.argmax(np.transpose(tflr.eval(X)),axis=0)
  bbox=(np.min(X, axis=0), np.max(X, axis=0))
  data.graph_surface(decfun, bbox, offset=0.5)

  # graph the data points
  data.graph_data(X, Y_, Y, special=[])
  # show the plot
  plt.show()