import theano as th
import theano.tensor as T
from theano.tensor.nnet import sigmoid
import numpy as np
from theano_toolkit import updates
import pickle

hidden_size = 100

def orthonormal_wts(n, m):
  nm = max(n, m)
  svd_u = np.linalg.svd(np.random.randn(nm, nm))[0]
  return svd_u.astype(th.config.floatX)[:n, :m]

def init_wts(*argv):
  return 0.4 * (np.random.rand(*argv) - 0.5)

def share(array, dtype=th.config.floatX, name=None):
  return th.shared(value=np.asarray(array, dtype=dtype), name=name)

class OutLayer:
  def __init__(self, input, in_size, n_classes):
    id = str(np.random.randint(0, 10000000))
    w = share(init_wts(in_size, n_classes), name=id+"w")
    b = share(np.zeros(n_classes), name=id+"b")
    eps = 0.0000001
    self.output = T.clip(T.nnet.softmax(T.dot(input, w) + b), eps, 1-eps)
    self.params = [w, b]

class SimpleLayer:
  def __init__(self, input, nin, nunits):
    id = str(np.random.randint(0, 10000000))
    wio = share(orthonormal_wts(nin, nunits), name="wio"+id)  # input to output
    wir = share(orthonormal_wts(nin, nunits), name="wir"+id)  # input to output
    wiu = share(orthonormal_wts(nin, nunits), name="wiu"+id)  # input to output
    woo = share(orthonormal_wts(nunits, nunits), name="woo"+id)  # output to output
    wou = share(orthonormal_wts(nunits, nunits), name="wou"+id)  # output to output
    wor = share(orthonormal_wts(nunits, nunits), name="wor"+id)  # output to output
    bo = share(np.zeros(nunits), name="bo"+id)
    bu = share(np.zeros(nunits), name="bu"+id)
    br = share(np.zeros(nunits), name="br"+id)
    h0 = share(np.zeros(nunits), name="h0"+id)

    def step(in_t, out_tm1):
      update_gate = sigmoid(T.dot(out_tm1, wou) + T.dot(in_t, wiu) + bu)
      reset_gate = sigmoid(T.dot(out_tm1, wor) + T.dot(in_t, wir) + br)
      new_val = T.tanh(T.dot(in_t, wio) + reset_gate * T.dot(out_tm1, woo) + bo)
      return update_gate * out_tm1 + (1 - update_gate) * new_val
    
    self.output, _ = th.scan(
      step, sequences=[input],
      outputs_info=[h0])

    self.params = [wio, woo, bo, wir, wiu, wor, wou, br, bu, h0]

class BiSimpleLayer():
  def __init__(self, input, nin, nunits):
    fwd = SimpleLayer(input, nin, nunits)
    bwd = SimpleLayer(input[::-1], nin, nunits)
    self.params = fwd.params + bwd.params
    self.output = T.concatenate([fwd.output, bwd.output[::-1]], axis=1)

class SimpleNet:
  def __init__(self, in_size, n_classes, only_test=False):
    np.random.seed(47)
    self.input = T.fmatrix()
    self.targets = T.ivector()
    self.targets2 = T.ivector()

    self.layer = BiSimpleLayer(self.input, in_size, hidden_size)
    self.layer2 = BiSimpleLayer(self.layer.output, 2*hidden_size, hidden_size)
    self.layer3 = BiSimpleLayer(self.layer2.output, 2*hidden_size, hidden_size)
    self.out_layer = OutLayer(self.layer3.output, 2*hidden_size, n_classes)
    self.output = self.out_layer.output
    self.out_layer2 = OutLayer(self.layer3.output, 2*hidden_size, n_classes)
    self.output2 = self.out_layer2.output
    self.cost = -T.mean(T.log(self.output)[T.arange(self.targets.shape[0]), self.targets]) 
    self.cost += -T.mean(T.log(self.output2)[T.arange(self.targets2.shape[0]), self.targets2]) 

