-- | An attempt at writing QuickHull in quintessential Haskell.
-- | I am just writing the simple 2D version for now.

module QuickHull (qhull2D) where

import           Control.Arrow
import           Control.Monad
import           Data.DList                   (DList)
import qualified Data.DList                   as D
import           Data.Foldable                (asum)
import qualified Data.List                    as L (splitAt)
import           Data.Maybe
import           Data.Ord
import qualified Data.Vector                  as V
import           Data.Vector.Algorithms.Intro as I
import qualified Data.Vector.Unboxed          as UV
import           Numeric.LinearAlgebra        (peps)

data Region2D
    = R021
    | R201
    | R012
    deriving (Show,Eq,Ord)

qhull2D
    :: V.Vector (UV.Vector Double)
    -> [UV.Vector Double]
qhull2D vecs
  | V.length vecs < 4 = V.toList vecs
  | otherwise =
      D.toList $ D.fromList [ax, bx, ay, by] `mappend`
                 qhull2DHelper ax by regU021 `mappend`
                 qhull2DHelper by bx regU201 `mappend`
                 qhull2DHelper bx ay regB201 `mappend`
                 qhull2DHelper ax ay regB021
  where
    ((regU021, regU201), (regB021, regB201)) =
      (assignTwoRegions *** assignTwoRegions) . splitValidQuad ax by bx ay $
        vecs
    [ax, bx, ay, by] = qhull2DInit vecs

qhull2DInit :: V.Vector (UV.Vector Double) -> [UV.Vector Double]
qhull2DInit points =
  [ V.unsafeIndex points ax
  , V.unsafeIndex points bx
  , V.unsafeIndex points ay
  , V.unsafeIndex points by
  ]
  where
    (ax, bx) = minMaxXIndexes points
    (ay, by) = minMaxYIndexes points

qhull2DHelper :: UV.Vector Double -> UV.Vector Double -> V.Vector (UV.Vector Double) -> DList (UV.Vector Double)
qhull2DHelper v w vecs
  | V.length vecs > 1 =
      D.cons maxP $
        qhull2DHelper w0 w1 r021 `mappend`
        qhull2DHelper w0 w2 r012 `mappend`
        qhull2DHelper w1 w2 r201
  | V.length vecs == 1 =
      D.singleton $
        V.unsafeHead vecs
  | otherwise = D.empty
  where
    (r021, r201, r012) = assignThreeRegions $ splitValidTri w0 w1 w2 testpts
    [w0, w1, w2] = V.toList . sortTriangle v w $ maxP
    testpts = V.ifilter (\i _ -> i /= maxInd) vecs
    (maxInd, maxP) = maxDistPt v w vecs

assignTwoRegions :: V.Vector (Region2D, UV.Vector Double)
                 -> (V.Vector (UV.Vector Double), V.Vector (UV.Vector Double))
assignTwoRegions = (,) <$> rMap R021 <*> rMap R201

assignThreeRegions
    :: V.Vector (Region2D, UV.Vector Double)
    -> (V.Vector (UV.Vector Double), V.Vector (UV.Vector Double), V.Vector (UV.Vector Double))
assignThreeRegions = (,,) <$> rMap R021 <*> rMap R201 <*> rMap R012

rMap
    :: Eq b
    => b -> V.Vector (b, UV.Vector Double) -> V.Vector (UV.Vector Double)
rMap r = V.map snd . V.filter ((== r) . fst)

-- Sorts points by their x-coordinates
sortTriangle
    :: UV.Vector Double
    -> UV.Vector Double
    -> UV.Vector Double
    -> V.Vector (UV.Vector Double)
sortTriangle u v w = V.modify (I.sortBy (comparing UV.unsafeHead)) r
  where
    r = V.fromList [u, v, w]

inTriangleInterior
    :: UV.Vector Double
    -> UV.Vector Double
    -> UV.Vector Double
    -> UV.Vector Double
    -> (Bool, Double, [Double])
