import matplotlib.pyplot as plt

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


# The path to a TF-Hub module.
MODULE_PATH = "https://tfhub.dev/deepmind/spiral/default-wgangp-celebahq64-gen-19steps/agent4/1"
# The folder containing `libmypaint` brushes.
BRUSHES_PATH = "the/path/to/libmypaint-brushes"

# Here, we create an environment.
env = libmypaint.LibMyPaint(episode_length=20,
                            canvas_width=64,
                            grid_width=32,
                            brush_type="classic/dry_brush",
                            brush_sizes=[1, 2, 4, 8, 12, 24],
                            use_color=True,
                            use_pressure=True,
                            use_alpha=False,
                            background="white",
                            brushes_basedir=BRUSHES_PATH)


# Now we load the agent from a snapshot.
initial_state, step = agent_utils.get_module_wrappers(MODULE_PATH)

# Everything is ready for sampling.
state = initial_state()
noise_sample = np.random.normal(size=(10,)).astype(np.float32)

time_step = env.reset()
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)

# Show the sample.
plt.close("all")
plt.imshow(time_step.observation["canvas"], interpolation="nearest")

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


# This where we're going to put our TF-Hub module.
TARGET_PATH = ...
# A path to a checkpoint of the trained model.
CHECKPOINT_PATH = ...

# We will need to create an environment in order to obtain the specifications
# for the agent's action and the observation.
env = libmypaint.LibMyPaint(...)

# Here, we wrap a Sonnet module constructor for our agent in a function.
# This is to avoid contaminating the default tensorflow graph.
def agent_ctor():
  return default_agent.Agent(action_spec=env.action_spec(),
                             input_shape=(64, 64),
                             grid_shape=(32, 32),
                             action_order="libmypaint")

# Finally, export a TF-Hub module. We need to specify which checkpoint to use
# to extract the weights for the agent. Since the variable names in the
# checkpoint may differ from the names in the Sonnet module produced by
# `agent_ctor`, we may also want to provide an appropriate name mapping
# function.
agent_utils.export_hub_module(agent_ctor=agent_ctor,
                              observation_spec=env.observation_spec(),
                              noise_dim=10,
                              module_path=TARGET_PATH,
                              checkpoint_path=CHECKPOINT_PATH,
                              name_transform_fn=lambda name: ...)