# Copyright 2018 The TensorFlow 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.
# ==============================================================================
"""Code to extract a series of positions + their next moves from an SGF.

Most of the complexity here is dealing with two features of SGF:
- Stones can be added via "play move" or "add move", the latter being used
  to configure L+D puzzles, but also for initial handicap placement.
- Plays don't necessarily alternate colors; they can be repeated B or W moves
  This feature is used to handle free handicap placement.
"""
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

from collections import namedtuple

import coords
import go
from go import Position, PositionWithContext
import numpy as np
import sgf
import utils


SGF_TEMPLATE = """(;GM[1]FF[4]CA[UTF-8]AP[Minigo_sgfgenerator]RU[{ruleset}]
SZ[{boardsize}]KM[{komi}]PW[{white_name}]PB[{black_name}]RE[{result}]
{game_moves})"""

PROGRAM_IDENTIFIER = "Minigo"


def translate_sgf_move_qs(player_move, q):
  return "{move}C[{q:.4f}]".format(
    move=translate_sgf_move(player_move), q=q)


def translate_sgf_move(player_move, comment):
  if player_move.color not in (go.BLACK, go.WHITE):
    raise ValueError("Can't translate color %s to sgf" % player_move.color)
  c = coords.to_sgf(player_move.move)
  color = 'B' if player_move.color == go.BLACK else 'W'
  if comment is not None:
    comment = comment.replace(']', r'\]')
    comment_node = "C[{}]".format(comment)
  else:
    comment_node = ""
  return ";{color}[{coords}]{comment_node}".format(
    color=color, coords=c, comment_node=comment_node)


def make_sgf(board_size,
  move_history,
  result_string,
  ruleset="Chinese",
  komi=7.5,
  white_name=PROGRAM_IDENTIFIER,
  black_name=PROGRAM_IDENTIFIER,
  comments=[]
):
  """Turn a game into SGF.

  Doesn't handle handicap games or positions with incomplete history.

  Args:
    move_history: iterable of PlayerMoves
    result_string: "B+R", "W+0.5", etc.
    comments: iterable of string/None. Will be zipped with move_history.
  """
  try:
    # Python 2
    from itertools import izip_longest
    zip_longest = izip_longest
  except ImportError:
    # Python 3
    from itertools import zip_longest

  boardsize = board_size
  game_moves = ''.join(translate_sgf_move(*z)
             for z in zip_longest(move_history, comments))
  result = result_string
  return SGF_TEMPLATE.format(**locals())


def sgf_prop(value_list):
  'Converts raw sgf library output to sensible value'
  if value_list is None:
    return None
  if len(value_list) == 1:
    return value_list[0]
  else:
    return value_list


def sgf_prop_get(props, key, default):
  return sgf_prop(props.get(key, default))


def handle_node(board_size, pos, node):
  'A node can either add B+W stones, play as B, or play as W.'
  props = node.properties
  black_stones_added = [coords.from_sgf(board_size,
    c) for c in props.get('AB', [])]
  white_stones_added = [coords.from_sgf(board_size,
    c) for c in props.get('AW', [])]
  if black_stones_added or white_stones_added:
    return add_stones(pos, black_stones_added, white_stones_added)
  # If B/W props are not present, then there is no move. But if it is present
  # and equal to the empty string, then the move was a pass.
  elif 'B' in props:
    black_move = coords.from_sgf(board_size, props.get('B', [''])[0])
    return pos.play_move(black_move, color=go.BLACK)
  elif 'W' in props:
    white_move = coords.from_sgf(board_size, props.get('W', [''])[0])
    return pos.play_move(white_move, color=go.WHITE)
  else:
    return pos


def add_stones(pos, black_stones_added, white_stones_added):
  working_board = np.copy(pos.board)
  go.place_stones(working_board, go.BLACK, black_stones_added)
  go.place_stones(working_board, go.WHITE, white_stones_added)
  new_position = Position(board=working_board, n=pos.n, komi=pos.komi,
              caps=pos.caps, ko=pos.ko, recent=pos.recent, to_play=pos.to_play)
  return new_position


def get_next_move(board_size, node):
  props = node.next.properties
  if 'W' in props:
    return coords.from_sgf(board_size, props['W'][0])
  else:
    return coords.from_sgf(board_size, props['B'][0])


def maybe_correct_next(pos, next_node):
  if (('B' in next_node.properties and not pos.to_play == go.BLACK) or
      ('W' in next_node.properties and not pos.to_play == go.WHITE)):
    pos.flip_playerturn(mutate=True)


def replay_sgf(board_size, sgf_contents):
  """
  Wrapper for sgf files, returning go.PositionWithContext instances.

  It does NOT return the very final position, as there is no follow up.
  To get the final position, call pwc.position.play_move(pwc.next_move)
  on the last PositionWithContext returned.

  Example usage:
  with open(filename) as f:
    for position_w_context in replay_sgf(f.read()):
      print(position_w_context.position)
  """
  collection = sgf.parse(sgf_contents)
  game = collection.children[0]
  props = game.root.properties
  assert int(sgf_prop(props.get('GM', ['1']))) == 1, "Not a Go SGF!"

  komi = 0
  if props.get('KM') != None:
    komi = float(sgf_prop(props.get('KM')))
  result = utils.parse_game_result(sgf_prop(props.get('RE')))

  pos = Position(komi=komi)
  current_node = game.root
  while pos is not None and current_node.next is not None:
    pos = handle_node(board_size, pos, current_node)
    maybe_correct_next(pos, current_node.next)
    next_move = get_next_move(board_size, current_node)
    yield PositionWithContext(pos, next_move, result)
    current_node = current_node.next


def replay_sgf_file(board_size, sgf_file):
  with open(sgf_file) as f:
    for pwc in replay_sgf(board_size, f.read()):
      yield pwc