#!/usr/bin/env python

# Copyright 2014-2016  Brno University of Technology (author: Karel Vesely)

# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#  http://www.apache.org/licenses/LICENSE-2.0
#
# THIS CODE IS PROVIDED *AS IS* BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, EITHER EXPRESS OR IMPLIED, INCLUDING WITHOUT LIMITATION ANY IMPLIED
# WARRANTIES OR CONDITIONS OF TITLE, FITNESS FOR A PARTICULAR PURPOSE,
# MERCHANTABLITY OR NON-INFRINGEMENT.
# See the Apache 2 License for the specific language governing permissions and
# limitations under the License.

# Generated Nnet prototype, to be initialized by 'nnet-initialize'.

from __future__ import division
from __future__ import print_function
import math, random, sys, re

###
### Parse options
###
from optparse import OptionParser
usage="%prog [options] <feat-dim> <num-leaves> <num-hid-layers> <num-hid-neurons> >nnet-proto-file"
parser = OptionParser(usage)

# Softmax related,
parser.add_option('--no-softmax', dest='with_softmax',
                   help='Do not put <SoftMax> in the prototype [default: %default]',
                   default=True, action='store_false');
parser.add_option('--block-softmax-dims', dest='block_softmax_dims',
                   help='Generate <BlockSoftmax> with dims D1:D2:D3 [default: %default]',
                   default="", type='string');
# Activation related,
parser.add_option('--activation-type', dest='activation_type',
                   help='Select type of activation function : (<Sigmoid>|<Tanh>|<ParametricRelu>) [default: %default]',
                   default='<Sigmoid>', type='string');
parser.add_option('--activation-opts', dest='activation_opts',
                   help='Additional options for protoype of activation function [default: %default]',
                   default='', type='string');
# Affine-transform related,
parser.add_option('--hid-bias-mean', dest='hid_bias_mean',
                   help='Set bias for hidden activations [default: %default]',
                   default=-2.0, type='float');
parser.add_option('--hid-bias-range', dest='hid_bias_range',
                   help='Set bias range for hidden activations (+/- 1/2 range around mean) [default: %default]',
                   default=4.0, type='float');
parser.add_option('--param-stddev-factor', dest='param_stddev_factor',
                   help='Factor to rescale Normal distriburtion for initalizing weight matrices [default: %default]',
                   default=0.1, type='float');
parser.add_option('--no-glorot-scaled-stddev', dest='with_glorot',
                   help='Generate normalized weights according to X.Glorot paper, but mapping U->N with same variance (factor sqrt(x/(dim_in+dim_out)))',
                   action='store_false', default=True);
parser.add_option('--no-smaller-input-weights', dest='smaller_input_weights',
                   help='Disable 1/12 reduction of stddef in input layer [default: %default]',
                   action='store_false', default=True);
parser.add_option('--no-bottleneck-trick', dest='bottleneck_trick',
                   help='Disable smaller initial weights and learning rate around bottleneck',
                   action='store_false', default=True);
parser.add_option('--max-norm', dest='max_norm',
                   help='Max radius of neuron-weights in L2 space (if longer weights get shrinked, not applied to last layer, 0.0 = disable) [default: %default]',
                   default=0.0, type='float');
parser.add_option('--affine-opts', dest='affine_opts',
                   help='Additional options for protoype of affine tranform [default: %default]',
                   default='', type='string');
# Topology related,
parser.add_option('--bottleneck-dim', dest='bottleneck_dim',
                   help='Make bottleneck network with desired bn-dim (0 = no bottleneck) [default: %default]',
                   default=0, type='int');
parser.add_option('--with-dropout', dest='with_dropout',
                   help='Add <Dropout> after the non-linearity of hidden layer.',
                   action='store_true', default=False);
parser.add_option('--dropout-opts', dest='dropout_opts',
                   help='Extra options for dropout [default: %default]',
                   default='', type='string');


(o,args) = parser.parse_args()
if len(args) != 4 :
  parser.print_help()
  sys.exit(1)

# A HACK TO PASS MULTI-WORD OPTIONS, WORDS ARE CONNECTED BY UNDERSCORES '_',
o.activation_opts = o.activation_opts.replace("_"," ")
o.affine_opts = o.affine_opts.replace("_"," ")
o.dropout_opts = o.dropout_opts.replace("_"," ")

(feat_dim, num_leaves, num_hid_layers, num_hid_neurons) = [int(i) for i in args];
### End parse options


# Check
assert(feat_dim > 0)
assert(num_leaves > 0)
assert(num_hid_layers >= 0)
assert(num_hid_neurons > 0)
if o.block_softmax_dims:
  assert(sum(map(int, re.split("[,:]", o.block_softmax_dims))) == num_leaves) # posible separators : ',' ':'