#    self.coster = th.function(inputs=[self.input, self.targets, self.targets2], outputs=self.cost) 
#    self.predict = th.function(inputs=[self.input], outputs=[self.output]) 
    self.tester = th.function(inputs=[self.input], outputs=[self.output, self.output2])

    self.layers = [self.layer, self.layer2, self.layer3, self.out_layer, self.out_layer2]

    self.params = [p for l in self.layers for p in l.params]
    total_pars = 0
    for p in self.params:
      shape = p.get_value().shape
      if len(shape) == 1:
        total_pars += shape[0]
      elif len(shape) == 2:
        total_pars += shape[0] * shape[1]
      else:
        assert(False)
    self.total_pars = total_pars
    if not only_test:
      self.lr = T.fscalar()
      inputs = [self.input]
      inputs.append(self.targets)
      inputs.append(self.targets2)
      inputs.append(self.lr)
      self.trainer = th.function(
          inputs=inputs,
          outputs=[self.cost, self.output, self.output2],
          updates=updates.momentum(self.params, (T.grad(self.cost, self.params)),
                                   learning_rate=self.lr, mu=0.8))
      self.cost_and_grad = th.function(
          inputs=inputs[:3],
          outputs=[self.cost] + T.grad(self.cost, self.params))

  def save(self, fn):
    with open(fn, "wb") as f:
      for p in self.params:
        pickle.dump(p.get_value(), f)

  def load(self, fn):
    with open(fn, "rb") as f:
      for p in self.params:
        p.set_value(pickle.load(f))

  def flatten_grad(self, to_flatten):
    ret = np.zeros(self.total_pars, dtype=th.config.floatX)
    cur_pos = 0
    for p in to_flatten:
      par = p.flatten()
      l = len(par)
      ret[cur_pos:cur_pos + l] = par
      cur_pos += l
    return ret

  def get_params(self):
    ret = np.zeros(self.total_pars, dtype=th.config.floatX)
    cur_pos = 0
    for p in self.params:
      par = p.get_value().flatten()
      l = len(par)
      ret[cur_pos:cur_pos + l] = par
      cur_pos += l
    return ret

  def set_params(self, params):
    cur_pos = 0
    params = np.asarray(params, dtype=th.config.floatX)
    for p in self.params:
      l = len(p.get_value().flatten())
      p.set_value(params[cur_pos:cur_pos+l].reshape(p.get_value().shape))
      cur_pos += l

class BatchLayer:
  def __init__(self, input, nin, nunits):
    id = str(np.random.randint(0, 10000000))
    wio = share(orthonormal_wts(nin, nunits), name="wio"+id)  # input to output
    wir = share(orthonormal_wts(nin, nunits), name="wir"+id)  # input to output
    wiu = share(orthonormal_wts(nin, nunits), name="wiu"+id)  # input to output
    woo = share(orthonormal_wts(nunits, nunits), name="woo"+id)  # output to output
    wou = share(orthonormal_wts(nunits, nunits), name="wou"+id)  # output to output
    wor = share(orthonormal_wts(nunits, nunits), name="wor"+id)  # output to output
    bo = share(np.zeros(nunits), name="bo"+id)
    bu = share(np.zeros(nunits), name="bu"+id)
    br = share(np.zeros(nunits), name="br"+id)
    h0 = share(np.zeros(nunits), name="h0"+id)

    def step(in_t, out_tm1):
      update_gate = sigmoid(T.dot(out_tm1, wou) + T.dot(in_t, wiu) + bu)
      reset_gate = sigmoid(T.dot(out_tm1, wor) + T.dot(in_t, wir) + br)
      new_val = T.tanh(T.dot(in_t, wio) + reset_gate * T.dot(out_tm1, woo) + bo)
      return update_gate * out_tm1 + (1 - update_gate) * new_val

    self.step_res = step(input[0], T.alloc(h0, input.shape[1]))
    self.output, _ = th.scan(
       step, sequences=[input],
       outputs_info=[T.alloc(h0, input.shape[1], nunits)])

    self.params = [wio, woo, bo, wir, wiu, wor, wou, br, bu, h0]

class BiBatchLayer():
  def __init__(self, input, nin, nunits):
    fwd = BatchLayer(input, nin, nunits)
    bwd = BatchLayer(input[::-1], nin, nunits)
    self.params = fwd.params + bwd.params
    self.output = T.concatenate([fwd.output, bwd.output[::-1]], axis=2)

