compute-constrained-bayes/src/jit_bayes.nim

import print
import strutils
import sequtils
import std/sugar

## Get sequences
let file_path = "../data/stripped"
proc getOEIS(): seq[seq[string]] =
  let f = open(file_path)
  var i = 0
  var line : string
  var seqs: seq[seq[string]]
  while f.read_line(line):
    if i > 3:
      let seq = split(line, ",")
      let l = seq.len
      let nums = seq[1..(l-2)] 
      seqs.add(nums)
    i = i + 1
  f.close()
  return seqs
var seqs = getOEIS()

## Sequence helpers
proc startsWithSubsequence(xs: seq[string], ys: seq[string]): bool = 
  if xs.len == 0:
    return true
  elif ys.len == 0:
    return false
  elif xs[0] == ys[0]:
    return startsWithSubsequence(xs[1..<xs.len], ys[1..<ys.len])
  else:
    return false

proc getSequencesWithStart(seqs: seq[seq[string]], start: seq[string]): seq[seq[string]] = 
  var continuations: seq[seq[string]]
  for seq in seqs:
    if startsWithSubsequence(start, seq):
      continuations.add(seq)
  return continuations

## Pretty print sequences

# var start = @["1", "2", "3", "4", "5"]
# var continuations = getSequencesWithStart(seqs, start)
# print continuations

## Do simple predictions
proc predictContinuation(seqs: seq[seq[string]], start: seq[string]): seq[string] = 
  let continuations = getSequencesWithStart(seqs, start)
  let l = start.len 
  let next_items = continuations.map(c => c[l])
  return next_items

var start = @["1", "2", "3", "4", "5"]
let sample_continuations = predictContinuation(seqs, start)
print(sample_continuations)