# Copyright 2019 The PlaNet Authors. All rights reserved.
#
# 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
#
# 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.

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import functools
import os

import numpy as np
import tensorflow as tf

from planet import control
from planet import models
from planet import networks
from planet import tools
from planet.scripts import tasks as tasks_lib


def default(config, params):
  config.debug = False
  config.zero_step_losses = tools.AttrDict(_unlocked=True)
  config = _data_processing(config, params)
  config = _model_components(config, params)
  config = _tasks(config, params)
  config = _loss_functions(config, params)
  config = _training_schedule(config, params)
  return config


def debug(config, params):
  with params.unlocked:
    params.collect_every = 20
    params.batch_shape = [5, 10]
    params.train_steps = 30
    params.test_steps = 30
    params.max_steps = 100 * (30 * 30)
    params.collect_every = 20
    params.num_seed_episodes = 2
  config = default(config, params)
  config.debug = True
  return config


def _data_processing(config, params):
  config.max_episodes = None
  config.scan_episodes_every = params.get('scan_episodes_every', 10)
  config.data_loader = params.get('data_loader', 'scan')
  config.batch_shape = params.get('batch_shape', (50, 50))  # training batch size, horizon
  config.num_chunks = params.get('num_chunks', 1)
  image_bits = params.get('image_bits', 8)
  config.preprocess_fn = functools.partial(
      tools.preprocess.preprocess, bits=image_bits)
  config.postprocess_fn = functools.partial(
      tools.preprocess.postprocess, bits=image_bits)
  config.open_loop_context = 5
  return config


def _model_components(config, params):
  network = getattr(networks, params.get('network', 'conv_ha'))
  config.encoder = network.encoder
  config.decoder = network.decoder
  config.heads = tools.AttrDict(image=config.decoder)
  size = params.get('model_size', 200)
  state_size = params.get('state_size', 30)
  model = params.get('model', 'rssm')
  if model == 'ssm':
    config.cell = functools.partial(
        models.SSM, state_size, size, params.mean_only,
        params.get('min_stddev', 1e-1))
  elif model == 'rssm':  # use rssm model
    config.cell = functools.partial(
        models.RSSM, state_size, size, size, params.mean_only,
        params.get('min_stddev', 1e-1))
  else:
    raise NotImplementedError("Unknown model '{}.".format(params.model))
  return config


def _tasks(config, params):
  tasks = params.get('tasks', ['cheetah_run'])
  if tasks == 'all':
    tasks = [
        'cartpole_balance', 'cartpole_swingup', 'finger_spin', 'cheetah_run',
        'cup_catch', 'walker_walk', 'pendulum', 'carla']
  tasks = [getattr(tasks_lib, name)(config, params) for name in tasks]
  def common_spaces_ctor(task, action_spaces):
    env = task.env_ctor()
    env = control.wrappers.SelectObservations(env, ['image'])
    env = control.wrappers.PadActions(env, action_spaces)
    return env
  if len(tasks) > 1:
    action_spaces = [task.env_ctor().action_space for task in tasks]
    for index, task in enumerate(tasks):
      env_ctor = functools.partial(common_spaces_ctor, task, action_spaces)
      tasks[index] = tasks_lib.Task(
          task.name, env_ctor, task.max_length, ['reward'])
  for name in tasks[0].state_components:
    config.heads[name] = networks.feed_forward
    config.zero_step_losses[name] = 1.0
  config.tasks = tasks
  return config


def _loss_functions(config, params):
  config.free_nats = params.get('free_nats', 2.0)
  config.stop_os_posterior_gradient = True
  config.zero_step_losses.image = params.get('image_loss_scale', 1.0)
  config.zero_step_losses.divergence = params.get('divergence_scale', 1.0)
  config.zero_step_losses.global_divergence = params.get('global_divergence_scale', 0.1)
  config.zero_step_losses.reward = params.get('reward_scale', 10.0)
  config.overshooting = params.get('overshooting', config.batch_shape[1] - 1)
  config.overshooting_losses = config.zero_step_losses.copy(_unlocked=True)
  config.overshooting_losses.reward = params.get(
      'overshooting_reward_scale', 100.0)
  del config.overshooting_losses['image']
  del config.overshooting_losses['global_divergence']
  config.optimizers = _define_optimizers(config, params)
  return config


