import sequtils, math, strutils, os, random
import nimPNG

type
  Algorithm* = enum
    none, thr, od, fs, jjn, atk, app, bn, bns, bnp, bnp2, by1, by2, by3, by4, ad1, ad2, ad3, ad4

proc dither(png: auto, algorithm: Algorithm, colours = 2) =
  if colours < 2: return
  let step = 255.0 / colours.float
  var steps: seq[float]
  for i in 1..<colours:
    steps.add (if steps.len == 0: 0.0 else: steps[^1]) + step
  template quantize(x: float): float =
    block:
      var s = 0
      for i, step in steps:
        if x > step: s = i + 1
        else: s = i; break
      if s == 0: 0.0
      elif s == steps.len: 255.0
      else: step * (s.float + 0.5)
  case algorithm:
  of none: return
  of od..atk:
    var errors = newSeqWith(png.height, newSeq[float](png.width))

    for y in 0..<png.height:
      for x in 0..<png.width:
        let colour = png.data[y*png.width + x].float + errors[y][x]
        var newColour = quantize(colour)
        #  if colour > 127:
        #    255.0
        #  else:
        #    0.0
        let error = colour - newColour
        case algorithm:
        of od:
          # Simple 1D error dissipation
          if x != png.width - 1:
            errors[y][x+1] = error
        of fs:
          # Floyd-Steinberg
          let errorStep = error / 16
          if x != png.width - 1:
            errors[y][x+1] += 7 * errorStep
            if y != png.height - 1:
              errors[y+1][x+1] += errorStep
          if y != png.height - 1:
            errors[y+1][x] += 5 * errorStep
            if x != 0:
              errors[y+1][x-1] += 3 * errorStep
        of jjn:
          # Jarvis, Judice, and Ninke
          let errorStep = error / 48
          if x != png.width - 1:
            errors[y][x+1] += 7 * errorStep
            if y != png.height - 1:
              errors[y+1][x+1] += 5 * errorStep
            if y < png.height - 2:
              errors[y+2][x+1] += 3 * errorStep
          if x < png.width - 2:
            errors[y][x+2] += 5 * errorStep
            if y != png.height - 1:
              errors[y+1][x+2] += 3 * errorStep
            if y < png.height - 2:
              errors[y+2][x+2] += errorStep
          if y != png.height - 1:
            errors[y+1][x] += 7 * errorStep
            if x != 0:
              errors[y+1][x-1] += 5 * errorStep
            if x > 1:
              errors[y+1][x-2] += 3 * errorStep
          if y < png.height - 2:
            errors[y+2][x] += 5 * errorStep
            if x != 0:
              errors[y+2][x-1] += 3 * errorStep
            if x > 1:
              errors[y+2][x-2] += 1 * errorStep
        of atk:
          # Atkinson
          let errorStep = error / 8
          if x != png.width - 1:
            errors[y][x+1] += errorStep
            if y != png.height - 1:
              errors[y+1][x+1] += errorStep
          if x < png.width - 2:
            errors[y][x+2] += errorStep
          if y != png.height - 1:
            errors[y+1][x] += errorStep
            if x != 0:
              errors[y+1][x-1] += errorStep
          if y < png.height - 2:
            errors[y+2][x] += errorStep
        else: discard
        png.data[y*png.width + x] = newColour.uint8
  of app:
    # Algorithm proposed by Liam Appelbe: https://liamappelbe.medium.com/dizzy-dithering-2ae76dbceba1
    # Doesn't seem to give as good results as shown in the article...
    var
      errors = newSeq[float](png.width * png.height)
    #  order = newSeq[int](png.width * png.height)
    #for i in 0..order.high:
    #  order[i] = i
    #for i in 0..20:
    #  order.shuffle()

    let
      m = nextPowerOfTwo(png.data.len)
      rand_odds = [rand(m) or 1, rand(m) or 1, rand(m) or 1, rand(m) or 1, rand(m) or 1]
      rand_anys = [rand(m), rand(m), rand(m), rand(m), rand(m)]
    for p in 0..<m:
      var pixel = p
      for i in 0..4:
        pixel = ((pixel * rand_odds[i]) mod m) xor rand_anys[i]
      if pixel > png.data.high: continue
    #for pixel in order:
      let colour = png.data[pixel].float + errors[pixel]
      var newColour = quantize(colour)
      #  if colour > 127:
      #    255.0
      #  else:
      #    0.0
      png.data[pixel] = newColour.uint8
      let error = colour - newColour
      template loop(body: untyped) =
        for y in -1..1:
          for x in -1..1:
            if x == 0 and y == 0: continue
            let idx {.inject.} = pixel + png.width * y + x
            if idx notin png.data.low..png.data.high: continue
            if png.data[idx] notin {0'u8, 255'u8}:
              let weight {.inject.} =
                if y == 0 or x == 0: 1.0 # Orthogonal
                else: 0.1 # Diagonal
              body

      var denom = 0.0
      loop:
        denom += weight
      if denom > 0:
        loop:
          errors[idx] += error * weight / denom + 0.5
  of bn..by4:
    var patres = loadPng(seq[uint8], "pattern_" & $algorithm & ".png", LCT_GREY, 8)
    if not patres.isOk:
      stderr.writeLine "Unable to load pattern_" & $algorithm & ".png"
      quit 1
    let pattern = patres.value

    for y in 0..<png.height:
      for x in 0..<png.width:
        var colour = png.data[y*png.width + x]
        var mask = pattern.data[(y mod pattern.height)*pattern.width + (x mod pattern.width)]
        let newColour = quantize(colour.float + (254.5 / (colours - 1).float) * ((mask.float / 255.0) - 0.5)).uint8
        #  if colour == 0 or colour < mask:
        #    0'u8
        #  else:
        #    255'u8
        png.data[y*png.width + x] = newColour
  else:
    for y in 0..<png.height:
      for x in 0..<png.width:
        var colour = png.data[y*png.width + x]
        let newColour =
          case algorithm:
          of thr:
            quantize(colour.float).uint8
            #if colour < 127: 0'u8
            #else: 255'u8
          of ad1..ad4:
            let
              c = 0
              mask =
                case algorithm:
                of ad1:
                  (((x xor y * 149) * 1234) and 511).float / 511.0
                of ad2:
                  ((((x + c * 17) xor y * 149) * 1234) and 511).float / 511.0
                of ad3:
                  (((x + y * 237) * 119) and 255).float / 255.0
                of ad4:
                  ((((x + c * 67) + y * 236) * 119) and 255).float / 255.0
                else: 0
            #(floor(colour.float / 255 + mask) * 255.0).uint8
            quantize(colour.float + (254.5 / (colours - 1).float) * (mask.float - 0.5)).uint8
          else: 0
        png.data[y*png.width + x] = newColour

when isMainModule:
  let
    file = paramStr(1)
    algorithm = paramStr(2).parseEnum[:Algorithm]
  var res = loadPng(seq[uint8], file, LCT_GREY, 8)
  if not res.isOk:
    stderr.writeLine "Unable to load image"
    quit 1
  var png = res.value
  echo png.width
  echo png.height

  if paramCount() > 2:
    dither(png, algorithm, paramStr(3).parseInt)
  else:
    dither(png, algorithm)

  let split = splitFile(file)
  discard savePng(split.dir / "dithered" / split.name & "_" & $algorithm & split.ext, png.data, LCT_GREY, 8, png.width, png.height)