{-# LANGUAGE BangPatterns          #-}
{-# LANGUAGE FlexibleContexts      #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE OverloadedLists       #-}
{-# LANGUAGE ViewPatterns          #-}

-- | Main module
module SatOpt (main) where
--module SatOpt (module SatOpt) where

import           SatOpt.FFT                    as SatOpt
import           SatOpt.Optimize               as SatOpt
import           SatOpt.Render                 as SatOpt
import           SatOpt.Utility                as SatOpt
-- GENERATE: import New.Module as SatOpt

import           Data.Colour                   (blend)
import           Data.Colour.RGBSpace
import           Data.Colour.RGBSpace.HSL
import           Data.Complex
import           Data.List
import           Data.Packed.Matrix            (mapMatrix)
import           Debug.Trace                   (trace)
import           Numeric.LinearAlgebra.HMatrix hiding (disp, matrix, size)
import qualified Numeric.LinearAlgebra.HMatrix as HM

data Field2D a = FMat  (Matrix a)
               | FFT   (Matrix a)
               | FDisp (Matrix a)
               | FFunc (ℝ -> ℝ -> a)

data HSource = HPointSources  [(ℝ, ℝ, ℂ)]
             | HCircleSources ℝ [ℂ]
             | HSum [HSource]

hsrcToField :: HSource -> Field2D ℂ
hsrcToField (HPointSources ps)   = FFunc $ pointSources ps
hsrcToField (HCircleSources r p) = FFunc $ circleOfSources r p
hsrcToField (HSum [h])           = hsrcToField h
hsrcToField (HSum (h:hs))        = fieldSum (hsrcToField h) (hsrcToField (HSum hs))

fieldSum :: Field2D ℂ -> Field2D ℂ -> Field2D ℂ
fieldSum (FMat m1)  (FMat m2)  = FMat (m1 + m2)
fieldSum (FFunc f1) (FFunc f2) = FFunc (\x y -> f1 x y + f2 x y)
fieldSum (FMat m1)  (FFunc f2) = fieldSum (FFunc f2) (FMat m1)
fieldSum (FFunc f1) (FMat m2)  = fieldSum (matField (FFunc f1)) (FMat m2)

data HDisp = HConstDisp ℂ
           | HFuncDisp (ℝ -> ℝ -> ℂ)

hdspToField :: HDisp -> Field2D ℂ
hdspToField (HConstDisp p) = constF p
hdspToField (HFuncDisp f) = FFunc f

data HelmConf = HelmConf { helmsrc :: HSource
                         , helmdsp :: HDisp
                         }

data State a = State { amplitude  :: Field2D ℂ
                     , dispersion :: Field2D ℂ
                     , sources    :: Field2D ℂ
                     , time       :: ℝ
                     , freq       :: ℝ
                     , conf       :: a
                     , configure  :: State a -> State a
                     , concrete   :: Bool
                     }

constF :: Element a => a -> Field2D a
constF x = matField (FFunc (\_ _ -> x))

defState :: State a
defState = State { amplitude  = constF (0 :+ 0)
                 , dispersion = constF (1 :+ 0.1)
                 , sources    = constF (0 :+ 0)
                 , time       = 0
                 , freq       = 1
                 , conf       = undefined
                 , configure  = const defState
                 , concrete   = False
                 }

fromFieldC :: Field2D ℂ -> Matrix ℂ
fromFieldC (FFunc f) = sample2DM f [0, 1 .. size - 1] [0, 1 .. size - 1]
fromFieldC (FMat m)  = m
fromFieldC (FFT m)   = ifft2dM m
fromFieldC (FDisp m) = sqrt $ (recip m) - (square normm)

matField :: Element a => Field2D a -> Field2D a
matField (FFunc f) = FMat $ sample2DM f [0, 1 .. size - 1] [0, 1 .. size - 1]
matField x         = x

fftMat :: Field2D ℂ -> Matrix ℂ
fftMat (FFT m)  = m
fftMat (FMat m) = fft2dM m
fftMat f        = fftMat $ matField f

dispMat :: Field2D ℂ -> Matrix ℂ
dispMat (FDisp m) = m
dispMat (FMat m)  = recip $ square m + square normm
dispMat f         = dispMat $ matField f

makeConcrete :: (State a) -> (State a)
makeConcrete st@(State { concrete = True }) = st
makeConcrete (st@(State { dispersion = d, sources = s, time = t, freq = f }))
  = st { amplitude = ampl
       , dispersion = FDisp disp
       , sources = FFT srcs
       , concrete = True }
  where
    disp = dispMat $ matField d
    srcs = fftMat $ matField s
    ampl = FMat $ timeShift f t $ ifft2dM $ srcs * disp

phaseShift :: Matrix ℝ -> Matrix ℂ -> Matrix ℂ
phaseShift p m = mags * zipMatrixWith (:+) (cos angles) (sin angles)
  where
    angles = p + (mapMatrix realPart m `arctan2` mapMatrix imagPart m)
    mags = m * conj m

timeShift :: ℝ -> ℝ -> Matrix ℂ -> Matrix ℂ
timeShift f t m = case matField $ constF (tau * f * t) of
                   FMat p -> phaseShift p m
                   _      -> error "matField output invalid"

size :: Num a => a
size = 256

tau :: ℝ
tau = 2 * pi

rootsOfUnity :: Int -> [ℂ]
rootsOfUnity i = take i $ iterate (rotate (tau / fromIntegral i)) (1 :+ 0)

pointsOnCircle :: Int -> ℝ -> [(ℝ, ℝ)]
pointsOnCircle i r = map (fromC . scaleC r) $ rootsOfUnity i
  where
    fromC (a :+ b) = (a, b)
    scaleC s c = (s :+ 0) * c

-- | Take a position with respect to the center and move its origin to the top right
center :: Floating a => a -> a
center x = x - (size / 2)

pointSources :: [(ℝ, ℝ, ℂ)] -> ℝ -> ℝ -> ℂ
pointSources ps x y = sum $ map (\(cx, cy, s) -> s * delta cx cy x y) ps

circleOfSources :: ℝ -> [ℂ] -> ℝ -> ℝ -> ℂ
circleOfSources rad cs = pointSources $ zipCs $ pointsOnCircle (length cs) rad
  where
    zipCs = map (\(z, (x, y)) -> (x, y, z)) . zip cs

--const $ FFunc testFunc

-- testFunc :: HelmConf -> Field2D ℂ
-- testFunc (HelmConf (HPointSources ps) _)   = FFunc $ pointSources ps
-- testFunc (HelmConf (HCircleSources r p) _) = FFunc $ circleOfSources r p
  -- where
  --   pol m p = mkPolar m ((pi/180)*p)
-- testFunc = pointSources [ (0.48, 0.48, 0.5 :+ 0)
--                         , (0.52, 0.52, (-0.5) :+ 0)
--                         ]

-- 0.1 [ pol 1 0
--                                , pol 1 60
--                                , pol 1 120
--                                , pol 1 120
--                                , pol 1 60
--                                , pol 1 0
--                                ]

normm :: Matrix ℂ
--im, jm :: Matrix ℂ
--im = sample2DM (\x _ -> x :+ 0) [0, 1 .. size - 1] [0, 1 .. size - 1]
--jm = sample2DM (\_ y -> y :+ 0) [0, 1 .. size - 1] [0, 1 .. size - 1]
normm = sample2DM (\x y -> norm x y :+ 0) [0, 1 .. size - 1] [0, 1 .. size - 1]

eps :: ℝ
eps = 5

delta :: ℝ -> ℝ -> ℝ -> ℝ -> ℂ
delta ((*size) -> cx) ((*size) -> cy) ((*size) -> x) ((*size) -> y)
  | norm (center x - cx) (center y - cy) < size * eps    = 1
  | otherwise                                                       = 0

rendCmp_ :: ℂ -> (ℝ, ℝ, ℝ)
rendCmp_ (polar -> (m, p)) = uncurryRGB (\x y z -> (x, y, z)) $ hsl a 1 m
  where
--    a = if p < 0 then 90 else 0
    a = (p + pi) * (180 / pi)

rendCmp :: Matrix ℂ -> (Matrix ℝ, Matrix ℝ, Matrix ℝ)
rendCmp = matrify . unzip3 . map (unzip3 . map rendCmp_) . toLists
  where
    matrify (x, y, z) = (fromLists x, fromLists y, fromLists z)

rotate :: ℝ -> ℂ -> ℂ
rotate r = uncurry mkPolar . (\(m, p) -> (m, p + r)) . polar

visual :: (ℝ, ℝ, ℝ, ℝ) -> ℂ -> ℂ
visual (xa, xb, ya, yb) (polar -> (m, p))
  = mkPolar (((yb - ya) * (m - xa)/(xb - xa)) + ya) p

visualM :: Matrix ℂ -> Matrix ℂ
visualM cs@(fromComplex -> (r, i)) = mapMatrix (visual (small, big, 0, 1)) cs
  where
    !mags  = flatten $ sqrt $ (r * r) + (i * i)
    !small = minElement mags
    !big   = maxElement mags

updateState :: State a -> State a
--updateState = id
updateState s@State { time = t } = s { time = t + 0.01, concrete = False }

data RenderMode = RendAmp | RendSrc | RendDsp | RendAmpMag | RendAmpPhs

renderState :: RenderMode -> State a -> RGBTrips
renderState r = case r of
  RendAmp    -> rend . concAmp
  RendAmpMag -> rend . concAmpMag
  RendAmpPhs -> rend . concAmpPhs
  RendSrc    -> rend . concSrc
  RendDsp    -> rend . concDsp
  where
  wrap f = fromFieldC . f . makeConcrete
  rend =  rendCmp . visualM
  concAmp    = wrap amplitude
  concAmpMag = mag . fromComplex . concAmp
  concAmpPhs = phs . fromComplex . concAmp
  concSrc    = wrap sources
  concDsp    = wrap dispersion
  mag (x, y) = toComplex $ ((x*x) + (y*y), y - y)
  phs (x, y) = toComplex $ ((y `arctan2` x), y - y)

changeRenderMode :: RenderMode -> Conf (State a) -> Conf (State a)
changeRenderMode r c = c { render = renderState r }

getConf :: Conf (State HelmConf) -> HelmConf
getConf (Conf { state = (State { conf = cn })}) = cn

changeConf :: (HelmConf -> HelmConf) -> Conf (State HelmConf) -> Conf (State HelmConf)
changeConf h c@(Conf { state = s@(State { configure = f, conf = cn })})
  = c { state = configure s' s' }
  where
    s' = s { conf = h cn }

relPhase (polar -> (_, p1)) (polar -> (_, p2)) = abs $ p1 - p2

--getPhase c@(HelmConf { helmsrc = HCircleSources _ [c1, c2] }) = Just $ relPhase c1 c2
--getPhase _ = Nothing
getPhase c@(HelmConf { helmsrc = HPointSources [_, (x, _, _)] }) = Just x
getPhase _ = Nothing

changePhase :: Double -> Conf (State HelmConf) -> Conf (State HelmConf)
changePhase p = changeConf shift
  where
    shift c@(getPhase -> m) = case m of
                               Just ph -> c { helmsrc = testSource $ ph + p }
                               Nothing -> c

keyState :: KMState -> Conf (State HelmConf) -> IO (Conf (State HelmConf))
keyState ""  c = return c
keyState "d" c = return $ changeRenderMode RendDsp c
keyState "s" c = return $ changeRenderMode RendSrc c
keyState "a" c = return $ changeRenderMode RendAmp c
keyState "m" c = return $ changeRenderMode RendAmpMag c
keyState "p" c = return $ changeRenderMode RendAmpPhs c
keyState "w" c = do
  let p = changePhase 0.5 c
  putStrLn $ "Shifted by 0.5 to " ++ show (getPhase $ getConf c)
  return p
keyState "q" c = do
  let p = changePhase (-0.5) c
  putStrLn $ "Shifted by -0.5 to " ++ show (getPhase $ getConf c)
  return p
keyState x   c = return $ trace x c

pol m p = mkPolar m (p * pi / 180)

testSource :: Double -> HSource
testSource p = HPointSources [ (size/2, size/2, (-1) :+ 0)
                             , (p, p, 1 :+ 0)
                             ]
-- testSource p = HSum [ HCircleSources 0.2  [ pol 1 0
--                                           , pol 1 0
--                                           , pol 1 0
--                                           , pol 1 0
--                                           ]
--                     , HCircleSources 0.15 [ pol 1 0
--                                           , pol 1 90
--                                           , pol 1 180
--                                           , pol 1 270
--                                           ]
--                     ]

dspfun :: ℝ -> ℝ -> ℂ
dspfun x y
  | outbox (center x) (center y) = 1.2 :+ 0
  | otherwise                    = 1 :+ 0
  where
    test = 1
    outbox a b
      | a < (-32) = True
      | a > 32    = True
      | b < (-32) = True
      | b > 32    = True
      | otherwise = False

testConf :: Double -> HelmConf
testConf p = HelmConf { helmsrc = testSource p
                      , helmdsp = HFuncDisp dspfun
                      }

-- testConf = HelmConf { helmsrc = HCircleSources 0.1 [ pol 1 0
--                                                    , pol 1 0
--                                                    , pol 1 0
--                                                    , pol 1 0
--                                                    , pol 1 0
--                                                    , pol 1 0 ]
--                     , helmdsp = HConstDisp 0.001
--                     }
--   where
--     pol m p = mkPolar m (p * pi / 180)

testConfigure st@(State { conf = HelmConf { helmdsp = d, helmsrc = s }})
  = st { sources = hsrcToField s
       , dispersion = hdspToField d
       , concrete = False
       }

initState :: State HelmConf
initState = defState { freq = 3
                     , conf = testConf 0
                     , configure = testConfigure
                     }

config :: Conf (State HelmConf)
config = changeConf id
         Conf { keyBinds = keyState
              , state    = initState
              , evolve   = updateState
              , render   = renderState RendAmp
              , canvas   = (size, size)
              }

-- | TODO
main :: IO ()
--main = optimize
main = runDisp config


{-# LANGUAGE ConstraintKinds #-}
{-# LANGUAGE FlexibleInstances #-}

module Math.Matrix
       ( DimVector (..)
       , DimMatrix (..)
       , GPConstraint
       , withMat
       , toMatrix
       , mulM
       , transposeM
       , mapMM
       , mapDiagonalM
       , invSM
       , substractMeanM
       , (+^^)
       , (-^^)
       , (*^^)
       , cholM
       , foldAllSM
       , sumAllSM
       , trace2SM
       , toDimMatrix
       , detSM
       , toMatrixM'
       , zipWithDim
       , zipWithArray
       , identM
       , delayMatrix
       , linearSolveM
       , vectorLength
       , randomMatrixD
       , matrixRowsNum
       , matrixColsNum
       , pinvSM
       , mulPM
       , eigSHM
       )
       where

import Universum hiding (Vector, transpose, map, natVal,
                         zipWith)

import           GHC.TypeLits hiding (someNatVal)

import           Data.Array.Repa
import           Data.Random.Normal (normals)
import           Data.Vector.Unboxed.Base (Unbox)
import           Data.Vinyl.TypeLevel (AllConstrained)
--import           Numeric.Dimensions
import           Numeric.LinearAlgebra.Repa hiding (Matrix, Vector)
import           System.Random (Random, RandomGen)

import PCA.Types
import PCA.Util hiding (randomMatrixD, zipWithArray)

newtype DimVector r (n :: Nat) a
  = DimVector { runVector :: Vector r a }
  deriving Eq

newtype DimMatrix r (y :: Nat) (x :: Nat) a
  = DimMatrix { getInternal :: Matrix r a}

instance Functor (DimVector D n) where
  fmap f (DimVector vector) = DimVector (smap f vector)

withMat
  :: Matrix D a
  -> (forall x y. (KnownNat x, KnownNat y) => DimMatrix D x y a -> k)
  -> k
withMat m f =
  let (Z :. x :. y) = extent m
  in
  case someNatVal (fromIntegral x) of
    SomeNat (Proxy :: Proxy m) -> case someNatVal (fromIntegral y) of
      SomeNat (Proxy :: Proxy n) -> f (DimMatrix @_ @m @n m)

mulM
  :: forall y1 x1 y2 x2 a.
  ( AllConstrained KnownNat [x1, x2, y1, y2]
  , Numeric a
  , x1 ~ y2
  )
  => DimMatrix D y1 x1 a
  -> DimMatrix D y2 x2 a
  -> DimMatrix D y1 x2 a
mulM (DimMatrix m1) (DimMatrix m2) = DimMatrix $ delay $ m1 `mulS` m2

transposeM
  :: (KnownNat y, KnownNat x)
  => DimMatrix D y x a
  -> DimMatrix D x y a
transposeM (DimMatrix m) = DimMatrix $ transpose m

mapMM
  ::
  ( KnownNat y
  , KnownNat x
  , Unbox a
  , Unbox b
  )
  => (a -> b)
  -> DimMatrix D y x a
  -> DimMatrix D y x b
mapMM f (DimMatrix m) =  DimMatrix $ map f m

mapDiagonalM
  ::
  ( KnownNat y
  , KnownNat x
  , Unbox a
  )
  => (a -> a)
  -> DimMatrix D y x a
  -> DimMatrix D y x a
mapDiagonalM f (DimMatrix m) = DimMatrix $ mapDiagonal f m

invSM
  ::
  ( KnownNat y
  , KnownNat x
  , Field a
  , Numeric a
  , y ~ x
  )
  => DimMatrix D y x a
  -> DimMatrix D y x a
invSM (DimMatrix m) = DimMatrix $ delay $ invS m

substractMeanM
  ::
  ( KnownNat y
  , KnownNat x
  )
  => DimMatrix D y x Double
  -> DimMatrix D y x Double
substractMeanM (DimMatrix m) = DimMatrix $ substractMean m

infixl 6 +^^, -^^
infixl 7 *^^

(+^^)
  :: forall y1 x1 y2 x2 a.
  ( AllConstrained KnownNat [x1, x2, y1, y2]
  , x1 ~ x2
  , y1 ~ y2
  , Num a
  )
  => DimMatrix D y1 x1 a
  -> DimMatrix D y2 x2 a
  -> DimMatrix D y2 x2 a
(+^^) (DimMatrix m1) (DimMatrix m2) = DimMatrix $ m1 +^ m2

(-^^)
  :: forall y1 x1 y2 x2 a.
  ( AllConstrained KnownNat [x1, x2, y1, y2]
  , x1 ~ x2
  , y1 ~ y2
  , Num a
  )
  => DimMatrix D y1 x1 a
  -> DimMatrix D y2 x2 a
  -> DimMatrix D y2 x2 a
(-^^) (DimMatrix m1) (DimMatrix m2) = DimMatrix $ m1 -^ m2

(*^^)
  :: forall y1 x1 y2 x2 a.
  ( AllConstrained KnownNat [x1, x2, y1, y2]
  , x1 ~ x2
  , y1 ~ y2
  , Num a
  )
  => DimMatrix D y1 x1 a
  -> DimMatrix D y2 x2 a
  -> DimMatrix D y2 x2 a
(*^^) (DimMatrix m1) (DimMatrix m2) = DimMatrix $ m1 *^ m2

cholM
  ::
  ( KnownNat y
  , KnownNat x
  , Field a
  , y ~ x
  )
  => DimMatrix D y x a
  -> DimMatrix D y x a
cholM (DimMatrix m) = DimMatrix $ delay $ chol $ trustSym $ computeS m

sumAllSM
  ::
  ( KnownNat y
  , KnownNat x
  , Num a)
  => DimMatrix D y x a
  -> a
sumAllSM (DimMatrix m) = sumAllS m

foldAllSM
  :: (KnownNat y, KnownNat x)
  => (Double -> Double -> Double)
  -> Double
  -> DimMatrix D y x Double
  -> Double
foldAllSM f initValue (DimMatrix m) = foldAllS f initValue m

detSM
  :: (KnownNat y, KnownNat x)
  => DimMatrix D y x Double
  -> Double
detSM (DimMatrix m) = detS m

trace2SM
  :: (KnownNat x, KnownNat y)
  => DimMatrix D x y Double
  -> Double
trace2SM (DimMatrix m) = trace2S $ computeS m

toDimMatrix
  ::
  ( Source r a
  , KnownNat m
  , KnownNat n
  )
  => DimVector r m a
  -> Int
  -> DimMatrix D m n a
toDimMatrix (DimVector arr) desiredSize =
  DimMatrix (toMatrix arr desiredSize)

toMatrixM'
  :: (Source r a, KnownNat n, KnownNat m)
  => DimVector r n a
  -> DimMatrix D n m a
toMatrixM' (DimVector arr) =
  DimMatrix $ fromFunction (Z :. dimension :. 1) generator
  where
    dimension = size . extent $ arr
    generator (Z :. rows :. _) = linearIndex arr rows

zipWithDim
  ::
  ( Source r1 a
  , Source r2 b
  , KnownNat m
  , KnownNat n
  )
  => (a -> b -> c)
  -> DimMatrix r1 m n a
  -> DimMatrix r2 m n b
  -> DimMatrix D m n c
zipWithDim f (DimMatrix mat1) (DimMatrix mat2) =
  DimMatrix $ zipWith f mat1 mat2

zipWithArray
  ::
  ( KnownNat n
  , KnownNat m
  )
  => (a -> b -> c)
  -> DimVector D n a
  -> DimMatrix D n m b
  -> DimMatrix D n m c
zipWithArray f (DimVector array1) (DimMatrix array2) =
  DimMatrix $ zipWith f (toMatrix array1 n) array2
  where
    (Z :. n :. _) = extent array2

identM
  :: forall m n a.
  ( KnownNat m
  , KnownNat n
  , GPConstraint a
  , m ~ n
  )
  => DimMatrix D m n a
identM =
  let dim = fromEnum $
            natVal @m @Proxy Proxy in DimMatrix $
            identD dim

delayMatrix
  ::
  ( KnownNat m
  , KnownNat n
  , Source r a
  )
  => DimMatrix r m n a
  -> DimMatrix D m n a
delayMatrix (DimMatrix matrix) = DimMatrix (delay matrix)

linearSolveM
  ::
  ( Field a
  , AllConstrained KnownNat '[m, n]
  )
  => DimMatrix D m m a
  -> DimMatrix D m n a
  -> Maybe (DimMatrix D m n a)
linearSolveM (DimMatrix mat1) (DimMatrix mat2) =
  case linearSolveS mat1 mat2 of
    Nothing  -> Nothing
    Just sol -> Just . delayMatrix . DimMatrix $ sol

vectorLength
  :: forall r m a. KnownNat m
  => DimVector r m a
  -> Int
vectorLength _ = fromEnum $ natVal (Proxy @m)

matrixRowsNum
  :: forall r m n a. KnownNat m
  => DimMatrix r m n a
  -> Int
matrixRowsNum _ = fromEnum $ natVal (Proxy @m)

matrixColsNum
  :: forall r m n a. KnownNat n
  => DimMatrix r m n a
  -> Int
matrixColsNum _ = fromEnum $ natVal (Proxy @n)

randomMatrixD
  :: forall a g m n.
  ( RandomGen g
  , Random a
  , Unbox a
  , Floating a
  , KnownNat m
  , KnownNat n
  )
  => g
  -> (Int, Int)
  -> DimMatrix D m n a
randomMatrixD gen (rows, cols) =
  let randomList = take (rows * cols) (normals gen) in
    DimMatrix . delay $ fromListUnboxed (Z :. rows :. cols) randomList

type GPConstraint a =
  ( Field a
  , Random a
  , Unbox a
  , Floating a
  , Eq a
  )
{-
hermToMatrixM
  :: Herm a
  -> DimMatrix D m n a
hermToMatrixM = undefined
-}
pinvSM
  ::
  ( KnownNat y
  , KnownNat x
  , Field a
  , Numeric a
  , y ~ x
  )
  => DimMatrix D y x a
  -> DimMatrix D y x a
pinvSM (DimMatrix m) = DimMatrix . delay $ pinvS m

mulPM
  :: forall y1 x1 y2 x2 a.
  ( AllConstrained KnownNat [x1, x2, y1, y2]
  , Numeric a
  , x1 ~ y2
  )
  => DimMatrix D y1 x1 a
  -> DimMatrix D y2 x2 a
  -> DimMatrix D y1 x2 a
mulPM (DimMatrix m) (DimMatrix n) =
  DimMatrix . delay . runIdentity $ m `mulP` n

eigSHM
  :: (Field a, Numeric a)
  => Herm a
  -> (DimVector D m Double, DimMatrix D m m a)
eigSHM hermM =
  (DimVector $ delay $ fst out, DimMatrix $ delay $ snd out)
  where
    out = eigSH hermM


{-# LANGUAGE BangPatterns     #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE Strict           #-}
module STC.InitialDistribution where

import           Array.UnboxedArray       as AU
import           Data.Array.Repa          as R
import           Data.Complex
import           Data.List                as L
import           Data.Vector.Generic      as VG
import           Data.Vector.Storable     as VS
import           DFT.Plan
import           FourierPinwheel
import           Pinwheel.FourierSeries2D
import           STC.Convolution
import           STC.DFTArray
import           STC.Point
import           STC.Utils
import           Utils.Distribution
import           Utils.List
import           Utils.Parallel
import Filter.Utils

{-# INLINE computeInitialDistribution #-}
computeInitialDistribution ::
     Int -> Int -> [Double] -> [Double] -> Double -> [Point] -> DFTArray
computeInitialDistribution rows cols phiFreqs rhoFreqs halfLogPeriod points =
  let (!minR, !maxR) = computeRange rows
      (!minC, !maxC) = computeRange cols
  in if L.any (\(Point _ _ _ s) -> s <= 0) points
       then error $ "computeInitialDistribution: initial scale <= 0."
       else DFTArray rows cols phiFreqs rhoFreqs .
            parMap
              rdeepseq
              (\(!rFreq, !thetaFreq) ->
                 VG.convert .
                 toUnboxedVector .
                 AU.accum (+) 0 ((minC, minR), (maxC, maxR)) .
                 L.map
                   (\(Point x y theta scale) ->
                      ( (round x, round y)
                      , cis $
                        (log scale) * (-rFreq) - (thetaFreq * theta * pi / 180))) $
                 points) $
            (,) <$> rhoFreqs <*> phiFreqs

{-# INLINE computeInitialDistribution' #-}
computeInitialDistribution' ::
     Int -> Int -> [Double] -> [Double] -> Double -> [Point] -> DFTArray
computeInitialDistribution' rows cols phiFreqs rhoFreqs  halfLogPeriod points =
  let (!minR, !maxR) = computeRange rows
      (!minC, !maxC) = computeRange cols
  in if L.any (\(Point _ _ _ s) -> s <= 0) points
       then error $ "computeInitialDistribution: initial scale <= 0."
       else DFTArray rows cols phiFreqs rhoFreqs .
            parMap
              rdeepseq
              (\(!thetaFreq) ->
                 VG.convert .
                 toUnboxedVector .
                 AU.accum (+) 0 ((minC, minR), (maxC, maxR)) .
                 L.map
                   (\(Point x y theta _) ->
                      ((round x, round y), cis $ -(thetaFreq * theta * pi / 180))) $
                 points) $
            phiFreqs

{-# INLINE computeInitialDistributionPowerMethod #-}
computeInitialDistributionPowerMethod ::
     Int -> Int -> [Double] -> [Double] -> Double -> [Point] -> DFTArray
computeInitialDistributionPowerMethod rows cols thetaFreqs rFreqs halfLogPeriod points =
  let (!minR, !maxR) = computeRange rows
      (!minC, !maxC) = computeRange cols
      !zeroVec = VG.replicate (rows * cols) 0
  in if L.any (\(Point _ _ _ s) -> s <= 0) points
       then error $ "computeInitialDistributionPowerMethod: initial scale <= 0."
       else DFTArray rows cols thetaFreqs rFreqs .
            L.map
              (\(!rFreq, !thetaFreq) ->
                 if rFreq == 0 && thetaFreq == 0
                   then VG.convert .
                        toUnboxedVector .
                        AU.accum (+) 0 ((minC, minR), (maxC, maxR)) .
                        L.map (\(Point x y theta scale) -> ((round x, round y), 1)) $
                        points
                   else zeroVec) $
            (,) <$> rFreqs <*> thetaFreqs

{-# INLINE computeInitialDistributionPowerMethod' #-}
computeInitialDistributionPowerMethod' ::
     Int -> Int -> [Double] -> [Double] -> [Point] -> DFTArray
computeInitialDistributionPowerMethod' rows cols phiFreqs rhoFreqs points =
  let (!minR, !maxR) = computeRange rows
      (!minC, !maxC) = computeRange cols
      !zeroVec = VG.replicate (rows * cols) 0
  in if L.any (\(Point _ _ _ s) -> s <= 0) points
       then error $ "computeInitialDistributionPowerMethod: initial scale <= 0."
       else DFTArray rows cols phiFreqs rhoFreqs .
            L.map
              (\(!thetaFreq) ->
                 if thetaFreq == 0
                   then VG.convert .
                        toUnboxedVector .
                        AU.accum (+) 0 ((minC, minR), (maxC, maxR)) .
                        L.map (\(Point x y theta scale) -> ((round x, round y), 1)) $
                        points
                   else zeroVec) $
            phiFreqs

{-# INLINE computeInitialDistributionPowerMethodSparse #-}
computeInitialDistributionPowerMethodSparse ::
     [Double] -> [Double] -> [Point] -> [R.Array U DIM2 (Complex Double)]
computeInitialDistributionPowerMethodSparse thetaFreqs rFreqs points =
  let !numThetaFreq = L.length thetaFreqs
      !numRFreq = L.length rFreqs
      !thetaCenter = div numThetaFreq 2
      !rCenter = div numRFreq 2
      !initArr =
        computeUnboxedS . fromFunction (Z :. numRFreq :. numThetaFreq) $ \(Z :. i :. j) ->
          if i == rCenter && j == thetaCenter
            then 1
            else 0
  in L.replicate (L.length points) initArr

{-# INLINE computeInitialDistributionPowerMethodSparse' #-}
computeInitialDistributionPowerMethodSparse' ::
     [Double] -> [Double] -> [Point] -> [R.Array U DIM1 (Complex Double)]
computeInitialDistributionPowerMethodSparse' thetaFreqs rFreqs points =
  let !numThetaFreq = L.length thetaFreqs
      !thetaCenter = div numThetaFreq 2
      !initArr =
        computeUnboxedS . fromFunction (Z :. numThetaFreq) $ \(Z :. j) ->
          if j == thetaCenter
            then 1
            else 0
      -- !initArr =
      --   computeUnboxedS .
      --   R.traverse (fromListUnboxed (Z :. numThetaFreq) thetaFreqs) id $ \f (Z :. thetaFreq) ->
      --     cis $ -(fromIntegral thetaFreq * 36 * pi / 180)
  in L.replicate (L.length points) initArr

computeInitialDistributionFull' ::
     Int -> Double -> Int -> Int -> [Point] -> DFTArray
computeInitialDistributionFull' numR2Freq period phiFreq rhoFreq points =
  let r2Freqs = L.map fromIntegral . getListFromNumber $ numR2Freq
      phiFreqs = L.map fromIntegral [-phiFreq .. phiFreq]
      rhoFreqs = L.map fromIntegral [-rhoFreq .. rhoFreq]
  in DFTArray
       numR2Freq
       numR2Freq
       phiFreqs
       rhoFreqs
       [ VG.fromList
         [ L.foldl'
           (\b (Point x y theta _) ->
              b +
              cis
                (-(angularFreq * theta * pi / 180 +
                   (freqX * x + freqY * y) * 2 * pi / period)))
           0
           points
         | freqY <- r2Freqs
         , freqX <- r2Freqs
         ]
       | angularFreq <- phiFreqs
       ]

computeInitialDistributionPowerMethodPinwheelBasis' ::
     Int -> Double -> Double -> Int -> Int -> [Point] -> DFTArray
computeInitialDistributionPowerMethodPinwheelBasis' numR2Freq sigma period phiFreq rhoFreq points =
  let r2Freqs = L.map fromIntegral . getListFromNumber $ numR2Freq
      phiFreqs = L.map fromIntegral [-phiFreq .. phiFreq]
      rhoFreqs = L.map fromIntegral [-rhoFreq .. rhoFreq]
      zeroVec = VG.replicate (numR2Freq ^ 2) 0
      envelope =
        VG.convert .
        toUnboxed . computeS . fromFunction (Z :. numR2Freq :. numR2Freq) $ \(Z :. i :. j) ->
          (exp $
           (fromIntegral $ (i - div numR2Freq 2) ^ 2 + (j - div numR2Freq 2) ^ 2) *
           (sigma ^ 2) /
           (-2)) *
          (sigma ^ 2) /
          (2 * pi) :+
          0
  in DFTArray
       numR2Freq
       numR2Freq
       phiFreqs
       rhoFreqs
       [ if angularFreq == 0
         then VG.zipWith (*) envelope . VG.fromList $
              [ L.foldl'
                (\b (Point x y _ _) ->
                   b + cis (-(freqX * x + freqY * y) * 2 * pi / period))
                0
                points
              | freqY <- r2Freqs
              , freqX <- r2Freqs
              ]
         else zeroVec
       | angularFreq <- phiFreqs
       ]

computeInitialDistributionPowerMethodPinwheelBasis ::
     Int -> Double -> Double -> Int -> Int -> [Point] -> DFTArray
computeInitialDistributionPowerMethodPinwheelBasis numR2Freq sigma period phiFreq rhoFreq points =
  let r2Freqs = L.map fromIntegral . getListFromNumber $ numR2Freq
      phiFreqs = L.map fromIntegral [-phiFreq .. phiFreq]
      rhoFreqs = L.map fromIntegral [-rhoFreq .. rhoFreq]
      zeroVec = VG.replicate (numR2Freq ^ 2) 0
      envelope =
        VG.convert .
        toUnboxed . computeS . fromFunction (Z :. numR2Freq :. numR2Freq) $ \(Z :. i :. j) ->
          (exp $
           (fromIntegral $ (i - div numR2Freq 2) ^ 2 + (j - div numR2Freq 2) ^ 2) *
           (sigma ^ 2) /
           (-2)) *
          (sigma ^ 2) /
          (2 * pi) :+
          0
      -- envelope =
      --   VG.convert . toUnboxed . computeUnboxedS $
      --   analyticalFourierCoefficients2 numR2Freq 1 0 0 sigma period (period * sqrt 2)
  in DFTArray
       numR2Freq
       numR2Freq
       phiFreqs
       rhoFreqs
       [ if angularFreq == 0 && radialFreq == 0
         then VG.zipWith (*) envelope .
              VG.fromList $
              [ L.foldl'
                (\b (Point x y _ _) ->
                   b + cis (-(freqX * x + freqY * y) * 2 * pi / period))
                0
                points
              | freqY <- r2Freqs
              , freqX <- r2Freqs
              ]
         else zeroVec
       | radialFreq <- rhoFreqs
       , angularFreq <- phiFreqs
       ]

computeInitialDistributionFourierPinwheel ::
     (VG.Vector vector (Complex Double), NFData (vector (Complex Double)))
  => Int
  -> Double
  -> Double
  -> Int
  -> Int
  -> Int
  -> Int
  -> [Point]
  -> FPArray (vector (Complex Double))
computeInitialDistributionFourierPinwheel numR2Freqs period periodEnv phiFreq rhoFreq thetaFreq rFreq points =
  let r2Freqs = L.map fromIntegral . getListFromNumber $ numR2Freqs
      thetaFreqs = L.map fromIntegral [-thetaFreq .. thetaFreq]
      rFreqs = L.map fromIntegral [-rFreq .. rFreq]
   in FPArray
        numR2Freqs
        numR2Freqs
        (2 * rFreq + 1)
        (2 * thetaFreq + 1)
        (2 * rhoFreq + 1)
        (2 * phiFreq + 1) .
      L.map
        (\radialFreq ->
           VG.concat .
           parMap
             rdeepseq
             (\angularFreq ->
                VG.convert . toUnboxed . computeS . makeFilter2D . fromListUnboxed (Z :. numR2Freqs :. numR2Freqs) $
                   [ L.foldl'
                     (\b (Point x y theta scale) ->
                        b +
                        cis
                          (-(angularFreq * theta * pi / 180 +
                             (freqX * x + freqY * y) * 2 * pi / period +
                             2 * pi / log periodEnv * log scale * radialFreq)))
                     0
                     points
                   | freqY <- r2Freqs
                   , freqX <- r2Freqs
                   ]) $
           thetaFreqs) $
      rFreqs

computeInitialDistributionPowerMethodFourierPinwheel ::
     Int
  -> Int
  -> Int
  -> Int
  -> Int
  -> VS.Vector (Complex Double)
  -> FPArray (VS.Vector (Complex Double))
computeInitialDistributionPowerMethodFourierPinwheel numR2Freqs phiFreq rhoFreq thetaFreq rFreq bias =
  let thetaFreqs = L.map fromIntegral [-thetaFreq .. thetaFreq]
      rFreqs = L.map fromIntegral [-rFreq .. rFreq]
      numThetaFreq = 2 * thetaFreq + 1
      zeroVec = VG.replicate (numThetaFreq * numR2Freqs ^ 2) 0
   in FPArray
        numR2Freqs
        numR2Freqs
        (2 * rFreq + 1)
        numThetaFreq
        (2 * rhoFreq + 1)
        (2 * phiFreq + 1) .
      L.map
        (\radialFreq ->
           if radialFreq == 0
             then bias
             else zeroVec) $
      rFreqs

computeInitialDistributionPowerMethodFourierPinwheelFull ::
     Int
  -> Int
  -> Int
  -> Int
  -> Int
  -> VS.Vector (Complex Double)
  -> FPArray (VS.Vector (Complex Double))
computeInitialDistributionPowerMethodFourierPinwheelFull numR2Freqs phiFreq rhoFreq thetaFreq rFreq bias =
  let numThetaFreq = 2 * thetaFreq + 1
      numRFreq = 2 * rFreq + 1
   in FPArray
        numR2Freqs
        numR2Freqs
        (2 * rFreq + 1)
        numThetaFreq
        (2 * rhoFreq + 1)
        (2 * phiFreq + 1) .
      parMap
        rdeepseq
        (\radialFreq ->
           VG.convert .
           toUnboxed .
           computeS .
           R.slice
             (fromUnboxed
                (Z :. numRFreq :. numThetaFreq :. numR2Freqs :. numR2Freqs) .
              VG.convert $
              bias) $
           (Z :. radialFreq :. All :. All :. All)) $
      [0 .. 2 * rFreq]


{-# LANGUAGE CPP                   #-}
{-# LANGUAGE DataKinds             #-}
{-# LANGUAGE ScopedTypeVariables   #-}
{-# LANGUAGE RecordWildCards       #-}
{-# LANGUAGE GADTs                 #-}
{-# LANGUAGE TypeOperators         #-}
{-# LANGUAGE TypeFamilies          #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE FlexibleInstances     #-}
{-# LANGUAGE FlexibleContexts      #-}
{-# LANGUAGE UndecidableInstances  #-}
{-|
Module      : Grenade.Layers.Convolution
Description : Convolution layer
Copyright   : (c) Huw Campbell, 2016-2017
License     : BSD2
Stability   : experimental

This module provides the Convolution layer, which is critical in many computer vision tasks.

-}
module Grenade.Layers.Convolution (
    Convolution (..)
  , Convolution' (..)
  , randomConvolution
  ) where

import           Control.Monad.Random hiding ( fromList )
import           Data.Maybe
import           Data.Proxy
import           Data.Serialize
import           Data.Singletons.TypeLits

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

import           Numeric.LinearAlgebra hiding ( uniformSample, konst )
import qualified Numeric.LinearAlgebra as LA
import           Numeric.LinearAlgebra.Static hiding ((|||), build, toRows)

import           Grenade.Core
import           Grenade.Layers.Internal.Convolution
import           Grenade.Layers.Internal.Update

-- | A convolution layer for a neural network.
--   This uses the im2col convolution trick popularised by Caffe, which essentially turns the
--   many, many, many, many loop convolution into a single matrix multiplication.
--
--   The convolution layer takes all of the kernels for the convolution, which are flattened
--   and then put into columns in the matrix.
--
--   The kernel size dictates which input and output sizes will "fit". Fitting the equation:
--   `out = (in - kernel) / stride + 1` for both dimensions.
--
--   One probably shouldn't build their own layer, but rather use the randomConvolution function.
data Convolution :: Nat -- Number of channels, for the first layer this could be RGB for instance.
                 -> Nat -- Number of filters, this is the number of channels output by the layer.
                 -> Nat -- The number of rows in the kernel filter
                 -> Nat -- The number of column in the kernel filter
                 -> Nat -- The row stride of the convolution filter
                 -> Nat -- The columns stride of the convolution filter
                 -> Type where
  Convolution :: ( KnownNat channels
                 , KnownNat filters
                 , KnownNat kernelRows
                 , KnownNat kernelColumns
                 , KnownNat strideRows
                 , KnownNat strideColumns
                 , KnownNat kernelFlattened
                 , kernelFlattened ~ (kernelRows * kernelColumns * channels))
              => !(L kernelFlattened filters) -- The kernel filter weights
              -> !(L kernelFlattened filters) -- The last kernel update (or momentum)
              -> Convolution channels filters kernelRows kernelColumns strideRows strideColumns

data Convolution' :: Nat -- Number of channels, for the first layer this could be RGB for instance.
                  -> Nat -- Number of filters, this is the number of channels output by the layer.
                  -> Nat -- The number of rows in the kernel filter
                  -> Nat -- The number of column in the kernel filter
                  -> Nat -- The row stride of the convolution filter
                  -> Nat -- The columns stride of the convolution filter
                  -> Type where
  Convolution' :: ( KnownNat channels
                  , KnownNat filters
                  , KnownNat kernelRows
                  , KnownNat kernelColumns
                  , KnownNat strideRows
                  , KnownNat strideColumns
                  , KnownNat kernelFlattened
                  , kernelFlattened ~ (kernelRows * kernelColumns * channels))
               => !(L kernelFlattened filters) -- The kernel filter gradient
               -> Convolution' channels filters kernelRows kernelColumns strideRows strideColumns

instance Show (Convolution c f k k' s s') where
  show (Convolution a _) = renderConv a
    where
      renderConv mm =
        let m  = extract mm
            ky = fromIntegral $ natVal (Proxy :: Proxy k)
            rs = LA.toColumns m
            ms = map (take ky) $ toLists . reshape ky <$> rs

            render n'  | n' <= 0.2  = ' '
                       | n' <= 0.4  = '.'
                       | n' <= 0.6  = '-'
                       | n' <= 0.8  = '='
                       | otherwise =  '#'

            px = (fmap . fmap . fmap) render ms
        in unlines $ foldl1 (zipWith (\a' b' -> a' ++ "   |   " ++ b')) $ px

randomConvolution :: ( MonadRandom m
                     , KnownNat channels
                     , KnownNat filters
                     , KnownNat kernelRows
                     , KnownNat kernelColumns
                     , KnownNat strideRows
                     , KnownNat strideColumns
                     , KnownNat kernelFlattened
                     , kernelFlattened ~ (kernelRows * kernelColumns * channels))
                  => m (Convolution channels filters kernelRows kernelColumns strideRows strideColumns)
randomConvolution = do
    s     <- getRandom
    let wN = uniformSample s (-1) 1
        mm = konst 0
    return $ Convolution wN mm

instance ( KnownNat channels
         , KnownNat filters
         , KnownNat kernelRows
         , KnownNat kernelColumns
         , KnownNat strideRows
         , KnownNat strideColumns
         , KnownNat (kernelRows * kernelColumns * channels)
         ) => UpdateLayer (Convolution channels filters kernelRows kernelColumns strideRows strideColumns) where
  type Gradient (Convolution channels filters kernelRows kernelColumns strideRows strideColumns) = (Convolution' channels filters kernelRows kernelColumns strideRows strideColumns)
  runUpdate LearningParameters {..} (Convolution oldKernel oldMomentum) (Convolution' kernelGradient) =
    let (newKernel, newMomentum) = descendMatrix learningRate learningMomentum learningRegulariser oldKernel kernelGradient oldMomentum
    in Convolution newKernel newMomentum

  createRandom = randomConvolution

instance ( KnownNat channels
         , KnownNat filters
         , KnownNat kernelRows
         , KnownNat kernelColumns
         , KnownNat strideRows
         , KnownNat strideColumns
         , KnownNat (kernelRows * kernelColumns * channels)
         ) => Serialize (Convolution channels filters kernelRows kernelColumns strideRows strideColumns) where
  put (Convolution w _) = putListOf put . toList . flatten . extract $ w
  get = do
      let f  = fromIntegral $ natVal (Proxy :: Proxy filters)
      wN    <- maybe (fail "Vector of incorrect size") return . create . reshape f . LA.fromList =<< getListOf get
      let mm = konst 0
      return $ Convolution wN mm

-- | A three dimensional image (or 2d with many channels) can have
--   an appropriately sized convolution filter run across it.
instance ( KnownNat kernelRows
         , KnownNat kernelCols
         , KnownNat filters
         , KnownNat strideRows
         , KnownNat strideCols
         , KnownNat inputRows
         , KnownNat inputCols
         , KnownNat outputRows
         , KnownNat outputCols
         , KnownNat channels
         , ((outputRows - 1) * strideRows) ~ (inputRows - kernelRows)
         , ((outputCols - 1) * strideCols) ~ (inputCols - kernelCols)
         , KnownNat (kernelRows * kernelCols * channels)
         , KnownNat (outputRows * filters)
         ) => Layer (Convolution channels filters kernelRows kernelCols strideRows strideCols) ('D3 inputRows inputCols channels) ('D3 outputRows outputCols filters) where

  type Tape (Convolution channels filters kernelRows kernelCols strideRows strideCols) ('D3 inputRows inputCols channels) ('D3 outputRows outputCols filters) = S ('D3 inputRows inputCols channels)

  runForwards (Convolution kernel _) (S3D input) =
    let ex = extract input
        ek = extract kernel
        ix = fromIntegral $ natVal (Proxy :: Proxy inputRows)
        iy = fromIntegral $ natVal (Proxy :: Proxy inputCols)
        kx = fromIntegral $ natVal (Proxy :: Proxy kernelRows)
        ky = fromIntegral $ natVal (Proxy :: Proxy kernelCols)
        sx = fromIntegral $ natVal (Proxy :: Proxy strideRows)
        sy = fromIntegral $ natVal (Proxy :: Proxy strideCols)
        ox = fromIntegral $ natVal (Proxy :: Proxy outputRows)
        oy = fromIntegral $ natVal (Proxy :: Proxy outputCols)

        c  = vid2col kx ky sx sy ix iy ex
        mt = c LA.<> ek
        r  = col2vid 1 1 1 1 ox oy mt
        rs = fromJust . create $ r
    in  (S3D input, S3D rs)
  runBackwards (Convolution kernel _) (S3D input) (S3D dEdy) =
    let ex = extract input
        ix = fromIntegral $ natVal (Proxy :: Proxy inputRows)
        iy = fromIntegral $ natVal (Proxy :: Proxy inputCols)
        kx = fromIntegral $ natVal (Proxy :: Proxy kernelRows)
        ky = fromIntegral $ natVal (Proxy :: Proxy kernelCols)
        sx = fromIntegral $ natVal (Proxy :: Proxy strideRows)
        sy = fromIntegral $ natVal (Proxy :: Proxy strideCols)
        ox = fromIntegral $ natVal (Proxy :: Proxy outputRows)
        oy = fromIntegral $ natVal (Proxy :: Proxy outputCols)

        c  = vid2col kx ky sx sy ix iy ex

        eo = extract dEdy
        ek = extract kernel

        vs = vid2col 1 1 1 1 ox oy eo

        kN = fromJust . create $ tr c LA.<> vs

        dW = vs LA.<> tr ek

        xW = col2vid kx ky sx sy ix iy dW
    in  (Convolution' kN, S3D . fromJust . create $ xW)

-- | A two dimentional image may have a convolution filter applied to it
instance ( KnownNat kernelRows
         , KnownNat kernelCols
         , KnownNat filters
         , KnownNat strideRows
         , KnownNat strideCols
         , KnownNat inputRows
         , KnownNat inputCols
         , KnownNat outputRows
         , KnownNat outputCols
         , ((outputRows - 1) * strideRows) ~ (inputRows - kernelRows)
         , ((outputCols - 1) * strideCols) ~ (inputCols - kernelCols)
         , KnownNat (kernelRows * kernelCols * 1)
         , KnownNat (outputRows * filters)
         ) => Layer (Convolution 1 filters kernelRows kernelCols strideRows strideCols) ('D2 inputRows inputCols) ('D3 outputRows outputCols filters) where
  type Tape (Convolution 1 filters kernelRows kernelCols strideRows strideCols) ('D2 inputRows inputCols) ('D3 outputRows outputCols filters) = S ('D3 inputRows inputCols 1)
  runForwards c (S2D input) =
    runForwards c (S3D input :: S ('D3 inputRows inputCols 1))

  runBackwards c tape grads =
    case runBackwards c tape grads of
      (c', S3D back :: S ('D3 inputRows inputCols 1)) ->  (c', S2D back)

-- | A two dimensional image may have a convolution filter applied to it producing
--   a two dimensional image if both channels and filters is 1.
instance ( KnownNat kernelRows
         , KnownNat kernelCols
         , KnownNat strideRows
         , KnownNat strideCols
         , KnownNat inputRows
         , KnownNat inputCols
         , KnownNat outputRows
         , KnownNat outputCols
         , ((outputRows - 1) * strideRows) ~ (inputRows - kernelRows)
         , ((outputCols - 1) * strideCols) ~ (inputCols - kernelCols)
         , KnownNat (kernelRows * kernelCols * 1)
         , KnownNat (outputRows * 1)
         ) => Layer (Convolution 1 1 kernelRows kernelCols strideRows strideCols) ('D2 inputRows inputCols) ('D2 outputRows outputCols) where
  type Tape (Convolution 1 1 kernelRows kernelCols strideRows strideCols) ('D2 inputRows inputCols) ('D2 outputRows outputCols) = S ('D3 inputRows inputCols 1)
  runForwards c (S2D input) =
    case runForwards c (S3D input :: S ('D3 inputRows inputCols 1)) of
      (tps, S3D back :: S ('D3 outputRows outputCols 1)) ->  (tps, S2D back)

  runBackwards c tape (S2D grads) =
    case runBackwards c tape (S3D grads :: S ('D3 outputRows outputCols 1)) of
      (c', S3D back :: S ('D3 inputRows inputCols 1)) -> (c', S2D back)

-- | A three dimensional image can produce a 2D image from a convolution with 1 filter
instance ( KnownNat kernelRows
         , KnownNat kernelCols
         , KnownNat strideRows
         , KnownNat strideCols
         , KnownNat inputRows
         , KnownNat inputCols
         , KnownNat outputRows
         , KnownNat outputCols
         , KnownNat channels
         , ((outputRows - 1) * strideRows) ~ (inputRows - kernelRows)
         , ((outputCols - 1) * strideCols) ~ (inputCols - kernelCols)
         , KnownNat (kernelRows * kernelCols * channels)
         , KnownNat (outputRows * 1)
         ) => Layer (Convolution channels 1 kernelRows kernelCols strideRows strideCols) ('D3 inputRows inputCols channels) ('D2 outputRows outputCols) where
  type Tape (Convolution channels 1 kernelRows kernelCols strideRows strideCols) ('D3 inputRows inputCols channels) ('D2 outputRows outputCols) = S ('D3 inputRows inputCols channels)
  runForwards c input =
    case runForwards c input of
      (tps, S3D back :: S ('D3 outputRows outputCols 1)) ->  (tps, S2D back)

  runBackwards c tape (S2D grads) =
    runBackwards c tape (S3D grads :: S ('D3 outputRows outputCols 1))


{-# LANGUAGE RankNTypes #-}
module Main where

import Codec.Picture (PixelRGBA8 (..), writePng)
import Control.Monad
import Data.Array.Accelerate
  ( Acc,
    All (..),
    Any (..),
    Array,
    DIM2,
    Elt,
    Exp,
    Generic,
    Lift (..),
    Matrix,
    Plain,
    Unlift,
    Vector,
    Z (..),
    (:.) (..),
  )
import qualified Data.Array.Accelerate as A
import Data.Array.Accelerate.Data.Complex
import Data.Array.Accelerate.LLVM.PTX
import Data.Array.Accelerate.Representation.Type
import Data.Array.Accelerate.Smart
import Data.Array.Accelerate.Sugar.Elt
import qualified Data.Complex as C
-- import Graphics.Amusing
import Graphics.Rasterific hiding (Vector)
import Graphics.Rasterific.Texture
import Prelude.Unicode
import Text.Printf

type Pendulum a b = (a, b)

zs :: Pendulum a b -> a
zs (zs,  _) = zs

ps :: Pendulum a b -> b
ps (_, ps) = ps

type Pair a = (a, a)

g :: (RealFrac a, A.RealFloat a, Elt a) => Exp a
g = 9.81

origin ::
  (RealFrac a, A.RealFloat a, Elt a) =>
  Exp (Pair (Complex a))
origin = A.constant (0 :+ 0, 0 :+ 0)

zero ::
  (RealFrac a, A.RealFloat a, Elt a) =>
  Exp (Pair a)
zero = A.constant (0, 0)

one ::
  (RealFrac a, A.RealFloat a, Elt a) =>
  Exp (Pair (Complex a))
one = A.constant (1 :+ 0, 1 :+ 0)

calc ::
  forall a.
  (RealFrac a, A.RealFloat a, Elt a) =>
  Exp a ->
  Exp a ->
  Exp a ->
  Exp (Pendulum (Pair (Complex a)) (Pair a)) ->
  Exp (Pendulum (Pair (Complex a)) (Pair a))
calc m1 m2 dt pendulum =
  let (zs, ps) = A.unlift pendulum
      zs :: Exp (Pair (Complex a))
      ps :: Exp (Pair a)
      (z1 :: Exp (Complex a), z3 :: Exp (Complex a)) = A.unlift zs
      (p1 :: Exp a, p3 :: Exp a) = A.unlift ps
      i1 = inertia z1 m1
      i3 = inertia z3 m2
      angleDiff = phase z1 - phase z3
      commonFactor = (magnitude z3 / i1)/(18 * i3 * cos angleDiff**2 -
                        8 * absSquared z3 * (3 * m2 + m1))
      ω1 = commonFactor * m1 *
             (3 * magnitude z1 * p3 * cos angleDiff - 2 * magnitude z3 * p1)
      ω3 = commonFactor/i3 *
             (3 * m1 * magnitude z1 * i3 * p1 * cos angleDiff - 2 * i1 * magnitude z3 * p3 * (3 * m2 + m1))
      v1 = cis(ω1 * dt)
      v3 = cis(ω3 * dt)
      pCommonFactor = 6 * magnitude (z1 / z3) * i3 * ω1 * ω3 * sin angleDiff
      p1' = (-pCommonFactor) - g * (m1/2 + m2) * magnitude z1 * cos(phase z1)
      p3' = pCommonFactor - g * m2 * magnitude z3 / 2 * cos(phase z3)
   in A.lift
        ( (z1 * v1, z3 * v3),
          (p1 + p1' * dt, p3 + p3' * dt)
        )
  where
    absSquared z = real (z * conjugate z)
    inertia z m = m * absSquared z / 12

inited ::
  (RealFrac a, A.RealFloat a, Elt a) =>
  Int ->
  Exp (Pair (Complex a)) ->
  Acc (Vector (Pendulum (Pair (Complex a)) (Pair a)))
inited frames zs = A.fill (A.constant $ Z :. 1) $ A.lift (zs, zero)

compute ::
  (RealFrac a, A.RealFloat a, Elt a) =>
  ( Exp (Pendulum (Pair (Complex a)) (Pair a)) ->
    Exp (Pendulum (Pair (Complex a)) (Pair a))
  ) ->
  Acc (Vector (Pendulum (Pair (Complex a)) (Pair a))) ->
  Acc (Vector (Pendulum (Pair (Complex a)) (Pair a)))
compute calc mat =
  let n = indexHead $ A.shape mat
      last = A.replicate (A.constant $ Z :. 1 :: Exp (Z :. Int)) $ A.slice mat (A.lift $ Z :. n - 1)
      new = A.map calc last
   in mat A.++ new

computation ::
  forall a.
  (RealFrac a, A.RealFloat a, Elt a) =>
  Int ->
  Acc (Vector (Pendulum (Pair (Complex a)) (Pair a))) ->
  ( Acc (Vector (Pendulum (Pair (Complex a)) (Pair a))) ->
    Acc (Vector (Pendulum (Pair (Complex a)) (Pair a)))
  ) ->
  Exp (Pair (Complex a)) ->
  Acc (Vector (Pair (Complex a)))
computation frames inited compute zs =
  A.map mapper ran
  where
    pred mat =
      let (A.unlift -> Z :. i :: Z :. Exp Int) = A.shape mat
          in A.unit $ i A.< A.constant frames

    mapper :: Exp (Pendulum (Pair (Complex a)) (Pair a)) -> Exp (Pair (Complex a))
    mapper (A.unlift -> pendulum :: Pendulum (Exp (Pair (Complex a))) (Exp (Pair a))) =
      let (z, _) = pendulum
       in z

    ran :: Acc (Vector (Pendulum (Pair (Complex a)) (Pair a)))
    ran = A.awhile pred compute inited

simulatePendulums ::
  (RealFrac a, A.RealFloat a, Elt a) =>
  (C.Complex a, a) ->
  (C.Complex a, a) ->
  a ->
  a ->
  [Pair (C.Complex a)]
simulatePendulums (z1, m1) (z2, m2) t dt = A.toList $ run computation'
  where
    frames = ceiling $ t / dt
    zs = A.constant (z1, z2)
    calc' = calc (A.constant m1) (A.constant m2) (A.constant dt)
    inited' = inited frames zs
    compute' = compute calc'
    computation' = computation frames inited' compute' zs

skip :: Int
skip = 2

paint :: Double -> (Int, Pair (Complex Double)) -> IO ()
paint m (i, (z1@(x1 C.:+ y1), z2@(x2 C.:+ y2))) =
  when (mod i skip == 0) $ do
    putStrLn $ "z1 = " ++ show z1
    putStrLn $ "z2 = " ++ show z2
    let side = 2000 :: Int
        width = 15
        x1' = realToFrac $ (fromIntegral side / 2) * (1 + x1 / m)
        y1' = realToFrac $ (fromIntegral side / 2) * (1 - y1 / m)
        x2' = realToFrac $ (fromIntegral side / 2) * (1 + (x2 + x1) / m)
        y2' = realToFrac $ (fromIntegral side / 2) * (1 - (y2 + y1) / m)
        bg = PixelRGBA8 0x0c 0x0a 0x20 255
        lineColor = PixelRGBA8 0xdf 0x85 0xff 255
        img = renderDrawing side side bg $
          withTexture (uniformTexture lineColor) $ do
            stroke width JoinRound (CapRound, CapRound) $
              line
                (V2 (fromIntegral side / 2) (fromIntegral side / 2))
                (V2 x1' y1')
                ++ line (V2 x1' y1') (V2 x2' y2')
    writePng (printf "results/double-pendulum-%05v.png" $ div i skip) img

main :: IO ()
main = do
  let z1 = C.cis (pi/4) :: Complex Double
      z2 = C.cis (pi/4) :: Complex Double
      list = zip [0 ..] $ simulatePendulums
                            (z1, 1) (z2, 1)
                            (50 * fromIntegral skip) (1 / (60 * fromIntegral skip))
  forM_ list $ paint $ C.magnitude z1 + C.magnitude z2


{-# LANGUAGE ForeignFunctionInterface #-}
-----------------------------------------------------------------------------
-- |
-- Module     : Foreign.BLAS.Zomplex
-- Copyright  : Copyright (c) 2010, Patrick Perry <patperry@gmail.com>
-- License    : BSD3
-- Maintainer : Patrick Perry <patperry@gmail.com>
-- Stability  : experimental
--

module Foreign.BLAS.Zomplex
    where

import Data.Complex( Complex )
import Foreign
import Foreign.BLAS.Types
import Foreign.C.Types

#include "config.h"
#include "f77_func-hsc.h"
#define la_int int

---------------------------- Level 1 Routines -------------------------------

foreign import ccall unsafe #f77_func zdotu
    zdotu :: Ptr (Complex Double) -> Ptr LAInt -> Ptr (Complex Double) -> Ptr LAInt -> Ptr (Complex Double) -> Ptr LAInt -> IO ()

foreign import ccall unsafe #f77_func zdotc
    zdotc :: Ptr (Complex Double) -> Ptr LAInt -> Ptr (Complex Double) -> Ptr LAInt -> Ptr (Complex Double) -> Ptr LAInt -> IO ()

foreign import ccall unsafe #f77_func dznrm2
    znrm2  :: Ptr LAInt -> Ptr (Complex Double) -> Ptr LAInt -> IO Double

foreign import ccall unsafe #f77_func dzasum
    zasum  :: Ptr LAInt -> Ptr (Complex Double) -> Ptr LAInt -> IO Double

foreign import ccall unsafe #f77_func izamax
    izamax_hidden :: Ptr LAInt -> Ptr (Complex Double) -> Ptr LAInt -> IO CInt
izamax :: Ptr LAInt -> Ptr (Complex Double) -> Ptr LAInt -> IO LAInt
izamax a b c = do
                res <- izamax_hidden a b c
                return $! LAInt res

foreign import ccall unsafe #f77_func zscal
    zscal  :: Ptr LAInt -> Ptr (Complex Double) -> Ptr (Complex Double) -> Ptr LAInt -> IO ()

foreign import ccall unsafe #f77_func zswap
    zswap  :: Ptr LAInt -> Ptr (Complex Double) -> Ptr LAInt -> Ptr (Complex Double) -> Ptr LAInt -> IO ()

foreign import ccall unsafe #f77_func zcopy
    zcopy  :: Ptr LAInt -> Ptr (Complex Double) -> Ptr LAInt -> Ptr (Complex Double) -> Ptr LAInt -> IO ()

foreign import ccall unsafe #f77_func zaxpy
    zaxpy  :: Ptr LAInt -> Ptr (Complex Double) -> Ptr (Complex Double) -> Ptr LAInt -> Ptr (Complex Double) -> Ptr LAInt -> IO ()

foreign import ccall unsafe #f77_func zrotg
    zrotg  :: Ptr (Complex Double) -> Ptr (Complex Double) -> Ptr (Complex Double) -> Ptr (Complex Double) -> IO ()

foreign import ccall unsafe #f77_func zdrot
    zdrot :: Ptr LAInt -> Ptr (Complex Double) -> Ptr LAInt -> Ptr (Complex Double) -> Ptr LAInt -> Ptr Double -> Ptr Double -> IO ()

---------------------------- Level 2 Routines -------------------------------

foreign import ccall unsafe #f77_func zgemv
    zgemv :: BLASTrans -> Ptr LAInt -> Ptr LAInt -> Ptr (Complex Double) -> Ptr (Complex Double) -> Ptr LAInt -> Ptr (Complex Double) -> Ptr LAInt -> Ptr (Complex Double) -> Ptr (Complex Double) -> Ptr LAInt -> IO ()

foreign import ccall unsafe #f77_func zgbmv
    zgbmv ::  BLASTrans -> Ptr LAInt -> Ptr LAInt -> Ptr LAInt -> Ptr LAInt -> Ptr (Complex Double) -> Ptr (Complex Double) -> Ptr LAInt -> Ptr (Complex Double) -> Ptr LAInt -> Ptr (Complex Double) -> Ptr (Complex Double) -> Ptr LAInt -> IO ()

foreign import ccall unsafe #f77_func ztrmv
    ztrmv ::  BLASUplo -> BLASTrans -> BLASDiag -> Ptr LAInt -> Ptr (Complex Double) -> Ptr LAInt -> Ptr (Complex Double) -> Ptr LAInt -> IO ()

foreign import ccall unsafe #f77_func ztpmv
    ztpmv ::  BLASUplo -> BLASTrans -> BLASDiag -> Ptr LAInt -> Ptr (Complex Double) -> Ptr (Complex Double) -> Ptr LAInt -> IO ()

foreign import ccall unsafe #f77_func ztpsv
    ztpsv ::  BLASUplo -> BLASTrans -> BLASDiag -> Ptr LAInt -> Ptr (Complex Double) -> Ptr (Complex Double) -> Ptr LAInt -> IO ()

foreign import ccall unsafe #f77_func ztbmv
    ztbmv ::  BLASUplo -> BLASTrans -> BLASDiag -> Ptr LAInt -> Ptr LAInt -> Ptr (Complex Double) -> Ptr LAInt -> Ptr (Complex Double) -> Ptr LAInt -> IO ()

foreign import ccall unsafe #f77_func ztrsv
    ztrsv ::  BLASUplo -> BLASTrans -> BLASDiag -> Ptr LAInt -> Ptr (Complex Double) -> Ptr LAInt -> Ptr (Complex Double) -> Ptr LAInt -> IO ()

foreign import ccall unsafe #f77_func ztbsv
    ztbsv ::  BLASUplo -> BLASTrans -> BLASDiag -> Ptr LAInt -> Ptr LAInt -> Ptr (Complex Double) -> Ptr LAInt -> Ptr (Complex Double) -> Ptr LAInt -> IO ()

foreign import ccall unsafe #f77_func zhemv
    zhemv ::  BLASUplo -> Ptr LAInt -> Ptr (Complex Double) -> Ptr (Complex Double) -> Ptr LAInt -> Ptr (Complex Double) -> Ptr LAInt -> Ptr (Complex Double) -> Ptr (Complex Double) -> Ptr LAInt -> IO ()

foreign import ccall unsafe #f77_func zhbmv
    zhbmv ::  BLASUplo -> Ptr LAInt -> Ptr LAInt -> Ptr (Complex Double) -> Ptr (Complex Double) -> Ptr LAInt -> Ptr (Complex Double) -> Ptr LAInt -> Ptr (Complex Double) -> Ptr (Complex Double) -> Ptr LAInt -> IO ()

foreign import ccall unsafe #f77_func zgeru
    zgeru  ::  Ptr LAInt -> Ptr LAInt -> Ptr (Complex Double) -> Ptr (Complex Double) -> Ptr LAInt -> Ptr (Complex Double) -> Ptr LAInt -> Ptr (Complex Double) -> Ptr LAInt -> IO ()

foreign import ccall unsafe #f77_func zgerc
    zgerc  ::  Ptr LAInt -> Ptr LAInt -> Ptr (Complex Double) -> Ptr (Complex Double) -> Ptr LAInt -> Ptr (Complex Double) -> Ptr LAInt -> Ptr (Complex Double) -> Ptr LAInt -> IO ()

foreign import ccall unsafe #f77_func zher
    zher  ::  BLASUplo -> Ptr LAInt -> Ptr Double -> Ptr (Complex Double) -> Ptr LAInt -> Ptr (Complex Double) -> Ptr LAInt -> IO ()

foreign import ccall unsafe #f77_func zher2
    zher2 ::  BLASUplo -> Ptr LAInt -> Ptr (Complex Double) -> Ptr (Complex Double) -> Ptr LAInt -> Ptr (Complex Double) -> Ptr LAInt -> Ptr (Complex Double) -> Ptr LAInt -> IO ()

foreign import ccall unsafe #f77_func zhpmv
    zhpmv ::  BLASUplo -> Ptr LAInt -> Ptr (Complex Double) -> Ptr (Complex Double) -> Ptr (Complex Double) -> Ptr LAInt -> Ptr (Complex Double) -> Ptr (Complex Double) -> Ptr LAInt -> IO ()

foreign import ccall unsafe #f77_func zhpr
    zhpr  ::  BLASUplo -> Ptr LAInt -> Ptr Double -> Ptr (Complex Double) -> Ptr LAInt -> Ptr (Complex Double) -> IO ()

foreign import ccall unsafe #f77_func zhpr2
    zhpr2 ::  BLASUplo -> Ptr LAInt -> Ptr (Complex Double) -> Ptr (Complex Double) -> Ptr LAInt -> Ptr (Complex Double) -> Ptr LAInt -> Ptr (Complex Double) -> IO ()

---------------------------- Level 3 Routines -------------------------------

foreign import ccall unsafe #f77_func zgemm
    zgemm  ::  BLASTrans -> BLASTrans -> Ptr LAInt -> Ptr LAInt -> Ptr LAInt -> Ptr (Complex Double) -> Ptr (Complex Double) -> Ptr LAInt -> Ptr (Complex Double) -> Ptr LAInt -> Ptr (Complex Double) -> Ptr (Complex Double) -> Ptr LAInt -> IO ()

foreign import ccall unsafe #f77_func zsymm
    zsymm  ::  BLASSide -> BLASUplo -> Ptr LAInt -> Ptr LAInt -> Ptr (Complex Double) -> Ptr (Complex Double) -> Ptr LAInt -> Ptr (Complex Double) -> Ptr LAInt -> Ptr (Complex Double) -> Ptr (Complex Double) -> Ptr LAInt -> IO ()

foreign import ccall unsafe #f77_func zhemm
    zhemm  ::  BLASSide -> BLASUplo -> Ptr LAInt -> Ptr LAInt -> Ptr (Complex Double) -> Ptr (Complex Double) -> Ptr LAInt -> Ptr (Complex Double) -> Ptr LAInt -> Ptr (Complex Double) -> Ptr (Complex Double) -> Ptr LAInt -> IO ()

foreign import ccall unsafe #f77_func ztrmm
    ztrmm  ::  BLASSide -> BLASUplo -> BLASTrans -> BLASDiag -> Ptr LAInt -> Ptr LAInt -> Ptr (Complex Double) -> Ptr (Complex Double) -> Ptr LAInt -> Ptr (Complex Double) -> Ptr LAInt -> IO ()

foreign import ccall unsafe #f77_func ztrsm
    ztrsm  ::  BLASSide -> BLASUplo -> BLASTrans -> BLASDiag -> Ptr LAInt -> Ptr LAInt -> Ptr (Complex Double) -> Ptr (Complex Double) -> Ptr LAInt -> Ptr (Complex Double) -> Ptr LAInt -> IO ()

foreign import ccall unsafe #f77_func zsyrk
    zsyrk  ::  BLASUplo -> BLASTrans -> Ptr LAInt -> Ptr LAInt -> Ptr (Complex Double) -> Ptr (Complex Double) -> Ptr LAInt -> Ptr (Complex Double) -> Ptr (Complex Double) -> Ptr LAInt -> IO ()

foreign import ccall unsafe #f77_func zsyr2k
    zsyr2k ::  BLASUplo -> BLASTrans -> Ptr LAInt -> Ptr LAInt -> Ptr (Complex Double) -> Ptr (Complex Double) -> Ptr LAInt -> Ptr (Complex Double) -> Ptr LAInt -> Ptr (Complex Double) -> Ptr (Complex Double) -> Ptr LAInt -> IO ()

foreign import ccall unsafe #f77_func zherk
    zherk  ::  BLASUplo -> BLASTrans -> Ptr LAInt -> Ptr LAInt -> Ptr (Complex Double) -> Ptr (Complex Double) -> Ptr LAInt -> Ptr (Complex Double) -> Ptr (Complex Double) -> Ptr LAInt -> IO ()

foreign import ccall unsafe #f77_func zher2k
    zher2k ::  BLASUplo -> BLASTrans -> Ptr LAInt -> Ptr LAInt -> Ptr (Complex Double) -> Ptr (Complex Double) -> Ptr LAInt -> Ptr (Complex Double) -> Ptr LAInt -> Ptr (Complex Double) -> Ptr (Complex Double) -> Ptr LAInt -> IO ()


{-# LANGUAGE DataKinds, FlexibleContexts, FlexibleInstances, GADTs #-}
{-# LANGUAGE MultiParamTypeClasses, TypeFamilies                   #-}
{-# OPTIONS_GHC -fno-warn-orphans #-}
-- | This Library provides some *dangerous* instances for @Double@s and @Complex@.
module Algebra.Instances () where
import Algebra.Scalar

import           AlgebraicPrelude
import           Control.DeepSeq       (NFData (..))
import           Control.Monad.Random  (Random (..), getRandom)
import           Control.Monad.Random  (getRandomR, runRand)
import           Data.Complex          (Complex (..))
import           Data.Convertible.Base (Convertible (..))
import qualified Data.Ratio            as P
import qualified Data.Vector           as DV
import           Data.Vector.Instances ()
import qualified Numeric.Algebra       as NA
import qualified Prelude               as P

instance Additive r => Additive (DV.Vector r) where
  (+) = DV.zipWith (+)

-- | These Instances are not algebraically right, but for the sake of convenience.
instance DecidableZero r => DecidableZero (Complex r) where
  isZero (a :+ b) = isZero a && isZero b

instance (NFData a) => NFData (Fraction a) where
  rnf a = rnf (numerator a) `seq` rnf (denominator a) `seq` ()

instance Additive r => Additive (Complex r) where
  (a :+ b) + (c :+ d) = (a + c) :+ (b + d)
instance Abelian r => Abelian (Complex r) where
instance (Group r, Semiring r) => Semiring (Complex r) where
instance (Group r, Rig r) => Rig (Complex r) where
  fromNatural = (:+ zero) . fromNatural
instance (Group r, Commutative r) => Commutative (Complex r) where
instance Ring r => Ring (Complex r) where
  fromInteger = (:+ zero) . fromInteger'
instance Group r => Group (Complex r) where
  (a :+ b) - (c :+ d) = (a - c) :+ (b - d)
  negate (a :+ b) = negate a :+ negate b
  times n (a :+ b) = times n a :+ times n b
instance LeftModule a r => LeftModule a (Complex r) where
  r .* (a :+ b) = (r .* a) :+ (r .* b)
instance RightModule a r => RightModule a (Complex r) where
  (a :+ b) *. r = (a *. r) :+ (b *. r)
instance Monoidal r => Monoidal (Complex r) where
  zero = zero :+ zero
instance (Group r, Monoidal r, Unital r) => Unital (Complex r) where
  one = one :+ zero
instance Additive Double where
  (+) = (P.+)
instance (Group r, Multiplicative r) => Multiplicative (Complex r) where
  (a :+ b) * (c :+ d) = (a*c - b*d) :+ (a*d + b*c)
instance LeftModule Natural Double where
  n .* d = fromIntegral n P.* d
instance RightModule Natural Double where
  d *. n = d P.* fromIntegral n
instance Monoidal Double where
  zero = 0
instance Unital Double where
  one = 1
instance Multiplicative Double where
  (*) = (P.*)
instance Commutative Double where
instance Group Double where
  (-) = (P.-)
  negate = P.negate
  subtract = P.subtract
  times n r = P.fromIntegral n P.* r
instance LeftModule Integer Double where
  n .* r = P.fromInteger n * r
instance RightModule Integer Double where
  r *. n = r * P.fromInteger n
instance Rig Double where
  fromNatural = P.fromInteger . fromNatural
instance Semiring Double where
instance Abelian Double where
instance Ring Double where
  fromInteger = P.fromInteger
instance DecidableZero Double where
  isZero 0 = True
  isZero _ = False

instance Division Double where
  recip = P.recip
  (/)   = (P./)

instance P.Integral r => Additive (P.Ratio r) where
  (+) = (P.+)

instance P.Integral r => Abelian (P.Ratio r)

instance P.Integral r => LeftModule Natural (P.Ratio r) where
  n .* r = fromIntegral n P.* r

instance P.Integral r => RightModule Natural (P.Ratio r) where
  r *. n = r P.* fromIntegral n

instance P.Integral r => LeftModule Integer (P.Ratio r) where
  n .* r = P.fromInteger n P.* r

instance P.Integral r => RightModule Integer (P.Ratio r) where
  r *. n = r P.* P.fromInteger n

instance P.Integral r => Group (P.Ratio r) where
  (-)    = (P.-)
  negate = P.negate
  subtract = P.subtract
  times n r = P.fromIntegral n P.* r

instance P.Integral r => Commutative (P.Ratio r)

instance (Semiring r, P.Integral r) => LeftModule (Scalar r) (P.Ratio r) where
  Scalar n .* r = (n P.% 1) * r

instance (Semiring r, P.Integral r) => RightModule (Scalar r) (P.Ratio r) where
  r *. Scalar n = r * (n P.% 1)

instance P.Integral r => Multiplicative (P.Ratio r) where
  (*) = (P.*)

instance P.Integral r => Unital (P.Ratio r) where
  one = 1

instance P.Integral r => Division (P.Ratio r) where
  (/) = (P./)
  recip = P.recip

instance P.Integral r => Monoidal (P.Ratio r) where
  zero = 0

instance P.Integral r => Semiring (P.Ratio r)

instance P.Integral r => Rig (P.Ratio r) where
  fromNatural = P.fromIntegral

instance P.Integral r => Ring (P.Ratio r) where
  fromInteger = P.fromInteger

instance P.Integral r => DecidableZero (P.Ratio r) where
  isZero 0 = True
  isZero _ = False

instance P.Integral r => DecidableUnits (P.Ratio r) where
  isUnit 0 = False
  isUnit _ = True
  recipUnit 0 = Nothing
  recipUnit n = Just (P.recip n)
  r ^? n
    | r == 0 = Just 1
    | r /= 0 = Just (r P.^^ n)
    | r == 0 && n P.> 0 = Just 0
    | otherwise = Nothing

instance Convertible (Fraction Integer) Double where
  safeConvert a = Right $ P.fromInteger (numerator a) P./ P.fromInteger (denominator a)

instance Convertible (Fraction Integer) (Complex Double) where
  safeConvert a = Right $ P.fromInteger (numerator a) P./ P.fromInteger (denominator a) :+ 0

instance (Random (Fraction Integer)) where
  random = runRand $ do
    i <- getRandom
    j <- getRandom
    return $ i % (P.abs j + 1)
  randomR (a, b) = runRand $ do
    j <- succ . P.abs <$> getRandom
    let g = foldl1 P.lcm  [denominator a, denominator b, j]
        lb = g * numerator a `quot` denominator a
        ub = g * numerator b `quot` denominator b
    i <- getRandomR (lb, ub)
    return $ i % g


{-# LANGUAGE PatternSynonyms #-}
{-# LANGUAGE BangPatterns #-}
module Polestar.Type where
import Data.Complex
import Data.Monoid
import Data.Ord

-- variable name (for debugging)
newtype Id = Id String deriving (Show)

-- ordering: s <: t implies s < t
data PrimType
  -- numeric types
  = PTyZero
  | PTyNat
  | PTyInt
  | PTyNNReal -- non-negative real number
  | PTyReal
  | PTyImaginary
  | PTyComplex
  -- boolean
  | PTyTrue
  | PTyFalse
  | PTyBool
  -- other
  | PTyUnit
  deriving (Eq,Ord,Show,Enum,Bounded)

data Type = TyPrim !PrimType
          | TyArr Type Type
          | TyAll Id !(Maybe Type) Type
          | TyRef !Int
          | TyTuple [Type] -- must have 2 or more elements
          | TyInter [Type] -- must have 2 or more elements
          deriving (Eq,Show)

pattern TyZero = TyPrim PTyZero
pattern TyNat = TyPrim PTyNat
pattern TyInt = TyPrim PTyInt
pattern TyNNReal = TyPrim PTyNNReal
pattern TyReal = TyPrim PTyReal
pattern TyImaginary = TyPrim PTyImaginary
pattern TyComplex = TyPrim PTyComplex
pattern TyTrue = TyPrim PTyTrue
pattern TyFalse = TyPrim PTyFalse
pattern TyBool = TyPrim PTyBool
pattern TyUnit = TyPrim PTyUnit
pattern TyPair a b = TyTuple [a,b]

-- compound canonical type
-- no pair of the components have subtyping relation
type CanonicalType = [ICanonicalType]

-- individual canonical type
data ICanonicalType = CTyPrim !PrimType
                    | CTyArr CanonicalType ICanonicalType
                    | CTyAll Id CanonicalType ICanonicalType
                    | CTyRef !Int
                    | CTyTuple [ICanonicalType]
                    deriving (Eq,Show)

pattern CTyTrue = CTyPrim PTyTrue
pattern CTyFalse = CTyPrim PTyFalse
pattern CTyBool = CTyPrim PTyBool

data Builtin
  -- unary
  = BNegate
  | BLogicalNot
  | BNatToInt
  | BNatToNNReal
  | BIntToNat -- max(n,0)
  | BIntToReal
  | BNNRealToReal
  | BRealToComplex
  | BMkImaginary
  | BImaginaryToComplex
  | BRealPart
  | BImagPart
  | BAbs
  | BSqrt
  | BExp
  | BExpm1
  | BLog
  | BLog1p
  | BSin
  | BCos
  | BTan
  | BSinh
  | BCosh
  | BTanh
  | BAsin
  | BAcos
  | BAtan
  | BAsinh
  | BAcosh
  | BAtanh
  -- TODO: factorial
  -- TODO: exp, log, expm1, log1p
  -- TODO: trigonometric functions and hyperbolic functions
  -- binary
  | BAdd
  | BSub
  | BMul
  | BDiv
  | BPow
  | BTSubNat -- truncated subtraction
  | BLt
  | BLe
  | BEqual
  | BMax
  | BMin
  | BIntDiv
  | BIntMod
  | BGcd
  | BLcm
  | BLogicalAnd
  | BLogicalOr
  -- TODO: binomial coefficients
  -- other primitives
  | BIterate
  -- | BUnsafeGlue
  deriving (Eq,Show,Enum,Bounded)

data PrimValue = PVZero
               | PVInt !Integer
               | PVReal !Double
               | PVImaginary !Double
               | PVComplex !(Complex Double)
               | PVBool !Bool
               | PVUnit
               | PVBuiltin !Builtin
               deriving (Eq,Show)

data Term = TmPrim !PrimValue             -- primitive value
          | TmAbs Id Type Term            -- lambda abstraction
          | TmTyAbs Id !(Maybe Type) Term -- bounded type abstraction
          | TmRef !Int                    -- variable (de Bruijn index)
          | TmApp Term Term               -- function application
          | TmTyApp Term Type             -- type application
          | TmLet Id Term Term            -- let-in
          | TmAlt Id [Type] Term          -- type alternation
          | TmIf Term Term Term           -- if-then-else
          | TmTuple [Term]                -- tuple
          | TmProj Term !Int              -- projection
          | TmCoerce Term Type            -- coercion
          | TmCoherentTuple [Term]        -- coherent tuple
          deriving (Eq,Show)

data Binding = VarBind Id CanonicalType
             | TyVarBind Id CanonicalType
             | AnonymousBind
             deriving (Eq,Show)

isUnary :: Builtin -> Bool
isUnary f = case f of
  BNegate -> True
  BLogicalNot -> True
  BNatToInt -> True
  BNatToNNReal -> True
  BIntToNat -> True
  BIntToReal -> True
  BNNRealToReal -> True
  BRealToComplex -> True
  BMkImaginary -> True
  BImaginaryToComplex -> True
  BRealPart -> True
  BImagPart -> True
  BAbs -> True
  BSqrt -> True
  BExp -> True
  BExpm1 -> True
  BLog -> True
  BLog1p -> True
  BSin -> True
  BCos -> True
  BTan -> True
  BSinh -> True
  BCosh -> True
  BTanh -> True
  BAsin -> True
  BAcos -> True
  BAtan -> True
  BAsinh -> True
  BAcosh -> True
  BAtanh -> True
  _ -> False

isBinary :: Builtin -> Bool
isBinary f = case f of
  BAdd -> True
  BSub -> True
  BMul -> True
  BDiv -> True
  BPow -> True
  BTSubNat -> True
  BLt -> True
  BLe -> True
  BEqual -> True
  BMax -> True
  BMin -> True
  BIntDiv -> True
  BIntMod -> True
  BGcd -> True
  BLcm -> True
  BLogicalAnd -> True
  BLogicalOr -> True
  _ -> False

isValue :: Term -> Bool
isValue t = case t of
  TmPrim _ -> True
  TmAbs _ _ _ -> True
  TmTyAbs _ _ _ -> True
  TmApp (TmPrim (PVBuiltin f)) x | isBinary f -> isValue x -- partial application
  TmTyApp (TmPrim (PVBuiltin BIterate)) ty -> True
  TmApp (TmTyApp (TmPrim (PVBuiltin BIterate)) ty) x -> isValue x
  TmApp (TmApp (TmTyApp (TmPrim (PVBuiltin BIterate)) ty) x) y -> isValue x && isValue y
  TmTuple xs -> all isValue xs
  TmCoherentTuple xs -> all isValue xs
  _ -> False

getCTypeFromContext :: [Binding] -> Int -> CanonicalType
getCTypeFromContext ctx i
  | 0 <= i && i < length ctx = case ctx !! i of
      VarBind _ ty -> ty
      b -> error ("TmRef: expected a variable binding, found " ++ show b)
  | otherwise = error "TmRef: index out of bounds"

getTypeFromContext :: [Binding] -> Int -> Type
getTypeFromContext ctx i = canonicalToOrdinary $ getCTypeFromContext ctx i

getCBoundFromContext :: [Binding] -> Int -> CanonicalType
getCBoundFromContext ctx i
  | i < length ctx = case ctx !! i of
                       TyVarBind _ b -> b
                       b -> error ("TyRef: expected a type variable binding, found " ++ show b)
  | otherwise = error "TyRef: index out of bounds"

getBoundFromContext :: [Binding] -> Int -> Maybe Type
getBoundFromContext ctx i = canonicalToOrdinaryM $ getCBoundFromContext ctx i

typeShift :: Int -> Int -> Type -> Type
-- typeShift 0 _ ty = ty
typeShift !delta = go
  where
    go !i ty = case ty of
      TyPrim _ -> ty
      TyArr u v -> TyArr (go i u) (go (i + 1) v)
      TyAll name b t -> TyAll name (typeShift delta i <$> b) (typeShift delta (i + 1) t)
      TyRef j | j >= i, j + delta >= 0 -> TyRef (j + delta)
              | j >= i, j + delta < 0 -> error "typeShift: negative index"
              | otherwise -> ty
      TyTuple tys -> TyTuple $ map (go i) tys
      TyInter tys -> TyInter $ map (go i) tys

typeShiftC :: Int -> Int -> CanonicalType -> CanonicalType
-- typeShiftC 0 _ = id
typeShiftC !delta !i = map (typeShiftI delta i)

typeShiftI :: Int -> Int -> ICanonicalType -> ICanonicalType
typeShiftI !delta = go
  where
    go :: Int -> ICanonicalType -> ICanonicalType
    go !i ty = case ty of
      CTyPrim _ -> ty
      CTyArr u v -> CTyArr (map (go i) u) (go (i + 1) v)
      CTyAll name b t -> CTyAll name (map (go i) b) (go (i + 1) t)
      CTyRef j | j >= i, j + delta >= 0 -> CTyRef (j + delta)
               | j >= i, j + delta < 0 -> error "typeShift: negative index"
               | otherwise -> ty
      CTyTuple tys -> CTyTuple $ map (go i) tys

typeSubstD :: Int -> Type -> Int -> Type -> Type
typeSubstD !depth s !i ty = case ty of
  TyPrim _ -> ty
  TyArr u v -> TyArr (typeSubstD depth s i u) (typeSubstD (depth + 1) s (i + 1) v)
  TyAll name b t -> TyAll name (typeSubstD depth s i <$> b) (typeSubstD (depth + 1) s (i + 1) t)
  TyRef j | j == i -> typeShift depth 0 s
          | j > i -> TyRef (j - 1)
          | otherwise -> ty
  TyTuple tys -> TyTuple $ map (typeSubstD depth s i) tys
  TyInter tys -> TyInter $ map (typeSubstD depth s i) tys

typeSubst = typeSubstD 0

typeSubstCD :: Int -> CanonicalType -> Int -> ICanonicalType -> CanonicalType
typeSubstCD !depth s !i ty = case ty of
  CTyPrim _ -> return ty
  CTyArr u v -> CTyArr (u >>= typeSubstCD depth s i) <$> typeSubstCD (depth + 1) s (i + 1) v
  CTyAll name b t -> CTyAll name (b >>= typeSubstCD depth s i) <$> typeSubstCD (depth + 1) s (i + 1) t
  CTyRef j | j == i -> typeShiftC depth 0 s
           | j > i -> return $ CTyRef (j - 1)
           | otherwise -> return ty
  CTyTuple tys -> CTyTuple <$> mapM (typeSubstCD depth s i) tys

typeSubstC = typeSubstCD 0

canonicalToOrdinary :: CanonicalType -> Type
canonicalToOrdinary [] = error "canonicalToOrdinary: Top type"
canonicalToOrdinary [t] = iCanonicalToOrdinary t
canonicalToOrdinary tys = TyInter $ map iCanonicalToOrdinary tys

canonicalToOrdinaryM :: CanonicalType -> Maybe Type
canonicalToOrdinaryM [] = Nothing
canonicalToOrdinaryM [t] = Just $ iCanonicalToOrdinary t
canonicalToOrdinaryM tys = Just $ TyInter $ map iCanonicalToOrdinary tys

iCanonicalToOrdinary :: ICanonicalType -> Type
iCanonicalToOrdinary (CTyPrim p) = TyPrim p
iCanonicalToOrdinary (CTyArr u v) = TyArr (canonicalToOrdinary u) (iCanonicalToOrdinary v)
iCanonicalToOrdinary (CTyAll name b t) = TyAll name (canonicalToOrdinaryM b) (iCanonicalToOrdinary t)
iCanonicalToOrdinary (CTyRef i) = TyRef i
iCanonicalToOrdinary (CTyTuple tys) = TyTuple $ map iCanonicalToOrdinary tys

instance Eq Id where
  _ == _ = True

{-
instance Ord ICanonicalType where
  compare (CTyPrim p) (CTyPrim p') = compare p p'
  compare (CTyArr s t) (CTyArr s' t') = compare (Down s) (Down s') <> compare t t'
  compare (CTyAll _ s t) (CTyAll _ s' t') = compare (Down s) (Down s') <> compare t t'
  compare (CTyRef i) (CTyRef i') = compare i i'
  compare (CTyTuple tys) (CTyTuple tys') = compare tys tys'
  compare (CTyPrim _) _ = LT
  compare _ (CTyPrim _) = GT
  compare (CTyArr _ _) _ = LT
  compare _ (CTyArr _ _) = GT
  compare (CTyAll _ _ _) _ = LT
  compare _ (CTyAll _ _ _) = GT
  compare (CTyRef _) _ = LT
  compare _ (CTyRef _) = GT
  -- compare (CTyTuple _) _ = LT
  -- compare _ (CTyTuple _) = GT
-}


module QIO.Circuit.Compile where

import qualified QIO.Circuit.Circuit as C

import Data.Complex
import QIO.QioSyn
import QIO.Qio

compileCircuit :: C.Circuit -> QIO [Bool]
compileCircuit circuit = do
  qs <- mkQbits (length $ C.qubits circuit) False
  applyU (gatesToU (C.gates circuit) qs)
  measQbits qs

gatesToU :: [C.Gate] -> [Qbit] -> U
gatesToU [] _      = mempty
gatesToU (g:gs) qs = (gateToU g qs) <> (gatesToU gs qs)

gateToU :: C.Gate -> [Qbit] -> U
gateToU (C.Had _ qis) qs         = condU uhad $ map (\i -> qs !! i) qis
gateToU (C.PX _ qis) qs          = condU unot $ map (\i -> qs !! i) qis
gateToU (C.PY _ qis) qs          = condU (\q -> rot q yRot) $ map (\i -> qs !! i) qis
                                   where
                                     yRot (False, True) = 0 :+ 1
                                     yRot (True, False) = 0 :+ (-1)
                                     yRot _             = 0
gateToU (C.PZ _ qis) qs          = condU (\q -> uphase q pi) $ map (\i -> qs !! i) qis
gateToU (C.Swap _ (q1:q2:[])) qs = swap (qs !! q1) (qs !! q2)
gateToU _ _                      = mempty

mkQbits :: Int -> Bool -> QIO [Qbit]
mkQbits n b = mkQbits' n b []
  where
    mkQbits' 0 _ qs = return qs
    mkQbits' n b qs = do
      q <- mkQbit b
      mkQbits' (n-1) b (qs ++ [q])

measQbits :: [Qbit] -> QIO [Bool]
measQbits qs = measQbits' qs []
  where
    measQbits' [] bs = return bs
    measQbits' (q:qs) bs = do
      b <- measQbit q
      measQbits' qs (bs ++ [b])

condU :: (Qbit -> U) -> [Qbit] -> U
condU u (ql:[]) = u ql
condU u (q:qs)  = cond q (\x -> if x then (condU u qs) else mempty)


{-# LANGUAGE AllowAmbiguousTypes        #-}
{-# LANGUAGE DeriveGeneric              #-}
{-# LANGUAGE FlexibleContexts           #-}
{-# LANGUAGE GADTs                      #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE InstanceSigs               #-}
{-# LANGUAGE LambdaCase                 #-}
{-# LANGUAGE ScopedTypeVariables        #-}
{-# LANGUAGE StandaloneDeriving         #-}
{-# LANGUAGE TypeApplications           #-}
{-# LANGUAGE TypeFamilies               #-}
{-# LANGUAGE TypeFamilyDependencies     #-}
{-# LANGUAGE TypeInType                 #-}
{-# LANGUAGE TypeOperators              #-}
{-# LANGUAGE UndecidableInstances       #-}
{-# LANGUAGE ViewPatterns               #-}

module Numeric.BLAS.HMatrix (
    HM(..)
  , HM'
  , flatten
  , unflatten
  ) where

import           Control.Applicative
import           Control.DeepSeq
import           Control.Monad
import           Control.Monad.Trans.State
import           Data.Finite.Internal
import           Data.Foldable
import           Data.Function
import           Data.Kind
import           Data.Maybe
import           Data.MonoTraversable
import           Data.Monoid                  (Endo(..))
import           Data.Singletons
import           Data.Singletons.Prelude
import           Data.Singletons.TypeLits
import           Data.Type.Product
import           GHC.Generics                 (Generic)
import           GHC.TypeLits
import           Numeric.BLAS hiding          (outer)
import           Numeric.LinearAlgebra.Static
import           Type.Family.List hiding      (Reverse)
import qualified Data.Vector                  as UV
import qualified Data.Vector.Generic.Sized    as VG
import qualified Data.Vector.Sized            as V
import qualified Data.Vector.Storable         as UVS
import qualified Data.Vector.Storable.Sized   as VS
import qualified Numeric.LinearAlgebra        as LA

type family HM' (s :: [Nat]) = (h :: Type) | h -> s where
    HM' '[]            = Double
    HM' '[n]           = R n
    HM' '[m,n]         = L m n
    HM' (n ': m ': ms) = V.Vector n (HM' (m ': ms))

newtype HM :: [Nat] -> Type where
    HM  :: { getHM :: HM' b }
        -> HM b
  deriving (Generic)

type instance Element (HM s) = Double

instance BLAS HM where

    transp (HM x) = HM (tr x)

    scal α (HM x)        = HM (konst α * x)
    axpy α (HM x) (HM y) = HM (konst α * x + y)
    dot    (HM x) (HM y) = x <.> y
    norm2  (HM x)        = norm_2 x
    asum   (HM x)        = norm_1 x

    iamax
        :: forall n. KnownNat n
        => HM '[n + 1]
        -> Finite (n + 1)
    iamax  (HM x)        = withKnownNat (SNat @n %:+ SNat @1) $
        Finite . fromIntegral
      . LA.maxIndex . extract
      $ abs x

    gemv α (HM a) (HM x) = \case
        Just (β, HM y) -> HM (konst α * (a #> x) + konst β * y)
        Nothing        -> HM (konst α * (a #> x))
    ger  α (HM x) (HM y) = \case
        Just (HM a) -> HM (konst α * outer x y + a)
        Nothing     -> HM (konst α * outer x y)
    gemm α (HM a) (HM b) = \case
        Just (β, HM c) -> HM (konst α * (a <> b) + konst β * c)
        Nothing        -> HM (konst α * (a <> b))

instance Tensor HM where
    type Scalar  HM = Double

    gen = \case
      SNil -> \f -> HM $ f Ø
      n@SNat `SCons` SNil -> \f -> HM . fromJust . create $
        LA.build (fromIntegral (fromSing n))
          (f . only . Finite . round)
      m@SNat `SCons` (n@SNat `SCons` SNil) -> \f -> HM . fromJust . create $
        LA.build (fromIntegral (fromSing m), fromIntegral (fromSing n))
          (\i j -> f (Finite (round i) :< Finite (round j) :< Ø))
      SNat `SCons` ns@(_ `SCons` (_ `SCons` _)) -> \f -> HM $
        V.generate_ $ \i -> getHM $ gen ns (\js -> f (i :< js))

    genA = \case
      SNil -> \f -> HM <$> f Ø
      n@SNat `SCons` SNil -> \f ->
        fmap (HM . vector) . traverse (f . only . Finite) $
          [0 .. fromSing n - 1]
      m@SNat `SCons` (n@SNat `SCons` SNil) -> \f ->
        fmap (HM . matrix) . traverse f $
          [ Finite i :< Finite j :< Ø | j <- [0 .. fromSing m - 1]
                                      , i <- [0 .. fromSing n - 1]
          ]
      SNat `SCons` ns@(_ `SCons` (_ `SCons` _)) -> \f -> fmap HM $
        sequenceA . V.generate_ $ \i -> getHM <$> genA ns (\js -> f (i :< js))

    tkonst s = HM . hkonst s

    tsum :: forall s. SingI s => HM s -> Double
    tsum = go sing . getHM
      where
        go :: Sing ns -> HM' ns -> Double
        go = \case
          SNil                                      -> id
          SNat `SCons` SNil                         -> LA.sumElements . extract
          SNat `SCons` (SNat `SCons` SNil)          -> LA.sumElements . extract
          SNat `SCons` ns@(_ `SCons` (_ `SCons` _)) -> sum . fmap (go ns)

    tmap :: forall s. SingI s => (Double -> Double) -> HM s -> HM s
    tmap = omap

    tzip
        :: forall s. SingI s
        => (Double -> Double -> Double)
        -> HM s -> HM s -> HM s
    tzip f (HM x0) (HM y0) = HM $ hzip f sing x0 y0

    tindex
        :: forall s. SingI s
        => Prod Finite s
        -> HM s
        -> Double
    tindex ix0 = go sing ix0 . getHM
      where
        go :: Sing ns -> Prod Finite ns -> HM' ns -> Double
        go = \case
          SNil -> \case
            Ø -> id
          SNat `SCons` SNil -> \case
            i :< Ø -> (`LA.atIndex` fromIntegral i) . extract
          SNat `SCons` (SNat `SCons` SNil) -> \case
            i :< j :< Ø -> (`LA.atIndex` (fromIntegral i, fromIntegral j)) . extract
          SNat `SCons` ns@(_ `SCons` (_ `SCons` _)) -> \case
            i :< js -> go ns js . (`V.index` i)

    tslice
        :: forall n m. ()
        => ProdMap Slice n m
        -> HM n
        -> HM m
    tslice sl0 (HM x0) = HM $ go sl0 x0
      where
        go  :: forall ns ms. ()
            => ProdMap Slice ns ms
            -> HM' ns
            -> HM' ms
        go = \case
          PMZ -> \x -> x
          PMS (Slice sL sC@SNat sR) PMZ -> \xs -> fromJust . create $
            let l = fromIntegral $ fromSing sL
                c = fromIntegral $ fromSing sC
            in  withKnownNat (sL %:+ sC %:+ sR) $
                  LA.subVector l c (extract xs)
          PMS (Slice sLy sCy@SNat sRy) (PMS (Slice sLx sCx@SNat sRx) PMZ) ->
              \xs -> fromJust . create $
            let lx = fromIntegral $ fromSing sLx
                ly = fromIntegral $ fromSing sLy
                cx = fromIntegral $ fromSing sCx
                cy = fromIntegral $ fromSing sCy
            in  withKnownNat (sLy %:+ sCy %:+ sRy) $
                withKnownNat (sLx %:+ sCx %:+ sRx) $
                  LA.subMatrix (ly, lx) (cy, cx) (extract xs)
          PMS (Slice sL sC@SNat _) pm@(PMS _ (PMS _ _)) -> \xs ->
            let l = fromIntegral $ fromSing sL
                c = fromIntegral $ fromSing sC
            in  fmap (go pm)
                  . fromJust . V.toSized
                  . UV.take c
                  . UV.drop l
                  . V.fromSized
                  $ xs

    tconv
        :: forall n m s. ()
        => Sing n
        -> ProdMap Conv m s
        -> HM (m >: n)      -- ^ mask
        -> HM s
        -> HM (s >: n)
    tconv sn@SNat pm0 (HM m0) (HM x0) = HM $ go pm0 m0 x0
      where
        go  :: forall ms ss. ()
            => ProdMap Conv ms ss
            -> HM' (ms >: n)
            -> HM' ss
            -> HM' (ss >: n)
        go = \case
          PMZ -> \m x -> konst x * m
          PMS (Conv sM sS str off) PMZ -> \m x ->
            undefined

    treshape
        :: (SingI s1, Product s1 ~ Product s2)
        => Sing s2
        -> HM s1
        -> HM s2
    treshape s = unflatten s . flatten

    tload
        :: Sing s
        -> V.Vector (Product s) Double
        -> HM s
    tload s = unflatten s . VG.convert

    textract
        :: SingI s
        => HM s
        -> V.Vector (Product s) Double
    textract = VG.convert . flatten

-- hconv
--     :: forall o ms ns. (KnownNat o)
--     => DoubleProd Sing ms ns
--     -> HM' (o ': ms)
--     -> HM' ns
--     -> HM' (o ': ns)
-- hconv = \case
--     DPZ -> \ms x -> konst x * ms
--     DPS SNat sn@SNat DPZ -> \m x -> fromJust . create $
--       let c = LA.conv2 (extract m) (LA.asRow (extract x))
--           o = fromInteger (natVal (Proxy @o))
--           left = o `div` 2
--       in  LA.subMatrix (0,left) (o, fromInteger (fromSing sn)) c
--     DPS smx@SNat snx@SNat (DPS smy@SNat sny@SNat DPZ) -> \m x ->
--       -- todo: vectorize with im2colV
--       flip fmap m $ \msk -> fromJust . create $
--         let c = LA.conv2 (extract msk) (extract x)
--             left = fromInteger (fromSing smx) `div` 2
--             top  = fromInteger (fromSing smy) `div` 2
--         in  LA.subMatrix (left, top) (fromInteger (fromSing snx), fromInteger (fromSing sny)) c
--     dp0@(DPS (SNat :: Sing m0) (SNat :: Sing n0) (DPS _ _ (DPS _ _ _) :: DoubleProd Sing ms0 ns0)) ->
--               \(ms :: V.Vector o (V.Vector m0 (HM' ms0))) (xs :: V.Vector n0 (HM' ns0)) ->
--       flip fmap ms $ \(msk :: V.Vector m0 (HM' ms0)) -> hconv1 dp0 msk xs

-- hconv1
--     :: forall m s. ()
--     => DoubleProd Sing m s
--     -> HM' m
--     -> HM' s
--     -> HM' s
-- hconv1 = \case
--     DPZ -> (*)
--     DPS sm@SNat sn@SNat DPZ -> \m x -> fromJust . create $
--       let c = LA.conv (extract m) (extract x)
--           left = fromInteger (fromSing sm) `div` 2
--       in  UVS.slice left (fromInteger (fromSing sn)) c
--     DPS smx@SNat snx@SNat (DPS smy@SNat sny@SNat DPZ) -> \m x -> fromJust . create $
--       let c = LA.conv2 (extract m) (extract x)
--           left = fromInteger (fromSing smx) `div` 2
--           top  = fromInteger (fromSing smy) `div` 2
--       in  LA.subMatrix (left, top) (fromInteger (fromSing snx), fromInteger (fromSing sny)) c
--     DPS (SNat :: Sing m0) (SNat :: Sing n0) dps@(DPS _ _ (DPS _ _ _) :: DoubleProd Sing ms0 ns0) ->
--               \(ms :: V.Vector m0 (HM' ms0)) (xs :: V.Vector n0 (HM' ns0)) ->
--       let s   :: Sing ns0
--           s   = prodSing $ secondDP dps
--           cl :: V.Vector n0 (V.Vector m0 (HM' ns0))
--           cl = im2colV (hkonst s 0) xs
--       in  fmap (hsum s . V.zipWith (hconv1 dps) ms) cl

flatten :: SingI s => HM s -> VS.Vector (Product s) Double
flatten = hflatten sing . getHM

unflatten :: Sing s -> VS.Vector (Product s) Double -> HM s
unflatten s = HM . hunflatten s

hflatten
    :: Sing s
    -> HM' s
    -> VS.Vector (Product s) Double
hflatten = \case
    SNil -> VS.singleton
    SNat `SCons` SNil -> fromJust . VS.toSized . extract
    sn@SNat `SCons` (sm@SNat `SCons` SNil) -> case sn %:* sm of
      SNat -> fromJust . VS.toSized . LA.flatten . extract
    SNat `SCons` ns@(_ `SCons` (_ `SCons` _)) ->
      VG.convert . V.concatMap (VG.convert . hflatten ns)

hunflatten
    :: Sing s
    -> VS.Vector (Product s) Double
    -> HM' s
hunflatten = \case
    SNil -> VS.head
    SNat `SCons` SNil -> fromJust . create . VS.fromSized
    SNat `SCons` (sm@SNat `SCons` SNil) -> fromJust . create . LA.reshape (fromInteger (fromSing sm)) . VS.fromSized
    sn@SNat `SCons` ns@(_ `SCons` (_ `SCons` _)) -> case sProduct ns of
      sp@SNat -> fromJust
            . V.fromList
            . evalState (replicateM (fromInteger (fromSing sn)) $
                           hunflatten ns . fromJust . VS.toSized <$> state (UVS.splitAt (fromInteger (fromSing sp)))
                        )
            . VS.fromSized

im2col
    :: forall m n o. (KnownNat m, KnownNat n, KnownNat o)
    => L n m
    -> L n o
im2col = undefined

im2colV
    :: forall m n a. (KnownNat m, KnownNat n)
    => a
    -> V.Vector n a
    -> V.Vector n (V.Vector m a)
im2colV pad (V.fromSized->v) = V.generate $ \i ->
      fromJust . V.toSized $ UV.slice i m padded
  where
    padded = UV.concat [UV.replicate left pad, v, UV.replicate right pad]
    m :: Int
    m  = fromIntegral $ natVal (Proxy @m)
    left  = m `div` 2
    right = m - left

hadd :: Sing s -> HM' s -> HM' s -> HM' s
hadd = \case
    SNil                                     -> (+)
    SNat `SCons` SNil                        -> (+)
    SNat `SCons` (SNat `SCons` SNil)         -> (+)
    SNat `SCons` s@(_ `SCons` (_ `SCons` _)) -> liftA2 (hadd s)

hsum :: Foldable f => Sing s -> f (HM' s) -> HM' s
hsum s = foldl' (hadd s) (hkonst s 0)

hkonst :: Sing s -> Double -> HM' s
hkonst = \case
    SNil                                     -> id
    SNat `SCons` SNil                        -> konst
    SNat `SCons` (SNat `SCons` SNil)         -> konst
    SNat `SCons` s@(_ `SCons` (_ `SCons` _)) -> pure . hkonst s

hzip
    :: (Double -> Double -> Double)
    -> Sing s
    -> HM' s
    -> HM' s
    -> HM' s
hzip f = go
  where
    go :: Sing t -> HM' t -> HM' t -> HM' t
    go = \case
      SNil                             -> f
      SNat `SCons` SNil                -> zipWithVector f
      SNat `SCons` (SNat `SCons` SNil) ->
        (\xs ys -> matrix (zipWith f xs ys))
           `on` (concat . LA.toLists . extract)
      SNat `SCons` ns@(_ `SCons` (_ `SCons` _)) ->
         V.zipWith (go ns)

-- firstDP
--     :: DoubleProd f as bs
--     -> Prod f as
-- firstDP = \case
--     DPZ        -> Ø
--     DPS x _ xs -> x :< firstDP xs

-- secondDP
--     :: DoubleProd f as bs
--     -> Prod f bs
-- secondDP = \case
--     DPZ        -> Ø
--     DPS _ x xs -> x :< secondDP xs

prodSing
    :: Prod Sing as
    -> Sing as
prodSing = \case
    Ø       -> SNil
    x :< xs -> x `SCons` prodSing xs

instance SingI s => MonoFunctor (HM s) where
    omap f = HM . go sing . getHM
      where
        go :: Sing ns -> HM' ns -> HM' ns
        go = \case
          SNil                                      -> f
          SNat `SCons` SNil                         -> dvmap f
          SNat `SCons` (SNat `SCons` SNil)          -> dmmap f
          SNat `SCons` ns@(_ `SCons` (_ `SCons` _)) -> fmap (go ns)

helems :: forall s. SingI s => HM s -> [Double]
helems = flip appEndo [] . go sing . getHM
  where
    go :: Sing ns -> HM' ns -> Endo [Double]
    go = \case
      SNil                                      -> \x -> Endo (x:)
      SNat `SCons` SNil                         -> Endo . (++) . LA.toList . extract
      SNat `SCons` (SNat `SCons` SNil)          -> foldMap (Endo . (++)) . LA.toLists . extract
      SNat `SCons` ns@(_ `SCons` (_ `SCons` _)) -> foldMap (go ns)

instance SingI s => MonoFoldable (HM s) where
    ofoldMap f     = foldMap f . helems
    ofoldr f z     = foldr f z . helems
    ofoldl' f z    = foldl' f z . helems
    otoList        = helems
    oall f         = all f . helems
    oany f         = any f . helems
    onull          = (== 0) . olength
    olength _      = fromIntegral (product (fromSing (sing @_ @s)))
    olength64      = fromIntegral . olength
    ocompareLength = ocompareLength . helems
    otraverse_ f   = traverse_ f . helems
    ofor_ x        = for_ (helems x)
    omapM_ f       = traverse_ f . helems
    oforM_ x       = for_ (helems x)
    ofoldlM f x    = foldlM f x . helems
    ofoldMap1Ex f  = ofoldMap1Ex f . helems
    ofoldr1Ex f    = ofoldr1Ex f . helems
    ofoldl1Ex' f   = ofoldl1Ex' f . helems
    headEx         = headEx . helems
    lastEx         = lastEx . helems
    maximumByEx f  = maximumByEx f . helems
    minimumByEx f  = minimumByEx f . helems

deriving instance NFData (HM' s)     => NFData (HM s)
deriving instance Show (HM' s)       => Show (HM s)
deriving instance Num (HM' s)        => Num (HM s)
deriving instance Fractional (HM' s) => Fractional (HM s)
deriving instance Floating (HM' s)   => Floating (HM s)