module Algebra where

import qualified Data.Text as T
import qualified Data.Map as M
import qualified Data.Vector as V
import Data.Maybe
import Statistics.LinearRegression

import Types

selectKeyData :: T.Text -> [Sample ItemsMap] -> [Sample Int]
selectKeyData key samples = [sample {sampleItems = select sample} | sample <- samples]
  where
    select sample = fromMaybe 0 $ M.lookup key (sampleItems sample)

prepareSystem :: [Sample Int] -> (V.Vector Double, V.Vector Double)
prepareSystem samples =
      (V.fromList $ map sampleTime samples,
       V.fromList $ map (fromIntegral . sampleItems) samples)

growCoefficient :: T.Text -> [Sample ItemsMap] -> Double
growCoefficient key samples =
  let (xs, ys) = prepareSystem $ selectKeyData key samples
      (_, beta) = linearRegression xs ys
  in  beta


module AI.Learning.CrossValidation where

import Control.Monad.Random
import Foreign.Storable (Storable)
import Numeric.LinearAlgebra
import qualified Data.List as L

import AI.Util.Matrix
import AI.Util.Util

----------------------
-- Cross Validation --
----------------------

class Indexable c where
    index :: c -> Index -> c
    nobs :: c -> Int

instance Storable a => Indexable (Vector a) where
    index = subRefVec
    nobs = dim

instance Element a => Indexable (Matrix a) where
    index = subRefRows
    nobs = rows

instance Indexable [a] where
    index = map . (!!)
    nobs = length

-- |Indexes are lists of 'Int'. Should refactor this to use something more
--  efficient.
type Index  = [Int]

-- |Type for cross-validation partition.
data CVPartition = CVPartition [(Index, Index)]

-- |Specify what type of cross-validation you want to do.
data CVType = LeaveOneOut
            | KFold Int

-- |Prediction function. A prediction function should take a training and a test
--  set, and use the training set to build a model whose performance is
--  evaluated on the test set, returning a final score as a 'Double'.
type PredFun a b = a        -- Training set predictors
                -> b        -- Training set target
                -> a        -- Test set predictors
                -> b        -- Test set target
                -> Double   -- Performance score

-- |Create a partition into test and training sets.
cvPartition :: RandomGen g => Int -> CVType -> Rand g CVPartition
cvPartition sz cvtype = case cvtype of
    KFold i     -> cvp sz i
    LeaveOneOut -> cvp sz sz

