{-# LANGUAGE OverloadedLists #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE NoMonomorphismRestriction #-}

module MathSpec where

import Data.ByteString.Builder qualified as ByteString.Builder
import Data.Colour (Colour)
import Data.Colour.Palette.BrewerSet qualified as Brewer
import Data.List qualified as List
import Data.Text.Short qualified as ShortText
import Diagrams.Backend.SVG
import Diagrams.Core.Envelope (envelopeS, envelopeV)
import Diagrams.Prelude hiding (dot, outer, radius)
import Diagrams.TwoD (rotateBy, translateX)
import Diagrams.TwoD.Factorization qualified as Factorization
import GHC.Conc qualified
import Graphics.Svg qualified as SVG
import Paths_worksheets qualified as Self
import System.Directory qualified as IO
import System.FilePath (joinPath, pathSeparator, (<.>), (</>))
import System.FilePath.Posix qualified as File
import System.IO qualified as IO
import Test.Syd (Spec, describe, doNotRandomiseExecutionOrder, goldenByteStringBuilderFile, it, liftIO, withoutRetries)
import Text.Blaze.Html5.Attributes qualified as HA
import Utils.Tests (goldenPath)
import Worksheets.Utils.Prelude
import Prelude (pi, sin, sqrt, (/), (^))

import Math.NumberTheory.Primes (Prime (..), factorise, primes)

data Config = Config
  { configName :: String
  , configUnit :: QDiagram SVG V2 Double Any
  , configUnitNested :: QDiagram SVG V2 Double Any
  , configColors :: [Colour Double]
  , configColors2 :: [Colour Double]
  , configPack :: Count -> Diag -> Diag -> Diag
  }
config1 =
  Config
    { configName = "config1"
    , configUnit = circle 1 # lw 1 # lc black # fc white
    , configUnitNested = circle 1 # lw 0 # fc white
    , configColors = Brewer.brewerSet Brewer.Blues 9 & List.reverse & List.drop 0
    , configColors2 = Brewer.brewerSet Brewer.Oranges 9 & List.reverse & List.drop 0
    , configPack = packRings
    }
config2 =
  Config
    { configName = "config2"
    , configUnit = circle 1 # lw 1 # lc black # fc white
    , configUnitNested = circle 1 # lw 0 # fc white
    , -- , configColors = Brewer.brewerSet Brewer.Blues 9 & List.drop 2 & evens
      configColors = Brewer.brewerSet Brewer.Blues 4
    , configColors2 = Brewer.brewerSet Brewer.Purples 4 & List.reverse
    , configPack = packRings
    }
config3 =
  Config
    { configName = "config3"
    , configUnit = circle 1 # lw 1 # lc black # fc white
    , configUnitNested = circle 1 # lw 0 # fc white
    , -- , configColors = Brewer.brewerSet Brewer.Blues 9 & List.drop 2 & evens
      configColors = Brewer.brewerSet Brewer.Oranges 4
    , configColors2 = Brewer.brewerSet Brewer.Reds 4 & List.reverse
    , configPack = packRings
    }
config4 =
  Config
    { configName = "config4"
    , configUnit = circle 1 # lw 1 # lc black # fc white
    , configUnitNested = circle 1 # lw 0 # fc white
    , configColors = Brewer.brewerSet Brewer.Blues 9 & List.reverse & List.drop 2
    , configColors2 = Brewer.brewerSet Brewer.Reds 9 & List.reverse & List.drop 2
    , configPack = packRings
    }

evens (x : xs) = x : odds xs
evens _ = []

odds (_ : xs) = evens xs
odds _ = []

spec :: HasCallStack => Spec
spec = do
  describe "Math" do
    describe "Arithmetic" do
      withoutRetries do
        doNotRandomiseExecutionOrder do
          forM_ ([{-config1,config2, config3,-} config4] & list) \config -> do
            describe (configName config) do
              -- forM_ (list [1,2,3,4,5,6,7,8,9, 44, 460, 554, 2025]) \n -> do
              --  numberSpec n (n & primeFactorsAscending) $
              --    nest2 config $ primeFactorsAscending n
              --  numberSpec n (n & primeFactorsDescending) $
              --    nest2 config $ primeFactorsDescending n
              -- forM_ ([10::Int ..10] & list) \n -> do
              --  let t = ("multtable."<>show n<>"x"<>show n)
              --  goldenDiagram t t $
              --    multtable config n
              -- describe "num" do
              --  forM_ ([1..100] & list) \n -> do
              --    numberSpec n [n] do
              --      packRings n
              --        (circle 1 # lw 0 # lc black # fc black)
              --        (hrule 2 # lc red # lw 1 `atop` circle 1 # lw 0 # fc white)
              describe "decomp" do
                threads <- liftIO GHC.Conc.listThreads
                traceShowM $ threads & List.length
                forM_ ([20 .. 20] & list) \n -> do
                  numberSpec "" n (primeFactorsDescending n) do
                    factorsDiag config n

-- numberSpec ".rot90" n (primeFactorsDescending n) do
--  factorsDiag config{configPack= \count outer inner -> packRings count outer inner & (# rotate (90@@deg)) } n
-- liftIO $ GHC.Conc.listThreads >>= mapM_ GHC.Conc.killThread

-- goldenDiagram "table" "table" $
--  table
-- describe "Primes" do
--  forM_ (primeDiagrams & List.take 100) \(prim, diag) ->
--    numberSpec (prim & unPrime & fromIntegral) diag

numberSpec :: String -> Int -> [Int] -> QDiagram SVG V2 Double Any -> _
numberSpec suffix num facts diag = do
  let title = show num
  let name = title <> "." <> List.intercalate "×" (show <$> facts) <> suffix
  goldenDiagram (title <> suffix) name diag

fd conf n = nest2 conf $ primeFactorsDescending n
factorDiagramDescending conf n = nest2 conf $ primeFactorsDescending n

-- Factorization.factorDiagram' @Double facts

goldenDiagram :: String -> String -> QDiagram SVG V2 Double Any -> _
goldenDiagram title name diag = do
  outPath <- goldenPath name "svg"
  it title do
    goldenByteStringBuilderFile outPath do
      -- ExplanationNote: factors from greatest to lowest
      return $
        diag
          & renderDia SVG (SVGOptions (mkWidth 1024) Nothing "" [] True)
          & SVG.renderBS
          & ByteString.Builder.lazyByteString

-----------------------------------------------------------------------------

-- |
-- Module      :  Diagrams.TwoD.Layout.RingPacking
-- Copyright   :  (c) 2013 Jan Van lent
-- License     :  BSD-style (see LICENSE)
-- Maintainer  :  jakke.vanlent+git@gmail.com
--
-- More compact versions of the factorization diagrams, as seen at
-- <http://mathlesstraveled.com/2012/10/05/factorization-diagrams/>
-- and
-- <http://mathlesstraveled.com/2012/11/05/more-factorization-diagrams/>
--
-- The compact layout is achieved by circle in circle packings based on concentric rings.
-- The resulting packings are not the most compact, but they have
-- more symmetry and the method is very simple.
-- The radius of a circle with area equal to that of $n$ unit circles
-- is equal to $\sqrt{n}$.
-- This can clearly not be achieved if the unit circles are not allowed to
-- overlap.
-- The compact layout of $n$ unit circles has a bounding radius that scales
-- as $\sqrt{8/7 n + O(1)} \approx 1.07 \sqrt{n + O(1)}$.
-- The infinite hexagonal packing is the best packing of identical circles
-- in the plane.
-- For this case, 10% extra area per circle (factor

-- $\sqrt{12}/\pi \approx 1.10$).
-- If we could pack $n$ circles and the 10% extra area perfectly into
-- a circle, it would have a radius of about $\sqrt{1.1 n}$ or

-- $1.05 \sqrt{n}$.
--
-- The bounding radius of the factorization diagrams scales as $O(n)$,
-- because numbers of the form $n=2^p$ are layed out in a linear fashion.
-- More compact diagrams are obtained by combining all identical factors.
-- E.g. use $72 = 2^3*3^2 = 8*9$ instead of $72 = 2*2*2*3*3$.
--
-- The main example is "allfactorisations.svg".
-- Prime numbers show up as a single compact diagram with only one color.
-- Powers of primes show up as a single, less compact diagram with as
-- many colors as there are factors.
-- For numbers with more than one distinct factor, the results for all
-- possible ordering of factors are shown.
--
-- Even quite big numbers still have reasonably compact factorization
-- diagrams as is shown by the example with 2012 and 2013 ("years.svg")

type Factor = Int
type Diag = QDiagram SVG V2 Double Any
nest2 :: Config -> [Factor] -> Diag
nest2 Config{..} facts =
  List.zip facts circs
    & List.foldr
      (\(fact, outer) inner -> configPack fact outer inner)
      configUnitNested
  where
    circs =
      [ circle 1 # lw 0 # fc col
      | col <- configColors
      ]
        & (\l -> if null l then [circle 1 # lw 0 # fc green] else l)
        & List.cycle
        & List.take (List.length facts + 1)
        & List.reverse

-- & List.tail
-- & (circle 1 # lw 0 # fc white :)

factorsDiag conf (n :: Int) =
  vcat $
    hcat
      [ txt (show facts)
      , nest2 conf facts
      ]
      : hcat
        [ alignedText 0 0 "=" # fontSize 100
        ]
      : [ hcat
          [ txt (show subfacts)
          , nest2 conf subfacts # scale radius
          , alignedText 0 0.5 (show radius) # fc green # fontSize 100
          ]
        | 1 < List.length facts
        , (fact, radius) :: (Int, _) <-
            List.zip facts $
              facts
                & List.scanl (\r fact -> r / (ringCountToRingRadius fact + 1)) 1
        , let subfacts = fact & primeFactorsDescending
        -- , let facts = factorDiagramDescending conf fact
        ]
  where
    txt t = alignedText 1 0.5 t # fontSize 100
    facts = n & primeFactorsDescending

packRings :: Count -> Diag -> Diag -> Diag
packRings count outer inner =
  [ packRing ringCount outer inner
  | ringCount <- count & partitionRings & List.reverse
  -- CorrectionNote: reverse to get the inner rings layered ontop of the outer rings
  ]
    & mconcat

-- & (# rotate (90@@deg))

-- `packRing count outer inner`
-- aligns `count` copies of the `inner` diagram along a ring,
-- with `outer` scaled to surround those `inner` diagrams.
packRing count outer inner = ringCircles <> ringCircum
  where
    ringRadius = count & ringCountToRingRadius
    innerRadius = 1
    ringRadiusScaled = ringRadius * innerRadius
    ringCircles =
      [ inner
        -- translate horizontaly on the ring
        # translateX ringRadiusScaled
        -- rotate along the ring
        # rotateBy ((fromIntegral i / fromIntegral count))
        -- scale down to keep it a unit circle
        # scale (1 / (ringRadiusScaled + innerRadius))
        <> alignedText 1 0.5 (show (1 / (ringRadiusScaled + innerRadius)))
        # fontSize 40
        # fc red
      | i <- [0 .. count - 1]
      ]
        & mconcat
    ringCircum = outer

-- | `ringCountToRingRadius n` is the necessary ringRadius
-- of a ring of `n` unit circles whose centers are on that ring.
--
-- See `RefOptimalPackingsForFilledRingsOfCircles`, figure 2.
ringCountToRingRadius :: Count -> Double
ringCountToRingRadius 0 = 0
ringCountToRingRadius 1 = 0
ringCountToRingRadius m = 1 / sin (pi / fromIntegral m)

type Count = Int

-- | `partitionRings n` is the list of number ringed circles
partitionRings :: Factor -> [Count]
partitionRings =
  List.unfoldr
    \case
      count
        | count == 0 -> Nothing
        | otherwise -> Just (outerCount, count - outerCount)
        where
          (_, outerCount) = countToOuterRingCount List.!! count

countToOuterRingCount :: [(Count, Count)]
countToOuterRingCount = 0 & List.unfoldr \n -> Just ((n, go n), n + 1)
  where
    go :: Count -> Count
    go 0 = 0
    go 1 = 1
    go n
      -- ExplanationNote:
      -- when the outer ring of `n-1` circles (eg. 6)
      -- can contain the outer ring of the remaining circles (eg. 1),
      -- the outer ring for `n` circles (eg. 7) can remain the same size.
      -- Otherwise add
      | inners `ringFitsIn` prevOuters = prevOuters
      | otherwise = prevOuters + 1
      where
        (_, prevOuters) = countToOuterRingCount List.!! (n - 1)
        (_, inners) = countToOuterRingCount List.!! (n - prevOuters `max` 0)

-- | `ringFitsIn inners outers` is `True` iif.
-- the circles on a ring of `inners` unit circles are disjoints or tangents
-- to the circles on a ring of `outers` unit circles.
ringFitsIn :: Count -> Count -> Bool
ringFitsIn inners outers
  | inners <= outers =
      2 <= ringCountToRingRadius outers - ringCountToRingRadius inners
  | otherwise = errorShow ("ringFitsIn" :: String, inners, outers)

{-
-- Return
partitionRings :: Factor -> [Count]
partitionRings =
  List.unfoldr
    \case
      fact | fact == 0 -> Nothing
           | otherwise -> Just (c, fact - c)
           where c = ringCount fact

ringCount :: Factor -> Count
ringCount = (fmap go [0::Factor ..] List.!!)
  where
  go 0 = 0
  go 1 = 1
  go n = outers + increment
    where
      outers = ringCount (n - 1)
      inners = ringCount (n - outers)
      increment = if inners `ringFitsIn` outers then 0 else 1
-}

---- equivalent definition
-- ringFitsIn 0 _ = True
-- ringFitsIn 1 m' = m' >= 6
-- ringFitsIn 2 m' = m' >= 10
-- ringFitsIn m m' = m' - m >= 7

nest :: (t1 -> t2 -> t3 -> t3) -> [t1] -> [t2] -> t3 -> t3 -> t3
-- nest _pack [] _ d0 d1 = d0
-- nest pack (n : ns) (b : bs) d0 d1 = pack n b (nest pack ns bs d1 d1)
nest pack ns bs d0 d1 = List.zip ns bs & List.foldr (\(n, b) acc -> pack n b acc) d1

-- factors = concatMap (uncurry $ flip replicate) . factorise
factors n =
  List.concat
    [ List.replicate (fromIntegral a) f
    | (unPrime -> f, a) <- factorise n
    ]

-- factors = map (uncurry (^)) . factorise
factors' n =
  [ f ^ a
  | (unPrime -> f, a) <- factorise n
  ]

-- number of prime factors
npf = sum . List.map snd . factorise

-- number of distinct prime factors
ndpf = length . factorise

bagSelect [] = []
bagSelect ((x, 1) : b) = (x, b) : [(y, (x, 1) : b') | (y, b') <- bagSelect b]
bagSelect ((x, n) : b) = (x, (x, n - 1) : b) : [(y, (x, n) : b') | (y, b') <- bagSelect b]

-- see also <http://hackage.haskell.org/package/multiset-comb>
bagPermutations [] = [[]]
bagPermutations b =
  [ x : ys
  | (x, b') <- bagSelect b
  , ys <- bagPermutations b'
  ]

factor1 = circle 1 # lw 1 # lc black # fc white

gdot = circle 1 # lw 0 # fc grey
gdots =
  [ circle 1 # lw 0 # fc col
  | col <- [grey, white]
  ]
    & List.cycle
rainbow = [white, red, orange, yellow, green, blue, indigo, violet] & list
colDotsBlues = Brewer.brewerSet Brewer.Blues 9 & List.reverse & List.drop 2
colDotsOranges = Brewer.brewerSet Brewer.Oranges 9 & List.reverse & List.drop 2
colDotsPuOr51 = Brewer.brewerSet Brewer.PuOr 11 & List.take 5 & List.reverse
colDotsPuOr52 = Brewer.brewerSet Brewer.PuOr 11 & List.reverse & List.take 5 & List.reverse

-- main = defaultMain (packRings 20 dots factor1)
-- main = defaultMain (packRings 7 (packRings 3 gdot factor1))
-- main = defaultMain (nest packRings [7, 5, 3] gdot factor1)
numlabel n = text (show n) <> circle 1
numbers ns =
  cat
    unit_Y
    [ numlabel n === packRings n (coldots colDotsBlues List.!! (fromIntegral $ npf (n) + 1)) factor1
    | n <- ns
    ]
primeDiagrams :: [(Prime Int, QDiagram SVG V2 Double Any)]
primeDiagrams =
  [ (n, numlabel n === packRings (unPrime n) gdot factor1)
  | n <- primes
  ]

factorDiagram pack n = nest pack (primeFactorsDescending n)
powerFactorDiagram pack n = nest pack $ n & factors' <&> fromIntegral & List.reverse

primeFactorsAscending = factors >>> fmap fromIntegral
primeFactorsDescending = primeFactorsAscending >>> List.reverse

multtable :: _ => Config -> Int -> QDiagram SVG V2 Double Any
multtable conf@Config{..} n =
  vcat
    [ hcat
      [ nest2 conf{configColors = dots} facts # scaleUToY 0.9 <> square 1
      | let rowFacts = row & primeFactorsAscending
      , let rowDots = configColors & List.take (List.length rowFacts + 1)
      , col <- list [1 .. n]
      , let colFacts = col & primeFactorsAscending
      , let colDots = configColors2 & List.take (List.length colFacts + 1)
      , let dots = rowDots <> colDots
      , let facts = rowFacts <> colFacts
      ]
    | row <- list [1 .. n]
    ]

coldots cols =
  [ circle 1 # lw 0 # fc col
  | col <- white : cols
  ]
    & List.cycle

pfd Config{..} n = powerFactorDiagram packRings n (coldots colDotsBlues) configUnit configUnitNested
factorisations ns = cat unit_Y [numlabel n === fd config1 n | n <- ns]
powerfactorisations ns = cat unit_Y [numlabel n === pfd config1 n | n <- ns]
table =
  vcat
    [ hcat
      [ fd config1 ((10 :: Int) * i + j + 1) # scaleUToY 0.8 <> square 1
      | j <- list [0 .. 9]
      ]
    | i <- list [0 .. 19]
    ]
powertable =
  vcat
    [ hcat
      [ pfd config1 ((10 :: Int) * i + j + 1) # scaleUToY 0.8 <> square 1
      | j <- list [0 .. 9]
      ]
    | i <- list [0 .. 19]
    ]
allfactorisations Config{..} ns =
  cat
    unit_Y
    [ numlabel n
      === cat
        unitX
        [ nest configPack (List.map (fromIntegral . unPrime) p) (coldots colDotsBlues) configUnit configUnitNested
        | p <- bagPermutations $ factorise n
        ]
    | n <- ns
    ]
allpowerfactorisations Config{..} ns =
  cat
    unit_Y
    [ numlabel n
      === cat
        unitX
        [ nest configPack (List.map fromIntegral p) (coldots colDotsBlues) configUnit configUnitNested
        | p <- List.permutations $ factors' n
        ]
    | n <- ns
    ]
years config =
  ( ( (text "2012" # fc white # fontSize 10)
        <> (factor1 # scale (sqrt 2012))
    )
      ||| (numbers [2012] # centerY)
      ||| (allfactorisations config $ list [2012 :: Int])
      # centerY
  )
    === ( ( (text "2013" # fc white # fontSize 10)
              <> (factor1 # scale (sqrt 2013))
          )
            ||| (numbers [2013] # centerY)
            ||| (allfactorisations config $ list [2013 :: Int])
            # centerY
        )

{-
--main = defaultMain allfactorisations
main = multiMain [ ("numbers", numbers [1..60]),
                   ("primes", primeDiagrams 60),
                   ("years", years),
                   ("factorisations", factorisations [1..60]),
                   ("powerfactorisations", powerfactorisations [1..60]),
                   ("table", table),
                   ("powertable", powertable),
                   ("allfactorisations", allfactorisations [1..60]),
                   ("allpowerfactorisations", allpowerfactorisations [1..60]) ]
-}

data RefOptimalPackingsForFilledRingsOfCircles
-- ^ Reference to [Optimal Packings for Filled Rings of Circles](https://doi.org/10.21136/AM.2020.0244-19)
-- > [
-- >   {
-- >     "DOI": "10.21136/AM.2020.0244-19",
-- >     "ISSN": "1572-9109",
-- >     "URL": "https://doi.org/10.21136/AM.2020.0244-19",
-- >     "abstract": "General circle packings are arrangements of circles on a given surface such that no two circles overlap except at tangent points. In this paper, we examine the optimal arrangement of circles centered on concentric annuli, in what we term rings. Our motivation for this is two-fold: first, certain industrial applications of circle packing naturally allow for filled rings of circles; second, any packing of circles within a circle admits a ring structure if one allows for irregular spacing of circles along each ring. As a result, the optimization problem discussed herein will be extended in a subsequent paper to a more general setting. With this framework in mind, we present properties of concentric rings that have common points of tangency, the exact solution for the optimal arrangement of filled rings along with its symmetry group, and applications to construction of aluminum-conductor steel reinforced cables.",
-- >     "author": [
-- >       {
-- >         "family": "Ekanayake",
-- >         "given": "Dinesh B."
-- >       },
-- >       {
-- >         "family": "Ranpatidewage",
-- >         "given": "Manjula Mahesh"
-- >       },
-- >       {
-- >         "family": "LaFountain",
-- >         "given": "Douglas J."
-- >       }
-- >     ],
-- >     "container-title": "Applications of Mathematics",
-- >     "id": "Ekanayake2020",
-- >     "issue": "1",
-- >     "issued": {
-- >       "date-parts": [
-- >         [
-- >           2020
-- >         ]
-- >       ]
-- >     },
-- >     "page": "1-22",
-- >     "title": "Optimal Packings for Filled Rings of Circles",
-- >     "type": "article-journal",
-- >     "volume": "65"
-- >   }
-- > ]