# Copyright 2019 Google LLC
#
# 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.
# pylint: disable=unused-import,unused-variable,g-doc-return-or-yield,broad-except,g-short-docstring-punctuation,g-no-space-after-docstring-summary,g-doc-args,g-importing-member,g-multiple-import,g-space-before-docstring-summary
"""opening_freqs.py -- Analyze SGFs for opening pattern frequency information.

Opening_freqs.py looks at subdirectories of SGFs and extracts joseki frequency
information to a sqlite DB. It tracks the first MAX_JOSEKI_LENGTH stones played
in each corner, moving to the next file when all four corners are 'finished'.
The 'sequences' are stored in SGF notation, i.e. "B[qq];W[pd]; ... ".   These
"SGF fragments" are not valid sgf files, but are easily embedded in various SGF
viewing tools, see opening_freqs_export.py for more.
"""

from collections import Counter, defaultdict
import datetime as dt
import functools
from hashlib import sha256
import multiprocessing as mp
import os
import random
import re
import sys

from absl import app
from absl import flags
import coords
from sgfmill import sgf, sgf_moves
import sqlite3
from tqdm import tqdm

FLAGS = flags.FLAGS

MIN_JOSEKI_LENGTH = 5
MAX_JOSEKI_LENGTH = 20

flags.DEFINE_string("in_dir", None, "sgfs here are parsed.")
flags.DEFINE_string("db_path", "joseki.db", "Path to josekidb")
flags.DEFINE_string(
    "run_name", "",
    "If the games should be associated with a 'run' key (e.g., 'v17')")
flags.DEFINE_integer("most_common", 3000,
                     "Number of most-common-joseki to save per hour")
flags.DEFINE_integer("threads", 12, "Number of threads to use.")
flags.DEFINE_float("sample_frac", 1.0,
                   "Fraction of sgfs in a directory to parse")


def sha(sequence):
  return sha256(bytes(sequence)).hexdigest()


SCHEMA = """

CREATE TABLE IF NOT EXISTS joseki (
  id integer primary key,
  seq text,
  length integer,
  num_tenukis integer
);

CREATE TABLE IF NOT EXISTS joseki_counts (
  id integer primary key,
  seq_id integer,
  hour text,
  run text,
  count integer,
  b_wins integer,
  example_sgf text,

  UNIQUE(seq_id, hour, run),
  FOREIGN KEY(seq_id) REFERENCES joseki(id) ON DELETE CASCADE
);

CREATE TABLE IF NOT EXISTS next_moves (
  id integer primary key,
  seq_id integer,         -- select distinct(next_move), sum(count) from next_moves where seq_id='...' and
  joseki_hour_id integer, -- select jc.hour, next_move, count from next_moves where seq='...' join on joseki_counts as jc where jc.id
  next_move text,         -- e.g. 'B[jj];'
  count integer,

  FOREIGN KEY(seq_id) REFERENCES joseki(id) ON DELETE CASCADE,
  FOREIGN KEY(joseki_hour_id) REFERENCES joseki_counts(id) ON DELETE CASCADE
);

CREATE TABLE IF NOT EXISTS game_counts_by_hour(
  id integer primary key,
  hour text,
  count integer
);

CREATE INDEX IF NOT EXISTS seq_index on joseki (seq);
CREATE INDEX IF NOT EXISTS nm_index on next_moves (seq_id);
CREATE INDEX IF NOT EXISTS jc_index on next_moves (joseki_hour_id);
CREATE INDEX IF NOT EXISTS jc_seq_index on joseki_counts (seq_id);
"""

COMPACT_NEXT_MOVES_TABLE_QUERY = """
DELETE FROM next_moves WHERE id IN (SELECT nm.id FROM next_moves AS nm JOIN joseki AS j ON nm.seq_id = j.id WHERE (j.seq || nm.next_
move) NOT IN (SELECT seq FROM joseki));
"""


def run_time_ranges(db):
  ts = lambda hr: int(dt.datetime.strptime(hr, "%Y-%m-%d-%H").timestamp())
  runs = {
      r[0]: (ts(r[1]), ts(r[2])) for r in db.execute("""
        select run, min(hour), max(hour) from joseki_counts group by 1;
        """).fetchall()
  }
  return runs


def build_run_time_transformers(ranges, buckets=250):
  """ Build a dict of functions to transform from a timestamp into a relative

    offset.  E.g.
    input: {'v17': (1234567890, 1235567879) ... }
    output: {'v17': lambda t: (t - min) * (1/max) ... }
  """

  funcs = {}

  def f(t, min_, max_):
    # return "%0.2f" % ((t-min_) * (1/(max_-min_)))
    key = (t - min_) * (1 / (max_ - min_))
    return "%0.3f" % (int(buckets * key) / (buckets / 100.0))

  for run, range_ in ranges.items():
    funcs[run] = functools.partial(f, min_=range_[0], max_=range_[1])

  return funcs


def move_to_corner(move):
  """Take a move ('color', (row,col)).

    where row, col are 0-indexed (from sgfmill)
    Figure out which corner it's in
    Transform it to be in the upper right corner;
    Returns [corners, move]
      corners: A list of the indices of the corners it is in (0-3)
      move: The transformed move
  """

  color, m = move
  if not m:
    return (None, move)
  y, x = m
  corners = []

  if (x >= 9 and y >= 9):
    corners.append(0)
  elif (x >= 9 and y <= 9):
    corners.append(1)
  elif (x <= 9 and y <= 9):
    corners.append(2)
  elif (x <= 9 and y >= 9):
    corners.append(3)

  y -= 9
  x -= 9
  y = abs(y)
  x = abs(x)
  y += 9
  x += 9

  return [corners, (color, (y, x))]


