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

hidden_size = 128

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.1 * (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 ConvLayer:
  def __init__(self, input, nin, nunits, width=7):
    id = str(np.random.randint(0, 10000000))
    wc = share(init_wts(nunits, nin, width, 1), name="wc"+id)
    bc = share(np.zeros(nunits), name="bc"+id)
    input = input.dimshuffle(1, 0, 2)
    inpr = input.reshape((input.shape[0], input.shape[1], input.shape[2], 1)).dimshuffle(0, 2, 1, 3)

    c1 = conv2d(input=inpr, filters=wc, border_mode='half') + bc.reshape((1, bc.shape[0], 1, 1))
    self.output = T.tanh(c1.flatten(3).dimshuffle(0, 2, 1).dimshuffle(1, 0, 2))

    self.params = [wc, bc]
  

class BatchLayer:
  def __init__(self, input, nin, nunits, width=7):
    id = str(np.random.randint(0, 10000000))
    wc = share(init_wts(nunits, nin, width, 1), name="wc"+id)
    bc = share(np.zeros(nunits), name="bc"+id)
    wf = share(init_wts(nunits, nin, width, 1), name="wf"+id)
    bf = share(np.zeros(nunits), name="bf"+id)
    h0 = share(np.zeros(nunits), name="h0"+id)
    input = input.dimshuffle(1, 0, 2)
    inpr = input.reshape((input.shape[0], input.shape[1], input.shape[2], 1)).dimshuffle(0, 2, 1, 3)

    c1 = conv2d(input=inpr, filters=wc, border_mode='half') + bc.reshape((1, bc.shape[0], 1, 1))
    c1 = T.tanh(c1.flatten(3).dimshuffle(0, 2, 1).dimshuffle(1, 0, 2))

    f1 = conv2d(input=inpr, filters=wf, border_mode='half') + bf.reshape((1, bf.shape[0], 1, 1))
    f1 = T.nnet.sigmoid(f1.flatten(3).dimshuffle(0, 2, 1).dimshuffle(1, 0, 2))
    self.dbg = f1

    def step(in_t, in_g, out_tm1):
#      return in_t
      return in_g * out_tm1 + (1 - in_g) * in_t

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


    self.params = [wc, bc, wf, bf, 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)
    self.dbg = fwd.dbg

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(49)
    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.layer2 = ConvLayer(self.layer.output, 2*hidden_size, 2*hidden_size)
    self.layer3 = BiBatchLayer(self.layer2.output, 2*hidden_size, hidden_size)
    self.layer4 = BiBatchLayer(self.layer3.output, 2*hidden_size, hidden_size)
#    self.layer4 = ConvLayer(self.layer3.output, 2*hidden_size, 2*hidden_size)
    self.out_layer = OutBatchLayer(self.layer4.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.layer4.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.layer4, 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 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.adam(self.grad_params, (T.grad(self.cost, self.grad_params)),
                               learning_rate=0.0001, epsilon=1e-6))

  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))

