beatsaber-cli

#!/usr/bin/env python3
#
# beatsaber-cli, by Shawn Presser. July 2021.
#
# requirements:
#
#   pip3 install ansi-styles ansi-escapes python-vlc 
#
# example usage:
#
#   beatsaber-cli ~/Downloads/beatsaber/dna/ExpertPlusStandard.dat
#
# get songs from https://bsaber.com/

#==============================================================================
# Functionality
#==============================================================================
import pdb
import sys
import os
import re

# utility funcs, classes, etc go here.

def asserting(cond):
    if not cond:
        pdb.set_trace()
    assert(cond)

def has_stdin():
    return not sys.stdin.isatty()

def reg(pat, flags=0):
    return re.compile(pat, re.VERBOSE | flags)

#==============================================================================
# Cmdline
#==============================================================================
import argparse

parser = argparse.ArgumentParser(formatter_class=argparse.RawTextHelpFormatter, 
    description="""
TODO
""")
     
parser.add_argument('-v', '--verbose',
    action="store_true",
    help="verbose output" )

args = None

#==============================================================================
# Main
#==============================================================================

import json
from collections import namedtuple

from ansi_styles import ansiStyles as styles

class Renderer:
  def __init__(self):
    self.cmds = []
    self(ansiEscapes.clearScreen)
    self(ansiEscapes.cursorHide)

  def __call__(self, msg):
    self.cmds.append(msg)

  def to(self, x, y):
    self(ansiEscapes.cursorTo(x, y))

  def flush(self):
    print(''.join(self.cmds), end='', flush=True)
    self.cmds.clear()

  def viewport(self, w, h):
    for line in range(h):
      self(ansiEscapes.cursorTo(w, line) + ansiEscapes.eraseStartLine)
    self(ansiEscapes.cursorTo(0, 0))


#arrows = "↑↓←→◜◝◟◞∎"
arrows = "↑↓←→↖↗↙↘∎"

class Note(namedtuple('Note', 'index time lineIndex lineLayer type cutDirection'.split())):
  @property
  def insignia(self):
    return arrows[self.cutDirection]

  @property
  def upwards(self):
    return "1↓←→11↙↘∎"[self.cutDirection] == "1"

  @property
  def downwards(self):
    return "↑1←→↖↗11∎"[self.cutDirection] == "1"

  @property
  def leftwards(self):
    return "↑↓1→1↗1↘∎"[self.cutDirection] == "1"

  @property
  def rightwards(self):
    return "↑↓←1↖1↙1∎"[self.cutDirection] == "1"

  @property
  def left(self):
    return "↑↓1→↖↗↙↘∎"[self.cutDirection] == "1"

  @property
  def right(self):
    return "↑↓←1↖↗↙↘∎"[self.cutDirection] == "1"

  @property
  def up(self):
    return "1↓←→↖↗↙↘∎"[self.cutDirection] == "1"

  @property
  def down(self):
    return "↑1←→↖↗↙↘∎"[self.cutDirection] == "1"


  def color(self, opacity=1.0):
    if self.type == 0:
      return (1.0, 0, 0, opacity)
    elif self.type == 1:
      return (0, 0, 1.0, opacity)
    else:
      return (96/255, 100/255, 105/255, opacity)

  @property
  def w(self):
    return 3

  @property
  def h(self):
    return 1

  def secs(self, bpm):
    return self.time / bpm * 60

  @property
  def x(self):
    return self.lineIndex

  @property
  def y(self):
    return 2 - self.lineLayer

class BeatsaberMap:
  def __init__(self, info='info.dat'):
    if isinstance(info, str):
      with open(info) as f:
        info = json.load(f)
    self.info = info

  def parse(self, file):
    if isinstance(file, str):
      with open(file) as f:
        file = json.load(f)
    notes = file['_notes']
    #notes = [list(sorted(x.items())) for x in notes]
    #notes = [[y[1] for y in x] for x in notes]
    notes = [Note(index=i, **{k.lstrip('_'): v for k, v in x.items()}) for i, x in enumerate(notes)]
    #return Track([(n.type, n.c, n.time / self.bpm * 60000, n.x, n.y) for n in notes])
    return Song(self.info, notes)


