import collections
import io
import re
import warnings
from datetime import datetime

import embodied
import numpy as np


def train_save(agent, env, replay, logger, args):

  logdir = embodied.Path(args.logdir)
  logdir.mkdirs()
  print('Logdir:', logdir)
  should_expl = embodied.when.Until(args.expl_until)
  should_train = embodied.when.Ratio(args.train_ratio / args.batch_steps)
  should_log = embodied.when.Clock(args.log_every)
  should_save = embodied.when.Clock(args.save_every)
  step = logger.step
  metrics = embodied.Metrics()
  print('Observation space:')
  for key, value in env.obs_space.items():
    if not key.startswith('log_'):
      print(f'  {key:<16} {value}')
  print('Action space:')
  for key, value in env.act_space.items():
    print(f'  {key:<16} {value}')

  timer = embodied.Timer()
  timer.wrap('agent', agent, ['policy', 'train', 'report', 'save'])
  timer.wrap('env', env, ['step'])
  timer.wrap('replay', replay, ['add', 'save'])
  timer.wrap('logger', logger, ['write'])

  nonzeros = set()
  def per_episode(ep):
    length = len(ep['reward']) - 1
    score = float(ep['reward'].astype(np.float64).sum())
    sum_abs_reward = float(np.abs(ep['reward']).astype(np.float64).sum())
    logger.add({
        'length': length,
        'score': score,
        'sum_abs_reward': sum_abs_reward,
        # TODO
        # 'reward_rate': (ep['reward'] - ep['reward'].min() >= 0.1).mean(),
        'reward_rate': (np.abs(ep['reward']) >= 0.5).mean(),
    }, prefix='episode')
    print(f'Episode has {length} steps and return {score:.1f}.')
    stats = {}
    for key in args.log_keys_video:
      if key in ep:
        stats[f'policy_{key}'] = ep[key]
    for key, value in ep.items():
      if not args.log_zeros and key not in nonzeros and (value == 0).all():
        continue
      nonzeros.add(key)
      if re.match(args.log_keys_sum, key):
        stats[f'sum_{key}'] = ep[key].sum()
      if re.match(args.log_keys_mean, key):
        stats[f'mean_{key}'] = ep[key].mean()
      if re.match(args.log_keys_max, key):
        stats[f'max_{key}'] = ep[key].max(0).mean()
    metrics.add(stats, prefix='stats')

  epsdir = embodied.Path(args.logdir) / 'saved_episodes'
  epsdir.mkdirs()
  print('Saving episodes:', epsdir)
  def save(ep):
    time = datetime.now().strftime("%Y%m%dT%H%M%S")
    uuid = str(embodied.uuid())
    score = str(np.round(ep['reward'].sum(), 1)).replace('-', 'm')
    length = len(ep['reward'])
    filename = epsdir / f'{time}-{uuid}-len{length}-rew{score}.npz'
    with io.BytesIO() as stream:
      np.savez_compressed(stream, **ep)
      stream.seek(0)
      filename.write(stream.read(), mode='wb')
    print('Saved episode:', filename)
  saver = embodied.Worker(save, 'thread')

  driver = embodied.Driver(env)
  driver.on_episode(lambda ep, worker: per_episode(ep))
  driver.on_episode(lambda ep, worker: saver(ep))
  driver.on_step(lambda tran, _: step.increment())
  driver.on_step(replay.add)

  train_fill = max(args.batch_steps, args.train_fill - len(replay))
  if train_fill:
    print(f'Fill train dataset ({train_fill} steps).')
    random_agent = embodied.RandomAgent(env.act_space)
    driver(random_agent.policy, steps=train_fill)
  logger.add(metrics.result())
  logger.write()

  dataset = iter(agent.dataset(replay.dataset))
  state = [None]  # To be writable from train step function below.
  assert args.pretrain > 0  # At least one step to initialize variables.
  for _ in range(args.pretrain):
    with timer.scope('dataset'):
      batch = next(dataset)
    _, state[0], _ = agent.train(batch, state[0])

  batch = [None]
  def train_step(tran, worker):
    for _ in range(should_train(step)):
      with timer.scope('dataset'):
        batch[0] = next(dataset)
      outs, state[0], mets = agent.train(batch[0], state[0])
      metrics.add(mets, prefix='train')
      if 'priority' in outs:
        replay.prioritize(outs['key'], outs['priority'])
    if should_log(step):
      agg = metrics.result()
      report = agent.report(batch[0])
      report = {k: v for k, v in report.items() if 'train/' + k not in agg}
      logger.add(agg)
      logger.add(report, prefix='report')
      logger.add(replay.stats, prefix='replay')
      logger.add(timer.stats(), prefix='timer')
      logger.write(fps=True)
  driver.on_step(train_step)

  checkpoint = embodied.Checkpoint(logdir / 'checkpoint.pkl')
  timer.wrap('checkpoint', checkpoint, ['save', 'load'])
  checkpoint.step = step
  checkpoint.agent = agent
  checkpoint.replay = replay
  checkpoint.load_or_save()
  should_save(step)  # Register that we jused saved.

  print('Start training loop.')
  policy = lambda *args: agent.policy(
      *args, mode='explore' if should_expl(step) else 'train')
  while step < args.steps:
    driver(policy, steps=1000)
    if should_save(step):
      checkpoint.save()
  logger.write()