def _training_schedule(config, params):
  config.train_steps = int(params.get('train_steps', 50000))
  config.test_steps = int(params.get('test_steps', 100))
  config.max_steps = int(params.get('max_steps', 6e7))
  config.train_log_every = config.train_steps
  config.train_checkpoint_every = None
  config.test_checkpoint_every = int(
      params.get('checkpoint_every', config.test_steps))
  config.savers = [tools.AttrDict(exclude=(r'.*_temporary.*',))]
  config.mean_metrics_every = config.train_steps // 10
  config.train_dir = os.path.join(params.logdir, 'train_episodes')
  config.test_dir = os.path.join(params.logdir, 'test_episodes')
  config.random_collects = _initial_collection(config, params)
  config.sim_collects = _active_collection(config, params)
  config.sim_summaries = tools.AttrDict(_unlocked=True)
  for task in config.tasks:
    for horizon in params.get('summary_horizons', [12]):
      name = 'summary-{}-cem-{}'.format(task.name, horizon)
      config.sim_summaries[name] = _define_simulation(
          task, config, params, horizon, 1)
  return config


def _define_optimizers(config, params):
  optimizers = tools.AttrDict(_unlocked=True)
  gradient_heads = params.get('gradient_heads', ['image', 'reward'])
  assert all(head in config.heads for head in gradient_heads)
  diagnostics = r'.*/head_(?!{})[a-z]+/.*'.format('|'.join(gradient_heads))
  kwargs = dict(
      optimizer_cls=functools.partial(tf.train.AdamOptimizer, epsilon=1e-4),
      learning_rate=params.get('learning_rate', 1e-3),
      schedule=functools.partial(tools.schedule.linear, ramp=10000),
      clipping=params.get('gradient_clipping', 1000.0))
  optimizers.main = functools.partial(
      tools.CustomOptimizer, include=r'.*', exclude=diagnostics, **kwargs)
  for name in config.heads:
    assert config.zero_step_losses.get(name), name
    # Diagnostic heads use separate optimizers to not interfere with the model.
    if name in gradient_heads:
      continue
    optimizers[name] = functools.partial(
        tools.CustomOptimizer, include=r'.*/head_{}/.*'.format(name), **kwargs)
  return optimizers


def _initial_collection(config, params):
  num_seed_episodes = params.get('num_seed_episodes', 5)
  sims = tools.AttrDict(_unlocked=True)
  for task in config.tasks:
    sims['train-' + task.name] = tools.AttrDict(
        task=task,
        save_episode_dir=config.train_dir,
        num_episodes=num_seed_episodes)
    sims['test-' + task.name] = tools.AttrDict(
        task=task,
        save_episode_dir=config.test_dir,
        num_episodes=num_seed_episodes)
  return sims


def _active_collection(config, params):
  sims = tools.AttrDict(_unlocked=True)
  batch_size = params.get('collect_batch_size', 1)
  for task in config.tasks:
    for index, horizon in enumerate(params.get('collect_horizons', [12])):
      sim = _define_simulation(task, config, params, horizon, batch_size)
      sim.unlock()
      sim.save_episode_dir = config.train_dir
      sim.steps_after = params.get('collect_every', 5000)
      sim.steps_every = params.get('collect_every', 5000)
      sim.exploration = tools.AttrDict(
          scale=params.get('exploration_noises', [0.3])[index],  # adding noise on action
          schedule=functools.partial(
              tools.schedule.linear,
              ramp=params.get('exploration_ramps', [0])[index]))
      sims['train-{}-cem-{}'.format(task.name, horizon)] = sim
      if params.get('collect_test', False):
        sim = sim.copy()
        sim.save_episode_dir = config.test_dir
        sims['test-{}-cem-{}'.format(task.name, horizon)] = sim
  return sims


def _define_simulation(task, config, params, horizon, batch_size):
  def objective(state, graph):
    return graph.heads['reward'](graph.cell.features_from_state(state)).mean()
  planner = functools.partial(
      control.planning.cross_entropy_method,
      amount=params.get('cem_amount', 1000),
      topk=params.get('cem_topk', 100),
      iterations=params.get('cem_iterations', 10),
      horizon=horizon)
  return tools.AttrDict(
      task=task,
      batch_size=batch_size,
      planner=planner,
      objective=objective)
