#!/usr/bin/env python3

from __future__ import print_function

import os
import sys
from argparse import ArgumentParser
import gzip
import xml.etree.ElementTree as etree


my_dir = os.path.dirname(os.path.abspath(__file__))
returnn_dir = os.path.dirname(my_dir)
sys.path.insert(0, returnn_dir)


class BlissItem:
  def __init__(self, segment_name, recording_filename, start_time, end_time, orth):
    """
    :param str segment_name:
    :param str recording_filename:
    :param float start_time:
    :param float end_time:
    :param str orth:
    """
    self.segment_name = segment_name
    self.recording_filename = recording_filename
    self.start_time = start_time
    self.end_time = end_time
    self.orth = orth

  def __repr__(self):
    keys = ["segment_name", "recording_filename", "start_time", "end_time", "orth"]
    return "BlissItem(%s)" % ", ".join(["%s=%r" % (key, getattr(self, key)) for key in keys])

  @property
  def delta_time(self):
    return self.end_time - self.start_time


def iter_bliss(filename):
  """
  :param str filename:
  :return: yields BlissItem
  :rtype: list[BlissItem]
  """
  corpus_file = open(filename, 'rb')
  if filename.endswith(".gz"):
    corpus_file = gzip.GzipFile(fileobj=corpus_file)

  context = iter(etree.iterparse(corpus_file, events=('start', 'end')))
  _, root = next(context) # get root element
  name_tree = [root.attrib["name"]]
  elem_tree = [root]
  count_tree = [0]
  recording_filename = None
  for event, elem in context:
    if elem.tag == "recording":
      recording_filename = elem.attrib["audio"] if event == "start" else None
    if event == 'end' and elem.tag == "segment":
      elem_orth = elem.find("orth")
      orth_raw = elem_orth.text  # should be unicode
      orth_split = orth_raw.split()
      orth = " ".join(orth_split)
      segment_name = "/".join(name_tree)
      yield BlissItem(
        segment_name=segment_name, recording_filename=recording_filename,
        start_time=float(elem.attrib["start"]), end_time=float(elem.attrib["end"]),
        orth=orth)
      root.clear()  # free memory
    if event == "start":
      count_tree[-1] += 1
      count_tree.append(0)
      elem_tree += [elem]
      elem_name = elem.attrib.get("name", None)
      if elem_name is None:
        elem_name = str(count_tree[-2])
      assert isinstance(elem_name, str)
      name_tree += [elem_name]
    elif event == "end":
      assert elem_tree[-1] is elem
      elem_tree = elem_tree[:-1]
      name_tree = name_tree[:-1]
      count_tree = count_tree[:-1]


def main():
  arg_parser = ArgumentParser()
  arg_parser.add_argument("bliss_filename")
  arg_parser.add_argument("--subset_segment_file")
  arg_parser.add_argument("--output_type", default="", help="e.g. segment_name")
  arg_parser.add_argument("--merge_swb_ab", action="store_true")
  arg_parser.add_argument("--sort_by_time", action="store_true")
  arg_parser.add_argument("--merge_segs_up_to_time", type=float)
  args = arg_parser.parse_args()
  subset_segment_list = None
  if args.subset_segment_file:
    subset_segment_list = set(open(args.subset_segment_file).read().splitlines())
  rec_filenames = set()
  items_by_rec = {}
  for bliss_item in iter_bliss(args.bliss_filename):
    if subset_segment_list and bliss_item.segment_name not in subset_segment_list:
      continue
    rec_name = bliss_item.recording_filename
    assert rec_name, "invalid item %r" % bliss_item
    if args.merge_swb_ab:
      rec_name = os.path.basename(rec_name)
      rec_name, _ = os.path.splitext(rec_name)
      rec_filenames.add(rec_name)
      assert rec_name[-1] in "AB"
      rec_name = rec_name[:-1]
    else:
      rec_filenames.add(rec_name)
    items_by_rec.setdefault(rec_name, []).append(bliss_item)
  assert items_by_rec
  if args.merge_swb_ab:
    if subset_segment_list:
      for key in list(items_by_rec.keys()):
        if key + "A" not in rec_filenames or key + "B" not in rec_filenames:
          del items_by_rec[key]
      assert items_by_rec, "rec_filenames %r" % (rec_filenames,)
    else:
      for key in items_by_rec.keys():
        assert key + "A" in rec_filenames
        assert key + "B" in rec_filenames
  for key, ls in items_by_rec.items():
    assert isinstance(ls, list)
    if args.sort_by_time:
      ls.sort(key=lambda item: item.start_time)
  if args.merge_segs_up_to_time:
    for key, ls in items_by_rec.items():
      i = 0
      while i < len(ls):
        j = i + 1
        dt = ls[i].delta_time
        while j < len(ls):
          if dt + ls[j].delta_time > args.merge_segs_up_to_time:
            break
          dt += ls[j].delta_time
          j += 1
        if j > i + 1:
          ls[i:j] = [BlissItem(
            segment_name=";".join([item.segment_name for item in ls[i:j]]),
            recording_filename=ls[i].recording_filename,  # might be wrong if merge_swb_ab...
            start_time=0.0, end_time=dt,  # dummy times
            orth=" ".join([item.orth for item in ls[i:j]]))]
        i += 1
  output_types = args.output_type.split(",")
  for key, ls in items_by_rec.items():
    assert isinstance(ls, list)
    for item in ls:
      assert isinstance(item, BlissItem)
      if not output_types:
        print(item)
      else:
        print(" ".join([str(getattr(item, key)) for key in output_types]))


if __name__ == "__main__":
  import better_exchook
  better_exchook.install()
  try:
    main()
  except BrokenPipeError:
    print("BrokenPipeError")
    sys.exit(1)