def extract_from_game(game_path):
  """
    In addition to the corner information returned, we also add the winner,
    returning a tuple of ((sequence_counts, nextmove_counts), winner)
    where `winner` == 1 for B win, 0 otherwise.
  """
  with open(game_path) as sgf_file:
    game_data = sgf_file.read().encode("utf-8")

  try:
    g = sgf.Sgf_game.from_bytes(game_data)
    _, moves = sgf_moves.get_setup_and_moves(g)
  except BaseException:
    print("bad file: ", game_path)
    return Counter(), {}

  return (extract_corners(moves), 1 if g.get_winner().lower() == "b" else 0)


def extract_corners(moves):
  """
    Takes a list of moves ('color', (row,col)) (e.g. from sgfmill)
    returns a tuple of sequence_counts, nextmove_counts:
      sequence_counts: a Counter("sequence", count) of all
      sequences/subsequences in all corners
                       of a given game_path
      nextmove_counts: a dict of {"sequence": Counter} objects of the next moves
      of given sequences.
  """
  corner_trackers = [[], [], [], []]
  stop = []

  for m in moves:
    corners, m = move_to_corner(m)
    if not corners:
      continue
    for c in corners:
      if c in stop:
        continue
      corner_trackers[c].append(m)
      if len(corner_trackers[c]) > MAX_JOSEKI_LENGTH:
        stop.append(c)
    if len(stop) == 4:
      break

  # we now have the sequences of the four corners of the board extracted, but
  # not 'canonical', i.e., the reflections across y=x would be distinct.
  # so, canonicalize them
  # this doesn't really work for when the position returns to a
  # symmetrical one, e.g.  4-4, knight approach, tenuki, knight approach the
  # other side = two sets of sequences even though the position is symmetric.

  sequence_counts = Counter()
  next_moves = defaultdict(Counter)
  for c in corner_trackers:
    seq = ""
    canonical = None
    for idx, (color, m) in enumerate(c):
      y, x = m
      if y < x and canonical is None:
        canonical = True
      elif y > x and canonical is None:
        x, y = y, x
        canonical = False
      elif not canonical:
        x, y = y, x
      next_move = sgf_format((color, (18 - y, x)))
      next_moves[seq][next_move] += 1
      seq += next_move
      sequence_counts[seq] += 1

  return sequence_counts, next_moves


def sgf_format(move):
  return "{}[{}];".format(move[0].upper(), coords.to_sgf(move[1]))


def analyze_dir(directory):
  """
    Parses all .sgfs in `directory` and updates database accordingly.

    `directory` is assumed to be a path ending in 'YYYY-mm-dd-HH'.
    e.g. /path/to/games/2019-07-01-00/
  """
  counts = Counter()
  next_moves = defaultdict(Counter)
  example_sgfs = {}
  b_wins = Counter()

  sgf_files = [
      os.path.join(directory, p)
      for p in os.listdir(directory)
      if p.endswith(".sgf")
  ]
  amt = int(len(sgf_files) * FLAGS.sample_frac)
  if FLAGS.sample_frac < 1:
    random.shuffle(sgf_files)

  hr = os.path.basename(directory.rstrip("/"))

  for path in sgf_files[:amt]:
    (corners, seq_nexts), b_win = extract_from_game(path)
    counts.update(corners)
    for seq, next_cts in seq_nexts.items():
      next_moves[seq].update(next_cts)
    for seq in corners:
      b_wins[seq] += b_win
      if seq not in example_sgfs:
        example_sgfs[seq] = os.path.join(hr, os.path.basename(path))

  db = sqlite3.connect(FLAGS.db_path, check_same_thread=False)
  with db:
    for seq, count in counts.most_common(FLAGS.most_common):
      cur = db.cursor()

      s_id = cur.execute("select id from joseki where seq=?", (seq,)).fetchone()
      if s_id:
        s_id = s_id[0]
      else:
        cur.execute(
            """ INSERT INTO joseki(seq, length, num_tenukis) VALUES(?, ?, ?) """,
            (seq, seq.count(";"), count_tenukis(seq)))
        s_id = cur.lastrowid

      cur.execute(
          """
                 INSERT INTO joseki_counts(seq_id, hour, count, run, b_wins, example_sgf) VALUES (?, ?, ?, ?, ?, ?)
                 """,
          (s_id, hr, count, FLAGS.run_name, b_wins[seq], example_sgfs[seq]))
      jc_id = cur.lastrowid

      for next_move, next_count in next_moves[seq].items():
        cur.execute(
            """
                     INSERT INTO next_moves(seq_id, joseki_hour_id, next_move, count) VALUES (?, ?, ?, ?)
                     """, (s_id, jc_id, next_move, next_count))

    db.commit()
  db.close()


def count_tenukis(seq):
  colors = "".join([c for c in seq if c in "BW"])
  # string.count only finds non-overlapping, and 'BWWW' is 2 tenukis
  # thus this weird regex expr.
  return (sum(1 for _ in re.finditer("(?=BB)", colors)) +
          sum(1 for _ in re.finditer("(?=WW)", colors)))


def main(_):
  """Entrypoint for absl.app"""
  db = sqlite3.connect(FLAGS.db_path)
  db.executescript(SCHEMA)
  db.close()

  root = FLAGS.in_dir
  dirs = [
      os.path.join(root, d)
      for d in os.listdir(root)
      if os.path.isdir(os.path.join(root, d))
  ]

  total = len(dirs)
  with mp.Pool(FLAGS.threads) as p:
    list(
        tqdm(
            p.imap_unordered(analyze_dir, dirs),
            desc="Extracting Joseki",
            total=total))


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