import copy
import os

import matplotlib.pyplot as plt
import numpy as np
from scipy import ndimage
import tensorflow as tf
import tensorflow_hub as hub

import spiral.agents.default as default_agent
import spiral.agents.utils as agent_utils
from spiral.environments import fluid
from spiral.environments import libmypaint


nest = tf.contrib.framework.nest

# Disable TensorFlow debug output.
tf.logging.set_verbosity(tf.logging.ERROR)

"""
Explanation: Boilerplate
End of explanation
"""


# The path to libmypaint brushes.
BRUSHES_BASEDIR = os.path.join(os.getcwd(), "..", "third_party/mypaint-brushes-1.3.0")
BRUSHES_BASEDIR = os.path.abspath(BRUSHES_BASEDIR)
# The path to a TF-Hub module.
MODULE_PATH = "https://tfhub.dev/deepmind/spiral/default-wgangp-celebahq64-gen-19steps/agent4/1"

"""
Explanation: Sample from the model
This section demonstrates the most basic usage of the package, i.e., sampling from a pre-trained model.
End of explanation
"""


env_settings = dict(
    episode_length=20,                 # Number of frames in each episode.
    canvas_width=64,                   # The width of the canvas in pixels.
    grid_width=32,                     # The width of the action grid.
    brush_type="classic/dry_brush",    # The type of the brush.
    brush_sizes=[1, 2, 4, 8, 12, 24],  # The sizes of the brush to use.
    use_color=True,                    # Color or black & white output?
    use_pressure=True,                 # Use pressure parameter of the brush?
    use_alpha=False,                   # Drop or keep the alpha channel of the canvas?
    background="white",                # Background could either be "white" or "transparent".
    brushes_basedir=BRUSHES_BASEDIR,   # The location of libmypaint brushes.
)
env = libmypaint.LibMyPaint(**env_settings)

"""
Explanation: Create the environment
First, we need to create the environment.
End of explanation
"""


initial_state, step = agent_utils.get_module_wrappers(MODULE_PATH)

"""
Explanation: Create the agent
We provide a convenience function get_module_wrappers that returns two python functions implementing the agent.
The first one, initial_state, is used to get the initial state of the agent (specifically, the state of the LSTM).
The second, step, takes an observation from the environment and performs a single agent step.
End of explanation
"""


state = initial_state()

"""
Explanation: Run sampling
End of explanation
"""


noise_sample = np.random.normal(size=(10,)).astype(np.float32)

time_step = env.reset()
actions = []
for t in range(19):
    time_step.observation["noise_sample"] = noise_sample
    action, state = step(time_step.step_type, time_step.observation, state)
    time_step = env.step(action)
    actions.append(action)
    
plt.close("all")
plt.figure(figsize=(5, 5))
plt.imshow(time_step.observation["canvas"], interpolation="nearest")

"""
Explanation: Running the following cell will return a sample from the model. You can execute it multiple times until you get the one that you like. In addition to showing the final state of the canvas we also record all the agent's actions (note actions variable) needed to reproduce it.
End of explanation
"""


# Let's make the canvas 8x8 times bigger.
SCALE_FACTOR = 8

# Patch the environments setting for higher resolution.
hires_env_settings = copy.deepcopy(env_settings)
hires_env_settings["canvas_width"] *= SCALE_FACTOR
hires_env_settings["brush_sizes"] = [
    s * SCALE_FACTOR for s in hires_env_settings["brush_sizes"]]

env = libmypaint.LibMyPaint(**hires_env_settings)

"""
Explanation: Manipulate the sample
Let's now do something the obtained sample. Since we have the corresponding sequence of actions we can re-render the image in higher resolution. To that end, we will need to create one more environment with modified settings.
Create a high-resolution version of the environment
End of explanation
"""


env.reset()
for t in range(19):
    time_step = env.step(actions[t])
    
plt.close("all")
plt.figure(figsize=(5, 5))
plt.imshow(time_step.observation["canvas"], interpolation="nearest")

"""
Explanation: Execute pre-recorded actions
End of explanation
"""


modified_actions = copy.deepcopy(actions)
for action in modified_actions:
    action["size"] = np.array(0, dtype=np.int32)
    action["pressure"] = np.array(2, dtype=np.int32)

env.reset()
for t in range(19):
    time_step = env.step(modified_actions[t])
    
plt.close("all")
plt.figure(figsize=(5, 5))
plt.imshow(time_step.observation["canvas"], interpolation="nearest")

"""
Explanation: Change the thickness of the brush strokes
In addition to changing the resolution of the images, let's in introduce some more subtle structural changes. We could, for example, change the thickness of all the strokes.
End of explanation
"""


# Patch the environments setting for a different brush.
pen_env_settings = copy.deepcopy(hires_env_settings)
pen_env_settings["brush_type"] = "classic/pen"

env = libmypaint.LibMyPaint(**pen_env_settings)

env.reset()
for t in range(19):
    time_step = env.step(modified_actions[t])
    
plt.close("all")
plt.figure(figsize=(5, 5))
plt.imshow(time_step.observation["canvas"], interpolation="nearest")

"""
Explanation: Change the brush type
Finally, let's re-render the image above using a different brush type.
End of explanation
"""


# The path to the shaders.
SHADERS_BASEDIR = os.path.join(os.getcwd(), "..", "third_party/paint/shaders")
SHADERS_BASEDIR = os.path.abspath(SHADERS_BASEDIR)
# The path to a TF-Hub module.
MODULE_PATH = "https://tfhub.dev/deepmind/spiral/default-fluid-gansn-celebahq64-gen-19steps/1"

env_settings = dict(
    episode_length=20,                               # Number of frames in each episode.
    canvas_width=256,                                # The width of the canvas in pixels.
    grid_width=32,                                   # The width of the action grid.
    brush_sizes=[2.5, 5.0, 10.0, 20.0, 40.0, 80.0],  # The sizes of the brush to use.
    shaders_basedir=SHADERS_BASEDIR,                 # The location of shaders.
)
env = fluid.FluidPaint(**env_settings)

initial_state, step = agent_utils.get_module_wrappers(MODULE_PATH)

state = initial_state()

noise_sample = np.random.normal(size=(10,)).astype(np.float32)

time_step = env.reset()
actions = []
for t in range(19):
    time_step.observation["noise_sample"] = noise_sample
    # The environment uses 256x256 canvas but the agent requires 64x64 input.
    ratio = 64 / 256
    time_step.observation["canvas"] = ndimage.zoom(
        time_step.observation["canvas"], [ratio, ratio, 1], order=1)
    action, state = step(time_step.step_type, time_step.observation, state)
    time_step = env.step(action)
    actions.append(action)
    
plt.close("all")
plt.figure(figsize=(5, 5))
plt.imshow(time_step.observation["canvas"], interpolation="nearest")

"""
Explanation: Fluid Paint environment demonstration
Fluid Paint environment works almost exactly like libmypaint.LibMyPaint. Below, we show how to obtain samples from the model trained in this environment.
End of explanation
"""


