-- | Interface to the FFTW routine
--
-- Ubuntu installation of C library:
--
--     sudo apt-get install libfftw3-dev
--
-- Basic usage:
-- <http://www.fftw.org/fftw3_doc/Complex-One_002dDimensional-DFTs.html#Complex-One_002dDimensional-DFTs>
--
-- Linker flags:
--
--     cc file.c -lm -lfftw3

module FFTW where

import qualified Prelude as P

import qualified Data.Complex as Complex

import qualified Test.QuickCheck as QC
import qualified Test.QuickCheck.Monadic as QC

import Language.Embedded.Expression (varExp)
  -- `varExp` can be used to make named constants in expressions

import Feldspar.Run
import Feldspar.Data.Vector

import DFT
import FFT_bench (printTime_def)

-- | Wrapper for the FFTW routine
--
-- This wrapper is mainly used for testing. It's not suitable for real code
-- because it constructs a plan every time it's called, and it silently
-- allocates an array for the output.
fftw :: DManifest (Complex Double) -> Run (DManifest (Complex Double))
fftw inp = do
    addInclude "<fftw3.h>"
    out  <- newArr (length inp)
    plan <- newObject "fftw_plan" False
    callProcAssign plan "fftw_plan_dft_1d"
      [ valArg (length inp)
      , iarrArg inp
      , arrArg out
      , valArg (varExp "FFTW_FORWARD"  :: Data Word32)
      , valArg (varExp "FFTW_ESTIMATE" :: Data Word32)
      ]
    callProc "fftw_execute" [objArg plan]
    freezeArr out

almostEq a b
    =    Complex.magnitude d P.< 1e-7
    P.&& Complex.phase d     P.< 1e-7
  where
    d = abs (a-b)

a ~= b = P.and $ P.zipWith almostEq a b

prop_fft_dft dft' fft' = QC.monadicIO $ do
    n   :: Int              <- QC.pick $ QC.choose (2,5)
    inp :: [Complex Double] <- QC.pick $ QC.vector (2 P.^ n)
    outd <- QC.run $ dft' inp
    outf <- QC.run $ fft' inp
    QC.assert (outd ~= outf)

-- | Compare 'fftw' against 'dft'
testFFTW =
    marshalledM (return . dft) $ \dft'  ->
    marshalledM fftw           $ \fftw' ->
      QC.quickCheck $ prop_fft_dft dft' fftw'
  where
    marshalledM = marshalled' def def
      { externalFlagsPre  = ["-Wno-incompatible-pointer-types"]
      , externalFlagsPost = ["-lm -lfftw3"]
      }

sizeOf_fftw_complex :: Data Length
sizeOf_fftw_complex = 16
  -- Checked on an x86_64 system

-- | Measure the time for 100 runs of 'fftw' (excluding initialization) for
-- arrays of the given size
benchmark n = do
    addInclude "<stdio.h>"
    addInclude "<string.h>"
    addInclude "<time.h>"
    addInclude "<fftw3.h>"

    addDefinition printTime_def

    inp  <- newObject "fftw_complex" True
    out  <- newObject "fftw_complex" True
    plan <- newObject "fftw_plan" False

    callProcAssign inp "fftw_malloc" [valArg (n*sizeOf_fftw_complex)]
    callProcAssign out "fftw_malloc" [valArg (n*sizeOf_fftw_complex)]
    callProc "memset"
      [ objArg inp
      , valArg (0 :: Data Index)
      , valArg (n*sizeOf_fftw_complex)
      ]

    callProcAssign plan "fftw_plan_dft_1d"
      [ valArg (n :: Data Word32)
      , objArg inp
      , objArg out
      , valArg (varExp "FFTW_FORWARD"  :: Data Word32)
      , valArg (varExp "FFTW_ESTIMATE" :: Data Word32)
          -- Change to `FFTW_MEASURE` to enable tuning
      ]

    start <- newObject "clock_t" False
    end   <- newObject "clock_t" False
    callProcAssign start "clock" []

    for (0,1,Excl 100) $ \(_ :: Data Index) ->
      callProc "fftw_execute" [objArg plan]

    callProcAssign end "clock" []
    callProc "printTime" [objArg start, objArg end]

    callProc "fftw_destroy_plan" [objArg plan]
    callProc "fftw_free" [objArg inp]
    callProc "fftw_free" [objArg out]

runBenchmark n = runCompiled'
    def
    def {externalFlagsPre = ["-O3"], externalFlagsPost = ["-lm","-lfftw3"]}
    (benchmark n)


