import argparse
import os
from base.proto import ReadProto,WriteProtoAsText
from mined_metric.builder.proto.data_loader_options_pb2 import ChumDataEntry, ChumDataEntrySet

kSimulationRange = 1 * 60

def main():
  parser = argparse.ArgumentParser('Analyzes the directory of chums generated from vnv_track_annotation_to_chum and generates chum_data_entries out of each')
  parser.add_argument(
      '--path',
      type=str,
      required=True,
      help='Path to the directory that includes the _vars generated from vnv_track_annotation_to_chum',
  )

  args = parser.parse_args()
  main_path = args.path
  names = os.listdir(main_path)
  check("_vars" in names, "Path doesn't have _vars directory in it")

  all_variants = os.listdir(os.path.join(main_path, "_vars"))
  chum_data_entry_set = ChumDataEntrySet()
  for idx,variant in enumerate(all_variants):
    if (idx % 1000 == 0): print(f"Done with {idx}/{len(all_variants)}")
    if (not variant.startswith("20")): continue
    # create the full chum
    timestamped_chum = variant.replace("A", "@").replace("P", ".")
    timestamp = float(timestamped_chum.split("@")[1])
    end_timestamp = timestamp + kSimulationRange
    full_chum = f"chum://{timestamped_chum}-{end_timestamp:.2f}?i={main_path}&v={variant}"
    entry = chum_data_entry_set.chum_data_entries.add()
    entry.chum_uri = full_chum
  out_path = os.path.join(main_path, "chum_entries.pbtxt")
  WriteProtoAsText(out_path, chum_data_entry_set)
  print(f"Wrote {len(chum_data_entry_set.chum_data_entries)} entries to {out_path}")

def check(condition, message):
  if not condition:
    raise RuntimeError(message)



if __name__ == "__main__":
  main()