{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE MultiWayIf #-}
{-# LANGUAGE ParallelListComp #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeFamilies #-}
{-# OPTIONS_GHC -fno-warn-name-shadowing #-}
{-# OPTIONS_GHC -fno-warn-incomplete-patterns #-}
module Language.Grappa.Inference.OldVariationalInference where

import Control.Monad ( forM, replicateM )
import Data.List ( transpose )

import Control.Lens ( (.~) )
import qualified Data.Colour as G
import qualified Data.Colour.Names as G
import qualified Data.Default.Class as G
import qualified Data.Matrix as Dm
import qualified Graphics.Rendering.Chart as G
import qualified Graphics.Rendering.Chart.Backend.Diagrams as G
import qualified Numeric.AD as Ad
import qualified Numeric.LinearAlgebra as La
import Numeric.LinearAlgebra ( Matrix, Vector, (<>), (<.>), (#>) )
import qualified Statistics.Distribution as S
import qualified Statistics.Distribution.StudentT as S
import qualified System.Random.MWC as MWC
import qualified Test.QuickCheck as Qc
import qualified Test.QuickCheck.Monadic as Qc
import qualified Test.QuickCheck.Property as Qc

import Language.Grappa.Distribution
import Language.Grappa.Inference.GradientDescent
import Language.Grappa.Rand.MWCRandM

import Debug.Trace

----------------------------------------------------------------
-- * BBVI algorithms
----------------------------------------------------------------
--
-- $bbviAlgorithms
-- Here is a summary of ways @conathan@ thought of to improve
-- BBVI. Only the first, simplest observation has been implemented, by
-- 'iteratedBbvi'.
--
-- When running BBVI to learn goal values @lambda@ that
-- are very far from the randomly chosen start values for
-- @lambda@, simply rerunning BBVI, starting from the previous
-- @lambda@, works well. For example, in
--
-- > runBbviTestUsingLazyRand 1
-- >   2      {- singleFactorBbvi2 -}
-- >   3      {- bbviTest3 -}
-- >   0.01   {- epsilon -}
-- >   100    {- num trials -}
-- >   [ 1000 {- mu -}
-- >   , 10 ] {- sigma -}
--
-- when lambda is initialized via
--
-- > replicateM 2 $ distSample (Normal 0 10)
--
-- the guesses are likely to be close to zero and far from 1000,
-- the goal value for @mu@.
--
-- The goal is to learn @lambda = [1000,10]@, and starting with
-- random lambda near zero, and then redefining @sampleLambda@ to
-- return the result, and iterating, we get the following sequence
-- of discovered lambda values:
--
-- @
-- sampleLambda = return $
--   -- [266.51484560478195,7.379162331856864]
--   -- [484.67445818993525,7.68279736309357]
--   -- [695.344938596016,8.239851183556492]
--   -- [874.483037696744,9.262297809983037]
--   -- [976.4159168688128,9.889695342218715]
--   -- [999.2564614475032,10.024071543947]
--   -- [999.9992140408281,9.999553691817553]
--   [999.99991729097,10.000438169812853]
-- @
--
-- As the result gets closer to the correct answer @[1000,10]@,
-- the incremental BBVI calls converge much more quickly. So, the
-- above corresponds to 8 runs of BBVI, but only take 3 or 4 times
-- as long as the first run.
--
-- With this improvement via iteration in mind, an obvious idea is
-- to simply implement a meta algorithm that does this iteration
-- for us: run BBVI in a loop, reusing the last @lambda@ to seed
-- the current run, until the @lambda@ computed by this meta
-- algorithm stabilizes (or some number of iterations are reached,
-- since for some functions, e.g. @sin@-like functions, we would
-- never converge).
--
-- And why does this meta algorithm work? Because when we restart
-- BBVI the learning rate / step size is large again, so we can
-- make lots of early progress before the decreasing learning rate
-- artificially slows us down.
--
-- This observation motivates a more general approach/goal: how
-- can we adjust the learning rate in a more intelligent way, so
-- that we don't slow down too early, but do slow down when we
-- reach a local optimum? ("Get while the gettin' is good" and
-- then stop.)
--
-- First, some ideas for detecting that we're not making progress
-- anymore:
--
-- 1. the difference between two consecutive values of lambda is
--    small. This is the approach the BBVI paper suggests. As
--    we've seen, it's misleading in conjunction with the
--    decreasing learning rates suggested by the paper.
--
-- 2. a whole group (tail of the sequence of lambdas) is
--    "clustered". The work here is defining "clustered". For
--    example, if the average distance (or max distance) between a
--    bunch of lambdas was close to the min distance (or average
--    distance), then the points are clustered.
--
-- 3.  the value of the objective function (ELBO) is not increasing
--     fast enough / in proportion to the effort. The work here
--     includes
--
--     - evaluating / estimating the objective function: turns out
--       this is quite similar to estimating the gradient of the
--       objective, and using a control variate should be similar
--       here. Indeed, we can define a control variate by adding a
--       multiple of the score function (as we do for estimating the
--       gradient itself.
--
--     - defining "fast enough / in proportion to the effort". For
--       example, we could define a "velocity of improvement" as the
--       ratio of the change in the objective to the step size.
--
-- Second, some ideas for ensuring we do keep making progress when
-- we can, by not decreasing the step size unnecessarily:
--
-- 1.  don't decrease the step size until the objective function stops
--     improving "fast enough", as determined by (3) above.
--
-- 2.  split the AdaGrad step size into two parts, a scalar learning
--     rate (@eta@ or @rho@ in the papers) and normalized AdaGrad
--     learning rate vector. That is, in place of
--
--     > diag(G^t)^{-1/2}
--
--     we use
--
--     > eta^t (d^t / || d^t ||_2)
--
--     for some @d^t@ based on @diag(G^t)^{-1/2}@.
--
--     While this approach punts a bit, by depending on a new
--     scalar learning rate @eta@, it simpler in that it decouples
--     the per component learning rate / emphasis part of AdaGrad,
--     a vector, from the scalar decrease in the step size. We can
--     use a notion of decreasing progress as above to wind down
--     the scalar learning rate.
--
--     For the @d^t@ we can simply use @diag(G^t)^{-1/2}@ itself,
--     or perhaps a forgetful version, e.g.
--
--     > diag(G^{t-k,t})^{-1/2}
--
--     for
--
--     > G^{t1,t2} := g^{t1} g^{t1}^T + ... + g^{t1} g^{t1}^T
--
--     i.e.
--
--     > G^t = G^{1,t}.
--
-- So, there are a lot of interesting things we could try ...

-- | The type of BBVI algorithms.
--
-- The @m@ is always monadic, but adding a @Monad m@ constraint on the
-- RHS here causes problems with "impredicative polymorphism" later,
-- when we try to make a list of @Bbvi m a@.
type Bbvi m a =
  R {-^ @epsilon@  -} ->
  Int {-^ @numSamples@ -} ->
  (Int -> [R] -> [R]) {-^ @projectLambda@ -} ->
  m [R] {-^ @sampleLambda@ -} ->
  ([R] -> m a) {-^ @sampleQ@  -} ->
  (a -> [R] -> m R) {-^ @evalLogQ@ -} ->
  (a -> [R] -> m [R]) {-^ @evalGradLoqQ@ -} ->
  (a -> m R) {-^ @evalLogP@ -} ->
  m [R]

-- | Black Box Variational Inference Algorithm 1.
--
-- From
-- http://www.cs.columbia.edu/~blei/papers/RanganathGerrishBlei2014.pdf.
--
-- We assume that @evalLogP@ is already (partially) applied to the
-- data @x@, unlike in the paper.
--
-- Here @epsilon@ is @0.01@ in the paper, and @numSamples@ is the
-- unspecified @S@ in the paper.
--
-- The @projectLambda@ function is used to project computed @lambda@
-- values back into the valid range of values for lambda.
--
-- === /Alternative approach/
--
-- As an alternative to defining a non-identity @projectLambda@, you
-- can "just" reparameterize to make @lambda@ range over all of
-- @R^n@. This is complicated and not suggested; see `bbviTest2` for
-- an example of reparameterizing.
--
-- You may be able to use functions from our 'Continuous' class to do
-- the transformations.
--
-- For example, normal distributions only support positive values for
-- sigma; we can work around this by using @fromRealLbInfty 0@ and
-- @toRealLbInfty 0@.
--
-- Alternatively, it might be possible to make the @q@-related
-- functions signal some kind of error when some part of lambda is out
-- of range. Not sure how to best recover here, but coming up with a
-- good way to recover could be much simpler than reparameterizing.
bbvi1 :: Monad m => Bbvi m a
bbvi1 epsilon numSamples projectLambda
  sampleLambda sampleQ evalLogQ evalGradLogQ evalLogP = do
  lambda <- projectLambda 1 <$> sampleLambda
  loop lambda 1
  where
    loop lambda t = do
      zs <- replicateM numSamples (sampleQ lambda)
      elboGradTerms <- forM zs $ \z -> do
        glq <- evalGradLogQ z lambda
        lq <- evalLogQ z lambda
        lp <- evalLogP z
        return $ (lp - lq) `scaleV` glq
      let elboGrad = (1 / fromIntegral numSamples) `scaleV`
                     sumV elboGradTerms
      let lambda' = projectLambda t $ lambda `addV` (rho t `scaleV` elboGrad)
      let delta = ell1 $ lambda `subV` lambda'
      if delta >= epsilon
      then loop lambda' (t + 1)
      else trace ("t = "++show t) $ return lambda'

    -- Paper says any sequence converging in @\ell^2@ and diverging in
    -- @\ell^1@ will do for @rho@ here.
    rho t = 1/fromIntegral t

-- | Single factor Black Box Variational Inference Algorithm 2,.
--
-- The n-factor algorithm can be implemented in terms of the single
-- factor algorithm, assuming the factors can be sampled
-- independently, which seems likely since their pdfs are independent
-- by assumption. The single factor Algorithm 2 is pretty similar to
-- Algorithm 1, except we use control variates to reduce the
-- variation.
--
-- See 'bbvi1' for more docs.
--
-- The use of @projectLambda@ in this implementation only corresponds
-- to the AdaGrad paper when the domain of lambda is a closed (hyper)
-- rectangle, but (partially) open rectangles should be close
-- enough. For non-rectangular domains, the projection function in the
-- AdaGrad paper takes @diagG@ into account via a "Mahalanobis norm"
-- (same as the @\ell^2@ norm for rectangular domains in this
-- context). I expect this doesn't matter much ...
--
-- TODO: change this into a stream algorithm, so that epsilon does not
-- have to be passed in, and instead the caller can just
--
-- > dropWhile (error > epsilon)
--
-- Chad suggested this approach earlier for coordinate descent.
singleFactorBbvi2 :: Monad m => Bbvi m a
singleFactorBbvi2 epsilon numSamples projectLambda
  sampleLambda sampleQ evalLogQ evalGradLogQ evalLogP = do
  lambda <- projectLambda 1 <$> sampleLambda
  let diagG = map (const 0) lambda
  loop diagG lambda 1
  where
    loop diagG lambda t = do
      zs <- replicateM numSamples (sampleQ lambda)
      hs <- forM zs $ \z -> do
        evalGradLogQ z lambda
      fs <- forM (zip zs hs) $ \(z, h) -> do
        lq <- evalLogQ z lambda
        lp <- evalLogP z
        return $ (lp - lq) `scaleV` h

      -- In the paper they say that they used 1000 for @numSamples@
      -- (called @S@ in the paper) and 100 estimates (i.e. 1/10th of
      -- the 1000) for the approximate Cov and Var in the computation
      -- of @\hat{a}^*_d@. So, we're generalizing this to using 1/10th
      -- of the samples.
      let numSamples' = if numSamples <= 10
                        then numSamples
                        else numSamples `div` 10
      let hs' = take numSamples' hs
      let fs' = take numSamples' fs
      let hs't = transpose hs'
      let fs't = transpose fs'
      let as = [ cov fCol hCol / var hCol | fCol <- fs't | hCol <- hs't ]

      let elboGradTerms = [ f `subV` aH
                          | f <- fs
                          | aH <- map (multV as) hs ]
      let elboGrad = (1 / fromIntegral numSamples) `scaleV` sumV elboGradTerms
      let diagG' = diagG `addV` (elboGrad `multV` elboGrad)
      -- Square root of pseudo inverse of @diagG'@.
      let learningRate = eta `scaleV`
                         [ if x == 0 then 0 else x**(-0.5) | x <- diagG' ]
      let dLambda = learningRate `multV` elboGrad
      let lambda' = projectLambda t $ lambda `addV` dLambda

      let delta = ell1 $ lambda `subV` lambda'
      if trace ("delta = "++show delta++"\ndiagG = "++show diagG'++"\ndLambda = "++show dLambda++"\nlambda = "++show lambda') $
         delta >= epsilon
      then loop diagG' lambda' (t + 1)
      else trace ("t = "++show t) $ return lambda'

    -- I think the AdaGrad paper said this parameter doesn't matter
    -- much (e.g. the theorem about the regret going to zero doesn't
    -- depend on an optimal choice for eta). I think this is the value
    -- they used though.
    eta = 2**0.5

    -- Note that @cov@ and @var@ are unnormalized -- i.e. not divided
    -- by the length of the input vectors -- since the normalization
    -- factors would just cancel in the calculation of @as@ as a ratio
    -- of @cov@ to @var@.
    var xs = cov xs xs
    cov xs ys = sum [ (x - meanXs) * (y - meanYs) | x <- xs | y <- ys ]
      where
        meanXs = mean xs
        meanYs = mean ys
    mean xs = sum xs / fromIntegral (length xs)

-- | Iterate BBVI until convergence or round limit is reached.
--
-- The @bbvi@ argument is expected to be a partially applied to all
-- arguments except for @sampleLambda@.
iterateBbvi ::
  Monad m =>
  Int ->
  R ->
  m [R] ->
  (m [R] -> m [R]) ->
  m [R]
iterateBbvi maxRounds epsilon sampleLambda bbvi = do
  result <- bbvi sampleLambda
  go 2 result
  where
    go round result
      | round >= maxRounds = return result
      | otherwise = do
          result' <- bbvi (return result)
          if ell1 (result `subV` result') <= epsilon
            then return result'
            else trace (show result') $ go (round + 1) result'

-- | Create an iterated version of an existing BBVI algorithm.
--
-- A helper to call @iterateBbvi@ for you, producing something of type
-- @Bbvi@' (useful for tests where I want to put BBVI algs in a list).
iteratedBbvi :: Monad m => Int -> Bbvi m a -> Bbvi m a
iteratedBbvi maxRounds bbvi epsilon numSamples projectLambda
  sampleLambda sampleQ evalLogQ evalGradLogQ evalLogP =
  iterateBbvi maxRounds epsilon sampleLambda
  (\sample -> bbvi epsilon numSamples projectLambda
              sample sampleQ evalLogQ evalGradLogQ evalLogP)

-- | A variety of BBVI algs for use in testing.
bbvis :: Monad m => [Bbvi m a]
bbvis =  [ bbvi1
         , singleFactorBbvi2
         , iteratedBbvi 100 bbvi1
         , iteratedBbvi 100 singleFactorBbvi2 ]

----------------------------------------------------------------
-- * Variational families
----------------------------------------------------------------
--
-- $variationalFamilies
-- The code needed to use a specific variational family (the @q \in
-- Q@) with a BBVI algorithm is independent of the target distribution
-- (the @p@) being approximated.

----------------------------------------------------------------
-- ** Normal

sampleNormal :: SampleableIn m Normal => [R] -> m R
sampleNormal [mu, sigma] = distSample (Normal mu sigma)

evalLogNormalDensity :: Monad m => R -> [R] -> m R
evalLogNormalDensity z [mu, sigma] = return $
  probToLogR $ distDensity (Normal mu sigma) z

evalGradLogNormalDensity :: Monad m => R -> [R] -> m [R]
evalGradLogNormalDensity z [mu, sigma] = return
  -- Calculated by hand from @normalDensity@ ...
  [ (z - mu) / sigma
  , (z - mu)**2 / sigma**3 - 1 / sigma ]

projectLambdaNormal :: Int -> [R] -> [R]
projectLambdaNormal = projectRectangle
  [ (OpenBound (-1/0), OpenBound (1/0)),
    (OpenBound 0, OpenBound (1/0)) ]

----------------------------------------------------------------
-- ** Multivariate normal
--
-- $mvNormal
-- Everywhere below @k@ is the dimension of the mv normal.

-- | Convert the flat list of params used by the bbvi algorithm into
-- the vector and matrix params used by the mv normal distribution.
paramsToDistParamsMVNormal :: Int -> [Double] -> (Vector Double, Matrix Double)
paramsToDistParamsMVNormal k muAndC =
  if length muAndC /= expectedLength
  then error $
       "paramsToDistParamsMVNormal: "++
       "the parameter list is the wrong length! Expected: "++
       show expectedLength++", actual: "++show (length muAndC)
  else (mu, c)
  where
    -- An upper triangular @k x k@ matrix has @k + 1@ choose 2
    -- potentially non-zero entries.
    expectedLength = k + ((k * (k+1)) `div` 2)

    (muList, cList) = splitAt k muAndC
    mu = La.fromList muList
    c = (k La.>< k) $ upperTriangular k cList

-- | Build full list of entries in upper triangular matrix from list
-- of its potentially non-zero entries.
--
-- Builds up the full list for the matrix by filling in zeros for
-- undefined entries. The undefined entries are the lower tringular,
-- i.e. where the row index is larger than the column index.
upperTriangular :: (Num a) => Int -> [a] -> [a]
upperTriangular k xs0 = zeroPad rcs xs0
  where
    rcs = [ (r, c) | r <- [0..(k-1)], c <- [0..(k-1)] ]
    zeroPad []           []         = []
    zeroPad ((r,c):rcs') xs@(x:xs') =
      if r > c
      then 0 : zeroPad rcs' xs
      else x : zeroPad rcs' xs'

sampleMVNormal :: SampleableIn m MVNormal => Int -> [R] -> m (Vector R)
sampleMVNormal k muAndC = distSample (MVNormal mu c)
  where
    (mu, c) = paramsToDistParamsMVNormal k muAndC

-- | The @k@ is the dimension of the mv normal.
evalLogMVNormalDensity ::
  Monad m => Int -> Vector R -> [R] -> m R
evalLogMVNormalDensity k z muAndC = return $
  probToLogR $ distDensity (MVNormal mu c) z
  where
    (mu, c) = paramsToDistParamsMVNormal k muAndC

-- | The log pdf is
--
-- > lpdf = -1/2 ( ln |C^T C| + (x-mu)^T(C^T C)^{-1}(x-mu) + k ln (2 pi) ) ,
--
-- where @|M|@ is the absolute value of the determinant of @M@.
--
-- The partials with respect to @mu@'s components are easy. The
-- product rule gives
--
-- > d_{mu_i} lpdf = -1/2 ( m^T(C^T C)^{-1}(x-mu) + (x-mu)^T(C^T C)^{-1}m )
--
-- for @m := - d_{mu_i} mu@, i.e. a vector with @m_j = - (Kronecker) delta i j@.
--
-- The partials with respect to @C@'s components are trickier, but
-- simplified by the fact that the determinant of @C@ is simple.
--
-- Let @d@ be some partial @d/dC_{ij}@. It comes down to calculating
--
-- > d ln |C^T C|
--
-- and
--
-- > d (x-mu)^T(C^T C)^{-1}(x-mu) .
--
--
-- For the first term we have @|C^T C|@ = @|C|^2@ and so
--
-- > ln |C^T C| = 2 ln |C|
--
-- and so
--
-- > d ln |C^T C| = (2/|C|) d |C| .
--
-- Since @C@ is upper triangular with positive diagonal, we have
--
-- > |C| = C_{11} .. C_{kk}
--
-- and so for @d = d/dC_{ij}@ we have
--
-- > d |C| = if i == j then |C|/C_{ij} else 0 .
--
-- and so
--
-- > d ln |C^T C| = if i == j then 2/C_{ij} else 0 .
--
--
-- For the second term we have
--
-- > (C^T C)^{-1} = (C^{-1}) (C^{-1})^T
--
-- and so
--
-- > d (C^T C)^{-1} = (d C^{-1}) (C^{-1})^T + C^{-1} (d C^{-1})^T ,
--
-- and so we just need the derivative of the inverse:
--
-- > d C^{-1} = - C^{-1} (d C) C^{-1}
--
-- see http://planetmath.org/derivativeofinversematrix.
--
-- Putting that all back together, with @B := C^{-1}@, we have
--
-- > d (C^T C)^{-1} = - (B dC B B^T + B B^T dC^T B^T)
-- >                = - B (dC B + (dC B)^T) B^T
--
-- and for @d = d/d(C_{ij})@ we have
--
-- > (dC)_{st} = delta i s * delta j t .
evalGradLogMVNormalDensity ::
  Monad m => Int -> Vector Double -> [R] -> m [R]
evalGradLogMVNormalDensity k x muAndC =
  return $ map (* (-1/2)) $ muGrads ++ cGrads
  where
    (mu, c) = paramsToDistParamsMVNormal k muAndC
    delta i j = if i == j then 1 else 0

    muGrads = [ d <.> cInvFun (x - mu) +
                (x - mu) <.> cInvFun d
              | i <- [ 0 .. k-1 ]
              , d <- [ La.fromList $ map (((-1) *) . delta i) [ 0 .. k-1 ] ] ]
    -- Efficient left multiplication of a vector by @(C^T C)^{-1}@.
    cInvFun = La.flatten . La.cholSolve c . La.asColumn

    cGrads = [ lnGrad i j + prodGrad i j
             | i <- [ 0 .. k-1 ]
             , j <- [ i .. k-1 ] ]
    lnGrad i j = if i == j then 2 / (c `La.atIndex` (i, j)) else 0
    prodGrad i j = (x - mu) <.> middle #> (x - mu)
      where
        middle = (-1) * b <> (dC i j <> b + La.tr (dC i j <> b)) <> La.tr b
    -- The inverse of @C@ as a matrix.
    (b, _) = La.invlndet c
    -- The derivative of @C@ w.r.t. @C_{ij}@.
    dC i j = La.assoc (k, k) 0 [ ((i, j), 1) ]

-- | The diagonal entries of the Cholesky matrix must be positive. The
-- other values are unconstrained.
--
-- The @k@ is the dimension of the mv normal.
projectLambdaMVNormal :: Int -> Int -> [R] -> [R]
projectLambdaMVNormal k = projectRectangle $ muBounds ++ cBounds
  where
    reals = (OpenBound (-1/0), OpenBound (1/0))
    positiveReals = (OpenBound 0, OpenBound (1/0))

    muBounds = replicate k reals
    cBounds = flip concatMap [ k, k - 1 .. 1 ] $ \rowLength ->
      take rowLength $ [ positiveReals ] ++ repeat reals

----------------------------------------------------------------
-- *** Test mv normal implementation.
--
-- $testMvNormal
--
-- The gradient we calculated by hand was quite complicated, so
-- compare it with the gradient we get from automatic differentiation
-- (would have been much faster in human time to just use automatic
-- differentiation in the first place, but now we have very high
-- confidence the gradient is correct).

-- | An overloaded (i.e. @Floating@) implementation of the log pdf.
--
-- This allows us to use automatic differentiation and compare the
-- result with the manual, by hand gradient
-- 'evalGradLogMVNormalDensity' above.
logMVNormalDensityOverloaded ::
  (Floating a, Eq a) => Int -> [a] -> [a] -> a
logMVNormalDensityOverloaded k xList muAndC =
  ((- 0.5) *) $
  2 * log detC +
  fromIntegral k * log (2 * pi) +
  -- The @Dm@ library uses 1-based indexing.
  prod Dm.! (1,1)
  where
    (muList, cList) = splitAt k muAndC
    x = Dm.fromList k 1 xList
    mu = Dm.fromList k 1 muList
    c = Dm.fromList k k $ upperTriangular k cList
    detC = foldr (*) 1 $ Dm.getDiag c
    sigma = Dm.transpose c * c
    Right sigmaInv = Dm.inverse sigma
    prod = (Dm.transpose $ x - mu) * sigmaInv * (x - mu)

-- | Use automatic differentiation to check the manually computed mv
-- normal gradient.
compareGrads :: Monad m => Int -> [R] -> [R] -> m [R]
compareGrads k x muAndC = do
  manual <- evalGradLogMVNormalDensity k (La.fromList x) muAndC
  -- AD needs the function to differentiate to overloaded / take
  -- @Floating@ arguments.
  let x' :: Floating a => [a]
      x' = map (fromRational . toRational) x
  let auto = Ad.grad (logMVNormalDensityOverloaded k x') muAndC
  -- The rescaling of errors is essential here: without it, for very
  -- large gradient entries (e.g. size 10^10), there can be large
  -- absolute differences between gradient entries. We care about
  -- large relative differences, which can't be explained away by
  -- rounding error. Dividing by 1 here avoids division by zero, and
  -- means we avoid spurious differences detected when we divide by
  -- very small numbers.
  return $ [ abs (m - a) / (abs m + abs a + 1) | m <- manual | a <- auto ]

-- | QuickCheck property asserting 'compareGrads' for given mv normal
-- dimension @k@.
prop_evalGradLogMVNormalDensity_correct :: Int -> Qc.Property
prop_evalGradLogMVNormalDensity_correct k =
  -- forAll' "k" Qc.arbitrary $ \(Qc.Positive (k :: Int)) ->
  forAll' "xList" (Qc.vector k) $ \xList ->
  forAll' "muList" (Qc.vector k) $ \muList ->
  forAll' "cList" (genCList k) $ \cList ->
  Qc.monadicIO $ do
    deltas <- Qc.run $ compareGrads k (i xList) (i $ muList ++ cList)
    Qc.monitor (Qc.counterexample $ "(Index, Delta): " ++ show (argMax deltas))
    Qc.assert $ all isSmall deltas
  where
    -- Return index of max entry, and max entry.
    argMax :: Ord a => [a] -> (Int, a)
    argMax (x:xs) = go 1 (0,x) xs
      where
        go _ p [] = p
        go j (i,x) (x':xs) = go (j+1) (if x < x' then (j, x') else (i, x)) xs

    i :: [Integer] -> [R]
    i = map fromInteger
    -- Generate a list representation of a Cholesky factor. The
    -- strategy is to recursively generate row tails from the diagonal
    -- entry onwards, with non-zero diagonal entries and arbitrary
    -- other entries.
    genCList :: Int -> Qc.Gen [Integer]
    genCList k = do
      rows <- forM [k,k-1..1] $ \len -> do
        Qc.Positive diagEntry <- Qc.arbitrary
        rest <- Qc.vectorOf (len-1) Qc.arbitrary
        return $ diagEntry : rest
      return $ concat rows

    isSmall delta = abs delta <= 0.001

    -- A version of @Qc.forAll@ with named quantifiers. This makes the
    -- counterexamples much easier to read.
    forAll' name gen pf =
      -- The library provided @Qc.forAll@ that I adapted to create
      -- this @forAll'@ has a @Qc.again@ here. But I don't understand
      -- why I need it, and more importantly, I get an import error,
      -- so leaving it out ...
      --
      -- Qc.again $
      Qc.MkProperty $
      gen >>= \x ->
        Qc.unProperty $
        Qc.counterexample (name++": "++show x) (pf x)

-- | Check the QuickCheck props @numTrials@ many times.
--
-- For larger @k@ this fails, but I think it's due to numerical
-- instability. For small @k@, e.g. less than or equal to 5, we can
-- pass hundreds of thousands of tests. For large @k@, e.g. @k = 30@,
-- we fail pretty fast.
checkProps :: Int -> Int -> IO ()
checkProps k numTrials = do
  qc $ prop_evalGradLogMVNormalDensity_correct k
  where
    qc = Qc.quickCheckWith (Qc.stdArgs { Qc.maxSuccess = numTrials })

----------------------------------------------------------------
-- * Tests
----------------------------------------------------------------

-- | Even simpler version of the 'bbv1Test2' below.
--
-- Here we approximate a normal by a normal. The standard deviation is
-- fixed, and we fit the mean. The input @mu@ and @sigma@ are the
-- parameters of the normal that we are trying to approximate; our
-- goal is to discover @mu@.
--
-- We make @mu@ the parameter to discover, because it ranges over all
-- reals and so no transformations are needed.
--
-- Returns the error in the discovered @mu@, and the discovered @mu@.
bbviTest1 ::
  (SampleableIn m Normal) =>
  Int -> R -> Int -> [R] -> m (R, [R])
bbviTest1 algNum epsilon numSamples [mu, sigma] = do
  [mu'] <- (bbvis !! (algNum - 1)) epsilon numSamples projectLambda
           sampleLambda sampleQ evalLogQ evalGradLogQ evalLogP
  let error' = ell1 $ [mu'] `subV` [mu]
  return (error', [mu'])
  where
    -- We can see that bbvi1 at least kind of works by seeding @lambda@
    -- with a value very close to the right answer.
    --
    -- sampleLambda = replicateM 1 $ distSample (Normal mu 0.1)
    sampleLambda = defaultSampleLambda 1
    sampleQ [mu'] = distSample (Normal mu' sigma)
    evalLogQ z [mu'] = return $
      probToLogR $ distDensity (Normal mu' sigma) z
    evalGradLogQ z [mu'] = return [ (z - mu') / sigma ]
    evalLogP z = return $
      probToLogR $ distDensity (Normal mu sigma) z
    projectLambda _ lambda = lambda

-- | Test 'bbv1' by approximating a linear combination of normals by a
-- normal ... except it turns out that a linear combination of normals
-- *is* a normal, and the only way to analytically weight its samples
-- is to know this, so I'm really just going to approximate a normal
-- by a normal. At least we can expect this to succeed :D
--
-- Formulas for linear combination of normals are here:
-- https://www.statlect.com/probability-distributions/normal-distribution-linear-combinations.
--
-- Note that we have to transform the standard deviation param sigma,
-- per the comments above in 'bbvi1'.
bbviTest2 ::
  (SampleableIn m Normal) =>
  Int -> R -> Int -> [R] -> m (R, [R])
bbviTest2 algNum epsilon numSamples [mu, sigma] = do
  [mu', sigma'] <- fRLambda <$>
    (bbvis !! (algNum - 1)) epsilon numSamples projectLambda
    sampleLambda sampleQ evalLogQ evalGradLogQ evalLogP
  -- In general, it might only make sense to compute the error in the
  -- real domain.
  let error' = ell1 $ [mu', sigma'] `subV` [mu, sigma]
  return (error', [mu', sigma'])
  where
    projectLambda _ lambda = lambda
    -- Since lambda is always a vector of reals, it's not clear why it
    -- should be a param to 'bbvi1'.
    --
    -- !!!: THE @fromRealLbInfty 0@ UNDERFLOWS TO ZERO FOR LARGE VALUES.
    --
    -- ???: does the delta in the output change when we increase
    -- numSampls???
    sampleLambda = defaultSampleLambda 2

    sampleQOrig = sampleNormal
    sampleQ = sampleQOrig . fRLambda

    evalLogQOrig = evalLogNormalDensity
    evalLogQ z = evalLogQOrig z . fRLambda

    -- Calculated by hand from @normalDensity@ ...
    evalGradLogQOrig = evalGradLogNormalDensity
    evalGradLogQ z lambda =
      -- Chain rule specialized to vars transformed independently
      -- (i.e. the derivative matrix of the transformation @fRLambda@
      -- is diagonal).
      zipWith (*) (fRLambda' lambda) <$>
      evalGradLogQOrig z (fRLambda lambda)

    evalLogP z = evalLogNormalDensity z [mu, sigma]

    -- Transformations from and to reals for sigma param of
    -- normal. Note: the 'Continuous' class transforms the support /
    -- output, not the params / input of a distribution.
    fRSigma = fromRealLbInfty 0
    fRSigma' = probToR . fromRealLbInfty' 0
    -- tRSigma = toRealLbInfty 0

    fRLambda [mu', sigma'R] = [mu', fRSigma sigma'R]
    -- Here @1@ is the derivative of the identity.
    fRLambda' [_mu', sigma'R] = [1, fRSigma' sigma'R]
    -- tRLambda [mu', sigma'] = [mu', tRSigma sigma']

-- | Like `bbviTest2`, but use a non-trivial `projectLambda`, instead
-- of a reparameterization.
bbviTest3 ::
  (SampleableIn m Normal) =>
  Int -> R -> Int -> [R] -> m (R, [R])
bbviTest3 algNum epsilon numSamples [mu, sigma] = do
  [mu', sigma'] <-
    (bbvis !! (algNum - 1)) epsilon numSamples projectLambda
    sampleLambda sampleQOrig evalLogQOrig evalGradLogQOrig evalLogP
  -- In general, it might only make sense to compute the error in the
  -- real domain.
  let error' = ell1 $ [mu', sigma'] `subV` [mu, sigma]
  return (error', [mu', sigma'])
  where
    -- Since lambda is always a vector of reals, it's not clear why it
    -- should be a param to 'bbvi1'.
    --
    -- !!!: THE @fromRealLbInfty 0@ UNDERFLOWS TO ZERO FOR LARGE VALUES.
    sampleLambda = defaultSampleLambda 2
    sampleQOrig = sampleNormal
    evalLogQOrig = evalLogNormalDensity
    evalGradLogQOrig = evalGradLogNormalDensity
    projectLambda = projectLambdaNormal

    evalLogP z = evalLogNormalDensity z [mu, sigma]

-- | Learn a normal that approximates a Student's t-distribution.
--
-- The parameters are @mu@ real and @nu@ positive real.
--
-- We want the mean of the learned normal to be the mean of the
-- t-distribution. I don't know what to expect for the std deviation
-- of the learned normal tho. The variance of the t-distribution is
-- @nu / (nu - 2)@ for @nu > 2@ and infinite or undefined otherwise.
--
-- See https://en.wikipedia.org/wiki/Student's_t-distribution.
bbviTest4 ::
  SampleableIn m Normal =>
  Int -> R -> Int -> [R] -> m (R, [R])
bbviTest4 algNum epsilon numSamples [mu, nu] = do
  [mu', sigma'] <-
    (bbvis !! (algNum - 1)) epsilon numSamples projectLambda
    sampleLambda sampleQ evalLogQ evalGradLogQ evalLogP
  -- In general, it might only make sense to compute the error in the
  -- real domain.
  let error' = ell1 $ [mu'] `subV` [mu]
  return (error', [mu', sigma'])
  where
    sampleQ = sampleNormal
    evalLogQ = evalLogNormalDensity
    evalGradLogQ = evalGradLogNormalDensity
    sampleLambda = defaultSampleLambda 2
    projectLambda = projectLambdaNormal

    evalLogP z = return $
      S.logDensity (S.studentT nu) (z - mu)

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

-- | Learn a multivariate normal that approximates a linear
-- regression.
--
-- The linear regression has parameters slope @m@, @y@-intercept @b@,
-- and error standard deviation @sigma@. The generative model
-- comprises a arbitrary distribution of the parameters @m@, @b@, and
-- @sigma@, say
--
-- > m ~ Uniform 0 10
-- > b ~ Uniform 0 10
-- > sigma ~ Uniform 0 10 ,
--
-- and a noisy linear distribution on the dependent variable @y@ given
-- the independent variable @x@, i.e.
--
-- > e ~ Normal 0 sigma
-- > y := m*x + b + e ,
--
-- or equivalently
--
-- > y ~ Normal (m*x + b) e .
--
-- So, the joint distribution given data
--
-- > xys = [(x1,y1),...,(xn,yn)]
--
-- is
--
-- > P[xys,m,b,sigma] = P[xys|m,b,sigma] * P[m,b,sigma] ,
--
-- where
--
-- > P[m,b,sigma] = 1
--
-- and
--
-- > P[xys|m,b,sigma] = p(x1,y1) * ... * p(xn,yn) ,
--
-- for
--
-- > p(x,y) =
-- > Pr[y = m*x + b + e] =
-- > Pr[e = y - m*x - b] =
-- > distDensity (Normal (m*x + b) sigma) y .
--
-- The @xysFlat@ argument is
--
-- > [x1,y1,x2,y2,...,xn,yn] .
bbviTest5 ::
  (SampleableIn m Normal, SampleableIn m MVNormal) =>
  Int -> R -> Int -> [R] -> m (R, [R])
bbviTest5 algNum epsilon numSamples (m:b:sigma:xysFlat) = do
  muAndC'@(m':b':sigma':_cList') <-
    (bbvis !! (algNum - 1)) epsilon numSamples projectLambda
    sampleLambda sampleQ evalLogQ evalGradLogQ evalLogP
  let error' = ell1 $ [m',b',sigma'] `subV` [m,b,sigma]
  return (error', muAndC')
  where
    -- Three dimensions, for @m@, @b@, and @sigma@.
    k = 3
    muSize = k
    cSize = (k * (k+1)) `div` 2

    sampleQ = sampleMVNormal k
    evalLogQ = evalLogMVNormalDensity k
    evalGradLogQ = evalGradLogMVNormalDensity k
    sampleLambda = defaultSampleLambda (muSize + cSize)
    projectLambda = projectLambdaMVNormal k

    xys = go xysFlat
      where
        go [] = []
        go (x:y:xysFlat') = (x,y) : go xysFlat'

    evalLogP mBSigma = return $ linRegLogDensity xys m' b' sigma'
      where
        [m', b', sigma'] = La.toList mBSigma

-- | A version of 'bbviTest5' that fixes the noise deviation sigma, so
-- that the mv normal only needs two params. The sigma passed in will
-- be used, and should correspond to the supplied data in @xysFlat@.
bbviTest7 ::
  (SampleableIn m Normal, SampleableIn m MVNormal) =>
  Int -> R -> Int -> [R] -> m (R, [R])
bbviTest7 algNum epsilon numSamples (m:b:sigma:xysFlat) = do
  muAndC'@(m':b':_cList') <-
    (bbvis !! (algNum - 1)) epsilon numSamples projectLambda
    sampleLambda sampleQ evalLogQ evalGradLogQ evalLogP
  let error' = ell1 $ [m',b'] `subV` [m,b]
  return (error', muAndC')
  where
    -- Two dimensions, for @m@, @b@.
    k = 2
    muSize = k
    cSize = (k * (k+1)) `div` 2

    sampleQ = sampleMVNormal k
    evalLogQ = evalLogMVNormalDensity k
    evalGradLogQ = evalGradLogMVNormalDensity k
    sampleLambda = defaultSampleLambda (muSize + cSize)
    projectLambda = projectLambdaMVNormal k

    xys = go xysFlat
      where
        go [] = []
        go (x:y:xysFlat') = (x,y) : go xysFlat'

    evalLogP mB = return $ linRegLogDensity xys m' b' sigma
      where
        [m', b'] = La.toList mB

linRegLogDensity ::
  [(R, R)] -> R -> R -> R -> R
linRegLogDensity xys m b sigma =
  -- An alternative approach:
  --
  -- The @abs sigma'@ is here to avoid using negative sigma
  -- values. Not sure what the right approach is here: we could
  -- use a different @projectLambda@ implementation that restricts
  -- to positive third component (i.e. the mean of the sigma
  -- values), but that won't stop them from being
  -- negative.
  {-
  sum [ logFromLogFloat (distDensity (Normal (m'*x + b') (abs sigma')) y)
      | (x, y) <- xys ]
  -}
  -- We penalize negative sigmas, by returning a very small value
  -- here.
  if sigma > 0
  then
    (if logProb < logVerySmallValue
     then trace ("XXX: logProb is very small: "++show logProb)
     else id)
    logProb
  else
    -- Is this value is really small enough? Could try @-1/0@.
    --
    -- Update: using negative infinity, or more generally very large
    -- negative values here, causes various crashes (for finite values
    -- an underflow in @LogFloat@, and for negative infinity a
    -- malformed Cholesky factor).
    logVerySmallValue
    where
      logVerySmallValue = -10e10
      logProb = sum [ probToLogR $ distDensity (Normal (m*x + b) sigma) y
                    | (x, y) <- xys ]

simpleLinRegData :: Int -> R -> R -> [(R, R)]
simpleLinRegData numPoints m b =
  [ (x, m*x+b) | x <- [ 1 .. fromIntegral numPoints ] ]

genLinRegDataIO :: Int -> R -> R -> R -> IO [(R, R)]
genLinRegDataIO numSamples m b sigma = do
  g <- MWC.create
  runRand g $ genLinRegData numSamples m b sigma

genLinRegData ::
  ( SampleableIn m Normal, SampleableIn m Uniform ) =>
  Int -> R -> R -> R -> m [(R, R)]
genLinRegData numSamples m b sigma = do
  replicateM numSamples $ do
    -- Use integral @x@s to make the output a little easier to inspect
    -- visually.
    x <- fromInteger . round <$> distSample (Uniform (-1000) 1000)
    y <- distSample (Normal (m*x + b) sigma)
    return (x,y)

-- scatterPlot :: [(R, R)] -> ???
-- | Based on http://indiana.edu/~ppaml/HakaruTutorial.html and
-- https://github.com/timbod7/haskell-chart/wiki/example%202.
--
-- Writes output to an SVG file, regardless of supplied file extension
-- :P The examples I referred to use the Cairo backend, but I'm using
-- the Diagrams backend, so that might be the issue. So, I changed
-- this function to supply the extension, to avoid confusion.
scatterPlot :: String -> String -> [(R, R)] -> IO (G.PickFn ())
scatterPlot title file xys =
  G.renderableToFile G.def (file++".svg") $ G.toRenderable layout
  where
    plot = G.plot_points_style .~ G.filledCircles 2 (G.opaque G.purple)
           $ G.plot_points_values .~ xys
           $ G.def

    layout = G.layout_title .~ title
             $ G.layout_plots .~ [G.toPlot plot]
             $ G.layout_x_axis . G.laxis_generate .~ G.scaledAxis G.def (-4, 4)
             $ G.layout_y_axis . G.laxis_generate .~ G.scaledAxis G.def (-4, 4)
             $ G.def

-- | https://en.wikipedia.org/wiki/File:MultivariateNormal.png
--
-- I can't figure out how to set the aspect ratio :P
wikipediaExample :: IO ()
wikipediaExample = do
  g <- MWC.create
  xyVectors <- runRand g $ replicateM numSamples (distSample (MVNormal mu c))
  let xys = [ (x, y) | v <- xyVectors, [x,y] <- [ La.toList v ] ]
  _ <- scatterPlot ("Wikipedia Test ("++show numSamples++" samples)") "/tmp/out" xys
  return ()
  where
    mu = La.fromList [0, 0]
    c = La.chol (La.trustSym $ (2 La.>< 2) [1, 0.6, 0.6, 2])
    numSamples = 10000

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

-- | Sanity test: match an mv normal to an mv normal.
--
-- This works pretty well with iterated bbvi 2, e.g. with
--
-- > runBbviTestUsingLazyRand 1 4 6 1 1000 [1,2,3,1,0,0,2,0,3]
--
-- Of course, this doesn't prove much.
bbviTest6 ::
  (SampleableIn m Normal, SampleableIn m MVNormal) =>
  Int -> R -> Int -> [R] -> m (R, [R])
bbviTest6 algNum epsilon numSamples muAndC = do
  muAndC' <-
    (bbvis !! (algNum - 1)) epsilon numSamples projectLambda
    sampleLambda sampleQ evalLogQ evalGradLogQ evalLogP
  let error' = ell1 $ muAndC' `subV` muAndC
  return (error', muAndC')
  where
    k = 3
    muSize = k
    cSize = (k * (k+1)) `div` 2

    sampleQ = sampleMVNormal k
    evalLogQ = evalLogMVNormalDensity k
    evalGradLogQ = evalGradLogMVNormalDensity k
    sampleLambda = defaultSampleLambda (muSize + cSize)
    projectLambda = projectLambdaMVNormal k

    evalLogP z = evalLogMVNormalDensity k z muAndC

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

runBbviTestUsingLazyRand ::
  Int -> Int -> Int -> R -> Int -> [R] -> IO [(R, [R])]
runBbviTestUsingLazyRand numTrials algNum testNum epsilon numSamples testParams = do
  g <- MWC.create
  runRand g $ replicateM numTrials $ test algNum epsilon numSamples testParams
  where
    test = [ bbviTest1
           , bbviTest2
           , bbviTest3
           , bbviTest4
           , bbviTest5 -- Mv normal for linear regression: learn m, b, sigma.
           , bbviTest6 -- Mv normal for mv normal.
           , bbviTest7 -- Mv normal for linear regression: learn m, b.
           ] !! (testNum - 1)

test5 :: IO [(R, [R])]
test5 = runBbviTestUsingLazyRand 1 3 5 0.1 1000
  ([m,b,sigma] ++ concat [ [x, y] | (x, y) <- simpleLinRegData 1000 m b ])
  where
    m = 1
    b = 2
    sigma = 0.001


{-# LANGUAGE FlexibleContexts, NamedFieldPuns #-}

module School.Unit.CostFunction
( SetupCost
, CostFunction(..)
, defSetupCost
) where

import Conduit ((.|), ConduitM, ZipSource(..), getZipSource, mapC, mapMC)
import Numeric.LinearAlgebra (Container, Matrix, Vector, add, cols, rows)
import School.Types.Slinky (Slinky(..), slinkyConcat)
import School.Unit.CostParams (CostParams)
import School.Unit.UnitActivation (UnitActivation(..))
import School.Unit.UnitForward (ForwardStack)
import School.Unit.UnitGradient (UnitGradient(..))
import School.Utils.LinearAlgebra (zeroMatrix)

type SetupCost a m = ConduitM () (Matrix a) m ()
                  -> ConduitM () (ForwardStack a) m ()

defSetupCost :: (Monad m) => SetupCost a m
defSetupCost source =
  source .| mapMC (\matrix -> return ([BatchActivation matrix], SNil))

type CompupteCost a = UnitActivation a
                   -> Slinky CostParams
                   -> Either String a

type DerivCost a = UnitActivation a
                -> Slinky CostParams
                -> Either String (UnitGradient a)

type PrepareCost m a =
  ConduitM (ForwardStack a) (ForwardStack a) m ()

data CostFunction a m =
  CostFunction { computeCost :: CompupteCost a
               , derivCost :: DerivCost a
               , prepareCost :: PrepareCost m a
               , setupCost :: SetupCost a m
               }

mappendCompute :: (Num a)
               => CompupteCost a
               -> CompupteCost a
               -> CompupteCost a
mappendCompute c1 c2 activations params1@(SNode _ params2) = do
  res1 <- c1 activations params1
  res2 <- c2 activations params2
  return $ res1 + res2
mappendCompute _ _ _ _ = Left "mappend costfunctions needs two cost params"

mappendDeriv :: (Container Vector a, Num a)
             => DerivCost a
             -> DerivCost a
             -> DerivCost a
mappendDeriv d1 d2 activations params1@(SNode _ params2) = do
  (BatchGradient res1) <- d1 activations params1
  (BatchGradient res2) <- d2 activations params2
  return $ BatchGradient $ add res1 res2
mappendDeriv _ _ _ _ = Left "mappend costfunctions needs two cost params"

combinator :: ForwardStack a -> ForwardStack a -> ForwardStack a
combinator (_, cParams1) (activations, cParams2) =
  (activations, slinkyConcat cParams1 cParams2)

mappendSetup :: (Monad m)
             => SetupCost a m
             -> SetupCost a m
             -> SetupCost a m
mappendSetup s1 s2 source = getZipSource $ combinator
                                       <$> (ZipSource $ s1 source)
                                       <*> (ZipSource $ s2 source)

instance (Container Vector a, Num a, Monad m) => Monoid (CostFunction a m) where
  mappend CostFunction { computeCost = c1
                       , derivCost = d1
                       , prepareCost = p1
                       , setupCost = s1
                       }
          CostFunction { computeCost = c2
                       , derivCost = d2
                       , prepareCost = p2
                       , setupCost = s2
                       } =
   CostFunction { computeCost = mappendCompute c1 c2
                , derivCost = mappendDeriv d1 d2
                , prepareCost = p1 .| p2
                , setupCost = mappendSetup s1 s2
                }
  mempty = CostFunction { computeCost
                        , derivCost
                        , prepareCost = mapC id
                        , setupCost = defSetupCost
                        } where
    computeCost _ _ = Right 0
    derivCost (BatchActivation g) _ =
      Right . BatchGradient $ zeroMatrix r c
      where r = rows g
            c = cols g
    derivCost _ _ = Left "mempty deriv cost expects batch activation"


{-# LANGUAGE RankNTypes #-}
module Main where

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

type Pendulum a b = (a, b)

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

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

type Pair a = (a, a)

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

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

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

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

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

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

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

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

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

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

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

skip :: Int
skip = 2

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

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


{-# LANGUAGE MultiParamTypeClasses #-}
module ChaosBox.Geometry.Class
  ( Boundary(..)
  , HasP2(..)
  , Intersects(..)
  , getP2
  , setP2
  , modifyP2
  )
where

import           ChaosBox.Geometry.P2
import           Control.Lens         (Lens', lens, (%~), (&), (.~), (^.))
import           Data.Complex
import           Linear.V2

-- | Class of objects that can be queried for points
class Boundary a where
  containsPoint :: a -> P2 -> Bool

class HasP2 a where
  _V2 :: Lens' a P2

instance HasP2 P2 where
  _V2 = _xy

instance HasP2 (Complex Double) where
  _V2 = lens (\(a :+ b) -> V2 a b) (\_ (V2 x y) -> x :+ y)

getP2 :: HasP2 a => a -> P2
getP2 = (^. _V2)

setP2 :: HasP2 a => a -> P2 -> a
setP2 x p2 = x & _V2 .~ p2

modifyP2 :: HasP2 a => a -> (P2 -> P2) -> a
modifyP2 x f = x & _V2 %~ f

class Intersects a b where
  intersectionPoints :: a -> b -> [P2]
  intersects :: a -> b -> Bool
  intersects a b = null (intersectionPoints a b)


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

module Foreign.BLAS.Zomplex
    where

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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


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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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


{-# LANGUAGE FlexibleContexts           #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE RecordWildCards            #-}
{-# LANGUAGE ScopedTypeVariables        #-}
{-# LANGUAGE StandaloneDeriving         #-}

module Main (main)
where

import           Phase
import           Random

import           Control.Monad                     as CM
import qualified Control.Monad.Log                 as ML
import           Control.Monad.Random
import qualified Control.Monad.State.Strict        as MS
import           Data.Complex
import qualified Data.List                         as L
import qualified Data.Massiv.Array                 as A
import qualified Data.Massiv.Array.Manifest.Vector as AMV
import qualified Data.MultiSet                     as M
import qualified Data.Vector.Unboxed               as V
import           Numeric
import           Prelude                           as P
import           System.Console.ANSI
import           System.Environment
import           System.TimeIt

-- | a vector of multiplication phases; should be kept sorted
type PhaseVector = A.Array A.U A.Ix1 Phase

-- | the massiv computation strategy to use throughout
strat :: A.Comp
strat = A.Par

fromVector :: Vector Phase -> PhaseVector
fromVector vec = AMV.fromVector' strat (A.Sz1 (V.length vec)) vec

-- | a tuple of phase vectors to be collimated
type PVTuple = (PhaseVector, PhaseVector)

-- what LoggingT couldn't
deriving instance (MonadRandom m) => MonadRandom (ML.LoggingT message m)

-- | the base-2 logarithm
log2 :: Integral a => a -> Double
log2 = logBase (2 :: Double) . fromIntegral

-- | for convenience
log2rnd :: Integral a => a -> Int
log2rnd  = round . log2

dot :: Num a => [a] -> [a] -> a
dot [] _          = 0
dot _ []          = 0
dot (x:xs) (y:ys) = x*y + dot xs ys

---------- COLLIMATION ----------

data CollimateInfo = SI { s     :: !Phase, lReq :: !Int,
                          lReal :: !Int,   kept :: !Bool }

instance Show CollimateInfo where
  show SI{..} =
    let lRatio = fromIntegral lReal / fromIntegral lReq :: Double
    in "(~2^" P.++ show (log2rnd s) P.++
       "," P.++ show lReq P.++
       "," P.++ show lReal P.++
       "," P.++ showFFloat (Just 2) lRatio
       (if not kept then " DISCARDED" else "") P.++ ")"

collimateInfoHandler :: MonadIO m
                     => Int     -- | number of bottom sieve layers to suppress
                     -> [Phase] -- | list of intervals sizes for sieve
                     -> CollimateInfo
                     -> m ()
collimateInfoHandler skip ss =
  let len = P.length ss in \si@SI{..} ->
    let i' = L.elemIndex s ss
    in maybe (return ())
       (\i -> when (i < len - skip) $ liftIO $ do
           setSGR [SetColor Foreground (if i < 7 then Vivid else Dull)
                   (toEnum (i `mod` 7 + 1))]
           putStrLn $ P.concat (P.replicate i "  ") P.++ show si
           setSGR [Reset]) i'

-- | all the mod-@n@ subset-sums of a vector. Results are not
-- necessarily sorted!
subsetSums :: Vector Phase -> Vector Phase
subsetSums v =
  V.generate (2^V.length v) (go 0 0) where
  go i acc s
    | i >= V.length v = acc
    | otherwise = go (i+1) (acc + if odd s then v V.! i else 0) (s `div` 2)

-- | choose a random phase multiplier from the given phase vector
randomElt :: MonadRandom m => PhaseVector -> m Phase
randomElt v =  (v A.!) <$> getRandomR (0, A.elemsCount v - 1)

-- | the first index (in a sorted vector) of an element that is /at
-- least/ @a@, or 'length v' if none exists
findAtLeast :: PhaseVector -> Phase -> Int
findAtLeast v a = go 0 (A.elemsCount v) where
  go i j | i == j         = i
         | (v A.! i) >= a = i
         -- don't use i+j `div` 2, to avoid overflow on big arrays
         | otherwise      = let k = i + ((j - i) `div` 2)
                            in if (v A.! k) >= a then go i k else go (k+1) j

-- | collimate phase vectors
collimate :: (MonadRandom m)
          => Phase              -- | desired interval size S (upper bound)
          -> PVTuple            -- | phase vectors to collimate
          -> m PhaseVector
collimate s (v1,v2) = do
  b1 <- randomElt v1
  b2 <- randomElt v2
  let q = (b1 + b2) `div` s
      qs = q*s
      start i1 = (i1, findAtLeast v2 (qs - (v1 A.! i1)))
      gen (i1,i2)
        | i1 >= A.elemsCount v1 = Nothing
        | i2 >= A.elemsCount v2 = gen $ start (i1+1)
        | v < qs + s            = Just (v `mod` s, (i1, i2+1))
        | otherwise             = gen $ start (i1+1)
        where v = (v1 A.! i1) + (v2 A.! i2)
      -- Create a Vector because massiv doesn't supported
      -- unknown-length unfoldr.
  return $ A.quicksort $ fromVector $ V.unfoldr gen (start 0)

---------- SIEVE ----------

data SieveState = SS { numQueries :: !Int, maxLength    :: !Int,
                       numNodes   :: !Int, numDiscarded :: !Int}
  deriving (Show, Eq)

newSieveState :: SieveState
newSieveState = SS 0 0 0 0

addQueries :: Int -> SieveState -> SieveState
addQueries i ss@SS{..} = ss { numQueries = numQueries + i }

updateLength :: Int -> SieveState -> SieveState
updateLength l ss@SS{..} = ss { maxLength = max l maxLength }

incrementNodes :: SieveState -> SieveState
incrementNodes ss@SS{..} = ss { numNodes = succ numNodes }

incrementDiscarded :: SieveState -> SieveState
incrementDiscarded ss@SS{..} = ss { numDiscarded = succ numDiscarded }

-- | Collimation sieve.
sieve :: (MonadRandom m, ML.MonadLog CollimateInfo m,
          MS.MonadState SieveState m)
      => Phase   -- | group order N
      -> Double  -- | threshold factor for long-enough phase vectors
      -> [Phase] -- | increasing interval sizes S_i for the i'th level
                 -- of the sieve; the final one should equal N
      -> Int     -- | desired phase vector length L
      -> m PhaseVector
sieve = sieve' True where       -- don't discard final output
  sieve' _ _ _ [] _ = error "sieve: empty list of interval sizes"
  sieve' alwaysKeep n threshold ss@(s:ss') l
    | s >= n = do
        let logl = log2rnd l
        v <- (A.quicksort . fromVector . V.map (`mod` n)) . subsetSums <$>
          -- create a Vector because massiv doesn't support monadic replicate
             V.replicateM logl (getRandomR (0,n-1))
        ML.logMessage (SI s l (A.elemsCount v) True)
        MS.modify'   incrementNodes
        MS.modify' $ addQueries logl
        MS.modify' $ updateLength (A.elemsCount v)
        return v
    | otherwise =
        let s' = P.head ss'
            z  = 1.5 * fromIntegral l * fromIntegral s' / fromIntegral s :: Double
            l' = ceiling $ sqrt z
        in do v1   <- sieve' False n threshold ss' l'
              v2   <- sieve' False n threshold ss' $
                      ceiling (z / fromIntegral (A.elemsCount v1))
              v    <- collimate s (v1,v2)
              let vlen = A.elemsCount v
                  -- check if vector is long enough to be useful
                  keep = alwaysKeep ||
                         fromIntegral vlen / fromIntegral l >= threshold
              ML.logMessage (SI s l vlen keep)
              MS.modify' incrementNodes
              if keep
                then MS.modify' (updateLength vlen) >> return v
                else MS.modify' incrementDiscarded >> sieve' False n threshold ss l

-- | histogram of (n, # multipliers appearing n times) in a phase
-- vector for a (final, small) interval [S]
histogram :: Int -> Vector Int -> [(Int,Int)]
histogram s v =
  let occurs = fmap snd $ M.toOccurList $ M.fromList $ V.toList v
      num    = P.length occurs
      hist   = L.sortOn fst $ M.toOccurList $ M.fromList occurs
  in if s == num then hist else (0, s - num) : hist

chi :: Double -> Complex Double
chi = cis . (2.0 * pi *)

square :: Num a => a -> a
square x = x*x

-- | partition a phase vector into one consisting of its unique phase
-- multipliers, and one containing all the leftovers.
puncture :: PhaseVector -> (PhaseVector, PhaseVector)
puncture v =
  let v' = AMV.toVector v
      u' = V.uniq v'
      gen i
        | i >= V.length v' = Nothing
        | v' V.! (i-1) == v' V.! i = Just (v' V.! i, i+1)
        | otherwise = gen $ i+1
  in (fromVector u', fromVector $ V.unfoldr gen 1)

unfold :: (b -> (a,b)) -> b -> [a]
unfold f b = let (a,b') = f b in a : unfold f b'

-- | probability assigned to \( w \), given by \( \theta = s/N - w/T
-- \), by the \( T \)-dimensional inverse-QFT of the given punctured
-- phase vector on \( [T] \).
probTheta :: Int                -- | range \( [T] \) of phase vector
          -> PhaseVector        -- | uniquified phase vector
          -> Double             -- | \( \theta \)
          -> Double
probTheta t v theta =
  square (magnitude $ A.sum $ A.map (chi . (* theta) . fromIntegral) v) /
  (fromIntegral t * fromIntegral (A.elemsCount v))

-- | probability that the \( T \)-dimensional inverse-QFT on the given
-- punctured phase vector on \( [T] \) will output some \( w \in
-- \lfloor sT/N \rceil - \{ - \lfloor z/2 \rfloor, \ldots, \lfloor
-- (z-1)/2 \rfloor \} \), given the (essentially uniform) shift \(
-- sT/N \bmod 1 \in [-1/2,1/2) \).  Assumes that (\ z \) is positive
-- and small enough not to double-count probabilities.
probClosest :: Int               -- | \( T \)
            -> Double            -- | \( sT/N \bmod 1 \in [-1/2,1/2) \)
            -> PhaseVector       -- | uniquified phase vector
            -> Int               -- | number \( z \) of closest to count
            -> Double
probClosest t shift v z =
  let thetas = A.makeArrayR A.D strat (A.Sz1 z)
               (\(A.Ix1 i) -> (shift + fromIntegral (i - z `div` 2))
                              / fromIntegral t)
  in A.sum $ A.map (probTheta t v) thetas

main :: IO ()
main = do
  args <- getArgs
  let (params,rest) = splitAt 3 args
      [logn :: Int, logl, logs] = read <$> params
      threshold = if null rest then 0.25 :: Double else read (head rest)
      n     = if logn > 0
              then 2^logn
                   -- from https://eprint.iacr.org/2019/498.pdf
              else 3 * 37 * 1407181 * 51593604295295867744293584889 *
                   31599414504681995853008278745587832204909
      l     = 2^logl
      s     = 2^logs
      ss    = P.takeWhile (< n)
              (iterate (\t -> (2 * t * fromIntegral l) `div` 3) $ fromIntegral s)
              P.++ [n]

  -- print schedule of log interval sizes
  putStrLn $ "log S's = " P.++ show (log2rnd <$> ss)

  -- run the sieve
  (time, (v, sieveState@SS{..})) <- timeItT $
    evalCryptoRandIO
    (flip ML.runLoggingT (collimateInfoHandler (min 5 (P.length ss - 1)) ss)
     (flip MS.runStateT newSieveState (sieve n threshold ss l)))

  -- print results

  -- parameters again, for convenience
  putStrLn $ "\n[log N, log L, log S_0] = " P.++
    show [logn, logl, logs]

  putStrLn $ "threshold = " P.++ showFFloat (Just 2) threshold ""

  -- log interval sizes again, for convenience
  putStrLn $ "\nlog S's = " P.++ show (log2rnd <$> ss)

  putStrLn $ "\nSieve summary = " P.++ show sieveState

  -- compute number of queries according to model
  let delta = fromIntegral numDiscarded / fromIntegral numNodes :: Double
      depth = P.length ss - 1
      l' = sqrt $ 1.5 * fromIntegral l * fromIntegral n /
           fromIntegral (ss P.!! (P.length ss - 2)) :: Double
      modelQueries = (2.0 / (1-delta))^depth * logBase 2 l'

  putStrLn $ "\nProbability of discarding = " P.++ showFFloat (Just 4) delta ""

  putStrLn $ (P.++) "\nNumber of queries actual/modeled = " $ (P.++)
    (show numQueries)                $ (P.++) "/" $
    showFFloat (Just 1) modelQueries $ (P.++) " ~= " $
    showFFloat (Just 2) (fromIntegral numQueries / modelQueries) ""

  -- compute probability of obtaining a regular state
  let hist = histogram s $ V.map fromIntegral $ AMV.toVector v
      num  = s * fst (P.head hist)
      den  = A.elemsCount v
      probRegular = fromIntegral num / fromIntegral den :: Double

  putStrLn $ "\nProbability of obtaining a regular state = " P.++
    show num P.++ "/" P.++ show den P.++ " ~= " P.++
    showFFloat (Just 3) probRegular ""

  -- compute probabilities of getting punctured regular states, and close
  shift :: Double <- getRandomR (-0.5, 0.5)
  let punctureds = takeWhile ((> 2^(logs - 6)) . A.elemsCount) $
                   unfold puncture v
      puncturedProbs :: [Double] =
        (/ fromIntegral den) . fromIntegral . A.elemsCount <$> punctureds
      closest = [1,2,4,8]       -- how many of the closest w to check
      closestProbs = map (<$> closest) (probClosest s shift <$> punctureds)
      totalClosestProbs = dot puncturedProbs <$> L.transpose closestProbs

  putStrLn $ (P.++) "\nProbability of obtaining each punctured state (1st, 2nd, ... attempt) =\n" $
    L.intercalate ", " $ mapM (showFFloat (Just 3)) puncturedProbs ""

  putStrLn $ (P.++) "\nProbability bounds for each punctured state for " $
    (P.++) (show closest) $ (P.++) " closest w = \n" $
    L.intercalate "\n" $ L.intercalate ", " <$>
    (mapM . mapM) (showFFloat (Just 3)) closestProbs ""

  putStrLn $ (P.++) "\nTotal probability bounds for " $
    (P.++) (show closest) $ (P.++) " closest w = \n" $
    L.intercalate ", " $ mapM (showFFloat (Just 3)) totalClosestProbs ""

  putStrLn $ (P.++) "\nLaTeX table row:\n" $
    showFFloat (Just 1) (log2 n)                          $ (P.++) " & " $
    showFFloat (Just 1) (log2 numQueries)                 $ (P.++) " & " $
    showFFloat (Just 1) (logBase 2 modelQueries)          $ (P.++) " & " $
    showFFloat (Just 1) (log2 maxLength)                  $ (P.++) " & " $
    (P.++) (show logl)                                    $ (P.++) " & " $
    (P.++) (show logs)                                    $ (P.++) " & " $
    showFFloat (Just 0) (probRegular * 100)               $ (P.++) " & " $
    showFFloat (Just 1) (fromIntegral logs * probRegular) $ (P.++) " & " $
    showFFloat (Just 2) threshold                         $ (P.++) " & " $
    showFFloat (Just 1) (delta * 100)                     $ (P.++) " & " $
    (P.++) (show depth)                                   $ (P.++) " & " $
    showFFloat (Just 1) (time / 3600) " \\\\"


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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

typeSubst = typeSubstD 0

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

typeSubstC = typeSubstCD 0

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

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

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

instance Eq Id where
  _ == _ = True

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


module Numeric.LinearAlgebra.Matrix
( module C
, module M44
, module M33
, Mat44
, Mat33
) where

import Numeric.LinearAlgebra.Matrix.Class as C
import Numeric.LinearAlgebra.Matrix.Mat44 as M44 hiding (Mat44(..))
import Numeric.LinearAlgebra.Matrix.Mat44 (Mat44)
import Numeric.LinearAlgebra.Matrix.Mat33 as M33 hiding (Mat33(..))
import Numeric.LinearAlgebra.Matrix.Mat33 (Mat33)


module QIO.Circuit.Compile where

import qualified QIO.Circuit.Circuit as C

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

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

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

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

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

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

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


{-# LANGUAGE ScopedTypeVariables #-}
module Main where

import Numeric.LinearAlgebra
import Control.DeepSeq
import Control.Concurrent
import Control.Monad
import Control.Exception
import System.Mem
import Control.Concurrent.Async
import System.Environment

main :: IO ()
main = do
  [matSizeS] <- getArgs
  let matSize = read matSizeS
  let gcs = do
        performMinorGC
        performMajorGC
        threadDelay 1000
  caps <- getNumCapabilities
  void $ concurrently gcs $ flip mapConcurrently [1..caps] $ \i0 -> do
    let matrices :: Int -> [Matrix Double]
        matrices i = if i == 10000
          then matrices (-10000)
          else let
            a :: Matrix Double = (matSize >< matSize) (take (matSize * matSize) [fromIntegral i..])
            in a : matrices (i+1)
    forM_ (matrices i0) $ \a -> do
      let (l, m, r) = svd a
      evaluate (force (l <> diagRect 0 m matSize matSize <> tr r))