File-To-Video-Encoder/cl.nim

import strutils
import sugar
import sequtils
import math
import arraymancer
import random
import itertools 
import simplepng
import nimPNG
import system
import times
import tables
import streams
import os
import json
import osproc
import tensorCeral
import strformat

proc itPad(pad, width: int, start = 0) : iterator() : seq[(int, int)] {.gcsafe.} =
  ## Creates a coordinate continum corresponding to the pad and resolution
  ## ie 0,0 0,1 0,2 1,0 1,1 1,2 2,0 2,1 2,2
  return iterator(): seq[(int,int)] =
    for x in countup(pad, width, pad+1):
        for y in countup(pad, width, pad+1):
            var temp : seq[(int,int)]
            for a in y-(pad) .. y:
                for b in x-(pad) .. x:
                    temp.add((a,b))
            yield temp

proc genColorMatrix*() : Tensor[byte] = 
  ##Generates a random color pallette 
  randomize(cpuTime().int64)
  return  collect(for x in 1 .. 256*3: (rand(0 .. 255)).byte).toTensor().reshape(256,3)

proc writeRandomImage(width, height : int, path : string) = 
  ##Writes a nonsense image
  var colorTensor = genColorMatrix()
  var p = initPixels(width, height)
  let chunky = collect(for x in chunked(collect(for x in colorTensor: x), 3): x)
  var n = 0
  for color in p.mitems:
    let
      r = chunky[n][0]
      g = chunky[n][1]
      b = chunky[n][2]
    color.setColor(r, g, b, 255)
    n+=1
  simplePNG(path, p)


proc fromPng*(input : string, pad = 0, total : int, multiple = true) : Tensor[byte]=
  ## Input = folder which it is in. eg ./out/0/
  ## uses input/1.png as a benchmark for resolution; it is intended to be used within the confines of the rest of the code
  ## total = the total amount files you want to scan
  ## if you want to use just one, multiple = false

  ## Returns a tensor of data, organized by blocks. First block, second block, third block, etc
  ## 
  var png : PNGResult[string]
  var size : int
  if multiple:
    png = loadPNG32(&"{input}1.png")
    size = png.width
  else:
    png = loadPNG32(input)
    size = png.width  

  var filesRead = 0
  var final : seq[byte]
  proc decode(x : seq[char]) : seq[byte] =
    #takes the data, to hex, then converts it to its, then cuts the alphachannel
    return x.map(x=>($x).toHex()).map(x=>byte(parseHexInt(x)))[0 .. 2]

  proc flatten(a : seq[seq[byte]]) : seq[byte] =
    for x in a:
      result.add(x)
             #out is in a seq[int] of rgba, so, this outs an array of them
                                              #decode, flatten, then tensor to get the data
                                                  #back to the "original" form 
  if pad == 0:
    return collect(for x in chunked(png.data,4): decode(x)).flatten().toTensor().reshape(size^2, 3)
  else:
    #it iterates to the end amount of files given in total
    #if false it'll only read 1.1
    proc decodeImage(png : PNGResult[string], final : var seq[byte]) =
      doAssert(png.height == size, "NON SQUARE SHAPE PADS NOT SUPPORTED")
      #IO operation to organize 
                          #one lines RGBA*vertical pad length, to get each b lock
      for bigchunk in chunked(png.data, (size*4)*(pad+1)):
                      #organize this into horizontal length blocks
                  #used as a vertical iterator
        let start = bigchunk.distribute(pad+1, spread=false)
        for x in countup(0, start[0].high, ((pad+1)*4)):
          for y in 0 .. start.high:
            #this seems inefficient, i should look at this again..
            #later.
            var a = start[y][x .. (((pad+1)*4)+x)-1]
            for delete in countup(3, a.high, 4): a.delete(delete-(int((delete+1)/4)-1))
            final.add(a.map(x=>x.byte))
    if multiple:
      for file in 1 .. total:
        let png = loadPNG32(input & $file & ".png")
        decodeImage(png, final)
        filesRead+=1
    else:
      decodeImage(png, final)
      filesRead = 1


  let mody = (((pad+1)^2)*3)
  let test = int(((size^2)*3)/mody)*(filesRead)
  return final.toTensor().reshape(test,mody)

iterator writeMultipleFiles(file : Stream, size, filesize, pad : int, outdir : string, reference : ref seq[byte]) : (proc(a: (byte,byte,byte)){.gcsafe.}, byte) {.gcsafe.} = 
  let byteSize = (pad+1)^2
  let imagesNumber = (((filesize+256) * byteSize) / size^2).ceil().int
  var written = 0
  #threading this would be so simple, thats kinda the point of the design.
  #previous form was a nightmare
  while true:
    var cbyte : byte
    var starty = 0
    for imageIndex in 1 .. imagesNumber:
      var image = initPixels(size,size)
      var it = itPad(pad, size)

      for current in it():
        if file.atEnd():
          break

        if starty < 256:
          cbyte = byte(starty)
          starty.inc()
        else: 
          cbyte = file.readChar.byte
        yield ((proc(input : (byte,byte,byte)) {.gcsafe.} =
          for x in current:
            image[x[0],x[1]].setColor(input[0],input[1],input[2], 255)
            reference[].add(input[0])
            reference[].add(input[1])
            reference[].add(input[2])
          ), 
      cbyte)
      simplePNG(outdir & "/in" & $written & ".png", image)
      written.inc()

    break
