module Triangles where

import Control.Arrow
import Control.Parallel.Strategies
import qualified Data.Array as A
import qualified Data.Colour as C
import qualified Data.Colour.Names as C
import Data.Colour.SRGB.Linear (Colour)
import qualified Data.Colour.SRGB.Linear as C
import qualified Data.Colour.SRGB.Linear as CL
import Data.Fixed
import qualified Data.Function as F
import Data.List
import qualified Data.List as L
import qualified Data.Map as M
import qualified Data.Massiv.Array as Ma
import qualified Data.Massiv.Array.IO as Ma
import Data.Maybe
import qualified Data.Ord as O
import Data.Ratio
import qualified Data.Set as S
import Data.Vector.Generic.Base (Vector)
import Data.Vector.Generic.Mutable (MVector)
import qualified Data.Vector.Unboxed as Vec
import Data.Vector.Unboxed.Deriving
import Debug.Trace (traceShow)
import qualified Debug.Trace
import qualified Debug.Trace as D
import Diagrams.Prelude
import Diagrams.Trail (trailPoints)
import Diagrams.TwoD
import qualified Diagrams.TwoD.Path.IntersectionExtras as I
import Diagrams.TwoD.Segment.Bernstein (listToBernstein)
import qualified Graphics.Color.Space as Co
import System.Random

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

-- from -0.05 to 1.05 so there aren't missing/elongated triangles at the edges
borderSize :: Double
borderSize = 0.05

randomPoints :: StdGen -> [(Double, Double)]
randomPoints = map (bimap toZeroToOneTuple toZeroToOneTuple) . randomRs ((0 :: Word, 0 :: Word), (maxBound, maxBound))
  where
    toZeroToOneTuple :: Word -> Double
    toZeroToOneTuple x = ((fromIntegral x / (fromIntegral (maxBound :: Word))) * (1 + (2 * borderSize))) - borderSize

combinations :: (Ord b, Floating b, NFData b) => [P2 b] -> [(P2 b, P2 b)]
combinations =
    sortOn (abs . uncurry distanceA)
        . S.toList -- deduplicate
        . S.fromList
        . filter (uncurry (/=))
        . concat
        . withStrategy (parListChunk 50 rdeepseq)
        . map (\(x : xs) -> take 10 . sortOn (abs . uncurry distanceA) . map (x,) $ xs)
        . init -- last output of tails is empty list
        . tails

toPlanarGraph :: forall n. (NFData n, Floating n, Ord n) => [P2 n] -> [(Point V2 n, Point V2 n)]
toPlanarGraph =
    removeIntersections
        . sortOn (abs . uncurry distanceA)
        . combinations
  where
    removeIntersections :: [(Point V2 n, Point V2 n)] -> [(Point V2 n, Point V2 n)]
    removeIntersections = foldl' addIfNoIntersection []

    addIfNoIntersection xs x
        | all (noIntersection x) xs = x : xs
        | otherwise = xs

    noIntersection l1 l2 = sharedEndPoint || (null $ intersectPointsT (uncurry toLocatedTrail l1) (uncurry toLocatedTrail l2))
      where
        sharedEndPoint = (< 4) . length . nub $ [fst l1, snd l1, fst l2, snd l2]

toLocatedTrail :: (TrailLike a) => Point (V a) (N a) -> Point (V a) (N a) -> Located a
toLocatedTrail p1 p2 = fromVertices [p1, p2] `at` p1

withinShape :: (RealFloat v) => Point V2 v -> [Point V2 v] -> Point V2 v -> Bool
withinShape pointInShape verticies candidate = all ((< quarterTurn) . fmap abs . uncurry signedAngleBetweenDirs) $ zip (shapeDirections pointInShape) (shapeDirections candidate)
  where
    shapeDirections p = map (dirBetween p) verticies

sortOnAngle center = sortOn (normalizeAngle . signedAngleBetweenDirs xDir . dirBetween center)

voroniDiagramCorners :: (RealFloat n) => Point V2 n -> [Point V2 n] -> [Point V2 n]
voroniDiagramCorners center midpoints =
    sortOnAngle center
        . filter isValidMidpoint
        . concat
        $ [intersectPointsT l0 l1 | l0 <- tangentTrails, l1 <- tangentTrails]
  where
    lessThanQuarterTurn candidate = candidate <= (10001 / 40000) @@ turn || candidate >= (29999 / 40000) @@ turn

    tangentTrails = map tangentTrail midpoints

    appendHead (x : xs) = xs ++ [x]

    isValidMidpoint candidate = all isNonObtuseMidpoint . filter (/= candidate) $ midpoints
      where
        isNonObtuseMidpoint m =
            lessThanQuarterTurn . normalizeAngle $ angleBetweenDirs (dirBetween m center) (dirBetween m candidate)

    tangentTrail midpoint = fromVertices [midpoint .-^ tangentVec, midpoint .+^ tangentVec]
      where
        -- implicitly uses the unit vector * 8 as an infinitely long vector
        tangentVec =
            scale 2
                . fromDirection
                . rotateBy (1 / 4)
                $ dirBetween midpoint center

