module Main where

import Control.Arrow
import qualified Control.Monad as M
import qualified Control.Monad.Parallel as MP
import Control.Parallel.Strategies
import qualified Data.Colour as C
import qualified Data.Colour.Names as CN
import Data.Colour.RGBSpace (uncurryRGB)
import qualified Data.Colour.SRGB as CL
import qualified Data.Colour.SRGB.Linear as CL
import Data.List
import qualified Data.Map as M
import qualified Data.Maybe as My
import qualified Data.Set as S
import qualified Data.Vector.Unboxed as Vec
import Debug.Trace
import qualified Debug.Trace as DT
import qualified Debug.Trace as T
import qualified Diagrams as DP
import Diagrams.Backend.SVG
import Diagrams.Backend.SVG.CmdLine
import Diagrams.Prelude
import qualified Diagrams.Prelude as D
import GHC.Generics
import Graphics.Image as Img hiding (map, zipWith)
import qualified Graphics.Image.ColorSpace as G
import qualified Graphics.Image.Interface as Int
import Options.Generic
import qualified System.Environment as Env
import System.Random
import System.Random.Internal
import System.Random.SplitMix
import Triangles (getTriangleAverageRGB)
import qualified Triangles as Tri

data Options = Options
  { numPoints :: Int
  , gen :: Maybe StdGen
  }

-- -- modify this to your liking
defaultOpts =
  Options
    { numPoints = 10
    , gen = Nothing
    }

-- CL.rgb might be the wrong fn...
tosRGB' :: (Ord b, Floating b) => Pixel G.RGB b -> CL.Colour b
tosRGB' (G.PixelRGB r g b) = CL.rgb r g b

convImage = Vec.map tosRGB' . Int.toVector

corners :: [(Double, Double)]
corners = (,) <$> [0, 1] <*> [0, 1]

shapeCircumference = Data.List.sum . map D.norm . loopOffsets . fromVertices

genImage' :: Image VU G.RGB Double -> StdGen -> Int -> QDiagram SVG V2 Double Any
genImage' image gen cornerCount =
  scaleY widthHeightRatio
    . reflectY
    . rectEnvelope (mkP2 0 0) (1 ^& 1)
    -- . atop visualizeGraph
    . mconcat
    . map drawVoroniRegion
    . sortOn shapeCircumference
    . withStrategy (parListChunk 50 rdeepseq)
    . map (uncurry Tri.voroniDiagramCorners)
    $ voroni
 where
  drawVoroniRegion shape =
    lw 0
      . fillColor (Tri.voroniRegionAverageColor img' dimensions shape)
      . strokeLocLoop
      . fromVertices
      $ shape

  widthHeightRatio :: Double
  widthHeightRatio = (fromIntegral . snd $ dimensions) / (fromIntegral . fst $ dimensions)

  img' = convImage image
  dimensions = uncurry (flip (,)) . Img.dims $ image

  singleVoroni = last voroni

  visualizeGraph :: QDiagram SVG V2 Double Any
  visualizeGraph =
    lc red
      . lw 1
      . position
      . map (\(x0, x1) -> (x0,) . strokeLine $ (x0 ~~ x1))
      . Tri.toPlanarGraph
      $ corners'

  voroni =
    Tri.findVoroniDiagram
      . Tri.toPlanarGraph
      $ corners'

  corners' = take cornerCount . map p2 $ corners ++ Tri.randomPoints gen

deriving instance Generic (CL.RGB a)
deriving instance NFData a => NFData (CL.RGB a)

toDimensionVector :: (Int.BaseArray arr cs e, Fractional n) => Image arr cs e -> SizeSpec V2 n
toDimensionVector image = Diagrams.Prelude.dims $ p2 (fromIntegral $ cols image, fromIntegral $ rows image) .-. p2 (0.0, 0.0)

data CLIOptions = CLIOptions
  { input :: FilePath
  , output :: FilePath
  , cornerCount :: Int
  }
  deriving (Generic)

instance ParseRecord CLIOptions

main :: IO ()
main = do
  CLIOptions{..} <- getRecord "image options"
  let (Options{gen = gen}) = defaultOpts
  gen' <- getStdGen -- pure . mkStdGen $ 2344
  let gens :: [StdGen] = map fst . iterate (split . snd) . split $ gen'
  print gen'
  image <- Img.readImageRGB VU input
  let nums = zip gens $ map show [0 .. 60]
  let dimVector = toDimensionVector image
  renderSVG output dimVector (genImage' image gen' cornerCount)
