import glob
import shutil
import os
import argparse
import re

parser = argparse.ArgumentParser()
parser.add_argument('-i', '--inputname', default="regular", help="Specify the png seqeunce prefix", nargs='?')
parser.add_argument('-s', '--startnum', default=0, help="Specify the start number", nargs='?') # nargs='?" is needed to use the default value
parser.add_argument('-e', '--endnum', default=1679, help="Specify the end number", nargs='?')
parser.add_argument('-st', '--step', default=1, help="Specify the frame step (render every step frame)", nargs='?')

args = parser.parse_args()
name = args.inputname
startnum = int(args.startnum)
endnum = int(args.endnum)
step = int(args.step)

def nums_to_ranges(nums):
  nums = sorted(set(nums))
  gaps = [[s, e] for s, e in zip(nums, nums[1:]) if s+step < e]
  edges = iter(nums[:1] + sum(gaps, []) + nums[-1:])
  return list(zip(edges, edges))

def hrs(size, decimal_places=2):
  for unit in ['B','KB','MB','GB','TB']:
    if size < 1024.0:
      break
    size /= 1024.0
  return f"{size:.{decimal_places}f} {unit}"

def getranges():
  names = []
  allpngs = glob.glob("*.png")
  if len(allpngs) == 0: # render everything if no pngs were found
    os.chdir(orig_dir)
    return [(self.startnum, self.endnum)]
  for png in allpngs:
    regex = r'(.*)\d{4,}\.png'
    prefix = re.search(regex, png).group(1)
    if prefix not in names:
      names.append(prefix)

  all_obtained = []
  all_missing = []
  filesizes = []

  for name in names:
    pngs = glob.glob(f"{name}*.png")
    obtained = []
    missing = []
    filesize = 0
    for ii, png in enumerate(pngs):

      regex = r'.*(\d{4,})\.png'
      val = int(re.search(regex, png).group(1))
      obtained.append(val)
      # filesize += os.stat(png).st_size
      filesize += os.path.getsize(png)

    missing = list(set(range(startnum, endnum+1, step)) - set(obtained))
    all_obtained.append(obtained)
    all_missing.append(missing)
    filesizes.append(hrs(filesize))

  res = all(ii == all_obtained[0] for ii in all_obtained) # check if all elements in list are identical
  for name, obtained, missing, size in list(zip(names, all_obtained, all_missing, filesizes)):
    print(name)
    print(f"Obtained: {nums_to_ranges(obtained)} : {size}")
    print(f"Missing:  {nums_to_ranges(missing)}")
    print("-"*80)

  if res:
    print("All ranges are identical!")
  else:
    print("Some frames are missing!")

getranges()