# Optionaly scale
def Glorot(dim1, dim2):
  if o.with_glorot:
    # 35.0 = magic number, gives ~1.0 in inner layers for hid-dim 1024dim,
    return 35.0 * math.sqrt(2.0/(dim1+dim2));
  else:
    return 1.0


###
### Print prototype of the network
###

# NO HIDDEN LAYER, ADDING BOTTLENECK!
# No hidden layer while adding bottleneck means:
# - add bottleneck layer + hidden layer + output layer
if num_hid_layers == 0 and o.bottleneck_dim != 0:
  assert(o.bottleneck_dim > 0)
  assert(num_hid_layers == 0)
  if o.bottleneck_trick:
    # 25% smaller stddev -> small bottleneck range, 10x smaller learning rate
    print("<LinearTransform> <InputDim> %d <OutputDim> %d <ParamStddev> %f <LearnRateCoef> %f" % \
     (feat_dim, o.bottleneck_dim, \
      (o.param_stddev_factor * Glorot(feat_dim, o.bottleneck_dim) * 0.75 ), 0.1))
    # 25% smaller stddev -> smaller gradient in prev. layer, 10x smaller learning rate for weigts & biases
    print("<AffineTransform> <InputDim> %d <OutputDim> %d <BiasMean> %f <BiasRange> %f <ParamStddev> %f <LearnRateCoef> %f <BiasLearnRateCoef> %f <MaxNorm> %f" % \
     (o.bottleneck_dim, num_hid_neurons, o.hid_bias_mean, o.hid_bias_range, \
      (o.param_stddev_factor * Glorot(o.bottleneck_dim, num_hid_neurons) * 0.75 ), 0.1, 0.1, o.max_norm))
  else:
    print("<LinearTransform> <InputDim> %d <OutputDim> %d <ParamStddev> %f" % \
     (feat_dim, o.bottleneck_dim, \
      (o.param_stddev_factor * Glorot(feat_dim, o.bottleneck_dim))))
    print("<AffineTransform> <InputDim> %d <OutputDim> %d <BiasMean> %f <BiasRange> %f <ParamStddev> %f <MaxNorm> %f" % \
     (o.bottleneck_dim, num_hid_neurons, o.hid_bias_mean, o.hid_bias_range, \
      (o.param_stddev_factor * Glorot(o.bottleneck_dim, num_hid_neurons)), o.max_norm))
  print("%s <InputDim> %d <OutputDim> %d %s" % (o.activation_type, num_hid_neurons, num_hid_neurons, o.activation_opts)) # Non-linearity
  # Last AffineTransform (10x smaller learning rate on bias)
  print("<AffineTransform> <InputDim> %d <OutputDim> %d <BiasMean> %f <BiasRange> %f <ParamStddev> %f <LearnRateCoef> %f <BiasLearnRateCoef> %f" % \
   (num_hid_neurons, num_leaves, 0.0, 0.0, \
    (o.param_stddev_factor * Glorot(num_hid_neurons, num_leaves)), 1.0, 0.1))
  # Optionaly append softmax
  if o.with_softmax:
    if o.block_softmax_dims == "":
      print("<Softmax> <InputDim> %d <OutputDim> %d" % (num_leaves, num_leaves))
    else:
      print("<BlockSoftmax> <InputDim> %d <OutputDim> %d <BlockDims> %s" % (num_leaves, num_leaves, o.block_softmax_dims))
  print("</NnetProto>")
  # We are done!
  sys.exit(0)

# NO HIDDEN LAYERS!
# Add only last layer (logistic regression)
if num_hid_layers == 0:
  print("<AffineTransform> <InputDim> %d <OutputDim> %d <BiasMean> %f <BiasRange> %f <ParamStddev> %f" % \
        (feat_dim, num_leaves, 0.0, 0.0, (o.param_stddev_factor * Glorot(feat_dim, num_leaves))))
  if o.with_softmax:
    if o.block_softmax_dims == "":
      print("<Softmax> <InputDim> %d <OutputDim> %d" % (num_leaves, num_leaves))
    else:
      print("<BlockSoftmax> <InputDim> %d <OutputDim> %d <BlockDims> %s" % (num_leaves, num_leaves, o.block_softmax_dims))
  print("</NnetProto>")
  # We are done!
  sys.exit(0)


# THE USUAL DNN PROTOTYPE STARTS HERE!
# Assuming we have >0 hidden layers,
assert(num_hid_layers > 0)