from ansi_escapes import ansiEscapes
import subprocess
import vlc

def dovlc(cmd):
  cmd = cmd.replace('\n', '\n  ')
  return subprocess.check_output(['osascript', '-e', f"""
tell application "VLC"
  {cmd}
end tell
""".strip()])


def lerp(a, b, t):
  return (b - a) * t + a


def lerprgb(a, b, t):
  return tuple([lerp(x, y, t) for x, y in zip(a, b)])

def wherebetween(lo, hi, t):
  return (t - lo) / (hi - lo)

def clamp(lo, hi, t):
  if t < lo:
    return lo
  if t > hi:
    return hi
  return t


def blend(src, dst=(0.0, 0.0, 0.0, 1.0)):
  sr, sg, sb, sa = src
  dr, dg, db, da = dst
  sr *= sa
  sg *= sa
  sb *= sa
  dr *= (1 - sa)
  dg *= (1 - sa)
  db *= (1 - sa)
  dr += sr * sa
  dg += sg * sa
  db += sb * sa
  return (
      clamp(0, 255, int(255*dr)),
      clamp(0, 255, int(255*dg)),
      clamp(0, 255, int(255*db))
      )


from glob import glob

class Song:
  def __init__(self, info, notes):
    self.info = info
    self.tracks = tuple(list(Track(self.bpm, tuple(note for note in notes if note.type == i)) for i in range(4)))
    self.viewport_offset = (5, 5)
    self.reset_viewport_offset()
    
  @property
  def bpm(self):
    return self.info['_beatsPerMinute']
    
  @property
  def song_path(self):
    return self.info['_songFilename']

  def reset_viewport_offset(self):
    size = os.get_terminal_size()
    prev_offset = self.viewport_offset
    self.viewport_offset = (
        size.columns // 2 - 3*2,
        size.lines // 2 - 2)
    return prev_offset != self.viewport_offset

  def resync(self, start):
    elapsed_vlc = self.player.get_time()
    while True:
      time.sleep(0.001)
      now = time.time()
      now_vlc = self.player.get_time()
      if elapsed_vlc != now_vlc:
        return start - (now - start) + (now_vlc / 1000)

  def resync_incremental(self, now, start, elapsed_vlc):
    now_vlc = self.player.get_time()
    if now_vlc != elapsed_vlc:
      start += (now - start) - (now_vlc / 1000)
      if self.reset_viewport_offset():
        self.renderer(ansiEscapes.clearScreen)
        self.renderer(ansiEscapes.cursorHide)
    return now, start, now_vlc

  def play(self):
    self.player = vlc.MediaPlayer(self.song_path)
    self.player.play()
    try:
      self.renderer = Renderer()
      self.renderer.flush()
      start = self.resync(time.time())
      elapsed_vlc = self.player.get_time()
      while self.player.is_playing():
        #time.sleep(0.005)
        now, start, elapsed_vlc = self.resync_incremental(time.time(), start, elapsed_vlc)
        elapsed = (now - start)
        notes = self.gather_notes(elapsed)
        self.render_notes(self.renderer, notes, elapsed)
    finally:
      self.player.stop()

  def last_note(self):
    won = None
    for kind, track in enumerate(self.tracks):
      for note in track.notes:
        if not won or note.time > won.time:
          won = note
    return won


  def gather_notes(self, elapsed):
    notes = []
    for kind, track in enumerate(self.tracks):
      notes.extend(track.at(elapsed))
    return notes

  def render_note(self, r, note, elapsed):
    til = note.secs(self.bpm) - elapsed
    before = 0.5 # notes show up this many seconds in advance
    linger = -0.05 # notes hang around this many seconds after dying
    if linger <= til <= before:
      nx = 0+note.w*note.x+self.viewport_offset[0]
      ny = 0+note.h*note.y+self.viewport_offset[1]
      c = note.insignia
      if til < 0:
        t = clamp(0.0, 1.0, wherebetween(0, linger, til))
      else:
        t = clamp(0.0, 1.0, wherebetween(0, before, til))
      # opacity = lerp(1.5, 0.5, 1 - (1 - t)**8) # good
      opacity = lerp(1.5, 0.5, 1 - (1 - t)**10)
      # opacity = 1.0
      rgb = note.color(opacity=opacity)
      bgColor = blend(rgb)
      # fgBlend = (1.0, 1.0, 1.0, opacity)
      fgBlend = (rgb[0] + 0.4, rgb[1] + 0.4, rgb[2] + 0.4, opacity)
      fgColor = blend(fgBlend, rgb)
      #fgColor = blend((1.0,1.0,1.0,opacity), (0.5,0.0,-0.5,1.0))
      bg, bge = styles.bgColor.ansi16m(*bgColor), styles.bgColor.close
      fg, fge = styles.color.ansi16m(*fgColor), styles.color.close
      e = '\033[0m'
      m = ''
      nx += 1
      if til > 0.0:
        # m += '\033[1m'
        m += '\033[3m' # italics
        # if note.upwards:
        #   m += '\033[4m' # underline
        r.to(nx, ny)
        r(f'{fg}{bg}{m}{c}{e}{fge}{bge}')
      elif til > 0:
        # m  += '\033[1m'
        #m += '\033[3m' # italics
        if note.upwards:
          m += '\033[4m'
        r.to(nx, ny)
        r(f'{bg}{fg}{m}{c}{e}{bg}{fge}{bge}')
      else:
        if not note.downwards and not note.upwards:
          m += '\033[9m'
        # m += '\033[3m' # italics
        dx = 0
        dy = 0
        if note.left:
          dx -= 1
        elif note.right:
          dx += 1
        elif note.up:
          dy -= 1
        elif note.down:
          dy += 1
        else:
          if note.rightwards:
            dx += 1
          if note.leftwards:
            dx -= 1
          if note.upwards:
            dy -= 1
            dx *= 2
          if note.downwards:
            dy += 1
            dx *= 2
          if not note.upwards and not note.downwards and not note.leftwards and not note.rightwards:
            dx += 1
        r.to(nx + dx, ny + dy)
        r(f'{m}{fg}{c}{fge}{e}')

  def render_notes(self, r, notes, elapsed):
    r.viewport(16 + self.viewport_offset[0], 8 + self.viewport_offset[1])
    songname = ellipsize(os.path.basename(os.path.dirname(os.path.realpath(self.song_path))))
    r(f'{songname}\n{int(elapsed/60.0):02d}m{int(elapsed%60.0):02d}s')
    for note in sorted(notes, key=lambda x: x.time):
      self.render_note(r, note, elapsed)
    r.flush()


def ellipsize(s, n=30):
  if len(s) > n:
    return s[:(n-3)] + '...'
  return s

 
import time

class Track:
  def __init__(self, bpm, notes):
    self.bpm = bpm
    self.notes = notes

  # def before(self, now):
  #   return tuple(n for n in self.notes if n.secs(self.bpm) <= now)

  def at(self, now):
    r = []
    for i, note in enumerate(self.notes):
      t = note.secs(self.bpm)
      if abs(now - t) < 2.0:
        r.append(note)
    return r



def run():
    if args.verbose:
        print(args)
    files = [os.path.realpath(file) for file in args.args]
    while True:
      # shuffle?
      for file in files:
        try:
          os.chdir(os.path.dirname(file) or '.')
          file = os.path.basename(file)
          bs = BeatsaberMap()
          song = bs.parse(file)
          song.play()
        except KeyboardInterrupt:
          time.sleep(1.0)

def main():
    try:
        global args
        if not args:
            args, leftovers = parser.parse_known_args()
            args.args = leftovers
        return run()
    except IOError:
        # http://stackoverflow.com/questions/15793886/how-to-avoid-a-broken-pipe-error-when-printing-a-large-amount-of-formatted-data
        try:
            sys.stdout.close()
        except IOError:
            pass
        try:
            sys.stderr.close()
        except IOError:
            pass

if __name__ == "__main__":
    main()