-- | Naive (very inefficient) algorithm to generate the irreducible (Dynkin) root systems
--
-- Based on <https://en.wikipedia.org/wiki/Root_system>

{-# LANGUAGE BangPatterns, FlexibleInstances, TypeSynonymInstances, FlexibleContexts #-}
module Math.Combinat.RootSystems where

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

import Control.Monad

import Data.Array

import Data.Set (Set)
import qualified Data.Set as Set

import Data.List
import Data.Ord

import Math.Combinat.Numbers.Primes
import Math.Combinat.Sets

--------------------------------------------------------------------------------
-- * Half-integers

-- | The type of half-integers (internally represented by their double)
--
-- TODO: refactor this into its own module
newtype HalfInt 
  = HalfInt Int  
  deriving (Eq,Ord)

half :: HalfInt
half = HalfInt 1

divByTwo :: Int -> HalfInt
divByTwo n = HalfInt n

mulByTwo :: HalfInt -> Int
mulByTwo (HalfInt n) = n

scaleBy :: Int -> HalfInt -> HalfInt
scaleBy k (HalfInt n) = HalfInt (k*n)

instance Show HalfInt where
  show (HalfInt n) = case divMod n 2 of
    (k,0) -> show k
    (_,1) -> show n ++ "/2"

instance Num HalfInt where
  fromInteger = HalfInt . (*2) . fromInteger
  a + b = divByTwo $ mulByTwo a + mulByTwo b
  a - b = divByTwo $ mulByTwo a - mulByTwo b
  a * b = case divMod (mulByTwo a * mulByTwo b) 4 of
            (k,0) -> HalfInt (2*k)
            (k,2) -> HalfInt (2*k+1)
            _     -> error "the result of multiplication is not a half-integer"
  negate = divByTwo . negate . mulByTwo
  signum = divByTwo . signum . mulByTwo
  abs    = divByTwo . abs    . mulByTwo

--------------------------------------------------------------------------------
-- * Vectors of half-integers

type HalfVec = [HalfInt]

instance Num HalfVec where
  fromInteger = error "HalfVec/fromInteger"
  (+) = safeZip (+)
  (-) = safeZip (-)
  (*) = safeZip (*)
  negate = map negate
  abs    = map abs
  signum = map signum

scaleVec :: Int -> HalfVec -> HalfVec  
scaleVec k = map (scaleBy k)

negateVec :: HalfVec -> HalfVec
negateVec = map negate

-- dotProd :: HalfVec -> HalfVec
-- dotProd xs ys = foldl' (+) 0 $ safeZip (*) xs ys

safeZip :: (a -> b -> c) -> [a] -> [b] -> [c]
safeZip f = go where
  go (x:xs) (y:ys) = f x y : go xs ys
  go []     []     = []
  go _      _      = error "safeZip: the lists do not have equal length"

--------------------------------------------------------------------------------
-- * Dynkin diagrams

data Dynkin
  = A !Int
  | B !Int
  | C !Int
  | D !Int
  | E6 | E7 | E8
  | F4
  | G2
  deriving (Eq,Show)

--------------------------------------------------------------------------------
-- * The roots of root systems

-- | The ambient dimension of (our representation of the) system (length of the vector)
ambientDim :: Dynkin -> Int
ambientDim d = case d of
  A n -> n+1   -- it's an n dimensional subspace of (n+1) dimensions
  B n -> n
  C n -> n
  D n -> n
  E6  -> 6
  E7  -> 8     -- sublattice of E8 ?
  E8  -> 8
  F4  -> 4
  G2  -> 3     -- it's a 2 dimensional subspace of 3 dimensions

simpleRootsOf :: Dynkin -> [HalfVec]
simpleRootsOf d = 

  case d of

    A n -> [ e i - e (i+1) | i <- [1..n]   ]

    B n -> [ e i - e (i+1) | i <- [1..n-1] ] ++ [e n]

    C n -> [ e i - e (i+1) | i <- [1..n-1] ] ++ [scaleVec 2 (e n)]

    D n -> [ e i - e (i+1) | i <- [1..n-1] ] ++ [e (n-1) + e n]

    E6  -> simpleRootsE6_123
    E7  -> simpleRootsE7_12
    E8  -> simpleRootsE8_even 

    F4  -> [ [ 1,-1, 0, 0]
           , [ 0, 1,-1, 0]
           , [ 0, 0, 1, 0]
           , [-h,-h,-h,-h]
           ]

    G2  -> [ [ 1,-1, 0]
           , [-1, 2,-1]
           ]

  where
    h = half
    n = ambientDim d

    e :: Int -> HalfVec
    e i = replicate (i-1) 0 ++ [1] ++ replicate (n-i) 0

positiveRootsOf :: Dynkin -> Set HalfVec
positiveRootsOf = positiveRoots . simpleRootsOf

negativeRootsOf :: Dynkin -> Set HalfVec
negativeRootsOf = Set.map negate . positiveRootsOf

allRootsOf :: Dynkin -> Set HalfVec
allRootsOf dynkin = Set.unions [  pos , neg ] where
  simple = simpleRootsOf dynkin
  pos   = positiveRoots simple
  neg   = Set.map negate pos

--------------------------------------------------------------------------------
-- * Positive roots

-- | Finds a vector, which is hopefully not orthognal to any root
-- (generated by the given simple roots), and has positive dot product with each of them.
findPositiveHyperplane :: [HalfVec] -> [Double]
findPositiveHyperplane vs = w where
  n  = length (head vs)
  w0 = map (fromIntegral . mulByTwo) (foldl1 (+) vs) :: [Double]
  w  = zipWith (+) w0 perturb
  perturb = map small $ map fromIntegral $ take n primes
  small :: Double -> Double
  small x = x / (10**10) 

positiveRoots :: [HalfVec] -> Set HalfVec
positiveRoots simples = Set.fromList pos where
  roots = mirrorClosure simples
  w     = findPositiveHyperplane simples
  pos   = [ r | r <- Set.toList roots , dot4 r > 0 ] where

  dot4 :: HalfVec -> Double
  dot4 a = foldl' (+) 0 $ safeZip (*) w $ map (fromIntegral . mulByTwo) a

basisOfPositives :: Set HalfVec -> [HalfVec]
basisOfPositives set = Set.toList (Set.difference set set2) where
  set2 = Set.fromList [ a + b | [a,b] <- choose 2 (Set.toList set) ]


--------------------------------------------------------------------------------
-- * Operations on half-integer vectors

-- | bracket b a = (a,b)/(a,a) 
bracket :: HalfVec -> HalfVec -> HalfInt
bracket b a = 
  case divMod (2*a_dot_b) (a_dot_a) of
    (n,0) -> divByTwo n
    _     -> error "bracket: result is not a half-integer"
  where
    a_dot_b = foldl' (+) 0 $ safeZip (*) (map mulByTwo a) (map mulByTwo b)
    a_dot_a = foldl' (+) 0 $ safeZip (*) (map mulByTwo a) (map mulByTwo a)

-- | mirror b a = b - 2*(a,b)/(a,a) * a
mirror :: HalfVec -> HalfVec -> HalfVec
mirror b a = b - scaleVec (mulByTwo $ bracket b a) a

-- | Cartan matrix of a list of (simple) roots
cartanMatrix :: [HalfVec] -> Array (Int,Int) Int
cartanMatrix list = array ((1,1),(n,n)) [ ((i,j), f i j) | i<-[1..n] , j<-[1..n] ] where
  n   = length list
  arr = listArray (1,n) list
  f !i !j = mulByTwo $ bracket (arr!j) (arr!i)

printMatrix :: Show a => Array (Int,Int) a -> IO ()
printMatrix arr = do
  let ((1,1),(n,m)) = bounds arr
      arr' = fmap show arr
  let ks   = [ 1 + maximum [ length (arr'!(i,j)) | i<-[1..n] ] | j<-[1..m] ]
  forM_ [1..n] $ \i -> do
    putStrLn $ flip concatMap [1..m] $ \j -> extendTo (ks!!(j-1)) $ arr' ! (i,j)
  where
    extendTo n s = replicate (n-length s) ' ' ++ s

--------------------------------------------------------------------------------
-- * Mirroring 

-- | We mirror stuff until there is no more things happening
-- (very naive algorithm, but seems to work)
mirrorClosure :: [HalfVec] -> Set HalfVec
mirrorClosure = go . Set.fromList where 
  
  go set 
    | n'  > n   = go set'
    | n'' > n   = go set''
    | otherwise = set
    where
      n   = Set.size set
      n'  = Set.size set'
      n'' = Set.size set''
      set'  = mirrorStep set
      set'' = Set.union set (Set.map negateVec set) 

mirrorStep :: Set HalfVec -> Set HalfVec
mirrorStep old = Set.union old new where
  new = Set.fromList [ mirror b a | [a,b] <- choose 2 $ Set.toList old ] 

--------------------------------------------------------------------------------
-- * E6, E7 and E8

-- | This is a basis of E6 as the subset of the even E8 root system
-- where the first three coordinates agree (they are consolidated 
-- into the first coordinate here)
simpleRootsE6_123:: [HalfVec]
simpleRootsE6_123 = roots where
  h = half
  roots =
    [ [-h,-h,-h,-h,-h,-h,-h,-h]
    , [ h, h, h, h, h, h,-h,-h]
    , [ 0, 0, 0, 0,-1, 0, 1, 0]
    , [ 0, 0, 0, 0, 0, 0,-1, 1]
    , [-h,-h,-h, h, h, h, h,-h]
    , [ 0, 0, 0,-1, 1, 0, 0, 0]
    ]

-- | This is a basis of E8 as the subset of the even E8 root system
-- where the first two coordinates agree (they are consolidated 
-- into the first coordinate here)
simpleRootsE7_12:: [HalfVec]
simpleRootsE7_12 = roots where
  h = half
  roots =
    [ [-h,-h,-h,-h,-h,-h,-h,-h]
    , [ h, h, h, h, h, h,-h,-h]
    , [ h, h,-h,-h,-h,-h, h, h]
    , [-h,-h, h, h,-h, h, h,-h]
    , [ 0, 0, 0,-1, 1, 0, 0, 0]
    , [ 0, 0,-1, 1, 0, 0, 0, 0]
    , [ 0, 0, 0, 0, 0, 0,-1, 1]
    ]

-- | This is a basis of E7 as the subset of the even E8 root system
-- for which the sum of coordinates sum to zero
simpleRootsE7_diag :: [HalfVec]
simpleRootsE7_diag = roots where
  roots = [ e i - e (i+1) | i <-[2..7] ] ++ [[h,h,h,h,-h,-h,-h,-h]]
  h = half
  n = 8

  e :: Int -> HalfVec
  e i = replicate (i-1) 0 ++ [1] ++ replicate (n-i) 0 

simpleRootsE8_even :: [HalfVec]
simpleRootsE8_even = roots where
  roots = [v1,v2,v3,v4,v5,v7,v8,v6]

  [v1,v2,v3,v4,v5,v6,v7,v8] = roots0
  roots0 = [ e i - e (i+1) | i <-[1..6] ] ++ [ e 6 + e 7 , replicate 8 (-h)  ]
    
  h = half
  n = 8

  e :: Int -> HalfVec
  e i = replicate (i-1) 0 ++ [1] ++ replicate (n-i) 0

simpleRootsE8_odd :: [HalfVec]
simpleRootsE8_odd = roots where
  roots = [ e i - e (i+1) | i <-[1..7] ] ++ [[-h,-h,-h,-h,-h , h,h,h]]
  h = half
  n = 8

  e :: Int -> HalfVec
  e i = replicate (i-1) 0 ++ [1] ++ replicate (n-i) 0 

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