class OutBatchLayer:
  def __init__(self, input, in_size, n_classes):
    id = str(np.random.randint(0, 10000000))
    w = share(init_wts(in_size, n_classes), name=id+"w")
    b = share(np.zeros(n_classes), name=id+"b")
    otmp = T.dot(input, w) + b
    e_x = T.exp(otmp - otmp.max(axis=2, keepdims=True))
    o2 = e_x / e_x.sum(axis=2, keepdims=True)
    eps = 0.00000001
    self.output = T.clip(o2, eps, 1-eps)
    self.params = [w, b]

class BatchNet:
  def __init__(self, in_size, n_classes, only_test=False, with_jacob=False, use_first=True,
use_second=True):
    np.random.seed(47)
    self.input = T.ftensor3()
    self.targets = T.imatrix()
    self.targets2 = T.imatrix()

    self.layer = BiBatchLayer(self.input.dimshuffle(1, 0, 2), in_size, hidden_size)
    self.layer2 = BiBatchLayer(self.layer.output, 2*hidden_size, hidden_size)
    self.layer3 = BiBatchLayer(self.layer2.output, 2*hidden_size, hidden_size)
    self.out_layer = OutBatchLayer(self.layer3.output, 2*hidden_size, n_classes)
    self.output = self.out_layer.output.dimshuffle(1, 0, 2)
    self.output_flat = T.reshape(self.output, (self.output.shape[0] * self.output.shape[1], -1))
    self.targets_flat = self.targets.flatten(ndim=1)
    self.cost = 0
    if use_first:
      self.cost = -T.mean(T.log(self.output_flat)[T.arange(self.targets_flat.shape[0]), self.targets_flat]) 
    
    self.out_layer2 = OutBatchLayer(self.layer3.output, 2*hidden_size, n_classes)
    self.output2 = self.out_layer2.output.dimshuffle(1, 0, 2)
    self.output2_flat = T.reshape(self.output2, (self.output2.shape[0] * self.output2.shape[1], -1))
    self.targets2_flat = self.targets2.flatten(ndim=1)
    if use_second:
      self.cost += -T.mean(T.log(self.output2_flat)[T.arange(self.targets2_flat.shape[0]), self.targets2_flat]) 

#    self.coster = th.function(inputs=[self.input, self.targets, self.targets2], outputs=[self.cost]) 
    self.predict = th.function(inputs=[self.input], outputs=[self.output]) 
    self.tester = th.function(inputs=[self.input], outputs=[self.output, self.output2])
    self.debug = th.function(inputs=[self.input], outputs=[self.layer.output, self.layer2.output,
    self.layer3.output])

    self.layers = [self.layer, self.layer2, self.layer3, self.out_layer, self.out_layer2]

    self.params = [p for l in self.layers for p in l.params]

    self.grad_params = [p for l in self.layers[:-2] for p in l.params]
    if use_first:
      self.grad_params += self.out_layer.params
    if use_second:
      self.grad_params += self.out_layer2.params

#    self.top_prob_sum = T.mean(T.log(T.max(self.output[0], axis=1))) +\
#                        T.mean(T.log(T.max(self.output2[0], axis=1)))
            
#    self.grad_in = th.function(inputs=[self.input], outputs=T.grad(self.top_prob_sum, self.input))

    if with_jacob:
      self.index = T.iscalar()
      self.j1 = T.grad(T.log(T.max(self.output[0][self.index])), self.input)
      self.j2 = T.grad(T.log(T.max(self.output2[0][self.index])), self.input)
      self.jacob = th.function(
        inputs=[self.input, self.index], outputs=[self.j1[0], self.j2[0]])
    if not only_test:
      self.lr = T.fscalar()
      inputs = [self.input]
      if use_first:
        inputs.append(self.targets)
      if use_second:
        inputs.append(self.targets2)
      inputs.append(self.lr)
      self.trainer = th.function(
          inputs=inputs,
          outputs=[self.cost, self.output, self.output2],
          updates=updates.momentum(self.grad_params, (T.grad(self.cost, self.grad_params)),
                                   learning_rate=self.lr, mu=0.8))

  def save(self, fn):
    with open(fn, "wb") as f:
      for p in self.params:
        pickle.dump(p.get_value(), f)

  def load(self, fn):
    with open(fn, "rb") as f:
      for p in self.params:
        p.set_value(pickle.load(f))