# Begin the prototype,
# First AffineTranform,
print("<AffineTransform> <InputDim> %d <OutputDim> %d <BiasMean> %f <BiasRange> %f <ParamStddev> %f <MaxNorm> %f %s" % \
      (feat_dim, num_hid_neurons, o.hid_bias_mean, o.hid_bias_range, \
       (o.param_stddev_factor * Glorot(feat_dim, num_hid_neurons) * \
        (math.sqrt(1.0/12.0) if o.smaller_input_weights else 1.0)), o.max_norm, o.affine_opts))
      # Note.: compensating dynamic range mismatch between input features and Sigmoid-hidden layers,
      # i.e. mapping the std-dev of N(0,1) (input features) to std-dev of U[0,1] (sigmoid-outputs).
      # This is done by multiplying with stddev(U[0,1]) = sqrt(1/12).
      # The stddev of weights is consequently reduced with scale 0.29,
print("%s <InputDim> %d <OutputDim> %d %s" % (o.activation_type, num_hid_neurons, num_hid_neurons, o.activation_opts))
if o.with_dropout:
  print("<Dropout> <InputDim> %d <OutputDim> %d %s" % (num_hid_neurons, num_hid_neurons, o.dropout_opts))


# Internal AffineTransforms,
for i in range(num_hid_layers-1):
  print("<AffineTransform> <InputDim> %d <OutputDim> %d <BiasMean> %f <BiasRange> %f <ParamStddev> %f <MaxNorm> %f %s" % \
        (num_hid_neurons, num_hid_neurons, o.hid_bias_mean, o.hid_bias_range, \
         (o.param_stddev_factor * Glorot(num_hid_neurons, num_hid_neurons)), o.max_norm, o.affine_opts))
  print("%s <InputDim> %d <OutputDim> %d %s" % (o.activation_type, num_hid_neurons, num_hid_neurons, o.activation_opts))
  if o.with_dropout:
    print("<Dropout> <InputDim> %d <OutputDim> %d %s" % (num_hid_neurons, num_hid_neurons, o.dropout_opts))

# Optionaly add bottleneck,
if o.bottleneck_dim != 0:
  assert(o.bottleneck_dim > 0)
  if o.bottleneck_trick:
    # 25% smaller stddev -> small bottleneck range, 10x smaller learning rate
    print("<LinearTransform> <InputDim> %d <OutputDim> %d <ParamStddev> %f <LearnRateCoef> %f" % \
     (num_hid_neurons, o.bottleneck_dim, \
      (o.param_stddev_factor * Glorot(num_hid_neurons, o.bottleneck_dim) * 0.75 ), 0.1))
    # 25% smaller stddev -> smaller gradient in prev. layer, 10x smaller learning rate for weigts & biases
    print("<AffineTransform> <InputDim> %d <OutputDim> %d <BiasMean> %f <BiasRange> %f <ParamStddev> %f <LearnRateCoef> %f <BiasLearnRateCoef> %f <MaxNorm> %f %s" % \
     (o.bottleneck_dim, num_hid_neurons, o.hid_bias_mean, o.hid_bias_range, \
      (o.param_stddev_factor * Glorot(o.bottleneck_dim, num_hid_neurons) * 0.75 ), 0.1, 0.1, o.max_norm, o.affine_opts))
  else:
    # Same learninig-rate and stddev-formula everywhere,
    print("<LinearTransform> <InputDim> %d <OutputDim> %d <ParamStddev> %f" % \
     (num_hid_neurons, o.bottleneck_dim, \
      (o.param_stddev_factor * Glorot(num_hid_neurons, o.bottleneck_dim))))
    print("<AffineTransform> <InputDim> %d <OutputDim> %d <BiasMean> %f <BiasRange> %f <ParamStddev> %f <MaxNorm> %f %s" % \
     (o.bottleneck_dim, num_hid_neurons, o.hid_bias_mean, o.hid_bias_range, \
      (o.param_stddev_factor * Glorot(o.bottleneck_dim, num_hid_neurons)), o.max_norm, o.affine_opts))
  print("%s <InputDim> %d <OutputDim> %d %s" % (o.activation_type, num_hid_neurons, num_hid_neurons, o.activation_opts))
  if o.with_dropout:
    print("<Dropout> <InputDim> %d <OutputDim> %d %s" % (num_hid_neurons, num_hid_neurons, o.dropout_opts))

# Last AffineTransform (10x smaller learning rate on bias)
print("<AffineTransform> <InputDim> %d <OutputDim> %d <BiasMean> %f <BiasRange> %f <ParamStddev> %f <LearnRateCoef> %f <BiasLearnRateCoef> %f" % \
      (num_hid_neurons, num_leaves, 0.0, 0.0, \
       (o.param_stddev_factor * Glorot(num_hid_neurons, num_leaves)), 1.0, 0.1))

# Optionaly append softmax
if o.with_softmax:
  if o.block_softmax_dims == "":
    print("<Softmax> <InputDim> %d <OutputDim> %d" % (num_leaves, num_leaves))
  else:
    print("<BlockSoftmax> <InputDim> %d <OutputDim> %d <BlockDims> %s" % (num_leaves, num_leaves, o.block_softmax_dims))

# We are done!
sys.exit(0)