# Copyright 2022 DeepMind Technologies Limited.
#
# 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.

"""Generates tfrecords containing converted replay data.

Takes as input a partition file, as generated by generate_partitions.py,
plus a ConverterSettings proto file (in text format), e.g.
alphastar_supervised_converter_settings.pbtxt, converting all replays listed in
the partition from .SC2Replay to .tfrecord.

Note that replays expect the correct version of the StarCraft2 binary. It is
recommended that replays are stored in version-specific directories and
processed a version at a time. The `SC2PATH` environment variable can be used
to control the active SC2 version for each run.
"""

import os
from typing import Mapping

from absl import app
from absl import flags
from absl import logging
from alphastar.unplugged.data import util
import chex
import dm_env
import jax
import numpy as np
from pysc2.env import converted_env
from pysc2.env.converter.proto import converter_pb2
from pysc2.lib import protocol as sc2_protocol
from pysc2.lib import remote_controller
from pysc2.lib import sc_process
from pysc2.lib.replay import replay_converter
import tensorflow as tf
import websocket

from google.protobuf import text_format

FLAGS = flags.FLAGS

flags.DEFINE_string(
    'sc2_replay_path', None, 'Path to .SC2Replay files.', required=True)
flags.DEFINE_string(
    'converted_path', None, 'Path to .tfrecord files.', required=True)
flags.DEFINE_string(
    'partition_file',
    None,
    'Path to partition file to be processed.',
    required=True)
flags.DEFINE_string(
    'converter_settings',
    None,
    'Converter settings to apply (pbtxt, using text_format).',
    required=True)
flags.DEFINE_integer(
    'max_retries', 3,
    'Maximum number of times to retry parsing of a replay (which may be foiled '
    'by websocket errors or similar) before moving on.')


def _generate_episode(replay_file_path: str, player_id: int, replay_data,
                      converter_settings) -> Mapping[str, chex.Array]:
  """Returns an episode as a tree of stacked converted observations."""

  for attempt in range(FLAGS.max_retries + 1):
    try:
      logging.info('Converting %s, player %d.%s', replay_file_path, player_id,
                   ' Attempt {attempt}.' if attempt else '')
      observations = replay_converter.converted_observation_stream(
          replay_data, player_id, converter_settings)
      episode = jax.tree_map(lambda *xs: np.stack(xs), *observations)
    except (
        sc_process.SC2LaunchError,
        remote_controller.ConnectError,
        sc2_protocol.ConnectionError,
        websocket.WebSocketTimeoutException,
    ) as e:
      logging.warning('Ephemeral(?) exception while converting replay %s: %s.',
                      replay_file_path, e)
    except AttributeError as e:
      logging.warning(
          'Giving up on converting replay %s for player %d due to '
          'exception: %s.', replay_file_path, player_id, e)
      return None
    else:
      return episode

  logging.warning(
      'Giving up on converting replay %s for player %d, max retries (%d) '
      'reached.', replay_file_path, player_id, FLAGS.max_retries)


def _step_type_array(length: int):
  step_type = np.full((length,), fill_value=dm_env.StepType.MID, dtype=np.int32)
  step_type[0] = dm_env.StepType.FIRST
  step_type[-1] = dm_env.StepType.LAST
  return step_type


def _process(replay_file_path: str, converter_settings, serializer):
  """Yields ({replay_hash}_{player_id}, serialized episode).

  0, 1 or 2 elements may be yielded, depending on whether the replay is
  successfully parsed and on the number of players present in the replay.

  Args:
    replay_file_path: Path to an .SC2Replay file.
    converter_settings: ConverterSettings proto.
    serializer: For encoding episodes in serialized form.
  """
  with tf.io.gfile.GFile(replay_file_path, 'rb') as f:
    replay_data = f.read()

  replay_hash = os.path.basename(replay_file_path).replace('.SC2Replay', '')
  for player_id in [1, 2]:   # Player IDs in a 2-player SC2 game.
    episode = _generate_episode(replay_file_path, player_id, replay_data,
                                converter_settings)
    if not episode:
      # Skip failed replays.
      continue

    length = episode['player'].shape[0]
    if length == 0:
      # Skip empty replays.
      continue

    episode = dict(observation=episode, step_type=_step_type_array(length))
    serialized_episode = serializer.encode(episode)
    if serialized_episode:
      episode_hash = f'{replay_hash}_{player_id}'
      yield episode_hash, serialized_episode


def _write(episode_hash: str, serialized_episode):
  """Writes a serialized episode to file using TFRecordWriter."""
  filename = episode_hash + '.tfrecord'
  tmp_name = os.path.join(FLAGS.converted_path, '~' + filename)
  final_name = os.path.join(FLAGS.converted_path, filename)
  tf.io.gfile.makedirs(os.path.dirname(tmp_name))
  tf.io.gfile.makedirs(os.path.dirname(final_name))

  # Write to temp file first, and move only after writing is finished,
  # to minimize the chances of partially-written files.
  with tf.io.TFRecordWriter(tmp_name, options=None) as writer:
    writer.write(serialized_episode)
  tf.io.gfile.rename(tmp_name, final_name, overwrite=True)


def main(argv):
  del argv

  with tf.io.gfile.GFile(FLAGS.converter_settings, 'r') as f:
    converter_settings = text_format.Parse(f.read(),
                                           converter_pb2.ConverterSettings())
  obs_spec, _ = converted_env.get_environment_spec(converter_settings)
  features = util.get_dataset_specs(obs_spec)
  serializer = util.TFExampleCoder(features=features, compress=True)

  with tf.io.gfile.GFile(FLAGS.partition_file, 'r') as f:
    replay_hashes = [l.strip() for l in f.readlines()]

  for h in replay_hashes:
    logging.info('Processing replay with hash %s.', h)
    for episode_hash, serialized_episode in _process(
        replay_file_path=os.path.join(FLAGS.sc2_replay_path, f'{h}.SC2Replay'),
        converter_settings=converter_settings,
        serializer=serializer):
      logging.info('Writing episode hash %s.', episode_hash)
      _write(episode_hash, serialized_episode)


if __name__ == '__main__':
  app.run(main)