proc rowToArray(a : Tensor[byte]) : (byte,byte,byte) {.gcsafe.} =
    (a[0, 0], a[0,1], a[0, 2])

proc encodeImage*(pad = 0, path : int, inputFile : string | Stream, pltePath : string | Tensor[byte]) : (Tensor[byte], Tensor[byte]) {.gcsafe.} = 
  var refout = new seq[byte]
  ##not for release

  #let rand =  genColorMatrix()
  #temp bebug matrix
  var colorPallet : Tensor[byte]
  when pltePath is string:
    if pltePath == "":
      colorPallet = genColorMatrix()
    else:
      colorPallet = deSerializeColorPallate(newFileStream(pltePath))
  else:
    
    colorPallet = pltePath

  let size = 600

  var file : Stream
  var fileSize : BiggestInt
  when inputFile is string:
    if inputFile == "" or inputFile == "-":
      while sizeof(stdin) == 0:
        continue
      file = newStringStream(newFileStream(stdin).readAll())
      #dirty code, I don't know else how to get the proper size of stdin
      fileSize = file.readAll().len()
      file.setPosition(0)
    else:
      file = (newFileStream(inputFile))
      fileSize =  getFileSize(inputFile)
  else:
    file = inputFile

  for buff in writeMultipleFiles(file, size, fileSize.int, pad, "in/" & $path, refout):
    #wMF yeilds a function to apply a color to a block, and a byte of data
    #this system allows for a greator flexabiliy of color choices and 
    #potential future complexity if needed
    let row = rowToArray(colorPallet[buff[1].int, 0 .. 2])
    buff[0](row)


  let mody = (((pad+1)^2)*3)  
  return (colorPallet, refout[].toTensor().reshape(int(refout[].len()/mody) , mody))

# proc decodeImagePerfect(pad : int, simple = true, path : string) {.gcsafe.} =
  #to be programmed 

proc outMatrix(pad : int, simple = true, path : int, multiple = true) : Tensor[byte]{.gcsafe.} =
  ##wrapper for fromPng
  let files = collect(for x in os.walkdir("out/" & $path) : x).len()
  var input = fromPng("out/" & $path & "/out", pad, files, multiple)
  return input


type 
  train* = ref object of RootObj
    pre* : (seq[uint8], Metadata)
    post* : (seq[uint8], Metadata) 
    colors* : (seq[uint8], Metadata)


proc convertFrames*(input : (int, string, string | Stream, bool, string)) {.gcsafe, thread.}=
  ## the main CL convert proc, everything put together
  ## int is the outdir in 'in/x' and 'out/x'
  ## string[0] is the name of the outdir, doesn't matter if the bool is falsy
  ## string[1] is a stream containing encoding data or a string to it
  ## bool is writey, false if you don't wise to write the output in the cerial format
  ## string[2] is the path of a plte file, if you want it, keep blank otherwise
  let
    doWrite = input[3]
    inputFile = input[2]
    path = input[0]
    pltePath = input[4]

  let t1 = cpuTime()
  let file = "in/" & $input[0]
  if fileExists(file):
    removeDir(file)
  discard existsOrCreateDir(file)
  let one = encodeImage(9, path, inputFile, pltepath)

  discard os.execShellCmd("./tovideo.sh " & $input[0])

  let outy  = outMatrix(9, true, 0)

  if doWrite:
    var outsize = [
      one[0].shape, one[1].shape, outy.shape
    ]

    var writtey = [toFlatSeq(one[0]), toFlatSeq(one[1]), toFlatSeq(outy)]
    var outputFile : FileStream
  
    stderr.writeline input

    if input[1] == "":
      outputFile = newFileStream(stdout)
    else:
      stderr.writeline input
      stderr.writeline "?"
      outputFile = newFileStream(input[1], fmWrite) 

    serializeTensor(outputFile, outsize, writtey)

var params {.global.} = commandLineParams()

proc printHelp() =
  stderr.writeline "This program encodes data!"
  stderr.writeline "use -h to display this!"
  stderr.writeline "-e to encode a file:"
  stderr.writeline "  cl -e [filetoencode] [tensorOutPath.bin] [pallet.plte]"
  stderr.writeline "  to use in pipe mode: cl -e - -"
  stderr.writeline "  to note release tensors: cl [input] x"
  stderr.writeline "-rf to encode a file forever, generating new pallets"
  stderr.writeline "  -cl -rf [file]"
  quit(1)
when isMainModule:  
  var channel = createShared(Channel[train], sizeof(Channel[train]))
  channel[].open()
  var files = (collect(for x in os.walkdir("./trainingdata"): x).len())
  params.setLen(4)
  
  if params[0] == "-e":
    
    let fileToEncode = params[1]
    let tensorOutPath = params[2]
    let pltePath = params[3]
    let doWrite = tensorOutPath != "x"
    if params[1] == "" and params[2] == "":
      printHelp()
    convertFrames((0, tensorOutPath, fileToEncode, doWrite, pltePath))
  elif params[0] == "-rf":
    let fileToEncode = params[1]
    if fileToEncode == "":
      printHelp()
    var cores = countProcessors()
    var trainThreads = newSeq[Thread[(int, string, string, bool, string)]](cores)
    while true:
      for x in 0 .. cores:
        createThread(trainThreads[x], convertFrames, (x, &"trainingdata/{files}.bin", fileToEncode, true, ""))
        files+=1
      joinThreads(trainThreads)
  elif params[0] == 
  "-h":
    printHelp()