module Diagram.Hasse where

import Control.Monad
import Data.Map qualified as Map
import Data.Maybe
import Data.Set qualified as Set
import Data.Typeable
import Diagram qualified as D
import Diagrams.Backend.SVG.CmdLine
import Diagrams.Prelude
import Diagrams.TwoD.Text qualified as Diag

hasseRow :: (V c ~ V2, Alignable c, HasOrigin c, Floating (N c), Juxtaposable c, Monoid' c) => (a -> c) -> [a] -> c
hasseRow f = centerX . hcat' (with & sep .~ 2) . map f

hasseDiagram
  :: forall vtx. (IsName vtx, Ord vtx)
  => (vtx -> QDiagram B V2 Double Any)
  -> D.Diagram vtx
  -> QDiagram B V2 Double Any
hasseDiagram f d = centerXY $ drawConnections setsD
  where
    setsD = vcat' (with & sep .~ fromIntegral n)
      . map (hasseRow f)
      . reverse
      $ subsets
    grades = D.verticesByGrade d
    subsets :: [[vtx]] = fmap (Set.toList . snd) $ Map.toList grades
    n :: Int = maximum $ Map.keys grades
    drawConnections = applyAll connections
    connections = concat $ zipWith connectSome subsets (drop 1 subsets)
    connectSome subs1 subs2 =
      [ connect'' p s1 s2
      | s1 <- subs1
      , s2 <- subs2
      , p <- maybeToList $ Map.lookup s1 <=< Map.lookup s2 $ D._diagram_down d
      ]
    connect'' p v1 v2 = withName v1 $ \b1 -> withName v2 $ \b2 ->
      beneath (boundaryFrom b1 unitY ~~ boundaryFrom b2 unit_Y) # polarity p # lw thick

polarity
  :: (V a ~ V2, Typeable (N a), Floating (N a), HasStyle a)
  => D.Polarity
  -> a
  -> a
polarity p = lc $ case p of
  D.Positive -> red
  D.Negative -> blue

node
  :: (Typeable n, RealFloat n, Renderable (Path V2 n) b, Renderable (Diag.Text n) b)
  => Char -> QDiagram b V2 n Any
node x = named x $ (scale 0.25 $ text ([x])) <> (unitSquare # lc black # fc grey # lw thin)

example :: D.Diagram Char
example =
  let
    x = do
        pure D.empty >>=
          -- Grade 0
          D.insert 'A' 0 mempty >>= D.insert 'B' 0 mempty >>= D.insert 'C' 0 mempty >>=
          -- Grade 1
          D.insert 'f' 1 (Map.fromList [('A', D.Negative), ('B', D.Positive)]) >>=
          D.insert 'g' 1 (Map.fromList [('A', D.Negative), ('B', D.Positive)]) >>=
          D.insert 'h' 1 (Map.fromList [('B', D.Negative), ('C', D.Positive)]) >>=
          D.insert 'k' 1 (Map.fromList [('A', D.Negative), ('C', D.Positive)]) >>=
          -- Grade 2
          D.insert 'ϴ' 2 (Map.fromList [('f', D.Negative), ('g', D.Positive)]) >>=
          D.insert 'φ' 2 (Map.fromList [('g', D.Negative), ('h', D.Negative), ('k', D.Positive)]) >>=
          D.insert 'ε' 2 (Map.fromList [('f', D.Negative), ('h', D.Negative), ('k', D.Positive)]) >>=
          -- Grade 3
          D.insert 'τ' 3 (Map.fromList [('ϴ', D.Positive), ('φ', D.Positive), ('ε', D.Negative)])
  in
    case x of
      Right a -> a
      Left err -> error $ "Invalid diagram: " <> show err

main :: IO ()
main = mainWith (scale 100 $ pad 1.1 $ hasseDiagram node example :: Diagram B)