text:
{-# LANGUAGE TemplateHaskell #-}
{-# LANGUAGE TypeApplications, DataKinds, GADTs #-}
{-# LANGUAGE FlexibleInstances, MultiParamTypeClasses, FunctionalDependencies #-}

module Lib
    ( someFunc
    ) where

import Data.Maybe ( fromMaybe )
import Data.Int (Int32)
import Linear.V1
import Linear.V3
import Linear.V4 hiding (vector)
import Linear.Matrix
import Linear.Metric (dot, norm)
import Graphics.UI.GLUT ( Key(..), SpecialKey(..), MouseButton(..), KeyState(..) )
import Vis
import Vis.Vis ( vis )
import Vis.Camera ( makeCamera, setCamera, cameraMotion, cameraKeyboardMouse, Camera )
import Control.Lens
import Control.Applicative
import Control.Monad (forever)
import Numeric.AD
import Numeric.LinearAlgebra ((#>), det, fromLists, fromList, toList, pinv, inv, vector, matrix, Vector, Matrix, scale, norm_2)
--import Debug.Trace

rotM33XYZ :: (Floating f) => f -> f -> f -> M33 f
rotM33XYZ r p y = rotM33Z y !*! rotM33Y p !*! rotM33X r

rotM33X :: (Floating f) => f -> M33 f
rotM33X a = V3 (V3 1 0 0) (V3 0 ca (-sa))  (V3 0 sa ca) where
    ca = cos a
    sa = sin a

rotM33Y :: (Floating f) => f -> M33 f
rotM33Y a = V3 (V3 ca 0 sa) (V3 0 1 0)  (V3 (-sa) 0 ca) where
    ca = cos a
    sa = sin a

rotM33Z :: (Floating f) => f -> M33 f
rotM33Z a = V3 (V3 ca (-sa) 0)  (V3 sa ca 0) (V3 0 0 1) where
    ca = cos a
    sa = sin a

printV4 :: (Show f) => V4 f -> IO ()
printV4 (V4 v1 v2 v3 v4) = print [v1, v2, v3, v4]

printV3 :: (Show f) => V3 f -> IO ()
printV3 (V3 v1 v2 v3) = print [v1, v2, v3]

printM33 :: (Show f) => M33 f -> IO ()
printM33 (V3 v1 v2 v3) = do
    putStrLn ""
    printV3 v1
    printV3 v2
    printV3 v3
    putStrLn ""

printM44 :: (Show f) => M44 f -> IO ()
printM44 (V4 v1 v2 v3 v4) = do
    putStrLn ""
    printV4 v1
    printV4 v2
    printV4 v3
    printV4 v4
    putStrLn ""

rotM44X :: (Floating f) => f -> M44 f
rotM44X a = identity & _m33 .~ rotM33X a

rotM44Z :: (Floating f) => f -> M44 f
rotM44Z a = identity & _m33 .~ rotM33Z a

transM44X :: (Floating f) => f -> M44 f
transM44X a =  identity & translation .~ V3 a 0 0

transM44Z :: (Floating f) => f -> M44 f
transM44Z a =  identity & translation .~ V3 0 0 a

twistX :: (Floating f) => f -> f -> M44 f
twistX d theta = transM44X d !*! rotM44X theta

twistZ :: (Floating f) => f -> f -> M44 f
twistZ d theta = transM44Z d !*! rotM44Z theta

skew :: (Floating f) => V3 f -> M33 f
skew (V3 x y z) = V3 (V3 0 (-z) y) (V3 z 0 (-x)) (V3 (-y) x 0)

invSkew :: (Floating f) => M33 f -> V3 f
invSkew (V3 (V3 _ _ y) (V3 z _ _) (V3 _ x _)) = V3 x y z

rodrigues :: (Eq f, Floating f) => V3 f -> M33 f
rodrigues r = identity !!* ct  !+! rrt !!* (1 - ct) !+! sr !!* st
    where
    theta = norm r
    ct    = cos theta
    st    = sin theta
    r'    = if theta == 0 then V3 1 0 0 else r / V3 theta theta theta
    sr    = skew r'
    mrt   = V1 r'
    mr    = transpose mrt
    rrt   = mr !*! mrt

invRodrigues :: (Floating f) => M33 f -> V3 f
invRodrigues m = V3 x' y' z' * V3 t t t
    where
    V3 x y z = invSkew (m - transpose m) / 2
    st = norm (V3 x y z)
    t  = asin st
    V3 x' y' z' = V3 x y z / V3 t t t

frame2xyzaxayaz :: (Floating f) => M44 f -> [f]
frame2xyzaxayaz m = [x, y, z, ax, ay, az]
    where
    V3 ax ay az = invRodrigues $ m ^. _m33
    V3 x y z = m ^. translation

class Renderable b f  | b -> f where
    renderFC :: Flavour -> Color -> b -> VisObject f
    render   :: b -> VisObject f
    render = renderFC Wireframe blue

-- M44 を render
instance (Floating f) => Renderable (M44 f) f where
    renderFC _ _ m = convfrm m $ axes (0.05, 10)
        where
        axes (size, aspectRatio) = VisObjects [ Arrow (size, aspectRatio) (V3 1 0 0) red
                                              , Arrow (size, aspectRatio) (V3 0 0 1) blue
                                              ]

class FrameConvertable b f | b -> f where
    convfrm :: M44 f -> b -> b

instance (Floating f) => FrameConvertable (VisObject f) f where
    convfrm m v = Trans t $ RotDcm rot v
        where
        rot = transpose $ m ^. _m33 -- NOTE: vis function use row vector??
        t   = m ^. translation

---- for rendering and collision detection
data LinkShape f = OBB (V3 f) (V3 f) (M33 f) -- center size/2 axis-vector
    deriving Show

instance (Floating f) => FrameConvertable (LinkShape f) f where
    convfrm m (OBB c a r) = OBB c' a r'
        where
        c' = (m !* point c) ^. _xyz
        r' = m^._m33 !*! r

instance (Floating f) => Renderable (LinkShape f) f where
    renderFC f color (OBB c b@(V3 bx by bz) m) = Trans c $ RotDcm rot $ Box (2*bx, 2*by, 2*bz) f color
        where
        rot = transpose $ m ^. _m33 -- NOTE: vis function use row vector??

hasIntersection :: (Ord f, Floating f) => LinkShape f -> LinkShape f -> Bool
hasIntersection (OBB c0 a@(V3 a0 a1 a2)  m0) (OBB c1 b@(V3 b0 b1 b2) m1) = not $ easy || i0 || i1 || i2 || i3 || i4 || i5 || i6 || i7 || i8 || i9 || i10 || i11 || i12 || i13 || i14
    where

    -- not intersection test using bounding sphere
    easy = norm a + norm b < norm vD

    -- not intersection test
    i0  = a0 + (b0*c00' + b1*c01' + b2*c02') < abs sA0D
    i1  = a1 + (b0*c10' + b1*c11' + b2*c12') < abs sA1D
    i2  = a2 + (b0*c20' + b1*c21' + b2*c22') < abs sA2D

    i3  = (a0*c00' + a1*c10' + a2*c20') + b0 < abs (vB0 `dot` vD)
    i4  = (a0*c01' + a1*c11' + a2*c21') + b1 < abs (vB1 `dot` vD)
    i5  = (a0*c02' + a1*c12' + a2*c22') + b2 < abs (vB2 `dot` vD)

    i6  = (a1*c20' + a2*c10') + (b1*c02' + b2*c01') < abs (c10*sA2D - c20*sA1D)
    i7  = (a1*c21' + a2*c11') + (b0*c02' + b2*c00') < abs (c11*sA2D - c21*sA1D)
    i8  = (a1*c22' + a2*c12') + (b0*c01' + b1*c00') < abs (c12*sA2D - c22*sA1D)

    i9  = (a0*c20' + a2*c00') + (b1*c12' + b2*c11') < abs (c20*sA0D - c00*sA2D)
    i10 = (a0*c21' + a2*c01') + (b0*c12' + b2*c10') < abs (c21*sA0D - c01*sA2D)
    i11 = (a0*c22' + a2*c02') + (b0*c11' + b1*c10') < abs (c22*sA0D - c02*sA2D)

    i12 = (a0*c10' + a1*c00') + (b1*c22' + b2*c21') < abs (c00*sA1D - c10*sA0D)
    i13 = (a0*c11' + a1*c01') + (b0*c22' + b2*c20') < abs (c01*sA1D - c11*sA0D)
    i14 = (a0*c12' + a1*c02') + (b0*c21' + b1*c20') < abs (c02*sA1D - c12*sA0D)

    vD = c1 - c0
    V3 vA0 vA1  vA2 = transpose m0
    V3 vB0 vB1  vB2 = transpose m1
    V3 (V3 c00 c01 c02) (V3 c10 c11 c12) (V3 c20 c21 c22)  = transpose m0 !*! m1

    ((c00', c01', c02'), (c10', c11', c12'), (c20', c21', c22'))  = ((abs c00, abs c01, abs c02), (abs c10, abs c11, abs c12), (abs c20, abs c21, abs c22))

    sA0D = vA0 `dot` vD
    sA1D = vA1 `dot` vD
    sA2D = vA2 `dot` vD

-- TODO
hasCollisions :: (Ord f, Floating f) => [LinkShape f] -> [Bool]
hasCollisions (x:xs) = any id (hasIntersection x <$> xs) : hasCollisions xs
hasCollisions [] = []

aabb :: (Ord f, Floating f) => (f, f) -> (f, f) -> (f, f) -> LinkShape f
aabb (x0,x1) (y0,y1) (z0,z1) = OBB c b identity
        where
        c = V3 ((x1+x0)/2) ((y1+y0)/2) ((z1+z0)/2)
        b = V3 ((x1-x0)/2) ((y1-y0)/2) ((z1-z0)/2)

type FrameTrans f = f -> M44 f

data DH f = DH {   __a :: f      -- [m]
               ,   __d :: f      -- [m]
               ,   __alpha :: f  -- [rad]
               } deriving Show

makeLenses ''DH

d1 :: (Floating f) => f
d1 = 0.089159

a2 :: (Floating f) => f
a2 = 0.425

a3 :: (Floating f) => f
a3 = 0.39225

d4 :: (Floating f) => f
d4 = 0.10915

d5 :: (Floating f) => f
d5 = 0.09465

d6 :: (Floating f) => f
d6 = 0.0823

ur5dh :: (Floating f) => [DH f]
ur5dh =
    [ DH 0             d1          (pi/2) -- J1
    , DH (-a2)         0           0      -- J2
    , DH (-a3)         0           0      -- J3
    , DH 0             d4          (pi/2) -- J4
    , DH 0             d5          (-pi/2)-- J5
    , DH 0             d6          0      -- J6(flange)
    ]

tcpOffsetZ :: (Floating f) => f
tcpOffsetZ = 0.15

-- how to draw link (from frame{i} to frame{i+1})
ur5link :: (Ord f, Floating f) => [LinkShape f]
ur5link =
    [ aabb (-r/2, r/2) (-d1, r/2) (-r/2, r/2) -- 1,2
    , aabb (-r/2, a2 + r/2) (-r/2, r/2) (delta+r/2, 3*r/2) -- 2,3
    , aabb (-r/2, a3 + r/2) (-r/2, r/2) (-r/2, r/2) -- 3,4
    , aabb (-r/2, r/2) (r/2 + delta-d4, r/2) (-r/2, r/2) -- 4,5
    , aabb (-r/2, r/2) (-r/2, d5 - r/2 - delta) (-r/2, r/2) -- 5,6
    , aabb (-r/2, r/2) (-r/2, r/2) (delta-d6+r, 0) -- 6,frange
    , aabb (-r/2, r/2) (-r/2, r/2) (delta, tcpOffsetZ)  -- tool
    ]
    where
    r = 0.05
    delta = 0.01

-- https://en.wikipedia.org/wiki/Denavit%E2%80%93Hartenberg_parameters
dr2ft :: (Floating f) => DH f -> FrameTrans f
dr2ft dh theta = twistZ (dh ^. _d) theta !*!  twistX (dh^. _a) (dh^. _alpha)

ur5fts :: (Floating f) => [FrameTrans f]
ur5fts = dr2ft <$> ur5dh

-- calc each frame in base-coordinate for each joint-angle
kinematics :: (Floating f) => [FrameTrans f] -> [f] -> [M44 f]
kinematics (e:es) (a:as) = t : ((t !*! ) <$> kinematics es as)
    where
    t = e a
kinematics [] [] = []
kinematics _ _ = error "invalid args"

joint2tcp :: (Floating f) => [FrameTrans f] -> [f] -> [f]
joint2tcp es js = frame2xyzaxayaz $ last ( kinematics es js) !*! transM44Z tcpOffsetZ

poslabel = ["x", "y", "z", "ax", "ay", "az"]

data World = World { _jPos    :: [Double]
                   , _tcpVel  :: [Double]
                   , _trajct   :: [V3 Double]
                   , _camera :: Camera
                   }

makeLenses ''World

world0 = World [0.1, 0.4, 0.4, 0.1, 0.1, 0.1] [0, 0, 0, 0, 0, 0] [] cameraState0
    where
    defaultCamera = Camera0 { phi0 = 60 , theta0 = 20 , rho0 = 7 }
    cameraState0 = makeCamera $ fromMaybe defaultCamera (optInitialCamera defaultOpts)

motionPolicy :: [Double] -> [Double] ->  [Double] -> [Double]
motionPolicy goal pos vel = toList acc
    where
    alpha = 0.001
    c     = 0.01
    beta  = 0.5
    acc :: Vector Double
    acc = alpha * s (vector goal - vector pos) - beta * vector vel
    s :: Vector Double -> Vector Double
    s v = scale (1 / h (norm_2 v)) v
    h :: Double -> Double
    h z = z + c * log (1 + exp (-2) * c * z)

mySimulateIO :: IO ()
mySimulateIO = vis defaultOpts 0.016 world0 simFun drawFun setCameraFun (Just kmCallback) (Just motionCallback) Nothing
  where
    delta = 0.001

    simFun (w,t) = return $ w  &~ do
                        let tcpp = joint2tcp ur5fts $ w ^. jPos
                        let tcpv = w ^. tcpVel
                        let goal = [0.192428507378705,-0.39791168805212945,-0.43347675593345414,0.7684417340393249,-0.40013839074582813,0.3969016751654795]
                        let tcpa = motionPolicy goal tcpp tcpv
                        let tcpv' = zipWith (+) tcpv  tcpa
                        let j    = fromLists (jacobian (joint2tcp $ ur5fts) (w ^. jPos) :: [[Double]])
                        let jvel = if det j > 0.0001 then toList $ pinv j #> fromList tcpv else [0, 0, 0, 0, 0, 0]
                        let jpos' = zipWith (+) (w ^. jPos) jvel
                        let pos  = joint2tcp ur5fts jpos'
                        jPos .= jpos'
                        tcpVel .= tcpv'
                        trajct %= (:) (V3 (pos !! 0) (pos !! 1) (pos !! 2))

    drawFun (w, t) = do
        let jpos = w ^. jPos
        let frames =  identity : kinematics ur5fts jpos
        let vframes = render <$> frames
        let ts     = tail frames ++ [last frames]
        let links  = zipWith convfrm ts ur5link
        let colors  = (\b -> if b then red else blue) <$> hasCollisions links
        let vlinks  = zipWith (renderFC Wireframe) colors links
        let pos = joint2tcp ur5fts jpos
        let info = [Text2d (show (l, j)) (0, 110  - i * 20) Fixed8By13 red | (l, i,j) <- zip3 poslabel [0..] pos]
        let infoj = [Text2d (show ("J", i, j)) (300, 110  - i * 20) Fixed8By13 red | (i,j) <- zip [0..] jpos]
        let points = Points (w ^. trajct) (Just 1.0) yellow
        print pos
        return (VisObjects $ vlinks ++ vframes ++ info ++ infoj ++ [points], Nothing)

    setCameraFun w = setCamera $ w ^. camera

    kmCallback w k0 k1 _ _ = case (k0,  k1) of
                         (Char 'h', Down)            -> w & tcpVel .~ [ delta, 0, 0, 0, 0, 0]
                         (Char 'l', Down)            -> w & tcpVel .~ [-delta, 0, 0, 0, 0, 0]
                         (Char 'j', Down)            -> w & tcpVel .~ [0,  delta, 0, 0, 0, 0]
                         (Char 'k', Down)            -> w & tcpVel .~ [0, -delta, 0, 0, 0, 0]
                         (Char 'n', Down)            -> w & tcpVel .~ [0, 0,  delta, 0, 0, 0]
                         (Char 'p', Down)            -> w & tcpVel .~ [0, 0, -delta, 0, 0, 0]
                         (SpecialKey KeyLeft,  Down) -> w & tcpVel .~ [0, 0, 0,  delta, 0, 0]
                         (SpecialKey KeyRight, Down) -> w & tcpVel .~ [0, 0, 0, -delta, 0, 0]
                         (SpecialKey KeyDown,  Down) -> w & tcpVel .~ [0, 0, 0, 0,  delta, 0]
                         (SpecialKey KeyUp,    Down) -> w & tcpVel .~ [0, 0, 0, 0, -delta, 0]
                         (Char ',', Down)            -> w & tcpVel .~ [0, 0, 0, 0, 0,  delta]
                         (Char '.', Down)            -> w & tcpVel .~ [0, 0, 0, 0, 0, -delta]
                         (Char 'a', Down)            -> world0 & camera .~ (w ^. camera)
                         (Char 's', Down)            -> w & tcpVel .~ [0, 0, 0, 0, 0, 0]
                         (_, _)           -> w & camera %~ (\x -> cameraKeyboardMouse x k0 k1)

    motionCallback w pos = w & camera %~ (`cameraMotion` pos)

someFunc :: IO ()
someFunc = mySimulateIO

main = mySimulateIO

max_stars_repo_name: SuperLazyK/navigation

max_issues_repo_name: SuperLazyK/navigation

max_forks_repo_name: SuperLazyK/navigation


text:
module ExampleSeeds where

import           Data.Complex
import           Seed                           ( Seed
                                                , createSeed
                                                )

circleSeed :: Seed
circleSeed = createSeed f n spaceScale timeScale trailLength
 where
  f           = cis . (*) (2 * pi)
  n           = 101
  spaceScale  = 300
  timeScale   = 1
  trailLength = 0.8

lineSeed :: Seed
lineSeed = createSeed f n spaceScale timeScale trailLength
 where
  f t = (cos (2 * pi * t) :+ 0) * (1 :+ 1)
  n           = 101
  spaceScale  = 300
  timeScale   = 0.5
  trailLength = 1

rodFunction :: Float -> Complex Float
rodFunction x | x < 0.5   = -1
              | x > 0.5   = 1
              | otherwise = 0

rodSeed :: Seed
rodSeed = createSeed rodFunction n spaceScale timeScale trailLength
 where
  n           = 501
  spaceScale  = 500
  timeScale   = 0.1
  trailLength = 0.01

cardioid :: (RealFloat a) => a -> a -> Complex a
cardioid a t = x :+ y
 where
  x = 2 * a * (1 - cos t) * cos t
  y = 2 * a * (1 - cos t) * sin t

cardioidSeed :: Seed
cardioidSeed = createSeed (cardioid 100 . (*) (2 * pi))
                          n
                          spaceScale
                          timeScale
                          trailLength
 where
  n           = 21
  spaceScale  = 1
  timeScale   = 0.1
  trailLength = 0.8

lissajous :: (RealFloat a) => a -> a -> a -> Complex a
lissajous a b t = cos (a * t) :+ sin (b * t)

infinitySeed :: Seed
infinitySeed = createSeed (lissajous 1 2 . (*) (2 * pi))
                          n
                          spaceScale
                          timeScale
                          trailLength
 where
  n           = 51
  spaceScale  = 300
  timeScale   = 0.1
  trailLength = 1

parabolaSeed :: Seed
parabolaSeed = createSeed (lissajous 2 1 . (*) (2 * pi))
                          n
                          spaceScale
                          timeScale
                          trailLength
 where
  n           = 51
  spaceScale  = 300
  timeScale   = 0.1
  trailLength = 1

heart :: (RealFloat a) => a -> Complex a
heart t = x t :+ y t
 where
  x t = 16 * sin t ^ 3
  y t = 13 * cos t - 5 * cos (2 * t) - 2 * cos (3 * t) - cos (4 * t)

heartSeed :: Seed
heartSeed = createSeed (heart . (*) (2 * pi))
                       n
                       spaceScale
                       timeScale
                       trailLength
 where
  n           = 21
  spaceScale  = 20
  timeScale   = 0.1
  trailLength = 1

ext: hs

max_stars_repo_name: juliagracedefoor/fourier-visualizer

max_issues_repo_name: juliagracedefoor/fourier-visualizer

max_forks_repo_name: juliagracedefoor/fourier-visualizer


text:
-- |
-- Module:      Math.NumberTheory.DirichletCharacters
-- Copyright:   (c) 2018 Bhavik Mehta
-- Licence:     MIT
-- Maintainer:  Bhavik Mehta <bhavikmehta8@gmail.com>
--
-- Implementation and enumeration of Dirichlet characters.
--

{-# LANGUAGE CPP                        #-}
{-# LANGUAGE DataKinds                  #-}
{-# LANGUAGE GADTs                      #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE KindSignatures             #-}
{-# LANGUAGE LambdaCase                 #-}
{-# LANGUAGE PatternSynonyms            #-}
{-# LANGUAGE ScopedTypeVariables        #-}
{-# LANGUAGE ViewPatterns               #-}

module Math.NumberTheory.DirichletCharacters
  (
  -- * Roots of unity
    RootOfUnity
  -- ** Conversions
  , toRootOfUnity
  , fromRootOfUnity
  , toComplex
  -- * An absorbing semigroup
  , OrZero, pattern Zero, pattern NonZero
  , orZeroToNum
  -- * Dirichlet characters
  , DirichletCharacter
  -- ** Construction
  , indexToChar
  , indicesToChars
  , characterNumber
  , allChars
  , fromTable
  -- ** Evaluation
  , eval
  , evalGeneral
  , evalAll
  -- ** Special Dirichlet characters
  , principalChar
  , isPrincipal
  , orderChar
  -- ** Real Dirichlet characters
  , RealCharacter
  , isRealCharacter
  , getRealChar
  , toRealFunction
  , jacobiCharacter
  -- ** Primitive characters
  , PrimitiveCharacter
  , isPrimitive
  , getPrimitiveChar
  , induced
  , makePrimitive
  , WithNat(..)
  -- * Debugging
  , validChar
  ) where

#if !MIN_VERSION_base(4,12,0)
import Control.Applicative                                 (liftA2)
#endif
import Data.Bits                                           (Bits(..))
import Data.Complex                                        (Complex(..), cis)
import Data.Foldable                                       (for_)
import Data.Functor.Identity                               (Identity(..))
import Data.List                                           (mapAccumL, foldl', sort, find, unfoldr)
import Data.Maybe                                          (mapMaybe, fromJust, fromMaybe)
#if MIN_VERSION_base(4,12,0)
import Data.Monoid                                         (Ap(..))
#endif
import Data.Proxy                                          (Proxy(..))
import Data.Ratio                                          ((%), numerator, denominator)
import Data.Semigroup                                      (Semigroup(..), Product(..))
import qualified Data.Vector as V
import qualified Data.Vector.Mutable as MV
import Data.Vector                                         (Vector, (!))
import GHC.TypeNats.Compat                                 (Nat, SomeNat(..), natVal, someNatVal)
import Numeric.Natural                                     (Natural)

import Math.NumberTheory.ArithmeticFunctions               (totient)
import Math.NumberTheory.Moduli.Chinese
import Math.NumberTheory.Moduli.Class                      (KnownNat, Mod, getVal)
import Math.NumberTheory.Moduli.Internal                   (isPrimitiveRoot', discreteLogarithmPP)
import Math.NumberTheory.Moduli.Multiplicative             (MultMod(..), isMultElement)
import Math.NumberTheory.Moduli.Singleton                  (Some(..), cyclicGroupFromFactors)
import Math.NumberTheory.Powers.Modular                    (powMod)
import Math.NumberTheory.Primes                            (Prime(..), UniqueFactorisation, factorise, nextPrime)
import Math.NumberTheory.Utils.FromIntegral                (wordToInt)
import Math.NumberTheory.Utils

-- | A Dirichlet character mod \(n\) is a group homomorphism from \((\mathbb{Z}/n\mathbb{Z})^*\)
-- to \(\mathbb{C}^*\), represented abstractly by `DirichletCharacter`. In particular, they take
-- values at roots of unity and can be evaluated using `eval`.
-- A Dirichlet character can be extended to a completely multiplicative function on \(\mathbb{Z}\)
-- by assigning the value 0 for \(a\) sharing a common factor with \(n\), using `evalGeneral`.
--
-- There are finitely many possible Dirichlet characters for a given modulus, in particular there
-- are \(\phi(n)\) characters modulo \(n\), where \(\phi\) refers to Euler's `totient` function.
-- This gives rise to `Enum` and `Bounded` instances.
newtype DirichletCharacter (n :: Nat) = Generated [DirichletFactor]

-- | The group (Z/nZ)^* decomposes to a product (Z/2^k0 Z)^* x (Z/p1^k1 Z)^* x ... x (Z/pi^ki Z)^*
-- where n = 2^k0 p1^k1 ... pi^ki, and the pj are odd primes, k0 possibly 0. Thus, a group
-- homomorphism from (Z/nZ)^* is characterised by group homomorphisms from each of these factor
-- groups. Furthermore, for odd p, we have (Z/p^k Z)^* isomorphic to Z / p^(k-1)*(p-1) Z, an
-- additive group, where an isomorphism is specified by a choice of primitive root.
-- Similarly, for k >= 2, (Z/2^k Z)^* is isomorphic to Z/2Z * (Z / 2^(k-2) Z) (and for k < 2
-- it is trivial).  (See `lambda` for this isomorphism).
-- Thus, to specify a Dirichlet character, it suffices to specify the value of generators
-- of each of these cyclic groups, when primitive roots are given. This data is given by a
-- DirichletFactor.
-- We have the invariant that the factors must be given in strictly increasing order, and the
-- generator is as given by `generator`, and are each non-trivial. These conditions are verified
-- using `validChar`.
data DirichletFactor = OddPrime { _getPrime :: Prime Natural
                                , _getPower :: Word
                                , _getGenerator :: Natural
                                , _getValue :: RootOfUnity
                                }
                     | TwoPower { _getPower2 :: Int -- this ought to be Word, but many applications
                                                    -- needed to use wordToInt, so Int is cleaner
                                                    -- Required to be >= 2
                                , _getFirstValue :: RootOfUnity
                                , _getSecondValue :: RootOfUnity
                                }
                     | Two

instance Eq (DirichletCharacter n) where
  Generated a == Generated b = a == b

instance Eq DirichletFactor where
  TwoPower _ x1 x2 == TwoPower _ y1 y2 = x1 == y1 && x2 == y2
  OddPrime _ _ _ x == OddPrime _ _ _ y = x == y
  Two              == Two              = True
  _ == _ = False

-- | A representation of <https://en.wikipedia.org/wiki/Root_of_unity roots of unity>: complex
-- numbers \(z\) for which there is \(n\) such that \(z^n=1\).
newtype RootOfUnity =
  RootOfUnity { -- | Every root of unity can be expressed as \(e^{2 \pi i q}\) for some
                -- rational \(q\) satisfying \(0 \leq q < 1\), this function extracts it.
                fromRootOfUnity :: Rational }
  deriving (Eq)

instance Show RootOfUnity where
  show (RootOfUnity q)
    | n == 0    = "1"
    | d == 1    = "-1"
    | n == 1    = "e^(πi/" ++ show d ++ ")"
    | otherwise = "e^(" ++ show n ++ "πi/" ++ show d ++ ")"
    where n = numerator (2*q)
          d = denominator (2*q)

-- | Given a rational \(q\), produce the root of unity \(e^{2 \pi i q}\).
toRootOfUnity :: Rational -> RootOfUnity
toRootOfUnity q = RootOfUnity ((n `rem` d) % d)
  where n = numerator q
        d = denominator q
        -- effectively q `mod` 1
  -- This smart constructor ensures that the rational is always in the range 0 <= q < 1.

-- | This Semigroup is in fact a group, so @'stimes'@ can be called with a negative first argument.
instance Semigroup RootOfUnity where
  RootOfUnity q1 <> RootOfUnity q2 = toRootOfUnity (q1 + q2)
  stimes k (RootOfUnity q) = toRootOfUnity (q * fromIntegral k)

instance Monoid RootOfUnity where
  mappend = (<>)
  mempty = RootOfUnity 0

-- | Convert a root of unity into an inexact complex number. Due to floating point inaccuracies,
-- it is recommended to avoid use of this until the end of a calculation. Alternatively, with
-- the [cyclotomic](http://hackage.haskell.org/package/cyclotomic-0.5.1) package, one can use
-- @[polarRat](https://hackage.haskell.org/package/cyclotomic-0.5.1/docs/Data-Complex-Cyclotomic.html#v:polarRat)
-- 1 . @'fromRootOfUnity' to convert to a cyclotomic number.
toComplex :: Floating a => RootOfUnity -> Complex a
toComplex (RootOfUnity t)
  | t == 1/2 = (-1) :+ 0
  | t == 1/4 = 0 :+ 1
  | t == 3/4 = 0 :+ (-1)
  | otherwise = cis . (2*pi*) . fromRational $ t

-- | For primes, define the canonical primitive root as the smallest such. For prime powers \(p^k\),
-- either the smallest primitive root \(g\) mod \(p\) works, or \(g+p\) works.
generator :: (Integral a, UniqueFactorisation a) => Prime a -> Word -> a
generator p k
  | k == 1 = modP
  | otherwise = if powMod modP (p'-1) (p'*p') == 1 then modP + p' else modP
  where p' = unPrime p
        modP = case cyclicGroupFromFactors [(p,k)] of
                 Just (Some cg) -> head $ filter (isPrimitiveRoot' cg) [2..p'-1]
                 _ -> error "illegal"

-- | Implement the function \(\lambda\) from page 5 of
-- https://www2.eecs.berkeley.edu/Pubs/TechRpts/1984/CSD-84-186.pdf
lambda :: Integer -> Int -> Integer
lambda x e = ((powMod x (2*modulus) largeMod - 1) `shiftR` (e+1)) .&. (modulus - 1)
  where modulus = bit (e-2)
        largeMod = bit (2*e - 1)

-- | For elements of the multiplicative group \((\mathbb{Z}/n\mathbb{Z})^*\), a Dirichlet
-- character evaluates to a root of unity.
eval :: DirichletCharacter n -> MultMod n -> RootOfUnity
eval (Generated ds) m = foldMap (evalFactor m') ds
  where m' = getVal $ multElement m

-- | Evaluate each factor of the Dirichlet character.
evalFactor :: Integer -> DirichletFactor -> RootOfUnity
evalFactor m =
  \case
    OddPrime (toInteger . unPrime -> p) k (toInteger -> a) b ->
      discreteLogarithmPP p k a (m `rem` p^k) `stimes` b
    TwoPower k s b -> (if testBit m 1 then s else mempty)
                   <> lambda (thingy k m) k `stimes` b
    Two -> mempty

thingy :: (Bits p, Num p) => Int -> p -> p
thingy k m = if testBit m 1
                then bit k - m'
                else m'
  where m' = m .&. (bit k - 1)

-- | A character can evaluate to a root of unity or zero: represented by @Nothing@.
evalGeneral :: KnownNat n => DirichletCharacter n -> Mod n -> OrZero RootOfUnity
evalGeneral chi t = case isMultElement t of
                      Nothing -> Zero
                      Just x -> NonZero $ eval chi x

-- | Give the principal character for this modulus: a principal character mod \(n\) is 1 for
-- \(a\) coprime to \(n\), and 0 otherwise.
principalChar :: KnownNat n => DirichletCharacter n
principalChar = minBound

mulChars :: DirichletCharacter n -> DirichletCharacter n -> DirichletCharacter n
mulChars (Generated x) (Generated y) = Generated (zipWith combine x y)
  where combine :: DirichletFactor -> DirichletFactor -> DirichletFactor
        combine Two Two = Two
        combine (OddPrime p k g n) (OddPrime _ _ _ m) =
          OddPrime p k g (n <> m)
        combine (TwoPower k a n) (TwoPower _ b m) =
          TwoPower k (a <> b) (n <> m)
        combine _ _ = error "internal error: malformed DirichletCharacter"

-- | This Semigroup is in fact a group, so @stimes@ can be called with a negative first argument.
instance Semigroup (DirichletCharacter n) where
  (<>) = mulChars
  stimes = stimesChar

instance KnownNat n => Monoid (DirichletCharacter n) where
  mempty = principalChar
  mappend = (<>)

stimesChar :: Integral a => a -> DirichletCharacter n -> DirichletCharacter n
stimesChar s (Generated xs) = Generated (map mult xs)
  where mult :: DirichletFactor -> DirichletFactor
        mult (OddPrime p k g n) = OddPrime p k g (s `stimes` n)
        mult (TwoPower k a b) = TwoPower k (s `stimes` a) (s `stimes` b)
        mult Two = Two

-- | We define `succ` and `pred` with more efficient implementations than
-- @`toEnum` . (+1) . `fromEnum`@.
instance KnownNat n => Enum (DirichletCharacter n) where
  toEnum = indexToChar . fromIntegral
  fromEnum = fromIntegral . characterNumber
  succ x = makeChar x (characterNumber x + 1)
  pred x = makeChar x (characterNumber x - 1)

  enumFromTo x y       = bulkMakeChars x [fromEnum x..fromEnum y]
  enumFrom x           = bulkMakeChars x [fromEnum x..]
  enumFromThenTo x y z = bulkMakeChars x [fromEnum x, fromEnum y..fromEnum z]
  enumFromThen x y     = bulkMakeChars x [fromEnum x, fromEnum y..]

instance KnownNat n => Bounded (DirichletCharacter n) where
  minBound = indexToChar 0
  maxBound = indexToChar (totient n - 1)
    where n = natVal (Proxy :: Proxy n)

-- | We have a (non-canonical) enumeration of dirichlet characters.
characterNumber :: DirichletCharacter n -> Integer
characterNumber (Generated y) = foldl' go 0 y
  where go x (OddPrime p k _ a) = x * m + numerator (fromRootOfUnity a * fromIntegral m)
          where p' = fromIntegral (unPrime p)
                m = p'^(k-1)*(p'-1)
        go x (TwoPower k a b)   = x' * 2 + numerator (fromRootOfUnity a * 2)
          where m = bit (k-2) :: Integer
                x' = x `shiftL` (k-2) + numerator (fromRootOfUnity b * fromIntegral m)
        go x Two = x

-- | Give the dirichlet character from its number.
-- Inverse of `characterNumber`.
indexToChar :: forall n. KnownNat n => Natural -> DirichletCharacter n
indexToChar = runIdentity . indicesToChars . Identity

-- | Give a collection of dirichlet characters from their numbers. This may be more efficient than
-- `indexToChar` for multiple characters, as it prevents some internal recalculations.
indicesToChars :: forall n f. (KnownNat n, Functor f) => f Natural -> f (DirichletCharacter n)
indicesToChars = fmap (Generated . unroll t . (`mod` m))
  where n = natVal (Proxy :: Proxy n)
        (Product m, t) = mkTemplate n

-- | List all characters for the modulus. This is preferred to using @[minBound..maxBound]@.
allChars :: forall n. KnownNat n => [DirichletCharacter n]
allChars = indicesToChars [0..m-1]
  where m = totient $ natVal (Proxy :: Proxy n)

-- | The same as `indexToChar`, but if we're given a character we can create others more efficiently.
makeChar :: Integral a => DirichletCharacter n -> a -> DirichletCharacter n
makeChar x = runIdentity . bulkMakeChars x . Identity

-- | Use one character to make many more: better than indicesToChars since it avoids recalculating
-- some primitive roots
bulkMakeChars :: (Integral a, Functor f) => DirichletCharacter n -> f a -> f (DirichletCharacter n)
bulkMakeChars x = fmap (Generated . unroll t . (`mod` m) . fromIntegral)
  where (Product m, t) = templateFromCharacter x

-- We assign each natural a unique Template, which can be decorated (eg in `unroll`) to
-- form a DirichletCharacter. A Template effectively holds the information carried around
-- in a DirichletFactor which depends only on the modulus of the character.
data Template = OddTemplate { _getPrime'     :: Prime Natural
                            , _getPower'     :: Word
                            , _getGenerator' :: !Natural
                            , _getModulus'   :: !Natural
                            }
              | TwoPTemplate { _getPower2'    :: Int
                             , _getModulus'   :: !Natural
                             } -- the modulus is derivable from the other values, but calculation
                               -- may be expensive, so we pre-calculate it
                               -- morally getModulus should be a prefactored but seems to be
                               -- pointless here
              | TwoTemplate

templateFromCharacter :: DirichletCharacter n -> (Product Natural, [Template])
templateFromCharacter (Generated t) = traverse go t
  where go (OddPrime p k g _) = (Product m, OddTemplate p k g m)
          where p' = unPrime p
                m = p'^(k-1)*(p'-1)
        go (TwoPower k _ _) = (Product (2*m), TwoPTemplate k m)
          where m = bit (k-2)
        go Two = (Product 1, TwoTemplate)

mkTemplate :: Natural -> (Product Natural, [Template])
mkTemplate = go . sort . factorise
  where go :: [(Prime Natural, Word)] -> (Product Natural, [Template])
        go ((unPrime -> 2, 1): xs) = (Product 1, [TwoTemplate]) <> traverse odds xs
        go ((unPrime -> 2, wordToInt -> k): xs) = (Product (2*m), [TwoPTemplate k m]) <> traverse odds xs
          where m = bit (k-2)
        go xs = traverse odds xs
        odds :: (Prime Natural, Word) -> (Product Natural, Template)
        odds (p, k) = (Product m, OddTemplate p k (generator p k) m)
          where p' = unPrime p
                m = p'^(k-1)*(p'-1)

-- the validity of the producted dirichletfactor list here requires the template to be valid
unroll :: [Template] -> Natural -> [DirichletFactor]
unroll t m = snd (mapAccumL func m t)
  where func :: Natural -> Template -> (Natural, DirichletFactor)
        func a (OddTemplate p k g n) = (a1, OddPrime p k g (toRootOfUnity $ (toInteger a2) % (toInteger n)))
          where (a1,a2) = quotRem a n
        func a (TwoPTemplate k n) = (b1, TwoPower k (toRootOfUnity $ (toInteger a2) % 2) (toRootOfUnity $ (toInteger b2) % (toInteger n)))
          where (a1,a2) = quotRem a 2
                (b1,b2) = quotRem a1 n
        func a TwoTemplate = (a, Two)

-- | Test if a given Dirichlet character is prinicpal for its modulus: a principal character mod
-- \(n\) is 1 for \(a\) coprime to \(n\), and 0 otherwise.
isPrincipal :: DirichletCharacter n -> Bool
isPrincipal chi = characterNumber chi == 0

-- | Induce a Dirichlet character to a higher modulus. If \(d \mid n\), then \(a \bmod{n}\) can be
-- reduced to \(a \bmod{d}\). Thus, the multiplicative function on \(\mathbb{Z}/d\mathbb{Z}\)
-- induces a multiplicative function on \(\mathbb{Z}/n\mathbb{Z}\).
--
-- >>> :set -XTypeApplications
-- >>> chi = indexToChar 5 :: DirichletCharacter 45
-- >>> chi2 = induced @135 chi
-- >>> :t chi2
-- Maybe (DirichletCharacter 135)
induced :: forall n d. (KnownNat d, KnownNat n) => DirichletCharacter d -> Maybe (DirichletCharacter n)
induced (Generated start) = if n `rem` d == 0
                            then Just (Generated (combine (snd $ mkTemplate n) start))
                            else Nothing
  where n = natVal (Proxy :: Proxy n)
        d = natVal (Proxy :: Proxy d)
        combine :: [Template] -> [DirichletFactor] -> [DirichletFactor]
        combine [] _ = []
        combine ts [] = map newFactor ts
        combine (t:xs) (y:ys) = case (t,y) of
                                  (TwoTemplate, Two) -> Two: combine xs ys
                                  (TwoTemplate, _) -> Two: combine xs (y:ys)
                                  (TwoPTemplate k _, Two) -> TwoPower k mempty mempty: combine xs ys
                                  (TwoPTemplate k _, TwoPower _ a b) -> TwoPower k a b: combine xs ys
                                  (TwoPTemplate k _, _) -> TwoPower k mempty mempty: combine xs (y:ys)
                                  (OddTemplate p k _ _, OddPrime q _ g a) | p == q -> OddPrime p k g a: combine xs ys
                                  (OddTemplate p k g _, OddPrime q _ _ _) | p < q -> OddPrime p k g mempty: combine xs (y:ys)
                                  _ -> error "internal error in induced: please report this as a bug"
        newFactor :: Template -> DirichletFactor
        newFactor TwoTemplate = Two
        newFactor (TwoPTemplate k _) = TwoPower k mempty mempty
        newFactor (OddTemplate p k g _) = OddPrime p k g mempty
        -- rest (p,k) = OddPrime p k (generator p k) mempty

-- | The <https://en.wikipedia.org/wiki/Jacobi_symbol Jacobi symbol> gives a real Dirichlet
-- character for odd moduli.
jacobiCharacter :: forall n. KnownNat n => Maybe (RealCharacter n)
jacobiCharacter = if odd n
                     then Just $ RealChar $ Generated $ map go $ snd $ mkTemplate n
                     else Nothing
  where n = natVal (Proxy :: Proxy n)
        go :: Template -> DirichletFactor
        go (OddTemplate p k g _) = OddPrime p k g $ toRootOfUnity ((toInteger k) % 2)
          -- jacobi symbol of a primitive root mod p over p is always -1
        go _ = error "internal error in jacobiCharacter: please report this as a bug"
          -- every factor of n should be odd

-- | A Dirichlet character is real if it is real-valued.
newtype RealCharacter n = RealChar { -- | Extract the character itself from a `RealCharacter`.
                                     getRealChar :: DirichletCharacter n
                                   }
                                   deriving Eq

-- | Test if a given `DirichletCharacter` is real, and if so give a `RealCharacter`.
isRealCharacter :: DirichletCharacter n -> Maybe (RealCharacter n)
isRealCharacter t@(Generated xs) = if all real xs then Just (RealChar t) else Nothing
  where real :: DirichletFactor -> Bool
        real (OddPrime _ _ _ a) = a <> a == mempty
        real (TwoPower _ _ b) = b <> b == mempty
        real Two = True

-- TODO: it should be possible to calculate this without eval/evalGeneral
-- and thus avoid using discrete log calculations: consider the order of m
-- inside each of the factor groups?
-- | Evaluate a real Dirichlet character, which can only take values \(-1,0,1\).
toRealFunction :: KnownNat n => RealCharacter n -> Mod n -> Int
toRealFunction (RealChar chi) m = case evalGeneral chi m of
                                    Zero -> 0
                                    NonZero t | t == mempty -> 1
                                    NonZero t | t == RootOfUnity (1 % 2) -> -1
                                    _ -> error "internal error in toRealFunction: please report this as a bug"
                                      -- A real character should not be able to evaluate to
                                      -- anything other than {-1,0,1}, so should not reach this branch

-- | Test if the internal DirichletCharacter structure is valid.
validChar :: forall n. KnownNat n => DirichletCharacter n -> Bool
validChar (Generated xs) = correctDecomposition && all correctPrimitiveRoot xs && all validValued xs
  where correctDecomposition = sort (factorise n) == map getPP xs
        getPP (TwoPower k _ _) = (two, fromIntegral k)
        getPP (OddPrime p k _ _) = (p, k)
        getPP Two = (two,1)
        correctPrimitiveRoot (OddPrime p k g _) = g == generator p k
        correctPrimitiveRoot _ = True
        validValued (TwoPower k a b) = a <> a == mempty && (bit (k-2) :: Integer) `stimes` b == mempty
        validValued (OddPrime (unPrime -> p) k _ a) = (p^(k-1)*(p-1)) `stimes` a == mempty
        validValued Two = True
        n = natVal (Proxy :: Proxy n)
        two = nextPrime 2

-- | Get the order of the Dirichlet Character.
orderChar :: DirichletCharacter n -> Integer
orderChar (Generated xs) = foldl' lcm 1 $ map orderFactor xs
  where orderFactor (TwoPower _ (RootOfUnity a) (RootOfUnity b)) = denominator a `lcm` denominator b
        orderFactor (OddPrime _ _ _ (RootOfUnity a)) = denominator a
        orderFactor Two = 1

-- | Test if a Dirichlet character is <https://en.wikipedia.org/wiki/Dirichlet_character#Primitive_characters_and_conductor primitive>.
isPrimitive :: DirichletCharacter n -> Maybe (PrimitiveCharacter n)
isPrimitive t@(Generated xs) = if all primitive xs then Just (PrimitiveCharacter t) else Nothing
  where primitive :: DirichletFactor -> Bool
        primitive Two = False
        -- for odd p, we're testing if phi(p^(k-1)) `stimes` a is 1, since this means the
        -- character can come from some the smaller modulus p^(k-1)
        primitive (OddPrime _ 1 _ a) = a /= mempty
        primitive (OddPrime (unPrime -> p) k _ a) = (p^(k-2)*(p-1)) `stimes` a /= mempty
        primitive (TwoPower 2 a _) = a /= mempty
        primitive (TwoPower k _ b) = (bit (k-3) :: Integer) `stimes` b /= mempty

-- | A Dirichlet character is primitive if cannot be 'induced' from any character with
-- strictly smaller modulus.
newtype PrimitiveCharacter n = PrimitiveCharacter { -- | Extract the character itself from a `PrimitiveCharacter`.
                                                    getPrimitiveChar :: DirichletCharacter n
                                                    }
                                                    deriving Eq

data WithNat (a :: Nat -> *) where
  WithNat :: KnownNat m => a m -> WithNat a

-- | This function also provides access to the new modulus on type level, with a KnownNat instance
makePrimitive :: DirichletCharacter n -> WithNat PrimitiveCharacter
makePrimitive (Generated xs) =
  case someNatVal (product mods) of
    SomeNat (Proxy :: Proxy m) -> WithNat (PrimitiveCharacter (Generated ys) :: PrimitiveCharacter m)
  where (mods,ys) = unzip (mapMaybe prim xs)
        prim :: DirichletFactor -> Maybe (Natural, DirichletFactor)
        prim Two = Nothing
        prim (OddPrime p' k g a) = case find works options of
                                     Nothing -> error "invalid character"
                                     Just (0,_) -> Nothing
                                     Just (i,_) -> Just (p^i, OddPrime p' i g a)
          where options = (0,1): [(i,p^(i-1)*(p-1)) | i <- [1..k]]
                works (_,phi) = phi `stimes` a == mempty
                p = unPrime p'
        prim (TwoPower k a b) = case find worksb options of
                                  Nothing -> error "invalid character"
                                  Just (2,_) | a == mempty -> Nothing
                                  Just (i,_) -> Just (bit i :: Natural, TwoPower i a b)
          where options = [(i, bit (i-2) :: Natural) | i <- [2..k]]
                worksb (_,phi) = phi `stimes` b == mempty

#if !MIN_VERSION_base(4,12,0)
newtype Ap f a = Ap { getAp :: f a }
  deriving (Eq, Functor, Applicative, Monad)

instance (Applicative f, Semigroup a) => Semigroup (Ap f a) where
  (<>) = liftA2 (<>)

instance (Applicative f, Semigroup a, Monoid a) => Monoid (Ap f a) where
  mempty = pure mempty
  mappend = (<>)
#endif

-- | Similar to Maybe, but with different Semigroup and Monoid instances.
type OrZero a = Ap Maybe a
pattern Zero :: OrZero a
pattern Zero = Ap Nothing

pattern NonZero :: a -> OrZero a
pattern NonZero x = Ap (Just x)
{-# COMPLETE Zero, NonZero #-}

-- | Interpret an `OrZero` as a number, taking the `Zero` case to be 0.
orZeroToNum :: Num a => (b -> a) -> OrZero b -> a
orZeroToNum _ Zero = 0
orZeroToNum f (NonZero x) = f x

-- | In general, evaluating a DirichletCharacter at a point involves solving the discrete logarithm
-- problem, which can be hard: the implementations here are around O(sqrt n).
-- However, evaluating a dirichlet character at every point amounts to solving the discrete
-- logarithm problem at every point also, which can be done together in O(n) time, better than
-- using a complex algorithm at each point separately. Thus, if a large number of evaluations
-- of a dirichlet character are required, `evalAll` will be better than `evalGeneral`, since
-- computations can be shared.
evalAll :: forall n. KnownNat n => DirichletCharacter n -> Vector (OrZero RootOfUnity)
evalAll (Generated xs) = V.generate (fromIntegral n) func
  where n = natVal (Proxy :: Proxy n)
        vectors = map mkVector xs
        func :: Int -> OrZero RootOfUnity
        func m = foldMap go vectors
          where go :: (Int, Vector (OrZero RootOfUnity)) -> OrZero RootOfUnity
                go (modulus,v) = v ! (m `mod` modulus)
        mkVector :: DirichletFactor -> (Int, Vector (OrZero RootOfUnity))
        mkVector Two = (2, V.fromList [Zero, mempty])
        mkVector (OddPrime p k (fromIntegral -> g) a) = (modulus, w)
          where
            p' = unPrime p
            modulus = fromIntegral (p'^k) :: Int
            w = V.create $ do
              v <- MV.replicate modulus Zero
              -- TODO: we're in the ST monad here anyway, could be better to use STRefs to manage
              -- this loop, the current implementation probably doesn't fuse well
              let powers = iterateMaybe go (1,mempty)
                  go (m,x) = if m' > 1
                                then Just (m', x<>a)
                                else Nothing
                    where m' = m*g `mod` modulus
              for_ powers $ \(m,x) -> MV.unsafeWrite v m (NonZero x)
              -- don't bother with bounds check since m was reduced mod p^k
              return v
        -- for powers of two we use lambda directly instead, since the generators of the cyclic
        -- groups aren't obvious; it's possible to get them though:
        -- 5^(lambda(5)^{-1} mod 2^(p-2)) mod 2^p
        mkVector (TwoPower k a b) = (modulus, w)
          where
            modulus = bit k
            w = V.generate modulus f
            f m
              | even m = Zero
              | otherwise = NonZero ((if testBit m 1 then a else mempty) <> lambda (toInteger m'') k `stimes` b)
              where m'' = thingy k m

-- somewhere between unfoldr and iterate
iterateMaybe :: (a -> Maybe a) -> a -> [a]
iterateMaybe f x = unfoldr (fmap (\t -> (t, f t))) (Just x)

-- | Attempt to construct a character from its table of values.
-- An inverse to `evalAll`, defined only on its image.
fromTable :: forall n. KnownNat n => Vector (OrZero RootOfUnity) -> Maybe (DirichletCharacter n)
fromTable v = if length v == fromIntegral n
                 then Generated <$> traverse makeFactor tmpl >>= check
                 else Nothing
  where n = natVal (Proxy :: Proxy n)
        n' = fromIntegral n :: Integer
        tmpl = snd (mkTemplate n)
        check :: DirichletCharacter n -> Maybe (DirichletCharacter n)
        check chi = if evalAll chi == v then Just chi else Nothing
        makeFactor :: Template -> Maybe DirichletFactor
        makeFactor TwoTemplate = Just Two
        makeFactor (TwoPTemplate k _) = TwoPower k <$> getValue (-1,bit k) <*> getValue (exp4 k, bit k)
        makeFactor (OddTemplate p k g _) = OddPrime p k g <$> getValue (toInteger g, toInteger (unPrime p)^k)
        getValue :: (Integer,Integer) -> Maybe RootOfUnity
        getValue (g,m) = getAp (v ! fromInteger (fromJust (chineseCoprime (g,m) (1,n' `quot` m)) `mod` n'))

exp4terms :: [Rational]
exp4terms = [4^k % product [1..k] | k <- [0..]]

-- For reasons that aren't clear to me, `exp4` gives the inverse of 1 under lambda, so it gives the generator
-- This is the same as https://oeis.org/A320814
-- In particular, lambda (exp4 n) n == 1 (for n >= 3)
-- I've verified this for 3 <= n <= 2000, so the reasoning in fromTable should be accurate for moduli below 2^2000
exp4 :: Int -> Integer
exp4 n = (`mod` bit n) $ sum $ map (`mod` bit n) $ map (\q -> numerator q * fromMaybe (error "error in exp4") (recipMod (denominator q) (bit n))) $ take n $ exp4terms

max_stars_repo_name: b-mehta/arithmoi

max_issues_repo_name: b-mehta/arithmoi

max_forks_repo_name: b-mehta/arithmoi


text:
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE Strict #-}
{-# LANGUAGE FlexibleContexts #-}

import CPVO.IO
import CPVO.Numeric
import CPVO.IO.DOS
import CPVO.IO.Fortran
import CPVO.IO.Reader.Ecalj.DOS
import CPVO.IO.Reader.Ecalj.Common
import CPVO.IO.Reader.Ecalj.MMOM
import CPVO.IO.Plot.Gnuplot.DOS
import CPVO.IO.Plot.Gnuplot.Common
import CPVO.IO.Plot.PDOS

import qualified Language.C.Inline as C

import qualified Control.Foldl as Fold
import System.Environment (getArgs)
import qualified Data.Text as T
import qualified Data.Text.IO as T
import qualified Data.Text.Read as T
import qualified Data.Text.Format as T
import Text.Printf
import Data.List.Split
import Data.List
import Data.Maybe
import Data.Either
-------------------------

import Numeric.LinearAlgebra
import Numeric.LinearAlgebra.Devel (readMatrix)
import Numeric.LinearAlgebra.Data hiding (find)
import Data.Char
import System.Process
import System.IO
-------------------------
import Text.PrettyPrint.Boxes hiding ((<>),cols,rows)
import qualified Text.PrettyPrint.Boxes as TB
import Data.List
-------------------------
import Text.Pandoc
import Control.Monad ((<=<),forM)

rydberg=1/13.605

testArgs = concat [ [ "GWtable01.tex"
                   , "Total and partial DOS at *E*~F~ in states/eV/unit-cell, QSGW method."
                   , "|*N*$_\\uparrow$(*E*~F~)|*N*$_\\downarrow$(*E*~F~)|*N*(*E*~F~)"
                     ]
                   , words "@{}lSSS@{} -9:6 25 T o flipSpin nico2o4 extendedNiCo2O4.normal/nico2o4.6G10 Ni:Ni#3d:6:7:9:8:10 Co:Co#3d:6:7:8:9:10 O:O#2p:3:4:5"
                   ]

main = do
  putStrLn "=======end Test========="

main1 = do
--    allArgs <- getArgs
    (invStat, ymax, xmin, xmax, ctrlAtoms, uniqAtoms, ctrlAtomicAOs,jdTable, jdHeads, foldernya, tailer, colAlign, texFile) <- readHeaderData testArgs
    totalDOS <- readTotalDOSText tailer foldernya
    let resTotY0 = map (\i -> getY0 $ totalDOS ¿ [0,i]) $ flipBy invStat [1,2]
    putStrLn "===========NF total DOS========="
    putStrLn $ show resTotY0
    pdosAtAll <- readPDOS invStat tailer foldernya ctrlAtomicAOs
    let y0AtAll = map (\(mp,b) -> (getY0 mp,b)) pdosAtAll
    let rs' = rendertable
              $ (:) jdHeads
              $ (:)  ("Total" : (map (showDouble 3) $ processCols $ map ((* rydberg) . last . toList) $ resTotY0))
              $ map (\a -> (((fst $ head a):(map (showDouble 3) $ processCols $ sumIt [0,0] $ map snd a))))
              $ groupBy (\a b -> fst a == fst b)
              $ (\[us,ds] -> zipWith (\(yu,(_,(j,_))) (yd,_) -> (j, [yu,yd])) us ds )
              $ groupBy (\(_,(s,_)) (_,(s',_)) ->  s == s')
              $ map (\(mp,b) -> ((* rydberg) $ last $ toList $ getY0 mp, b)) pdosAtAll
    let rs = unlines [ rs' , jdTable ]
    res' <- markdownToTex rs
    let res = T.replace "\\}" "}"
            $ T.replace "\\{" "{"
            $ T.pack $ unlines  [ "\\begin{longtable}[]{" ++ colAlign ++ "}"
                                , unlines $ tail $ lines $ T.unpack res'
                                ]
    T.writeFile texFile res
    putStrLn $ texFile ++ "=============done:t1dosEf.hs=========="
      where
        processCols [u,d] = [u,d,u+d]
        sumIt x [] = x
        sumIt [u0,d0] ([up,dn]:xs) = sumIt [u0+up,d0+dn] xs

max_stars_repo_name: hasanalrasyid/cpvoh

max_issues_repo_name: hasanalrasyid/cpvoh

max_forks_repo_name: hasanalrasyid/cpvoh


text:
{-|
 Module      : Data.Boltzmann.System.Utils
 Description : Various routines for combinatorial systems.
 Copyright   : (c) Maciej Bendkowski, 2017-2020

 License     : BSD3
 Maintainer  : maciej.bendkowski@tcs.uj.edu.pl
 Stability   : experimental

 The following module implements subroutines for combinatorial systems based
 on the excellent paper of Carine Pivoteau, Bruno Salvy, and Michèle Soria:

 Algorithms for combinatorial structures: Well-founded systems and Newton iterations.
 Journal of Combinatorial Theory, Series A 119 (2012) p. 1711–1773.
 -}
module Data.Boltzmann.System.Utils
    ( isEmptyAtZero
    , zeroCoordinates
    , wellFoundedAtZero
    , polynomial
    ) where

import           Prelude hiding ((<>))

import           Data.Map (Map)
import qualified Data.Map.Strict as M
import           Numeric.LinearAlgebra hiding (size)

import           Data.Boltzmann.System

-- | Non-strict multiplication.
(.*) :: (Num a, Eq a) => a -> a -> a
0 .* _ = 0
_ .* 0 = 0
1 .* x = x
x .* 1 = x
x .* y = x * y

-- | Evaluates the given system Y = H(Z, Y) at coordinates
--   Z = 0 and Y = 0 and checks if the outcome system is empty,
--   i.e. H(0, 0) = 0 or not.
isEmptyAtZero :: (Eq a, Num a) => System a -> Bool
isEmptyAtZero sys = all isZeroType ts
    where ts = M.toList $ defs sys

          isZeroType (_, consL) = all isZeroCons consL
          -- Note: A given constructor can give a non-zero coordinate
          --       iff it has no arguments and no positive weight. List
          --       arguments are irrelevant as they yield powers of unity.
          isZeroCons cons = not $ weight cons == 0 && null (args cons)

evalAtZero :: Num a => Arg -> a
evalAtZero (Type _) = 0
evalAtZero (List _) = 1

evalAtZeroL :: (Eq a, Num a) => [Arg] -> a
evalAtZeroL xs = foldl (.*) 1 $ map evalAtZero xs

derivTypeAtZero :: (Eq a, Eq b, Num a, Num b)
                => String -> [Cons b] -> a

derivTypeAtZero dt consL =
    sum $ map (derivConsAtZero dt) consL

derivConsAtZero :: (Eq a, Eq b, Num a, Num b)
                => String -> Cons a -> b

derivConsAtZero dt cons
    | weight cons /= 0 = 0
    | otherwise = derivConsAtZeroL dt (args cons)

derivConsAtZeroL :: (Eq a, Num a)
                 => String -> [Arg] -> a

derivConsAtZeroL _ [] = 1
derivConsAtZeroL dt (x:xs) = lhs + rhs
    where lhs = evalAtZero x .* derivConsAtZeroL dt xs
          rhs = derivAtZero dt x .* evalAtZeroL xs

derivAtZero :: Num a
            => String -> Arg -> a

derivAtZero dt arg
    | argName arg /= dt = 0
    | otherwise = 1

-- | Computes the Jacobian matrix of the given system
--   and evaluates it at zero-valued coordinates.
jacobian :: (Eq a, Num a) => System a -> Matrix Double
jacobian sys = (n><n) [idx i j | i <- [0..n-1], j <- [0..n-1]]
    where n = size sys
          ts = defs sys
          idx i j =
              let typ = snd $ M.elemAt i ts
                  dt  = fst $ M.elemAt j ts
                  in derivTypeAtZero dt typ

-- | Checks whether the given matrix is zero or not.
isZero :: Matrix Double -> Bool
isZero mat = all (== 0) $ concat (toLists mat)

-- | Checks whether the given Jacobian matrix nilpotent is or not.
isNilpotent :: Matrix Double -> Bool
isNilpotent mat = isZero $ power mat m
    where m = rows mat

-- | Squares the given matrix.
square :: Matrix Double -> Matrix Double
square m = m <> m

-- | Fast matrix exponentiation.
power :: Integral a
      => Matrix Double -> a -> Matrix Double

power m 1 = m
power m n
  | odd n = m <> square (power m $ n-1)
  | otherwise = square (power m $ n `div` 2)

data Coordinate = Zero | Z deriving (Eq)

-- | Maps coordinates into booleans.
coord :: Coordinate -> Bool
coord Zero = True
coord Z    = False

-- | Gives the constructor list of a given type.
constrList :: System a -> String -> [Cons a]
constrList sys = (M.!) (defs sys)

-- | Detects whether the given system, satisfying H(0,0) = 0
--   and having a nilpotent Jacobian matrix, admits zero coordinates
--   in its solution. See also the 0-coord subroutine of Pivoteau et al.
zeroCoordinates :: System a -> Bool
zeroCoordinates sys = zeroCoordinates' sys m f
    where f = M.fromSet (const Zero) $ types sys
          m = size sys

zeroCoordinates' :: (Eq b, Num b)
                 => System a -> b -> Map String Coordinate -> Bool

zeroCoordinates' _ 0 f = elem Zero $ M.elems f
zeroCoordinates' sys m f = zeroCoordinates' sys (m-1) f'
    where xs = map (\k -> (k, mapF $ constrList sys k)) (M.keys f)
          f' = M.fromAscList xs

          mapF consL
              | all (coord . mapF') consL = Zero
              | otherwise                 = Z

          mapF' cons
              | any (coord . mapF'') (args cons) = Zero
              | otherwise                        = Z

          mapF'' (List _) = Z
          mapF'' (Type t) =
              case t `M.lookup` f of
                  Just Zero -> Zero
                  _         -> Z

-- | Checks whether the given system is well-founded at zero. Note that
--   the input system is assumed to satisfy H(0,0) = 0. See also the
--   isWellFoundedAt0 subroutine of Pivoteau et al.
wellFoundedAtZero :: (Eq a, Num a)
                  => System a -> Bool

wellFoundedAtZero sys
    | isNilpotent (jacobian sys) = not $ zeroCoordinates sys
    | otherwise                  = False

-- | Yields an infinite stream of successive evaluation iterations
--   in form of Y[m+1] = H(Z, Y[m]); starting with the given Y[0] and z.
evals :: System Int -> Vector Double -> Double -> [Vector Double]
evals sys ys z = ys : evals sys (eval sys ys z) z

-- | Checks whether the given system, assumed to satisfy H(0,0) = 0
--   and having a nilpotent Jacobian, encodes an implicit polynomial
--   species. See also the isPolynomial subroutine of Pivoteau et al.
polynomial :: System Int -> Bool
polynomial sys = hs !! m  == hs !! (m+1)
    where m   = size sys
          hs  = evals sys vec 1 -- note: numerical evaluation.
          vec = vector $ replicate m 0

max_stars_repo_name: timsears/boltzmann-brain

max_issues_repo_name: timsears/boltzmann-brain

max_forks_repo_name: timsears/boltzmann-brain


text:
{-# LANGUAGE CPP             #-}
{-# LANGUAGE TemplateHaskell #-}

{-|
Module      : Main
Description : Main module running the GMM
Copyright   : (c) Julian Kopka Larsen, 2015
Stability   : experimental

-}
module Main (
    main
) where

import           Control.Monad          (forM, unless)
import           Data.List              (transpose, (\\))
import           FileParsing
import           GHC.Arr (range)
import           Graphics.EasyPlot
import           Numeric.LinearAlgebra
-- import           Numeric.Statistics.PCA
import           System.CPUTime
import           System.Exit            (exitFailure)
import           System.IO              (readFile)
import           Text.Printf
import           System.Random (mkStdGen)
import Debug.Trace

import           MCMC
import           Partition
import           Math
import           Distributions (lNormW, dlNormW,lnormalInvWishartSS)

list2Touple (a:b:_) = (a,b)
list2Touple _ = (0,0)

toTouples :: Num a => [[a]] -> [(a,a)]
toTouples = map list2Touple

--to2D = toTouples . toLists
plotClusters :: X -> Partition -> [Graph2D Double Double]
plotClusters x p = map toplot $ range (0,((k p)-1))
                      where toplot i = Data2D [Title "d"] [] (toTouples $ map toList $ select x p i)

qtr x = trace ("value: " ++ show x) x

data State = State (Partition,[Sstat]) deriving Show

instance Sampleable State X where
    condMove x g (State (p,s)) = State (newPar, update)
                  where newPar = move g p
                        update = ss' x p s newPar

    llikelihood _ (State (_,s)) = sum $ map (\(i,m,v) -> lnormalInvWishartSS i m v) s
    llikDiff _ (State (_,s')) (State (_,s))
                = f s' - f s
                where f s = sum $ map (\(i,m,v) -> lnormalInvWishartSS i m v) s

ss :: X -> Partition -> [Sstat]
ss x part = zip3 csizes scatters means
            where csizes = componentSizes part
                  scatters = map scatterMatrix xs
                  means = map meanv xs
                  xs = groups x part

ss' :: X -> Partition -> [Sstat] -> Partition -> [Sstat]
ss' x par s npar = map go $ zip3 (p par) s (p npar)
        where go (c1, s, c2)
               | c1 == c2 = s
               | otherwise = updateSstat x c1 s c2

ssComponent :: X -> [Int] -> Sstat
ssComponent x []= (0, konst 0 (d,d), konst 0 d)
                where  d = dim $ head x
ssComponent x c = (length c, scatterMatrix xs, meanv xs)
                where  xs = group x c

updateSstat :: X -> [Int] -> Sstat -> [Int] -> Sstat
updateSstat x c1 s c2 = (s <-> removed) <+> added
                where removed =  (ssComponent x (c1\\c2))
                      added =   (ssComponent x (c2\\c1))

gmmElement :: X -> Int -> Int -> Int -> Partition
gmmElement x seed k num = p
          where State (p,_) = getElement gen x start num
                gen = mkStdGen seed
                start = State (startP, ss x startP)
                startP = naiveFromNk n k
                n = length x

time :: IO t -> IO t
time a = do
    start <- getCPUTime
    v <- a
    end   <- getCPUTime
    let diff = (fromIntegral (end - start)) / (10^12)
    printf "Computation time: %0.3f sec\n" (diff :: Double)
    return v

-- Main
main = do
    --contents <- readFile "../../Data/AccidentData.csv"
    --contents <- readFile "../../Data/2Clusters.csv"
    --contents <- readFile "../../Data/synthetic.6clust.csv"
    contents <- readFile "../../Data/synth2c2d.csv"

    let num_Components = 2
        num_Samples = 4000
        stdData = p2NormList contents
        result = gmmElement stdData 1 num_Components
        p a = plot X11 $ plotClusters stdData a

    --Plot Data
    putStrLn "Starting..."
    time $ p (result num_Samples)
    putStrLn "Done."

max_stars_repo_name: juliankopkalarsen/FpStats

max_issues_repo_name: juliankopkalarsen/FpStats

max_forks_repo_name: juliankopkalarsen/FpStats


text:
{-# LANGUAGE CPP #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE UndecidableInstances #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE DeriveFunctor #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE DefaultSignatures #-}
{-# LANGUAGE Trustworthy #-}
{-# LANGUAGE TypeOperators #-}
{-# OPTIONS_GHC -fenable-rewrite-rules #-}
----------------------------------------------------------------------
-- |
-- Copyright   :  (c) Edward Kmett 2011-2014
-- License     :  BSD3
--
-- Maintainer  :  ekmett@gmail.com
-- Stability   :  experimental
--
-- Representable endofunctors over the category of Haskell types are
-- isomorphic to the reader monad and so inherit a very large number
-- of properties for free.
----------------------------------------------------------------------

module Data.Functor.Rep
  (
  -- * Representable Functors
    Representable(..)
  , tabulated
  -- * Wrapped representable functors
  , Co(..)
  -- * Default definitions
  -- ** Functor
  , fmapRep
  -- ** Distributive
  , distributeRep
  , collectRep
  -- ** Apply/Applicative
  , apRep
  , pureRep
  , liftR2
  , liftR3
  -- ** Bind/Monad
  , bindRep
  -- ** MonadFix
  , mfixRep
  -- ** MonadZip
  , mzipRep
  , mzipWithRep
  -- ** MonadReader
  , askRep
  , localRep
  -- ** Extend
  , duplicatedRep
  , extendedRep
  -- ** Comonad
  , duplicateRep
  , extendRep
  , extractRep
  -- ** Comonad, with user-specified monoid
  , duplicateRepBy
  , extendRepBy
  , extractRepBy
  -- ** WithIndex
  , imapRep
  , ifoldMapRep
  , itraverseRep

  -- ** Generics
  , GRep
  , gindex
  , gtabulate
  , WrappedRep(..)
  ) where

import Control.Applicative
import Control.Applicative.Backwards
import Control.Arrow ((&&&))
#if __GLASGOW_HASKELL__ >= 708
import Data.Coerce
#endif
import Control.Comonad
import Control.Comonad.Trans.Class
import Control.Comonad.Trans.Traced
import Control.Comonad.Cofree
import Control.Monad.Trans.Identity
import Control.Monad.Reader
#if MIN_VERSION_base(4,4,0)
import Data.Complex
#endif
import Data.Distributive
import Data.Foldable (Foldable(fold))
import Data.Functor.Bind
import Data.Functor.Identity
import Data.Functor.Compose
import Data.Functor.Extend
import Data.Functor.Product
import Data.Functor.Reverse
import qualified Data.Monoid as Monoid
import Data.Profunctor.Unsafe
import Data.Proxy
import Data.Sequence (Seq)
import qualified Data.Sequence as Seq
import Data.Semigroup hiding (Product)
import Data.Tagged
import Data.Traversable (Traversable(sequenceA))
import Data.Void
import GHC.Generics hiding (Rep)
import Prelude hiding (lookup)

-- | A 'Functor' @f@ is 'Representable' if 'tabulate' and 'index' witness an isomorphism to @(->) x@.
--
-- Every 'Distributive' 'Functor' is actually 'Representable'.
--
-- Every 'Representable' 'Functor' from Hask to Hask is a right adjoint.
--
-- @
-- 'tabulate' . 'index'  ≡ id
-- 'index' . 'tabulate'  ≡ id
-- 'tabulate' . 'return' ≡ 'return'
-- @

class Distributive f => Representable f where
  -- | If no definition is provided, this will default to 'GRep'.
  type Rep f :: *
  type Rep f = GRep f

  -- |
  -- @
  -- 'fmap' f . 'tabulate' ≡ 'tabulate' . 'fmap' f
  -- @
  --
  -- If no definition is provided, this will default to 'gtabulate'.
  tabulate :: (Rep f -> a) -> f a
  default tabulate :: (Generic1 f, GRep f ~ Rep f, GTabulate (Rep1 f))
                   => (Rep f -> a) -> f a
  tabulate = gtabulate

  -- | If no definition is provided, this will default to 'gindex'.
  index    :: f a -> Rep f -> a
  default index :: (Generic1 f, GRep f ~ Rep f, GIndex (Rep1 f))
                => f a -> Rep f -> a
  index = gindex

-- | A default implementation of 'Rep' for a datatype that is an instance of
-- 'Generic1'. This is usually composed of 'Either', tuples, unit tuples, and
-- underlying 'Rep' values. For instance, if you have:
--
-- @
-- data Foo a = MkFoo a (Bar a) (Baz (Quux a)) deriving ('Functor', 'Generic1')
-- instance 'Representable' Foo
-- @
--
-- Then you'll get:
--
-- @
-- 'GRep' Foo = Either () (Either ('WrappedRep' Bar) ('WrappedRep' Baz, 'WrappedRep' Quux))
-- @
--
-- (See the Haddocks for 'WrappedRep' for an explanation of its purpose.)
type GRep f = GRep' (Rep1 f)

-- | A default implementation of 'tabulate' in terms of 'GRep'.
gtabulate :: (Generic1 f, GRep f ~ Rep f, GTabulate (Rep1 f))
          => (Rep f -> a) -> f a
gtabulate = to1 . gtabulate'

-- | A default implementation of 'index' in terms of 'GRep'.
gindex :: (Generic1 f, GRep f ~ Rep f, GIndex (Rep1 f))
       => f a -> Rep f -> a
gindex = gindex' . from1

type family GRep' (f :: * -> *) :: *
class GTabulate f where
  gtabulate' :: (GRep' f -> a) -> f a
class GIndex f where
  gindex' :: f a -> GRep' f -> a

type instance GRep' (f :*: g) = Either (GRep' f) (GRep' g)
instance (GTabulate f, GTabulate g) => GTabulate (f :*: g) where
  gtabulate' f = gtabulate' (f . Left) :*: gtabulate' (f . Right)
instance (GIndex f, GIndex g) => GIndex (f :*: g) where
  gindex' (a :*: _) (Left  i) = gindex' a i
  gindex' (_ :*: b) (Right j) = gindex' b j

type instance GRep' (f :.: g) = (WrappedRep f, GRep' g)
instance (Representable f, GTabulate g) => GTabulate (f :.: g) where
  gtabulate' f = Comp1 $ tabulate $ fmap gtabulate' $ fmap (curry f) WrapRep
instance (Representable f, GIndex g) => GIndex (f :.: g) where
  gindex' (Comp1 fg) (i, j) = gindex' (index fg (unwrapRep i)) j

type instance GRep' Par1 = ()
instance GTabulate Par1 where
  gtabulate' f = Par1 (f ())
instance GIndex Par1 where
  gindex' (Par1 a) () = a

type instance GRep' (Rec1 f) = WrappedRep f
#if __GLASGOW_HASKELL__ >= 708
-- Using coerce explicitly here seems a bit more readable, and
-- likely a drop easier on the simplifier.
instance Representable f => GTabulate (Rec1 f) where
  gtabulate' = coerce (tabulate :: (Rep f -> a) -> f a)
                 :: forall a . (WrappedRep f -> a) -> Rec1 f a
instance Representable f => GIndex (Rec1 f) where
  gindex' = coerce (index :: f a -> Rep f -> a)
                 :: forall a . Rec1 f a -> WrappedRep f -> a
#else
instance Representable f => GTabulate (Rec1 f) where
  gtabulate' = Rec1 #. tabulate .# (. WrapRep)
instance Representable f => GIndex (Rec1 f) where
  gindex' = (. unwrapRep) #. index .# unRec1
#endif

type instance GRep' (M1 i c f) = GRep' f
instance GTabulate f => GTabulate (M1 i c f) where
  gtabulate' = M1 #. gtabulate'
instance GIndex f => GIndex (M1 i c f) where
  gindex' = gindex' .# unM1

-- | On the surface, 'WrappedRec' is a simple wrapper around 'Rep'. But it plays
-- a very important role: it prevents generic 'Representable' instances for
-- recursive types from sending the typechecker into an infinite loop. Consider
-- the following datatype:
--
-- @
-- data Stream a = a :< Stream a deriving ('Functor', 'Generic1')
-- instance 'Representable' Stream
-- @
--
-- With 'WrappedRep', we have its 'Rep' being:
--
-- @
-- 'Rep' Stream = 'Either' () ('WrappedRep' Stream)
-- @
--
-- If 'WrappedRep' didn't exist, it would be:
--
-- @
-- 'Rep' Stream = Either () (Either () (Either () ...))
-- @
--
-- An infinite type! 'WrappedRep' breaks the potentially infinite loop.
newtype WrappedRep f = WrapRep { unwrapRep :: Rep f }

{-# RULES
"tabulate/index" forall t. tabulate (index t) = t #-}

-- | 'tabulate' and 'index' form two halves of an isomorphism.
--
-- This can be used with the combinators from the @lens@ package.
--
-- @'tabulated' :: 'Representable' f => 'Iso'' ('Rep' f -> a) (f a)@
tabulated :: (Representable f, Representable g, Profunctor p, Functor h)
          => p (f a) (h (g b)) -> p (Rep f -> a) (h (Rep g -> b))
tabulated = dimap tabulate (fmap index)
{-# INLINE tabulated #-}

-- * Default definitions

fmapRep :: Representable f => (a -> b) -> f a -> f b
fmapRep f = tabulate . fmap f . index

pureRep :: Representable f => a -> f a
pureRep = tabulate . const

bindRep :: Representable f => f a -> (a -> f b) -> f b
bindRep m f = tabulate $ \a -> index (f (index m a)) a

mfixRep :: Representable f => (a -> f a) -> f a
mfixRep = tabulate . mfix . fmap index

mzipWithRep :: Representable f => (a -> b -> c) -> f a -> f b -> f c
mzipWithRep f as bs = tabulate $ \k -> f (index as k) (index bs k)

mzipRep :: Representable f => f a -> f b -> f (a, b)
mzipRep as bs = tabulate (index as &&& index bs)

askRep :: Representable f => f (Rep f)
askRep = tabulate id

localRep :: Representable f => (Rep f -> Rep f) -> f a -> f a
localRep f m = tabulate (index m . f)

apRep :: Representable f => f (a -> b) -> f a -> f b
apRep f g = tabulate (index f <*> index g)

distributeRep :: (Representable f, Functor w) => w (f a) -> f (w a)
distributeRep wf = tabulate (\k -> fmap (`index` k) wf)

collectRep :: (Representable f, Functor w) => (a -> f b) -> w a -> f (w b)
collectRep f w = tabulate (\k -> (`index` k) . f <$> w)

duplicateRepBy :: Representable f => (Rep f -> Rep f -> Rep f) -> f a -> f (f a)
duplicateRepBy plus w = tabulate (\m -> tabulate (index w . plus m))

extendRepBy :: Representable f => (Rep f -> Rep f -> Rep f) -> (f a -> b) -> f a -> f b
extendRepBy plus f w = tabulate (\m -> f (tabulate (index w . plus m)))

extractRepBy :: Representable f => (Rep f) -> f a -> a
extractRepBy = flip index

duplicatedRep :: (Representable f, Semigroup (Rep f)) => f a -> f (f a)
duplicatedRep = duplicateRepBy (<>)

extendedRep :: (Representable f, Semigroup (Rep f)) => (f a -> b) -> f a -> f b
extendedRep = extendRepBy (<>)

duplicateRep :: (Representable f, Monoid (Rep f)) => f a -> f (f a)
duplicateRep = duplicateRepBy mappend

extendRep :: (Representable f, Monoid (Rep f)) => (f a -> b) -> f a -> f b
extendRep = extendRepBy mappend

extractRep :: (Representable f, Monoid (Rep f)) => f a -> a
extractRep = extractRepBy mempty

imapRep :: Representable r => (Rep r -> a -> a') -> (r a -> r a')
imapRep f xs = tabulate (f <*> index xs)

ifoldMapRep :: forall r m a. (Representable r, Foldable r, Monoid m)
            => (Rep r -> a -> m) -> (r a -> m)
ifoldMapRep ix xs = fold (tabulate (\(i :: Rep r) -> ix i $ index xs i) :: r m)

itraverseRep :: forall r f a a'. (Representable r, Traversable r, Applicative f)
             => (Rep r -> a -> f a') -> (r a -> f (r a'))
itraverseRep ix xs = sequenceA $ tabulate (ix <*> index xs)

-- * Instances

instance Representable Proxy where
  type Rep Proxy = Void
  index Proxy = absurd
  tabulate _ = Proxy

instance Representable Identity where
  type Rep Identity = ()
  index (Identity a) () = a
  tabulate f = Identity (f ())

instance Representable (Tagged t) where
  type Rep (Tagged t) = ()
  index (Tagged a) () = a
  tabulate f = Tagged (f ())

instance Representable m => Representable (IdentityT m) where
  type Rep (IdentityT m) = Rep m
  index = index .# runIdentityT
  tabulate = IdentityT #. tabulate

instance Representable ((->) e) where
  type Rep ((->) e) = e
  index = id
  tabulate = id

instance Representable m => Representable (ReaderT e m) where
  type Rep (ReaderT e m) = (e, Rep m)
  index (ReaderT f) (e,k) = index (f e) k
  tabulate = ReaderT . fmap tabulate . curry

instance (Representable f, Representable g) => Representable (Compose f g) where
  type Rep (Compose f g) = (Rep f, Rep g)
  index (Compose fg) (i,j) = index (index fg i) j
  tabulate = Compose . tabulate . fmap tabulate . curry

instance Representable w => Representable (TracedT s w) where
  type Rep (TracedT s w) = (s, Rep w)
  index (TracedT w) (e,k) = index w k e
  tabulate = TracedT . unCo . collect (Co #. tabulate) . curry

instance (Representable f, Representable g) => Representable (Product f g) where
  type Rep (Product f g) = Either (Rep f) (Rep g)
  index (Pair a _) (Left i)  = index a i
  index (Pair _ b) (Right j) = index b j
  tabulate f = Pair (tabulate (f . Left)) (tabulate (f . Right))

instance Representable f => Representable (Cofree f) where
  type Rep (Cofree f) = Seq (Rep f)
  index (a :< as) key = case Seq.viewl key of
      Seq.EmptyL -> a
      k Seq.:< ks -> index (index as k) ks
  tabulate f = f Seq.empty :< tabulate (\k -> tabulate (f . (k Seq.<|)))

instance Representable f => Representable (Backwards f) where
  type Rep (Backwards f) = Rep f
  index = index .# forwards
  tabulate = Backwards #. tabulate

instance Representable f => Representable (Reverse f) where
  type Rep (Reverse f) = Rep f
  index = index .# getReverse
  tabulate = Reverse #. tabulate

instance Representable Monoid.Dual where
  type Rep Monoid.Dual = ()
  index (Monoid.Dual d) () = d
  tabulate f = Monoid.Dual (f ())

instance Representable Monoid.Product where
  type Rep Monoid.Product = ()
  index (Monoid.Product p) () = p
  tabulate f = Monoid.Product (f ())

instance Representable Monoid.Sum where
  type Rep Monoid.Sum = ()
  index (Monoid.Sum s) () = s
  tabulate f = Monoid.Sum (f ())

#if MIN_VERSION_base(4,4,0)
instance Representable Complex where
  type Rep Complex = Bool
  index (r :+ i) key = if key then i else r
  tabulate f = f False :+ f True
#endif

instance Representable U1 where
  type Rep U1 = Void
  index U1 = absurd
  tabulate _ = U1

instance (Representable f, Representable g) => Representable (f :*: g) where
  type Rep (f :*: g) = Either (Rep f) (Rep g)
  index (a :*: _) (Left  i) = index a i
  index (_ :*: b) (Right j) = index b j
  tabulate f = tabulate (f . Left) :*: tabulate (f . Right)

instance (Representable f, Representable g) => Representable (f :.: g) where
  type Rep (f :.: g) = (Rep f, Rep g)
  index (Comp1 fg) (i, j) = index (index fg i) j
  tabulate = Comp1 . tabulate . fmap tabulate . curry

instance Representable Par1 where
  type Rep Par1 = ()
  index (Par1 a) () = a
  tabulate f = Par1 (f ())

instance Representable f => Representable (Rec1 f) where
  type Rep (Rec1 f) = Rep f
  index = index .# unRec1
  tabulate = Rec1 #. tabulate

instance Representable f => Representable (M1 i c f) where
  type Rep (M1 i c f) = Rep f
  index = index .# unM1
  tabulate = M1 #. tabulate

newtype Co f a = Co { unCo :: f a } deriving Functor

instance Representable f => Representable (Co f) where
  type Rep (Co f) = Rep f
  tabulate = Co #. tabulate
  index = index .# unCo

instance Representable f => Apply (Co f) where
  (<.>) = apRep

instance Representable f => Applicative (Co f) where
  pure = pureRep
  (<*>) = apRep

instance Representable f => Distributive (Co f) where
  distribute = distributeRep
  collect = collectRep

instance Representable f => Bind (Co f) where
  (>>-) = bindRep

instance Representable f => Monad (Co f) where
  return = pure
  (>>=) = bindRep

#if defined(__GLASGOW_HASKELL__) && __GLASGOW_HASKELL__ >= 704
instance (Representable f, Rep f ~ a) => MonadReader a (Co f) where
  ask = askRep
  local = localRep
#endif

instance (Representable f, Semigroup (Rep f)) => Extend (Co f) where
  extended = extendedRep

instance (Representable f, Monoid (Rep f)) => Comonad (Co f) where
  extend = extendRep
  extract = extractRep

instance ComonadTrans Co where
  lower (Co f) = f

liftR2 :: Representable f => (a -> b -> c) -> f a -> f b -> f c
liftR2 f fa fb = tabulate $ \i -> f (index fa i) (index fb i)

liftR3 :: Representable f => (a -> b -> c -> d) -> f a -> f b -> f c -> f d
liftR3 f fa fb fc = tabulate $ \i -> f (index fa i) (index fb i) (index fc i)

max_stars_repo_name: bgamari/adjunctions

max_issues_repo_name: bgamari/adjunctions

max_forks_repo_name: bgamari/adjunctions


text:
{-# LANGUAGE UndecidableInstances,ScopedTypeVariables,FlexibleContexts,RankNTypes,NoMonomorphismRestriction #-}
import Numeric.LinearAlgebra hiding ((><))
import Control.Applicative
import Foreign.Cholesky
import Linear.Metric
import Data.Functor.Compose
import qualified Data.Foldable as F
import Space.Class
import Linear.Vector
import Linear.V1
import Data.Profunctor
import Data.Monoid
import Linear.V2
import Test.QuickCheck
import Data.Distributive
import Linear.V3
import Linear.V4
import Linear.Quaternion
import Linear.Matrix
import Foreign.Cholesky
import Vectorization
import Control.Lens
import Kalman

import Test.Tasty
import Test.Tasty.SmallCheck as SC
import Test.Tasty.QuickCheck as QC
import Test.Tasty.HUnit

-- Aided Inertial Navigation -- Example 5.4

pe0 = (V2 (V2 100 0) (V2 0 100))
xe0 = V2 0 (0 :: Double)
se0 = (xe0,sqrtM pe0)
sc0 = (xe0,pe0)

sqrtM = unRight . potrf

qr0 = (V2 (V2 0.1 0 ) (V2 0 0.1))

he0 :: V2 Double -> V1 Double
he0 (V2 x _) = V1 x
he1 (V2 _ y) = V1 y
he2 (V2 x y) = V1 (x*0.7 + 0.3*y)
he3 (V2 x y) = V1 (x*0.5 + y*0.5)

f t a (V2 x y) = (V2 y a )
m0 = (he0,re0,ye0)
m1 = (he1,re1,ye1)
m2 = (he2,re2,ye2)
m3 = (he3,re3,ye3)

mlist = [m0,m1,m2,m3]

re0 :: V1 (V1 Double)
re0 = 1
re1 = 1
re2 = 1
re3 = 1

hn x = (distribute $ V4 (he0 x)(he1 x) (he2 x) (he3 x))^. _x

rn = V4
       (V4 1 0 0  0 )
       (V4 0 1 0 0 )
       (V4 0 0 1 0)
       (V4 0 0 0 1)

V1 yn = distribute (V4 ye0 ye1 ye2 ye3)

ye0 = V1 10.24
ye1 = V1 21.20
ye2 = V1 13.91
ye3 = V1 14.84

measure' (h,r,y) = sqrtMeasure h r y
sef = foldr measure' se0 (reverse mlist)

main = defaultMain tests
tests :: TestTree
tests = testGroup "Tests" [properties ]
properties = testGroup "Property" [quick] -- ,small]

quick = testGroup "(checked by QuickCheck)"
    [ QC.testProperty "x == sqrt x * sqrt x"  propSquare
    , QC.testProperty "measureState " propState
    , QC.testProperty "measureCovariance" propCovariance
    , QC.testProperty "propMulti" propMulti
    ]
{-
small = testGroup "(checked by SmallCheck)"
  [ SC.testProperty "x == sqrt x * sqrt x"  propSquare
  ]
-}

newtype Square f a = Square {unSquare :: (f (f a) )}

instance (Traversable f ,Show (f (f a)) ) => Show (Square f a ) where
    show  = show . unSquare

instance (R f,Floating a , Fractional a ,Num a , Arbitrary a, Arbitrary (f (f a)) ) => Arbitrary (Square f a ) where
    arbitrary = fmap (Square . squareT)  tgen
        where tgen = arbitrary :: Gen (f (f a))

instance (Num a ,Arbitrary a )=> Arbitrary (V2 a)  where
    arbitrary = V2 <$> (((+1).abs) <$> arbitrary ) <*> (((+1).abs) <$>  arbitrary)

instance (Num a ,Arbitrary a )=> Arbitrary (V1 a)  where
    arbitrary = fmap (pure . (+1) . abs ) arbitrary

squareNeg :: (Num a, R f) => f a -> f (f a)
squareNeg v = mult v (fmap negate v)

propCovariance (Square re0 ) ye1 s1 (Square s2) =  nearEqual (snd sc2) (squareT $ snd se2)
    where
        se2 = sqrtMeasure he1 (sqrtM re0) ye1 (s1,sqrtM s2)
        sc2 = measure he1 re0 ye1 (s1,s2)

propSquare :: Square V2 Double -> Bool
propSquare (Square x)  = nearEqual x (distribute sq !*! sq)
    where sq = sqrtM x

propMulti s1 (Square s2) =  nearEqualP t (fst sn1) (fst sn2)   &&   nearEqualP (t+3) (snd sn1) (snd sn2)
    where
        t = 7
        sn2 = sqrtMeasure hn (sqrtM rn) yn (s1,sqrtM s2)
        m1 (h,y) = sqrtMeasure h (sqrtM re0) y
        l = reverse [(he0,ye0),(he1,ye1),(he2,ye2),(he3,ye3)]
        sn1 = foldr m1 (s1,sqrtM s2) l

propState (Square re0 ) ye1 s1 (Square s2) = nearEqual (fst sc2) (fst se2) -- && nearEqual (snd sc2) (squareT $ snd se2)
    where
        se2 = sqrtMeasure he1 (sqrtM re0) ye1 (s1,sqrtM s2)
        sc2 = measure he1 re0 ye1 (s1,s2)

nearEqualP t x y = abs (x - y) < eps*(abs x) && abs (y - x) < eps*(abs y)
    where eps = 10^t*2.220446049250313e-16

nearEqual x y = abs (x - y) < eps*(abs x) && abs (y - x) < eps*(abs y)
    where eps = 1000000*2.220446049250313e-16

max_stars_repo_name: massudaw/mtk

max_issues_repo_name: massudaw/mtk

max_forks_repo_name: massudaw/mtk


text:
{-# LANGUAGE GADTs #-}
{-# LANGUAGE StandaloneDeriving #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE TupleSections #-}
{-# LANGUAGE InstanceSigs #-}
{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE DeriveGeneric #-}
{-# LANGUAGE DeriveFunctor #-}
{-# LANGUAGE OverloadedLists #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}

module Data.Semigroup.Free where

import GHC.Generics
import GHC.Read
import Control.Applicative
import Data.Functor.Classes
import Data.Functor.Identity
import Data.Semigroup
import Data.List.NonEmpty (NonEmpty(..))
import Data.Coerce
import Data.List (sortOn)
import Control.Monad
import Control.Monad.Free (Free(..), MonadFree(..), iterA)
import Control.Monad.Free.Church (F(..))
import qualified Control.Monad.Free.Church as F
import Data.Pair

import Criterion
import Criterion.Main
import Criterion.Types
import Data.Proxy (Proxy(..), asProxyTypeOf)
import Statistics.Types

import Weigh (Weigh, func, mainWith, wgroup)

-- | Size notes:
--
-- @
--  So a single element weighs ~400 bytes.
--    Supposing a MArray has a 1-byte overhead per element,
--    could we get a ~2 byte overhead per element if we stored every 400 elements in a MArray?
--    Ehh.. I don't know if the price of copying arrays will offset the cost..
--
--  What if we only access one cycle element at a time?
--    In that case, it's probably better to simply carry around a list and cons onto it.
--
--  What if we have a predictable cycle-size?
--    In that case, it's probably best to use pre-allocated memory
--    Possibly a pool of smaller pre-allocated chunks
-- @
--
sizeNotes :: ()
sizeNotes = ()

-- | Note:
--
-- @
--  On one hand, it seems nice to preserve the commutativity of the GappedIxCycle's
--  On the other hand, it takes O(min n m) time to zip two GappedIxCycle's together.
--  On the other hand, I'm not sure it needs to be anything other than a []/Seq/Set of period-1 cycles whose LCM is the given period. (that would give O(1) (<>))
-- @
--
notes :: ()
notes = ()

---- A strict finite collection of single-element cycles with:
---- O(1) `return`, `cons`, `snoc`, and `mappend`.
----
---- Pretty much everything else is much slower than lists.
---- (This data type's main features are its strictness and ability to `mappend`
---- two structures in constant time.)
----
---- Isomorphic to a strict version of:
----
---- @
----  GappedCycle a == Free Pair (Int, a)
----  wrap :: Pair (Free Pair (Int, a)) -> Free Pair (Int, a)
---- @
----
---- Which is effectively "pairing for free", exactly what we want.
----
---- Take a look at this:
----
---- @
----  F f a = F { runF :: forall r. (a -> r) -> (f r -> r) -> r }
----  F Pair a = F { runF :: forall r. (a -> r) -> (Pair r -> r) -> r }
----  F Pair a = F { runF :: forall r. (a -> r) -> (r -> r -> r) -> r }
----  F Pair a = F { runF :: forall r. (a -> r) -> (r -> r -> r) -> r }
----  F Pair a = F { runF :: forall r. Semigroup r => (a -> r) -> r }
----  F Pair a = FreeSemigroup a
---- @
----
----
--data GappedCycle a where
--  GCPure :: !Int -> !a -> GappedCycle a -- Period, Value
--  GCMappend :: !(GappedCycle a) -> !(GappedCycle a) -> GappedCycle a
--
--
-- |
--
-- @
--  GappedCycle a == Free Pair (Int, a)
-- @
--
-- data IntGappedCycle where
--   IGCPure :: !Int -> !Int -> IntGappedCycle -- Period, Value
--   IGCMappend :: !IntGappedCycle -> !IntGappedCycle -> IntGappedCycle
--

-- fun idea to test collection of bijections:
-- - make graph of bijections
-- - randomly traverse the graph until loop
-- - test equality with looped

-- If there are no edges, fails with an error.
-- Otherwise:
-- - randomly select a vertex with positive degree
-- - randomly select an adjacent vertex
-- - if the randomly selected vertex has been visited already, return all previous values
-- - otherwise add to collection of visited
-- randomCycleFrom :: Graph a -> m [a]

-- Interesting:
--   Store an array as a Set and a multi-permutation (index table) -- for fast updates, order, and no-repeated values
--   Or an ordered array and a (muti)-permutation -- read/write is fast, bsearch is fast, most combinations/alterations are very slow.

-- | `Alternative` without `empty`
class Applicative f => Semialternative f where
  (<!>) :: f a -> f a -> f a

  saconcat :: NonEmpty (f a) -> f a
  saconcat ~(x :| xs) =
    case xs of
      [] -> x
      ~(y:ys) -> x <!> saconcat (y :| ys)

  saconcatMap :: (a -> f b) -> NonEmpty a -> f b
  saconcatMap f = saconcat . fmap f

instance Semialternative NonEmpty where
  (<!>) = (<>)

-- | Wrap a `Semialternative` to provide a `Semigroup` instance
newtype WrapSemialternative f a = WrapSemialternative { runWrapSemialternative :: f a }

instance Semialternative f => Semigroup (WrapSemialternative f a) where
  {-# INLINE (<>) #-}
  (<>) :: WrapSemialternative f a -> WrapSemialternative f a -> WrapSemialternative f a
  (<>) x y = coerce $ coerce x <!> (coerce `asTypeOf` runWrapSemialternative) y

-- | A free `Semigroup` is a `Semialternative` that we can `foldMap1` over
-- (`foldMap` specialized to non-empty inputs) and convert to any other
-- `SemigroupFree` (they're all isomorphic to each other).
class Semialternative f => SemigroupFree f where
  foldMap1 :: Semigroup s => (a -> s) -> f a -> s
  foldMap1 f xs = runFSemigroup (semigroupFree xs) f

  semigroupFree :: SemigroupFree g => f a -> g a
  semigroupFree = runWrapSemialternative . foldMap1 (WrapSemialternative . pure)

  {-# MINIMAL foldMap1 | semigroupFree #-}

instance SemigroupFree NonEmpty where
  foldMap1 f ~(x :| xs) = loop x xs
    where
      loop x [] = f x
      loop x ~(y:ys) = f x <> loop x ys

-- | From the triangle benchmakrs, this is the fastest.
-- How fast? @~100x@ faster than `NonEmpty`, taking
-- @~180 ms@ to construct and destruct a triangle with
-- around a million elements.
--
-- For a little comparison, the average human blink of an eye
-- is around @100-400 ms@ (according to WolframAlpha). I can
-- process a trangular collection of a million elements in less
-- then half the average time that it takes to blink an eye!
--
-- That's @~180 ns@ per element..
--
-- My computer has a @3.1 GHz@ processor. Assuming total utilization
-- of a single core, we get a cycle-duration of @~0.323 ns@. That
-- gives us an estimate of @557@ cycles per element.
--
-- That's not the greatest. I'll have to check CPU utilization for that benchmark to see if it's skewing my quick math.
-- (267.2%?, not sure what to make of that. need to run without criterion to remove overhead)
-- Though it's still not bad for a pure, polymorphic data type.
--
-- All in all, this data type is plenty fast for its primary purpose:
-- quickly accumulating cycle-pieces when calculating the final cycle.
data FreeSemigroup a where
  PureFreeSemigroup :: a -> FreeSemigroup a
  AppendFreeSemigroup :: FreeSemigroup a -> FreeSemigroup a -> FreeSemigroup a
  deriving (Eq, Ord, Read, Show, Functor, Generic, Generic1)

instance Semigroup (FreeSemigroup a) where
  {-# INLINE (<>) #-}
  (<>) = AppendFreeSemigroup

instance Applicative FreeSemigroup where
  {-# INLINE pure #-}
  pure = PureFreeSemigroup

  PureFreeSemigroup f <*> xs = f <$> xs
  f <*> PureFreeSemigroup x = ($ x) <$> f
  ~(AppendFreeSemigroup fs gs) <*> ~(AppendFreeSemigroup xs ys) = AppendFreeSemigroup (fs <*> xs) (gs <*> ys)

instance Semialternative FreeSemigroup where
  {-# INLINE (<!>) #-}
  (<!>) = AppendFreeSemigroup

instance SemigroupFree FreeSemigroup where
  semigroupFree (PureFreeSemigroup x) = pure x
  semigroupFree (AppendFreeSemigroup xs ys) = semigroupFree xs <!> semigroupFree ys

-- | `semigroupFree` specialized to `FreeSemigroup`
{-# INLINE toFreeSemigroup #-}
toFreeSemigroup :: SemigroupFree f => f a -> FreeSemigroup a
toFreeSemigroup = semigroupFree

-- | It looks like this sucks around the [2^12..2^16] elements range
data FreeSemigroup' a where
  PureFreeSemigroup' :: !a -> FreeSemigroup' a
  AppendFreeSemigroup' :: !(FreeSemigroup' a) -> !(FreeSemigroup' a) -> FreeSemigroup' a
  deriving (Eq, Ord, Read, Show, Functor, Generic, Generic1)

instance Semigroup (FreeSemigroup' a) where
  {-# INLINE (<>) #-}
  (<>) = AppendFreeSemigroup'

instance Applicative FreeSemigroup' where
  {-# INLINE pure #-}
  pure !x = PureFreeSemigroup' x

  PureFreeSemigroup' f <*> xs = f <$> xs
  f <*> PureFreeSemigroup' x = ($ x) <$> f
  ~(AppendFreeSemigroup' fs gs) <*> ~(AppendFreeSemigroup' xs ys) = AppendFreeSemigroup' (fs <*> xs) (gs <*> ys)

instance Semialternative FreeSemigroup' where
  {-# INLINE (<!>) #-}
  (<!>) = AppendFreeSemigroup'

instance SemigroupFree FreeSemigroup' where
  semigroupFree (PureFreeSemigroup' x) = pure x
  semigroupFree (AppendFreeSemigroup' xs ys) = semigroupFree xs <!> semigroupFree ys

-- | `Free`-based `FreeSemigroup`
newtype FreePS a = FreePS { runFreePS :: Free Pair a } deriving (Eq, Ord, Show, Read, Functor, Applicative, Monad)

deriving instance MonadFree Pair FreePS

instance Semialternative FreePS where
  (<!>) = fmap (wrap . Pair) . (,)

instance SemigroupFree FreePS where
  semigroupFree = iterA (uncurry (<!>) . getPair) . runFreePS

-- | `F`-based `FreeSemigroup`
newtype FPS a = FPS { runFPS :: F Pair a } deriving (Functor, Applicative, Monad)

deriving instance MonadFree Pair FPS

instance Semialternative FPS where
  (<!>) = fmap (wrap . Pair) . (,)

instance SemigroupFree FPS where
  semigroupFree = F.iter (uncurry (<!>) . getPair) . fmap pure . runFPS

-- | Partially applied `foldMap1` wrapped in a @newtype@.
newtype FSemigroup a = FSemigroup { runFSemigroup :: forall s. Semigroup s => (a -> s) -> s } deriving (Functor)

instance Eq a => Eq (FSemigroup a) where
  {-# INLINE (==) #-}
  x == y = toFreeSemigroup x == toFreeSemigroup y

instance Ord a => Ord (FSemigroup a) where
  {-# INLINE compare #-}
  compare x y = toFreeSemigroup x `compare` toFreeSemigroup y

instance Show a => Show (FSemigroup a) where
  {-# INLINE show #-}
  show = show . toFreeSemigroup

instance Semigroup (FSemigroup a) where
  {-# INLINE (<>) #-}
  FSemigroup x <> FSemigroup y = FSemigroup $ \cnv -> y cnv <> y cnv

instance Applicative FSemigroup where
  {-# INLINE pure #-}
  pure x = FSemigroup $ \cnv -> cnv x

  {-# INLINE liftA2 #-}
  liftA2 f (FSemigroup xs) (FSemigroup ys) = FSemigroup $ \cnv -> ys . xs $ fmap cnv . f

instance Semialternative FSemigroup where
  {-# INLINE (<!>) #-}
  FSemigroup x <!> FSemigroup y = FSemigroup $ \cnv -> x cnv <> y cnv

instance SemigroupFree FSemigroup where
  {-# INLINE foldMap1 #-}
  foldMap1 cnv (FSemigroup xs) = xs cnv

-- | `FSemigroup` with explicitly provided @(`<>`)@.
--
-- Slightly faster in triangle benchmarks.
newtype FS a = FS
  { runFS :: forall s. (s -> s -> s) -> (a -> s) -> s
  } deriving (Functor)

instance Eq a => Eq (FS a) where
  {-# INLINE (==) #-}
  x == y = toFreeSemigroup x == toFreeSemigroup y

instance Ord a => Ord (FS a) where
  {-# INLINE compare #-}
  compare x y = toFreeSemigroup x `compare` toFreeSemigroup y

instance Show a => Show (FS a) where
  {-# INLINE show #-}
  show = show . toFreeSemigroup

instance Semigroup (FS a) where
  {-# INLINE (<>) #-}
  FS x <> FS y = FS $ \append cnv -> x append cnv `append` y append cnv

instance Applicative FS where
  {-# INLINE pure #-}
  pure x = FS $ \append cnv -> cnv x

  {-# INLINE liftA2 #-}
  liftA2 f (FS xs) (FS ys) =
    FS $ \append cnv -> ys append . xs (liftA2 append) $ fmap cnv . f

instance Semialternative FS where
  {-# INLINE (<!>) #-}
  FS x <!> FS y = FS $ \append cnv -> x append cnv `append` y append cnv

instance SemigroupFree FS where
  {-# INLINE foldMap1 #-}
  foldMap1 cnv (FS xs) = xs (<>) cnv

-- | Lift function application to strict function application
{-# INLINE liftS2 #-}
liftS2 :: (a -> b -> c) -> a -> b -> c
liftS2 f x y = f x' y'
  where
    !x' = x
    !y' = y

-- | Strict `FSemigroup` with explicitly provided @(`<>`)@.
newtype FS' a = FS'
  { runFS' :: forall s. (s -> s -> s) -> (a -> s) -> s
  } deriving (Functor)

instance Eq a => Eq (FS' a) where
  {-# INLINE (==) #-}
  x == y = toFreeSemigroup x == toFreeSemigroup y

instance Ord a => Ord (FS' a) where
  {-# INLINE compare #-}
  compare x y = toFreeSemigroup x `compare` toFreeSemigroup y

instance Show a => Show (FS' a) where
  {-# INLINE show #-}
  show = show . toFreeSemigroup

instance Semigroup (FS' a) where
  {-# INLINE (<>) #-}
  FS' x <> FS' y = FS' $ \append cnv -> x append cnv `append` y append cnv

instance Applicative FS' where
  {-# INLINE pure #-}
  pure x = let !x' = x in FS' $ \_ cnv -> cnv x'

  {-# INLINE liftA2 #-}
  liftA2 f (FS' xs) (FS' ys) =
    FS' $ \append cnv -> ys append . xs (liftA2 append) $ fmap cnv . f

instance Semialternative FS' where
  {-# INLINE (<!>) #-}
  FS' x <!> FS' y = FS' $ \append cnv -> x append cnv `append` y append cnv

instance SemigroupFree FS' where
  {-# INLINE foldMap1 #-}
  foldMap1 cnv (FS' xs) = xs (liftS2 (<>)) (cnv $!)

-- | A concatenated "triangle" of values:
--
-- @
--  saconcat [1] <> saconcat [1, 2] <> saconcat [1, 2, 3] <> .. <> saconcat [1..n]
-- @
--
makeTriangle :: Semialternative f => Int -> f Int
makeTriangle n | n < 1 = error "makeTriangle: non-positive input"
makeTriangle 1 = pure 1
makeTriangle n = makeTriangle (n - 1) <!> makeLine n

-- | Conveniently construct and `Sum` a triangle of `SemigroupFree` values
-- for benchmarking.
destructTriangle :: SemigroupFree f => Proxy f -> Int -> Int
destructTriangle prxy = getSum . foldMap1 Sum . (`asProxyTypeOf` proxyApInt prxy) . makeTriangle

-- | A concatenated line of values
--
-- @
--  saconcat [1..n]
-- @
makeLine :: Semialternative f => Int -> f Int
makeLine n = saconcatMap pure [1..n]

-- | Conveniently construct and `Sum` a line of `SemigroupFree` values
-- for benchmarking.
destructLine :: SemigroupFree f => Proxy f -> Int -> Int
destructLine prxy = getSum . foldMap1 Sum . (`asProxyTypeOf` proxyApInt prxy) . makeLine

-- | Apply `Int` to a type inside a @proxy@
proxyApInt :: proxy g -> Proxy (g Int)
proxyApInt _ = Proxy

-- | Benchmark `destructTriangle`
benchTriangle :: SemigroupFree f => Proxy f -> Int -> Benchmarkable
benchTriangle prxy !x = nf (destructTriangle prxy) x

-- | Benchmark `destructLine`
benchLine :: SemigroupFree f => Proxy f -> Int -> Benchmarkable
benchLine prxy !x = nf (destructLine prxy) x

-- | Weigh `destructTriangle`
weighTriangle :: SemigroupFree f => String -> Proxy f -> Int -> Weigh ()
weighTriangle name prxy !sz = func desc' (destructTriangle prxy) sz
  where
    desc' = unwords [name, "destructTriangle, size:", show sz]

-- | Weigh `destructLine`
weighLine :: SemigroupFree f => String -> Proxy f -> Int -> Weigh ()
weighLine name prxy !sz = func desc' (destructLine prxy) sz
  where
    desc' = unwords [name, "destructLine    , size:", show sz]

-- | `weighTriangle` and `weighLine` for powers of @2@
-- in experimentally derived ranges.
weighPow2s :: SemigroupFree f => String -> Proxy f -> Weigh ()
weighPow2s name prxy = wgroup name $ do
  forM_ ([5..15] :: [Int]) $ weighLine name prxy . (2 ^)
  forM_ ([1..10] :: [Int]) $ weighTriangle name prxy . (2 ^)

type Est = Estimate ConfInt Double

reportEstimate :: Report -> Est
reportEstimate = anMean . reportAnalysis

estimate :: Benchmarkable -> IO Est
estimate = fmap reportEstimate . benchmark'

quickEstimate :: Benchmarkable -> IO Est
quickEstimate = fmap reportEstimate . benchmarkWith' quickConfig

quickConfig :: Config
quickConfig =
  defaultConfig {confInterval = cl90, timeLimit = 1, resamples = 100}

-- | `benchTriangle` and `benchLine` for powers of @2@
-- in experimentally derived ranges.
benchPow2s :: SemigroupFree f => String -> Proxy f -> Benchmark
benchPow2s name prxy =
  bgroup name $
  ((\n -> (name ++ " Triangle: " ++ show n) `bench` benchTriangle prxy n) .
   (2 ^) <$>
   [1 .. 10]) ++
  ((\n -> (name ++ " Line: " ++ show n) `bench` benchLine prxy n) . (2 ^) <$>
   [5 .. 15])

-- | `sequenceA` a benchmark over all of the `SemigroupFree` implementations
benchable :: Applicative f => (forall g. SemigroupFree g => String -> Proxy g -> f a) -> f [a]
benchable f = sequenceA
  [ f "NonEmpty      " (Proxy @NonEmpty)
  , f "FreeSemigroup " (Proxy @FreeSemigroup)
  , f "FreeSemigroup'" (Proxy @FreeSemigroup')
  , f "FreePS        " (Proxy @FreePS)
  , f "FPS           " (Proxy @FPS)
  , f "FSemigroup    " (Proxy @FSemigroup)
  , f "FS            " (Proxy @FS)
  , f "FS'           " (Proxy @FS')
  ]

weighFS :: Weigh ()
weighFS = void $ benchable weighPow2s

runWeighFS :: IO ()
-- runWeighFS = mainWith $ weighPow2s "FS" (Proxy @FS)
runWeighFS = mainWith weighFS

benchCompare :: Int -> (forall g. SemigroupFree g => Proxy g -> Benchmarkable) -> IO ()
benchCompare n f = do
  results <- sortOn (estPoint . snd) <$> benchable (\n prxy -> (n, ) <$> estimate (f prxy))
  let resultNames = fst <$> results
  appendFile "orderings.txt" . unlines $ show n : resultNames ++ [""]
  mapM_ (\(n, e) -> putStrLn n >> putStrLn ("  " ++ show e)) results

benchFS :: [Benchmark]
benchFS = runIdentity $ benchable (fmap Identity . benchPow2s)

runBenchFS :: IO ()
runBenchFS =
  forM_ ([8,8+40..1414] :: [Int]) $ \i -> do
    putStrLn ""
    putStrLn $
      unwords
        [ "Benchmarking Triangle with input:"
        , show i
        , "(size"
        , show (sum [n | n <- [1 .. i]])
        , ")"
        ]
    benchCompare i (`benchTriangle` i)

max_stars_repo_name: michaeljklein/find-cycles

max_issues_repo_name: michaeljklein/find-cycles

max_forks_repo_name: michaeljklein/find-cycles