-- |Helper function for 'cvPartition'.
cvp :: RandomGen g => Int -> Int -> Rand g CVPartition
cvp n k = do
    is <- go i (k - i) idx
    return . CVPartition $ map (\i -> (idx L.\\ i, i)) is
    where
        go 0 0 idx = return []

        go 0 j idx = do
            (is, idx') <- selectMany' s idx
            iss        <- go 0 (j-1) idx'
            return (is:iss)

        go i j idx = do
            (is, idx') <- selectMany' (s+1) idx
            iss        <- go (i-1) j idx'
            return (is:iss)

        s   = n `div` k
        i   = n `mod` k
        idx = [0 .. n-1]

-- |Perform k-fold cross-validation. Given a 'CVPartition' containing a list
--  of training and test sets, we repeatedly fit a model on the training set
--  and test its performance on the test set/
kFoldCV_ :: (Indexable a, Indexable b) =>
            CVPartition
         -> PredFun a b
         -> a
         -> b
         -> [Double]
kFoldCV_ (CVPartition partition) predfun x y = map go partition
    where
        go (trainIdx,testIdx) = predfun xTrain yTrain xTest yTest
            where
                xTrain = x `index` trainIdx
                yTrain = y `index` trainIdx
                xTest  = x `index` testIdx
                yTest  = y `index` testIdx

-- |Perform k-fold cross-validation, randomly generating the training and
--  test sets first.
kFoldCV :: (RandomGen g, Indexable a, Indexable b) =>
           CVType           -- What type of cross-validation?
        -> PredFun a b      -- Prediction function
        -> a                -- Predictors
        -> b                -- Targets
        -> Rand g [Double]  -- List of scores
kFoldCV cvtype predfun x y = if nobs x /= nobs y
    then error "Inconsistent dimensions -- KFOLDCV"
    else do
        cp <- cvPartition (nobs x) cvtype
        return (kFoldCV_ cp predfun x y)

---------------
-- Old Stuff --
---------------

-- |Model builder. A model builder takes a training set of regressors and
--  targets, and constructs a function that makes predictions from an out-
--  of-sample set of regressors.
type ModelBuilder = Matrix Double   -- Training set regressors
                 -> Vector Double   -- Training set target
                 -> Matrix Double   -- Out-of-sample regressors
                 -> Vector Double   -- Predictions

-- |Evaluation function. An evaluation function takes a vector of targets and
--  a vector of predictions, and returns a score corresponding to how closely
--  the predictions match the target.
type EvalFun = Vector Double    -- Target
            -> Vector Double    -- Predictions
            -> Double           -- Score (e.g. MSE, MCR, likelihood)


-- Filename: fft.hs
-- Created by: Daniel Winograd-Cort
-- Created on: unknown
-- Last Modified by: Daniel Winograd-Cort
-- Last Modified on: 12/12/2013

-- This module requires the array and pure-fft packages.

{-# LANGUAGE Arrows #-}
module Euterpea.IO.MUI.FFT where
import FRP.UISF
import Control.Arrow.Operations
import Numeric.FFT (fft)
import Data.Complex
import Data.Map (Map)
import qualified Data.Map as Map

-- | Alternative for working with Math.FFT instead of Numeric.FFT
--import qualified Math.FFT as FFT
--import Data.Array.IArray
--import Data.Array.CArray
--myFFT n lst = elems $ (FFT.dft) (listArray (0, n-1) lst)

--------------------------------------
-- Fast Fourier Transform
--------------------------------------

-- | Returns n samples of type b from the input stream at a time,
--   updating after k samples.  This function is good for chunking
--   data and is a critical component to fftA
quantize :: ArrowCircuit a => Int -> Int -> a b (SEvent [b])
quantize n k = proc d -> do
    rec (ds,c) <- delay ([],0) -< (take n (d:ds), c+1)
    returnA -< if c >= n && c `mod` k == 0 then Just ds else Nothing

-- | Converts the vector result of a dft into a map from frequency to magnitude.
--   One common use is:
--      fftA >>> arr (fmap $ presentFFT clockRate)
presentFFT :: Double -> [Double] -> Map Double Double
presentFFT clockRate a = Map.fromList $ zipWith (curry mkAssoc) [0..] a where
    mkAssoc (i,c) = (freq, c) where
        samplesPerPeriod = fromIntegral (length a)
        freq = i * (clockRate / samplesPerPeriod)

-- | Given a quantization frequency (the number of samples between each
--   successive FFT calculation) and a fundamental period, this will decompose
--   the input signal into its constituent frequencies.
--   NOTE: The fundamental period must be a power of two!
fftA :: ArrowCircuit a => Int -> Int -> a Double (SEvent [Double])
fftA qf fp = proc d -> do
    carray <- quantize fp qf -< d :+ 0
    returnA -< fmap (map magnitude . take (fp `div` 2) . fft) carray


{-# LANGUAGE MagicHash, UnboxedTuples #-}

module VecUtil (unpack2, unpack3, tangentBitangent, vec3Of, vec2Of) where

import Numeric.Vector

unpack2 :: Vec2d -> (Double, Double)
unpack2 v = case unpackV2# v of
              (# x, y #) -> (x, y)

unpack3 :: Vec3d -> (Double, Double, Double)
unpack3 v = case unpackV3# v of
              (# x, y, z #) -> (x, y, z)

tangentBitangent :: Vec3d -> (Vec3d, Vec3d)
tangentBitangent normal =
  -- Cross product with an arbitrary vector to produce the tangent. The if is for the rare case that
  -- the normal is the arbitrary vector. Is there a better approach?
  let tangent = normalized $ normal × (if normal == (vec3 1 0 0) then vec3 0 1 0 else vec3 1 0 0)
      bitangent = normal × tangent -- Don't need to normalize, magnitude is mathematically 1
  in (tangent, bitangent)

vec3Of :: Double -> Vec3d
vec3Of = realToFrac

vec2Of :: Double -> Vec2d
vec2Of = realToFrac


{-# LANGUAGE TemplateHaskell #-}
module Turtle (main) where
import Data.Complex
import Ros.Node
import Ros.Topic (cons, repeatM)
import Ros.TopicUtil (tee, filterBy, everyNew, interruptible, gate, share)
import Ros.Turtlesim.Pose
import Ros.Turtlesim.Velocity
import Ros.Logging
import System.IO (hFlush, stdout)

-- A type synonym for a 2D point.
type Point = Complex Float

-- A Topic of user-supplied waypoint trajectories.
getTraj :: Topic IO [Point]
getTraj = repeatM (do putStr "Enter waypoints: " >> hFlush stdout
                      $(logInfo "Waiting for new traj")
                      (map (uncurry (:+)) . read) `fmap` getLine)

wrapAngle theta
  | theta < 0    = theta + 2 * pi
  | theta > 2*pi = theta - 2 * pi
  | otherwise    = theta

angleDiff x y = wrapAngle (x' + pi - y') - pi
  where x' = if x < 0 then x + 2*pi else x
        y' = if y < 0 then y + 2*pi else y

-- Produce a unit value every time a goal is reached.
arrivalTrigger :: Topic IO Point -> Topic IO Pose -> Topic IO ()
arrivalTrigger goals poses = fmap (const ()) $
                             filterBy (fmap arrived goals) (fmap p2v poses)
  where arrived goal pose = magnitude (goal - pose) < 1.5
        p2v (Pose x y _ _ _) = x :+ y

-- Navigate to a goal given a current pose estimate.
navigate :: (Point, Pose) -> Velocity
navigate (goal, pos) = Velocity (min 2 (magnitude v)) angVel
  where v        = goal - (x pos :+ y pos)
        thetaErr = (angleDiff (phase v) (theta pos)) * (180 / pi)
        angVel   = signum thetaErr * (min 2 (abs thetaErr))

main = runNode "HaskellBTurtle" $
       do enableLogging (Just Warn)
          poses <- subscribe "/turtle1/pose"
          (t1,t2) <- liftIO . tee . interruptible $ getTraj
          let goals = gate t1 (cons () (arrivalTrigger t2 poses))
          advertise "/turtle1/command_velocity" $
                    fmap navigate (everyNew goals poses)


{-|
    Module      : Vocoder
    Description : Phase vocoder
    Copyright   : (c) Celina Pawlińska, 2020
                      Marek Materzok, 2021
    License     : BSD2

This module implements the phase vocoder algorithms.
The implementation is designed to be used directly or to be integrated
into some convenient abstraction (streaming or FRP).
-}
module Vocoder (
      Moduli,
      Phase,
      PhaseInc,
      Frame,
      Window,
      HopSize,
      Length,
      STFTFrame,
      FFTOutput,
      VocoderParams,
      vocoderParams,
      vocFrameLength,
      vocInputFrameLength,
      vocFreqFrameLength,
      vocHopSize,
      vocWindow,
      doFFT,
      doIFFT,
      analysisBlock,
      analysisStep,
      analysisStage,
      synthesisBlock,
      synthesisStep,
      synthesisStage,
      zeroPhase,
      volumeCoeff,
      frameFromComplex,
      frameToComplex,
      addFrames
    ) where

import Data.List
import Data.Complex
import Data.Fixed
import Data.Tuple
import Control.Arrow
import Numeric.FFT.Vector.Invertible as FFT
import Numeric.FFT.Vector.Plan as FFTp
import qualified Data.Vector.Storable as V

-- | Complex moduli of FFT frames. Represent signal amplitudes.
type Moduli = V.Vector Double

-- | Complex arguments of FFT frames. Represent signal phases.
type Phase = V.Vector Double

-- | Phase increments. Represent the deviation of the phase difference
-- between successive frames from the expected difference for the center
-- frequencies of the FFT bins.
type PhaseInc = V.Vector Double

-- | Time domain frame.
type Frame = V.Vector Double

-- | Sampled STFT window function.
type Window = Frame

-- | Offset between successive STFT frames, in samples.
type HopSize = Int

-- | Size in samples.
type Length = Int

-- | STFT processing unit.
type STFTFrame = (Moduli, PhaseInc)

-- | Frequency domain frame.
type FFTOutput = V.Vector (Complex Double)

-- | Type of FFT plans for real signals.
type FFTPlan = FFTp.Plan Double (Complex Double)

-- | Type of IFFT plans for real signals.
type IFFTPlan = FFTp.Plan (Complex Double) Double

-- | Configuration parameters for the phase vocoder algorithm.
data VocoderParams = VocoderParams{
    -- | FFT plan used in analysis stage.
    vocFFTPlan  :: FFTPlan,
    -- | FFT plan used in synthesis stage.
    vocIFFTPlan :: IFFTPlan,
    -- | STFT hop size.
    vocHopSize :: HopSize,
    -- | Window function used during analysis and synthesis.
    vocWindow :: Window
    -- TODO thread safety?
}

-- | FFT frequency frame length.
vocFreqFrameLength :: VocoderParams -> Length
vocFreqFrameLength par = planOutputSize $ vocFFTPlan par

-- | FFT frame length. Can be larger than `vocInputFrameLength` for zero-padding.
vocFrameLength :: VocoderParams -> Length
vocFrameLength par = planInputSize $ vocFFTPlan par

-- | STFT frame length.
vocInputFrameLength :: VocoderParams -> Length
vocInputFrameLength par = V.length $ vocWindow par

-- | Create a vocoder configuration.
vocoderParams :: Length -> HopSize -> Window -> VocoderParams
vocoderParams len hs wnd = VocoderParams (plan dftR2C len) (plan dftC2R len) hs wnd

-- | Apply a window function on a time domain frame.
applyWindow :: Window -> Frame -> Frame
applyWindow = V.zipWith (*)

-- | Change the vector indexing so that the sample at the middle has the number 0.
-- This is done so that the FFT of the window has zero phase, and therefore does not
-- introduce phase shifts in the signal.
rewind :: (V.Storable a) => V.Vector a -> V.Vector a
rewind vec = uncurry (V.++) $ swap $ V.splitAt (V.length vec `div` 2) vec

-- | Zero-pad the signal symmetrically from both sides.
addZeroPadding :: Length
    -> Frame
    -> Frame
addZeroPadding len v
    | diff < 0  = error $ "addZeroPadding: input is " ++ (show diff) ++ " samples longer than target length"
    | diff == 0 = v
    | otherwise = res
    where
    l = V.length v
    diff = len - l
    halfdiff = diff - (diff `div` 2)
    res = (V.++) ((V.++) (V.replicate halfdiff 0) v) (V.replicate (diff-halfdiff) 0)

-- | Perform FFT processing, which includes the actual FFT, rewinding, zero-paddding
-- and windowing.
doFFT :: VocoderParams -> Frame -> FFTOutput
doFFT par =
    FFT.execute (vocFFTPlan par) . rewind . addZeroPadding (vocFrameLength par) . applyWindow (vocWindow par)

-- | Perform analysis on a sequence of frames. This consists of FFT processing
-- and performing analysis on frequency domain frames.
analysisStage :: Traversable t => VocoderParams -> Phase -> t Frame -> (Phase, t STFTFrame)
analysisStage par ph = mapAccumL (analysisBlock par) ph

-- | Perform FFT transform and frequency-domain analysis.
analysisBlock :: VocoderParams -> Phase -> Frame -> (Phase, STFTFrame)
analysisBlock par prev_ph vec = analysisStep (vocHopSize par) (vocFrameLength par) prev_ph (doFFT par vec)

-- | Analyze a frequency domain frame. Phase from a previous frame must be supplied.
-- It returns the phase of the analyzed frame and the result.
analysisStep :: HopSize -> Length -> Phase -> FFTOutput -> (Phase, STFTFrame)
analysisStep h eN prev_ph vec =
    (ph,(mag,ph_inc))
    where
    (mag, ph) = frameFromComplex vec
    ph_inc = V.imap (calcPhaseInc eN h) $ V.zipWith (-) ph prev_ph

-- | Wraps an angle (in radians) to the range [-pi : pi].
wrap :: Double -> Double
wrap e = (e+pi) `mod'` (2*pi) - pi

calcPhaseInc :: Length -> HopSize -> Int -> Double -> Double
calcPhaseInc eN hop k ph_diff =
    (omega + wrap (ph_diff - omega)) / fromIntegral hop
    where
    omega = (2*pi*fromIntegral k*fromIntegral hop) / fromIntegral eN

-- | Perform synthesis on a sequence of frames. This consists of performing
-- synthesis and IFFT processing.
synthesisStage :: Traversable t => VocoderParams -> Phase -> t STFTFrame -> (Phase, t Frame)
synthesisStage par ph frs = mapAccumL (synthesisBlock par) ph frs

-- | Perform frequency-domain synthesis and IFFT transform.
synthesisBlock :: VocoderParams -> Phase -> STFTFrame -> (Phase, Frame)
synthesisBlock par ph fr = (id *** doIFFT par) $ synthesisStep (vocHopSize par) ph fr

-- | Synthesize a frequency domain frame. Phase from the previously synthesized frame
-- must be supplied. It returns the phase of the synthesized frame and the result.
synthesisStep :: HopSize -> Phase -> STFTFrame -> (Phase, FFTOutput)
synthesisStep hop ph (mag, ph_inc) =
    (new_ph, frameToComplex (mag, new_ph))
    where
    new_ph = V.zipWith (+) ph $ V.map (* fromIntegral hop) ph_inc

-- | Perform IFFT processing, which includes the actual IFFT, rewinding, removing padding
-- and windowing.
doIFFT :: VocoderParams -> FFTOutput -> Frame
doIFFT par =
    applyWindow (vocWindow par) . cutCenter (vocInputFrameLength par) . rewind . FFT.execute (vocIFFTPlan par)

-- | Cut the center of a time domain frame, discarding zero padding.
cutCenter :: (V.Storable a) => Length -> V.Vector a -> V.Vector a
cutCenter len vec = V.take len $ V.drop ((V.length vec - len) `div` 2) vec

-- | Zero phase for a given vocoder configuration.
-- Can be used to initialize the synthesis stage.
zeroPhase :: VocoderParams -> Phase
zeroPhase par = V.replicate (vocFreqFrameLength par) 0

-- | An amplitude change coefficient for the processing pipeline.
-- Can be used to ensure that the output has the same volume as the input.
volumeCoeff :: VocoderParams -> Double
volumeCoeff par = fromIntegral (vocHopSize par) / V.sum (V.map (**2) $ vocWindow par)

-- | Converts frame representation to complex numbers.
frameToComplex :: STFTFrame -> FFTOutput
frameToComplex = uncurry $ V.zipWith mkPolar

-- | Converts frame representation to magnitude and phase.
frameFromComplex :: FFTOutput -> STFTFrame
frameFromComplex = V.map magnitude &&& V.map phase

-- | Adds STFT frames.
addFrames :: STFTFrame -> STFTFrame -> STFTFrame
addFrames f1 f2 = frameFromComplex $ V.zipWith (+) (frameToComplex f1) (frameToComplex f2)


{-# LANGUAGE FlexibleContexts          #-}
{-# LANGUAGE FlexibleInstances         #-}
{-# LANGUAGE GADTs                     #-}
{-# LANGUAGE MultiParamTypeClasses     #-}
{-# LANGUAGE NoMonomorphismRestriction #-}
{-# LANGUAGE OverloadedStrings         #-}
{-# LANGUAGE PolyKinds                 #-}
{-# LANGUAGE RankNTypes                #-}
{-# LANGUAGE ScopedTypeVariables       #-}
{-# OPTIONS_GHC -fwarn-incomplete-patterns #-}

module Math.HMatrixUtils where

import qualified Knit.Effect.Logger            as KL
import qualified Knit.Report                   as K
import qualified Data.Text                     as T

import qualified Numeric.LinearAlgebra         as LA
import           Numeric.LinearAlgebra.Data     ( Matrix
                                                , R
                                                , Vector
                                                )

textSize :: (LA.Container c e, Show (LA.IndexOf c)) => c e -> T.Text
textSize = T.pack . show . LA.size

checkEqualVectors
  :: K.Member (KL.Logger KL.LogEntry) effs
  => T.Text
  -> T.Text
  -> Vector R
  -> Vector R
  -> K.Sem effs ()
checkEqualVectors nA nB vA vB = if LA.size vA == LA.size vB
  then return ()
  else
    KL.logLE KL.Error
    $  "Unequal vector length. length("
    <> nA
    <> ")="
    <> textSize vA
    <> " and length("
    <> nB
    <> ")="
    <> textSize vB

checkMatrixVector
  :: K.Member (KL.Logger KL.LogEntry) effs
  => T.Text
  -> T.Text
  -> Matrix R
  -> Vector R
  -> K.Sem effs ()
checkMatrixVector nA nB mA vB = if snd (LA.size mA) == LA.size vB
  then return ()
  else
    KL.logLE KL.Error
    $  "Bad matrix * vector lengths. dim("
    <> nA
    <> ")="
    <> textSize mA
    <> " and length("
    <> nB
    <> ")="
    <> textSize vB

checkVectorMatrix
  :: K.Member (KL.Logger KL.LogEntry) effs
  => T.Text
  -> T.Text
  -> Vector R
  -> Matrix R
  -> K.Sem effs ()
checkVectorMatrix nA nB vA mB = if LA.size vA == fst (LA.size mB)
  then return ()
  else
    KL.logLE KL.Error
    $  "Bad vector * matrix lengths. length("
    <> nA
    <> ")="
    <> textSize vA
    <> " and dim("
    <> nB
    <> ")="
    <> textSize mB


{-# OPTIONS_GHC -fno-warn-dodgy-exports #-}
{- | an incomplete extension of dimensional-tf to work with hmatrix and ad
Haddocks are organized to follow hmatrix

note: for subscripting, use the 'HNat' while the units still use 'N.NumType'

TODO:

*  Friendly syntax for introduction (more matD)

*  Friendly syntax for elimination (matD as pattern)

*  A pretty-printer that multiplies out the dimensions at each place?
   The current show instance makes you mentally multiply out row and column units
   (which helps you to see the big picture, but may be more work in other cases)

*  check that all types that could could be inferred are

*  default columns/rows to dimensionless?

*  missing operations (check export list comments)
-}
module DimMat.Internal  where

import Foreign.Storable (Storable)
import GHC.Exts (Constraint)

-- import Data.Type.Equality (type (==))
import qualified Prelude as P
import Prelude (Double)
import Numeric.Units.Dimensional hiding (Mul, Div)
import Numeric.Units.Dimensional.Prelude hiding (Mul, Div)
import qualified Numeric.Units.Dimensional as D

import qualified Numeric.LinearAlgebra as H
import qualified Numeric.LinearAlgebra.LAPACK as H

import Text.PrettyPrint.ANSI.Leijen
import Data.List (transpose)

import Data.HList.CommonMain hiding (MapFst)

-- * Replacements for Dimensional classes

-- | a version of Numeric.Units.Dimensional.'D.Mul' which
-- requires the arguments to include the 'D.Dim' type constructor
class (D.Mul a b c) => Mul a b c

instance (D.Mul a b c,
       a ~ Dim l m t i th n j,
       b ~ Dim l' m' t' i' th' n' j',
       c ~ Dim l'' m'' t'' i'' th'' n'' j'') =>
       Mul a b c

-- | a version of Numeric.Units.Dimensional.'D.Div' which
-- requires the arguments to include the 'D.Dim' type constructor
class (D.Div a b c) => Div a b c

instance (D.Div a b c,
       a ~ Dim l m t i th n j,
       b ~ Dim l' m' t' i' th' n' j',
       c ~ Dim l'' m'' t'' i'' th'' n'' j'') =>
       Div a b c

-- * AD
{- $ad
TODO: gradients, hessians, etc.

Types for derivative towers can see hlist's @HList\/Data\/HList\/broken\/Lazy.hs@,
but laziness doesn't really make much sense if the @take@ that is eventually used
to get a finite list for printing etc.

Complications include the fact that AD.grad needs a traversable,
but hmatrix stuff is not traversable (due needing Storable). In ipopt-hs
I got around this problem by copying data. Perhaps that is the solution?

> BROKEN
> grad :: (Num a, AreRecips i iinv, H.Element a, Storable a,
>           MapMultEq o iinv r) =>
>         (forall s. (AD.Mode s, H.Container H.Vector (AD.AD s a),
>                     Storable (AD.AD s a), H.Field (AD.AD s a))
>                 => DimMat '[i] (AD.AD s a)
>                 -> Quantity o (AD.AD s a))
>      -> DimMat '[i] a
>      -> DimMat '[r] a
> grad f (DimVec x) = DimMat (H.fromLists [AD.grad (unQty . f . DimVec . H.fromList) (H.toList x)])
>     where unQty (Dimensional a) = a
-}

#ifdef WITH_AD
{- |
>>> let ke velocity = velocity*velocity*(1*~kilo gram)
>>> diff ke (3 *~ (metre/second))
6.0 m^-1 kg^-1 s
-}

diff :: (Num a) =>
        (forall s. AD.Mode s => Dimensional v x (AD.AD s a)
                             -> Dimensional v y (AD.AD s a))
        -> Dimensional v x a -> Dimensional v (Div x y) a
diff f z = Dimensional $ AD.diff (unD . f . Dimensional) (unD z)
    where unD (Dimensional a) = a
#endif

-- * GADT for linear algebra with units

{- | Generalization of 'Dimensional' to matrices and vectors. Units
in each coordinate are known at compile-time. This wraps up HMatrix.

[@ 'DMat' @] the units at coordinate ij are @(r1 ': r)_i (DOne ': c)_j@

[@ 'DVec' @] the units at coordinate i are @(r1 ': r)_i@

[@ DScal @] is the same as Dimensional
-}
data D (sh :: ( *, [[ * ]])) e where
  DMat :: (H.Container H.Matrix e, H.Field e)
    => H.Matrix e -> D '(r1,[r, c]) e
  DVec :: (H.Container H.Vector e, H.Field e)
    => H.Vector e -> D '(r1, '[r]) e
  DScal :: (H.Field e) => e -> D '(r1,'[]) e

instance (Show a, PPUnits sh) => Pretty (D sh a) where
    pretty (DVec v) = case ppUnits (Proxy :: Proxy sh) of
        [rs] -> vcat
             [ dullgreen (string label) <+> string (show e)
                | (e,label) <- H.toList v `zip` pad rs ]
    pretty (DMat m) = case ppUnits (Proxy :: Proxy sh) of
        [rs,cs] ->
            vcat $
            map (hsep . onHead dullgreen) $
            transpose $ map (onHead dullgreen . map string . pad) $
            zipWith (\a b -> a:b)
                ((show (H.rows m) ++ "><"++ show (H.cols m)) : cs) $
            transpose $
            zipWith (\r e -> r : map show e) rs (H.toLists m)
        where
            onHead f (x:xs) = f x : xs
            onHead _ [] = []

instance Pretty (D sh a) => Show (D sh a) where
    showsPrec p x = showsPrec p (pretty x)

-- ** pretty instance
pad :: [String] -> [String]
pad [] = []
pad xs = let
    w = maximum (map length xs)
    in map (\x -> take w $ x ++ replicate w ' ') xs

class PPUnits (sh :: k) where
    ppUnits :: Proxy sh -> [[String]]

instance forall (r1 :: *) (r::[*]) (c :: [*]) l m t i th n j.
      (PPUnits [r1 ': r, c], Show (Quantity r1 Int), PPUnits' c, PPUnits' r,
       r1 ~ Dim l m t i th n j) => PPUnits '(r1, [r,c]) where
    ppUnits _ = ppUnits (Proxy :: Proxy [r1 ': r, DOne ': c])

instance (PPUnits' x, PPUnits xs) => PPUnits (x ': xs) where
    ppUnits _ = ppUnits' (Proxy :: Proxy x) : ppUnits (Proxy :: Proxy xs)
instance PPUnits '[] where
    ppUnits _ = []

class PPUnits' (sh :: [ * ]) where
    ppUnits' :: Proxy sh -> [String]
instance (PPUnits' xs) => PPUnits' (DOne ': xs) where
    ppUnits' _ = "1" : ppUnits' (Proxy :: Proxy xs)
instance (ShowDimSpec x, PPUnits' xs) => PPUnits' (x ': xs) where
    ppUnits' _ = showDimSpec (Proxy :: Proxy x) : ppUnits' (Proxy :: Proxy xs)
instance PPUnits' '[] where
    ppUnits' _ = []

class ShowDimSpec a where
    showDimSpec :: Proxy a -> String

instance (Show (Quantity a Int), Dim l m t i th n j ~ a) => ShowDimSpec a where
    showDimSpec _ = case drop 2 $ show (1 *~ (Dimensional 1 :: Unit a Int)) of
          "" -> "1"
          x -> x

-- * Constraints
{- $justification
A major theme in this library is that type inference goes in whichever direction
it can: in ordinary haskell it is very common for argument types to be determined
by the result types. For example see any code that uses 'Num' or 'Read'.

When we use type families, things look more convenient:

@
data Nat = S Nat | Z

type family Add (a :: Nat) (b :: Nat) :: Nat
type instance Add Z b = b
type instance Add (S a) b = Add a (S b)
@

But ghc is unable to deduce things like @a ~ Z@ given evidence such as @Add Z a ~ Z@.
One way around this is to use @ConstraintKinds@:

@
type AddT (a :: Nat) (b :: Nat) (c :: Nat)
          = (Add a b ~ c, Sub c a ~ b, Sub c b ~ a)
@

Which leads to functions like  @f :: AddT a b ab => ... @. This is bad for a couple reasons:

* the right-hand-side of the type can only mention type variables on the
  left-hand-side.
* the left hand side can only bind type variables. Working around this
  leads to many auxiliary type families such as Fst, Snd and Head, or
  leads to a @type family AddT@

So below many constraints expressed as classes, since they have less
of those limitations.

-}

-- | @a*b = c@ when any are lists
class MultEq (a :: k1) (b :: k2) (c :: k3)

-- instance (Zip3 MultEq aas bbs ccs) => MultEq aas bbs ccs
instance Zip3 Mul as bs cs => MultEq as bs cs
instance Zip1 Mul as b  c  => MultEq as b  c
instance Zip1 Mul bs a  c  => MultEq a  bs c
instance (Zip1 Mul cs aInv b,
          Mul a aInv DOne) => MultEq a  b  cs

instance (SameLength as bs,
        Zip2 Mul as bs c)  => MultEq as bs c
instance (SameLength as cs,
          Zip2 Div as cs bInv,
          Mul b bInv DOne) => MultEq as b  cs
instance (Zip2 Div bs cs aInv,
          SameLength bs cs,
          Mul a aInv DOne) => MultEq a  bs cs
instance Mul a b c => MultEq a b c

-- ** Zip

class (SameLength a b, SameLength b c) =>
    Zip3
      (op :: k -> k -> k -> Constraint)
      (a :: [k])
      (b :: [k])
      (c :: [k])

instance (SameLength aas bbs,
          SameLength ccs bbs,
          op a b c,
          (a ': as) ~ aas,
          (b ': bs) ~ bbs,
          (c ': cs) ~ ccs,
          Zip3 op as bs cs) => Zip3 op aas bbs ccs
instance Zip3 op '[] '[] '[]

class (SameLength a b) =>
    Zip2
      (op :: k -> k -> k -> Constraint)
      (a :: [k])
      (b :: [k])
      (c ::  k)

instance (SameLength aas bbs,
      op a b c,
      (a ': as) ~ aas,
      (b ': bs) ~ bbs,
      Zip2 op as bs c) => Zip2 op aas bbs c

instance Zip2 op '[] '[] c

class Zip1
      (op :: k -> k -> k -> Constraint)
      (a :: [k])
      (b ::  k)
      (c ::  k)

instance ((a ': as) ~ aas,
    op a b c,
    Zip1 op as b c) => Zip1 op aas b c

instance Zip1 op '[] b c

{- | given @ijs :: [[Quantity a]]@ (except the : and [] constructors are
actually (,) and (), ie. a HList that doesn't use the HList constructors),
calculate a @DimMat rowUnits colUnits@, where the outer product of rowUnits
and colUnits gives the units at each index in the ijs.  The first element
of colUnits is DOne.
-}
class (SameLength a ab) => Outer a b ab
instance Outer '[] b '[]
instance (SameLength aas ccs,
          (a ': as) ~ aas,
          (c ': cs) ~ ccs,
          MultEq a b c,
          Outer as b cs)
  => Outer aas b ccs

-- * DimMatFromTuple
class DimMatFromTuple ijs r1 r c e

type family TupleToHListU (a :: *) :: [*]
type instance TupleToHListU (a, b) = () ': TupleToHListU b
type instance TupleToHListU () = '[]

type family TuplesToHListU (a :: *) :: [[*]]
type instance TuplesToHListU (a, b) = TupleToHListU a ': TuplesToHListU b
type instance TuplesToHListU () = '[]

instance (Outer (r1 ': r) (DOne ': c) ijs',
      DMFromTuple1 e ijs ijs',
      SameLength (TuplesToHListU ijs) ijs') => DimMatFromTuple ijs r1 r c e

-- | helper for 'DimMatFromTuple'
type family DMFromTuple1 e b (b' :: [[*]]) :: Constraint
type family DMFromTuple2 e b (b' :: [*]) :: Constraint
type family DMFromTuple3 e b b' :: Constraint

type instance DMFromTuple3 e (Quantity b e') b' = (e ~ e', b ~ b')
type instance DMFromTuple1 e (x, xs) (x' ': xs') = (TupleToHListU x `SameLength` x',
                                                    DMFromTuple2 e x x', DMFromTuple1 e xs xs')
type instance DMFromTuple1 e () xs = (xs ~ '[])
type instance DMFromTuple2 e (x, xs) (x' ': xs') = (DMFromTuple3 e x x', DMFromTuple2 e xs xs')
type instance DMFromTuple2 e () xs = (xs ~ '[])

-- | just for types produced by the matD quasiquote
toDM :: DimMatFromTuple ijs r1 r c e => ijs -> Proxy (D '(r1, [r, c]) e)
toDM _ = Proxy

-- | @InnerCxt t a b = t ~ 'H.dot' a b@
type family InnerCxt (t :: k) (a :: [k]) (b :: [k]) :: Constraint
type instance InnerCxt t (a ': as) (b ': bs) = (MultEq a b t, InnerCxt t as bs)
type instance InnerCxt t '[] '[] = ()

class (SameLength a b, InnerCxt c a b) => Inner (a :: [*]) (b :: [*]) (c :: *)

instance (SameLength aas bbs, InnerCxt c aas bbs) => Inner aas bbs c

-- | @ProdEq a b@ is @product a ~ b@
class ProdEq a b
instance (ProdEq as b', Mul a b' b) => ProdEq (a ': as) b
instance (dOne ~ DOne) => ProdEq '[] dOne

-- | @RecipEq a aInv@ is @a*aInv ~ DOne@ (or a list of DOne)
class RecipEq (a :: k) (aInv :: k)
instance (MultEq as aInvs DOne) => RecipEq as aInvs

-- | @AtEq a n b m c@ calculates @(At a n \`Mult\` At b m) ~ c@,
-- but also can infer part of the `a` if the `b` and `c` are known
type family AtEq2 (a :: [k]) (n :: HNat) (b :: [k]) (m :: HNat) (c :: k) :: Constraint
type instance AtEq2  (a ': as) HZero (b ': bs) HZero c = (MultEq a b c)
type instance AtEq2  (a ': as) (HSucc n) bs m c = AtEq2 as n bs m c
type instance AtEq2  as HZero (b ': bs) (HSucc m) c = AtEq2 as HZero bs m c

type family AtEq (a :: [k]) (n :: HNat) (b :: k) :: Constraint
type instance AtEq (a ': as) HZero b = (a ~ b)
type instance AtEq (a ': as) (HSucc n) b = AtEq as n b

-- | Data.Packed.Vector.'H.@>'
(@>) :: (HNat2Integral i, AtEq r i ri, MultEq r1 ri u)
    => D '(r1,'[r]) a
    -> Proxy i
    -> Quantity u a
DVec v @> i = Dimensional (v H.@> hNat2Integral i)

-- | Data.Packed.Matrix.'H.@@>'
(@@>) :: (HNat2Integral i, HNat2Integral j, AtEq2 (r1 ': r) i (DOne ': c) j ty)
    => D '(r1, [r,c]) a
    -> (Proxy i, Proxy j)
    -> Quantity ty a
DMat m @@> (i,j) = Dimensional (m H.@@> (hNat2Integral i,hNat2Integral j))

pnorm :: (AllEq r1 r, AllEq DOne c)
         => H.NormType -> D '(r1, [r, c]) a -> Quantity r1 (H.RealOf a)
pnorm normType (DMat a) = Dimensional (H.pnorm normType a)

{- | @AllEq a xs@ is like @all (all (a==)) xs@, @all (a ==) xs@, @a == xs@:
whichever amount of [ ] is peeled off before making the comparison (with ~)
-}
class AllEq (a :: k1) (xs :: k2)

instance (a ~ x, AllEq a xs) => AllEq a (x ': xs)
instance AllEq a '[]
instance AllEq '[] xs
instance (AllEq a xs, AllEq as xs) => AllEq (a ': as) xs

{- | @c = a `dot` b@ is one of:

> c_ij = sum_j a_ij b_jk
> c_k  = sum_j a_j  b_jk
> c_i  = sum_j a_ij b_j
> c    = sum_j a_j  b_j

-}
class Dot a b c {- | a b -> c -} where
    dot :: H.Element e => D a e -> D b e -> D c e

instance
    ( MultEq (a ': ra) b (c ': rc),
      MultEq ca rb b,
      shA ~ '(a,[ra, ca]),
      shB ~ '(b, rb ': cb),
      shC ~ '(c, rc ': cb))
    => Dot shA shB shC where
    dot (DMat a) (DMat b) = DMat (H.multiply a b)
    dot (DMat a) (DVec b) = DVec (H.mXv a b)
    {-
    dot (DVec a) (DMat b) = DVec (H.vXm a b)
    dot (DVec a) (DVec b) = DScal (H.dot a b)
    -}

class DotS s a b

instance (rest' ~ rest, MultEq s (a ': as) (b ': bs))
    => DotS s '(a, as ': rest) '(b, bs ': rest')

class Trans a b where
    trans :: D a e -> D b e

instance (a ~ '(r1, [x,y]),
          b ~ '(r1, [y,x]))
    => Trans a b where
    trans (DMat a) = DMat (H.trans a)

{- | type for a pseudoinverse (and inverse):

The single instance comes from looking at inverses from a 2x2 matrix (let's call A):

> a b
> c d

and the inverse * determinant of the original

>  d  -b
> -c   a

In the product A * A^-1 the diagonal is dimensionless ('DOne').

That happens if the row and column type-level unit lists are reciprocals of
eachother ('AreRecips'), so the constraint on the instance of PInv encodes
this exactly (plus some constraints requiring that sh and sh' are at least
1x1)
-}
class PInv a b where
  pinv :: D a e -> D b e

type family LengthSndTwo (a :: k) :: Constraint
type instance LengthSndTwo '(a, as) = SameLength as '[(), ()]

type AreRecips a b = MultEq a b DOne
instance (MultEq ra cb a,
          MultEq ca rb b,
          AreRecips a b,
          bSh ~ '(b, [rb, cb]),
          aSh ~ '(a, [ra, ca]))
    => PInv aSh bSh where
  pinv (DMat a) = DMat (H.pinv a)

inv :: (PInv a b,
        b ~ '(t1, [t2, t3]))
  => D a e -> D b e
inv (DMat a) = DMat (H.inv a)

pinvTol :: (PInv a b,
           b ~ '(t1, [t2, t3]) )
  => Double -> D a e -> D b e
pinvTol tol (DMat a) = DMat (H.pinvTol tol a)

class Det a b where
    det :: D a e -> Quantity b e

instance (SameLength r c,
       ProdEq (r1 ': r) (pr :: *),
       ProdEq c (pc :: *),
       MultEq pr pc b,
       a ~ '(r1, [r, c])) =>
    Det a b where
  det (DMat a) = Dimensional (H.det a)

{- | Numeric.LinearAlgebra.Algorithms.'H.expm'

@y t = expm (scale t a) \`multiply\` y0@ solves the DE @y' = Ay@ where y0 is the
value of y at time 0

-}
expm :: (AreRecips r c)
    => D '(r1, [r, c]) a
    -> D '(r1, [r, c])  a
expm (DMat a) = DMat (H.expm a)

{- | Numeric.Container.'H.scale'

-}
class Scale a b c where
  scale :: Quantity a e -> D b e -> D c e

fromQty :: H.Field e => Quantity a e -> D '(a, '[]) e
fromQty (Dimensional a) = DScal a

toQty :: D '(a, '[]) e -> Quantity a e
toQty (DScal a) = Dimensional a

instance (MultEq a (r1 ': r) (r1' ': r'),
      b ~ '(r1, r ': rs),
      c ~ '(r1', r' ': rs)) =>
  Scale a b c where
  scale (Dimensional t) (DMat a) = DMat (H.scale t a)
  scale (Dimensional t) (DVec a) = DVec (H.scale t a)

{- | Numeric.Container.'H.scaleRecip'
-}
class ScaleRecip a b c where
  scaleRecip :: D '(a, '[]) e -> D b e -> D c e

class ScaleRecip1 (bool :: Bool) a b c where
  scaleRecip1 :: Proxy bool -> D '(a, '[]) e -> D b e -> D c e

instance
  (ScaleRecipCxt r1 r1' r r' rs rs' a b c
   , rs' ~ '[ t1 ]
   ) => ScaleRecip1 True a b c where
  scaleRecip1 _ (DScal t) (DMat a) = DMat (H.scaleRecip t a)

instance
  (ScaleRecipCxt r1 r1' r r' rs rs' a b c,
   rs' ~ '[]) =>
  ScaleRecip1 False a b c where
  scaleRecip1 _ (DScal t) (DVec a) = DVec (H.scaleRecip t a)

instance (HEq (HLength bs) (HSucc (HSucc HZero)) bool1,
          HEq (HLength cs) (HSucc (HSucc HZero)) bool1,
    ScaleRecip1 bool1 a '(b, bs) '(c, cs)) => ScaleRecip a '(b, bs) '(c, cs) where
  scaleRecip = scaleRecip1 (Proxy :: Proxy bool1)

type ScaleRecipCxt (r1 :: *) (r1' :: *) r r' rs rs' (a :: *) b c =
  (MultEq a (r1' ': r') (r1 ': r) ,
   MultEq rs rs' DOne,
   b ~ '(r1, r ': rs),
   c ~ '(r1', r' ': rs'))

-- | a shortcut for @scaleRecip (DScal 1)@
recipMat :: forall b c e. (H.Field e, ScaleRecip DOne b c) => D b e -> D c e
recipMat m = scaleRecip (DScal 1 :: D '(DOne, '[]) e) m

liftH2 ::
  ( forall h f. (H.Container f e, h ~ f e) => h -> h -> h) ->
    D a e -> D a e -> D a e
liftH2 f (DMat a) (DMat b) = DMat (f a b)
liftH2 f (DVec a) (DVec b) = DVec (f a b)

add a b = liftH2 H.add a b
sub a b = liftH2 H.sub a b

mulMat :: ( MultEq as bs cs, MultEq a b c, cs ~ '[t1 , t2] )
  => D '(a,as) e -> D '(b,bs) e -> D '(c,cs) e
mulMat (DMat a) (DMat b) = DMat (H.mul a b)

mulVec :: ( MultEq as bs cs, MultEq a b c, cs ~ '[t1] )
  => D '(a,as) e -> D '(b,bs) e -> D '(c,cs) e
mulVec (DVec a) (DVec b) = DVec (H.mul a b)

divideMat :: ( MultEq as cs bs, MultEq a c b, cs ~ '[t1 , t2] )
  => D '(a,as) e -> D '(b,bs) e -> D '(c,cs) e
divideMat (DMat a) (DMat b) = DMat (H.divide a b)

divideVec :: ( MultEq as cs bs, MultEq a c b, cs ~ '[t1] )
  => D '(a,as) e -> D '(b,bs) e -> D '(c,cs) e
divideVec (DVec a) (DVec b) = DVec (H.divide a b)

arctan2 :: (bs ~ MapMapConst DOne as) => D '(a,as) e -> D '(a,as) e -> D '(b,bs) e
arctan2 (DMat a) (DMat b) = DMat (H.arctan2 a b)
arctan2 (DVec a) (DVec b) = DVec (H.arctan2 a b)

equal :: D a e -> D a e -> Bool
equal (DMat a) (DMat b) = H.equal a b
equal (DVec a) (DVec b) = H.equal a b

{- | @cmap f m@ gives a matrix @m'@

@f@ is applied to

-}
class CMap f a b e e' where
    cmap :: f -> D a e -> D b e'

-- | Maybe there's a way to implement in terms of the real cmap (possibly
-- unsafeCoerce?)
instance
    (ToHLists sh e xs,
     FromHLists xs' sh' e',
     SameLength xs xs',
     HMapCxt HList (HMap f) xs xs') =>
    CMap f sh sh' e e' where
  cmap f m = fromHLists (HMap f `hMap` (toHLists m :: HList xs) :: HList xs')

type family AppendEq' (a :: [k]) (b :: [k]) (ab :: [k]) :: Constraint
type instance AppendEq' (a ': as) b (a' ': abs) = (a ~ a', AppendEq' as b abs)
type instance AppendEq' '[] b abs = (b ~ abs)

-- | a bit overkill?
--  @a ++ b = ab@
type AppendEq a b ab =
   (ab ~ HAppendR a b,
    AppendEq' a b ab,
    SameLength (DropPrefix a ab) b,
    SameLength (DropPrefix b ab) a)

type instance HAppendR (x ': xs) ys = x ': HAppendR xs ys
type instance HAppendR '[] ys = ys

type family DropPrefix (a :: [k]) (ab :: [k]) :: [k]
type instance DropPrefix (a ': as) (a' ': abs) = DropPrefix as abs
type instance DropPrefix '[] bs = bs

{- | the slightly involved type here exists because
ci1 and ci2 both start with DOne, but ci2's contribution
to ci3 does not need a DOne at the start. Another way to
read the constraints here is:

> map (*rem) (a11 : ri) = b11 : bi
> ci3 = ci1 ++ map (*rem) ci2

The same idea happens with vconcat.
-}
hconcat ::
    ( MultEq (rem :: *) a b,
      MultEq rem ra rb,
      MultEq rem (DOne ': cb) cb',
      AppendEq ca cb' cc ) =>
    D '(a, [ra,ca]) e -> D '(b, [rb, cb]) e -> D '(a, [ra, cc]) e
hconcat (DMat a) (DMat b) = DMat (H.fromBlocks [[a, b]])

vconcat :: (AppendEq ra (b ': rb) rc) =>
    D '(a, '[ra,ca]) e -> D '(b, '[rb,ca]) e -> D '(a, '[rc,ca]) e
vconcat (DMat a) (DMat b) = DMat (H.fromBlocks [[a],[b]])

rank, rows, cols :: D t a -> Int
rank (DMat a) = H.rank a
rows (DMat a) = H.rows a
cols (DMat a) = H.cols a

-- | H.'H.rows' except type-level
rowsNT :: D '(a, r ': c) e -> Proxy (HLength (a ': ri))
rowsNT _ = Proxy

-- | H.'H.cols' except type-level
colsNT :: D '(a, r ': c ': cs) e -> Proxy (HLength (DOne ': c))
colsNT _ = Proxy

-- | (m `hasRows` n) constrains the matrix/vector @m@ to have @n@ rows
hasRows :: (SameLength (HReplicateR n ()) r, -- forwards
            HLength r ~ n -- backwards
    ) => D '(a, ra ': ca) e -> Proxy (n :: HNat) -> D '(a, ra ': ca) e
hasRows x _ = x

-- | (m `hasRows` n) constrains the matrix/vector @m@ to have @n@ rows
hasCols :: (SameLength (HReplicateR n ()) ci, -- forwards
            HLength ci ~ n -- backwards
    ) => D '(a, ra ': ca ': rest) e -> Proxy (n :: HNat) -> D '(a, ra ': ca ': rest) e
hasCols x _ = x

-- | H.'H.scalar'
class (MapConst '[] as ~ as) => Scalar as where
    scalar :: D '(a, '[]) e -> D '(a, as) e

instance Scalar '[ '[] ] where
  scalar (DScal a) = DVec (H.scalar a)

instance Scalar '[ '[], '[] ] where
  scalar (DScal a) = DMat (H.scalar a)

{- | Numeric.Container.'H.konst', but the size is determined by the type.

>>> let n = hSucc (hSucc hZero) -- 2
>>> konst ((1::Double) *~ second) `hasRows` n `hasCols` n
2><2 1   1
s    1.0 1.0
s    1.0 1.0

-}
konst :: forall e a ra ca.
    (H.Field e,
     HNat2Integral (HLength (a ': ra)),
     HNat2Integral (HLength (DOne ': ca)),
     AllEq DOne ca,
     AllEq a ra)
    => D '(a, '[]) e -> D '(a, '[ra, ca]) e
konst (DScal a) = DMat (H.konst a
    (hNat2Integral (Proxy :: Proxy (HLength (a ': ra))),
     hNat2Integral (Proxy :: Proxy (HLength (DOne ': ca)))))

-- | identity matrix. The size is determined by the type.
ident :: forall ones e.
    (H.Field e, HNat2Integral (HLength (DOne ': ones))) =>
    D '(DOne, [ones, ones]) e
ident = DMat (H.ident (hNat2Integral (Proxy :: Proxy (HLength (DOne ': ones)))))

-- | zero matrix. The size and dimension is determined by the type.
zeroes :: forall c a r e. (H.Field e,
                        HNat2Integral (HLength (a ': r)),
                        HNat2Integral (HLength (DOne ': c)))
    => D '(a, '[r, c]) e
zeroes = DMat (H.konst 0
        (hNat2Integral (Proxy :: Proxy (HLength (a ': r))),
         hNat2Integral (Proxy :: Proxy (HLength (DOne ': c)))))

type family CanAddConst (a :: k) (m :: [[k]]) :: Constraint
type instance CanAddConst a [as, ones] = (AllEq a as, AllEq '[] ones)
type instance CanAddConst a '[as] = (AllEq a as)

addConstant :: (H.Field e, CanAddConst a sh)
    => D '(a, '[]) e
    -> D '(a, sh) e
    -> D '(a, sh) e
addConstant (DScal a) (DMat b) = DMat (H.addConstant a b)
addConstant (DScal a) (DVec b) = DVec (H.addConstant a b)

conj :: D sh a -> D sh a
conj (DMat a) = DMat (H.conj a)
conj (DVec a) = DVec (H.conj a)

-- | conjugate transpose
ctrans x = conj . trans $ x

diag :: (MapConst DOne r ~ c, SameLength r c)
  => D '(a, '[r]) t -> D '(a, '[r,c]) t
diag (DVec a) = DMat (H.diag a)

-- | 'H.blockDiag'. The blocks should be provided as:
--
-- @blockDiag $ 'hBuild' m1 m2 m3@
--
-- only available if hmatrix >= 0.15
diagBlock :: (HMapOut UnDimMat (b ': bs) (H.Matrix e),
              Num e, H.Field e, AppendShOf b bs (D '(a, sh) e),
              sh ~ '[r,c])
  => HList (b ': bs)
  -> D '(a, sh) e
diagBlock pairs = DMat (H.diagBlock (hMapOut UnDimMat pairs))

data UnDimMat = UnDimMat
instance (D sh a ~ x, H.Matrix a ~ y) => ApplyAB UnDimMat x y where
        applyAB _ (DMat x) = x

class DiagBlock (bs :: [*]) t

-- | @AppendShOf a [b,c,d] aas@ makes aas have the type of a matrix that
-- has a,b,c,d along the diagonal
class AppendShOf a (as :: [*]) aas
instance
 (e ~ f, f ~ g,
  AppendShOf (D ab e) ds z,
  AppendDims a b ab,

  -- constraints to force D in the type
  x ~ D a e,
  y ~ D b f,
  z ~ D c g) =>
  AppendShOf x (y ': ds) z
instance (x ~ z) => AppendShOf x '[] z

class AppendDims (a :: (*, [[*]])) (b :: (*, [[*]])) (c :: (*, [[*]]))
    | a b -> c, c a -> b, c b -> a
instance (c ~ a, AppendEq ra (b ': rb) rc, AppendEq ca cb cc) =>
  AppendDims '(a, [ra,ca]) '(b, [rb,cb]) '(c, [rc,cc])
-- how to handle vectors?
--type instance AppendDims '(a, '[ra]) '(b, '[rb]) = '(a, '[HAppendR ra (b ': rb)])

class ToHList sh e result where
    toHList :: D sh e -> HList result

-- | given a vector like @x = DimMat '[units] e@ this does something like
-- @[ (x \@> i) | i <- [1 .. n] ]@, if we had comprehensions for HLists
instance
    (HListFromList e e1,
     SameLength result e1,
     HMapCxt HList AddDimensional e1 result,
     -- HMapAddDimensional result e1,
     ToHListRow (r ': rs) e result) =>
   ToHList '(r, '[rs]) e result where
  toHList (DVec v) = case hListFromList (H.toList v) :: HList e1 of
      e1 -> hMap AddDimensional e1

class HMapAddDimensional (a :: [*]) (b :: [*]) | a -> b
instance HMapAddDimensional '[] '[]
instance (HMapAddDimensional as bs,
    dta ~ Quantity t b) => HMapAddDimensional (dta ': as) (b ': bs)

class HListFromList e e' | e' -> e where
        hListFromList :: [e] -> HList e'

instance HListFromList e '[e] where
        hListFromList (e : _) = HCons e HNil
        hListFromList _ = error "HListFromList: not enough input"
instance (e ~ e', HListFromList e es) => HListFromList e (e' ': es) where
        hListFromList (e : es) = e `HCons` hListFromList es

type family ToHListRow (a :: [*]) e (b :: [*]) :: Constraint
type instance ToHListRow (a ': as) e (b ': bs) = (Quantity a e ~ b, ToHListRow as e bs)

data AddDimensional = AddDimensional
instance (Quantity t x ~ y) => ApplyAB AddDimensional x y where
        applyAB _ x = Dimensional x

class FromHList list sh e where
  fromHList :: HList list -> D sh e

instance
    (H.Field e,
     HMapOut RmDimensional list e,
     ToHListRow (r ': rs) e list) =>
  FromHList list '(r, '[rs]) e where
  fromHList xs = DVec (H.fromList (hMapOut RmDimensional xs))

data RmDimensional = RmDimensional
instance (x ~ Quantity d y) => ApplyAB RmDimensional x y where
        applyAB _ (Dimensional a) = a

class FromHLists lists sh e where
  fromHLists :: HList lists -> D sh e

-- | [[Dim e unit]] -> DimMat units e
instance
  (ToHListRows' (r1 ': r) c e lists,
   HMapOut (HMapOutWith RmDimensional) lists [e],
   H.Field e) =>
  FromHLists lists '(r1, [r,c]) e where
    fromHLists xs = DMat (H.fromLists (hMapOut (HMapOutWith RmDimensional) xs))

newtype HMapOutWith f = HMapOutWith f
instance (HMapOut f l e, es ~ [e], HList l ~ hl) => ApplyAB (HMapOutWith f) hl es where
    applyAB (HMapOutWith f) = hMapOut f

class ToHListRows' (r :: [*]) (c :: [*]) (e :: *) (rows :: [*])
instance ToHListRows' '[] c e '[]

instance (ToHListRows' r c e rows,
          MultEq r c c',
          HMapCxt HList (AddQty e) c' row')
  => ToHListRows' (r1 ': r) c e (hListRow ': rows)

data AddQty u
instance (qty ~ Quantity u e) => ApplyAB (AddQty u) e qty where
    applyAB _ = Dimensional

class ToHLists sh e xs where
    toHLists :: D sh e -> HList xs

-- | DimMat units e -> [[Dim e unit]]
instance
    (HListFromList e e1,
     HListFromList (HList e1) e2,
     HMapCxt HList (HMap AddDimensional) e2 xs,
     ToHListRows' ri ci e xs,
     SameLength e2 xs,
     (r1 ': r) ~ ri, (DOne ': c) ~ ci )
  => ToHLists '(r1, [r,c]) e xs where
  toHLists (DMat m) = case hListFromList (map hListFromList (H.toLists m) :: [HList e1]) :: HList e2 of
    e2 -> hMap (HMap AddDimensional) e2

{- still bad

class PairsToList a t where
        pairsToList :: a -> [H.Matrix t]
instance PairsToList () t where
        pairsToList _ = []
instance (PairsToList b t, t' ~ t) => PairsToList (DimMat sh t',b) t where
        pairsToList (DimMat a,b) = a : pairsToList b

class EigV (sh :: [[ [DimSpec *] ]])
           (eigenValue  :: [[DimSpec *]])
           (eigenVector :: [[[DimSpec *]]])

instance
  ( SameLengths [r,c,r',c',rinv,cinv,eigval,erinv],
    -- ZipWithMul r c eigval,
    MapConst '[] r ~ eigval,
    PInv [r',c'] [rinv,cinv],
    -- AreRecips r' cinv,
    -- AreRecips c' rinv,
    cinv ~ c,
    c ~ ('[] ': _1),
    c' ~ ('[] ': _2),
    ZipWithMul eigval rinv erinv,
    MultiplyCxt [r',c'] erinv r,
    sh ~ [r,c],
    sh' ~ [r',c'])
    =>  EigV sh eigval sh'
-- | when no eigenvectors are needed
type family EigE (sh :: [[ [DimSpec *] ]]) (eigenValue  :: [ [DimSpec *] ]) :: Constraint
type instance EigE [r,c] eigval = ( SameLengths [r,c,eigval], ZipWithMul r c eigval)

{- $eigs

The Hmatrix eig factors A into P and D where A = P D inv(P) and D is diagonal.

The units for eigenvalues can be figured out:

>               _____
>      -1       |  c
> P D P  = A =  |r
>               |

>       _______
>       |   d
> P   = |c
>       |

>       _______
>       |   -1
>       |  c
>  -1   |
> P   = | -1
>       |d

So we can see that the dimension labeled `d-1` in P inverse is actually the
same `c` in `A`. The actual units of `d` don't seem to matter because the
`inv(d)` un-does any units that the `d` adds. So `d` can be all DOne. But
another choice, such as 1/c would be more appropriate, since then you can
expm your eigenvectors (not that that seems to be something people do)?

To get the row-units of A to match up, sometimes `D` will have units.
The equation ends up as D/c = r

Please ignore the type signatures on 'eig' 'eigC' etc. instead look at the type of
'wrapEig' 'wrapEigOnly' together with the hmatrix documentation (linked).

Perhaps the convenience definitions `eig m = wrapEig H.eig m` should be in
another module.
-}

{-
-- | 'wrapEig' H.'H.eig'
eig m = wrapEig H.eig m
-- | 'wrapEig' H.'H.eigC'
eigC m = wrapEig H.eigC m
-- | 'wrapEig' H.'H.eigH'
eigH m = wrapEig H.eigH m
-- | 'wrapEig' H.'H.eigH''
eigH' m = wrapEig H.eigH' m
-- | 'wrapEig' H.'H.eigR'
eigR m = wrapEig H.eigR m
-- | 'wrapEig' H.'H.eigS'
eigS m = wrapEig H.eigS m
-- | 'wrapEig' H.'H.eigS''
eigS' m = wrapEig H.eigS' m
-- | 'wrapEig' H.'H.eigSH'
eigSH m = wrapEig H.eigSH m
-- | 'wrapEig' H.'H.eigSH''
eigSH' m = wrapEig H.eigSH' m

-- | 'wrapEigOnly' H.'H.eigOnlyC'
eigOnlyC m = wrapEigOnly H.eigOnlyC m
-- | 'wrapEigOnly' H.'H.eigOnlyH'
eigOnlyH m = wrapEigOnly H.eigOnlyH m
-- | 'wrapEigOnly' H.'H.eigOnlyR'
eigOnlyR m = wrapEigOnly H.eigOnlyR m
-- | 'wrapEigOnly' H.'H.eigOnlyS'
eigOnlyS m = wrapEigOnly H.eigOnlyS m
-- | 'wrapEigOnly' H.'H.eigenvalues'
eigenvalues m = wrapEigOnly H.eigenvalues m
-- | 'wrapEigOnly' H.'H.eigenvaluesSH'
eigenvaluesSH m = wrapEigOnly H.eigenvaluesSH m
-- | 'wrapEigOnly' H.'H.eigenvaluesSH''
eigenvaluesSH' m = wrapEigOnly H.eigenvaluesSH' m
-}

wrapEig :: (c' ~ ('[] ': _1),
            EigV [r,c] eigVal [r',c'],
    H.Field y, H.Field z)
    => (H.Matrix x -> (H.Vector y, H.Matrix z)) ->
    DimMat [r,c] x ->
    (DimMat '[eigVal] y, DimMat [r',c'] z)
wrapEig hmatrixFun (DimMat a) = case hmatrixFun a of
    (e,v) -> (DimVec e, DimMat v)

wrapEigOnly :: (EigE [r,c] eigVal, H.Field y)
    => (H.Matrix x -> H.Vector y) ->
    DimMat [r,c] x -> DimMat '[eigVal] y
wrapEigOnly hmatrixFun (DimMat a) = case hmatrixFun a of
    (e) -> DimVec e

-}

-- | @\\a xs -> map (map (const a)) xs@
type family MapMapConst (a::k) (xs :: [[l]]) :: [[k]]
type instance MapMapConst a (x ': xs) = MapConst a x ': MapMapConst a xs
type instance MapMapConst a '[] = '[]

-- | @\\a xs -> map (const a) xs@
type family MapConst (a :: k) (xs :: [l]) :: [k]
type instance MapConst a (x ': xs) = a ': MapConst a xs
type instance MapConst a '[] = '[]

-- | convert from (a,(b,(c,(d,())))) to '[a,b,c,d]
type family FromPairs (a :: *) :: [*]
type instance FromPairs (a,b) = a ': FromPairs b
type instance FromPairs () = '[]

{-
-- | @map fst@
type family MapFst (a :: *) :: [*]
type instance MapFst ((a,_t) , as) = a ': MapFst as
type instance MapFst () = '[]

-- | @\\a xs -> map (/a) xs@
type family MapDiv (a :: k) (xs :: [k]) :: [k]
type instance MapDiv a (x ': xs) = (x @- a) ': MapDiv a xs
type instance MapDiv a '[] = '[]

type family UnDQuantity (a :: [*]) :: [ [*] ]
type instance UnDQuantity (x ': xs) = UnDQuantity1 x ': UnDQuantity xs
type instance UnDQuantity '[] = '[]

type family UnDQuantity1 (a :: *) :: [*]
type instance UnDQuantity1 (Unit t x) = x

type family DimMatFromTuple ijs :: * -> *
type instance DimMatFromTuple ijs =
        DimMat [UnDQuantity (MapFst ijs),
               '[] ': MapDiv (UnDQuantity1 (Fst (Fst ijs)))
               (UnDQuantity (FromPairs (Snd (Fst ijs))))]
type family Append (a :: [k]) (b :: [k]) :: [k]
type instance Append (a ': as) b = a ': Append as b
type instance Append '[] b = b
-}


{-# LANGUAGE GADTs #-}
{-# LANGUAGE StandaloneDeriving #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE MultiParamTypeClasses #-}

module Control.Quantum
    ( Amplitude
    , WaveFunction
    , Quantum
    , quantize
    , MonadQuantum
    , runQuantum
    , runQuantum'
    , measure
    , measurements
    )
    where

import qualified Data.Map  as M
import           Data.Map (Map)
import           Data.Monoid
import           Control.Monad
import           Control.Applicative
import           Control.Monad.Random as R
import           Data.Complex

nearZero :: Double -> Bool
nearZero x = abs x <= 1e-12

type Amplitude = Complex Double

-- | The Wave function is the distribution of amplitudes over the
-- possible states of the underlying data type
type WaveFunction a = [(a, Amplitude)]

class (Monad m) => MonadQuantum m where
    quantize  :: (Ord a) => WaveFunction a -> m a
    quantize' ::            WaveFunction a -> m a
    condition :: Bool -> m ()

liftQ  :: (Ord b, MonadQuantum m) => (a -> b) -> m a -> m b
liftQ f q1 = q1 >>= always . f

liftQ2 :: (Ord c, MonadQuantum m) => (a -> b -> c) -> m a -> m b -> m c
liftQ2 f p1 p2 = do x1 <- p1
                    x2 <- p2
                    always (f x1 x2)

liftQ3 :: (Ord d, MonadQuantum m) => (a -> b -> c -> d) -> m a -> m b -> m c -> m d
liftQ3 f p1 p2 p3 = do x1 <- p1
                       x2 <- p2
                       x3 <- p3
                       always (f x1 x2 x3)

always a = quantize [(a, 1.0:+0.0)]
always' a = quantize' [(a, 1.0:+0.0)]

-- | Describes the quantized version of an arbitrary data type.
-- The underlying data type defines the basis that the state is measured in.
data Quantum a where
    Quantum :: Ord a => Map (Maybe a) Amplitude -> Quantum a
    QuantumAny :: [(Maybe a, Amplitude)]  -> Quantum a

noState  :: (Ord a) => Quantum a
noState  = Quantum (M.singleton Nothing (1.0:+0.0))

noState' :: Quantum a
noState' = QuantumAny [(Nothing, 1.0 :+ 0.0)]

deriving instance (Show a) => Show (Quantum a)

instance Functor Quantum where
    fmap = liftM

instance Applicative Quantum where
    pure  = return
    (<*>) = ap

instance Monad Quantum where
    return  = always'

    m >>= f = if unitary then next
                         else error "Non-unitary transformation applied to quantum state!"
        where
            unitary = nearZero $ l2norm (map snd $ toList' next) - 1.0
            next = collect [multAmpl q (go a) | (a, q) <- toList' m]
            go = maybe noState' f

multAmpl :: Amplitude -> Quantum a -> Quantum a
multAmpl q (Quantum x) = Quantum $ M.map (* conjugate q) x
multAmpl q (QuantumAny x) = QuantumAny [ (a, q * r) | (a, r) <- x ]

toList' :: Quantum a -> [(Maybe a, Amplitude)]
toList' (Quantum x) = M.toList x
toList' (QuantumAny x) = x

toList :: (Ord a) => Quantum a -> [(Maybe a, Amplitude)]
toList (Quantum x) = M.toList x
toList (QuantumAny x) = merge x

collect :: [Quantum a] -> Quantum a
collect [ ]        = QuantumAny []
collect [x]        = x
collect (Quantum x:t) = case collect t of
                        Quantum y -> Quantum (M.unionWith (+) x y)
                        QuantumAny y -> Quantum (M.unionWith (+) x (M.fromList y))
collect (QuantumAny x:t) = case collect t of
                        Quantum y -> Quantum (M.unionWith (+) (M.fromList x) y)
                        QuantumAny y -> QuantumAny (x ++ y)

merge :: (Ord a) => WaveFunction a -> WaveFunction a
merge = M.toList . M.fromListWith (+)

instance MonadQuantum Quantum where
    quantize    = Quantum . M.fromListWith (+) . normalize
    quantize'   = QuantumAny . normalize
    condition test = if test then always () else noState

instance (Ord a, Monoid a) => Monoid (Quantum a) where
    mempty  = always mempty
    mappend = liftQ2 mappend

normalize :: WaveFunction a -> [(Maybe a, Amplitude)]
normalize xs = map (\(a, q) -> (Just a, q / total)) xs
    where
        total = sum $ zipWith (*) (map conjugate ampl) ampl
        ampl  = map snd xs

-- | Remove all impossible states from the quantum state and renormalize
collapse :: [(Maybe a, Amplitude)] -> WaveFunction a
collapse xs = [ (x, q / (norm:+0)) | (Just x, q) <- xs ]
    where
        norm = l2norm [ q | (Just x, q) <- xs ]

l2norm :: [Amplitude] -> Double
l2norm xs = sqrt $ sum $ map (\x -> (magnitude x)**2) xs

-- | Extracts the wave function from a quantum state.
runQuantum :: (Ord a) => Quantum a -> WaveFunction a
runQuantum = collapse . toList

runQuantum' :: Quantum a -> WaveFunction a
runQuantum' = collapse . toList'

-- | Measures a quantum state. Returns a single realization of the underlying data type
-- with probability equal to the squared magnitude of the amplitude of that realization.
-- Note that if this Module would be backed by a real quantum processor, this would be
-- the only valid way to extract information from the Monad.
measure :: (Ord a) => Quantum a -> IO a
measure q = do
  x <- evalRandIO $ R.fromList $ map (\(a, b) -> (a, toRational b)) $ measurements q
  return x

-- | Converts a quantum state into a probability state through measurement.
-- Equivalent to performing repeated measurements on equally prepared quantum states
-- and noting the frequency of each possible realization.
measurements :: (Ord a) => Quantum a -> [(a, Double)]
measurements = f . runQuantum
    where f xs = [(x, (magnitude q)**2) | (x, q) <- xs]


{-| Bindings to FFTW.

Example usage:

> import Foreign.Marshal.Array
> import Data.Complex
> import Foreign.Storable.Complex
> import FFTW
>
> main = do
>     inA  <- fftwAllocComplex 1024
>     outA <- fftwAllocComplex 1024
>
>     plan <- planDFT1d 1024 inA outA Forward fftwEstimate
>
>     pokeArray inA $ map (:+ 0) [0..1023]
>     execute plan
>     res <- peekArray 1024 outA
>
>     fftwFree inA
>     fftwFree outA
>
>     print res
-}

module Numeric.FFTW (
    -- * Memory allocation functions
    fftwMalloc,
    fftwFree,
    fftwFreePtr,
    fftwAllocReal,
    fftwAllocComplex,

    -- * FFT planning flags
    Direction(..),
    Flag(),
    -- ** Planning rigor flags
    fftwMeasure,
    fftwExhaustive,
    fftwPatient,
    fftwEstimate,
    fftwWisdomOnly,
    -- ** Algorithm restriction flags
    fftwDestroyInput,
    fftwUnaligned,
    fftwPreserveInput,

    -- * FFT planning
    FFTWPlan,
    planDFT1d,
    planDFTR2C1d,

    -- * FFT execution
    execute,
    executeDFT,
    executeDFTR2C
    ) where

import Foreign.C.Types
import Foreign.Ptr
import Data.Word
import Data.Complex
import Control.Monad
import Data.Bits
import Data.Monoid
import Data.Semigroup as Sem

#include <fftw3.h>

foreign import ccall unsafe "fftw_malloc"
    c_fftwMalloc :: CUInt -> IO (Ptr a)

-- | Like malloc, but ensures that the pointer obeys the alignment restrictions of FFTW (e.g. for SIMD acceleration). You probably want to use 'fftwAllocReal' or 'fftwAllocComplex' instead.
fftwMalloc :: Word32 -- ^ size
           -> IO (Ptr a)
fftwMalloc = c_fftwMalloc . fromIntegral

-- | Free a pointer returned by 'fftwMalloc', 'fftwAllocReal', or 'fftwAllocComplex'
foreign import ccall unsafe "fftw_free"
    fftwFree :: Ptr a -- ^ the pointer to be freed
             -> IO ()

-- | A function pointer to @fftwFree@.
foreign import ccall unsafe "&fftw_free"
    fftwFreePtr :: FunPtr (Ptr a -> IO ())

foreign import ccall unsafe "fftw_alloc_real"
    c_fftwAllocReal :: CUInt -> IO (Ptr CDouble)

-- | Allocates an array of Doubles. It ensures that the pointer obeys the alignment restrictions of FFTW (e.g. for SIMD acceleration).
fftwAllocReal :: Word32 -- ^ size
              -> IO (Ptr CDouble)
fftwAllocReal = c_fftwAllocReal . fromIntegral

foreign import ccall unsafe "fftw_alloc_complex"
    c_fftwAllocComplex :: CUInt -> IO (Ptr (Complex CDouble))

-- | Allocates an array of complex Doubles (i.e. the c type "double complex"). It ensures that the pointer obeys the alignment restrictions of FFTW (e.g. for SIMD acceleration).
fftwAllocComplex :: Word32 -- ^ size
                 -> IO (Ptr (Complex CDouble))
fftwAllocComplex = c_fftwAllocComplex . fromIntegral

-- | The direction of the transform: Forward for a normal transform, Backward for an inverse transform
data Direction = Forward
               | Backward

dirToInt :: Direction -> CInt
dirToInt Forward  = #const FFTW_FORWARD
dirToInt Backward = #const FFTW_BACKWARD

-- | FFTW planner flags. These flags affect the planning process. They can be combined using the 'Monoid' instance. See the FFTW flag documentation: <http://www.fftw.org/doc/Planner-Flags.html>.
newtype Flag = Flag {unFlag :: CUInt}

instance Sem.Semigroup Flag where
  (Flag x) <> (Flag y) = Flag (x .|. y)

instance Monoid Flag where
    mempty   = Flag 0
#if !(MIN_VERSION_base(4,11,0))
    mappend  = (Sem.<>)
#endif

fftwMeasure, fftwExhaustive, fftwPatient, fftwEstimate, fftwWisdomOnly :: Flag
fftwEstimate       = Flag #const FFTW_ESTIMATE
fftwMeasure        = Flag #const FFTW_MEASURE
fftwPatient        = Flag #const FFTW_PATIENT
fftwExhaustive     = Flag #const FFTW_EXHAUSTIVE
fftwWisdomOnly     = Flag #const FFTW_WISDOM_ONLY

fftwDestroyInput, fftwUnaligned, fftwPreserveInput :: Flag
fftwDestroyInput   = Flag #const FFTW_DESTROY_INPUT
fftwUnaligned      = Flag #const FFTW_UNALIGNED
fftwPreserveInput  = Flag #const FFTW_PRESERVE_INPUT

data CFFTWPlan

-- | A @FFTWPlan i o@ contains all of the information necessary to perform a transform from an input array of type @i@ to an output array of type @o@, including pointers to the input and output arrays.
newtype FFTWPlan i o = FFTWPlan (Ptr CFFTWPlan)

foreign import ccall unsafe "fftw_plan_dft_1d"
    c_planDFT1d :: CInt -> Ptr (Complex CDouble) -> Ptr (Complex CDouble) -> CInt -> CUInt -> IO (Ptr CFFTWPlan)

--This appears to be missing from the fft package on Hackage
-- | Create a plan for a 1 dimensional complex to complex DFT. The plan stores pointers to the input and output arrays, and these will be used if you 'execute' the plan in the future. They are required even if you intend to specify different input and output arrays in the future (i.e. using 'executeDFT')
planDFT1d :: Int                   -- ^ size
          -> Ptr (Complex CDouble) -- ^ input pointer
          -> Ptr (Complex CDouble) -- ^ output pointer
          -> Direction             -- ^ direction
          -> Flag                  -- ^ planner flags
          -> IO (FFTWPlan (Complex CDouble) (Complex CDouble))
planDFT1d n inp out sign flags = liftM FFTWPlan $ c_planDFT1d (fromIntegral n) inp out (dirToInt sign) (unFlag flags)

foreign import ccall unsafe "fftw_plan_dft_r2c_1d"
    c_planDFTR2C1d :: CInt -> Ptr CDouble -> Ptr (Complex CDouble) -> CUInt -> IO (Ptr CFFTWPlan)

--This appears to be missing from the fft package on Hackage
-- | Create a plan for a 1 dimensional real to complex DFT. The plan stores pointers to the input and output arrays, and these will be used if you 'execute' the plan in the future. They are required even if you intend to specify different input and output arrays in the future (i.e. using 'executeDFTR2C')
planDFTR2C1d :: Int                   -- ^ size
             -> Ptr CDouble           -- ^ input pointer
             -> Ptr (Complex CDouble) -- ^ output pointer
             -> Flag                  -- ^ planner flags
             -> IO (FFTWPlan CDouble (Complex CDouble))
planDFTR2C1d n inp out flags = liftM FFTWPlan $ c_planDFTR2C1d (fromIntegral n) inp out (unFlag flags)

foreign import ccall unsafe "fftw_execute"
    c_execute :: Ptr CFFTWPlan -> IO ()

-- | Execute a plan. Performs an FFT. The input and output arrays are stored within the plan so do not need to be given.
execute :: FFTWPlan i o -- ^ the plan to execute
        -> IO ()
execute (FFTWPlan p) = c_execute p

foreign import ccall unsafe "fftw_execute_dft"
    c_executeDFT :: Ptr CFFTWPlan -> Ptr (Complex CDouble) -> Ptr (Complex CDouble) -> IO ()

-- | Execute a complex to complex DFT but on different input and output arrays to those specified when the plan was created.
executeDFT :: FFTWPlan (Complex CDouble) (Complex CDouble) -- ^ the plan to execute
           -> Ptr (Complex CDouble)                        -- ^ input pointer
           -> Ptr (Complex CDouble)                        -- ^ output pointer
           -> IO ()
executeDFT (FFTWPlan p) inp out = c_executeDFT p inp out

foreign import ccall unsafe "fftw_execute_dft_r2c"
    c_executeDFTR2C :: Ptr CFFTWPlan -> Ptr CDouble -> Ptr (Complex CDouble) -> IO ()

-- | Execute a real to complex DFT but on different input and output arrays to those specified when the plan was created.
executeDFTR2C :: FFTWPlan CDouble (Complex CDouble) -- ^ the plan to execute
              -> Ptr CDouble                        -- ^ input pointer
              -> Ptr (Complex CDouble)              -- ^ output pointer
              -> IO ()
executeDFTR2C (FFTWPlan p) inp out = c_executeDFTR2C p inp out


module COMPLEX11 where

import Data.Complex
import Numeric

x = 0 :+ 3
y = 1 + x


module IllusoryContourShapePinwheelBasis where

import           Control.Monad                  as M
import           Control.Monad.Parallel         as MP
import qualified Data.Array.Accelerate          as A
import           Data.Array.Accelerate.LLVM.PTX
import           Data.Array.Repa                as R
import           Data.Binary
import           Data.Complex
import           Data.List                      as L
import           Data.Vector.Storable           as VS
import           Data.Vector.Unboxed            as VU
import           FokkerPlanck
import           Foreign.CUDA.Driver            as CUDA
import           FourierMethod.BlockMatrixAcc
import           FourierMethod.FourierSeries2D
import           Image.IO
import           Pinwheel.FourierSeries2D
import           STC
import           System.Directory
import           System.Environment
import           System.FilePath
import           Text.Printf
import           Utils.Array
import           Utils.Time
import           Utils.List
import Data.Array.IArray as IA
import FourierPinwheel
import Utils.Distribution
import FourierPinwheel.GaussianEnvelopePinwheel
import Filter.Utils

main = do
  args@(deviceIDsStr:numPointsStr:deltaStr:thresholdStr:numPointsReconStr:deltaReconStr:numOrientationStr:numScaleStr:thetaSigmaStr:scaleSigmaStr:tauStr:numTrailsStr:deltaTStr:poissonWeightStr:numR2FreqStr:periodR2Str:phiFreqsStr:rhoFreqsStr:thetaFreqsStr:scaleFreqsStr:initDistStr:initScaleStr:histFilePath:histFilePathCorner:stdR2Str:stdThetaStr:stdRStr:numBatchR2Str:numBatchR2FreqsStr:numBatchOriStr:batchSizeStr:sStr:writeFlagStr:numIterationStr:shape2DStr:radiusStr:numThreadStr:_) <-
    getArgs
  let deviceIDs = read deviceIDsStr :: [Int]
      numPoints = read numPointsStr :: Int
      delta = read deltaStr :: Double
      threshold = read thresholdStr :: Double
      numPointsRecon = read numPointsReconStr :: Int
      deltaRecon = read deltaReconStr :: Double
      numOrientation = read numOrientationStr :: Int
      numScale = read numScaleStr :: Int
      thetaSigma = read thetaSigmaStr :: Double
      scaleSigma = read scaleSigmaStr :: Double
      tau = read tauStr :: Double
      numTrails = read numTrailsStr :: Int
      deltaT = read deltaTStr :: Double
      poissonWeight = read poissonWeightStr :: Double
      numR2Freq = read numR2FreqStr :: Int
      periodR2 = read periodR2Str :: Double
      phiFreq = read phiFreqsStr :: Int
      phiFreqs = L.map fromIntegral [-phiFreq .. phiFreq]
      rhoFreq = read rhoFreqsStr :: Int
      rhoFreqs = L.map fromIntegral [-rhoFreq .. rhoFreq]
      thetaFreq = read thetaFreqsStr :: Int
      thetaFreqs = L.map fromIntegral [-thetaFreq .. thetaFreq]
      scaleFreq = read scaleFreqsStr :: Int
      scaleFreqs = L.map fromIntegral [-scaleFreq .. scaleFreq]
      initScale = read initScaleStr :: Double
      initDist = read initDistStr :: [(Double, Double, Double, Double)]
      initPoints = L.map (\(x, y, t, s) -> Point x y t s) initDist
      initSource = [L.head initPoints]
      initSink = [L.last initPoints]
      numThread = read numThreadStr :: Int
      folderPath = "output/test/IllusoryContourShapePinwheelBasis"
      stdR2 = read stdR2Str :: Double
      stdTheta = read stdThetaStr :: Double
      stdR = read stdRStr :: Double
      numBatchR2 = read numBatchR2Str :: Int
      numBatchR2Freqs = read numBatchR2FreqsStr :: Int
      numBatchOri = read numBatchOriStr :: Int
      batchSize = read batchSizeStr :: Int
      s = read sStr :: Double
      writeFlag = read writeFlagStr :: Bool
      numIteration = read numIterationStr :: Int
      shape2Ds = read shape2DStr :: [Points Shape2D]
      radius = read radiusStr :: Double
      periodEnv = periodR2 ^ 2 / 4 --  * sqrt 2 --  ^ 2 * 2
  createDirectoryIfMissing True folderPath
  flag <- doesFileExist histFilePath
  flagCorner <- doesFileExist histFilePathCorner
  hist <-
    case (getShape . L.head $ shape2Ds) of
      Circle _ _ ->
        if flag
          then do
            printCurrentTime "read coefficients from file"
            decodeFile histFilePath
          else do
            printCurrentTime "Start computing coefficients..."
            initialise []
            devs <- M.mapM device deviceIDs
            ctxs <- M.mapM (\dev -> CUDA.create dev []) devs
            ptxs <- M.mapM createTargetFromContext ctxs
            sampleCartesian
              histFilePath
              folderPath
              ptxs
              numPoints
              periodEnv
              delta
              numOrientation
              initScale
              thetaSigma
              tau
              threshold
              s
              phiFreq
              rhoFreq
              thetaFreq
              scaleFreq
              stdR2
      KoffkaCross _ _ ->
        if flagCorner
          then do
            printCurrentTime "read coefficients from file"
            decodeFile histFilePathCorner
          else do
            printCurrentTime "Start computing coefficients..."
            initialise []
            devs <- M.mapM device deviceIDs
            ctxs <- M.mapM (\dev -> CUDA.create dev []) devs
            ptxs <- M.mapM createTargetFromContext ctxs
            sampleCartesianCorner
              histFilePathCorner
              folderPath
              ptxs
              numPoints
              periodEnv
              delta
              numOrientation
              initScale
              thetaSigma
              tau
              threshold
              s
              phiFreq
              rhoFreq
              thetaFreq
              scaleFreq
              stdR2
              poissonWeight
  printCurrentTime "Done"
  printCurrentTime "Start Convloution.."
  plan <-
    makePlan
      folderPath
      emptyPlan
      numPointsRecon
      numPointsRecon
      numR2Freq
      (2 * thetaFreq + 1)
      (2 * scaleFreq + 1)
      (2 * phiFreq + 1)
      (2 * rhoFreq + 1)
  printCurrentTime "Compute DFT Plan done."
  let coefficients = getNormalizedHistogramArr hist
      asteriskGaussianVec =
        case getShape (L.head shape2Ds) of
          Circle _ _ ->
            gaussianPinwheel
              numR2Freq
              periodR2
              stdR2
              10
              thetaFreq
              scaleFreq
              periodEnv
              stdTheta
              stdR
          KoffkaCross _ _ ->
            gaussianPinwheel1
              numR2Freq
              periodR2
              stdR2
              10
              thetaFreq
              scaleFreq
              periodEnv
              stdTheta
              stdR
  harmonicsArray <-
    createHarmonics
      numR2Freq
      phiFreq
      rhoFreq
      thetaFreq
      scaleFreq
      (-s)
      periodR2
      periodEnv
      coefficients
  M.mapM_
    (\shape2D -> do
       let points =
             L.map (\(x, y) -> Point (x) (y) 0 1) .
             getShape2DIndexList' . makeShape2D $
             shape2D
       printCurrentTime (show points)
       (bias, dftBias) --Full
                          -- (dftBias, bias) <- -- Discrete
          <-
         case getShape shape2D of
           Circle _ _ ->
             computeBiasFourierPinwheelFull
               plan
               numR2Freq
               thetaFreq
               scaleFreq
               (-s)
               periodR2
               periodEnv
               radius
               stdTheta
               stdR
               stdR2
               asteriskGaussianVec
               points
           KoffkaCross _ _ ->
             computeBiasFourierPinwheelKoffkaCorss
               plan
               numR2Freq
               thetaFreq
               scaleFreq
               (-s)
               periodR2
               periodEnv
               radius
               stdTheta
               stdR
               stdR2
               asteriskGaussianVec
               points
       let initDist =
             computeInitialDistributionPowerMethodFourierPinwheelFull
               numR2Freq
               phiFreq
               rhoFreq
               thetaFreq
               scaleFreq
               bias -- dftBias
       computeContourFourierPinwheel
         plan
         folderPath
         writeFlag
         harmonicsArray
         dftBias -- bias
         numIteration
         numBatchR2
         numPointsRecon
         deltaRecon
         periodR2
         periodEnv
         initDist
         (show . getShape $ shape2D)
         deviceIDs)
    shape2Ds


{-# LANGUAGE FlexibleContexts #-}

module LearnSpec
  ( spec
  ) where

import           Control.Monad
import qualified Data.Set                as Set
import           Data.Tuple.Extra
import           Debug.Trace
import           Foundation.Monad
import           Layers
import           Learn
import           Mnist
import           Numeric.LinearAlgebra
import           Test.Hspec
import           Test.Hspec.QuickCheck   (prop)
import           Test.QuickCheck
import           Test.QuickCheck.Gen
import           Test.QuickCheck.Monadic

runTest :: IO ()
runTest = hspec spec

spec :: Spec
spec = do
  describe "hotone" propsHotone
  describe "convertTests" propsConvertTests
  describe "convertTrains" propsConvertTrains
  describe "initNN" propsInitNN

propsHotone =
  prop "length" $
  forAll (genNM 100) $ \(n, m) ->
    let v = toList $ hotone n m :: [R]
        (a, k:c) = splitAt m v
        x = sum $ a ++ c
     in (length v, x, k) `shouldBe` (n, 0, 1)

propsConvertTests =
  prop "div 255" $
  forAll genMnistData $ \d ->
    let r = convertTests d
        MnistData ns = d
        (ts, is) = second concat $ unzip $ map (second $ concat . toLists) ns
        (ps, vs) = second concat $ unzip $ map (second toList) r
        x1 = map fromIntegral ts
        x2 = map (round . (* 255)) vs
     in (x1, is) `shouldBe` (ps, x2)

propsConvertTrains = do
  prop "length" $
    forAll genMnistData $ \d ->
      let r = convertTrains (length ns `div` 10) d
          MnistData ns = d
          (xs, ys) = unzip $ map countRows r
          countRows (TrainBatch a) = both rows a
          a = sum xs
          b = length ns
       in (xs, a) `shouldBe` (ys, b)
  prop "size" $
    forAll
      (do d <- genMnistData
          s <- choose (2, 10)
          return (s, d)) $ \(s, d) ->
      let r = convertTrains s d
          MnistData ns = d
          (w, c) = size $ snd $ head ns
          xs = unzip $ map getSize r
          getSize (TrainBatch a) = both size a
          (a, b) = both (maximum . Set.toList . Set.fromList) xs
       in (a, b) `shouldBe` ((s, 10), (s, w * c))
  prop "div 255" $
    forAll genMnistData $ \d ->
      let r = convertTrains (length ns `div` 10) d
          MnistData ns = d
          xs = map fromIntegral $ concatMap (concat . toLists . snd) ns
          ys = map (round . (* 255)) $ concatMap flatMs r
          flatMs (TrainBatch (_, ms)) = concat $ toLists ms
       in xs `shouldBe` ys

propsInitNN =
  it "size of layers" $ do
    r <- initNN ReLUForward [16, 50, 8, 10]
    let ForwardNN lA SoftmaxWithCrossForward = r
    let JoinedForwardLayer lB (AffineForward m3 b3) = lA
    let JoinedForwardLayer lC ReLUForward = lB
    let JoinedForwardLayer lD (AffineForward m2 b2) = lC
    let JoinedForwardLayer (AffineForward m1 b1) ReLUForward = lD
    let ms = map size [m1, m2, m3]
    let bs = map size [b1, b2, b3]
    (ms, bs) `shouldBe` ([(16, 50), (50, 8), (8, 10)], [50, 8, 10])

genNM :: Int -> Gen (Int, Int)
genNM x = do
  n <- choose (1, x)
  m <- choose (0, n - 1)
  return (n, m)

genMnistData :: Gen MnistData
genMnistData = do
  len <- choose (10, 100)
  nRows <- choose (10, 20)
  nCols <- choose (10, 20)
  ms <- vectorOf len $ genMatrixZ 255 nRows nCols
  vs <- vectorOf len $ choose (0, 9)
  return $ MnistData $ zip vs ms

genMatrixZ :: Int -> Int -> Int -> Gen (Matrix Z)
genMatrixZ value nRows nCols = do
  zss <- vectorOf nRows $ vectorOf nCols $ fromIntegral <$> choose (0, value)
  return $ fromLists zss


{-# LANGUAGE ConstraintKinds     #-}
{-# LANGUAGE DataKinds           #-}

{-# LANGUAGE DeriveTraversable   #-}
{-# LANGUAGE FlexibleContexts    #-}
{-# LANGUAGE GADTs               #-}
{-# LANGUAGE RankNTypes          #-}
{-# LANGUAGE ScopedTypeVariables #-}

{-# OPTIONS_GHC -fno-warn-missing-signatures #-}
module Test.Grenade.Recurrent.Layers.LSTM.Reference where

import           Data.Reflection
import           GHC.TypeLits                  (KnownNat)
import           Numeric.AD.Internal.Reverse   (Tape)
import           Numeric.AD.Mode.Reverse
import qualified Numeric.LinearAlgebra         as H
import qualified Numeric.LinearAlgebra.Static  as S

import           Grenade.Recurrent.Layers.LSTM (LSTMWeights (..))
import qualified Grenade.Recurrent.Layers.LSTM as LSTM
import           Grenade.Types

--
-- This module contains a set of list only versions of
-- an LSTM layer which can be used with the AD library.
--
-- Using this, we can check to make sure that our fast
-- back propagation implementation is correct.
--

-- | List only matrix deriving functor
newtype Matrix a = Matrix {
    matrixWeights         :: [[a]]
  } deriving (Functor, Foldable, Traversable, Eq, Show)

-- | List only vector deriving functor
newtype Vector a = Vector {
    vectorWeights         :: [a]
  } deriving (Functor, Foldable, Traversable, Eq, Show)

-- | List only LSTM weights
data RefLSTM a = RefLSTM
    { refLstmWf :: Matrix a -- Weight Forget     (W_f)
    , refLstmUf :: Matrix a -- Cell State Forget (U_f)
    , refLstmBf :: Vector a -- Bias Forget       (b_f)
    , refLstmWi :: Matrix a -- Weight Input      (W_i)
    , refLstmUi :: Matrix a -- Cell State Input  (U_i)
    , refLstmBi :: Vector a -- Bias Input        (b_i)
    , refLstmWo :: Matrix a -- Weight Output     (W_o)
    , refLstmUo :: Matrix a -- Cell State Output (U_o)
    , refLstmBo :: Vector a -- Bias Output       (b_o)
    , refLstmWc :: Matrix a -- Weight Cell       (W_c)
    , refLstmBc :: Vector a -- Bias Cell         (b_c)
    } deriving (Functor, Foldable, Traversable, Eq, Show)

lstmToReference :: (KnownNat a, KnownNat b) => LSTM.LSTMWeights a b -> RefLSTM RealNum
lstmToReference lw =
    RefLSTM
      { refLstmWf = Matrix . H.toLists . S.extract $ lstmWf lw -- Weight Forget     (W_f)
      , refLstmUf = Matrix . H.toLists . S.extract $ lstmUf lw -- Cell State Forget (U_f)
      , refLstmBf = Vector . H.toList  . S.extract $ lstmBf lw -- Bias Forget       (b_f)
      , refLstmWi = Matrix . H.toLists . S.extract $ lstmWi lw -- Weight Input      (W_i)
      , refLstmUi = Matrix . H.toLists . S.extract $ lstmUi lw -- Cell State Input  (U_i)
      , refLstmBi = Vector . H.toList  . S.extract $ lstmBi lw -- Bias Input        (b_i)
      , refLstmWo = Matrix . H.toLists . S.extract $ lstmWo lw -- Weight Output     (W_o)
      , refLstmUo = Matrix . H.toLists . S.extract $ lstmUo lw -- Cell State Output (U_o)
      , refLstmBo = Vector . H.toList  . S.extract $ lstmBo lw -- Bias Output       (b_o)
      , refLstmWc = Matrix . H.toLists . S.extract $ lstmWc lw -- Weight Cell       (W_c)
      , refLstmBc = Vector . H.toList  . S.extract $ lstmBc lw -- Bias Cell         (b_c)
      }

runLSTM :: Floating a => RefLSTM a -> Vector a -> Vector a -> (Vector a, Vector a)
runLSTM rl cell input =
    let -- Forget state vector
        f_t = sigmoid   $ refLstmBf rl #+ refLstmWf rl #> input #+ refLstmUf rl #> cell
        -- Input state vector
        i_t = sigmoid   $ refLstmBi rl #+ refLstmWi rl #> input #+ refLstmUi rl #> cell
        -- Output state vector
        o_t = sigmoid   $ refLstmBo rl #+ refLstmWo rl #> input #+ refLstmUo rl #> cell
        -- Cell input state vector
        c_x = fmap tanh $ refLstmBc rl #+ refLstmWc rl #> input
        -- Cell state
        c_t = f_t #* cell #+ i_t #* c_x
        -- Output (it's sometimes recommended to use tanh c_t)
        h_t = o_t #* c_t
    in (c_t, h_t)

runLSTMback :: forall a. Floating a => Vector a -> Vector a -> RefLSTM a -> RefLSTM a
runLSTMback cell input =
  grad f
    where
  f :: forall s. Reifies s Tape => RefLSTM (Reverse s a) -> Reverse s a
  f net =
    let cell'   = fmap auto cell
        input'  = fmap auto input
        (cells, forwarded) = runLSTM net cell' input'
    in  sum forwarded + sum cells

runLSTMbackOnInput :: forall a. Floating a => Vector a -> RefLSTM a -> Vector a -> Vector a
runLSTMbackOnInput cell net =
  grad f
    where
  f :: forall s. Reifies s Tape => Vector (Reverse s a) -> Reverse s a
  f input =
    let cell'   = fmap auto cell
        net'    = fmap auto net
        (cells, forwarded) = runLSTM net' cell' input
    in  sum forwarded + sum cells

runLSTMbackOnCell :: forall a. Floating a => Vector a -> RefLSTM a -> Vector a -> Vector a
runLSTMbackOnCell input net =
  grad f
    where
  f :: forall s. Reifies s Tape => Vector (Reverse s a) -> Reverse s a
  f cell =
    let input'  = fmap auto input
        net'    = fmap auto net
        (cells, forwarded) = runLSTM net' cell input'
    in  sum forwarded + sum cells

-- | Helper to multiply a matrix by a vector
matMult :: Num a => Matrix a -> Vector a -> Vector a
matMult (Matrix m) (Vector v) = Vector result
  where
    lrs = map length m
    l   = length v
    result = if all (== l) lrs
             then map (\r -> sum $ zipWith (*) r v) m
             else error $ "Matrix has rows of length " ++ show lrs ++
                          " but vector is of length " ++ show l

(#>) :: Num a => Matrix a -> Vector a -> Vector a
(#>) = matMult
infixr 8 #>

(#+) :: Num a => Vector a -> Vector a -> Vector a
(#+) (Vector as) (Vector bs) = Vector $ zipWith (+) as bs
infixl 6 #+

(#-) :: Num a => Vector a -> Vector a -> Vector a
(#-) (Vector as) (Vector bs) = Vector $ zipWith (-) as bs
infixl 6 #-

(#*) :: Num a => Vector a -> Vector a -> Vector a
(#*) (Vector as) (Vector bs) = Vector $ zipWith (*) as bs
infixl 7 #*

sigmoid :: (Functor f, Floating a) => f a -> f a
sigmoid xs = (\x -> 1 / (1 + exp (-x))) <$> xs


{-# LANGUAGE BangPatterns        #-}
{-# LANGUAGE CPP                 #-}
{-# LANGUAGE DataKinds           #-}
{-# LANGUAGE FlexibleContexts    #-}
{-# LANGUAGE GADTs               #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeOperators       #-}

module Grenade.Utils.Conversion
  ( -- Column major mode helper functions
    NrColumns
  , Row
  , Col
  , ColumnMajorModeIndex
  , columnMajorModeIndex
  , columnMajorModeRowCol

  , toLayerShape
  -- 1D conversions
  , toS1D
  , toS1DV
  , fromS1D
  , fromS1DV
  -- 2D converions
  , toS2D
  , toS2DV
  , fromS2D
  , fromS2DV
  -- further conversions
  , toRows
  , toCols
  , toColumnsS2D
  , toRowsS2D
  , fromRowMajorVectorToSD1
  , fromRowMajorVectorToSD1V
  , fromRowMajorVectorToSD2
  , fromRowMajorVectorToSD2V
  , fromColumnMajorVectorToSD2V
  ) where

import           Data.Maybe                   (fromJust, fromMaybe)
import           Data.Proxy
import           Data.Singletons.TypeLits
import qualified Data.Vector.Storable         as V
import           Foreign
import qualified Numeric.LinearAlgebra        as LA
import qualified Numeric.LinearAlgebra.Static as LAS
import           System.IO.Unsafe             (unsafePerformIO)
import           Unsafe.Coerce                (unsafeCoerce)

import           Grenade.Core.Shape
import           Grenade.Types
import           Grenade.Utils.Vector

-- Column major mode helper functions

type NrColumns = Int
type Row = Int
type Col = Int
type ColumnMajorModeIndex = Int

columnMajorModeIndex :: NrColumns -> Row -> Col -> ColumnMajorModeIndex
columnMajorModeIndex cols row col = row * cols + col

columnMajorModeRowCol :: NrColumns -> ColumnMajorModeIndex -> (Row, Col)
columnMajorModeRowCol cols idx = (idx `div` cols, idx `mod` cols)

-- import           Debug.Trace

-- test =
--   toCols $
--   (fromRowMajorVectorToSD2 (V.fromList [0..9]) :: S ('D2 2 5))
--   (fromRowMajorVectorToSD1 (V.fromList [0..9]) :: S ('D1 10))

-- testShapeV = fromRowMajorVectorToSD2V (V.fromList [0..9]) :: S ('D2 2 5)
-- testShape = fromRowMajorVectorToSD2 (V.fromList [0..9]) :: S ('D2 2 5)

-- | Convert the Shape to a list of rows.
toRows :: S x -> [V.Vector RealNum]
toRows x =
  case x of
    (S1DV v)  -> [v]
    (S1D v)   -> [unsafeCoerce v]
    (S2D m)   -> map LAS.extract . LAS.toRows $ m
    (S2DV {}) -> toRowsS2D x
    (S3D m)   -> map LAS.extract . LAS.toRows $ m
{-# INLINE toRows #-}

-- | Convert the Shape to a list of columns.
toCols :: S x -> [V.Vector RealNum]
toCols x =
  case x of
    S1DV v    -> map V.singleton . V.toList $ v
    S1D v     -> map V.singleton . V.toList $ (unsafeCoerce v :: V.Vector RealNum)
    (S2DV {}) -> toColumnsS2D x
    (S2D m)   -> map LAS.extract . LAS.toColumns $ m
    (S3D m)   -> map LAS.extract . LAS.toColumns $ m

-- | Converts the given vector to the correct layer shape.
toLayerShape :: S i -> S x -> S x
toLayerShape x y = case (x, y) of
  (S1D{}, S1DV{}) -> toS1D y
  (S1DV{}, S1D{}) -> fromS1D y
  (S2D{}, S2DV{}) -> toS2D y
  (S2DV{}, S2D{}) -> fromS2D y
  _               -> y
{-# INLINE toLayerShape #-}

-- 1D conversion

-- | Convert to S1D.
toS1D :: S ('D1 l) -> S ('D1 l)
toS1D (S1DV vec) = S1D $ (fromMaybe err . LAS.create $ vec)
  where
    err = error $ "wrong length of vector with " ++ show (V.length vec) ++ " in toS1D "
toS1D x@S1D{} = x
{-# INLINE toS1D #-}

-- | Convert from S1DV. This is the same as @toS1D@.
fromS1DV :: S ('D1 l) -> S ('D1 l)
fromS1DV = toS1D
{-# INLINE fromS1DV #-}

-- | Convert from S1D to S1DV.
fromS1D :: S ('D1 l) -> S ('D1 l)
fromS1D (S1D vec) = S1DV (LAS.extract vec)
fromS1D x@S1DV{}  = x
{-# INLINE fromS1D #-}

-- | Convert to S1DV (from S1D). Thsi is the same as @fromS1D@.
toS1DV :: S ('D1 l) -> S ('D1 l)
toS1DV = fromS1D
{-# INLINE toS1DV #-}

-- 2D conversions

-- | Convert from vector representation.
toS2D :: forall i j . S ('D2 i j) -> S ('D2 i j)
toS2D mat@S2DV{} = S2D $ LAS.matrix (concatMap V.toList $ toRows mat)
toS2D x@S2D{}    = x
{-# INLINE toS2D #-}

-- | Convert from S2DV. This is the same as @toS2D@.
fromS2DV :: S ('D2 i j) -> S ('D2 i j)
fromS2DV = toS2D
{-# INLINE fromS2DV #-}

-- | Convert from S2D to S2DV.
fromS2D :: S ('D2 i j) -> S ('D2 i j)
fromS2D (S2D mat) = S2DV . V.concat . map LAS.extract . LAS.toColumns $ mat
fromS2D x@S2DV{}  = x
{-# INLINE fromS2D #-}

-- | Convert to S2DV. This is the same as @fromS2D@.
toS2DV :: S ('D2 i j) -> S ('D2 i j)
toS2DV = fromS2D
{-# INLINE toS2DV #-}

-- test = LAS.create $ NLA.reshape 5 $ (V.fromList [0..9]) :: Maybe (LAS.L 2 5)
-- Just (matrix
--  [ 0.0, 1.0, 2.0, 3.0, 4.0
--  , 5.0, 6.0, 7.0, 8.0, 9.0 ] :: L 2 5)

-- | Efficiently extract the columns of the shape. Returns a list of j vectors.
toColumnsS2D :: forall i j . S ('D2 i j) -> [V.Vector RealNum]
toColumnsS2D (S2D mat) = map LAS.extract . LAS.toColumns $ mat
toColumnsS2D (S2DV vec) = map (\idx -> V.slice idx m vec) [0,m .. (V.length vec - m)]
  where
    m = fromIntegral $ natVal (Proxy :: Proxy i)
{-# INLINE toColumnsS2D #-}

-- | Efficiently extract the rows of the shape. Returns a list of i vectors.
toRowsS2D :: forall i j . S ('D2 i j) -> [V.Vector RealNum]
toRowsS2D (S2D mat) = map LAS.extract . LAS.toRows $ mat
toRowsS2D (S2DV vec) =
  unsafePerformIO $
  V.unsafeWith vec $ \from ->
    flip mapM [0 .. m - 1] $ \row -> do
      vec' <- createVector n
      V.unsafeWith vec' $ \to -> do
        let go (-1) = return ()
            go !col = do
              let idx = col * m + row
              x <- peekElemOff from idx
              pokeElemOff to col x
              go (col - 1)
        go (n - 1)
        return vec'
  where
    m = fromIntegral $ natVal (Proxy :: Proxy i)
    n = fromIntegral $ natVal (Proxy :: Proxy j)
{-# INLINE toRowsS2D #-}

-- | Convert from a row major vector to @SD1@.
fromRowMajorVectorToSD1 :: forall l . (KnownNat l) => V.Vector RealNum -> S ('D1 l)
fromRowMajorVectorToSD1 vec
  | V.length vec /= l = error $ "cannot create Vector R " ++ show l ++ " from vector with length " ++ show (V.length vec) ++ " in fromRowMajorVectorToSD1"
  | otherwise = S1D (unsafeCoerce vec)
  where l = fromIntegral $ natVal (Proxy :: Proxy l)
{-# INLINE fromRowMajorVectorToSD1 #-}

-- | Convert from a row major vector to @SD1V@.
fromRowMajorVectorToSD1V :: forall l . (KnownNat l) => V.Vector RealNum -> S ('D1 l)
fromRowMajorVectorToSD1V = S1DV
{-# INLINE fromRowMajorVectorToSD1V #-}

-- | Convert from a row major vector to @SD2@.
fromRowMajorVectorToSD2 :: forall i j . (KnownNat i, KnownNat j) => V.Vector RealNum -> S ('D2 i j)
fromRowMajorVectorToSD2 vec
  | V.length vec /= m * n = error $ "cannot create matrix L " ++ show (m,n) ++ " from vector length " ++ show (V.length vec) ++ " in fromRowMajorVectorToSD2"
  | otherwise = toS2D $ fromRowMajorVectorToSD2V vec
                -- in S2D $ unsafeCoerce v'
  where
    m = fromIntegral $ natVal (Proxy :: Proxy i)
    n = fromIntegral $ natVal (Proxy :: Proxy j)
{-# INLINE fromRowMajorVectorToSD2 #-}

-- toRowMajorVector :: (KnownNat i, KnownNat j) => S ('D2 i j) -> V.Vector RealNum
-- toRowMajorVector = V.concat . toColumnsS2D

-- | Convert from a row major vector to @SD2V@.
fromRowMajorVectorToSD2V :: forall i j . (KnownNat i, KnownNat j) => V.Vector RealNum -> S ('D2 i j)
fromRowMajorVectorToSD2V vec = unsafePerformIO $ do
  vec' <- createVector (V.length vec)
  V.unsafeWith vec $ \from ->
    V.unsafeWith vec' $ \to ->  do
    let go (-1) = return ()
        go !k = do
          -- let idx = nRow * m + nCol
              -- nCol = k `div` n
              -- nRow = k `mod` n
          let idx = columnMajorModeIndex n row col
              col = k - row * n
              row = k `mod` n
          x <- peekElemOff from k
          pokeElemOff to idx x
          go (k-1)
    go (V.length vec - 1)
  return $ S2DV vec'
  where
    m = fromIntegral $ natVal (Proxy :: Proxy i)
    n = fromIntegral $ natVal (Proxy :: Proxy j)
{-# INLINE fromRowMajorVectorToSD2V #-}

-- | Convert from a row major vector to @SD2V@.
fromColumnMajorVectorToSD2V :: forall i j . (KnownNat i, KnownNat j) => V.Vector RealNum -> S ('D2 i j)
fromColumnMajorVectorToSD2V = S2DV
{-# INLINE fromColumnMajorVectorToSD2V #-}

test =
  -- (\x -> fromColumnMajorVectorToSD2V x :: S ('D2 5 2)) $
  -- V.concat $
  -- toRowsS2D $
  -- ttoRowMajorVector $
  fromS2D $
  -- (\x -> fromRowMajorVectorToSD2V x :: S ('D2 5 2) ) $ V.concat $
  -- toRowsS2D $
  -- ((\(S2D x) -> S2DV $ V.concat $ map LAS.extract . LAS.toColumns $ x :: S ('D2 2 5)) )
  -- (fromRowMajorVectorToSD2V (V.fromList [0..9]) :: S ('D2 2 5))
  (fromRowMajorVectorToSD2 (V.fromList [0..9]) :: S ('D2 5 2))
  -- (fromRowMajorVectorToSD1 (V.fromList [0..9]) :: S ('D1 10))


module LabelMatrix where

-- import Data.Function    ((&))
import Foreign.Storable (Storable)

import           Data.Map              (Map)
-- import qualified Data.Map              as M
import qualified Data.Vector  as V
-- import qualified Data.Vector.Storable  as VS
import           Numeric.LinearAlgebra hiding (fromList)

{-
Matrices with labeled rows and columns.

Construction functions handle uniqueness & correct length of labels.

Provides similar convenience functions to pandas DataFrames. Contrary to the latter, all
elements in the matrix contain only one type.
-}

--bla = ((12><12) [0..] :: Matrix I) #> fromList [0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0] :: Vector I

data LabelMatrix a r c =
    LabelMatrix
        { matrix :: Matrix a
        , rows   :: V.Vector r
        , cols   :: V.Vector c
        } deriving (Show)

fromMatrix :: Matrix a -> [r] -> [c] -> LabelMatrix a r c
fromMatrix = undefined

fromList :: Storable a => [a] -> [r] -> [c] -> LabelMatrix a r c
fromList content rowLabels columnLabels =
    let
        m = length rowLabels
        n = length columnLabels
        mat = (m><n) content
    in
        LabelMatrix mat (V.fromList rowLabels) (V.fromList columnLabels)

{-
v = m ¿ [0] -- extract column(s)

flatten v -- makes it into a vector, especially nice if the matrix is nx1 anyway.
-}

groupBy :: LabelMatrix a r c -> r -> Map a (LabelMatrix a r c)
groupBy = groupByRow

groupByRow :: LabelMatrix a r c -> r -> Map a (LabelMatrix a r c)
groupByRow = undefined

groupByColumn :: LabelMatrix a r c -> c -> Map a (LabelMatrix a r c)
groupByColumn = undefined


{-# LANGUAGE CPP                   #-}
#if defined(__GLASGOW_HASKELL__) && __GLASGOW_HASKELL__ >= 702
{-# LANGUAGE Trustworthy
           , DefaultSignatures     #-}
#define USE_GHC_GENERICS
#endif
{-# LANGUAGE RankNTypes
           , TypeFamilies
           , KindSignatures
           , DeriveGeneric
           , FlexibleInstances
           , FlexibleContexts
           , BangPatterns
           , MultiParamTypeClasses #-}
{-# OPTIONS_GHC -fno-warn-orphans  #-}

-------------------------------------------------------------------------------
-- |
-- Module      : OGLS.Engine.Math.Vectors
-- description : 3D and 4D vectors arithmetic
-- Copyright   : (c) Adam, 2017
-- License     : MIT
-- Maintainer  : awkure@protonmail.ch
-- Stability   : unstable
-- Portability : POSIX
-------------------------------------------------------------------------------

module OGLS.Engine.Math.Vectors
    ( (^$+), (^$-), (^$*)
    , dot, step
    ) where

#ifdef USE_GHC_GENERICS
import GHC.Generics               ( Generic (..) )
#endif
import Control.Applicative        ( liftA2       )
import Control.Monad              ( liftM2       )
import Control.Parallel           ( pseq         )
import Data.Complex               ( Complex (..) )
import Data.Orphans               (              )
import Control.Lens hiding ( (<.>) )

import OGLS.Engine.Math.Instances

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

default ()

class Vector a v where
    -- not yet implemented

(^$+) :: forall (t :: * -> *) a. (Applicative t, Num a) => t a -> t a -> t a
(^$-) :: forall (t :: * -> *) a. (Applicative t, Num a) => t a -> t a -> t a
(^$*) :: forall (t :: * -> *) a. (Applicative t, Num a) => t a -> t a -> t a

newtype E t = E { el :: forall x. Lens' (t x) x }

data Vec3 a = Vec3 { v3x :: {-# UNPACK #-} !a
                   , v3y :: {-# UNPACK #-} !a
                   , v3z :: {-# UNPACK #-} !a
                   } deriving ( Eq, Ord, Show, Generic )

instance Functor Vec3 where
    fmap f (Vec3 x y z) = Vec3 (f x) (f y) (f z)

instance Applicative Vec3 where
    pure x = Vec3 x x x
    Vec3 f g h <*> Vec3 x y z = Vec3 (f x) (g y) (h z)

data Vec4 a = Vec4 { v4x :: {-# UNPACK #-} !a
                   , v4y :: {-# UNPACK #-} !a
                   , v4z :: {-# UNPACK #-} !a
                   , v4w :: {-# UNPACK #-} !a
                   } deriving ( Eq, Ord, Show )

(^$+) = liftA2 (+)
{-# INLINE (^$+) #-}
(^$-) = liftA2 (-)
{-# INLINE (^$-) #-}
(^$*) = liftA2 (*)
{-# INLINE (^$*) #-}

-- | Left scalar product vector
(*^^) :: forall (f :: * -> *) a. (Functor f, Num a) => a -> f a -> f a
(*^^) a = fmap (a*)
{-# INLINE (*^^) #-}

-- | Right scalar product vector
(^^*) :: forall (f :: * -> *) a. (Functor f, Num a) => f a -> a -> f a
f ^^* a = fmap (*a) f
{-# INLINE (^^*) #-}

{- Needs Additive class and instances
basis :: (Traversable t, Num a) => [t a]
basis = basisFor (zero :: Additive v => v Int)

basisFor :: (Traversable t, Num a) => t b -> [t a]
basisFor = \t ->
   ifoldMapOf traversed ?? t $ \i _ ->
     return $ iover  traversed ?? t $ \j _ ->
         if i == j then 1 else 0
{-# INLINABLE basisFor #-}
-}

-- | Linear interpolation
lerp :: forall (f :: * -> *) a. (Functor f, Applicative f, Num a) => a -> f a -> f a -> f a
lerp alpha u v = alpha *^^ u ^$+ (1 - alpha) *^^ v
{-# INLINE lerp #-}

-- | Dot product
dot :: forall (t :: * -> *) a. (Applicative t, Foldable t, Num a) => t a -> t a -> a
dot v1 v2 = sum $ v1 ^$* v2

step :: forall (t :: * -> *) a. (Applicative t, Num a, Ord a) => t a -> t a -> t a
step = liftA2 (\a b -> if b < a then 0 else 1)

#ifdef HLINT
{-# ANN module "HLint: ignore Redundant lambda" #-}
#endif