{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE Strict #-}
module STC.CompletionField where

import           Control.Monad.Parallel as MP
import           Data.Array.Repa        as R
import           Data.Complex
import           Data.List              as L
import           Data.Vector.Storable   as VS
import           DFT.Plan
import           Filter.Utils
import           STC.Convolution
import           STC.DFTArray
import           STC.Plan
import           Utils.Array
import           Utils.Parallel

{-# INLINE completionField #-}
completionField :: DFTPlan -> DFTArray -> DFTArray -> IO DFTArray
completionField plan source@(DFTArray rows cols thetaFreqs rFreqs _) sink = do
  let numThetaFreq = L.length thetaFreqs
      numRFreq = L.length rFreqs
      sourceFilter =
        VS.convert .
        toUnboxed .
        computeUnboxedS .
        R.backpermute
          (Z :. numRFreq :. numThetaFreq :. cols :. rows)
          (\(Z :. a :. b :. c :. d) ->
             (Z :. (makeFilterHelper numRFreq a) :.
              (makeFilterHelper numThetaFreq b) :.
              c :.
              d)) .
        pad [rows, cols, numThetaFreq, numRFreq] 0 .
        dftArrayToRepa $
        source
      dftID = DFTPlanID DFT1DG [numRFreq, numThetaFreq, cols, rows] [0, 1]
  dftSource <- dftExecute plan dftID sourceFilter
  dftSink <- dftExecute plan dftID . VS.concat . getDFTArrayVector $ sink
  arr <-
    fmap
      (fromUnboxed (Z :. numRFreq :. numThetaFreq :. cols :. rows) . VS.convert) .
    dftExecute
      plan
      (DFTPlanID IDFT1DG [numRFreq, numThetaFreq, cols, rows] [0, 1]) .
    VS.zipWith (*) dftSource $
    dftSink
  return $ repaToDFTArray thetaFreqs rFreqs $ arr

{-# INLINE completionField' #-}
completionField' :: DFTPlan -> DFTArray -> DFTArray -> IO DFTArray
completionField' plan source@(DFTArray rows cols thetaFreqs rFreqs _) sink = do
  let numThetaFreq = L.length thetaFreqs
      sourceFilter =
        VS.convert .
        toUnboxed .
        computeUnboxedS .
        R.backpermute
          (Z :. (1 :: Int) :. numThetaFreq :. cols :. rows)
          (\(Z :. _ :. b :. c :. d) ->
             (Z :. (0 :: Int) :. (makeFilterHelper numThetaFreq b) :. c :. d)) .
        dftArrayToRepa $
        source
      dftID = DFTPlanID DFT1DG [1, numThetaFreq, cols, rows] [0, 1]
  dftSource <- dftExecute plan dftID sourceFilter
  dftSink <- dftExecute plan dftID . VS.concat . getDFTArrayVector $ sink
  arr <-
    fmap
      (fromUnboxed (Z :. (1 :: Int) :. numThetaFreq :. cols :. rows) .
       VS.convert) .
    dftExecute plan (DFTPlanID IDFT1DG [1, numThetaFreq, cols, rows] [0, 1]) .
    VS.zipWith (*) dftSource $
    dftSink
  return $ repaToDFTArray thetaFreqs rFreqs $ arr

{-# INLINE completionFieldRepa #-}
completionFieldRepa ::
     (R.Source s1 (Complex Double), R.Source s2 (Complex Double))
  => DFTPlan
  -> R.Array s1 DIM4 (Complex Double)
  -> R.Array s2 DIM4 (Complex Double)
  -> IO (R.Array U DIM4 (Complex Double))
completionFieldRepa plan source sink = do
  print . extent $ source
  let (Z :. numRFreq :. numThetaFreq :. cols :. rows) = extent source
      dftID = DFTPlanID DFT1DG [numRFreq, numThetaFreq, cols, rows] [0, 1]
  sourceFilter <-
    fmap (VS.convert . toUnboxed) .
    computeUnboxedP .
    R.backpermute
      (extent source)
      (\(Z :. a :. b :. c :. d) ->
         Z :. makeFilterHelper numRFreq a :. makeFilterHelper numThetaFreq b :.
          c :.
          d) $
    source
  dftSource <- dftExecute plan dftID sourceFilter
  sinkU <- computeUnboxedP . delay $ sink
  dftSink <- dftExecute plan dftID . VS.convert . toUnboxed $ sinkU
  fmap (fromUnboxed (extent source) . VS.convert) .
    dftExecute
      plan
      (DFTPlanID IDFT1DG [numRFreq, numThetaFreq, cols, rows] [0, 1]) .
    VS.zipWith (*) dftSource $
    dftSink

-- {-# INLINE completionField #-}
-- completionField :: DFTPlan -> DFTArray -> DFTArray -> IO DFTArray
-- completionField plan source@(DFTArray rows cols thetaFreqs rFreqs _) sink = do
--   let numThetaFreq = L.length thetaFreqs
--       numRFreq = L.length rFreqs
--       sourceFilter =
--         VS.convert .
--         toUnboxed .
--         computeUnboxedS .
--         R.backpermute
--           (Z :. numRFreq :. numThetaFreq :. cols :. rows)
--           (\(Z :. a :. b :. c :. d) ->
--              (Z :. (makeFilterHelper numRFreq a) :.
--               (makeFilterHelper numThetaFreq b) :.
--               c :.
--               d)) .
--         dftArrayToRepa $
--         source
--       dftID = DFTPlanID DFT1DG [numRFreq, numThetaFreq, cols, rows] [0, 1]
--   dftSource <- dftExecute plan dftID $ sourceFilter
--   dftSink <- dftExecute plan dftID . VS.concat . getDFTArrayVector $ sink
--   arr <-
--     fmap
--       (fromUnboxed (Z :. numRFreq :. numThetaFreq :. cols :. rows) . VS.convert) .
--     dftExecute
--       plan
--       (DFTPlanID IDFT1DG [numRFreq, numThetaFreq, cols, rows] [0, 1]) .
--     VS.zipWith (*) dftSource $
--     dftSink
--   return $ repaToDFTArray thetaFreqs rFreqs $ arr


#!/usr/bin/env stack
-- stack --resolver lts-15.04 runghc --package reanimate
{-# LANGUAGE OverloadedStrings #-}
module Main where

import           Codec.Picture.Types
import           Control.Lens                  ()
import           Control.Monad
import           Data.Function
import           Data.List
import           Data.List.NonEmpty            (NonEmpty)
import qualified Data.List.NonEmpty            as NE
import           Data.Maybe
import           Data.Ratio
import qualified Data.Text                     as T
import           Data.Tuple
import qualified Data.Vector                   as V
import           Debug.Trace
import           Linear.Matrix                 hiding (trace)
import           Linear.Metric
import           Linear.V2
import           Linear.V3
import           Linear.Vector
import           Numeric.LinearAlgebra         hiding (polar, scale, (<>))
import qualified Numeric.LinearAlgebra         as Matrix
import           Numeric.LinearAlgebra.HMatrix hiding (polar, scale, (<>))
import           Reanimate
import           Reanimate.Animation
import           Reanimate.Math.Balloon
import           Reanimate.Math.Common
import           Reanimate.Math.Triangulate
import           Reanimate.Math.Polygon
import           Reanimate.Math.EarClip
import           Reanimate.Math.SSSP
import           Reanimate.Math.Render
import           Reanimate.Math.Visibility
import           Reanimate.Math.Compatible
import           Reanimate.Morph.Common
import           Reanimate.PolyShape           (svgToPolygons)

p1 = pScale 2 $ pAtCenter $ pAddPoints (0+2) (pSetOffset shape13 0)
-- p1 = pSetOffset (addPoints 2 shape13) 0
-- p2 = pScale 0.5 $ centerPolygon shape20
-- p1 = centerPolygon shape2
p2 = pScale 2 $ pAtCenter $ pAddPoints 0 (pSetOffset shape14 0)

main :: IO ()
main = reanimate $ sceneAnimation $ do
  bg <- newSpriteSVG $ mkBackground "black"
  spriteZ bg (-1)
  -- play $ drawVisibleFrom triangle
  -- play $ drawVisibleFrom shape1
  -- play $ drawVisibleFrom shape2
  -- play $ drawVisibleFrom shape3
  -- play $ drawVisibleFrom shape4
  -- play $ drawVisibleFrom shape5
  -- play $ drawVisibleFrom shape6
  -- play $ drawOverlap shape20
  -- play $ drawSSSP triangle naive
  -- play $ drawSSSPFast shape5
  -- fork $ play $ mapA (translate (-3) 0) $ drawSSSPVisibilityFast shape2
  -- fork $ play $ mapA (translate (3) 0) $ drawSSSPVisibilityFast shape7
  -- play $ drawOverlap $ fst $ split1Link (fst (split1Link shape7 3 7)) 0 2
  -- fork $ play $ mapA (translate (-4) 0) $ drawSSSPVisibilityFast $ pSetOffset (addPoints 0 shape14) 0
  -- fork $ play $ mapA (translate (4) 0) $ drawSSSPVisibilityFast $ pSetOffset (addPoints 2 shape13) 0
  fork $ play $ staticFrame 1 $
    translate (4) 0 $ mkGroup
    [ withFillColor "grey" $ polygonShape p1
    , polygonNumDots p1 ]
  fork $ play $ staticFrame 1 $
    let -- pOrigin = (addPoints 10 $ pSetOffset shape14 0)
        pOrigin = pScale 2 $ pAtCenter $ pAddPoints (0+2) (pSetOffset shape13 0)
        --pOrigin = (addPoints 2 shape13)
        p1 = pScale 1 $ pSetOffset pOrigin 0
        p2 = snd $ split1Link p1 1 11 0
        p3 = fst $ split1Link p2 0 2 1
        -- p4 = fst $ split2Link p3 2 5
        p4 = snd $ split1Link p3 0 9 1
        p5 = fst $ split1Link p4 1 3 0
        p6 = fst $ split2Link p5 0 2
        p7 = fst $ split1Link p6 1 3 0
        p8 = fst $ split2Link p7 0 2
        p9 = fst $ split1Link p7 1 3 0
        p = p4
    in mkGroup
    [ withFillColor "grey" $ polygonShape p
    , polygonNumDots p ]
  -- newSpriteSVG $
  --   let p1 = pSetOffset shape14 0
  --       V2 x y = realToFrac <$> steiner2Link p1 2 6
  --   in translate (-4) 0 $
  --     translate x y $ withFillColor "red" $
  --     mkCircle 0.1
  -- newSpriteSVG $
  --   let p1 = pSetOffset shape14 0
  --   in translate (-4) 0 $
  --     mkGroup
  --     [ mkGroup
  --       [ withStrokeColor "purple" $
  --         mkLinePath [(x1,y1),(x2,y2)]
  --       , translate x2 y2 $ withFillColor "red" $ mkCircle 0.1 ]
  --     | (eA, eB) <- take 1 $ steiner2Edges p1 4 11
  --     , let V2 x1 y1 = realToFrac <$> eA
  --           V2 x2 y2 = realToFrac <$> eB
  --     ]
  -- play $ staticFrame 1 $
  --   let p1 = pSetOffset shape14 0
  --   in mkGroup
  --   [ withFillColor "grey" $ polygonShape p1
  --   , polygonNumDots p1
  --   , drawWindowOverlap p1 1 6
  --   -- , withStrokeColor "red" $ drawWindow (pSetOffset p1 2)
  --   -- , withStrokeColor "blue" $ drawWindow (pSetOffset p1 6)
  --   ]
  -- play $ drawCompatible (pSetOffset (addPoints 2 shape13) 0) (pSetOffset shape14 0)

  -- play $ drawSSSP shape2 naive
  -- play $ drawSSSP shape3 naive
  -- play $ drawSSSP shape4 naive
  -- play $ drawSSSP shape5 naive
  -- play $ drawSSSP (pScale 0.5 $ winding 10) (\p -> sssp p (dual (earClip p)))
  -- play $ drawSSSP shape1 (\p -> sssp p (dual (earClip p)))
  -- play $ drawTriangulation (pCycle shape5 0.2910962834555265) earClip'
  -- play $ drawTriangulation shape5 earClip'
  -- play $ mkAnimation 1 $ \t ->
  --   let p = shape4
  --   in polygonNumDots (pCycle p t)
  -- play $ setDuration 20 $ drawSSSPVisibility $ pScale 1 $ shape7
  -- let shapeI = head $ svgToPolygons 0.1 $ scale 8 $ center $ latex "I"
  -- play $ animate $ \_ -> withFillColor "grey" $ polygonNumDots shapeI
  -- play $ drawTriangulation (pScale 0.5 $ winding 10) earClip'
  -- play $ staticFrame 1 $
  --   mkGroup
  --   [ withFillColor "grey" $ polygonShape shape13
  --   , withFillColor "grey" $ polygonDots shape13 ]
  -- let p = balloonP origin
  --     origin = centerPolygon $ pScale 1 $ shiftLongestDiameter shapeI
  --     mkB = balloon (scale 8 $ center $ latex "C")
  --     inf = unsafeSVGToPolygon 0.1 $ (scale 8 $ center $ latex "$\\infty$")
  --     cShape = shiftLongestDiameter $ unsafeSVGToPolygon 0.01 $ (scale 8 $ center $ latex "C")
  -- _ <- newSpriteSVG $
  --   translate (0) (3) $ withFillColor "white" $
  --   center $ latex $ T.pack $ show (isSimple inf)
  -- play $ pauseAtEnd 1 $ mkAnimation 3 $ \t ->
  --   let inflated = p t in
  --   translate (0) (0) $ withFillColor "white" $ withStrokeColor "red" $
  --   withStrokeWidth (defaultStrokeWidth*0) $
    -- polygonShape inf
    -- renderTriangulation inf (earClip inf)
    -- mkB t
    -- scale 3 $ -- translate (-2.2) (-1) $
    -- mkGroup
    -- [ mkGroup []
    -- , withStrokeWidth (defaultStrokeWidth*0.05) $
    --   renderDual (pRing shape22) (dual (polygonOffset shape22) $ polygonTriangulation shape22)
    -- , polygonNumDots shape22
    --   -- renderTriangulation cShape (polygonTriangulation cShape)
    -- ]
    -- mkGroup
    -- [ -- translate (-2) 0 $ withFillColor "white" $ polygonShape $ balloonP (min 1 $ t+0.1) origin
    --   if False then mkGroup [] else translate (0) 0 $ mkGroup
    --   [ withFillColor "white" $ polygonShape inflated
    --   -- , polygonNumDots inflated
    --   ]
    -- -- , translate (2) 0 $ mkGroup
    -- --   [ withFillColor "grey" $ polygonShape origin
    -- --   , polygonNumDots origin
    -- --   ]
    -- ]
  -- play $ mapA (withStrokeWidth (defaultStrokeWidth*0.2)) $ drawTriangulation cShape earClip'
  -- play $ staticFrame 1 $ renderTriangulation shape3 earClip
  -- play $ staticFrame 1 $ renderTriangulation shape4 earClip
  -- play $ staticFrame 1 $ renderTriangulation shape5 earClip
  -- play $ staticFrame 1 $ renderTriangulation shape6 earClip
  -- let p = deoverlapPolygon shape23
  -- newSpriteSVG $ translate (-3) 0 $ scale 2 $ center $ mkGroup
  --   [ withFillColor "grey" $ polygonShape p
  --   , polygonNumDots p
  --   ]
  -- wait 1
  -- play $ mapA (scale 2 . withStrokeWidth (defaultStrokeWidth*0.2)) $
  --   drawTriangulation p (earClip')
  return ()

drawVisibility :: Polygon -> Animation
drawVisibility p' = mkAnimation 5 $ \t ->
  let p = pCycle p' (t::Double) in
  centerUsing (polygonShape p) $
  mkGroup
  [ withFillColor "grey" $ polygonShape p
  , withFillColor "grey" $ polygonDots p
  , withFillColor "white" $ mkLinePathClosed
    [ (x,y) | V2 x y <- visibility (map (fmap realToFrac) $ V.toList $ polygonPoints p) ]
  , let V2 x y = fmap realToFrac $ pAccess p 0 in
    translate x y $ withFillColor "red" $ mkCircle 0.1
  -- , withFillColor "blue" $ latex $ T.pack $ show (t)
  ]

drawSSSPVisibility :: Polygon -> Animation
drawSSSPVisibility p' = mkAnimation 5 $ \t ->
  let p = pSetOffset p' (round $ t*(fromIntegral $ pSize p'-1))
      vis = ssspVisibility p in
  centerUsing (polygonShape p) $
  mkGroup
  [ withFillColor "grey" $ polygonShape p
  -- , withFillColor "grey" $ polygonDots p
  -- , withFillColor "white" $ polygonShape vis
  , let V2 x y = fmap realToFrac $ pAccess p 0 in
    translate x y $ withFillColor "red" $ mkCircle 0.1
  -- , withFillColor "blue" $ latex $ T.pack $ show (t)
  ]

drawSSSPVisibilityFast :: Polygon -> Animation
drawSSSPVisibilityFast p' = mkAnimation 5 $ \t ->
  let root = min (pSize p-1) $ (floor $ t*(fromIntegral $ pSize p))
      p = pSetOffset p' root
      vis = ssspVisibility p in
  -- centerUsing (polygonShape p) $
  mkGroup
  [ withFillColor "grey" $ polygonShape p
  , withFillColor "white" $ polygonShape vis
  , withFillColor "grey" $ polygonNumDots p
  , let V2 x y = fmap realToFrac $ pAccess p 0 in
    translate x y $ withFillColor "red" $ mkCircle 0.09
  -- , withFillColor "blue" $ latex $ T.pack $ show (t)
  ]

drawCompatible :: Polygon -> Polygon -> Animation
drawCompatible a b = sceneAnimation $ do
  newSpriteSVG $ translate (-3) 0 $ mkGroup
    [ withFillColor "grey" $ polygonShape a
    , withFillColor "grey" $ polygonNumDots a
    ]
  newSpriteSVG $ translate (3) 0 $ mkGroup
    [ withFillColor "grey" $ polygonShape b
    , withFillColor "grey" $ polygonNumDots b
    ]
  let compat = compatiblyTriangulateP a b
  forM_ compat $ \(l, r) -> do
    fork $ play $ staticFrame 1 $
      translate (-3) 0 $ withStrokeColor "white" $ withStrokeWidth (defaultStrokeWidth*0.2) $
      withFillOpacity 0 $ polygonShape l
    fork $ play $ staticFrame 1 $
      translate (3) 0 $ withStrokeColor "white" $ withStrokeWidth (defaultStrokeWidth*0.2) $
      withFillOpacity 0 $ polygonShape r
  -- forM_ compat $ \(l, r) -> waitOn $ do
  --   fork $ play $ staticFrame 1 $
  --     translate (-3) 0 $
  --     withFillColor "white" $ polygonShape l
  --   fork $ play $ staticFrame 1 $
  --     translate (3) 0 $
  --     withFillColor "white" $ polygonShape r

drawOverlap :: Polygon -> Animation
drawOverlap p' = mkAnimation 5 $ \t ->
  let p = pCycle p' (t::Double)
      vis = ssspVisibility p
      vis' = ssspVisibility p'
      mWins = ssspWindows p
      oWins = ssspWindows p'
      -- (left, right) = split1Link p' 0 3
      -- (left, right) = split2Link p' 0 2
      -- (left, right) = splitNLink p' 0 [(TwoLink,2),(OneLink,3)]
      -- (left, right) = splitNLink p' 0 [(OneLink,5),(TwoLink,3)]
      -- sPoly = if t < 0.5 then left else right
      in
  centerUsing (polygonShape p) $
  mkGroup
  [ withFillColor "grey" $ polygonShape p
  -- , withFillColor "grey" $ polygonDots p
  , withFillOpacity 0.5 $ withFillColor "lightgreen" $ polygonShape vis
  -- , withFillOpacity 0.5 $ withFillColor "cyan" $ polygonShape vis'
  , let V2 x y = fmap realToFrac $ pAccess p 0 in
    translate x y $ withFillColor "red" $ mkCircle 0.1
  , let V2 x y = fmap realToFrac $ pAccess p' 0 in
    translate x y $ withFillColor "red" $ mkCircle 0.1
  -- , withFillColor "blue" $ latex $ T.pack $ show (t)
  , mkGroup
    [ withStrokeColor "red" $
      mkLine (x1,y1) (x2,y2)
    | (a,b) <- mWins
    , let V2 x1 y1 = realToFrac <$> a
          V2 x2 y2 = realToFrac <$> b
    ]
  -- , mkGroup
  --   [ withStrokeColor "blue" $
  --     mkLine (x1,y1) (x2,y2)
  --   | (a,b) <- oWins
  --   , let V2 x1 y1 = realToFrac <$> a
  --         V2 x2 y2 = realToFrac <$> b
  --   ]
  , mkGroup
    [ withStrokeColor "green" $mkGroup
      [ mkLine (x1,y1) (x2,y2)
      , mkLine (x1',y1') (x2',y2') ]
    | (a,b) <- mWins
    , (i,j) <- oWins
    , a == i || a == j || b == i || b == j
    , not (sort [a,b] == sort [i,j])
    , let V2 x1 y1 = realToFrac <$> a
          V2 x2 y2 = realToFrac <$> b
          V2 x1' y1' = realToFrac <$> i
          V2 x2' y2' = realToFrac <$> j
    ]
  -- , mkGroup
  --   [ let V2 x y = realToFrac <$> link in
  --     translate x y $
  --     withFillColor "red" $
  --     mkCircle 0.1
  --   | link <- [steiner2Link p' 0 2, steiner2Link p' 5 3] ]
  -- , withFillColor "blue" $ polygonShape sPoly
  ]

drawWindow :: Polygon -> SVG
drawWindow p =
  let mWins = ssspWindows p
      in
  mkGroup
  [ mkGroup
    [ mkLine (x1,y1) (x2,y2)
    | (a,b) <- mWins
    , let V2 x1 y1 = realToFrac <$> a
          V2 x2 y2 = realToFrac <$> b
    ]
  ]

drawWindowOverlap :: Polygon -> Int -> Int -> SVG
drawWindowOverlap p a b =
  let aWins = ssspWindows (pAdjustOffset p a)
      bWins = ssspWindows (pAdjustOffset p b)
      in
  mkGroup
  [ mkGroup $
    [ withStrokeColor "green" $mkGroup
      [ mkLine (x1,y1) (x2,y2)
      , mkLine (x1',y1') (x2',y2') ]
    | (a,b) <- aWins
    , (i,j) <- bWins
    , a == i || a == j || b == i || b == j
    , not (sort [a,b] == sort [i,j])
    , let V2 x1 y1 = realToFrac <$> a
          V2 x2 y2 = realToFrac <$> b
          V2 x1' y1' = realToFrac <$> i
          V2 x2' y2' = realToFrac <$> j
    ]
  ]

drawSSSP :: Polygon -> (Polygon -> SSSP) -> Animation
drawSSSP p gen = mkAnimation 5 $ \t -> centerUsing outline $
  let p' = pCycles p !! (round $ t*(fromIntegral $ pSize p-1)) in
  mkGroup
  [ outline
  , renderSSSP p' (gen p')
  -- , let V2 x y = fmap realToFrac $ pAccess (pCycle p t) 0 in
  --   translate x y $ withFillColor "red" $ mkCircle 0.1
  , withFillColor "grey" $ polygonNumDots $ p'
  ]
  where
    outline =
      withFillColor "grey" $ mkLinePathClosed
        [ (x,y) | V2 x y <- map (fmap realToFrac) (V.toList (polygonPoints p)  ++ [pAccess p 0]) ]

{-# INLINE drawSSSPFast #-}
drawSSSPFast :: Polygon -> Animation
drawSSSPFast p = mkAnimation 5 $ \t -> centerUsing outline $
  let root = (round $ t*(fromIntegral $ pSize p-1))
      d = dual root triangulation
      sTree = sssp (pRing p) d in
  mkGroup
  [ outline
  , renderSSSP p sTree
  -- , let V2 x y = fmap realToFrac $ pAccess (pCycle p t) 0 in
  --   translate x y $ withFillColor "red" $ mkCircle 0.1
  , withFillColor "grey" $ polygonNumDots $ p
  ]
  where
    triangulation = earClip $ pRing p
    outline =
      withFillColor "grey" $ mkLinePathClosed
        [ (x,y) | V2 x y <- map (fmap realToFrac) (V.toList (polygonPoints p) ++ [pAccess p 0]) ]

drawVisibleFrom :: Polygon -> Animation
drawVisibleFrom p = mkAnimation 5 $ \t -> centerUsing (polygonShape p) $ mkGroup
  [ withFillColor "grey" $ polygonShape p
  , withFillColor "grey" $ polygonDots p
  , renderVisibleFrom (pCycle p t)
  , let V2 x y = fmap realToFrac $ pAccess (pCycle p t) 0 in
    translate x y $ withFillColor "red" $ mkCircle 0.1
  ]

renderVisibleFrom :: Polygon -> SVG
renderVisibleFrom p = withStrokeColor "white" $ withFillColor "white" $ mkGroup
  [ mkGroup
    [ mkLine (ax,ay) (bx,by)
    , translate bx by $ mkCircle 0.1 ]
  | i <- visibilityArray (pRing p) V.! 0
  , let V2 ax ay = fmap realToFrac $ pAccess p 0
        V2 bx by = fmap realToFrac $ pAccess p i ]

drawTriangulation :: Polygon -> (Ring Rational -> [Triangulation]) -> Animation
drawTriangulation p gen = sceneAnimation $ do
  forM_ (gen $ pRing p) $ \t -> play $ staticFrame 1 $ renderTriangulation p t

renderDual :: Ring Rational -> Dual -> SVG
renderDual ring d = case d of
    Dual (a,b,c) l r -> mkGroup
      [ withFillColor "blue" $ mkTrig a b c
      , worker c a l
      , worker b c r
      ]
  where
    mkTrig a b c =
      let V2 x1 y1 = realToFrac <$> ringAccess ring a
          V2 x2 y2 = realToFrac <$> ringAccess ring b
          V2 x3 y3 = realToFrac <$> ringAccess ring c
      in mkLinePathClosed [ (x1, y1), (x2,y2), (x3,y3) ]
    worker p1 p2 EmptyDual = mkGroup []
    worker p1 p2 (NodeDual x l r) = mkGroup
      [ mkTrig p1 p2 x
      , worker x p2 l
      , worker p1 x r
      ]
-- data Dual = Dual (Int,Int,Int) -- (a,b,c)
--                   DualTree -- borders ca
--                   DualTree -- borders bc
--   deriving (Show)
--
-- data DualTree
--   = EmptyDual
--   | NodeDual Int -- axb triangle, a and b are from parent.
--       DualTree -- borders xb
--       DualTree -- borders ax
--   deriving (Show)


{-# LANGUAGE CPP                        #-}
{-# LANGUAGE DataKinds                  #-}
{-# LANGUAGE EmptyDataDecls             #-}
{-# LANGUAGE FlexibleContexts           #-}
{-# LANGUAGE FlexibleInstances          #-}
{-# LANGUAGE FunctionalDependencies     #-}
{-# LANGUAGE GADTs                      #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE KindSignatures             #-}
{-# LANGUAGE MultiParamTypeClasses      #-}
{-# LANGUAGE Rank2Types                 #-}
{-# LANGUAGE ScopedTypeVariables        #-}
{-# LANGUAGE TypeFamilies               #-}
{-# LANGUAGE TypeOperators              #-}
{-# LANGUAGE ViewPatterns               #-}

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

{- |
Module      :  Numeric.LinearAlgebra.Static
Copyright   :  (c) Alberto Ruiz 2014
License     :  BSD3
Stability   :  experimental

Experimental interface with statically checked dimensions.

See code examples at http://dis.um.es/~alberto/hmatrix/static.html.

-}

module Numeric.LinearAlgebra.Static(
    -- * Vector
       ℝ, R,
    vec2, vec3, vec4, (&), (#), split, headTail,
    vector,
    linspace, range, dim,
    -- * Matrix
    L, Sq, build,
    row, col, (|||),(===), splitRows, splitCols,
    unrow, uncol,
    tr,
    eye,
    diag,
    blockAt,
    matrix,
    -- * Complex
    ℂ, C, M, Her, her, 𝑖,
    toComplex,
    fromComplex,
    complex,
    real,
    imag,
    sqMagnitude,
    magnitude,
    -- * Products
    (<>),(#>),(<.>),
    -- * Linear Systems
    linSolve, (<\>),
    -- * Factorizations
    svd, withCompactSVD, svdTall, svdFlat, Eigen(..),
    withNullspace, withOrth, qr, chol,
    -- * Norms
    Normed(..),
    -- * Random arrays
    Seed, RandDist(..),
    randomVector, rand, randn, gaussianSample, uniformSample,
    -- * Misc
    mean, meanCov,
    Disp(..), Domain(..),
    withVector, withMatrix, exactLength, exactDims,
    toRows, toColumns, withRows, withColumns,
    Sized(..), Diag(..), Sym, sym, mTm, unSym, (<·>)
) where

import           Control.Arrow               ((***))
import qualified Data.Bifunctor              as BF (first)
import           Data.Proxy                  (Proxy (..))
import           Data.Type.Equality          ((:~:) (Refl))
import           GHC.TypeLits
import           Internal.Static
import           Numeric.LinearAlgebra       hiding (C, Konst (..), R, build, chol, col,
                                              complex, diag, disp, dot, eig, eigSH,
                                              eigenvalues, eigenvaluesSH, fromComplex,
                                              fromList, gaussianSample, linspace, mTm,
                                              magnitude, matrix, meanCov, qr, rand, randn,
                                              randomVector, range, real, row, size, svd,
                                              sym, takeDiag, toColumns, toComplex, toRows,
                                              unSym, uniformSample, vector, ( #> ), (<.>),
                                              (<>), (<\>), (===), (|||))
import qualified Numeric.LinearAlgebra       as LA
import qualified Numeric.LinearAlgebra.Devel as LA
import           Text.Printf
#if MIN_VERSION_base(4,11,0)
import           Prelude                     hiding ((<>))
#endif

ud1 :: R n -> Vector ℝ
ud1 (R (Dim v)) = v

infixl 4 &
(&) :: forall n . KnownNat n
    => R n -> ℝ -> R (n+1)
u & x = u # (konst x :: R 1)

infixl 4 #
(#) :: forall n m . (KnownNat n, KnownNat m)
    => R n -> R m -> R (n+m)
(R u) # (R v) = R (vconcat u v)

vec2 :: ℝ -> ℝ -> R 2
vec2 a b = R (gvec2 a b)

vec3 :: ℝ -> ℝ -> ℝ -> R 3
vec3 a b c = R (gvec3 a b c)

vec4 :: ℝ -> ℝ -> ℝ -> ℝ -> R 4
vec4 a b c d = R (gvec4 a b c d)

vector :: KnownNat n => [ℝ] -> R n
vector = fromList

matrix :: (KnownNat m, KnownNat n) => [ℝ] -> L m n
matrix = fromList

linspace :: forall n . KnownNat n => (ℝ,ℝ) -> R n
linspace (a,b) = v
  where
    v = mkR (LA.linspace (size v) (a,b))

range :: forall n . KnownNat n => R n
range = v
  where
    v = mkR (LA.linspace d (1,fromIntegral d))
    d = size v

dim :: forall n . KnownNat n => R n
dim = v
  where
    v = mkR (scalar (fromIntegral $ size v))

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

ud2 :: L m n -> Matrix ℝ
ud2 (L (Dim (Dim x))) = x

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

diag :: KnownNat n => R n -> Sq n
diag = diagR 0

eye :: KnownNat n => Sq n
eye = diag 1

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

blockAt :: forall m n . (KnownNat m, KnownNat n) =>  ℝ -> Int -> Int -> Matrix Float -> L m n
blockAt x r c a = res
  where
    z = scalar x
    z1 = LA.konst x (r,c)
    z2 = LA.konst x (max 0 (m'-(ra+r)), max 0 (n'-(ca+c)))
    ra = min (rows a) . max 0 $ m'-r
    ca = min (cols a) . max 0 $ n'-c
    sa = subMatrix (0,0) (ra, ca) a
    (m',n') = size res
    res = mkL $ fromBlocks [[z1,z,z],[z,sa,z],[z,z,z2]]

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

row :: R n -> L 1 n
row = mkL . asRow . ud1

--col :: R n -> L n 1
col v = tr . row $ v

unrow :: L 1 n -> R n
unrow = mkR . head . LA.toRows . ud2

--uncol :: L n 1 -> R n
uncol v = unrow . tr $ v

infixl 2 ===
(===) :: (KnownNat r1, KnownNat r2, KnownNat c) => L r1 c -> L r2 c -> L (r1+r2) c
a === b = mkL (extract a LA.=== extract b)

infixl 3 |||
-- (|||) :: (KnownNat r, KnownNat c1, KnownNat c2) => L r c1 -> L r c2 -> L r (c1+c2)
a ||| b = tr (tr a === tr b)

type Sq n  = L n n
--type CSq n = CL n n

type GL = forall n m . (KnownNat n, KnownNat m) => L m n
type GSq = forall n . KnownNat n => Sq n

isKonst :: forall m n . (KnownNat m, KnownNat n) => L m n -> Maybe (ℝ,(Int,Int))
isKonst s@(unwrap -> x)
    | singleM x = Just (x `atIndex` (0,0), (size s))
    | otherwise = Nothing

isKonstC :: forall m n . (KnownNat m, KnownNat n) => M m n -> Maybe (ℂ,(Int,Int))
isKonstC s@(unwrap -> x)
    | singleM x = Just (x `atIndex` (0,0), (size s))
    | otherwise = Nothing

infixr 8 <>
(<>) :: forall m k n. (KnownNat m, KnownNat k, KnownNat n) => L m k -> L k n -> L m n
(<>) = mulR

infixr 8 #>
(#>) :: (KnownNat m, KnownNat n) => L m n -> R n -> R m
(#>) = appR

infixr 8 <·>
(<·>) :: KnownNat n => R n -> R n -> ℝ
(<·>) = dotR

infixr 8 <.>
(<.>) :: KnownNat n => R n -> R n -> ℝ
(<.>) = dotR

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

class Diag m d | m -> d
  where
    takeDiag :: m -> d

instance KnownNat n => Diag (L n n) (R n)
  where
    takeDiag x = mkR (LA.takeDiag (extract x))

instance KnownNat n => Diag (M n n) (C n)
  where
    takeDiag x = mkC (LA.takeDiag (extract x))

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

toComplex :: KnownNat n => (R n, R n) -> C n
toComplex (r,i) = mkC $ LA.toComplex (ud1 r, ud1 i)

fromComplex :: KnownNat n => C n -> (R n, R n)
fromComplex (C (Dim v)) = let (r,i) = LA.fromComplex v in (mkR r, mkR i)

complex :: KnownNat n => R n -> C n
complex r = mkC $ LA.toComplex (ud1 r, LA.konst 0 (size r))

real :: KnownNat n => C n -> R n
real = fst . fromComplex

imag :: KnownNat n => C n -> R n
imag = snd . fromComplex

sqMagnitude :: KnownNat n => C n -> R n
sqMagnitude c = let (r,i) = fromComplex c in r**2 + i**2

magnitude :: KnownNat n => C n -> R n
magnitude = sqrt . sqMagnitude

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

linSolve :: (KnownNat m, KnownNat n) => L m m -> L m n -> Maybe (L m n)
linSolve (extract -> a) (extract -> b) = fmap mkL (LA.linearSolve a b)

(<\>) :: (KnownNat m, KnownNat n, KnownNat r) => L m n -> L m r -> L n r
(extract -> a) <\> (extract -> b) = mkL (a LA.<\> b)

svd :: (KnownNat m, KnownNat n) => L m n -> (L m m, R n, L n n)
svd (extract -> m) = (mkL u, mkR s', mkL v)
  where
    (u,s,v) = LA.svd m
    s' = vjoin [s, z]
    z = LA.konst 0 (max 0 (cols m - LA.size s))

svdTall :: (KnownNat m, KnownNat n, n <= m) => L m n -> (L m n, R n, L n n)
svdTall (extract -> m) = (mkL u, mkR s, mkL v)
  where
    (u,s,v) = LA.thinSVD m

svdFlat :: (KnownNat m, KnownNat n, m <= n) => L m n -> (L m m, R m, L n m)
svdFlat (extract -> m) = (mkL u, mkR s, mkL v)
  where
    (u,s,v) = LA.thinSVD m

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

class Eigen m l v | m -> l, m -> v
  where
    eigensystem :: m -> (l,v)
    eigenvalues :: m -> l

newtype Sym n = Sym (Sq n) deriving Show

sym :: KnownNat n => Sq n -> Sym n
sym m = Sym $ (m + tr m)/2

mTm :: (KnownNat m, KnownNat n) => L m n -> Sym n
mTm x = Sym (tr x <> x)

unSym :: Sym n -> Sq n
unSym (Sym x) = x

𝑖 :: Sized ℂ s c => s
𝑖 = konst iC

newtype Her n = Her (M n n)

her :: KnownNat n => M n n -> Her n
her m = Her $ (m + LA.tr m)/2

instance (KnownNat n) => Disp (Sym n)
  where
    disp n (Sym x) = do
        let a = extract x
        let su = LA.dispf n a
        printf "Sym %d" (cols a) >> putStr (dropWhile (/='\n') $ su)

instance (KnownNat n) => Disp (Her n)
  where
    disp n (Her x) = do
        let a = extract x
        let su = LA.dispcf n a
        printf "Her %d" (cols a) >> putStr (dropWhile (/='\n') $ su)

instance KnownNat n => Eigen (Sym n) (R n) (L n n)
  where
    eigenvalues (Sym (extract -> m)) =  mkR . LA.eigenvaluesSH . LA.trustSym $ m
    eigensystem (Sym (extract -> m)) = (mkR l, mkL v)
      where
        (l,v) = LA.eigSH . LA.trustSym $ m

instance KnownNat n => Eigen (Sq n) (C n) (M n n)
  where
    eigenvalues (extract -> m) = mkC . LA.eigenvalues $ m
    eigensystem (extract -> m) = (mkC l, mkM v)
      where
        (l,v) = LA.eig m

chol :: KnownNat n => Sym n -> Sq n
chol (extract . unSym -> m) = mkL $ LA.chol $ LA.trustSym m

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

withNullspace
    :: forall m n z . (KnownNat m, KnownNat n)
    => L m n
    -> (forall k . (KnownNat k) => L n k -> z)
    -> z
withNullspace (LA.nullspace . extract -> a) f =
    case someNatVal $ fromIntegral $ cols a of
       Nothing                       -> error "static/dynamic mismatch"
       Just (SomeNat (_ :: Proxy k)) -> f (mkL a :: L n k)

withOrth
    :: forall m n z . (KnownNat m, KnownNat n)
    => L m n
    -> (forall k. (KnownNat k) => L n k -> z)
    -> z
withOrth (LA.orth . extract -> a) f =
    case someNatVal $ fromIntegral $ cols a of
       Nothing                       -> error "static/dynamic mismatch"
       Just (SomeNat (_ :: Proxy k)) -> f (mkL a :: L n k)

withCompactSVD
    :: forall m n z . (KnownNat m, KnownNat n)
    => L m n
    -> (forall k . (KnownNat k) => (L m k, R k, L n k) -> z)
    -> z
withCompactSVD (LA.compactSVD . extract -> (u,s,v)) f =
    case someNatVal $ fromIntegral $ LA.size s of
       Nothing                       -> error "static/dynamic mismatch"
       Just (SomeNat (_ :: Proxy k)) -> f (mkL u :: L m k, mkR s :: R k, mkL v :: L n k)

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

qr :: (KnownNat m, KnownNat n) => L m n -> (L m m, L m n)
qr (extract -> x) = (mkL q, mkL r)
  where
    (q,r) = LA.qr x

-- use qrRaw?

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

split :: forall p n . (KnownNat p, KnownNat n, p<=n) => R n -> (R p, R (n-p))
split (extract -> v) = ( mkR (subVector 0 p' v) ,
                         mkR (subVector p' (LA.size v - p') v) )
  where
    p' = fromIntegral . natVal $ (undefined :: Proxy p) :: Int

headTail :: (KnownNat n, 1<=n) => R n -> (ℝ, R (n-1))
headTail = ((!0) . extract *** id) . split

splitRows :: forall p m n . (KnownNat p, KnownNat m, KnownNat n, p<=m) => L m n -> (L p n, L (m-p) n)
splitRows (extract -> x) = ( mkL (takeRows p' x) ,
                             mkL (dropRows p' x) )
  where
    p' = fromIntegral . natVal $ (undefined :: Proxy p) :: Int

splitCols :: forall p m n. (KnownNat p, KnownNat m, KnownNat n, KnownNat (n-p), p<=n) => L m n -> (L m p, L m (n-p))
splitCols = (tr *** tr) . splitRows . tr

toRows :: forall m n . (KnownNat m, KnownNat n) => L m n -> [R n]
toRows (LA.toRows . extract -> vs) = map mkR vs

withRows
    :: forall n z . KnownNat n
    => [R n]
    -> (forall m . KnownNat m => L m n -> z)
    -> z
withRows (LA.fromRows . map extract -> m) f =
    case someNatVal $ fromIntegral $ LA.rows m of
       Nothing                       -> error "static/dynamic mismatch"
       Just (SomeNat (_ :: Proxy m)) -> f (mkL m :: L m n)

toColumns :: forall m n . (KnownNat m, KnownNat n) => L m n -> [R m]
toColumns (LA.toColumns . extract -> vs) = map mkR vs

withColumns
    :: forall m z . KnownNat m
    => [R m]
    -> (forall n . KnownNat n => L m n -> z)
    -> z
withColumns (LA.fromColumns . map extract -> m) f =
    case someNatVal $ fromIntegral $ LA.cols m of
       Nothing                       -> error "static/dynamic mismatch"
       Just (SomeNat (_ :: Proxy n)) -> f (mkL m :: L m n)

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

build
  :: forall m n . (KnownNat n, KnownNat m)
    => (ℝ -> ℝ -> ℝ)
    -> L m n
build f = r
  where
    r = mkL $ LA.build (size r) f

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

withVector
    :: forall z
     . Vector ℝ
    -> (forall n . (KnownNat n) => R n -> z)
    -> z
withVector v f =
    case someNatVal $ fromIntegral $ LA.size v of
       Nothing                       -> error "static/dynamic mismatch"
       Just (SomeNat (_ :: Proxy m)) -> f (mkR v :: R m)

-- | Useful for constraining two dependently typed vectors to match each
-- other in length when they are unknown at compile-time.
exactLength
    :: forall n m . (KnownNat n, KnownNat m)
    => R m
    -> Maybe (R n)
exactLength v = do
    Refl <- sameNat (Proxy :: Proxy n) (Proxy :: Proxy m)
    return $ mkR (unwrap v)

withMatrix
    :: forall z
     . Matrix ℝ
    -> (forall m n . (KnownNat m, KnownNat n) => L m n -> z)
    -> z
withMatrix a f =
    case someNatVal $ fromIntegral $ rows a of
       Nothing -> error "static/dynamic mismatch"
       Just (SomeNat (_ :: Proxy m)) ->
           case someNatVal $ fromIntegral $ cols a of
               Nothing -> error "static/dynamic mismatch"
               Just (SomeNat (_ :: Proxy n)) ->
                  f (mkL a :: L m n)

-- | Useful for constraining two dependently typed matrices to match each
-- other in dimensions when they are unknown at compile-time.
exactDims
    :: forall n m j k . (KnownNat n, KnownNat m, KnownNat j, KnownNat k)
    => L m n
    -> Maybe (L j k)
exactDims m = do
    Refl <- sameNat (Proxy :: Proxy m) (Proxy :: Proxy j)
    Refl <- sameNat (Proxy :: Proxy n) (Proxy :: Proxy k)
    return $ mkL (unwrap m)

randomVector
    :: forall n . KnownNat n
    => Seed
    -> RandDist
    -> R n
randomVector s d = mkR (LA.randomVector s d
                          (fromInteger (natVal (Proxy :: Proxy n)))
                       )

rand
    :: forall m n . (KnownNat m, KnownNat n)
    => IO (L m n)
rand = mkL <$> LA.rand (fromInteger (natVal (Proxy :: Proxy m)))
                       (fromInteger (natVal (Proxy :: Proxy n)))

randn
    :: forall m n . (KnownNat m, KnownNat n)
    => IO (L m n)
randn = mkL <$> LA.randn (fromInteger (natVal (Proxy :: Proxy m)))
                         (fromInteger (natVal (Proxy :: Proxy n)))

gaussianSample
    :: forall m n . (KnownNat m, KnownNat n)
    => Seed
    -> R n
    -> Sym n
    -> L m n
gaussianSample s (extract -> mu) (Sym (extract -> sigma)) =
    mkL $ LA.gaussianSample s (fromInteger (natVal (Proxy :: Proxy m)))
                            mu (LA.trustSym sigma)

uniformSample
    :: forall m n . (KnownNat m, KnownNat n)
    => Seed
    -> R n    -- ^ minimums of each row
    -> R n    -- ^ maximums of each row
    -> L m n
uniformSample s (extract -> mins) (extract -> maxs) =
    mkL $ LA.uniformSample s (fromInteger (natVal (Proxy :: Proxy m)))
                           (zip (LA.toList mins) (LA.toList maxs))

meanCov
    :: forall m n . (KnownNat m, KnownNat n, 1 <= m)
    => L m n
    -> (R n, Sym n)
meanCov (extract -> vs) = mkR *** (Sym . mkL . LA.unSym) $ LA.meanCov vs

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

class Domain field vec mat | mat -> vec field, vec -> mat field, field -> mat vec
  where
    mul :: forall m k n. (KnownNat m, KnownNat k, KnownNat n) => mat m k -> mat k n -> mat m n
    app :: forall m n . (KnownNat m, KnownNat n) => mat m n -> vec n -> vec m
    dot :: forall n . (KnownNat n) => vec n -> vec n -> field
    cross :: vec 3 -> vec 3 -> vec 3
    diagR ::  forall m n k . (KnownNat m, KnownNat n, KnownNat k) => field -> vec k -> mat m n
    dvmap :: forall n. KnownNat n => (field -> field) -> vec n -> vec n
    dmmap :: forall n m. (KnownNat m, KnownNat n) => (field -> field) -> mat n m -> mat n m
    outer :: forall n m. (KnownNat m, KnownNat n) => vec n -> vec m -> mat n m
    zipWithVector :: forall n. KnownNat n => (field -> field -> field) -> vec n -> vec n -> vec n
    det :: forall n. KnownNat n => mat n n -> field
    invlndet :: forall n. KnownNat n => mat n n -> (mat n n, (field, field))
    expm :: forall n. KnownNat n => mat n n -> mat n n
    sqrtm :: forall n. KnownNat n => mat n n -> mat n n
    inv :: forall n. KnownNat n => mat n n -> mat n n

instance Domain ℝ R L
  where
    mul = mulR
    app = appR
    dot = dotR
    cross = crossR
    diagR = diagRectR
    dvmap = mapR
    dmmap = mapL
    outer = outerR
    zipWithVector = zipWithR
    det = detL
    invlndet = invlndetL
    expm = expmL
    sqrtm = sqrtmL
    inv = invL

instance Domain ℂ C M
  where
    mul = mulC
    app = appC
    dot = dotC
    cross = crossC
    diagR = diagRectC
    dvmap = mapC
    dmmap = mapM'
    outer = outerC
    zipWithVector = zipWithC
    det = detM
    invlndet = invlndetM
    expm = expmM
    sqrtm = sqrtmM
    inv = invM

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

mulR :: forall m k n. (KnownNat m, KnownNat k, KnownNat n) => L m k -> L k n -> L m n

mulR (isKonst -> Just (a,(_,k))) (isKonst -> Just (b,_)) = konst (a * b * fromIntegral k)

mulR (isDiag -> Just (0,a,_)) (isDiag -> Just (0,b,_)) = diagR 0 (mkR v :: R k)
  where
    v = a' * b'
    n = min (LA.size a) (LA.size b)
    a' = subVector 0 n a
    b' = subVector 0 n b

mulR (isDiag -> Just (0,a,_)) (extract -> b) = mkL (asColumn a * takeRows (LA.size a) b)

mulR (extract -> a) (isDiag -> Just (0,b,_)) = mkL (takeColumns (LA.size b) a * asRow b)

mulR a b = mkL (extract a LA.<> extract b)

appR :: (KnownNat m, KnownNat n) => L m n -> R n -> R m
appR (isDiag -> Just (0, w, _)) v = mkR (w * subVector 0 (LA.size w) (extract v))
appR m v                          = mkR (extract m LA.#> extract v)

dotR :: KnownNat n => R n -> R n -> ℝ
dotR (extract -> u) (extract -> v) = LA.dot u v

crossR :: R 3 -> R 3 -> R 3
crossR (extract -> x) (extract -> y) = vec3 z1 z2 z3
  where
    z1 = x!1*y!2-x!2*y!1
    z2 = x!2*y!0-x!0*y!2
    z3 = x!0*y!1-x!1*y!0

outerR :: (KnownNat m, KnownNat n) => R n -> R m -> L n m
outerR (extract -> x) (extract -> y) = mkL (LA.outer x y)

mapR :: KnownNat n => (ℝ -> ℝ) -> R n -> R n
mapR f (unwrap -> v) = mkR (LA.cmap f v)

zipWithR :: KnownNat n => (ℝ -> ℝ -> ℝ) -> R n -> R n -> R n
zipWithR f (extract -> x) (extract -> y) = mkR (LA.zipVectorWith f x y)

mapL :: (KnownNat n, KnownNat m) => (ℝ -> ℝ) -> L n m -> L n m
mapL f = overMatL' (LA.cmap f)

detL :: KnownNat n => Sq n -> ℝ
detL = LA.det . unwrap

invlndetL :: KnownNat n => Sq n -> (L n n, (ℝ, ℝ))
invlndetL = BF.first mkL . LA.invlndet . unwrap

expmL :: KnownNat n => Sq n -> Sq n
expmL = overMatL' LA.expm

sqrtmL :: KnownNat n => Sq n -> Sq n
sqrtmL = overMatL' LA.sqrtm

invL :: KnownNat n => Sq n -> Sq n
invL = overMatL' LA.inv

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

mulC :: forall m k n. (KnownNat m, KnownNat k, KnownNat n) => M m k -> M k n -> M m n

mulC (isKonstC -> Just (a,(_,k))) (isKonstC -> Just (b,_)) = konst (a * b * fromIntegral k)

mulC (isDiagC -> Just (0,a,_)) (isDiagC -> Just (0,b,_)) = diagR 0 (mkC v :: C k)
  where
    v = a' * b'
    n = min (LA.size a) (LA.size b)
    a' = subVector 0 n a
    b' = subVector 0 n b

mulC (isDiagC -> Just (0,a,_)) (extract -> b) = mkM (asColumn a * takeRows (LA.size a) b)

mulC (extract -> a) (isDiagC -> Just (0,b,_)) = mkM (takeColumns (LA.size b) a * asRow b)

mulC a b = mkM (extract a LA.<> extract b)

appC :: (KnownNat m, KnownNat n) => M m n -> C n -> C m
appC (isDiagC -> Just (0, w, _)) v = mkC (w * subVector 0 (LA.size w) (extract v))
appC m v                           = mkC (extract m LA.#> extract v)

dotC :: KnownNat n => C n -> C n -> ℂ
dotC (extract -> u) (extract -> v) = LA.dot u v

crossC :: C 3 -> C 3 -> C 3
crossC (extract -> x) (extract -> y) = mkC (LA.fromList [z1, z2, z3])
  where
    z1 = x!1*y!2-x!2*y!1
    z2 = x!2*y!0-x!0*y!2
    z3 = x!0*y!1-x!1*y!0

outerC :: (KnownNat m, KnownNat n) => C n -> C m -> M n m
outerC (extract -> x) (extract -> y) = mkM (LA.outer x y)

mapC :: KnownNat n => (ℂ -> ℂ) -> C n -> C n
mapC f (unwrap -> v) = mkC (LA.cmap f v)

zipWithC :: KnownNat n => (ℂ -> ℂ -> ℂ) -> C n -> C n -> C n
zipWithC f (extract -> x) (extract -> y) = mkC (LA.zipVectorWith f x y)

mapM' :: (KnownNat n, KnownNat m) => (ℂ -> ℂ) -> M n m -> M n m
mapM' f = overMatM' (LA.cmap f)

detM :: KnownNat n => M n n -> ℂ
detM = LA.det . unwrap

invlndetM :: KnownNat n => M n n -> (M n n, (ℂ, ℂ))
invlndetM = BF.first mkM . LA.invlndet . unwrap

expmM :: KnownNat n => M n n -> M n n
expmM = overMatM' LA.expm

sqrtmM :: KnownNat n => M n n -> M n n
sqrtmM = overMatM' LA.sqrtm

invM :: KnownNat n => M n n -> M n n
invM = overMatM' LA.inv

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

diagRectR :: forall m n k . (KnownNat m, KnownNat n, KnownNat k) => ℝ -> R k -> L m n
diagRectR x v
    | m' == 1 = mkL (LA.diagRect x ev m' n')
    | m'*n' > 0 = r
    | otherwise = matrix []
  where
    r = mkL (asRow (vjoin [scalar x, ev, zeros]))
    ev = extract v
    zeros = LA.konst x (max 0 ((min m' n') - LA.size ev))
    (m',n') = size r

diagRectC :: forall m n k . (KnownNat m, KnownNat n, KnownNat k) => ℂ -> C k -> M m n
diagRectC x v
    | m' == 1 = mkM (LA.diagRect x ev m' n')
    | m'*n' > 0 = r
    | otherwise = fromList []
  where
    r = mkM (asRow (vjoin [scalar x, ev, zeros]))
    ev = extract v
    zeros = LA.konst x (max 0 ((min m' n') - LA.size ev))
    (m',n') = size r

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

mean :: (KnownNat n, 1<=n) => R n -> ℝ
mean v = v <·> (1/dim)

test :: (Bool, IO ())
test = (ok,info)
  where
    ok =   extract (eye :: Sq 5) == ident 5
           && (unwrap .unSym) (mTm sm :: Sym 3) == tr ((3><3)[1..]) LA.<> (3><3)[1..]
           && unwrap (tm :: L 3 5) == LA.matrix 5 [1..15]
           && thingS == thingD
           && precS == precD
           && withVector (LA.vector [1..15]) sumV == sumElements (LA.fromList [1..15])

    info = do
        print $ u
        print $ v
        print (eye :: Sq 3)
        print $ ((u & 5) + 1) <·> v
        print (tm :: L 2 5)
        print (tm <> sm :: L 2 3)
        print thingS
        print thingD
        print precS
        print precD
        print $ withVector (LA.vector [1..15]) sumV
        splittest

    sumV w = w <·> konst 1

    u = vec2 3 5

    𝕧 x = vector [x] :: R 1

    v = 𝕧 2 & 4 & 7

    tm :: GL
    tm = lmat 0 [1..]

    lmat :: forall m n . (KnownNat m, KnownNat n) => ℝ -> [ℝ] -> L m n
    lmat z xs = r
      where
        r = mkL . reshape n' . LA.fromList . take (m'*n') $ xs ++ repeat z
        (m',n') = size r

    sm :: GSq
    sm = lmat 0 [1..]

    thingS = (u & 1) <·> tr q #> q #> v
      where
        q = tm :: L 10 3

    thingD = vjoin [ud1 u, 1] LA.<.> tr m LA.#> m LA.#> ud1 v
      where
        m = LA.matrix 3 [1..30]

    precS = (1::Float) + (2::Float) * ((1 :: R 3) * (u & 6)) <·> konst 2 #> v
    precD = 1 + 2 * vjoin[ud1 u, 6] LA.<.> LA.konst 2 (LA.size (ud1 u) +1, LA.size (ud1 v)) LA.#> ud1 v

splittest
    = do
    let v = range :: R 7
        a = snd (split v) :: R 4
    print $ a
    print $ snd . headTail . snd . headTail $ v
    print $ first (vec3 1 2 3)
    print $ second (vec3 1 2 3)
    print $ third (vec3 1 2 3)
    print $ (snd $ splitRows eye :: L 4 6)
 where
    first v = fst . headTail $ v
    second v = first . snd . headTail $ v
    third v = first . snd . headTail . snd . headTail $ v

instance (KnownNat n', KnownNat m') => Testable (L n' m')
  where
    checkT _ = test

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

instance KnownNat n => Normed (R n)
  where
    norm_0 v = norm_0 (extract v)
    norm_1 v = norm_1 (extract v)
    norm_2 v = norm_2 (extract v)
    norm_Inf v = norm_Inf (extract v)

instance (KnownNat m, KnownNat n) => Normed (L m n)
  where
    norm_0 m = norm_0 (extract m)
    norm_1 m = norm_1 (extract m)
    norm_2 m = norm_2 (extract m)
    norm_Inf m = norm_Inf (extract m)

mkSym f = Sym . f . unSym
mkSym2 f x y = Sym (f (unSym x) (unSym y))

instance KnownNat n =>  Num (Sym n)
  where
    (+) = mkSym2 (+)
    (*) = mkSym2 (*)
    (-) = mkSym2 (-)
    abs = mkSym abs
    signum = mkSym signum
    negate = mkSym negate
    fromInteger = Sym . fromInteger

instance KnownNat n => Fractional (Sym n)
  where
    fromRational = Sym . fromRational
    (/) = mkSym2 (/)

instance KnownNat n => Floating (Sym n)
  where
    sin   = mkSym sin
    cos   = mkSym cos
    tan   = mkSym tan
    asin  = mkSym asin
    acos  = mkSym acos
    atan  = mkSym atan
    sinh  = mkSym sinh
    cosh  = mkSym cosh
    tanh  = mkSym tanh
    asinh = mkSym asinh
    acosh = mkSym acosh
    atanh = mkSym atanh
    exp   = mkSym exp
    log   = mkSym log
    sqrt  = mkSym sqrt
    (**)  = mkSym2 (**)
    pi    = Sym pi

instance KnownNat n => Additive (Sym n) where
    add = (+)

instance KnownNat n => Transposable (Sym n) (Sym n) where
    tr  = id
    tr' = id

instance KnownNat n => Transposable (Her n) (Her n) where
    tr          = id
    tr' (Her m) = Her (tr' m)


module LispType
    (
    LispVal(..)
  , Array
    ) where
import Data.Ratio
import Data.Complex
import Data.Array

{-
台数的データ型の一例
LispVal型の変数が持つことのできる値の集合を定めている
選択肢のそれぞれ(コンストラクタ、と呼ばれ、| で区切られる)は、
、コンストラクタのタグとそのコンストラクタが持つことのできるデータの型を含む
この例では、LispValは次のどれか
  1. Atom - そのアトムの示す文字列を格納します。
  2. List - 他のLispValのリストを保持します(Haskellのリストは角括弧で表されます)。properリストとも呼ばれます。
  3. DottedList - Schemeの(a b . c)を表し、improperリストとも呼ばれます。これは最後以外全ての要素のリストを持ち、最後の要素を別に格納します。
  4. Number - Haskellの整数を保持します。
  5. String - Haskellの文字列を保持します。
  6. Bool - Haskellの真偽値を保持します。
-}
data LispVal =  Atom        String
              | String      String
              | Character   Char
              | Bool        Bool
              | Number      Integer
              | Float       Double
              | Complex    (Complex Double)
              | Ratio       Rational
              | Vector     (Array Int LispVal)
              | List       [LispVal]
              | DottedList [LispVal] LispVal


module ForMNIST
 ( makeImageMatrix
 , loadImage
 , loadLabel
 , datasize
 , loadImageTest
 , loadLabelTest
 ) where

-------------------------------------------------------------
--module
-------------------------------------------------------------

import qualified Data.ByteString.Lazy as B
import Numeric.LinearAlgebra
import Tools

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

-------------------------------------------------------------
--MNISTのデータを回収する。
-------------------------------------------------------------
datasize = 10000

--与えられたリスト[a]から指定した要素数(num)ごとに分割してリストのリスト[[a]]を作る。
divideEqually :: Int -> [a] -> [[a]]
divideEqually _ [] = []
divideEqually num xs = xs1 : (divideEqually num xs2)
    where (xs1, xs2) = splitAt num xs

--画像イメージのBytestring(xs)から指定した長さを切り取ってリスト[Double]を作成する。
convertImage :: B.ByteString -> [Double]
convertImage xs = fmap toEnum $ fmap fromEnum . B.unpack $ B.drop byte xs
    where byte = 16  + 784*60000 -784 * datasize

--リストのリスト[[a]]からMatrixのリスト[[Matrix]]を作成する。
makeImageMatrix :: [[Double]] -> [Matrix Double]
makeImageMatrix [] = []
makeImageMatrix (x:xs) = (row x) : (makeImageMatrix xs)

loadImage1 :: B.ByteString -> [[Double]]
loadImage1 xs = divideEqually 784 (convertImage xs)

--画像イメージのBytestring(xs)から指定した長さを切り取ってMatrixのリスト[[Matrix]]を作成する。
loadImage :: B.ByteString -> [Matrix Double]
loadImage xs = makeImageMatrix $ divideEqually 784 (convertImage xs)

--ビジュアル化用なのでIntで出力する。
convertImage2 :: B.ByteString -> [Int]
convertImage2 xs = fmap fromEnum . B.unpack $ B.drop byte xs
    where byte = 16  + 784*10000 -784 * datasize

--ラベルのBytestring(xs)から指定した長さを切り取ってリスト[Int]を作成する。
loadLabel :: B.ByteString -> [Int]
loadLabel xs = fmap fromEnum . B.unpack $ B.drop byte xs
    where byte = 8  + 1*60000 - 1 * datasize

loadLabelTest :: B.ByteString -> [Int]
loadLabelTest xs = fmap fromEnum . B.unpack $ B.drop byte xs
    where byte = 8  + 1*10000 - 1 * 1000

convertImageTest :: B.ByteString -> [Double]
convertImageTest xs = fmap toEnum $ fmap fromEnum . B.unpack $ B.drop byte xs
    where byte = 16  + 784*10000 -784 * 1000

loadImageTest :: B.ByteString -> [Matrix Double]
loadImageTest xs = makeImageMatrix $ divideEqually 784 (convertImageTest xs)

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


{-# LANGUAGE DataKinds            #-}
{-# LANGUAGE FlexibleContexts     #-}
{-# LANGUAGE FlexibleInstances    #-}
{-# LANGUAGE GADTs                #-}
{-# LANGUAGE PatternSynonyms      #-}
{-# LANGUAGE RankNTypes           #-}
{-# LANGUAGE ScopedTypeVariables  #-}
{-# LANGUAGE StandaloneDeriving   #-}
{-# LANGUAGE TypeFamilies         #-}
{-# LANGUAGE TypeOperators        #-}
{-# LANGUAGE UndecidableInstances #-}
{-# OPTIONS_GHC -Wall #-}
-- | This module is wrongly named.
module VectorSpace (
    LinMap (.., LI),
    lmul,
    HasDim(Dim, dimDict),
    toRawMatrix,
    evalL,
    L (..),
    linear,
    VectorSpace (..),
    toVector,
    fromVector,
) where

import Data.Constraint       (Dict (..), withDict, (:-))
import Data.Proxy            (Proxy (..))
import GHC.TypeLits
import Overloaded.Categories

import qualified Control.Category      as C
import qualified Data.Constraint.Nat   as C
import qualified Numeric.LinearAlgebra as L

-- import qualified Numeric.LinearAlgebra.Static as LS

data LinMap a b where
    LZ :: LinMap a b
    LD :: Double -> LinMap a a
    LH :: LinMap a b -> LinMap a c -> LinMap a (b, c)
    LV :: LinMap a c -> LinMap b c -> LinMap (a, b) c
    LA :: LinMap a b -> LinMap a b -> LinMap a b

deriving instance Show (LinMap a b)

pattern LI :: forall a b. () => (b ~ a) => LinMap a b
pattern LI = LD 1

lmul :: Double -> LinMap a b -> LinMap a b
lmul _ LZ       = LZ
lmul k (LD x)   = LD (k * x)
lmul k (LH f g) = LH (lmul k f) (lmul k g)
lmul k (LV f g) = LV (lmul k f) (lmul k g)
lmul k (LA f g) = LA (lmul k f) (lmul k g)

lcomp :: LinMap b c -> LinMap a b -> LinMap a c
lcomp LZ       _        = LZ
lcomp _        LZ       = LZ
lcomp (LD k)   h        = lmul k h
lcomp h        (LD k)   = lmul k h
lcomp (LA f g) h        = LA (lcomp f h) (lcomp g h)
lcomp f        (LA g h) = LA (lcomp f g) (lcomp f h)
lcomp (LH f g) h        = LH (lcomp f h) (lcomp g h)
lcomp h        (LV f g) = LV (lcomp h f) (lcomp h g)
lcomp (LV f g) (LH u v) = LA (lcomp f u) (lcomp g v)

instance Category LinMap where
    id  = LI
    (.) = lcomp

instance CategoryWith1 LinMap where
    type Terminal LinMap = ()
    terminal = LZ

instance CartesianCategory LinMap where
    type Product LinMap = (,)
    proj1  = LV C.id LZ
    proj2  = LV LZ C.id
    fanout = LH

instance CategoryWith0 LinMap where
    type Initial LinMap = ()
    initial = LZ

instance CocartesianCategory LinMap where
    type Coproduct LinMap = (,)
    inl   = LH C.id LZ
    inr   = LH LZ C.id
    fanin = LV

instance BicartesianCategory LinMap where
    distr = LH
        (LH (LV (LV LI LZ) LZ) (LV LZ LI))
        (LH (LV (LV LZ LI) LZ) (LV LZ LI))

newtype L a b = L (forall r. LinMap r a -> LinMap r b)

lfst :: LinMap a (b, c) -> LinMap a b
lfst (LA f g) = LA (lfst f) (lfst g)
lfst (LH f _) = f
lfst (LV f g) = LV (lfst f) (lfst g)
lfst LZ       = LZ
lfst (LD k)   = LV (LD k) LZ

lsnd :: LinMap a (b, c) -> LinMap a c
lsnd (LH _ g) = g
lsnd (LA f g) = LA (lsnd f) (lsnd g)
lsnd (LV f g) = LV (lsnd f) (lsnd g)
lsnd LZ       = LZ
lsnd (LD k)   = LV LZ (LD k)

linitial :: LinMap r () -> LinMap r a
linitial _ = LZ

linear :: Double -> L a a
linear k = L $ lmul k

-- lmult :: Double -> Double -> LinMap r (a, a) -> LinMap r a
-- lmult x y (LH f g) = LA (LK y f) (LK x g)
-- lmult x y (LV f g) = LV (lmult x y f) (lmult x y g)
-- lmult x y (LA f g) = LA (lmult x y f) (lmult x y g)
-- lmult x y (LK k f) = LK k (lmult x y f)
-- lmult _ _ LZ       = LZ
-- lmult x y LI       = LV (LK y LI) (LK x LI)

instance Category L where
    id = L id
    L f . L g = L (f . g)

instance CategoryWith1 L where
    type Terminal L = ()

    terminal = L (\_ -> LZ)

instance CartesianCategory L where
    type Product L = (,)

    proj1 = L lfst
    proj2 = L lsnd

    fanout (L f) (L g) = L $ \x -> LH (f x) (g x)

instance CategoryWith0 L where
    type Initial L = ()

    initial = L linitial

-- Is this correct?
instance CocartesianCategory L where
    type Coproduct L = (,)

    inl = L $ \f -> LH f LZ
    inr = L $ \g -> LH LZ g

    fanin (L f) (L g) = L $ \x -> LA (f (lfst x)) (g (lsnd x))

class HasDim a where
    type Dim a :: Nat

    dimDict :: Proxy a -> Dict (KnownNat (Dim a))

    splitPair :: (a ~ (b, c)) => (Dict (HasDim b), Dict (HasDim c))
    splitPair = error "impossible: splitPair"

instance HasDim () where
    type Dim () = 0
    dimDict _ = Dict

instance HasDim Double where
    type Dim Double = 1
    dimDict _ = Dict

instance (HasDim a, HasDim b) => HasDim (a, b) where
    type Dim (a, b) = Dim a + Dim b

    dimDict _ =
        withDimDict (Proxy :: Proxy a) $
        withDimDict (Proxy :: Proxy b) $
        withDict (C.plusNat :: (KnownNat (Dim a), KnownNat (Dim b)) :- KnownNat (Dim a + Dim b))
        Dict

    splitPair = (Dict, Dict)

withDimDict :: HasDim a => Proxy a -> (KnownNat (Dim a) => r) -> r
withDimDict p = withDict (dimDict p)

dim :: forall a. HasDim a => Proxy a -> Int
dim p = withDimDict p $ fromInteger $ natVal (Proxy :: Proxy (Dim a))

toRawMatrix :: forall a b. (HasDim a, HasDim b) => LinMap a b -> L.Matrix Double
toRawMatrix LZ       = (dim (Proxy :: Proxy a) L.>< dim (Proxy :: Proxy b)) (repeat 0)
toRawMatrix (LD k)   = L.scale k (L.ident (dim (Proxy :: Proxy a)))
toRawMatrix (LA f g) = L.add (toRawMatrix f) (toRawMatrix g)
toRawMatrix (LH f g) = go splitPair f g where
    go :: (Dict (HasDim x), Dict (HasDim y)) -> LinMap a x -> LinMap a y -> L.Matrix Double
    go (Dict, Dict) f' g' = toRawMatrix f' L.||| toRawMatrix g'
toRawMatrix (LV f g) = go splitPair f g where
    go :: (Dict (HasDim x), Dict (HasDim y)) -> LinMap x b -> LinMap y b -> L.Matrix Double
    go (Dict, Dict) f' g' = toRawMatrix f' L.=== toRawMatrix g'

evalL :: (HasDim a, HasDim b) => L a b -> L.Matrix Double
evalL (L f) = toRawMatrix (f (LD 1))

-- toStaticMatrix :: forall a b. (HasDim a, HasDim b) => LinMap a b -> LS.L (Dim a) (Dim b)
-- toStaticMatrix LZ =
--     withDimDict (Proxy :: Proxy a) $
--     withDimDict (Proxy :: Proxy b) 0
-- toStaticMatrix LI =
--     withDimDict (Proxy :: Proxy a) LS.eye
-- toStaticMatrix (LA f g) =
--     withDimDict (Proxy :: Proxy a) $
--     withDimDict (Proxy :: Proxy b) $
--     L.add (toStaticMatrix f) (toStaticMatrix g)
-- toStaticMatrix (LK k f) =
--     withDimDict (Proxy :: Proxy a) $
--     withDimDict (Proxy :: Proxy b) $
--     toStaticMatrix f LS.<> LS.diag (LS.konst k)
-- toStaticMatrix (LH f g) = go splitPair f g where
--     go :: forall x y. (x,y) ~ b => (Dict (HasDim x), Dict (HasDim y)) -> LinMap a x -> LinMap a y -> LS.L (Dim a) (Dim x + Dim y)
--     go (Dict, Dict) f' g' =
--         withDimDict (Proxy :: Proxy a) $
--         withDimDict (Proxy :: Proxy b) $
--         withDimDict (Proxy :: Proxy x) $
--         withDimDict (Proxy :: Proxy y) $
--         toStaticMatrix f' LS.||| toStaticMatrix g'
-- toStaticMatrix (LV f g) = go splitPair f g where
--     go :: forall x y. (x,y) ~ a => (Dict (HasDim x), Dict (HasDim y)) -> LinMap x b -> LinMap y b -> LS.L (Dim x + Dim y) (Dim b)
--     go (Dict, Dict) f' g' =
--         withDimDict (Proxy :: Proxy a) $
--         withDimDict (Proxy :: Proxy b) $
--         withDimDict (Proxy :: Proxy x) $
--         withDimDict (Proxy :: Proxy y) $
--         toStaticMatrix f' LS.=== toStaticMatrix g'

-------------------------------------------------------------------------------
-- Vector space
-------------------------------------------------------------------------------

class HasDim a => VectorSpace a where
    toVector' :: a -> [Double] -> [Double]

    fromVector' :: [Double] -> (a -> [Double] -> r) -> r

toVector :: VectorSpace a => a -> [Double]
toVector x = toVector' x []

fromVector :: VectorSpace a => [Double] -> a
fromVector ds = fromVector' ds const

instance VectorSpace Double where
    toVector' d = (d :)

    fromVector' []     k = k 0 []
    fromVector' (d:ds) k = k d ds

instance (VectorSpace a, VectorSpace b) => VectorSpace (a, b) where
    toVector' (a, b) = toVector' a . toVector' b

    fromVector' xs k =
        fromVector' xs $ \a ys ->
        fromVector' ys $ \b zs ->
        k (a, b) zs


{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE TupleSections #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TemplateHaskell #-}

module Day15 where

import Control.Lens
import Control.Lens.TH

import qualified Data.ByteString as BS
import Data.Attoparsec.ByteString (Parser)
import qualified Data.Attoparsec.ByteString as DAB
import qualified Data.Attoparsec.ByteString.Char8 as DABC8
import Data.Attoparsec.ByteString.Char8 (decimal, inClass, digit, endOfLine, char, anyChar, letter_ascii, notChar)

import Data.Map.Strict (Map)
import qualified Data.Map.Strict as Map

import qualified Data.List as DL
import qualified Data.List.Split as DLS

import Debug.Trace

import Control.Arrow

import Data.Array.IArray (Array)
import qualified Data.Array.IArray as IA

import Data.Set (Set)
import qualified Data.Set as Set
import Data.Either (isRight)
import Data.Maybe (mapMaybe, catMaybes)
import Data.Bool (bool)

import Data.Vector (Vector, (!))
import qualified Data.Vector as Vec

import Data.Function (on)

import Control.Monad.State.Strict (State)
import qualified Control.Monad.State.Strict as State

import qualified Control.Monad as CM

import qualified Control.Foldl as L

import Data.Bifunctor (bimap)

import Data.Complex
import Data.Functor

import Data.Word
import Data.Bits
import Data.Semigroup
import Data.List.NonEmpty (NonEmpty(..))
import qualified Data.List.NonEmpty as DLNE

type Input = Integer

solver input fTurn = State.evalState (code 6 (fromIntegral $ length input + 1)) $ -- turns start at length input + 1
	Map.fromList $ zip input $ fmap (,Nothing) [1..] --Insert the seed input with their order
	where
	-- Monadic loop. Maybe avoid custom recursion
	-- It returns the fTurn-th value
	code :: Integer -> Integer -> State (Map Integer (Integer, Maybe Integer)) Integer
	code prev turn
		| turn > fTurn = pure prev
		| otherwise =
			-- Check is the previous value is repeated
			State.get >>= \m -> case Map.lookup prev m of
				Just (_, Nothing) -> do --not repeated
					State.modify (Map.insertWith inserter 0 (turn, Nothing))
					code 0 (turn + 1)
				Just (l, Just ll) -> do -- repeated
					State.modify (Map.insertWith inserter (l-ll) (turn, Nothing))
					code (l - ll) (turn + 1)

	-- Update the map keeping the previous position
	-- DONE: Use a tupple (Integer, Maybe Integer)
	-- DONE: avoid error (removed safety check)
	inserter :: (Integer, Maybe Integer) -> (Integer, Maybe Integer) -> (Integer, Maybe Integer)
	inserter (new, Nothing) (prev, _) = (new, Just prev)
		-- | new > prev =  (new, Just prev)
		-- | otherwise = error $ show (prev, new)

runSolution :: FilePath -> IO ()
runSolution _ = do
	putStrLn "**Day 15**"
	print $ solver input 30000000
	where
	input :: [Input]
	--input = [0,3,6]
	input = [0,13,1,16,6,17]


{-# LANGUAGE CPP                       #-}
{-# LANGUAGE BangPatterns              #-}
{-# LANGUAGE DataKinds                 #-}
{-# LANGUAGE ExistentialQuantification #-}
{-# LANGUAGE FlexibleContexts          #-}
{-# LANGUAGE GADTs                     #-}
{-# LANGUAGE RankNTypes                #-}
{-# LANGUAGE ScopedTypeVariables       #-}
{-# LANGUAGE TypeFamilies              #-}

import           Control.Applicative
import           Control.Monad.Random
import           Options.Applicative

import           Data.Function                (on)
import           Data.List

import           Graphics.Image               hiding (map, on)

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

import           Grenade
import           Grenade.Utils.ImageNet

data ResNetOptions = ResNetOptions FilePath         -- onnx file
                                   FilePath         -- input image

pResnet :: Parser ResNetOptions
pResnet = ResNetOptions <$> argument str (metavar "onnx") <*> argument str (metavar "image")

-- loadResNetImage :: FilePath -> IO (Maybe (S ('D3 224 224 3)))
loadResNetImage :: FilePath -> IO (Maybe (S ('D3 224 224 3)))
loadResNetImage path = do
  img <- readImageRGB VU path
  displayImage img
  return $ do
    guard $ dims img == (224, 224)
    let [img_red, img_green, img_blue] = toImagesX img
        [reds, greens, blues]          = map (map (\(PixelX y) -> doubleToRealNum y) . concat . toLists) [img_red, img_green, img_blue]

        redM   = H.dmmap (\a -> (a - 0.485) / 0.229) . H.tr $ H.fromList reds   :: H.L 224 224
        greenM = H.dmmap (\a -> (a - 0.456) / 0.224) . H.tr $ H.fromList greens :: H.L 224 224
        blueM  = H.dmmap (\a -> (a - 0.406) / 0.225) . H.tr $ H.fromList blues  :: H.L 224 224
        mat    = redM H.=== greenM H.=== blueM

    return (S3D mat)

main :: IO ()
main = do
    ResNetOptions netPath imgPath <- execParser (info (pResnet <**> helper) idm)
    res <- loadResNet netPath

    inputM <- loadResNetImage imgPath

    case (res, inputM) of
      (Right net, Just input)  -> do
        let S1D y = runNet net input
            tops  = getTop 5 $ LA.toList $ H.extract y
        mapM_ (\(i :: Int, x) -> print $ show i ++ ": " ++ (show . getLabel) x) $ zip [1..] tops
      (Left err, _) -> print err
      _             -> error "Could not load image"
  where
    getTop :: Ord a => Int -> [a] -> [Int]
    getTop n xs = map fst $ take n $ sortBy (flip compare `on` snd) $ zip [0..] xs


{-# LANGUAGE AllowAmbiguousTypes   #-}
{-# LANGUAGE DataKinds             #-}
{-# LANGUAGE DeriveGeneric         #-}
{-# LANGUAGE DerivingStrategies    #-}
{-# LANGUAGE DerivingVia           #-}
{-# LANGUAGE FlexibleContexts      #-}
{-# LANGUAGE FlexibleInstances     #-}
{-# LANGUAGE GADTs                 #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE RecordWildCards       #-}
{-# LANGUAGE ScopedTypeVariables   #-}
{-# LANGUAGE TemplateHaskell       #-}
{-# LANGUAGE TupleSections         #-}
{-# LANGUAGE TypeApplications      #-}
{-# LANGUAGE UndecidableInstances  #-}
{-# LANGUAGE ViewPatterns          #-}
{-# OPTIONS_GHC -fno-warn-orphans  #-}

import           Control.DeepSeq
import           Control.Lens hiding                          ((<.>))
import           Control.Monad.Primitive
import           Control.Monad.Trans.Maybe
import           Data.Bitraversable
import           Data.Default
import           Data.IDX
import           Data.Mutable
import           Data.Traversable
import           Data.Tuple
import           GHC.Generics                                 (Generic)
import           GHC.TypeLits
import           Numeric.Backprop hiding                      (auto)
import           Numeric.LinearAlgebra.Static.Backprop hiding ((<>))
import           Numeric.OneLiner
import           Numeric.Opto hiding                          ((<.>))
import           Numeric.Opto.Backprop
import           Numeric.Opto.Run.Simple
import           Options.Applicative
import           System.FilePath hiding                       ((<.>))
import           Text.Printf
import qualified Data.Vector.Generic                          as VG
import qualified Numeric.LinearAlgebra                        as HM
import qualified Numeric.LinearAlgebra.Static                 as H
import qualified System.Random.MWC                            as MWC

data OMode = OSingle
           | OParallel Bool Int

data Opts = Opts
    { oDataDir :: FilePath
    , oReport  :: Int
    , oBatch   :: Int
    , oMode    :: OMode
    }

parseOMode :: Parser OMode
parseOMode = subparser
    ( command "single" (info (pure OSingle) (progDesc "Single-threaded"))
   <> command "parallel" (info (uncurry OParallel <$> parseParallel) (progDesc "Parallel"))
    )
  where
    parseParallel :: Parser (Bool, Int)
    parseParallel = (,)
        <$> switch (long "chunked" <> help "Chunked mode")
        <*> option auto
                ( long "split"
               <> short 's'
               <> help "Number of items per thread"
               <> metavar "INT"
               <> showDefault
               <> value 750
                )

parseOpts :: Parser Opts
parseOpts = Opts
    <$> strArgument ( help "Data directory (containing uncompressed MNIST data set)"
                   <> metavar "DIR"
                    )
    <*> option auto
          ( long "report"
         <> short 'r'
         <> help "Report frequency (in batches)"
         <> metavar "INT"
         <> showDefault
         <> value 2500
          )
    <*> option auto
          ( long "batch"
         <> short 'b'
         <> help "Batching amount"
         <> metavar "INT"
         <> showDefault
         <> value 1
          )
    <*> (parseOMode <|> pure OSingle)

data Net = N { _weights1 :: !(L 250 784)
             , _bias1    :: !(R 250)
             , _weights2 :: !(L 10 250)
             , _bias2    :: !(R 10)
             }
  deriving (Generic)
  deriving (Num, Fractional, Floating) via (GNum Net)
makeLenses ''Net

instance Linear Double Net
instance Mutable q Net
instance LinearInPlace q Double Net

logistic :: Floating a => a -> a
logistic x = 1 / (1 + exp (-x))

softMax
    :: Reifies s W
    => BVar s (R 10)
    -> BVar s (R 10)
softMax x = expx / konst (norm_1V expx)
  where
    expx = exp x

crossEntropy
    :: Reifies s W
    => BVar s (R 10)
    -> BVar s (R 10)
    -> BVar s Double
crossEntropy targ res = -(log res <.> targ)

runNet
    :: Reifies s W
    => BVar s Net
    -> BVar s (R 784)
    -> BVar s (R 10)
runNet n x = z
  where
    y = logistic $ (n ^^. weights1) #> x + (n ^^. bias1)
    z = softMax  $ (n ^^. weights2) #> y + (n ^^. bias2)

netErr
    :: Reifies s W
    => BVar s (R 784)
    -> BVar s (R 10)
    -> BVar s Net
    -> BVar s Double
netErr x targ n = crossEntropy targ (runNet n x)

-- *********************************************
-- Plumbing for running the network on real data
-- *********************************************

main :: IO ()
main = MWC.withSystemRandom $ \g -> do
    Opts{..} <- execParser $ info (parseOpts <**> helper)
        ( fullDesc
       <> progDesc "Run optimizer samples on MNIST"
       <> header "opto-neural - opto runner on MNIST data set"
        )

    Just train <- loadMNIST (oDataDir </> "train-images-idx3-ubyte")
                            (oDataDir </> "train-labels-idx1-ubyte")
    Just test  <- loadMNIST (oDataDir </> "t10k-images-idx3-ubyte")
                            (oDataDir </> "t10k-labels-idx1-ubyte")
    putStrLn "Loaded data."
    net0 <- MWC.uniformR (-0.5, 0.5) g

    let o :: Opto (PrimState IO) (R 784, R 10) Net
        o = adam def $
              bpGradSample $ \(x, y) -> netErr (constVar x) (constVar y)

        runTest chnk net = printf "Error: %.2f%%" ((1 - score) * 100)
          where
            score = testNet chnk net

        ro = def { roBatch = oBatch
                 }
        so = def { soTestSet   = Just test
                 , soEvaluate  = runTest
                 , soSkipSamps = oReport
                 }

    case oMode of
      OSingle       -> simpleRunner so train SOSingle ro net0 o g
      OParallel c s -> do
        let po = def { poSplit = s }
        if c
          then simpleRunner so train (SOParallel   po) ro net0 o g
          else simpleRunner so train (SOParChunked po) ro net0 o g

testNet :: [(R 784, R 10)] -> Net -> Double
testNet xs n = sum (map (uncurry test) xs) / fromIntegral (length xs)
  where
    test x (H.extract->t)
        | HM.maxIndex t == HM.maxIndex (H.extract r) = 1
        | otherwise                                  = 0
      where
        r = evalBP (`runNet` constVar x) n

loadMNIST
    :: FilePath
    -> FilePath
    -> IO (Maybe [(R 784, R 10)])
loadMNIST fpI fpL = runMaybeT $ do
    i <- MaybeT          $ decodeIDXFile       fpI
    l <- MaybeT          $ decodeIDXLabelsFile fpL
    d <- MaybeT . return $ labeledIntData l i
    MaybeT . return $ for d (bitraverse mkImage mkLabel . swap)
  where
    mkImage = H.create . VG.convert . VG.map (\i -> fromIntegral i / 255)
    mkLabel n = H.create $ HM.build 10 (\i -> if round i == n then 1 else 0)

instance KnownNat n => MWC.Variate (R n) where
    uniform g = H.randomVector <$> MWC.uniform g <*> pure H.Uniform
    uniformR (l, h) g = (\x -> x * (h - l) + l) <$> MWC.uniform g

instance (KnownNat m, KnownNat n) => MWC.Variate (L m n) where
    uniform g = H.uniformSample <$> MWC.uniform g <*> pure 0 <*> pure 1
    uniformR (l, h) g = (\x -> x * (h - l) + l) <$> MWC.uniform g

instance MWC.Variate Net where
    uniform g = N <$> MWC.uniform g
                  <*> MWC.uniform g
                  <*> MWC.uniform g
                  <*> MWC.uniform g
    uniformR (l, h) g = (\x -> x * (h - l) + l) <$> MWC.uniform g

instance NFData Net

instance Backprop Net


{-# LANGUAGE CPP #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE UndecidableInstances #-}

-----------------------------------------------------------------------------
-- |
-- Module      :  Numeric.ContainerBoot
-- Copyright   :  (c) Alberto Ruiz 2010
-- License     :  GPL-style
--
-- Maintainer  :  Alberto Ruiz <aruiz@um.es>
-- Stability   :  provisional
-- Portability :  portable
--
-- Module to avoid cyclyc dependencies.
--
-----------------------------------------------------------------------------

module Numeric.ContainerBoot (
    -- * Basic functions
    ident, diag, ctrans,
    -- * Generic operations
    Container(..),
    -- * Matrix product and related functions
    Product(..),
    mXm,mXv,vXm,
    outer, kronecker,
    -- * Element conversion
    Convert(..),
    Complexable(),
    RealElement(),

    RealOf, ComplexOf, SingleOf, DoubleOf,

    IndexOf,
    module Data.Complex,
    -- * Experimental
    build', konst'
) where

import Data.Packed
import Data.Packed.ST as ST
import Numeric.Conversion
import Data.Packed.Internal
import Numeric.GSL.Vector
import Data.Complex
import Control.Monad(ap)

import Numeric.LinearAlgebra.LAPACK(multiplyR,multiplyC,multiplyF,multiplyQ)

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

type family IndexOf (c :: * -> *)

type instance IndexOf Vector = Int
type instance IndexOf Matrix = (Int,Int)

type family ArgOf (c :: * -> *) a

type instance ArgOf Vector a = a -> a
type instance ArgOf Matrix a = a -> a -> a

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

-- | Basic element-by-element functions for numeric containers
class (Complexable c, Fractional e, Element e) => Container c e where
    -- | create a structure with a single element
    scalar      :: e -> c e
    -- | complex conjugate
    conj        :: c e -> c e
    scale       :: e -> c e -> c e
    -- | scale the element by element reciprocal of the object:
    --
    -- @scaleRecip 2 (fromList [5,i]) == 2 |> [0.4 :+ 0.0,0.0 :+ (-2.0)]@
    scaleRecip  :: e -> c e -> c e
    addConstant :: e -> c e -> c e
    add         :: c e -> c e -> c e
    sub         :: c e -> c e -> c e
    -- | element by element multiplication
    mul         :: c e -> c e -> c e
    -- | element by element division
    divide      :: c e -> c e -> c e
    equal       :: c e -> c e -> Bool
    --
    -- element by element inverse tangent
    arctan2     :: c e -> c e -> c e
    --
    -- | cannot implement instance Functor because of Element class constraint
    cmap        :: (Element b) => (e -> b) -> c e -> c b
    -- | constant structure of given size
    konst       :: e -> IndexOf c -> c e
    -- | create a structure using a function
    --
    -- Hilbert matrix of order N:
    --
    -- @hilb n = build (n,n) (\\i j -> 1/(i+j+1))@
    build       :: IndexOf c -> (ArgOf c e) -> c e
    --build       :: BoundsOf f -> f -> (ContainerOf f) e
    --
    -- | indexing function
    atIndex     :: c e -> IndexOf c -> e
    -- | index of min element
    minIndex    :: c e -> IndexOf c
    -- | index of max element
    maxIndex    :: c e -> IndexOf c
    -- | value of min element
    minElement  :: c e -> e
    -- | value of max element
    maxElement  :: c e -> e
    -- the C functions sumX/prodX are twice as fast as using foldVector
    -- | the sum of elements (faster than using @fold@)
    sumElements :: c e -> e
    -- | the product of elements (faster than using @fold@)
    prodElements :: c e -> e

    -- | A more efficient implementation of @cmap (\\x -> if x>0 then 1 else 0)@
    --
    -- @> step $ linspace 5 (-1,1::Double)
    -- 5 |> [0.0,0.0,0.0,1.0,1.0]@

    step :: RealElement e => c e -> c e

    -- | Element by element version of @case compare a b of {LT -> l; EQ -> e; GT -> g}@.
    --
    -- Arguments with any dimension = 1 are automatically expanded:
    --
    -- @> cond ((1>\<4)[1..]) ((3>\<1)[1..]) 0 100 ((3>\<4)[1..]) :: Matrix Double
    -- (3><4)
    -- [ 100.0,   2.0,   3.0,  4.0
    -- ,   0.0, 100.0,   7.0,  8.0
    -- ,   0.0,   0.0, 100.0, 12.0 ]@

    cond :: RealElement e
         => c e -- ^ a
         -> c e -- ^ b
         -> c e -- ^ l
         -> c e -- ^ e
         -> c e -- ^ g
         -> c e -- ^ result

    -- | Find index of elements which satisfy a predicate
    --
    -- @> find (>0) (ident 3 :: Matrix Double)
    -- [(0,0),(1,1),(2,2)]@

    find :: (e -> Bool) -> c e -> [IndexOf c]

    -- | Create a structure from an association list
    --
    -- @> assoc 5 0 [(2,7),(1,3)] :: Vector Double
    -- 5 |> [0.0,3.0,7.0,0.0,0.0]@

    assoc :: IndexOf c        -- ^ size
          -> e                -- ^ default value
          -> [(IndexOf c, e)] -- ^ association list
          -> c e              -- ^ result

    -- | Modify a structure using an update function
    --
    -- @> accum (ident 5) (+) [((1,1),5),((0,3),3)] :: Matrix Double
    -- (5><5)
    --  [ 1.0, 0.0, 0.0, 3.0, 0.0
    --  , 0.0, 6.0, 0.0, 0.0, 0.0
    --  , 0.0, 0.0, 1.0, 0.0, 0.0
    --  , 0.0, 0.0, 0.0, 1.0, 0.0
    --  , 0.0, 0.0, 0.0, 0.0, 1.0 ]@

    accum :: c e              -- ^ initial structure
          -> (e -> e -> e)    -- ^ update function
          -> [(IndexOf c, e)] -- ^ association list
          -> c e              -- ^ result

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

instance Container Vector Float where
    scale = vectorMapValF Scale
    scaleRecip = vectorMapValF Recip
    addConstant = vectorMapValF AddConstant
    add = vectorZipF Add
    sub = vectorZipF Sub
    mul = vectorZipF Mul
    divide = vectorZipF Div
    equal u v = dim u == dim v && maxElement (vectorMapF Abs (sub u v)) == 0.0
    arctan2 = vectorZipF ATan2
    scalar x = fromList [x]
    konst = constantD
    build = buildV
    conj = id
    cmap = mapVector
    atIndex = (@>)
    minIndex     = round . toScalarF MinIdx
    maxIndex     = round . toScalarF MaxIdx
    minElement  = toScalarF Min
    maxElement  = toScalarF Max
    sumElements  = sumF
    prodElements = prodF
    step = stepF
    find = findV
    assoc = assocV
    accum = accumV
    cond = condV condF

instance Container Vector Double where
    scale = vectorMapValR Scale
    scaleRecip = vectorMapValR Recip
    addConstant = vectorMapValR AddConstant
    add = vectorZipR Add
    sub = vectorZipR Sub
    mul = vectorZipR Mul
    divide = vectorZipR Div
    equal u v = dim u == dim v && maxElement (vectorMapR Abs (sub u v)) == 0.0
    arctan2 = vectorZipR ATan2
    scalar x = fromList [x]
    konst = constantD
    build = buildV
    conj = id
    cmap = mapVector
    atIndex = (@>)
    minIndex     = round . toScalarR MinIdx
    maxIndex     = round . toScalarR MaxIdx
    minElement  = toScalarR Min
    maxElement  = toScalarR Max
    sumElements  = sumR
    prodElements = prodR
    step = stepD
    find = findV
    assoc = assocV
    accum = accumV
    cond = condV condD

instance Container Vector (Complex Double) where
    scale = vectorMapValC Scale
    scaleRecip = vectorMapValC Recip
    addConstant = vectorMapValC AddConstant
    add = vectorZipC Add
    sub = vectorZipC Sub
    mul = vectorZipC Mul
    divide = vectorZipC Div
    equal u v = dim u == dim v && maxElement (mapVector magnitude (sub u v)) == 0.0
    arctan2 = vectorZipC ATan2
    scalar x = fromList [x]
    konst = constantD
    build = buildV
    conj = conjugateC
    cmap = mapVector
    atIndex = (@>)
    minIndex     = minIndex . fst . fromComplex . (zipVectorWith (*) `ap` mapVector conjugate)
    maxIndex     = maxIndex . fst . fromComplex . (zipVectorWith (*) `ap` mapVector conjugate)
    minElement  = ap (@>) minIndex
    maxElement  = ap (@>) maxIndex
    sumElements  = sumC
    prodElements = prodC
    step = undefined -- cannot match
    find = findV
    assoc = assocV
    accum = accumV
    cond = undefined -- cannot match

instance Container Vector (Complex Float) where
    scale = vectorMapValQ Scale
    scaleRecip = vectorMapValQ Recip
    addConstant = vectorMapValQ AddConstant
    add = vectorZipQ Add
    sub = vectorZipQ Sub
    mul = vectorZipQ Mul
    divide = vectorZipQ Div
    equal u v = dim u == dim v && maxElement (mapVector magnitude (sub u v)) == 0.0
    arctan2 = vectorZipQ ATan2
    scalar x = fromList [x]
    konst = constantD
    build = buildV
    conj = conjugateQ
    cmap = mapVector
    atIndex = (@>)
    minIndex     = minIndex . fst . fromComplex . (zipVectorWith (*) `ap` mapVector conjugate)
    maxIndex     = maxIndex . fst . fromComplex . (zipVectorWith (*) `ap` mapVector conjugate)
    minElement  = ap (@>) minIndex
    maxElement  = ap (@>) maxIndex
    sumElements  = sumQ
    prodElements = prodQ
    step = undefined -- cannot match
    find = findV
    assoc = assocV
    accum = accumV
    cond = undefined -- cannot match

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

instance (Container Vector a) => Container Matrix a where
    scale x = liftMatrix (scale x)
    scaleRecip x = liftMatrix (scaleRecip x)
    addConstant x = liftMatrix (addConstant x)
    add = liftMatrix2 add
    sub = liftMatrix2 sub
    mul = liftMatrix2 mul
    divide = liftMatrix2 divide
    equal a b = cols a == cols b && flatten a `equal` flatten b
    arctan2 = liftMatrix2 arctan2
    scalar x = (1><1) [x]
    konst v (r,c) = reshape c (konst v (r*c))
    build = buildM
    conj = liftMatrix conj
    cmap f = liftMatrix (mapVector f)
    atIndex = (@@>)
    minIndex m = let (r,c) = (rows m,cols m)
                     i = (minIndex $ flatten m)
                 in (i `div` c,i `mod` c)
    maxIndex m = let (r,c) = (rows m,cols m)
                     i = (maxIndex $ flatten m)
                 in (i `div` c,i `mod` c)
    minElement = ap (@@>) minIndex
    maxElement = ap (@@>) maxIndex
    sumElements = sumElements . flatten
    prodElements = prodElements . flatten
    step = liftMatrix step
    find = findM
    assoc = assocM
    accum = accumM
    cond = condM

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

-- | Matrix product and related functions
class Element e => Product e where
    -- | matrix product
    multiply :: Matrix e -> Matrix e -> Matrix e
    -- | dot (inner) product
    dot        :: Vector e -> Vector e -> e
    -- | sum of absolute value of elements (differs in complex case from @norm1@)
    absSum     :: Vector e -> RealOf e
    -- | sum of absolute value of elements
    norm1      :: Vector e -> RealOf e
    -- | euclidean norm
    norm2      :: Vector e -> RealOf e
    -- | element of maximum magnitude
    normInf    :: Vector e -> RealOf e

instance Product Float where
    norm2      = toScalarF Norm2
    absSum     = toScalarF AbsSum
    dot        = dotF
    norm1      = toScalarF AbsSum
    normInf    = maxElement . vectorMapF Abs
    multiply = multiplyF

instance Product Double where
    norm2      = toScalarR Norm2
    absSum     = toScalarR AbsSum
    dot        = dotR
    norm1      = toScalarR AbsSum
    normInf    = maxElement . vectorMapR Abs
    multiply = multiplyR

instance Product (Complex Float) where
    norm2      = toScalarQ Norm2
    absSum     = toScalarQ AbsSum
    dot        = dotQ
    norm1      = sumElements . fst . fromComplex . vectorMapQ Abs
    normInf    = maxElement . fst . fromComplex . vectorMapQ Abs
    multiply = multiplyQ

instance Product (Complex Double) where
    norm2      = toScalarC Norm2
    absSum     = toScalarC AbsSum
    dot        = dotC
    norm1      = sumElements . fst . fromComplex . vectorMapC Abs
    normInf    = maxElement . fst . fromComplex . vectorMapC Abs
    multiply = multiplyC

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

-- synonym for matrix product
mXm :: Product t => Matrix t -> Matrix t -> Matrix t
mXm = multiply

-- matrix - vector product
mXv :: Product t => Matrix t -> Vector t -> Vector t
mXv m v = flatten $ m `mXm` (asColumn v)

-- vector - matrix product
vXm :: Product t => Vector t -> Matrix t -> Vector t
vXm v m = flatten $ (asRow v) `mXm` m

{- | Outer product of two vectors.

@\> 'fromList' [1,2,3] \`outer\` 'fromList' [5,2,3]
(3><3)
 [  5.0, 2.0, 3.0
 , 10.0, 4.0, 6.0
 , 15.0, 6.0, 9.0 ]@
-}
outer :: (Product t) => Vector t -> Vector t -> Matrix t
outer u v = asColumn u `multiply` asRow v

{- | Kronecker product of two matrices.

@m1=(2><3)
 [ 1.0,  2.0, 0.0
 , 0.0, -1.0, 3.0 ]
m2=(4><3)
 [  1.0,  2.0,  3.0
 ,  4.0,  5.0,  6.0
 ,  7.0,  8.0,  9.0
 , 10.0, 11.0, 12.0 ]@

@\> kronecker m1 m2
(8><9)
 [  1.0,  2.0,  3.0,   2.0,   4.0,   6.0,  0.0,  0.0,  0.0
 ,  4.0,  5.0,  6.0,   8.0,  10.0,  12.0,  0.0,  0.0,  0.0
 ,  7.0,  8.0,  9.0,  14.0,  16.0,  18.0,  0.0,  0.0,  0.0
 , 10.0, 11.0, 12.0,  20.0,  22.0,  24.0,  0.0,  0.0,  0.0
 ,  0.0,  0.0,  0.0,  -1.0,  -2.0,  -3.0,  3.0,  6.0,  9.0
 ,  0.0,  0.0,  0.0,  -4.0,  -5.0,  -6.0, 12.0, 15.0, 18.0
 ,  0.0,  0.0,  0.0,  -7.0,  -8.0,  -9.0, 21.0, 24.0, 27.0
 ,  0.0,  0.0,  0.0, -10.0, -11.0, -12.0, 30.0, 33.0, 36.0 ]@
-}
kronecker :: (Product t) => Matrix t -> Matrix t -> Matrix t
kronecker a b = fromBlocks
              . splitEvery (cols a)
              . map (reshape (cols b))
              . toRows
              $ flatten a `outer` flatten b

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

class Convert t where
    real    :: Container c t => c (RealOf t) -> c t
    complex :: Container c t => c t -> c (ComplexOf t)
    single  :: Container c t => c t -> c (SingleOf t)
    double  :: Container c t => c t -> c (DoubleOf t)
    toComplex   :: (Container c t, RealElement t) => (c t, c t) -> c (Complex t)
    fromComplex :: (Container c t, RealElement t) => c (Complex t) -> (c t, c t)

instance Convert Double where
    real = id
    complex = comp'
    single = single'
    double = id
    toComplex = toComplex'
    fromComplex = fromComplex'

instance Convert Float where
    real = id
    complex = comp'
    single = id
    double = double'
    toComplex = toComplex'
    fromComplex = fromComplex'

instance Convert (Complex Double) where
    real = comp'
    complex = id
    single = single'
    double = id
    toComplex = toComplex'
    fromComplex = fromComplex'

instance Convert (Complex Float) where
    real = comp'
    complex = id
    single = id
    double = double'
    toComplex = toComplex'
    fromComplex = fromComplex'

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

type family RealOf x

type instance RealOf Double = Double
type instance RealOf (Complex Double) = Double

type instance RealOf Float = Float
type instance RealOf (Complex Float) = Float

type family ComplexOf x

type instance ComplexOf Double = Complex Double
type instance ComplexOf (Complex Double) = Complex Double

type instance ComplexOf Float = Complex Float
type instance ComplexOf (Complex Float) = Complex Float

type family SingleOf x

type instance SingleOf Double = Float
type instance SingleOf Float  = Float

type instance SingleOf (Complex a) = Complex (SingleOf a)

type family DoubleOf x

type instance DoubleOf Double = Double
type instance DoubleOf Float  = Double

type instance DoubleOf (Complex a) = Complex (DoubleOf a)

type family ElementOf c

type instance ElementOf (Vector a) = a
type instance ElementOf (Matrix a) = a

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

class Build f where
    build' :: BoundsOf f -> f -> ContainerOf f

type family BoundsOf x

type instance BoundsOf (a->a) = Int
type instance BoundsOf (a->a->a) = (Int,Int)

type family ContainerOf x

type instance ContainerOf (a->a) = Vector a
type instance ContainerOf (a->a->a) = Matrix a

instance (Element a, Num a) => Build (a->a) where
    build' = buildV

instance (Element a, Num a) => Build (a->a->a) where
    build' = buildM

buildM (rc,cc) f = fromLists [ [f r c | c <- cs] | r <- rs ]
    where rs = map fromIntegral [0 .. (rc-1)]
          cs = map fromIntegral [0 .. (cc-1)]

buildV n f = fromList [f k | k <- ks]
    where ks = map fromIntegral [0 .. (n-1)]

----------------------------------------------------
-- experimental

class Konst s where
    konst' :: Element e => e -> s -> ContainerOf' s e

type family ContainerOf' x y

type instance ContainerOf' Int a = Vector a
type instance ContainerOf' (Int,Int) a = Matrix a

instance Konst Int where
    konst' = constantD

instance Konst (Int,Int) where
    konst' k (r,c) = reshape c $ konst' k (r*c)

--------------------------------------------------------
-- | conjugate transpose
ctrans :: (Container Vector e, Element e) => Matrix e -> Matrix e
ctrans = liftMatrix conj . trans

-- | Creates a square matrix with a given diagonal.
diag :: (Num a, Element a) => Vector a -> Matrix a
diag v = diagRect 0 v n n where n = dim v

-- | creates the identity matrix of given dimension
ident :: (Num a, Element a) => Int -> Matrix a
ident n = diag (constantD 1 n)

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

findV p x = foldVectorWithIndex g [] x where
    g k z l = if p z then k:l else l

findM p x = map ((`divMod` cols x)) $ findV p (flatten x)

assocV n z xs = ST.runSTVector $ do
        v <- ST.newVector z n
        mapM_ (\(k,x) -> ST.writeVector v k x) xs
        return v

assocM (r,c) z xs = ST.runSTMatrix $ do
        m <- ST.newMatrix z r c
        mapM_ (\((i,j),x) -> ST.writeMatrix m i j x) xs
        return m

accumV v0 f xs = ST.runSTVector $ do
        v <- ST.thawVector v0
        mapM_ (\(k,x) -> ST.modifyVector v k (f x)) xs
        return v

accumM m0 f xs = ST.runSTMatrix $ do
        m <- ST.thawMatrix m0
        mapM_ (\((i,j),x) -> ST.modifyMatrix m i j (f x)) xs
        return m

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

condM a b l e t = reshape (cols a'') $ cond a' b' l' e' t'
  where
    args@(a'':_) = conformMs [a,b,l,e,t]
    [a', b', l', e', t'] = map flatten args

condV f a b l e t = f a' b' l' e' t'
  where
    [a', b', l', e', t'] = conformVs [a,b,l,e,t]