findVoroniDiagram :: (RealFloat n) => [(Point V2 n, Point V2 n)] -> [(Point V2 n, [Point V2 n])]
findVoroniDiagram =
    M.toList
        . M.mapWithKey
            ( \key ->
                L.sortOn (normalizeAngle . signedAngleBetweenDirs xDir . dirBetween key)
                    . map (pointBetween key)
                    . S.toList
            )
        . adjacencyMapOf
  where
    pointBetween p0 p1 = p0 .+^ ((p1 .-. p0) ^/ 2)

findTriangles :: (Ord b) => [(b, b)] -> S.Set (S.Set b)
findTriangles edges = S.unions . S.map threeCyclesOf . M.keysSet $ adjacencyMap
  where
    threeCyclesOf node =
        S.unions
            . S.map (\x -> S.map (\y -> S.fromList [node, x, y]) $ S.delete node . S.intersection originalNodeNeighbors . (M.!) adjacencyMap $ x)
            $ originalNodeNeighbors
      where
        originalNodeNeighbors = fromMaybe S.empty (adjacencyMap M.!? node)

    adjacencyMap = adjacencyMapOf edges

adjacencyMapOf edges = M.fromListWith S.union . map (second S.singleton) $ (edges ++ edgesReversed)
  where
    edgesReversed = map (\(a, b) -> (b, a)) edges

triangleAdjacencyMap :: (Ord b) => S.Set (S.Set b) -> M.Map b (S.Set (S.Set b))
triangleAdjacencyMap s = M.fromListWith S.union . concatMap (\s' -> map (,S.singleton s') . S.toList $ s') $ S.toList s

getPointsInTriangle :: p -> S.Set (P2 Int) -> [(Int, Int)]
getPointsInTriangle image pts =
    S.toList . S.unions . map S.fromList $
        [ ptsBtween (makeLine p1 p3) (makeLine p1 p2),
          ptsBtween (makeLine p1 p3) (makeLine p2 p3),
          ptsBtween (makeLine p1 p2) (makeLine p2 p3)
        ]
  where
    [p1, p2, p3] = sortOn fst . map unp2 . S.toList $ pts

blendEqually :: (Ord a, Floating a) => [Colour a] -> Colour a
blendEqually colors = C.affineCombo (map (fraction,) colors) C.white
  where
    fraction = 1.0 / (fromIntegral . length $ colors)

voroniRegionAverageColor :: (Integral a, Integral b) => Ma.Image Ma.S (Co.SRGB 'Co.Linear) Double -> (a, b) -> [P2 Double] -> Colour Double
voroniRegionAverageColor image (x', y') =
    blendEqually
        . concatMap (getColorsInTriangle image (x', y'))
        . filter ((== 3) . S.size)
        . map (S.fromList . take 3)
        . tails
        . L.nub
        . map scaleToImageCoords
  where
    scaleToImageCoords :: P2 Double -> P2 Int
    scaleToImageCoords p = round <$> p2 (fromIntegral x' * p ^. _x, fromIntegral y' * p ^. _y)

    scaleToUnitCoords :: P2 Int -> P2 Double
    scaleToUnitCoords p = p2 ((fromIntegral x' / (fromIntegral $ p ^. _x)), fromIntegral y' / (fromIntegral $ p ^. _y))

getColorsInTriangle :: Ma.Image Ma.S (Co.SRGB 'Co.Linear) Double -> (a, b) -> S.Set (P2 Int) -> [Colour Double]
getColorsInTriangle image (x', y') triangle = mapMaybe index' points
  where
    points :: [(Int, Int)]
    points = getPointsInTriangle image triangle

    index' :: (Int, Int) -> Maybe (C.Colour Double)
    index' (x, y) = toColour . Ma.pixelColor <$> Ma.index image (Ma.Ix2 y x)

getTriangleAverageRGB :: Ma.Image Ma.S (Co.SRGB 'Co.Linear) Double -> (a, b) -> S.Set (P2 Int) -> Colour Double
getTriangleAverageRGB image (x', y') triangle = blendEqually $ getColorsInTriangle image (x', y') triangle

ptsBtween :: LineMXB -> LineMXB -> [(Int, Int)]
ptsBtween l1 l2 = concatMap rasterLine . noSingletons $ [startingX .. endingX]
  where
    startingX = max (startX l1) (startX l2)
    endingX = min (endX l1) (endX l2)

    rasterLine x = map (x,) $ range' (yAt l1 x) (yAt l2 x)

noSingletons :: [a] -> [a]
noSingletons [x] = []
noSingletons l = l

range' :: Int -> Int -> [Int]
range' a b = [(min a b) .. (max a b)]

yAt :: LineMXB -> Int -> Int
yAt (LineMXB {..}) x = round $ (m * (fromIntegral x)) + b

makeLine :: (Int, Int) -> (Int, Int) -> LineMXB
makeLine (x1, y1) (x2, y2) =
    LineMXB
        { m = slope,
          b = (fromIntegral y1) - (slope * (fromIntegral x1)),
          startX = min x1 x2,
          endX = max x1 x2
        }
  where
    slope =
        if x1 /= x2
            then (fromIntegral $ y1 - y2) % (fromIntegral $ x1 - x2)
            else fromIntegral . ceiling $ ((10.0 :: Double) ** 100.0)

data LineMXB = LineMXB
    { m :: Rational,
      b :: Rational,
      startX :: Int,
      endX :: Int
    }
    deriving (Show, Ord, Eq)