module Main where

import Control.Arrow
import qualified Control.Monad as M
import qualified Control.Monad.Parallel as MP
import Control.Monad.Zip (MonadZip (mzipWith))
import Control.Parallel.Strategies
import qualified Data.Bifunctor as Bi
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.List as L
import qualified Data.Map as M
import qualified Data.Massiv.Array as M
import qualified Data.Massiv.Array as Ma
import qualified Data.Massiv.Array.IO 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 D
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 qualified Graphics.Color.Space as Co
import qualified MinDistanceSample as MDS
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

toColour :: Fractional a => Co.Color (Co.SRGB Co.Linear) a -> Colour a
toColour (Co.ColorSRGB r g b) = CL.rgb r g b

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

shapeCircumference :: [Point V2 Double] -> Double
shapeCircumference = Data.List.sum . map D.norm . loopOffsets . fromVertices

genImage :: M.Image M.S (Co.SRGB 'Co.Linear) Double -> V2 Double -> Double -> StdGen -> QDiagram SVG V2 Double Any
genImage image dimensionsVec minDistance gen =
    scaleX widthHeightRatio
        . reflectY
        . rectEnvelope (mkP2 0 0) (1 ^& 1)
        . mconcat
        . map drawVoroniRegion
        . sortOn shapeCircumference
        . withStrategy (parListChunk 50 rdeepseq)
        . map (uncurry Tri.voroniDiagramCorners)
        $ voroni
  where
    drawVoroniRegion shape =
        lw 0
            . fillColor (Tri.voroniRegionAverageColor image dimensions shape)
            . strokeLocLoop
            . fromVertices
            $ shape

    widthHeightRatio :: Double
    widthHeightRatio = (dimensionsVec ^. _x) / (dimensionsVec ^. _y)

    dimensions :: (Int, Int)
    dimensions = (ceiling $ dimensionsVec ^. _x, ceiling $ dimensionsVec ^. _y)

    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'

    averageSideSize = (fromIntegral (uncurry (+) dimensions)) / 2

    padding = (/ 10) <$> V2 1 1

    corners' :: [P2 Double]
    corners' = map (p2 . Bi.first (/ widthHeightRatio) . unp2 . (.-^ ((/ 2) <$> padding))) . MDS.randomPoints ((mkP2 widthHeightRatio 1) .+^ padding) (minDistance * widthHeightRatio) $ gen

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

toDimensionVector :: (Ma.Size r, Fractional a) => Ma.Array r Ma.Ix2 e -> V2 a
toDimensionVector image =
    p2 (fromIntegral $ cols, fromIntegral $ rows) .-. p2 (0.0, 0.0)
  where
    (M.Sz2 rows cols) = Ma.size image

data CLIOptions = CLIOptions
    { input :: FilePath
    , output :: FilePath
    , minDistance :: Double
    }
    deriving (Generic)

instance ParseRecord CLIOptions

main :: IO ()
main = do
    CLIOptions{..} <- getRecord "image options"
    gen' <- getStdGen -- for consistency, swap with something like: pure . mkStdGen $ 2344
    print gen'
    image :: M.Image M.S (Co.SRGB 'Co.Linear) Double <- M.readImageAuto input
    let dims = toDimensionVector image
    renderSVG output (Diagrams.Prelude.dims dims) (genImage image dims minDistance gen')