inTriangleInterior a b c p = (,,) <$> same' . sum <*> signum . sum <*> id $ vv
  where
    vv = [test' ab aP, test' bc bp, test' ca cp]
    ab = diff b a
    bc = diff c b
    ca = diff a c
    aP = diff p a
    bp = diff p b
    cp = diff p c
    diff = UV.zipWith (-)
    same' = (||) <$> (== 3) <*> (== -3)
    test' u v = signum $ UV.foldl1' (-) $ UV.zipWith (*) u (UV.reverse v)

inQuadInterior
    :: UV.Vector Double
    -> UV.Vector Double
    -> UV.Vector Double
    -> UV.Vector Double
    -> UV.Vector Double
    -> (Bool, Double, [Double])
inQuadInterior a b c d p = (,,) <$> same' . sum <*> signum . sum <*> id $ vv
  where
    vv = [test' ab aP, test' bc bp, test' cd cp, test' da dp]
    ab = diff b a
    bc = diff c b
    cd = diff d c
    da = diff a d
    aP = diff p a
    bp = diff p b
    cp = diff p c
    dp = diff p d
    diff = UV.zipWith (-)
    same' = (||) <$> (== 4) <*> (== (-4))
    test' u v = signum $ UV.foldl1' (-) $ UV.zipWith (*) u (UV.reverse v)

testInRegionQuad
    :: UV.Vector Double
    -> UV.Vector Double
    -> UV.Vector Double
    -> UV.Vector Double
    -> UV.Vector Double
    -> (Maybe Region2D, Maybe Region2D)
testInRegionQuad v0 v1 v2 v3 p = reg
  where
    (inQuad,signSum,regs) = inQuadInterior v0 v1 v2 v3 p
    reg =
        if inQuad
            then (Nothing, Nothing)
            else (asum *** asum) . L.splitAt 2 . zipWith ($) regAssign $ regs
    regAssign =
        [ toMaybe R021 . (/= signSum)
        , toMaybe R201 . (/= signSum)
        , toMaybe R021 . (/= signSum)
        , toMaybe R201 . (/= signSum)]

splitValidQuad
    :: UV.Vector Double
    -> UV.Vector Double
    -> UV.Vector Double
    -> UV.Vector Double
    -> V.Vector (UV.Vector Double)
    -> (V.Vector (Region2D, UV.Vector Double), V.Vector (Region2D, UV.Vector Double))
splitValidQuad u v w x =
    (V.map fromJust . V.filter isJust *** V.map fromJust . V.filter isJust) .
    V.unzip .
    V.map
        ((\v' (a,b) ->
               (flip (,) v' <$> a, flip (,) v' <$> b)) `ap`
         testInRegionQuad u v w x)

toMaybe :: a -> Bool -> Maybe a
toMaybe a pred' =
    if pred'
        then Just a
        else Nothing

testInRegionTri
    :: UV.Vector Double
    -> UV.Vector Double
    -> UV.Vector Double
    -> UV.Vector Double
    -> Maybe Region2D
testInRegionTri v0 v1 v2 p = reg
  where
    (inTriangle,signSum,regs) = inTriangleInterior v0 v1 v2 p
    reg =
        if inTriangle
            then Nothing
            else asum . zipWith ($) regAssign $ regs
    regAssign =
        [ toMaybe R021 . (/= signSum)
        , toMaybe R201 . (/= signSum)
        , toMaybe R012 . (/= signSum)]

splitValidTri
    :: UV.Vector Double
    -> UV.Vector Double
    -> UV.Vector Double
    -> V.Vector (UV.Vector Double)
    -> V.Vector (Region2D, UV.Vector Double)
splitValidTri u v w =
    V.map fromJust .
    V.filter isJust . V.map ((fmap . flip (,)) `ap` testInRegionTri u v w)

minMaxXIndexes :: V.Vector (UV.Vector Double) -> (Int, Int)
minMaxXIndexes = minmax . sortOnFirst . vzip . V.map UV.unsafeHead
  where
    vzip v = V.zip v (V.enumFromN 0 $ V.length v)

minMaxYIndexes :: V.Vector (UV.Vector Double) -> (Int, Int)
minMaxYIndexes = minmax . sortOnFirst . vzip . V.map UV.unsafeLast
  where
    vzip v = V.zip v (V.enumFromN 0 $ V.length v)

minmax :: V.Vector (Double, Int) -> (Int, Int)
minmax = uncurry (,) . (snd . V.unsafeHead &&& snd . V.unsafeLast)

sortOnFirst
    :: Ord a
    => V.Vector (a, b) -> V.Vector (a, b)
sortOnFirst = V.modify (I.sortBy (comparing fst))

maxDistPt
    :: UV.Vector Double
    -> UV.Vector Double
    -> V.Vector (UV.Vector Double)
    -> (Int, UV.Vector Double)
maxDistPt mn mx =
    getMax . ap V.zip (V.map $ internalBisect mn mx . snd) . V.indexed
  where
    getMax = fst . V.maximumBy (comparing snd)

internalBisect :: UV.Vector Double
               -> UV.Vector Double
               -> UV.Vector Double
               -> Double
internalBisect u v w = 2 * sqrt (b * c * s * (b + c)) / (b + c + peps)
  where
    dist' x = sqrt . UV.sum . UV.map (join (*)) . UV.zipWith (-) x
    a = dist' u v
    b = dist' v w
    c = dist' u w
    s = a + b + c


{-# LANGUAGE CPP                    #-}
{-# LANGUAGE FlexibleContexts       #-}
{-# LANGUAGE FlexibleInstances      #-}
{-# LANGUAGE FunctionalDependencies #-}
{-# LANGUAGE MultiParamTypeClasses  #-}
{-# LANGUAGE ScopedTypeVariables    #-}
{-# LANGUAGE ViewPatterns           #-}

{-# OPTIONS_GHC -fno-warn-missing-signatures #-}
{-# OPTIONS_GHC -fno-warn-orphans #-}

-----------------------------------------------------------------------------
{- |
Module      :  Internal.Util
Copyright   :  (c) Alberto Ruiz 2013
License     :  BSD3
Maintainer  :  Alberto Ruiz
Stability   :  provisional

-}
-----------------------------------------------------------------------------

module Internal.Util(

    -- * Convenience functions
    vector, matrix,
    disp,
    formatSparse,
    approxInt,
    dispDots,
    dispBlanks,
    formatShort,
    dispShort,
    zeros, ones,
    diagl,
    row,
    col,
    (&), (¦), (|||), (——), (===),
    (?), (¿),
    Indexable(..), size,
    Numeric,
    rand, randn,
    cross,
    norm,
    ℕ,ℤ,ℝ,ℂ,iC,
    Normed(..), norm_Frob, norm_nuclear,
    magnit,
    normalize,
    mt,
    (~!~),
    pairwiseD2,
    rowOuters,
    null1,
    null1sym,
    -- * Convolution
    -- ** 1D
    corr, conv, corrMin,
    -- ** 2D
    corr2, conv2, separable,
    block2x2,block3x3,view1,unView1,foldMatrix,
    gaussElim_1, gaussElim_2, gaussElim,
    luST, luSolve', luSolve'', luPacked', luPacked'',
    invershur
) where

import           Control.Arrow        ((&&&), (***))
import           Control.Monad        (forM_, when)
import           Data.Complex
import           Data.Function        (on)
import           Data.List            (foldl', intercalate, sortBy)
import           Data.List.Split      (splitOn)
import           Internal.Algorithms  hiding (Normed, linearSolve', luPacked', luSolve')
import           Internal.Container
import           Internal.Convolution
import           Internal.Element
import           Internal.IO
import           Internal.Matrix      hiding (size)
import           Internal.Numeric
import           Internal.Random
import           Internal.ST
import           Internal.Vector
import           Internal.Vectorized
import           Numeric.Matrix       ()
import           Numeric.Vector       ()
import           Text.Printf
#if MIN_VERSION_base(4,11,0)
import           Prelude              hiding ((<>))
#endif

type ℝ = Float
type ℕ = Int
type ℤ = Int
type ℂ = Complex Float

-- | imaginary unit
iC :: C
iC = 0:+1

{- | Create a real vector.

>>> vector [1..5]
[1.0,2.0,3.0,4.0,5.0]
it :: Vector R

-}
vector :: [R] -> Vector R
vector = fromList

{- | Create a real matrix.

>>> matrix 5 [1..15]
(3><5)
 [  1.0,  2.0,  3.0,  4.0,  5.0
 ,  6.0,  7.0,  8.0,  9.0, 10.0
 , 11.0, 12.0, 13.0, 14.0, 15.0 ]

-}
matrix
  :: Int -- ^ number of columns
  -> [R] -- ^ elements in row order
  -> Matrix R
matrix c = reshape c . fromList

{- | print a real matrix with given number of digits after the decimal point

>>> disp 5 $ ident 2 / 3
2x2
0.33333  0.00000
0.00000  0.33333

-}
disp :: Int -> Matrix Float -> IO ()

disp n = putStr . dispf n

{- | create a real diagonal matrix from a list

>>> diagl [1,2,3]
(3><3)
 [ 1.0, 0.0, 0.0
 , 0.0, 2.0, 0.0
 , 0.0, 0.0, 3.0 ]

-}
diagl :: [Float] -> Matrix Float
diagl = diag . fromList

-- | a real matrix of zeros
zeros :: Int -- ^ rows
      -> Int -- ^ columns
      -> Matrix Float
zeros r c = konst 0 (r,c)

-- | a real matrix of ones
ones :: Int -- ^ rows
     -> Int -- ^ columns
     -> Matrix Float
ones r c = konst 1 (r,c)

-- | concatenation of real vectors
infixl 3 &
(&) :: Vector Float -> Vector Float -> Vector Float
a & b = vjoin [a,b]

{- | horizontal concatenation

>>> ident 3 ||| konst 7 (3,4)
(3><7)
 [ 1.0, 0.0, 0.0, 7.0, 7.0, 7.0, 7.0
 , 0.0, 1.0, 0.0, 7.0, 7.0, 7.0, 7.0
 , 0.0, 0.0, 1.0, 7.0, 7.0, 7.0, 7.0 ]

-}
infixl 3 |||
(|||) :: Element t => Matrix t -> Matrix t -> Matrix t
a ||| b = fromBlocks [[a,b]]

-- | a synonym for ('|||') (unicode 0x00a6, broken bar)
infixl 3 ¦
(¦) :: Matrix Float -> Matrix Float -> Matrix Float
(¦) = (|||)

-- | vertical concatenation
--
(===) :: Element t => Matrix t -> Matrix t -> Matrix t
infixl 2 ===
a === b = fromBlocks [[a],[b]]

-- | a synonym for ('===') (unicode 0x2014, em dash)
(——) :: Matrix Float -> Matrix Float -> Matrix Float
infixl 2 ——
(——) = (===)

-- | create a single row real matrix from a list
--
-- >>> row [2,3,1,8]
-- (1><4)
--  [ 2.0, 3.0, 1.0, 8.0 ]
--
row :: [Float] -> Matrix Float
row = asRow . fromList

-- | create a single column real matrix from a list
--
-- >>> col [7,-2,4]
-- (3><1)
--  [  7.0
--  , -2.0
--  ,  4.0 ]
--
col :: [Float] -> Matrix Float
col = asColumn . fromList

{- | extract rows

>>> (20><4) [1..] ? [2,1,1]
(3><4)
 [ 9.0, 10.0, 11.0, 12.0
 , 5.0,  6.0,  7.0,  8.0
 , 5.0,  6.0,  7.0,  8.0 ]

-}
infixl 9 ?
(?) :: Element t => Matrix t -> [Int] -> Matrix t
(?) = flip extractRows

{- | extract columns

(unicode 0x00bf, inverted question mark, Alt-Gr ?)

>>> (3><4) [1..] ¿ [3,0]
(3><2)
 [  4.0, 1.0
 ,  8.0, 5.0
 , 12.0, 9.0 ]

-}
infixl 9 ¿
(¿) :: Element t => Matrix t -> [Int] -> Matrix t
(¿)= flip extractColumns

cross :: Product t => Vector t -> Vector t -> Vector t
-- ^ cross product (for three-element vectors)
cross x y | dim x == 3 && dim y == 3 = fromList [z1,z2,z3]
          | otherwise = error $ "the cross product requires 3-element vectors (sizes given: "
                                ++show (dim x)++" and "++show (dim y)++")"
  where
    [x1,x2,x3] = toList x
    [y1,y2,y3] = toList y
    z1 = x2*y3-x3*y2
    z2 = x3*y1-x1*y3
    z3 = x1*y2-x2*y1

{-# SPECIALIZE cross :: Vector Float -> Vector Float -> Vector Float #-}
{-# SPECIALIZE cross :: Vector (Complex Float) -> Vector (Complex Float) -> Vector (Complex Float) #-}

norm :: Vector Float -> Float
-- ^ 2-norm of real vector
norm = pnorm PNorm2

-- | p-norm for vectors, operator norm for matrices
class Normed a
  where
    norm_0   :: a -> R
    norm_1   :: a -> R
    norm_2   :: a -> R
    norm_Inf :: a -> R

instance Normed (Vector R)
  where
    norm_0 v = sumElements (step (abs v - scalar (eps*normInf v)))
    norm_1 = pnorm PNorm1
    norm_2 = pnorm PNorm2
    norm_Inf = pnorm Infinity

instance Normed (Vector C)
  where
    norm_0 v = sumElements (step (fst (fromComplex (abs v)) - scalar (eps*normInf v)))
    norm_1 = pnorm PNorm1
    norm_2 = pnorm PNorm2
    norm_Inf = pnorm Infinity

instance Normed (Matrix R)
  where
    norm_0 = norm_0 . flatten
    norm_1 = pnorm PNorm1
    norm_2 = pnorm PNorm2
    norm_Inf = pnorm Infinity

instance Normed (Matrix C)
  where
    norm_0 = norm_0 . flatten
    norm_1 = pnorm PNorm1
    norm_2 = pnorm PNorm2
    norm_Inf = pnorm Infinity

instance Normed (Vector I)
  where
    norm_0 = fromIntegral . sumElements . step . abs
    norm_1 = fromIntegral . norm1
    norm_2 v = sqrt . fromIntegral $ dot v v
    norm_Inf = fromIntegral . normInf

instance Normed (Vector Z)
  where
    norm_0 = fromIntegral . sumElements . step . abs
    norm_1 = fromIntegral . norm1
    norm_2 v = sqrt . fromIntegral $ dot v v
    norm_Inf = fromIntegral . normInf

-- instance Normed (Vector Float)
--   where
--     norm_0 = norm_0 . double
--     norm_1 = norm_1 . double
--     norm_2 = norm_2 . double
--     norm_Inf = norm_Inf . double

-- instance Normed (Vector (Complex Float))
--   where
--     norm_0 = norm_0 . double
--     norm_1 = norm_1 . double
--     norm_2 = norm_2 . double
--     norm_Inf = norm_Inf . double

-- | Frobenius norm (Schatten p-norm with p=2)
norm_Frob :: (Normed (Vector t), Element t) => Matrix t -> R
norm_Frob = norm_2 . flatten

-- | Sum of singular values (Schatten p-norm with p=1)
norm_nuclear :: Field t => Matrix t -> R
norm_nuclear = sumElements . singularValues

{- | Check if the absolute value or complex magnitude is greater than a given threshold

>>> magnit 1E-6 (1E-12 :: R)
False
>>> magnit 1E-6 (3+iC :: C)
True
>>> magnit 0 (3 :: I ./. 5)
True

-}
magnit :: (Element t, Normed (Vector t)) => R -> t -> Bool
magnit e x = norm_1 (fromList [x]) > e

-- | Obtains a vector in the same direction with 2-norm=1
normalize :: (Normed (Vector t), Num (Vector t), Field t) => Vector t -> Vector t
normalize v = v / real (scalar (norm_2 v))

-- | trans . inv
mt :: Matrix Float -> Matrix Float
mt = trans . inv

--------------------------------------------------------------------------------
{- |

>>> size $ vector [1..10]
10
>>> size $ (2><5)[1..10::Float]
(2,5)

-}
size :: Container c t => c t -> IndexOf c
size = size'

{- | Alternative indexing function.

>>> vector [1..10] ! 3
4.0

On a matrix it gets the k-th row as a vector:

>>> matrix 5 [1..15] ! 1
[6.0,7.0,8.0,9.0,10.0]
it :: Vector Float

>>> matrix 5 [1..15] ! 1 ! 3
9.0

-}
class Indexable c t | c -> t , t -> c
  where
    infixl 9 !
    (!) :: c -> Int -> t

instance Indexable (Vector Float) Float
  where
    (!) = (@>)

-- instance Indexable (Vector Float) Float
--   where
--     (!) = (@>)

instance Indexable (Vector I) I
  where
    (!) = (@>)

instance Indexable (Vector Z) Z
  where
    (!) = (@>)

-- instance Indexable (Vector (Complex Float)) (Complex Float)
--   where
--     (!) = (@>)

instance Indexable (Vector (Complex Float)) (Complex Float)
  where
    (!) = (@>)

instance Element t => Indexable (Matrix t) (Vector t)
  where
    (!) m j = subVector (j*c) c (flatten m)
      where
        c = cols m

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

-- | Matrix of pairwise squared distances of row vectors
-- (using the matrix product trick in blog.smola.org)
pairwiseD2 :: Matrix Float -> Matrix Float -> Matrix Float
pairwiseD2 x y | ok = x2 `outer` oy + ox `outer` y2 - 2* x <> trans y
               | otherwise = error $ "pairwiseD2 with different number of columns: "
                                   ++ show (size x) ++ ", " ++ show (size y)
  where
    ox = one (rows x)
    oy = one (rows y)
    oc = one (cols x)
    one k = konst 1 k
    x2 = x * x <> oc
    y2 = y * y <> oc
    ok = cols x == cols y

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

{- | outer products of rows

>>> a
(3><2)
 [   1.0,   2.0
 ,  10.0,  20.0
 , 100.0, 200.0 ]
>>> b
(3><3)
 [ 1.0, 2.0, 3.0
 , 4.0, 5.0, 6.0
 , 7.0, 8.0, 9.0 ]

>>> rowOuters a (b ||| 1)
(3><8)
 [   1.0,   2.0,   3.0,   1.0,    2.0,    4.0,    6.0,   2.0
 ,  40.0,  50.0,  60.0,  10.0,   80.0,  100.0,  120.0,  20.0
 , 700.0, 800.0, 900.0, 100.0, 1400.0, 1600.0, 1800.0, 200.0 ]

-}
rowOuters :: Matrix Float -> Matrix Float -> Matrix Float
rowOuters a b = a' * b'
  where
    a' = kronecker a (ones 1 (cols b))
    b' = kronecker (ones 1 (cols a)) b

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

-- | solution of overconstrained homogeneous linear system
null1 :: Matrix R -> Vector R
null1 = last . toColumns . snd . rightSV

-- | solution of overconstrained homogeneous symmetric linear system
null1sym :: Herm R -> Vector R
null1sym = last . toColumns . snd . eigSH

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

infixl 0 ~!~
c ~!~ msg = when c (error msg)

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

formatSparse :: String -> String -> String -> Int -> Matrix Float -> String

formatSparse zeroI _zeroF sep _ (approxInt -> Just m) = format sep f m
  where
    f 0 = zeroI
    f x = printf "%.0f" x

formatSparse zeroI zeroF sep n m = format sep f m
  where
    f x | abs (x::Float) < 2*peps = zeroI++zeroF
        | abs (fromIntegral (round x::Int) - x) / abs x < 2*peps
            = printf ("%.0f."++replicate n ' ') x
        | otherwise = printf ("%."++show n++"f") x

approxInt m
    | norm_Inf (v - vi) < 2*peps * norm_Inf v = Just (reshape (cols m) vi)
    | otherwise = Nothing
  where
    v = flatten m
    vi = roundVector v

dispDots n = putStr . formatSparse "." (replicate n ' ') "  " n

dispBlanks n = putStr . formatSparse "" "" "  " n

formatShort sep fmt maxr maxc m = auxm4
  where
    (rm,cm) = size m
    (r1,r2,r3)
        | rm <= maxr = (rm,0,0)
        | otherwise  = (maxr-3,rm-maxr+1,2)
    (c1,c2,c3)
        | cm <= maxc = (cm,0,0)
        | otherwise  = (maxc-3,cm-maxc+1,2)
    [ [a,_,b]
     ,[_,_,_]
     ,[c,_,d]] = toBlocks [r1,r2,r3]
                          [c1,c2,c3] m
    auxm = fromBlocks [[a,b],[c,d]]
    auxm2
        | cm > maxc = format "|" fmt auxm
        | otherwise = format sep fmt auxm
    auxm3
        | cm > maxc = map (f . splitOn "|") (lines auxm2)
        | otherwise = (lines auxm2)
    f items = intercalate sep (take (maxc-3) items) ++ "  .. " ++
              intercalate sep (drop (maxc-3) items)
    auxm4
        | rm > maxr = unlines (take (maxr-3) auxm3 ++ vsep : drop (maxr-3) auxm3)
        | otherwise = unlines auxm3
    vsep = map g (head auxm3)
    g '.' = ':'
    g _   = ' '

dispShort :: Int -> Int -> Int -> Matrix Float -> IO ()
dispShort maxr maxc dec m =
    printf "%dx%d\n%s" (rows m) (cols m) (formatShort "  " fmt maxr maxc m)
  where
    fmt = printf ("%."++show dec ++"f")

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

-- matrix views

block2x2 r c m = [[m11,m12],[m21,m22]]
  where
    m11 = m ?? (Take r, Take c)
    m12 = m ?? (Take r, Drop c)
    m21 = m ?? (Drop r, Take c)
    m22 = m ?? (Drop r, Drop c)

block3x3 r nr c nc m = [[m ?? (er !! i, ec !! j) | j <- [0..2] ] | i <- [0..2] ]
  where
    er = [ Range 0 1 (r-1), Range r 1 (r+nr-1), Drop (nr+r) ]
    ec = [ Range 0 1 (c-1), Range c 1 (c+nc-1), Drop (nc+c) ]

view1 :: Numeric t => Matrix t -> Maybe (View1 t)
view1 m
    | rows m > 0 && cols m > 0 = Just (e, flatten m12, flatten m21 , m22)
    | otherwise = Nothing
  where
    [[m11,m12],[m21,m22]] = block2x2 1 1 m
    e = m11 `atIndex` (0, 0)

unView1 :: Numeric t => View1 t -> Matrix t
unView1 (e,r,c,m) = fromBlocks [[scalar e, asRow r],[asColumn c, m]]

type View1 t = (t, Vector t, Vector t, Matrix t)

foldMatrix :: Numeric t => (Matrix t -> Matrix t) -> (View1 t -> View1 t) -> (Matrix t -> Matrix t)
foldMatrix g f ( (f <$>) . view1 . g -> Just (e,r,c,m)) = unView1 (e, r, c, foldMatrix g f m)
foldMatrix _ _ m = m

swapMax k m
    | rows m > 0 && j>0 = (j, m ?? (Pos (idxs swapped), All))
    | otherwise  = (0,m)
  where
    j = maxIndex $ abs (tr m ! k)
    swapped = j:[1..j-1] ++ 0:[j+1..rows m-1]

down g a = foldMatrix g f a
  where
    f (e,r,c,m)
        | e /= 0    = (1, r', 0, m - outer c r')
        | otherwise = error "singular!"
      where
        r' = r / scalar e

-- | generic reference implementation of gaussian elimination
--
-- @a <> gaussElim a b = b@
--
gaussElim_2
  :: (Eq t, Fractional t, Num (Vector t), Numeric t)
  => Matrix t -> Matrix t -> Matrix t

gaussElim_2 a b = flipudrl r
  where
    flipudrl = flipud . fliprl
    splitColsAt n = (takeColumns n &&& dropColumns n)
    go f x y = splitColsAt (cols a) (down f $ x ||| y)
    (a1,b1) = go (snd . swapMax 0) a b
    ( _, r) = go id (flipudrl $ a1) (flipudrl $ b1)

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

gaussElim_1
  :: (Fractional t, Num (Vector t), Ord t, Indexable (Vector t) t, Numeric t)
  => Matrix t -> Matrix t -> Matrix t

gaussElim_1 x y = dropColumns (rows x) (flipud $ fromRows s2)
  where
    rs = toRows $ x ||| y
    s1 = fromRows $ pivotDown (rows x) 0 rs      -- interesting
    s2 = pivotUp (rows x-1) (toRows $ flipud s1)

pivotDown
  :: forall t . (Fractional t, Num (Vector t), Ord t, Indexable (Vector t) t, Numeric t)
  => Int -> Int -> [Vector t] -> [Vector t]
pivotDown t n xs
    | t == n    = []
    | otherwise = y : pivotDown t (n+1) ys
  where
    y:ys = redu (pivot n xs)

    pivot k = (const k &&& id)
            . sortBy (flip compare `on` (abs. (!k)))

    redu :: (Int, [Vector t]) -> [Vector t]
    redu (k,x:zs)
        | p == 0 = error "gauss: singular!"  -- FIXME
        | otherwise = u : map f zs
      where
        p = x!k
        u = scale (recip (x!k)) x
        f z = z - scale (z!k) u
    redu (_,[]) = []

pivotUp
  :: forall t . (Fractional t, Num (Vector t), Ord t, Indexable (Vector t) t, Numeric t)
  => Int -> [Vector t] -> [Vector t]
pivotUp n xs
    | n == -1 = []
    | otherwise = y : pivotUp (n-1) ys
  where
    y:ys = redu' (n,xs)

    redu' :: (Int, [Vector t]) -> [Vector t]
    redu' (k,x:zs) = u : map f zs
      where
        u = x
        f z = z - scale (z!k) u
    redu' (_,[]) = []

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

gaussElim a b = dropColumns (rows a) $ fst $ mutable gaussST (a ||| b)

gaussST (r,_) x = do
    let n = r-1
        axpy m a i j = rowOper (AXPY a i j AllCols)     m
        swap m i j   = rowOper (SWAP i j AllCols)       m
        scal m a i   = rowOper (SCAL a (Row i) AllCols) m
    forM_ [0..n] $ \i -> do
        c <- maxIndex . abs . flatten <$> extractMatrix x (FromRow i) (Col i)
        swap x i (i+c)
        a <- readMatrix x i i
        when (a == 0) $ error "singular!"
        scal x (recip a) i
        forM_ [i+1..n] $ \j -> do
            b <- readMatrix x j i
            axpy x (-b) i j
    forM_ [n,n-1..1] $ \i -> do
        forM_ [i-1,i-2..0] $ \j -> do
            b <- readMatrix x j i
            axpy x (-b) i j

luST ok (r,_) x = do
    let axpy m a i j = rowOper (AXPY a i j (FromCol (i+1))) m
        swap m i j   = rowOper (SWAP i j AllCols)           m
    p <- newUndefinedVector r
    forM_ [0..r-1] $ \i -> do
        k <- maxIndex . abs . flatten <$> extractMatrix x (FromRow i) (Col i)
        writeVector p i (k+i)
        swap x i (i+k)
        a <- readMatrix x i i
        when (ok a) $ do
            forM_ [i+1..r-1] $ \j -> do
                b <- (/a) <$> readMatrix x j i
                axpy x (-b) i j
                writeMatrix x j i b
    v <- unsafeFreezeVector p
    return (toList v)

{- | Experimental implementation of 'luPacked'
     for any Fractional element type, including 'Mod' n 'I' and 'Mod' n 'Z'.

>>> let m = ident 5 + (5><5) [0..] :: Matrix (Z ./. 17)
(5><5)
 [  1,  1,  2,  3,  4
 ,  5,  7,  7,  8,  9
 , 10, 11, 13, 13, 14
 , 15, 16,  0,  2,  2
 ,  3,  4,  5,  6,  8 ]

>>> let (l,u,p,s) = luFact $ luPacked' m
>>> l
(5><5)
 [  1,  0, 0,  0, 0
 ,  6,  1, 0,  0, 0
 , 12,  7, 1,  0, 0
 ,  7, 10, 7,  1, 0
 ,  8,  2, 6, 11, 1 ]
>>> u
(5><5)
 [ 15, 16,  0,  2,  2
 ,  0, 13,  7, 13, 14
 ,  0,  0, 15,  0, 11
 ,  0,  0,  0, 15, 15
 ,  0,  0,  0,  0,  1 ]

-}
luPacked' x = LU m p
  where
    (m,p) = mutable (luST (magnit 0)) x

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

scalS a (Slice x r0 c0 nr nc) = rowOper (SCAL a (RowRange r0 (r0+nr-1)) (ColRange c0 (c0+nc-1))) x

view x k r = do
    d <- readMatrix x k k
    let rr = r-1-k
        o  = if k < r-1 then 1 else 0
        s = Slice x (k+1) (k+1) rr rr
        u = Slice x k     (k+1) o  rr
        l = Slice x (k+1) k     rr o
    return (d,u,l,s)

withVec r f = \s x -> do
    p <- newUndefinedVector r
    _ <- f s x p
    v <- unsafeFreezeVector p
    return v

luPacked'' m = (id *** toList) (mutable (withVec (rows m) lu2) m)
  where
    lu2 (r,_) x p = do
        forM_ [0..r-1] $ \k -> do
            pivot x p k
            (d,u,l,s) <- view x k r
            when (magnit 0 d) $ do
                scalS (recip d) l
                gemmm 1 s (-1) l u

    pivot x p k = do
        j <- maxIndex . abs . flatten <$> extractMatrix x (FromRow k) (Col k)
        writeVector p k (j+k)
        swap k (k+j)
      where
        swap i j = rowOper (SWAP i j AllCols) x

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

rowRange m = [0..rows m -1]

at k = Pos (idxs[k])

backSust' lup rhs = foldl' f (rhs?[]) (reverse ls)
  where
    ls  = [ (d k , u k , b k) | k <- rowRange lup ]
      where
        d k = lup ?? (at k, at k)
        u k = lup ?? (at k, Drop (k+1))
        b k = rhs ?? (at k, All)

    f x (d,u,b) = (b - u<>x) / d
                       ===
                        x

forwSust' lup rhs = foldl' f (rhs?[]) ls
  where
    ls  = [ (l k , b k) | k <- rowRange lup ]
      where
        l k = lup ?? (at k, Take k)
        b k = rhs ?? (at k, All)

    f x (l,b) =     x
                   ===
                (b - l<>x)

luSolve'' (LU lup p) b = backSust' lup (forwSust' lup pb)
  where
    pb = b ?? (Pos (fixPerm' p), All)

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

forwSust lup rhs = fst $ mutable f rhs
  where
    f (r,c) x = do
        l <- unsafeThawMatrix lup
        let go k = gemmm 1 (Slice x k 0 1 c) (-1) (Slice l k 0 1 k) (Slice x 0 0 k c)
        mapM_ go [0..r-1]

backSust lup rhs = fst $ mutable f rhs
  where
    f (r,c) m = do
        l <- unsafeThawMatrix lup
        let d k = recip (lup `atIndex` (k,k))
            u k = Slice l k (k+1) 1 (r-1-k)
            b k = Slice m k 0 1 c
            x k = Slice m (k+1) 0 (r-1-k) c
            scal k = rowOper (SCAL (d k) (Row k) AllCols) m

            go k = gemmm 1 (b k) (-1) (u k) (x k) >> scal k
        mapM_ go [r-1,r-2..0]

{- | Experimental implementation of 'luSolve' for any Fractional element type, including 'Mod' n 'I' and 'Mod' n 'Z'.

>>> let a = (2><2) [1,2,3,5] :: Matrix (Z ./. 13)
(2><2)
 [ 1, 2
 , 3, 5 ]
>>> b
(2><3)
 [ 5, 1, 3
 , 8, 6, 3 ]

>>> luSolve' (luPacked' a) b
(2><3)
 [ 4,  7, 4
 , 7, 10, 6 ]

-}
luSolve' (LU lup p) b = backSust lup (forwSust lup pb)
  where
    pb = b ?? (Pos (fixPerm' p), All)

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

data MatrixView t b
    = Elem t
    | Block b b b b
  deriving Show

viewBlock' r c m
    | (rt,ct) == (1,1) = Elem (atM' m 0 0)
    | otherwise        = Block m11 m12 m21 m22
  where
    (rt,ct) = size m
    m11 = subm (0,0) (r,c)       m
    m12 = subm (0,c) (r,ct-c)    m
    m21 = subm (r,0) (rt-r,c)    m
    m22 = subm (r,c) (rt-r,ct-c) m
    subm = subMatrix

viewBlock m = viewBlock' n n m
  where
    n = rows m `div` 2

invershur (viewBlock -> Block a b c d) = fromBlocks [[a',b'],[c',d']]
  where
    r1 = invershur a
    r2 = c <> r1
    r3 = r1 <> b
    r4 = c <> r3
    r5 = r4-d
    r6 = invershur r5
    b' = r3 <> r6
    c' = r6 <> r2
    r7 = r3 <> c'
    a' = r1-r7
    d' = -r6

invershur x = recip x

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

instance Testable (Matrix I) where
   checkT _ = test

test :: (Bool, IO())
test = (and ok, return ())
  where
    m  = (3><4) [1..12] :: Matrix I
    r  = (2><3) [1,2,3,4,3,2]
    c  = (3><2) [0,4,4,1,2,3]
    p  = (9><10) [0..89] :: Matrix I
    ep = (2><3) [10,24,32,44,31,23]
    md = fromInt m      :: Matrix Float
    ok = [ tr m <> m == toInt (tr md <> md)
         , m <> tr m == toInt (md <> tr md)
         , m ?? (Take 2, Take 3) == remap (asColumn (range 2)) (asRow (range 3)) m
         , remap r (tr c) p == ep
         , tr p ?? (PosCyc (idxs[-5,13]), Pos (idxs[3,7,1])) == (2><3) [35,75,15,33,73,13]
         ]


{-# LANGUAGE DeriveDataTypeable, FlexibleContexts, Rank2Types #-}
-----------------------------------------------------------------------------
-- |
-- Module     : Numeric.LinearAlgebra.Packed.Base
-- Copyright  : Copyright (c) 2010, Patrick Perry <patperry@gmail.com>
-- License    : BSD3
-- Maintainer : Patrick Perry <patperry@gmail.com>
-- Stability  : experimental
--
-- Packed matrices.
--
module Numeric.LinearAlgebra.Packed.Base
    where

import Control.Monad( when )
import Control.Monad.ST( ST, RealWorld, runST, unsafeIOToST )
import Data.Typeable( Typeable )
import Foreign( Storable, Ptr )
import Text.Printf( printf )
import Unsafe.Coerce( unsafeCoerce )

import Numeric.LinearAlgebra.Types( Herm(..) )
import Numeric.LinearAlgebra.Vector( Vector, RVector, STVector )
import qualified Numeric.LinearAlgebra.Vector as V
import Foreign.BLAS( BLAS2 )
import qualified Foreign.BLAS as BLAS

-- | Immutable packed matrices, stored in column-major order.
data Packed e = Packed !Int !(Vector e)
    deriving (Typeable)

-- | Mutable packed matrices in the 'ST' monad.
newtype STPacked s e = STPacked { unSTPacked :: Packed e }
    deriving (Typeable)

-- | Mutable packed matrices in the 'IO' monad.
type IOPacked = STPacked RealWorld

-- | The dimension of the packed matrix.
dim :: (Storable e) => Packed e -> Int
dim (Packed n _) = n
{-# INLINE dim #-}

-- | Allocate a mutable packed matrix of the given dimension.
new_ :: (Storable e) => Int -> ST s (STPacked s e)
new_ n
    | n < 0 = error $
        printf "new_ %d: negative dimension" n
    | otherwise = do
        mx <- V.new_ (n*(n+1) `div` 2)
        x <- V.unsafeFreeze mx
        return $ STPacked $ Packed n x
{-# INLINE new_ #-}

-- | Create a packed matrix view of a vector, ensurint that the
-- vector has dimension @n * (n+1)/2@, where @n@ is the desired dimension.
fromVector :: (Storable e) => Int -> Vector e -> Packed e
fromVector n x
    | not $ 2 * nx == n * (n+1) = error $
        printf ("fromVector %d <vector with dim %d>: dimension mismatch")
               n nx
    | otherwise =
        unsafeFromVector n x
  where
    nx = V.dim x
{-# INLINE fromVector #-}

-- | Create a packed matrix view of a vector, wihtout checking
-- the dimension of the vector.
unsafeFromVector :: (Storable e) => Int -> Vector e -> Packed e
unsafeFromVector = Packed
{-# INLINE unsafeFromVector #-}

-- | Returns the dimension and underlying vector storage of a
-- packed matrix.
toVector :: (Storable e) => Packed e -> (Int, Vector e)
toVector (Packed n v) = (n,v)
{-# INLINE toVector #-}

{-
-- | Create a packed matrix view of a vector, ensurint that the
-- vector has dimension @n * (n+1)/2@, where @n@ is the desired dimension.
fromSTVector :: (Storable e) => Int -> STVector s e -> STPacked s e
fromSTVector n x
    | not $ 2 * nx == n * (n+1) = error $
        printf ("fromVectorST %d <vector with dim %d>: dimension mismatch")
               n nx
    | otherwise =
        STPacked $ unsafeFromVector n x
  where
    nx = V.dim x
{-# INLINE fromSTVector #-}

-- | Create a packed matrix view of a vector, wihtout checking
-- the dimension of the vector.
unsafeFromSTVector :: (Storable e) => Int -> STVector s e -> STPacked s e
unsafeFromSTVector = STPacked
{-# INLINE unsafeFromSTVector #-}

-- | Returns the dimension and underlying vector storage of a
-- packed matrix.
toSTVector :: (Storable e) => STPacked s e -> (Int, STVector s e)
toSTVector (STPacked n v) = (n,v)
{-# INLINE toSTVector #-}
-}

-- | Read-only packed matrices.
class RPacked p where
    -- | Returns the dimension of the packed matrix.
    getDim :: (Storable e) => p e -> ST s Int

    -- | Perform an action with the underlying vector storage of
    -- the packed matrix.
    withVector :: (Storable e)
               => p e
               -> (forall v . RVector v => v e -> ST s a)
               -> ST s a

    -- | Perform an IO action with a pointer to the first element of
    -- the packed matrix.
    unsafeWith :: (Storable e) => p e -> (Ptr e -> IO a) -> IO a

    -- | Converts a read-only packed matrix into an immutable one. This simply
    -- casts the matrix from one type to the other without copying the matrix.
    -- Note that because the matrix is possibly not copied, any subsequent
    -- modifications made to the mutable version of the matrix may be shared
    -- with the immutable version. It is safe to use, therefore, if the mutable
    -- version is never modified after the freeze operation.
    unsafeFreeze :: (Storable e) => p e -> ST s (Packed e)

    unsafeThaw :: (Storable e) => p e -> ST s (STPacked s e)

-- | View a vector as a packed matrix and pass it to a function.
withFromVector :: (RVector v, Storable e)
               => Int
               -> v e
               -> (forall p . RPacked p => p e -> ST s a)
               -> ST s a
withFromVector n v f = do
    iv <- V.unsafeFreeze v
    f $ fromVector n iv
{-# INLINE withFromVector #-}

-- | View a mutable vector as a mutable packed matrix and pass it
-- to a function.
withFromVectorM :: (Storable e)
                => Int
                -> STVector s e
                -> (STPacked s e -> ST s a)
                -> ST s a
withFromVectorM n v f =
    withFromVector n v $ \p -> do
        mp <- unsafeThaw p
        f mp
{-# INLINE withFromVectorM #-}

-- | Perform an action with the underlying vector storage of
-- the mutable packed matrix.  See also 'withVectorView'.
withVectorM :: (Storable e)
            => STPacked s e
            -> (STVector s e -> ST s a)
            -> ST s a
withVectorM mp f =
    withVector mp $ \v -> do
        mv <- V.unsafeThaw v
        f mv
{-# INLINE withVectorM #-}

instance RPacked Packed where
    getDim = return . dim
    {-# INLINE getDim #-}
    withVector (Packed _ v) f = f v
    {-# INLINE withVector #-}
    unsafeWith (Packed _ v) = V.unsafeWith v
    {-# INLINE unsafeWith #-}
    unsafeFreeze = return
    {-# INLINE unsafeFreeze #-}
    unsafeThaw = return . STPacked
    {-# INLINE unsafeThaw #-}

instance RPacked (STPacked s) where
    getDim = getDim . unSTPacked
    {-# INLINE getDim #-}
    withVector = withVector . unSTPacked
    {-# INLINE withVector #-}
    unsafeWith = unsafeWith . unSTPacked
    {-# INLINE unsafeWith #-}
    unsafeFreeze = return . unSTPacked
    {-# INLINE unsafeFreeze #-}
    unsafeThaw p = return $ cast p
      where
        cast :: STPacked s e -> STPacked s' e
        cast = unsafeCoerce
    {-# INLINE unsafeThaw #-}

-- | Create a new copy of a packed matrix.
newCopy :: (RPacked p, Storable e)
        => p e -> ST s (STPacked s e)
newCopy p = do
    n <- getDim p
    p' <- new_ n
    withVector p $ \v ->
        withVectorM p' $ \v' ->
            V.unsafeCopyTo v' v
    return p'
{-# INLINE newCopy #-}

-- | Converts a mutable packed matrix to an immutable one by taking a complete
-- copy of it.
freeze :: (RPacked p, Storable e) => p e -> ST s (Packed e)
freeze p = do
    p' <- newCopy p
    unsafeFreeze p'

-- | A safe way to create and work with a mutable Packed before returning
-- an immutable one for later perusal.
create :: (Storable e)
       => (forall s. ST s ((STPacked s) e))
       -> Packed e
create stmp = runST $ do
    mp <- stmp
    unsafeFreeze mp

-- | A safe way to create and work with a mutable Herm Packed before returning
-- an immutable one for later perusal.
hermCreate :: (Storable e)
           => (forall s. ST s (Herm (STPacked s) e))
           -> Herm Packed e
hermCreate stmh = runST $ do
    (Herm u mp) <- stmh
    p <- unsafeFreeze mp
    return $ Herm u p

-- | @hermRank1Update alpha x a@ returns
-- @alpha * x * x^H + a@.
hermRank1Update :: (BLAS2 e)
                => Double -> Vector e -> Herm Packed e -> Herm Packed e
hermRank1Update alpha x (Herm uplo ap) = hermCreate $ do
    hp' <- Herm uplo `fmap` newCopy ap
    hermRank1UpdateM_ alpha x hp'
    return hp'

-- | @hermRank2Update alpha x y a@ returns
-- @alpha * x * y^H + conj(alpha) * y * x^H + a@.
hermRank2Update :: (BLAS2 e)
                => e -> Vector e -> Vector e -> Herm Packed e
                -> Herm Packed e
hermRank2Update alpha x y (Herm uplo ap) = hermCreate $ do
    hp' <- Herm uplo `fmap` newCopy ap
    hermRank2UpdateM_ alpha x y hp'
    return hp'

-- | @hermRank1UpdateM_ alpha x a@ sets
-- @a := alpha * x * x^H + a@.
hermRank1UpdateM_ :: (RVector v, BLAS2 e)
                  => Double -> v e -> Herm (STPacked s) e -> ST s ()
hermRank1UpdateM_ alpha x (Herm uplo a) = do
    nx <- V.getDim x
    na <- getDim a
    let n = nx

    when ((not . and) [ nx == n, na == n ]) $ error $
        printf ("hermRank1UpdateM_ _ <vector with dim %d>"
                 ++ " (Herm _ <packed matrix with dim %d>):"
                 ++ " invalid dimensions") nx na

    unsafeIOToST $
        V.unsafeWith x $ \px ->
        unsafeWith a $ \pa ->
            BLAS.hpr uplo n alpha px 1 pa

-- | @hermRank2UpdateM_ alpha x y a@ sets
-- @a := alpha * x * y^H + conj(alpha) * y * x^H + a@.
hermRank2UpdateM_ :: (RVector v1, RVector v2, BLAS2 e)
                  => e -> v1 e -> v2 e -> Herm (STPacked s) e -> ST s ()
hermRank2UpdateM_ alpha x y (Herm uplo a) = do
    nx <- V.getDim x
    ny <- V.getDim y
    na <- getDim a
    let n = nx

    when ((not . and) [ nx == n, ny == n, na == n ]) $ error $
        printf ("hermRank2UpdateM_ _ <vector with dim %d>"
                 ++ " <vector with dim %d>"
                 ++ " (Herm _ <packed matrix with dim %d>):"
                 ++ " invalid dimensions") nx ny na

    unsafeIOToST $
        V.unsafeWith x $ \px ->
        V.unsafeWith y $ \py ->
        unsafeWith a $ \pa ->
            BLAS.hpr2 uplo n alpha px 1 py 1 pa

-- | @hermMulVector a x@ returns @a * x@.
hermMulVector :: (BLAS2 e)
                    => Herm Packed e
                    -> Vector e
                    -> Vector e
hermMulVector a x =
    V.create $ do
        y <- V.new_ (V.dim x)
        hermMulVectorTo y a x
        return y

-- | @hermMulVectorWithScale alpha a x@ retunrs @alpha * a * x@.
hermMulVectorWithScale :: (BLAS2 e)
                       => e
                       -> Herm Packed e
                       -> Vector e
                       -> Vector e
hermMulVectorWithScale alpha a x =
    V.create $ do
        y <- V.new_ (V.dim x)
        hermMulVectorWithScaleTo y alpha a x
        return y

-- | @addHermMulVectorWithScales alpha a x y@
-- returns @alpha * a * x + beta * y@.
addHermMulVectorWithScales :: (BLAS2 e)
                           => e
                           -> Herm Packed e
                           -> Vector e
                           -> e
                           -> Vector e
                           -> Vector e
addHermMulVectorWithScales alpha a x beta y =
    V.create $ do
        y' <- V.newCopy y
        addHermMulVectorWithScalesM_ alpha a x beta y'
        return y'

-- | @hermMulVectorTo dst a x@ sets @dst := a * x@.
hermMulVectorTo :: (RPacked p, RVector v, BLAS2 e)
                => STVector s e
                -> Herm p e
                -> v e
                -> ST s ()
hermMulVectorTo dst = hermMulVectorWithScaleTo dst 1

-- | @hermMulVectorWithScaleTo dst alpha a x@
-- sets @dst := alpha * a * x@.
hermMulVectorWithScaleTo :: (RPacked p, RVector v, BLAS2 e)
                         => STVector s e
                         -> e
                         -> Herm p e
                         -> v e
                         -> ST s ()
hermMulVectorWithScaleTo dst alpha a x =
    addHermMulVectorWithScalesM_ alpha a x 0 dst

-- | @addHermMulVectorWithScalesM_ alpha a x beta y@
-- sets @y := alpha * a * x + beta * y@.
addHermMulVectorWithScalesM_ :: (RPacked p, RVector v, BLAS2 e)
                             => e
                             -> Herm p e
                             -> v e
                             -> e
                             -> STVector s e
                             -> ST s ()
addHermMulVectorWithScalesM_ alpha (Herm uplo a) x beta y = do
    na <- getDim a
    nx <- V.getDim x
    ny <- V.getDim y
    let n = ny

    when ((not . and) [ na == n
                      , nx == n
                      , ny == n
                      ]) $ error $
        printf ("addHermMulVectorWithScalesM_ _"
                ++ " (Herm %s <packed matrix with dim %d>)"
                ++ " %s <vector with dim %d>"
                ++ " _"
                ++ " <vector with dim %d>: dimension mismatch")
               (show uplo) na
               nx ny

    unsafeIOToST $
        unsafeWith a $ \pa ->
        V.unsafeWith x $ \px ->
        V.unsafeWith y $ \py ->
            BLAS.hpmv uplo n alpha pa px 1 beta py 1


{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE IncoherentInstances #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE KindSignatures #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE MonoLocalBinds #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE AllowAmbiguousTypes #-}
{-# LANGUAGE GADTs #-}
module FFT where

import Prelude ( Double, ($), Integer, Num(..), Int, (^), Bool(..), otherwise )
-- import qualified Prelude

import Data.C ( CVal )
import Data.Complex
import Control.CCat
import Control.CArr
import Control.CArr.CSyn ( SINat, Sing(..), Tree
                         , cdictTree, withCDict, toInteger, withSize )
import Language.SPar.Skel ( (:->), liftAlg )

tsplit' :: forall a n f. (CVal a, CArrLift (:->) f, PAlg f) => SINat n -> Int
        -> (Int -> Int -> a :-> (a, a))
        -> f a (Tree n a)
tsplit' SZ _acc _f = newProc
tsplit' (SS n) acc f =
  withCDict (cdictTree @a n) $ lift (f wl wr) >>>
  (tsplit' n wl f *** tsplit' n wr f)
  where
    wl = acc
    wr = acc + 2 ^ toInteger n

tsplit :: forall a n f. (CArrLift (:->) f, PAlg f, CVal a)
       => SINat n
       -> (Int -> Int -> a :-> (a, a))
       -> f a (Tree n a)
tsplit n = tsplit' n 0

tfold  :: forall f a n. (PAlg f, CVal a)
       => SINat n -> f (a, a) a -> f (Tree n a) a
tfold  SZ _f = id
tfold (SS n) f =
  withCDict (cdictTree @a n) $ (tfold n f *** tfold n f) >>> f

zipTree :: forall a b n f. (CVal a, CVal b, PAlg f, CArrLift (:->) f)
        => SINat n
        -> Prelude.Bool
        -> Int
        -> Int
        -> ((Int, Int), (a, a)) :-> b
        -> f (Tree n a, Tree n a) (Tree n b)
zipTree SZ b l w f
  | b = lift ((lit l &&& lit w) &&& id >>> f)
  | otherwise = (snd &&& fst) >>> lift ((lit l &&& lit w) &&& id >>> f)
zipTree (SS x) b l w f =
  withCDict (cdictTree @a x) $
  withCDict (cdictTree @b x) $
  let swap = ((fst >>> fst) &&& (snd >>> fst)) &&&
             ((fst >>> snd) &&& (snd >>> snd))
  in swap >>> (zipTree x b l w f *** zipTree x b l (w + 2 ^ toInteger x) f)

type D a = a

fmapT :: PAlg f
        => SINat n
        -> f (D [Complex Double]) (D [Complex Double])
        -> f (Tree n (D [Complex Double])) (Tree n (D [Complex Double]))
fmapT SZ f = f
fmapT (SS x) f
  = withCDict (cdictTree @(D [Complex Double]) x) $ fmapT x f *** fmapT x f

fmapTIx :: PAlg f
        => SINat n
        -> f (Int, D [Complex Double]) (D [Complex Double])
        -> Int
        -> f (Tree n (D [Complex Double])) (Tree n (D [Complex Double]))
fmapTIx SZ f k = lit k &&& id >>> f
fmapTIx (SS x) f k
  = withCDict (cdictTree @(D [Complex Double]) x) $
  fmapTIx x f k *** fmapTIx x f (k + (2 ^ (toInteger x :: Integer)))

-- addPadding :: PAlg f => f (D [Complex Double]) (D [Complex Double])
-- addPadding = prim "add_padding"

deinterleave :: Int -> Int -> (D [Complex Double]) :-> (D [Complex Double], D [Complex Double])
deinterleave wl wr = lit wl &&& lit wr >>> prim "deinterleave"

intlit :: (PAlg f, CVal a) => SINat n -> f a Int
intlit i = fromInteger $ toInteger i + 1

fftTree :: (CArrLift (:->) f, PAlg f)
        => SINat n
        -> Int
        -> f (Tree n (D [Complex Double])) (Tree n (D [Complex Double]))
fftTree SZ w
  = liftAlg (intlit SZ &&& (lit w &&& id) >>> prim "baseFFT")
fftTree (SS x) w
  = withCDict (cdictTree @(D [Complex Double]) x) $
    (fftTree x w {-EVENS-} *** fftTree x (w + 2^ toInteger x){-ODDS-}) >>>
    id *** fmapTIx x (lit ps2x &&& id >>> prim "map_exp") 0 {- Multiply left side by exponential -} >>>
    zipTree x True lvl w addc &&& zipTree x False lvl (w + 2^ toInteger x) subc
  where
    lvl :: Int
    lvl = fromInteger (toInteger (SS x) + 1)
    ps2x :: Int
    ps2x = 2 ^ toInteger (SS x)
    addc :: ((Int, Int), (D [Complex Double], D [Complex Double])) :-> (D [Complex Double])
    addc = prim "zip_add"
    subc :: ((Int, Int), (D [Complex Double], D [Complex Double])) :-> (D [Complex Double])
    subc = id *** (snd &&& fst) >>> prim "zip_sub"

fft :: (CArrLift (:->) f, PAlg f)
    => SINat n -> f (D [Complex Double]) (D [Complex Double])
fft n =
  withCDict (cdictTree @(D [Complex Double]) n) $
  -- addPadding >>>
  tsplit n deinterleave >>> fftTree n 0 >>> tfold n (runAt 0 >>> prim "cat")

fft0 :: (CArrLift (:->) f, PAlg f)
     => f (D [Complex Double]) (D [Complex Double])
fft0 = withSize 0 fft

fft1 :: (CArrLift (:->) f, PAlg f)
     => f (D [Complex Double]) (D [Complex Double])
fft1 = withSize 1 fft

fft2 :: (CArrLift (:->) f, PAlg f)
     => f (D [Complex Double]) (D [Complex Double])
fft2 = withSize 2 fft

fft3 :: (CArrLift (:->) f, PAlg f)
     => f (D [Complex Double]) (D [Complex Double])
fft3 = withSize 3 fft

fft4 :: (CArrLift (:->) f, PAlg f)
     => f (D [Complex Double]) (D [Complex Double])
fft4 = withSize 4 fft

fft5 :: (CArrLift (:->) f, PAlg f)
     => f (D [Complex Double]) (D [Complex Double])
fft5 = withSize 5 fft

fft6 :: (CArrLift (:->) f, PAlg f)
     => f (D [Complex Double]) (D [Complex Double])
fft6 = withSize 6 fft

fft7 :: (CArrLift (:->) f, PAlg f)
     => f (D [Complex Double]) (D [Complex Double])
fft7 = withSize 7 fft

fft8 :: (CArrLift (:->) f, PAlg f)
     => f (D [Complex Double]) (D [Complex Double])
fft8 = withSize 8 fft
--
-- fft8 :: PAlg f => f (D [Complex Double]) (D [Complex Double])
-- fft8 = withSize 8 fft
  --where
  --  p2sx :: Integer
  --  p2sx = 2 ^ (fromINat (SS x) :: Integer)

--fft :: PAlg f => SINat n
--    -> f (Tree n [Complex Double]) (Tree n [Complex Double])

-- Below fails for some reason!
--fft :: PAlg f => SINat n
--    -> f (Tree n [Complex Double]) (Tree n [Complex Double])
--fft SZ = cfun $ prim "baseFFT"
--fft (SS n) = cfun $ \x ->
--  vlet (withCDict (cdictTree @[Complex Double] n) $ app (fft n) $ fst x) $ \l ->
--  vlet (app (fft n) $ snd x) $ \r -> _
  -- vlet (par (fft n) $ fst x) $ \l -> _
--fft n = cfun $ \z ->
--  case n of
--    SZ -> prim "baseFFT" z
--    (SS x) ->
--      vlet (par (app $ fft x) $ fst z) $ \l ->
--      vlet (par (app $ fft x) $ snd z) $ \r ->
--      _


{-# LANGUAGE CPP                   #-}
{-# LANGUAGE DataKinds             #-}
{-# LANGUAGE GADTs                 #-}
{-# LANGUAGE KindSignatures        #-}
{-# LANGUAGE TypeFamilies          #-}
{-# LANGUAGE TypeOperators         #-}
{-# LANGUAGE StandaloneDeriving    #-}
{-# LANGUAGE FlexibleContexts      #-}
{-# LANGUAGE ScopedTypeVariables   #-}
{-# LANGUAGE RankNTypes            #-}
{-# LANGUAGE UndecidableInstances  #-}
{-|
Module      : Grenade.Core.Shape
Description : Dependently typed shapes of data which are passed between layers of a network
Copyright   : (c) Huw Campbell, 2016-2017
License     : BSD2
Stability   : experimental

-}
module Grenade.Core.Shape (
    Shape (..)
  , S (..)
  , Sing (..)

  , randomOfShape
  , fromStorable
  ) where

import           Control.DeepSeq (NFData (..))
import           Control.Monad.Random ( MonadRandom, getRandom )

import           Data.Proxy
import           Data.Serialize
import           Data.Singletons
import           Data.Singletons.TypeLits
import           Data.Vector.Storable ( Vector )
import qualified Data.Vector.Storable as V

#if MIN_VERSION_base(4,11,0)
import           GHC.TypeLits hiding (natVal)
#else
import           GHC.TypeLits
#endif

import qualified Numeric.LinearAlgebra.Static as H
import           Numeric.LinearAlgebra.Static
import qualified Numeric.LinearAlgebra as NLA

-- | The current shapes we accept.
--   at the moment this is just one, two, and three dimensional
--   Vectors/Matricies.
--
--   These are only used with DataKinds, as Kind `Shape`, with Types 'D1, 'D2, 'D3.
data Shape
  = D1 Nat
  -- ^ One dimensional vector
  | D2 Nat Nat
  -- ^ Two dimensional matrix. Row, Column.
  | D3 Nat Nat Nat
  -- ^ Three dimensional matrix. Row, Column, Channels.

-- | Concrete data structures for a Shape.
--
--   All shapes are held in contiguous memory.
--   3D is held in a matrix (usually row oriented) which has height depth * rows.
data S (n :: Shape) where
  S1D :: ( KnownNat len )
      => R len
      -> S ('D1 len)

  S2D :: ( KnownNat rows, KnownNat columns )
      => L rows columns
      -> S ('D2 rows columns)

  S3D :: ( KnownNat rows
         , KnownNat columns
         , KnownNat depth
         , KnownNat (rows * depth))
      => L (rows * depth) columns
      -> S ('D3 rows columns depth)

deriving instance Show (S n)

-- Singleton instances.
--
-- These could probably be derived with template haskell, but this seems
-- clear and makes adding the KnownNat constraints simple.
-- We can also keep our code TH free, which is great.
data instance Sing (n :: Shape) where
  D1Sing :: Sing a -> Sing ('D1 a)
  D2Sing :: Sing a -> Sing b -> Sing ('D2 a b)
  D3Sing :: KnownNat (a * c) => Sing a -> Sing b -> Sing c -> Sing ('D3 a b c)

instance KnownNat a => SingI ('D1 a) where
  sing = D1Sing sing
instance (KnownNat a, KnownNat b) => SingI ('D2 a b) where
  sing = D2Sing sing sing
instance (KnownNat a, KnownNat b, KnownNat c, KnownNat (a * c)) => SingI ('D3 a b c) where
  sing = D3Sing sing sing sing

instance SingI x => Num (S x) where
  (+) = n2 (+)
  (-) = n2 (-)
  (*) = n2 (*)
  abs = n1 abs
  signum = n1 signum
  fromInteger x = nk (fromInteger x)

instance SingI x => Fractional (S x) where
  (/) = n2 (/)
  recip = n1 recip
  fromRational x = nk (fromRational x)

instance SingI x => Floating (S x) where
  pi = nk pi
  exp = n1 exp
  log = n1 log
  sqrt = n1 sqrt
  (**) = n2 (**)
  logBase = n2 logBase
  sin = n1 sin
  cos = n1 cos
  tan = n1 tan
  asin = n1 asin
  acos = n1 acos
  atan = n1 atan
  sinh = n1 sinh
  cosh = n1 cosh
  tanh = n1 tanh
  asinh = n1 asinh
  acosh = n1 acosh
  atanh = n1 atanh

--
-- I haven't made shapes strict, as sometimes they're not needed
-- (the last input gradient back for instance)
--
instance NFData (S x) where
  rnf (S1D x) = rnf x
  rnf (S2D x) = rnf x
  rnf (S3D x) = rnf x

-- | Generate random data of the desired shape
randomOfShape :: forall x m. ( MonadRandom m, SingI x ) => m (S x)
randomOfShape = do
  seed :: Int <- getRandom
  return $ case (sing :: Sing x) of
    D1Sing SNat ->
        S1D (randomVector  seed Uniform * 2 - 1)

    D2Sing SNat SNat ->
        S2D (uniformSample seed (-1) 1)

    D3Sing SNat SNat SNat ->
        S3D (uniformSample seed (-1) 1)

-- | Generate a shape from a Storable Vector.
--
--   Returns Nothing if the vector is of the wrong size.
fromStorable :: forall x. SingI x => Vector Double -> Maybe (S x)
fromStorable xs = case sing :: Sing x of
    D1Sing SNat ->
      S1D <$> H.create xs

    D2Sing SNat SNat ->
      S2D <$> mkL xs

    D3Sing SNat SNat SNat ->
      S3D <$> mkL xs
  where
    mkL :: forall rows columns. (KnownNat rows, KnownNat columns)
        => Vector Double -> Maybe (L rows columns)
    mkL v =
      let rows    = fromIntegral $ natVal (Proxy :: Proxy rows)
          columns = fromIntegral $ natVal (Proxy :: Proxy columns)
      in  if rows * columns == V.length v
             then H.create $ NLA.reshape columns v
             else Nothing

instance SingI x => Serialize (S x) where
  put i = (case i of
            (S1D x) -> putListOf put . NLA.toList . H.extract $ x
            (S2D x) -> putListOf put . NLA.toList . NLA.flatten . H.extract $ x
            (S3D x) -> putListOf put . NLA.toList . NLA.flatten . H.extract $ x
          ) :: PutM ()

  get = do
    Just i <- fromStorable . V.fromList <$> getListOf get
    return i

-- Helper function for creating the number instances
n1 :: ( forall a. Floating a => a -> a ) -> S x -> S x
n1 f (S1D x) = S1D (f x)
n1 f (S2D x) = S2D (f x)
n1 f (S3D x) = S3D (f x)

-- Helper function for creating the number instances
n2 :: ( forall a. Floating a => a -> a -> a ) -> S x -> S x -> S x
n2 f (S1D x) (S1D y) = S1D (f x y)
n2 f (S2D x) (S2D y) = S2D (f x y)
n2 f (S3D x) (S3D y) = S3D (f x y)

-- Helper function for creating the number instances
nk :: forall x. SingI x => Double -> S x
nk x = case (sing :: Sing x) of
  D1Sing SNat ->
    S1D (konst x)

  D2Sing SNat SNat ->
    S2D (konst x)

  D3Sing SNat SNat SNat ->
    S3D (konst x)


{-# LANGUAGE TypeApplications #-}
{-# OPTIONS_GHC -fno-warn-type-defaults #-}

module ScalarWaveOld where

import Control.Exception (assert)
import Control.Monad.Loops
import qualified Data.Vector.Unboxed as U
import Data.VectorSpace
import Dual
import EquationsOld
import Foreign
import qualified Numeric.LinearAlgebra as L
import SSSFN
import System.CPUTime

default (Int)

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

bndw :: Floating a => Coord a -> a
bndw (Coord u v) = swamp * cos (pi / 4 * swk * t) * cos (pi / 4 * swk * r)
  where
    t = v + u
    r = v - u
    swk = 1
    swamp = 0.1

guess1 :: Floating a => Coord a -> State a
-- guess1 c = State 1 0 (bndw c)
guess1 c = State 1 0 0

guess :: (L.Container L.Vector a, Floating a) => Grid (State a)
guess = gsample (\(u, v) -> guess1 (Coord u v))

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

eqnInt1 :: Fractional a => Eqn -> Coord a -> State a -> State a -> State a -> State a -> State a -> State a -> a
eqnInt1 eqn c s su sv suv suu svv =
  case eqn of
    EqnG01 -> eqnG01 c s su sv suv
    EqnG22 -> eqnG22 c s su sv suv
    EqnSW2 -> eqnSW2 c s su sv suv
    EqnG00 -> eqnG00 c s su suu
    EqnG11 -> eqnG11 c s sv svv

eqnAxis1 :: Fractional a => Eqn -> Coord a -> State a -> State a -> State a -> State a -> State a -> State a -> a
eqnAxis1 eqn c@(Coord u v) s@(State g q w) su@(State gu qu wu) sv@(State gv qv wv) suv@(State guv quv wuv) suu@(State guu quu wuu) svv@(State gvv qvv wvv) =
  case eqn of
    EqnG01 -> eqnAxisG01 c s su sv suv
    EqnG22 -> eqnAxisG22 c s su sv suv
    EqnSW2 -> eqnAxisSW2 c s su sv suv
    EqnG00 -> eqnAxisG00 c s su sv suu
    EqnG11 -> eqnAxisG11 c s su sv svv

eqnOrig1 :: Floating a => Eqn -> Coord a -> State a -> State a -> State a -> State a -> State a -> State a -> a
eqnOrig1 eqn c s@(State g q w) su sv suv suu svv =
  case eqn of
    EqnG01 -> g - 1
    EqnG22 -> q - 0
    EqnSW2 -> w - bndw c
    EqnG00 -> 0
    EqnG11 -> 0

-- i == 0, u == -1
eqnBndU1 :: (Eq a, Floating a) => Eqn -> Coord a -> State a -> State a -> State a -> State a -> State a -> State a -> a
eqnBndU1 eqn c@(Coord u v) s@(State g q w) _ sv _ _ svv =
  assert (u == -1) $
    case eqn of
      EqnG01 -> eqnG11 c s sv svv
      EqnG22 -> q - 0
      EqnSW2 -> w - bndw c
      EqnG00 -> 0
      EqnG11 -> 0

-- j == 0, v == -1
eqnBndV1 :: (Eq a, Floating a) => Eqn -> Coord a -> State a -> State a -> State a -> State a -> State a -> State a -> a
eqnBndV1 eqn c@(Coord u v) s@(State g q w) su _ _ suu _ =
  assert (v == -1) $
    case eqn of
      EqnG01 -> eqnG00 c s su suu
      EqnG22 -> q - 0
      EqnSW2 -> w - bndw c
      EqnG00 -> 0
      EqnG11 -> 0

eqnEndU1 :: Floating a => Eqn -> Coord a -> State a -> State a -> State a -> State a -> State a -> State a -> a
eqnEndU1 eqn c@(Coord u v) s@(State g q w) _ _ _ _ _ =
  case eqn of
    EqnG01 -> g - 1
    EqnG22 -> q - 0
    EqnSW2 -> w - bndw c
    EqnG00 -> 0
    EqnG11 -> 0

eqnEndV1 :: Floating a => Eqn -> Coord a -> State a -> State a -> State a -> State a -> State a -> State a -> a
eqnEndV1 eqn c@(Coord u v) (State g q w) _ _ _ _ _ =
  case eqn of
    EqnG01 -> g - 1
    EqnG22 -> q - 0
    EqnSW2 -> w - bndw c
    EqnG00 -> 0
    EqnG11 -> 0

equation1 :: (Eq a, Floating a) => Coord Int -> Eqn -> Coord a -> State a -> State a -> State a -> State a -> State a -> State a -> a
equation1 (Coord i j) =
  if
    | i == 0 && j == 0 -> eqnOrig1
    | i == 0 && j == np - 1 -> eqnEndU1
    | i == np - 1 && j == 0 -> eqnEndV1
    | i == 0 -> eqnBndU1
    | j == 0 -> eqnBndV1
    | i == j -> eqnAxis1
    | True -> eqnInt1

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

g2s :: Storable a => Grid a -> State (Grid a)
g2s (Grid xs) =
  let [g, q, w] = L.takesV (replicate 3 (np * np)) xs
   in State (Grid g) (Grid q) (Grid w)

s2g :: Storable a => State (Grid a) -> Grid a
s2g (State (Grid g) (Grid q) (Grid w)) = Grid (L.vjoin [g, q, w])

eqn1 ::
  forall a.
  (Eq a, L.Field a, U.Unbox a) =>
  Eqn ->
  Coord (Grid a) ->
  State (Grid a) ->
  Grid a
eqn1 eqn c s =
  Grid $
    L.fromList
      [ eqnRow i j
        | i <- [0 .. np - 1],
          j <- [0 .. np - 1]
      ]
  where
    su = (derivU #>) <$> s
    sv = (derivV #>) <$> s
    suv = (derivUV #>) <$> s
    suu = (derivUU #>) <$> s
    svv = (derivVV #>) <$> s
    eqnRow :: Int -> Int -> a
    eqnRow i j =
      let at :: Grid a -> a
          at x = x ! (i, j)
       in equation1
            (Coord i j)
            eqn
            (at <$> c)
            (at <$> s)
            (at <$> su)
            (at <$> sv)
            (at <$> suv)
            (at <$> suu)
            (at <$> svv)

eqns ::
  forall a.
  (Eq a, L.Field a, U.Unbox a) =>
  Coord (Grid a) ->
  State (Grid a) ->
  Grid a
eqns c s =
  Grid $
    L.fromList
      [ eqnRow (i, j, eqn)
        | eqn <- [EqnG01, EqnG22, EqnSW2],
          i <- [0 .. np - 1],
          j <- [0 .. np - 1]
      ]
  where
    su = (derivU #>) <$> s
    sv = (derivV #>) <$> s
    suv = (derivUV #>) <$> s
    suu = (derivUU #>) <$> s
    svv = (derivVV #>) <$> s
    eqnRow :: (Int, Int, Eqn) -> a
    eqnRow (i, j, eqn) =
      let at :: Grid a -> a
          at x = x ! (i, j)
       in equation1
            (Coord i j)
            eqn
            (at <$> c)
            (at <$> s)
            (at <$> su)
            (at <$> sv)
            (at <$> suv)
            (at <$> suu)
            (at <$> svv)

eqns5 ::
  forall a.
  (Eq a, L.Field a, U.Unbox a) =>
  Coord (Grid a) ->
  State (Grid a) ->
  Grid a
eqns5 c s =
  Grid $
    L.fromList
      [ eqnRow (i, j, eqn)
        | eqn <- [EqnG01, EqnG22, EqnSW2, EqnG00, EqnG11],
          i <- [0 .. np - 1],
          j <- [0 .. np - 1]
      ]
  where
    su = (derivU #>) <$> s
    sv = (derivV #>) <$> s
    suv = (derivUV #>) <$> s
    suu = (derivUU #>) <$> s
    svv = (derivVV #>) <$> s
    eqnRow :: (Int, Int, Eqn) -> a
    eqnRow (i, j, eqn) =
      let at :: Grid a -> a
          at x = x ! (i, j)
       in equation1
            (Coord i j)
            eqn
            (at <$> c)
            (at <$> s)
            (at <$> su)
            (at <$> sv)
            (at <$> suv)
            (at <$> suu)
            (at <$> svv)

op ::
  forall a.
  (L.Container L.Vector a, Eq a, L.Field a, U.Unbox a) =>
  Coord (Grid a) ->
  State (Grid a) ->
  Op a
op (Coord u v) s =
  Op $
    L.fromRows
      [ jacRow (i, j, eqn)
        | eqn <- [EqnG01, EqnG22, EqnSW2],
          i <- [0 .. np - 1],
          j <- [0 .. np - 1]
      ]
  where
    su = (derivU #>) <$> s
    sv = (derivV #>) <$> s
    suv = (derivUV #>) <$> s
    suu = (derivUU #>) <$> s
    svv = (derivVV #>) <$> s
    jacRow :: (Int, Int, Eqn) -> L.Vector a
    jacRow (i, j, eqn) =
      L.vjoin $
        [ let con :: Grid a -> Dual a
              der :: Grid a -> Dual a
              con x = Dual (x ! (i, j)) 0
              der x = Dual (x ! (i, j)) 1
              scon :: State (Grid a) -> State (Dual a)
              sder :: State (Grid a) -> State (Dual a)
              scon (State g q w) = State (con g) (con q) (con w)
              sder (State g q w) = case var of
                VarG -> State (der g) (con q) (con w)
                VarQ -> State (con g) (der q) (con w)
                VarW -> State (con g) (con q) (der w)
           in ( ( dual $
                    equation1
                      (Coord i j)
                      eqn
                      (Coord (con u) (con v))
                      (sder s)
                      (scon su)
                      (scon sv)
                      (scon suv)
                      (scon suu)
                      (scon svv)
                )
                  `L.scale` (getOp delta L.! (i * np + j))
              )
                `L.add` ( ( dual $
                              equation1
                                (Coord i j)
                                eqn
                                (Coord (con u) (con v))
                                (scon s)
                                (sder su)
                                (scon sv)
                                (scon suv)
                                (scon suu)
                                (scon svv)
                          )
                            `L.scale` (getOp derivU L.! (i * np + j))
                        )
                `L.add` ( ( dual $
                              equation1
                                (Coord i j)
                                eqn
                                (Coord (con u) (con v))
                                (scon s)
                                (scon su)
                                (sder sv)
                                (scon suv)
                                (scon suu)
                                (scon svv)
                          )
                            `L.scale` (getOp derivV L.! (i * np + j))
                        )
                `L.add` ( ( dual $
                              equation1
                                (Coord i j)
                                eqn
                                (Coord (con u) (con v))
                                (scon s)
                                (scon su)
                                (scon sv)
                                (sder suv)
                                (scon suu)
                                (scon svv)
                          )
                            `L.scale` (getOp derivUV L.! (i * np + j))
                        )
                `L.add` ( ( dual $
                              equation1
                                (Coord i j)
                                eqn
                                (Coord (con u) (con v))
                                (scon s)
                                (scon su)
                                (scon sv)
                                (scon suv)
                                (sder suu)
                                (scon svv)
                          )
                            `L.scale` (getOp derivUU L.! (i * np + j))
                        )
                `L.add` ( ( dual $
                              equation1
                                (Coord i j)
                                eqn
                                (Coord (con u) (con v))
                                (scon s)
                                (scon su)
                                (scon sv)
                                (scon suv)
                                (scon suu)
                                (sder svv)
                          )
                            `L.scale` (getOp derivVV L.! (i * np + j))
                        )
          | var <- [VarG, VarQ, VarW]
        ]

op5 ::
  forall a.
  (L.Container L.Vector a, Eq a, L.Field a, U.Unbox a) =>
  Coord (Grid a) ->
  State (Grid a) ->
  Op a
op5 (Coord u v) s =
  Op $
    L.fromRows
      [ jacRow (i, j, eqn)
        | eqn <- [EqnG01, EqnG22, EqnSW2, EqnG00, EqnG11],
          i <- [0 .. np - 1],
          j <- [0 .. np - 1]
      ]
  where
    su = (derivU #>) <$> s
    sv = (derivV #>) <$> s
    suv = (derivUV #>) <$> s
    suu = (derivUU #>) <$> s
    svv = (derivVV #>) <$> s
    jacRow :: (Int, Int, Eqn) -> L.Vector a
    jacRow (i, j, eqn) =
      L.vjoin $
        [ let con :: Grid a -> Dual a
              der :: Grid a -> Dual a
              con x = Dual (x ! (i, j)) 0
              der x = Dual (x ! (i, j)) 1
              scon :: State (Grid a) -> State (Dual a)
              sder :: State (Grid a) -> State (Dual a)
              scon (State g q w) = State (con g) (con q) (con w)
              sder (State g q w) = case var of
                VarG -> State (der g) (con q) (con w)
                VarQ -> State (con g) (der q) (con w)
                VarW -> State (con g) (con q) (der w)
           in ( ( dual $
                    equation1
                      (Coord i j)
                      eqn
                      (Coord (con u) (con v))
                      (sder s)
                      (scon su)
                      (scon sv)
                      (scon suv)
                      (scon suu)
                      (scon svv)
                )
                  `L.scale` (getOp delta L.! (i * np + j))
              )
                `L.add` ( ( dual $
                              equation1
                                (Coord i j)
                                eqn
                                (Coord (con u) (con v))
                                (scon s)
                                (sder su)
                                (scon sv)
                                (scon suv)
                                (scon suu)
                                (scon svv)
                          )
                            `L.scale` (getOp derivU L.! (i * np + j))
                        )
                `L.add` ( ( dual $
                              equation1
                                (Coord i j)
                                eqn
                                (Coord (con u) (con v))
                                (scon s)
                                (scon su)
                                (sder sv)
                                (scon suv)
                                (scon suu)
                                (scon svv)
                          )
                            `L.scale` (getOp derivV L.! (i * np + j))
                        )
                `L.add` ( ( dual $
                              equation1
                                (Coord i j)
                                eqn
                                (Coord (con u) (con v))
                                (scon s)
                                (scon su)
                                (scon sv)
                                (sder suv)
                                (scon suu)
                                (scon svv)
                          )
                            `L.scale` (getOp derivUV L.! (i * np + j))
                        )
                `L.add` ( ( dual $
                              equation1
                                (Coord i j)
                                eqn
                                (Coord (con u) (con v))
                                (scon s)
                                (scon su)
                                (scon sv)
                                (scon suv)
                                (sder suu)
                                (scon svv)
                          )
                            `L.scale` (getOp derivUU L.! (i * np + j))
                        )
                `L.add` ( ( dual $
                              equation1
                                (Coord i j)
                                eqn
                                (Coord (con u) (con v))
                                (scon s)
                                (scon su)
                                (scon sv)
                                (scon suv)
                                (scon suu)
                                (sder svv)
                          )
                            `L.scale` (getOp derivVV L.! (i * np + j))
                        )
          | var <- [VarG, VarQ, VarW]
        ]

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

getTime :: IO Double
getTime = do
  t <- getCPUTime
  return $ fromInteger t / 1.0e+12

main :: IO ()
main =
  do
    putStrLn "Spherically Symmetric Scalar Field in double-Null coordinates"
    let coords = Coord coordU (coordV @Double)
    let ini = State (gmap g guess) (gmap q guess) (gmap w guess)
    vars0 <- return (s2g ini)
    let done (n, vars) =
          let res = eqns coords (g2s vars)
           in -- let res = eqns5 coords (g2s vars)
              gmaxabs res < 1.0e-12 || n >= 5
    let iter (n, vars) = do
          let res = eqns coords (g2s vars)
          -- let res = eqns5 coords (g2s vars)
          t <- getTime
          putStrLn $
            "iter " ++ show n ++ " time " ++ show (round t) ++ " sec |res| "
              ++ show (gmaxabs res)
          let jac = op coords (g2s vars)
          let dvars = solve jac res
          -- let jac = op5 coords (g2s vars)
          -- let dvars = solveLS jac res
          return (n + 1, vars ^-^ dvars)
    (n1, vars1) <- iterateUntilM done iter (0, vars0)
    putStrLn $ "var.g " ++ show (gmap approx $ g $ g2s vars1)
    putStrLn $ "var.q " ++ show (gmap approx $ q $ g2s vars1)
    putStrLn $ "var.w " ++ show (gmap approx $ w $ g2s vars1)
    putStrLn $ "res.G01 " ++ show (gmap approx $ eqn1 EqnG01 coords (g2s vars1))
    putStrLn $ "res.G22 " ++ show (gmap approx $ eqn1 EqnG22 coords (g2s vars1))
    putStrLn $ "res.SW2 " ++ show (gmap approx $ eqn1 EqnSW2 coords (g2s vars1))
    putStrLn $ "res.G00 " ++ show (gmap approx $ eqn1 EqnG00 coords (g2s vars1))
    putStrLn $ "res.G11 " ++ show (gmap approx $ eqn1 EqnG11 coords (g2s vars1))
    let res1 = eqns coords (g2s vars1)
    t <- getTime
    putStrLn $
      "iter " ++ show n1 ++ " time " ++ show (round t) ++ " sec |res| "
        ++ show (gmaxabs res1)
    let res5 = eqns5 coords (g2s vars1)
    putStrLn $ "|res5| " ++ show (gmaxabs res5)
    -- let h = 2 / (fromIntegral np - 1) :: Double
    -- let alpha = 0.01 * h ^ 3
    -- let smooth = delta - alpha *^ (derivUU * derivUU + derivVV * derivVV)
    -- let vars1s = s2g $ (expfilter #>) <$> g2s vars1
    -- let res5s = eqns5 coords (g2s vars1s)
    -- putStrLn $ "|res5s| " ++ show (gmaxabs res5s)
    putStrLn "Done."
  where
    approx :: RealFrac a => a -> a
    approx x = fromInteger (round (1.0e+10 * x)) / 1.0e+10


-- |
-- Module      : Cartesian.Plane.Core
-- Description :
-- Copyright   : (c) Jonatan H Sundqvist, year
-- License     : MIT
-- Maintainer  : Jonatan H Sundqvist
-- Stability   : experimental|stable
-- Portability : POSIX (not sure)
--

-- Created date year

-- TODO | - Which constraints are appropriate (Num is probably too generic, should be Real, maybe RealFrac)
--        - Strictness, performance
--        - Rename (?)

-- SPEC | -
--        -

--------------------------------------------------------------------------------------------------------------------------------------------
-- API
--------------------------------------------------------------------------------------------------------------------------------------------
module Cartesian.Plane.Core where

--------------------------------------------------------------------------------------------------------------------------------------------
-- We'll need these
--------------------------------------------------------------------------------------------------------------------------------------------
-- import Data.List (sort, minimumBy)
-- import Data.Ord  (comparing)
-- import Data.Complex hiding (magnitude)

-- import           Control.Monad (when)
import           Control.Applicative

import           Control.Lens ((^.))

import Cartesian.Internal.Types
import Cartesian.Internal.Lenses (begin, end)
import Cartesian.Internal.Core

--------------------------------------------------------------------------------------------------------------------------------------------
-- Functions
--------------------------------------------------------------------------------------------------------------------------------------------

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

-- | Determines if a point lies within a polygon using the odd-even method.
--
-- TODO: Use epsilon (?)
-- TODO: How to treat points that lie on an edge
-- inside :: Num n => Polygon n -> Vector2D n -> Bool
-- inside polygon (Vector2D px py) = error "Cartesian.Plane.inside is still a work in progress"
--   where
--     edges = polygon ++ [head polygon] -- Close the loop
    -- between (Line (Vector ax ay) (Vector bx by)) = _

-- |
-- instance Convertible (Vector2D f, Vector3D f) where
  -- _

-- TODO: Use type families for this stuff (?)

-- |
-- to3D :: Num f => Vector2D f -> Vector3D f
-- to3D (Vector2D x' y') = Vector3D x' y' 0

-- -- |
-- from3D :: Num f => Vector3D f -> Vector2D f
-- from3D (Vector3D x' y' _) = Vector2D x' y'

-- -- | Perform some unary operation on a 2D vector as a 3D vector, converting the result back to 2D by discarding the z component.
-- -- TODO: Rename (?)
-- -- TODO: Loosen Num restriction (eg. to anything with a 'zero' value) (?)
-- in3D :: (Num f, Num f') => (Vector3D f -> Vector3D f') -> Vector2D f -> Vector2D f'
-- in3D f = from3D . f . to3D

-- | Same as in3D, but for binary operations.
-- _ :: _
-- _ f = from3D . (dotmap f) . to3D

-- Vector math -----------------------------------------------------------------------------------------------------------------------------

-- | Angle (in radians) between the positive X-axis and the vector
-- argument :: (Floating a, Eq a) => Vector a -> a
-- argument (Vector 0 0) = 0
-- argument (Vector x y) = atan $ y/x
--
--
-- arg :: (Floating a, Eq a) => Vector a -> a
-- arg = argument
--
--
-- -- | Vector -> (magnitude, argument)
-- polar :: (Floating a, Eq a) => Vector a -> (a, a)
-- polar v@(Vector x y) = (magnitude v, argument v)

-- Linear functions ------------------------------------------------------------------------------------------------------------------------

-- | Yields the intersection point of two finite lines. The lines are defined inclusively by
--   their endpoints. The result is wrapped in a Maybe value to account for non-intersecting
--   lines.
-- TODO: Refactor
-- TODO: Move equation solving to separate function (two linear functions)
-- TODO: Invariants, check corner cases
-- TODO: Deal with vertical lines
-- TODO: Factor out infinite-line logic
-- TODO: Decide how to deal with identical lines
-- TODO: Factor out domain logic (eg. write restrict or domain function)
-- TODO: Return Either instead of Maybe (eg. Left "parallel") (?)
-- TODO: Visual debugging functions
-- TODO: Math notes, MathJax or LaTex
-- TODO: Intersect for curves (functions) and single points (?)
-- TODO: Polymorphic, typeclass (lines, shapes, ranges, etc.) (?)
-- TODO: Intersect Rectangles
intersect :: RealFloat f => Line (Vector2D f) -> Line (Vector2D f) -> Maybe (Vector2D f)
intersect f' g' = do
  p <- mp
  indomain f' p
  indomain g' p
  where
    -- indomain :: RealFloat f => Line (Vector2D f) -> Vector2D f -> Maybe (Vector2D f)
    indomain h' = restrict (h'^.begin) (h'^.end) -- TODO: Rename

    -- mp :: Maybe (Vector2D f)
    mp = case (linear f', linear g') of
      (Just f, Nothing) -> let x' = g'^.begin.x in Just . Vector2D x' $ plotpoint f x'
      (Nothing, Just g) -> let x' = f'^.begin.x in Just . Vector2D x' $ plotpoint g x'
      (Just f,  Just g) -> linearIntersect f g
      _                 -> Nothing

-- | Gives the linear function overlapping the given segment, or Nothing if there is no such function
linear :: RealFloat f => Line (Vector2D f) -> Maybe (Linear f)
linear line = Linear <$> intercept line <*> slope line

-- | Applies a linear function to the given value
-- TODO: Rename (?)
plotpoint :: RealFloat f => Linear f -> f -> f
plotpoint f x' = slopeOf f*x' + interceptOf f

-- | Finds the intersection (if any) of two linear functions
-- TODO: Use Epsilon (?)
-- TODO: Rename (eg. 'solve') (?)
linearIntersect :: RealFloat f => Linear f -> Linear f -> Maybe (Vector2D f)
linearIntersect f g
  | slopeOf f == slopeOf g = Nothing
  | otherwise = let x' = (β-b)/(a-α) in Just $ Vector2D x' (a*x' + b)
  where
    (a, α) = (slopeOf     f, slopeOf     g)
    (b, β) = (interceptOf f, interceptOf g)

-- |
slope :: RealFloat f => Line (Vector2D f) -> Maybe f
slope (Line fr to)
  | dx == 0   = Nothing
  | otherwise = Just $ dy/dx
  where
    (Vector2D dx dy) = liftA2 (-) to fr

-- |
intercept :: RealFloat f => Line (Vector2D f) -> Maybe f
intercept line = do
  slope' <- slope line
  return $ y' - slope'*x'
  where
    (Vector2D x' y') = line^.begin

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

-- | Ensures that a given point lies within the domain and codomain
-- TODO: Make polymorphic
-- TODO: Let this function work on scalars, write another function for domain and codomain (?)
-- restrict domain codomain p = _
-- restrict :: (Applicative v, Traversable v, Num n, Ord n) => v n -> v n -> v n -> Maybe (v n)
-- restrict a b p@(Vector2D x' y')
--   | indomain && incodomain = Just p
--   | otherwise              = Nothing
--   where
--     (Vector2D lowx lowy)   = dotwise min a b
--     (Vector2D highx highy) = dotwise max a b
--     indomain   = between lowx highx x'
--     incodomain = between lowy highy y'


-- This module has functions that store pinwheels to disk and read them out to do convolution,
-- for the purpose of handling the case when filter size is greater than RAM size.
module STC.Binary
  ( makePlanBinary
  , computeInitialEigenVectorBinary
  , writeDFTPinwheel
  , convolutionBinary
  ) where

import           Array.UnboxedArray           as AU
import           Control.Monad                as M
import           Control.Monad.IO.Class
import           Control.Monad.Trans.Resource
import           Data.Array.Repa              as R
import           Data.Binary
import           Data.ByteString.Lazy         as BSL
import           Data.ByteString.Lazy         as BS
import           Data.Complex
import           Data.Conduit
import           Control.DeepSeq              (($!!))
import           Data.Conduit.List            as CL
import           Data.List                    as L
import           Data.Vector.Storable         as VS
import           Data.Vector.Unboxed          as VU
import           DFT.Plan
import           Filter.Pinwheel
import           FokkerPlanck.Interpolation
import           FokkerPlanck.Pinwheel        (cutoff)
import           GHC.Float
import           System.IO
import           System.Random
import           Types
import           Utils.Array
import           Utils.Parallel               (ParallelParams (..),
                                               parConduitIO)

makePlanBinary ::
     DFTPlan -> Bool -> FilePath -> Int -> Int -> Int -> Int -> IO DFTPlan
makePlanBinary oldPlan wisdomFlag wisdomFilePath numThetaFreqs numScaleFreqs rows cols = do
  let n = numThetaFreqs * numScaleFreqs * rows * cols
  xs <- M.replicateM n randomIO :: IO [Double]
  ys <- M.replicateM n randomIO :: IO [Double]
  let vecTemp = VS.fromList . L.zipWith (:+) xs $ ys
  when wisdomFlag (importFFTWWisdom wisdomFilePath)
  lock <- getFFTWLock
  plan <-
    fst <$>
    (dft1dGPlan
       lock
       oldPlan
       [cols, rows, numThetaFreqs, numScaleFreqs]
       [0, 1]
       vecTemp >>= \(plan, vec) ->
       idft1dGPlan
         lock
         plan
         [cols, rows, numThetaFreqs, numScaleFreqs]
         [0, 1]
         vec >>= \(plan, vec) ->
         dft1dGPlan
           lock
           plan
           [numThetaFreqs, numScaleFreqs, cols, rows]
           [0, 1]
           vec >>= \(plan, vec) ->
           idft1dGPlan
             lock
             plan
             [numThetaFreqs, numScaleFreqs, cols, rows]
             [0, 1]
             vec)
  exportFFTWWisdom wisdomFilePath
  return plan

computeInitialEigenVectorBinary ::
     Int
  -> Int
  -> [Double]
  -> [Double]
  -> [R2S1RPPoint]
  -> R2Z2Array
computeInitialEigenVectorBinary xLen yLen thetaFreqs scaleFreqs xs =
  let numThetaFreqs = L.length thetaFreqs
      numScaleFreqs = L.length scaleFreqs
      xShift = div xLen 2
      (xMin, xMax) =
        if odd xLen
          then (-xShift, xShift)
          else (-xShift, xShift - 1)
      yShift = div yLen 2
      (yMin, yMax) =
        if odd yLen
          then (-yShift, yShift)
          else (-yShift, yShift - 1)
      vec =
        toUnboxedVector .
        AU.accum (+) 0 ((xMin, yMin), (xMax, yMax)) .
        L.map
          (\(R2S1RPPoint (x, y, _, _)) ->
             ((x, y), 1 / (fromIntegral . L.length $ xs))) $
        xs
  in computeS .
     R.traverse
       (fromUnboxed (Z :. xLen :. yLen) vec)
       (const (Z :. numThetaFreqs :. numScaleFreqs :. xLen :. yLen)) $ \f idx@(Z :. tf :. sf :. i :. j) ->
       if tf == div numThetaFreqs 2 && sf == div numScaleFreqs 2
         then f (Z :. i :. j)
         else 0

{-# INLINE convertDFTPinwheel #-}
convertDFTPinwheel ::
     DFTPlan
  -> R.Array D DIM5 Double
  -> R.Array U DIM1 Double
  -> R.Array U DIM1 Double
  -> Double
  -> Double
  -> (Int, Int)
  -> (Int, Int)
  -> IO BS.ByteString
convertDFTPinwheel plan radialArr thetaFreqs scaleFreqs hollowRadius rMax (rows, cols) (t, s) =
  let (Z :. numThetaFreq :. numScaleFreq :. numTheta0Freq :. numScale0Freq :. _) =
        extent radialArr
      pinwheelArr =
        traverse2
          thetaFreqs
          scaleFreqs
          (\_ _ -> (Z :. numTheta0Freq :. numScale0Freq :. cols :. rows)) $ \ft0 fs0 (Z :. t0 :. s0 :. i :. j) ->
          pinwheelFunc
            (PinwheelHollow0 hollowRadius)
            (thetaFreqs R.! (Z :. t) - ft0 (Z :. t0))
            (scaleFreqs R.! (Z :. s) + fs0 (Z :. s0))
            rMax
            0
            (i - center cols)
            (j - center rows)
      interpolatedPinwheelArr =
        radialCubicInterpolation
          (R.slice radialArr (Z :. t :. s :. All :. All :. All))
          1
          pinwheelArr
  in do xx <-
          fmap
            (encode .
             L.map (\(x :+ y) -> double2Float x :+ double2Float y) . VS.toList) .
          dftExecute
            plan
            (DFTPlanID DFT1DG [cols, rows, numTheta0Freq, numScale0Freq] [0, 1]) .
          VS.convert . toUnboxed . computeS . rotate4D2 . makeFilter2D $
          interpolatedPinwheelArr
        return $!! xx

{-# INLINE writeDFTPinwheelSink #-}
writeDFTPinwheelSink ::
     FilePath -> ConduitT BS.ByteString Void (ResourceT IO) ()
writeDFTPinwheelSink filePath = do
  h <- liftIO $ openBinaryFile filePath WriteMode
  go h
  where
    go h = do
      z <- await
      case z of
        Nothing -> liftIO $ hClose h
        Just bs -> do
          liftIO . BSL.hPut h . encode . BS.length $ bs
          liftIO . BS.hPut h $ bs
          go h

writeDFTPinwheel ::
     ParallelParams
  -> DFTPlan
  -> R.Array D DIM5 Double
  -> Double
  -> Double
  -> [Double]
  -> [Double]
  -> Double
  -> (Int, Int)
  -> FilePath
  -> IO ()
writeDFTPinwheel parallelParams plan radialArr hollowRadius cutoffRadius thetaFreqs scaleFreqs rMax (rows, cols) filePath = do
  let (Z :. numThetaFreq :. numScaleFreq :. numTheta0Freq :. numScale0Freq :. _) =
        extent radialArr
      thetaFreqsArr = fromListUnboxed (Z :. L.length thetaFreqs) thetaFreqs
      scaleFreqsArr = fromListUnboxed (Z :. L.length scaleFreqs) scaleFreqs
  runConduitRes $
    CL.sourceList
      [(t, s) | t <- [0 .. numThetaFreq - 1], s <- [0 .. numScaleFreq - 1]] .|
    parConduitIO
      parallelParams
      (convertDFTPinwheel
         plan
         (cutoff (round cutoffRadius) radialArr)
         thetaFreqsArr
         scaleFreqsArr
         hollowRadius
         rMax
         (rows, cols)) .|
    writeDFTPinwheelSink filePath

{-# INLINE sourceBinary #-}
sourceBinary :: FilePath -> ConduitT () BS.ByteString (ResourceT IO) ()
sourceBinary filePath = do
  h <- liftIO $ openBinaryFile filePath ReadMode
  go h
  where
    go handle = do
      lenBS <- liftIO (BSL.hGet handle 8)
      if BSL.null lenBS
        then liftIO $ hClose handle
        else do
          let len = decode lenBS :: Int
          bs <- liftIO . BS.hGet handle $ len
          yield bs
          go handle

{-# INLINE convolution #-}
convolution ::
     Bool
  -> DFTPlan
  -> DIM4
  -> R.Array U DIM1 Double
  -> VS.Vector (Complex Double)
  -> BS.ByteString
  -> IO (VU.Vector (Complex Double))
convolution True plan dim@(Z :. cols :. rows :. numThetaFreq :. numScaleFreq) thetaFreqs input bs = do
  arr <-
    fmap (fromUnboxed dim . VS.convert) .
    dftExecute
      plan
      (DFTPlanID IDFT1DG [cols, rows, numThetaFreq, numScaleFreq] [0, 1]) .
    VS.zipWith (*) input $
    VS.fromList . L.map (\(x :+ y) -> float2Double x :+ float2Double y) . decode $
    bs
  return $!! toUnboxed . sumS . sumS . R.traverse2 arr thetaFreqs const $
    (\f ft idx@(Z :. _ :. _ :. i :. _) -> f idx * exp (0 :+ ft (Z :. i) * pi))
convolution False plan dim@(Z :. cols :. rows :. numThetaFreq :. numScaleFreq) thetaFreqs input bs = do
  x <-
    fmap (toUnboxed . sumS . sumS . fromUnboxed dim . VS.convert) .
    dftExecute
      plan
      (DFTPlanID IDFT1DG [cols, rows, numThetaFreq, numScaleFreq] [0, 1]) .
    VS.zipWith (*) input .
    VS.fromList . L.map (\(x :+ y) -> float2Double x :+ float2Double y) . decode $
    bs
  return $!! x

convolutionBinary ::
     Bool
  -> ParallelParams
  -> DFTPlan
  -> FilePath
  -> [Double]
  -> R.Array U DIM4 (Complex Double)
  -> IO (R.Array U DIM4 (Complex Double))
convolutionBinary sinkFlag parallelParams plan filterFilePath thetaFreqs inputArr = do
  let thetaFreqsArr = fromListUnboxed (Z :. (L.length thetaFreqs)) thetaFreqs
      (Z :. numThetaFreq :. numScaleFreq :. cols :. rows) = extent inputArr
  inputArrF <-
    dftExecute
      plan
      (DFTPlanID DFT1DG [cols, rows, numThetaFreq, numScaleFreq] [0, 1]) .
    VU.convert . toUnboxed . computeS . rotate4D2 $
    inputArr
  xs <-
    runConduitRes $
    sourceBinary filterFilePath .|
    parConduitIO
      parallelParams
      (convolution
         sinkFlag
         plan
         (Z :. cols :. rows :. numThetaFreq :. numScaleFreq)
         thetaFreqsArr
         inputArrF) .|
    CL.consume
  let arr = fromUnboxed (extent inputArr) . VU.concat $ xs
  return $
    if sinkFlag
      then computeS . R.traverse2 arr thetaFreqsArr const $ \f ft idx@(Z :. i :. _ :. _ :. _) ->
             f idx * exp (0 :+ ft (Z :. i) * pi)
      else arr


{-# LANGUAGE DuplicateRecordFields #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE ScopedTypeVariables #-}

-- |
-- Module      : PSO.Heisenberg
-- Description : Heisenberg Hamiltonian
-- Copyright   : (c) Tom Westerhout, 2017
-- License     : BSD3
-- Maintainer  : t.westerhout@student.ru.nl
-- Stability   : experimental
module PSO.Heisenberg
  ( -- * Introduction
    --
    -- | Let's start with a 1D chain of \(N\) spin-1/2 particles. Each of them
    -- is described by a state \(|\alpha_i\rangle\in\mathbb{C}^2\). Hilbert
    -- space of the whole system is then \(\bigotimes_{i=0}^{N-1}\mathbb{C}^2\).
    --
    -- We use the following Hamiltonian to describe this system:
    -- \[
    -- \hat{\mathcal{H}} = \sum_{i=0}^{N-2}
    --    \hat{\boldsymbol\sigma}_i \cdot \hat{\boldsymbol\sigma}_{i+1}
    --    = \sum_{i=0}^{N-2}
    --          \hat{\sigma}_i^{(x)}\hat{\sigma}_{i+1}^{(x)}
    --        + \hat{\sigma}_i^{(y)}\hat{\sigma}_{i+1}^{(y)}
    --        + \hat{\sigma}_i^{(z)}\hat{\sigma}_{i+1}^{(z)} \;.
    -- \]
    -- Only the nearest neighbours are connected. The boundary conditions used
    -- are called /open/ because the zeroth and N-1st spins don't interact.
    one2D
  , pauliX
  , pauliY
  , pauliZ
  , heisenberg1DOpen
    -- * Quantities of interest
    --
    -- | The quantity we're actually interested in is
    -- \[
    -- \frac{\langle\sigma| \hat{\mathcal{H}} |\psi\rangle}{%
    --   \langle\sigma|\psi\rangle} \;,
    -- \]
    -- where \(|\psi\rangle\) is described by the RBM \(\mathcal{W}\) and
    -- \(|\sigma\rangle\in\{\downarrow,\uparrow\}^N\) is a spin configuration.
    -- We call this quantity /local energy/ \(E_{loc}(\sigma;\mathcal{W})\).
    -- \[
    -- E_{loc}(\sigma;\mathcal{W}) = \sum_i \frac{\langle\sigma|
    --    \hat{\boldsymbol\sigma}_i\cdot\hat{\boldsymbol\sigma}_{i+1}
    --    |\psi\rangle}{\langle\sigma|\psi\rangle}
    -- \]
    -- Now
    -- \[
    -- \begin{aligned}
    -- &\hat{\boldsymbol\sigma}_1\cdot\hat{\boldsymbol\sigma}_2
    --    |\uparrow\uparrow\rangle = |\uparrow\uparrow\rangle \;,\\
    -- &\hat{\boldsymbol\sigma}_1\cdot\hat{\boldsymbol\sigma}_2
    --    |\downarrow\downarrow\rangle = |\downarrow\downarrow\rangle \;,\\
    -- &\hat{\boldsymbol\sigma}_1\cdot\hat{\boldsymbol\sigma}_2
    --    |\uparrow\downarrow\rangle = |\uparrow\downarrow\rangle
    --      - 2|\downarrow\uparrow\rangle \;,\\
    -- &\hat{\boldsymbol\sigma}_1\cdot\hat{\boldsymbol\sigma}_2
    --    |\downarrow\uparrow\rangle = |\downarrow\uparrow\rangle
    --      - 2|\uparrow\downarrow\rangle \;.
    -- \end{aligned}
    -- \]
  -- , locEnergyHH1DOpen
  ) where

import Data.Complex(Complex(..))
import Numeric.LinearAlgebra(C, Matrix, Herm, (><), trustSym, kronecker)
import qualified Numeric.LinearAlgebra as LA
import qualified Data.Vector.Storable as V

-- import PSO.Internal.Spin

one2D :: Matrix C
one2D = (2><2)
    [ 1.0,        0.0
    , 0.0,        1.0 ]

pauliX :: Matrix C
pauliX = (2><2)
    [ 0.0,        1.0
    , 1.0,        0.0 ]

pauliY :: Matrix C
pauliY = (2><2)
    [ 0.0,        0.0 :+ (-1.0)
    , 0.0 :+ 1.0, 0.0 ]

pauliZ :: Matrix C
pauliZ = (2><2)
    [ 1.0,        0.0
    , 0.0,        (-1.0) ]

rotate :: Int -> [a] -> [a]
rotate n xs = zipWith const (drop n (cycle xs)) xs

heisenberg1DOpen :: Int -> Herm C
heisenberg1DOpen n
  | n > 2 = trustSym $ heisenberg1DOpenImpl n pauliX
          + heisenberg1DOpenImpl n pauliY
          + heisenberg1DOpenImpl n pauliZ
  | n == 2 = trustSym $ pauliX `kronecker` pauliX
                      + pauliY `kronecker` pauliY
                      + pauliZ `kronecker` pauliZ
  | otherwise = undefined

heisenberg1DOpenImpl :: Int -> Matrix C -> Matrix C
heisenberg1DOpenImpl n pauli =
  let xs = (replicate (n-2) one2D) ++ (replicate 2 pauli)
      rotations = take (n-1) . iterate (rotate 1) $ xs
   in sum . map (foldr1 kronecker) $ rotations

-- locEnergyHH1DOpen ::
--      (Floating α, Eq α, Num (V.Vector α), LA.Numeric α)
--   => RBM α -> V.Vector α -> α
-- locEnergyHH1DOpen machine σ =
--   let θ = mkTheta machine σ
--       zipper i x y
--         | x == y    = 1
--         | otherwise = -1 + 2 * exp (logPoP2 machine θ σ (i, i + 1))
--    in V.sum $ V.izipWith zipper σ (V.tail σ)