# 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
#
# https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

!pip install -q lucid>=0.2.3

from __future__ import print_function
import io
import string
import numpy as np
import PIL
import base64
from glob import glob

import matplotlib.pylab as pl

import tensorflow as tf
from tensorflow.contrib import slim

from IPython.display import clear_output, Image, display, HTML

import moviepy.editor as mpy
from moviepy.video.io.ffmpeg_writer import FFMPEG_VideoWriter


from google.colab import files

from lucid.modelzoo import vision_models
from lucid.misc.io import show, save, load
from lucid.optvis import objectives
from lucid.optvis import render
from lucid.misc.tfutil import create_session

model = vision_models.InceptionV1()
model.load_graphdef()

def composite_activation(x):
  x = tf.atan(x)
  # Coefficients computed by:
  #   def rms(x):
  #     return np.sqrt((x*x).mean())
  #   a = np.arctan(np.random.normal(0.0, 1.0, 10**6))
  #   print(rms(a), rms(a*a))
  return tf.concat([x/0.67, (x*x)/0.6], -1)


def composite_activation_unbiased(x):
  x = tf.atan(x)
  # Coefficients computed by:
  #   a = np.arctan(np.random.normal(0.0, 1.0, 10**6))
  #   aa = a*a
  #   print(a.std(), aa.mean(), aa.std())
  return tf.concat([x/0.67, (x*x-0.45)/0.396], -1)


def relu_normalized(x):
  x = tf.nn.relu(x)
  # Coefficients computed by:
  #   a = np.random.normal(0.0, 1.0, 10**6)
  #   a = np.maximum(a, 0.0)
  #   print(a.mean(), a.std())
  return (x-0.40)/0.58


def image_cppn(
    size,
    num_output_channels=3,
    num_hidden_channels=24,
    num_layers=8,
    activation_fn=composite_activation,
    normalize=False):
  r = 3.0**0.5  # std(coord_range) == 1.0
  coord_range = tf.linspace(-r, r, size)
  y, x = tf.meshgrid(coord_range, coord_range, indexing='ij')
  net = tf.expand_dims(tf.stack([x, y], -1), 0)  # add batch dimension

  with slim.arg_scope([slim.conv2d], kernel_size=1, activation_fn=None):
    for i in range(num_layers):
      in_n = int(net.shape[-1])
      net = slim.conv2d(
          net, num_hidden_channels,
          # this is untruncated version of tf.variance_scaling_initializer
          weights_initializer=tf.random_normal_initializer(0.0, np.sqrt(1.0/in_n)),
      )
      if normalize:
        net = slim.instance_norm(net)
      net = activation_fn(net)

    rgb = slim.conv2d(net, num_output_channels, activation_fn=tf.nn.sigmoid,
                      weights_initializer=tf.zeros_initializer())
  return rgb

from lucid.misc.io.serialize_array import _normalize_array

def render_story(obj_str, lr=0.004, step_n=512,
                 normalize=False,
                 activation_fn=composite_activation):
  sess = create_session()

  # Set up optimization problem
  size = 224
  t_size = tf.placeholder_with_default(size, [])
  T = render.make_vis_T(
      model, obj_str, 
      param_f=lambda: image_cppn(
          t_size, normalize=normalize, activation_fn=activation_fn),
      transforms=[],
      optimizer=tf.train.AdamOptimizer(lr),
  )
  tf.global_variables_initializer().run()

  # Prepare video writer and filenames
  subst = {ord(':'):'_', ord('/'):'_'}
  out_name = 'xy2rgb_' + obj_str.translate(subst)
  video_fn = out_name + '.mp4'
  writer = FFMPEG_VideoWriter(video_fn, (size, size), 60.0)

  # Optimization loop
  try:
    for i in range(step_n):
      _, loss, img = sess.run([T("vis_op"), T("loss"), T("input")])
      writer.write_frame(_normalize_array(img))
      if i > 0 and i % 50 == 0:
        clear_output()
        print("%d / %d  score: %f"%(i, step_n, loss))
        show(img)
  except KeyboardInterrupt:
    pass
  finally:
    writer.close()

  # Show the resulting video
  clear_output()
  display(mpy.ipython_display(video_fn, height=400))

  # Save trained variables
  train_vars = sess.graph.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES)
  params = np.array(sess.run(train_vars), object)
  save(params, out_name + '.npy')

  # Save final image
  final_img = T("input").eval({t_size: 400})
  save(final_img, out_name+'.jpg', quality=90)

render_story('mixed4c_pool_reduce_pre_relu:5')

render_story('softmax0_pre_activation/matmul:316')

render_story('mixed4b_3x3_pre_relu:77')

render_story('mixed4d_3x3_bottleneck_pre_relu:114')

render_story('mixed4e_3x3_pre_relu:120')

render_story('head0_bottleneck_pre_relu:0')

render_story('mixed4d_3x3_bottleneck_pre_relu:139')

def interpolate_params(param1, param2, duration=5.0, size=224):

  def frame(t):
    t = t / duration
    t = (1.0-np.cos(2.0*np.pi*t))/2.0       # looping & easing
    params = param1*(1.0-t) + param2*t      # blending
    params *= 1.0 + t*(1.0-t)               # exaggerating
    img = render_params(params, size=size)
    return _normalize_array(img)

  clip = mpy.VideoClip(frame, duration=duration)
  clip.write_videofile('tmp.mp4', fps=30.0)
  display(mpy.ipython_display('tmp.mp4', height=400))