{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE TupleSections #-}

{-MCCOMMENT

simple glich maker
use: `stack run image`
Accept multiple input

example: stack run img/klee.png
output to: stack run img/klee.png.out.png

value of `seed` can be changed for different output on same image

image seed is calculated with image data

Haha, this is taking other image and do stuff with it, not really creating new ones

I found some random images to test

There is 3 effect, move, reverse, drag
move: move some pixels (some of r/g/b) to other place
reverse: invert color
drag: drag some pixels to the left side, make all pixels on the road the same

line is a special version of move, move a horizontal line horizontally

* some file format might not be supported, like 32bit png and HDR image
-}

module Glitch where

import Codec.Picture (Pixel (readPixel, writePixel), imageData, imageWidth, pixelMap, readImage, savePngImage)
import Codec.Picture.Types (ColorConvertible (promoteImage), ColorSpaceConvertible (convertImage), DynamicImage (ImageCMYK16, ImageCMYK8, ImageRGB16, ImageRGB8, ImageRGBA16, ImageRGBA8, ImageRGBF, ImageY16, ImageY8, ImageYA16, ImageYA8, ImageYCbCr8, ImageYF), Image, MutableImage (mutableImageHeight, mutableImageWidth), Pixel16 (..), Pixel8 (..), PixelCMYK16, PixelCMYK8, PixelRGB16 (..), PixelRGB8 (..), PixelRGBA16 (..), PixelRGBA8 (..), PixelYA16 (..), PixelYA8 (..), PixelYCbCr8, pixelMap, unsafeFreezeImage, unsafeThawImage)
import Control.Monad.ST (RealWorld)
import Data.List (nub, permutations, subsequences)
import System.Environment (getArgs)
import System.Random (StdGen, mkStdGen)
import System.Random.Stateful (AtomicGenM, newAtomicGenM, uniformM, uniformRM)

seed :: Int
seed = 1145141919810

type I = MutableImage RealWorld PixelRGBA16

getCoordList :: (Int, Int) -> (Int, Int) -> [(Int, Int)]
getCoordList (_, _) (_, 0) = []
getCoordList (fromX, fromY) (x, y) = map (,fromY) [fromX .. fromX + x - 1] ++ getCoordList (fromX, fromY + 1) (x, y -1)

readPixels :: I -> [(Int, Int)] -> IO [(PixelRGBA16, Int, Int)]
readPixels _ [] = return []
readPixels img ((x, y) : cs) = do
  p <- readPixel img x y
  rest <- readPixels img cs
  return ((p, x, y) : rest)

writePixels :: I -> Bool -> (Bool, Bool, Bool) -> [(PixelRGBA16, Int, Int)] -> IO ()
writePixels _ _ _ [] = return ()
writePixels img add (moveR, moveG, moveB) ((p, x, y) : cs) = do
  let PixelRGBA16 newR newG newB _ = p
  PixelRGBA16 currentR currentG currentB t <- readPixel img x y
  let r = if moveR then if add then newR + currentR else newR else currentR
  let g = if moveG then if add then newG + currentG else newG else currentG
  let b = if moveB then if add then newB + currentB else newB else currentB
  writePixel img x y (PixelRGBA16 r g b t)
  writePixels img add (moveR, moveG, moveB) cs

getPart :: (Int, Int) -> Int -> (Int, Int)
getPart (x, y) n = (x `div` n, y `div` n)

movePart :: I -> AtomicGenM StdGen -> IO ()
movePart img gen = do
  putStrLn "Move Pixel"
  let width = mutableImageWidth img
  let height = mutableImageHeight img
  let (partX, partY) = getPart (width, height) 7

  -- how big
  sizeX <- uniformRM (0, partX * 2) gen
  sizeY <- uniformRM (0, partY * 2) gen
  -- from where
  fromX <- uniformRM (0, width - sizeX) gen
  fromY <- uniformRM (0, height - sizeY) gen
  -- to where
  diffX <- uniformRM (- partX, partX) gen
  diffY <- uniformRM (- partY, partY) gen
  -- move what
  (moveR, mG, mB) <- uniformM gen
  -- list
  let coordList = getCoordList (fromX, fromY) (sizeX, sizeY)
  pixels <- readPixels img coordList
  let pixelsFit = map (\(p, x, y) -> (p, (x + diffX + width) `mod` width, (y + diffY + height) `mod` height)) pixels
  -- write
  (add1, add2) <- uniformM gen
  writePixels img (add1 && add2) (moveR, mG, mB) pixelsFit

line :: I -> AtomicGenM StdGen -> IO ()
line img gen = do
  putStrLn "Line"
  let width = mutableImageWidth img
  let height = mutableImageHeight img
  let (partX, partY) = getPart (width, height) 16
  -- line width
  sizeY <- uniformRM (0, partY) gen
  -- from where
  fromY <- uniformRM (0, height - sizeY) gen
  -- move length
  diffX <- uniformRM (- partX * 4, partX * 4) gen
  -- move what
  (moveR, moveG, moveB) <- uniformM gen
  -- list
  let coordList = getCoordList (0, fromY) (width, sizeY)
  pixels <- readPixels img coordList
  let pixelsFit = map (\(p, x, y) -> (p, (x + diffX + width) `mod` width, y)) pixels
  -- write
  add <- uniformM gen
  writePixels img add (moveR, moveG, moveB) pixelsFit

reversePixel :: (Bool, Bool, Bool) -> PixelRGBA16 -> PixelRGBA16
reversePixel (reverseR, reverseG, reverseB) (PixelRGBA16 r g b a) = PixelRGBA16 (if reverseR then - r else r) (if reverseG then - g else g) (if reverseB then - b else b) a

reverseI :: I -> AtomicGenM StdGen -> IO ()
reverseI img gen = do
  putStrLn "Reverse"
  let width = mutableImageWidth img
  let height = mutableImageHeight img
  let (partX, partY) = getPart (width, height) 8
  -- how big
  sizeX <- uniformRM (0, partX) gen
  sizeY <- uniformRM (0, partY) gen
  -- where
  fromX <- uniformRM (0, width - sizeX) gen
  fromY <- uniformRM (0, height - sizeY) gen
  -- reverse what
  (moveR, moveG, moveB) <- uniformM gen
  -- list
  let coordList = getCoordList (fromX, fromY) (sizeX, sizeY)
  pixels <- readPixels img coordList
  let pixelsFit = map (\(p, x, y) -> (reversePixel (moveR, moveG, moveB) p, x, y)) pixels
  -- write
  writePixels img False (moveR, moveG, moveB) pixelsFit

genSimonN :: Int -> Int -> Int -> Int -> Int -> AtomicGenM StdGen -> Int -> IO [Int]
genSimonN 0 _ _ _ _ _ _ = return []
genSimonN n step min max result gen p
  | result < min = genSimonN n (abs step) max min (result + abs step) gen p
  | result > max = genSimonN n (- (abs step)) max min (result - abs step) gen p
  | otherwise = do
    rest <- genSimonN (n - 1) step min max (result + step) gen p
    runAway <- uniformRM (0, p) gen
    away <- uniformRM (min, max) gen
    return ((if runAway == (0 :: Int) then away else result) : rest)

drag :: Int -> I -> AtomicGenM StdGen -> IO ()
drag i img gen = do
  putStrLn "Drag"
  let width = mutableImageWidth img
  let height = mutableImageHeight img
  let (partX, partY) = getPart (width, height) 13
  -- how big
  xDiff <- uniformRM (partX * min i 3, partX * i) gen
  xSize <- uniformRM (0, partX * 2) gen
  ySize <- uniformRM (0, partY) gen
  -- where
  yStart <- uniformRM (0, height - ySize) gen
  xStart <- uniformRM (0, width - xSize) gen
  -- list
  let ys = [yStart .. yStart + ySize - 1]
  step <- uniformM gen
  xs <- genSimonN ySize (if step then 1 else -1) xStart (xStart + xSize) (xStart + xSize - 5) gen 4
  let xys = zip xs ys
  pxys <- readPixels img xys
  let xTargets = map ((\(x, maxX) -> [x .. maxX]) . (\x -> (x, min (width -1) (x + xDiff)))) xs
  let xpTargets = zip xTargets pxys
  let targets = concatMap (\(xs, (p, _, y)) -> map (p,,y) xs) xpTargets
  writePixels img False (True, True, True) targets

shuffle :: AtomicGenM StdGen -> [a] -> IO [a]
shuffle _ [] = return []
shuffle gen xs = do
  n <- uniformRM (0, length xs -1) gen
  let (xs1, x : xs2) = splitAt n xs
  rest <- shuffle gen (xs1 ++ xs2)
  return (x : rest)

d :: AtomicGenM StdGen -> [AtomicGenM StdGen -> IO ()] -> IO ()
d _ [] = return ()
d gen (f : fs) = do
  f gen
  d gen fs

pixelSum :: I -> Int -> Int -> IO Int
pixelSum _ (-1) _ = return 0
pixelSum img x (-1) = do
  let height = mutableImageHeight img
  pixelSum img (x - 1) (height - 1)
pixelSum img x y = do
  (PixelRGBA16 r g b a) <- readPixel img x y
  rest <- pixelSum img x (y - 1)
  return (sum (rest : map fromIntegral [r, g, b, a]))

doStuff :: I -> IO ()
doStuff img = do
  let width = mutableImageWidth img
  let height = mutableImageHeight img
  putStrLn ("x y: " ++ show width ++ " " ++ show height)
  imageSum <- pixelSum img (width - 1) (height - 1)
  let iSeed = imageSum * seed
  putStrLn ("Image Seed: " ++ show iSeed)
  atomicGen <- newAtomicGenM (mkStdGen iSeed)
  putStrLn "Glitch"
  -- move pixels
  moveN <- uniformRM (8, 24) atomicGen
  let moves = replicate moveN (movePart img)
  -- line
  lineN <- uniformRM (2, 9) atomicGen
  let lines = replicate lineN (line img)
  -- reverse
  reverseN <- uniformRM (12, 26) atomicGen
  let reverses = replicate reverseN (reverseI img)
  -- drag
  dragN <- uniformRM (1, 3) atomicGen
  let drags = replicate dragN (drag 6 img)
  smallDragN <- uniformRM (4, 8) atomicGen
  let smallDrags = replicate smallDragN (drag 1 img)
  -- execute
  let pending = moves ++ lines ++ reverses ++ drags ++ smallDrags
  order <- shuffle atomicGen pending
  d atomicGen order

{-WETT-}

m :: [String] -> IO ()
m [] = return ()
m (path : ps) = do
  putStrLn ("Read " ++ path)
  img <- readImageRGBA16 path
  img <- unsafeThawImage img
  doStuff img
  let outPath = path ++ ".out.png"
  putStrLn ("Save " ++ outPath)
  img <- unsafeFreezeImage img
  savePngImage outPath (ImageRGBA16 img)
  m ps

main :: IO ()
main = do
  putStrLn "Hi, it is Glich"
  putStrLn "input some image file, output .out.png file"
  putStrLn ("Seed: " ++ show seed)
  paths <- getArgs
  putStrLn ("Pending: " ++ show (length paths))
  m paths

{-TTEW-}

-- Modified Code from https://github.com/fumieval/JuicyPixels-util/blob/master/Codec/Picture/RGBA8.hs
-- Added some codec support
-- now I know what is not supported
-- Maybe I can do pull request when I have time
-- And i can use RGB16

fromDynamicImage :: DynamicImage -> Image PixelRGBA16
fromDynamicImage (ImageY8 img) = promoteImage ((promoteImage :: Image Pixel8 -> Image PixelRGBA8) img)
fromDynamicImage (ImageYA8 img) = promoteImage ((promoteImage :: Image PixelYA8 -> Image PixelRGBA8) img)
fromDynamicImage (ImageRGB8 img) = promoteImage img
fromDynamicImage (ImageRGBA8 img) = promoteImage img
fromDynamicImage (ImageY16 img) = promoteImage img
fromDynamicImage (ImageYF _) = error "Unsupported format: ImageYF"
fromDynamicImage (ImageYA16 img) = promoteImage img
fromDynamicImage (ImageRGB16 img) = promoteImage img
fromDynamicImage (ImageRGBF _) = error "Unsupported format: ImageRGBF"
fromDynamicImage (ImageRGBA16 img) = img
fromDynamicImage (ImageYCbCr8 img) = promoteImage ((convertImage :: Image PixelYCbCr8 -> Image PixelRGB8) img)
fromDynamicImage (ImageCMYK8 img) = promoteImage ((convertImage :: Image PixelCMYK8 -> Image PixelRGB8) img)
fromDynamicImage (ImageCMYK16 img) = promoteImage ((convertImage :: Image PixelCMYK16 -> Image PixelRGB16) img)

readImageRGBA16 :: FilePath -> IO (Image PixelRGBA16)
readImageRGBA16 path = readImage path >>= either fail (return . fromDynamicImage)