{-# OPTIONS_GHC -Wall -Wno-orphans #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TemplateHaskell #-}
module TestZDD (zddTestGroup) where

import Control.DeepSeq
import Control.Monad
import Control.Monad.ST
import Data.IntMap (IntMap)
import qualified Data.IntMap as IntMap
import Data.IntSet (IntSet)
import qualified Data.IntSet as IntSet
import Data.IORef
import Data.List
import qualified Data.Map.Strict as Map
import Data.Proxy
import Data.Set (Set)
import qualified Data.Set as Set
import Data.Vector (Vector)
import Data.Word
import qualified GHC.Exts as Exts
import Statistics.Distribution
import Statistics.Distribution.ChiSquared (chiSquared)
import System.IO.Unsafe
import qualified System.Random.MWC as Rand
import Test.QuickCheck.Function (apply)
import Test.QuickCheck.Instances.Vector ()
import Test.Tasty
import Test.Tasty.HUnit
import Test.Tasty.QuickCheck
import Test.Tasty.TH

import Data.DecisionDiagram.BDD.Internal.ItemOrder (OrderedItem (..))
import Data.DecisionDiagram.ZDD (ZDD (..), ItemOrder (..))
import qualified Data.DecisionDiagram.ZDD as ZDD

import Utils

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

instance ZDD.ItemOrder a => Arbitrary (ZDD a) where
  arbitrary = do
    vars <- liftM (sortBy (ZDD.compareItem (Proxy :: Proxy a)) . IntSet.toList) arbitrary
    let f vs n = oneof $
          [ return ZDD.empty
          , return ZDD.base
          ]
          ++
          [ do v <- elements vs
               let vs' = dropWhile (\v' -> compareItem (Proxy :: Proxy a) v' v  /= GT) vs
               p0 <- f vs' (n `div` 2)
               p1 <- f vs' (n `div` 2) `suchThat` (/= ZDD.empty)
               return (ZDD.Branch v p0 p1)
          | n > 0, not (null vs)
          ]
    sized (f vars)

  shrink (ZDD.Empty) = []
  shrink (ZDD.Base) = []
  shrink (ZDD.Branch x p0 p1) =
    [p0, p1]
    ++
    [ ZDD.Branch x p0' p1'
    | (p0', p1') <- shrink (p0, p1), p1' /= ZDD.empty
    ]

arbitraryMember :: ZDD.ItemOrder a => ZDD a -> Gen IntSet
arbitraryMember zdd = do
  (seed :: Vector Word32) <- arbitrary
  return $ runST $ do
    gen <- Rand.initialize seed
    ZDD.uniformM zdd gen

-- ------------------------------------------------------------------------
-- Union
-- ------------------------------------------------------------------------

prop_union_unitL :: Property
prop_union_unitL =
  forAllItemOrder $ \(_ :: Proxy o) ->
    forAll arbitrary $ \(a :: ZDD o) ->
      ZDD.empty `ZDD.union` a === a

prop_union_unitR :: Property
prop_union_unitR =
  forAllItemOrder $ \(_ :: Proxy o) ->
    forAll arbitrary $ \(a :: ZDD o) ->
      a `ZDD.union` ZDD.empty === a

prop_union_comm :: Property
prop_union_comm =
  forAllItemOrder $ \(_ :: Proxy o) ->
    forAll arbitrary $ \(a :: ZDD o, b) ->
      a `ZDD.union` b === b `ZDD.union` a

prop_union_assoc :: Property
prop_union_assoc =
  forAllItemOrder $ \(_ :: Proxy o) ->
    forAll arbitrary $ \(a :: ZDD o, b, c) ->
      a `ZDD.union` (b `ZDD.union` c) === (a `ZDD.union` b) `ZDD.union` c

prop_union_idempotent :: Property
prop_union_idempotent =
  forAllItemOrder $ \(_ :: Proxy o) ->
    forAll arbitrary $ \(a :: ZDD o) ->
      a `ZDD.union` a === a

-- ------------------------------------------------------------------------
-- Intersection
-- ------------------------------------------------------------------------

prop_intersection_comm :: Property
prop_intersection_comm =
  forAllItemOrder $ \(_ :: Proxy o) ->
    forAll arbitrary $ \(a :: ZDD o, b) ->
      a `ZDD.intersection` b === b `ZDD.intersection` a

prop_intersection_assoc :: Property
prop_intersection_assoc =
  forAllItemOrder $ \(_ :: Proxy o) ->
    forAll arbitrary $ \(a :: ZDD o, b, c) ->
      a `ZDD.intersection` (b `ZDD.intersection` c) === (a `ZDD.intersection` b) `ZDD.intersection` c

prop_intersection_idempotent :: Property
prop_intersection_idempotent =
  forAllItemOrder $ \(_ :: Proxy o) ->
    forAll arbitrary $ \(a :: ZDD o) ->
      a `ZDD.intersection` a === a

prop_intersection_emptyL :: Property
prop_intersection_emptyL =
  forAllItemOrder $ \(_ :: Proxy o) ->
    forAll arbitrary $ \(a :: ZDD o) ->
      ZDD.empty `ZDD.intersection` a === ZDD.empty

prop_intersection_emptyR :: Property
prop_intersection_emptyR =
  forAllItemOrder $ \(_ :: Proxy o) ->
    forAll arbitrary $ \(a :: ZDD o) ->
      a `ZDD.intersection` ZDD.empty === ZDD.empty

prop_intersection_distL :: Property
prop_intersection_distL =
  forAllItemOrder $ \(_ :: Proxy o) ->
    forAll arbitrary $ \(a :: ZDD o, b, c) ->
      a `ZDD.intersection` (b `ZDD.union` c) === (a `ZDD.intersection` b) `ZDD.union` (a `ZDD.intersection` c)

prop_intersection_distR :: Property
prop_intersection_distR =
  forAllItemOrder $ \(_ :: Proxy o) ->
    forAll arbitrary $ \(a :: ZDD o, b, c) ->
      (a `ZDD.union` b) `ZDD.intersection` c === (a `ZDD.intersection` c) `ZDD.union` (b `ZDD.intersection` c)

-- ------------------------------------------------------------------------
-- Difference
-- ------------------------------------------------------------------------

prop_difference_cancel :: Property
prop_difference_cancel =
  forAllItemOrder $ \(_ :: Proxy o) ->
    forAll arbitrary $ \(a :: ZDD o) ->
      a ZDD.\\ a === ZDD.empty

prop_difference_unit :: Property
prop_difference_unit =
  forAllItemOrder $ \(_ :: Proxy o) ->
    forAll arbitrary $ \(a :: ZDD o) ->
      a ZDD.\\ ZDD.empty === a

prop_union_difference :: Property
prop_union_difference =
  forAllItemOrder $ \(_ :: Proxy o) ->
    forAll arbitrary $ \(a :: ZDD o, b, c) ->
      (a `ZDD.union` b) ZDD.\\ c === (a ZDD.\\ c) `ZDD.union` (b ZDD.\\ c)

prop_difference_union :: Property
prop_difference_union =
  forAllItemOrder $ \(_ :: Proxy o) ->
    forAll arbitrary $ \(a :: ZDD o, b, c) ->
      a ZDD.\\ (b `ZDD.union` c) === (a ZDD.\\ b) ZDD.\\ c

-- ------------------------------------------------------------------------
-- Non-superset
-- ------------------------------------------------------------------------

prop_nonSuperset :: Property
prop_nonSuperset =
  forAllItemOrder $ \(_ :: Proxy o) ->
    forAll arbitrary $ \(a :: ZDD o, b) ->
      let a' = ZDD.toSetOfIntSets a
          b' = ZDD.toSetOfIntSets b
          p xs = and [not (IntSet.isSubsetOf ys xs) | ys <- Set.toList b']
       in ZDD.toSetOfIntSets (a `ZDD.nonSuperset` b) === Set.filter p a'

prop_nonSuperset_cancel :: Property
prop_nonSuperset_cancel =
  forAllItemOrder $ \(_ :: Proxy o) ->
    forAll arbitrary $ \(a :: ZDD o) ->
      a `ZDD.nonSuperset` a === ZDD.empty

prop_nonSuperset_unit :: Property
prop_nonSuperset_unit =
  forAllItemOrder $ \(_ :: Proxy o) ->
    forAll arbitrary $ \(a :: ZDD o) ->
      a `ZDD.nonSuperset` ZDD.empty === a

prop_union_nonSuperset :: Property
prop_union_nonSuperset =
  forAllItemOrder $ \(_ :: Proxy o) ->
    forAll arbitrary $ \(a :: ZDD o, b, c) ->
      (a `ZDD.union` b) `ZDD.nonSuperset` c === (a `ZDD.nonSuperset` c) `ZDD.union` (b `ZDD.nonSuperset` c)

prop_nonSuperset_union :: Property
prop_nonSuperset_union =
  forAllItemOrder $ \(_ :: Proxy o) ->
    forAll arbitrary $ \(a :: ZDD o, b, c) ->
      a `ZDD.nonSuperset` (b `ZDD.union` c) === (a `ZDD.nonSuperset` b) `ZDD.nonSuperset` c

prop_nonSuperset_difference :: Property
prop_nonSuperset_difference =
  forAllItemOrder $ \(_ :: Proxy o) ->
    forAll arbitrary $ \(a :: ZDD o, b) ->
      let c = a `ZDD.nonSuperset` b
          d = a ZDD.\\ b
       in counterexample (show (c, d)) $ c `ZDD.isSubsetOf` d

-- ------------------------------------------------------------------------
-- Minimal hitting sets
-- ------------------------------------------------------------------------

isHittingSetOf :: IntSet -> Set IntSet -> Bool
isHittingSetOf s g = all (\e -> not (IntSet.null (s `IntSet.intersection` e))) g

prop_minimalHittingSetsKnuth_isHittingSet :: Property
prop_minimalHittingSetsKnuth_isHittingSet =
  forAllItemOrder $ \(_ :: Proxy o) ->
    forAll arbitrary $ \(a :: ZDD o) ->
      let b = ZDD.minimalHittingSetsKnuth a
          a' = ZDD.toSetOfIntSets a
          b' = ZDD.toSetOfIntSets b
       in all (`isHittingSetOf` a') b'

prop_minimalHittingSetsImai_isHittingSet :: Property
prop_minimalHittingSetsImai_isHittingSet =
  forAllItemOrder $ \(_ :: Proxy o) ->
    forAll arbitrary $ \(a :: ZDD o) ->
      let b = ZDD.minimalHittingSetsImai a
          a' = ZDD.toSetOfIntSets a
          b' = ZDD.toSetOfIntSets b
       in all (`isHittingSetOf` a') b'

prop_minimalHittingSetsToda_isHittingSet :: Property
prop_minimalHittingSetsToda_isHittingSet =
  forAllItemOrder $ \(_ :: Proxy o) ->
    forAll arbitrary $ \(a :: ZDD o) ->
      let b = ZDD.minimalHittingSetsToda a
          a' = ZDD.toSetOfIntSets a
          b' = ZDD.toSetOfIntSets b
       in all (`isHittingSetOf` a') b'

prop_minimalHittingSetsKnuth_duality :: Property
prop_minimalHittingSetsKnuth_duality =
  forAllItemOrder $ \(_ :: Proxy o) ->
    forAll arbitrary $ \(a :: ZDD o) ->
      let b = ZDD.minimalHittingSetsKnuth a
       in ZDD.minimalHittingSetsKnuth (ZDD.minimalHittingSetsKnuth b) === b

prop_minimalHittingSetsImai_duality :: Property
prop_minimalHittingSetsImai_duality =
  forAllItemOrder $ \(_ :: Proxy o) ->
    forAll arbitrary $ \(a :: ZDD o) ->
      let b = ZDD.minimalHittingSetsImai a
       in ZDD.minimalHittingSetsImai (ZDD.minimalHittingSetsImai b) === b

prop_minimalHittingSetsToda_duality :: Property
prop_minimalHittingSetsToda_duality =
  forAllItemOrder $ \(_ :: Proxy o) ->
    forAll arbitrary $ \(a :: ZDD o) ->
      let b = ZDD.minimalHittingSetsToda a
       in ZDD.minimalHittingSetsToda (ZDD.minimalHittingSetsToda b) === b

prop_minimalHittingSets_Imai_equal_Knuth :: Property
prop_minimalHittingSets_Imai_equal_Knuth =
  forAllItemOrder $ \(_ :: Proxy o) ->
    forAll arbitrary $ \(a :: ZDD o) ->
      ZDD.minimalHittingSetsImai a === ZDD.minimalHittingSetsKnuth a

prop_minimalHittingSets_Toda_equal_Knuth :: Property
prop_minimalHittingSets_Toda_equal_Knuth =
  forAllItemOrder $ \(_ :: Proxy o) ->
    forAll arbitrary $ \(a :: ZDD o) ->
      ZDD.minimalHittingSetsToda a === ZDD.minimalHittingSetsKnuth a

-- ------------------------------------------------------------------------
-- Misc
-- ------------------------------------------------------------------------

prop_empty :: Property
prop_empty =
  forAllItemOrder $ \(_ :: Proxy o) ->
    ZDD.toSetOfIntSets (ZDD.empty :: ZDD o) === Set.empty

prop_base :: Property
prop_base =
  forAllItemOrder $ \(_ :: Proxy o) ->
    ZDD.toSetOfIntSets (ZDD.base :: ZDD o) === Set.singleton IntSet.empty

prop_singleton :: Property
prop_singleton =
  forAllItemOrder $ \(_ :: Proxy o) ->
    forAll arbitrary $ \xs ->
      let a :: ZDD o
          a = ZDD.singleton xs
       in counterexample (show a) $ ZDD.toSetOfIntSets a === Set.singleton xs

prop_subsets_member :: Property
prop_subsets_member =
  forAllItemOrder $ \(_ :: Proxy o) ->
    forAll arbitrary $ \xs ->
      let a :: ZDD o
          a = ZDD.subsets xs
       in counterexample (show a) $ forAll (subsetOf xs) $ \ys ->
            ys `ZDD.member` a

prop_subsets_member_empty :: Property
prop_subsets_member_empty =
  forAllItemOrder $ \(_ :: Proxy o) ->
    forAll arbitrary $ \xs ->
      let a :: ZDD o
          a = ZDD.subsets xs
       in counterexample (show a) $ IntSet.empty `ZDD.member` a

prop_subsets_member_itself :: Property
prop_subsets_member_itself =
  forAllItemOrder $ \(_ :: Proxy o) ->
    forAll arbitrary $ \xs ->
      let a :: ZDD o
          a = ZDD.subsets xs
       in counterexample (show a) $ xs `ZDD.member` a

prop_subsets_size :: Property
prop_subsets_size =
  forAllItemOrder $ \(_ :: Proxy o) ->
    forAll arbitrary $ \xs ->
      let a :: ZDD o
          a = ZDD.subsets xs
       in counterexample (show a) $ ZDD.size a === (2 :: Integer) ^ (IntSet.size xs)

prop_subsetsAtLeast :: Property
prop_subsetsAtLeast =
  forAllItemOrder $ \(_ :: Proxy o) ->
    forAll arbitrarySmallIntMap $ \(xs :: IntMap Integer) ->
      forAll arbitrary $ \k ->
        let a :: ZDD o
            a = ZDD.subsetsAtLeast xs k
         in counterexample (show a) $
              if ZDD.null a then
                property (k > sum [max 0 w | (_,w) <- IntMap.toList xs])
              else
                forAll (arbitraryMember a) $ \ys ->
                  (ys `IntSet.isSubsetOf` IntMap.keysSet xs)
                  .&&.
                  sum [xs IntMap.! y | y <- IntSet.toList ys] >= k

prop_subsetsAtMost :: Property
prop_subsetsAtMost =
  forAllItemOrder $ \(_ :: Proxy o) ->
    forAll arbitrarySmallIntMap $ \(xs :: IntMap Integer) ->
      forAll arbitrary $ \k ->
        let a :: ZDD o
            a = ZDD.subsetsAtMost xs k
         in counterexample (show a) $
              if ZDD.null a then
                property (k < sum [min 0 w | (_,w) <- IntMap.toList xs])
              else
                forAll (arbitraryMember a) $ \ys ->
                  (ys `IntSet.isSubsetOf` IntMap.keysSet xs)
                  .&&.
                  sum [xs IntMap.! y | y <- IntSet.toList ys] <= k

prop_subsetsExactly :: Property
prop_subsetsExactly =
  forAllItemOrder $ \(_ :: Proxy o) ->
    forAll arbitrarySmallIntMap $ \(xs :: IntMap Integer) ->
      forAll arbitrary $ \k ->
        let a :: ZDD o
            a = ZDD.subsetsExactly xs k
         in counterexample (show a) $
              if ZDD.null a then
                property True
              else
                forAll (arbitraryMember a) $ \ys ->
                  (ys `IntSet.isSubsetOf` IntMap.keysSet xs)
                  .&&.
                  sum [xs IntMap.! y | y <- IntSet.toList ys] === k

prop_subsetsExactly_2 :: Property
prop_subsetsExactly_2 =
  forAllItemOrder $ \(_ :: Proxy o) ->
    forAll arbitrarySmallIntMap $ \(xs :: IntMap Integer) ->
      forAll (gen xs) $ \(ys, k) ->
        let a :: ZDD o
            a = ZDD.subsetsExactly xs k
         in counterexample (show a) $ ys `ZDD.member` a
  where
    gen xs = do
      ys <- sublistOf (IntMap.toList xs)
      return (IntSet.fromList [y | (y,_) <- ys], sum [w | (_,w) <- ys])

prop_subsetsExactlyIntegral :: Property
prop_subsetsExactlyIntegral =
  forAllItemOrder $ \(_ :: Proxy o) ->
    forAll arbitrarySmallIntMap $ \(xs :: IntMap Integer) ->
      forAll arbitrary $ \k ->
        (ZDD.subsetsExactlyIntegral xs k :: ZDD o) === ZDD.subsetsExactly xs k

prop_combinations_are_combinations :: Property
prop_combinations_are_combinations =
  forAllItemOrder $ \(_ :: Proxy o) ->
    forAll arbitrary $ \xs ->
      forAll arbitrary $ \(NonNegative k) ->
        let a :: ZDD o
            a = ZDD.combinations xs k
         in counterexample (show a) $
              not (ZDD.null a)
              ==>
              (forAll (arbitraryMember a) $ \ys -> (ys `IntSet.isSubsetOf` xs) .&&. (IntSet.size ys === k))

prop_combinations_size :: Property
prop_combinations_size =
  forAllItemOrder $ \(_ :: Proxy o) ->
    forAll arbitrary $ \xs ->
      forAll arbitrary $ \(NonNegative k) ->
        let a :: ZDD o
            a = ZDD.combinations xs k
            n = toInteger $ IntSet.size xs
         in counterexample (show a) $ ZDD.size a === (product [(n - toInteger k + 1)..n] `div` (product [1..toInteger k]))

case_toList_lazyness :: Assertion
case_toList_lazyness = do
  let xss :: ZDD ZDD.AscOrder
      xss = ZDD.subsets (IntSet.fromList [1..128])
  deepseq (take 100 (Exts.toList xss)) $ return ()

prop_toList_sorted :: Property
prop_toList_sorted =
  forAllItemOrder $ \(_ :: Proxy o) ->
    forAll arbitrary $ \(xss :: ZDD o) ->
      let yss :: [[OrderedItem o]]
          yss = map (sort . map OrderedItem . IntSet.toList) $ take 100 $ Exts.toList xss
       in yss ===  sort yss

prop_toList_fromList :: Property
prop_toList_fromList =
  forAllItemOrder $ \(_ :: Proxy o) ->
    forAll arbitrary $ \xss ->
      let a :: ZDD o
          a = Exts.fromList xss
          f = Set.fromList
       in counterexample (show a) $ f (Exts.toList a) === f xss

prop_fromList_toList :: Property
prop_fromList_toList =
  forAllItemOrder $ \(_ :: Proxy o) ->
    forAll arbitrary $ \(a :: ZDD o) ->
      let xss = Exts.toList a
       in counterexample (show xss) $ Exts.fromList xss === a

prop_toSetOfIntSets_fromSetOfIntSets :: Property
prop_toSetOfIntSets_fromSetOfIntSets =
  forAllItemOrder $ \(_ :: Proxy o) ->
    forAll arbitrary $ \xss ->
      let a :: ZDD o
          a = ZDD.fromSetOfIntSets xss
       in counterexample (show a) $ ZDD.toSetOfIntSets a === xss

prop_fromSetOfIntSets_toSetOfIntSets :: Property
prop_fromSetOfIntSets_toSetOfIntSets =
  forAllItemOrder $ \(_ :: Proxy o) ->
    forAll arbitrary $ \(a :: ZDD o) ->
      let xss = ZDD.toSetOfIntSets a
       in counterexample (show xss) $ ZDD.fromSetOfIntSets xss === a

prop_insert :: Property
prop_insert =
  forAllItemOrder $ \(_ :: Proxy o) ->
    forAll arbitrary $ \(a :: ZDD o) ->
      forAll arbitrary $ \xs ->
        ZDD.toSetOfIntSets (ZDD.insert xs a) === Set.insert xs (ZDD.toSetOfIntSets a)

prop_insert_idempotent :: Property
prop_insert_idempotent =
  forAllItemOrder $ \(_ :: Proxy o) ->
    forAll arbitrary $ \(a :: ZDD o) ->
      forAll arbitrary $ \xs ->
        let b = ZDD.insert xs a
         in counterexample (show b) $ ZDD.insert xs b === b

prop_delete :: Property
prop_delete =
  forAllItemOrder $ \(_ :: Proxy o) ->
    forAll arbitrary $ \(a :: ZDD o) ->
      forAll (subsetOf (ZDD.flatten a)) $ \xs ->
        ZDD.toSetOfIntSets (ZDD.delete xs a) === Set.delete xs (ZDD.toSetOfIntSets a)

prop_delete_idempotent :: Property
prop_delete_idempotent =
  forAllItemOrder $ \(_ :: Proxy o) ->
    forAll arbitrary $ \(a :: ZDD o) ->
      forAll (subsetOf (ZDD.flatten a)) $ \xs ->
        let b = ZDD.delete xs a
         in counterexample (show b) $ ZDD.delete xs b === b

prop_mapInsert :: Property
prop_mapInsert =
  forAllItemOrder $ \(_ :: Proxy o) ->
    forAll arbitrary $ \(a :: ZDD o) ->
      forAll arbitrary $ \x ->
        ZDD.toSetOfIntSets (ZDD.mapInsert x a) === Set.map (IntSet.insert x) (ZDD.toSetOfIntSets a)

prop_mapInsert_idempotent :: Property
prop_mapInsert_idempotent =
  forAllItemOrder $ \(_ :: Proxy o) ->
    forAll arbitrary $ \(a :: ZDD o) ->
      forAll arbitrary $ \x ->
        let b = ZDD.mapInsert x a
         in counterexample (show b) $ ZDD.mapInsert x b === b

prop_mapDelete :: Property
prop_mapDelete =
  forAllItemOrder $ \(_ :: Proxy o) ->
    forAll arbitrary $ \(a :: ZDD o) ->
      forAll arbitrary $ \x ->
        ZDD.toSetOfIntSets (ZDD.mapDelete x a) === Set.map (IntSet.delete x) (ZDD.toSetOfIntSets a)

prop_mapDelete_idempotent :: Property
prop_mapDelete_idempotent =
  forAllItemOrder $ \(_ :: Proxy o) ->
    forAll arbitrary $ \(a :: ZDD o) ->
      forAll arbitrary $ \x ->
        let b = ZDD.mapDelete x a
         in counterexample (show b) $ ZDD.mapDelete x b === b

prop_change :: Property
prop_change =
  forAllItemOrder $ \(_ :: Proxy o) ->
    forAll arbitrary $ \(a :: ZDD o) ->
    forAll arbitrary $ \x ->
      let f xs
            | IntSet.member x xs = IntSet.delete x xs
            | otherwise = IntSet.insert x xs
       in ZDD.toSetOfIntSets (ZDD.change x a) === Set.map f (ZDD.toSetOfIntSets a)

prop_change_involution :: Property
prop_change_involution =
  forAllItemOrder $ \(_ :: Proxy o) ->
    forAll arbitrary $ \(a :: ZDD o) ->
      forAll arbitrary $ \x ->
        ZDD.change x (ZDD.change x a) === a

prop_subset1 :: Property
prop_subset1 =
  forAllItemOrder $ \(_ :: Proxy o) ->
    forAll arbitrary $ \(a :: ZDD o) ->
    forAll arbitrary $ \x ->
      ZDD.toSetOfIntSets (ZDD.subset1 x a) === Set.map (IntSet.delete x) (Set.filter (IntSet.member x) (ZDD.toSetOfIntSets a))

prop_subset0 :: Property
prop_subset0 =
  forAllItemOrder $ \(_ :: Proxy o) ->
    forAll arbitrary $ \(a :: ZDD o) ->
    forAll arbitrary $ \x ->
      ZDD.toSetOfIntSets (ZDD.subset0 x a) === Set.filter (IntSet.notMember x) (ZDD.toSetOfIntSets a)

prop_member_1 :: Property
prop_member_1 =
  forAllItemOrder $ \(_ :: Proxy o) ->
    forAll arbitrary $ \(a :: ZDD o) ->
      conjoin [counterexample (show xs) (ZDD.member xs a) | xs <- ZDD.toListOfIntSets a]

prop_member_2 :: Property
prop_member_2 =
  forAllItemOrder $ \(_ :: Proxy o) ->
    forAll arbitrary $ \(a :: ZDD o) ->
      forAll (subsetOf (ZDD.flatten a)) $ \s2 ->
        (s2 `ZDD.member` a) === (s2 `Set.member` ZDD.toSetOfIntSets a)

prop_size :: Property
prop_size =
  forAllItemOrder $ \(_ :: Proxy o) ->
    forAll arbitrary $ \(a :: ZDD o) ->
      ZDD.size a === Set.size (ZDD.toSetOfIntSets a)

prop_null_size :: Property
prop_null_size =
  forAllItemOrder $ \(_ :: Proxy o) ->
    forAll arbitrary $ \(a :: ZDD o) ->
      ZDD.null a === (ZDD.size a == (0 :: Int))

prop_isSubsetOf_refl :: Property
prop_isSubsetOf_refl =
  forAllItemOrder $ \(_ :: Proxy o) ->
    forAll arbitrary $ \(a :: ZDD o) ->
      a `ZDD.isSubsetOf` a

prop_isSubsetOf_empty :: Property
prop_isSubsetOf_empty =
  forAllItemOrder $ \(_ :: Proxy o) ->
    forAll arbitrary $ \(a :: ZDD o) ->
      ZDD.empty `ZDD.isSubsetOf` a

prop_isSubsetOf_and_isProperSubsetOf :: Property
prop_isSubsetOf_and_isProperSubsetOf =
  forAllItemOrder $ \(_ :: Proxy o) ->
    forAll arbitrary $ \(a :: ZDD o, b) ->
      (a `ZDD.isSubsetOf` b) === (a `ZDD.isProperSubsetOf` b || a == b)

prop_isProperSubsetOf_not_refl :: Property
prop_isProperSubsetOf_not_refl =
  forAllItemOrder $ \(_ :: Proxy o) ->
    forAll arbitrary $ \(a :: ZDD o) ->
      not (ZDD.isProperSubsetOf a a)

prop_disjoint :: Property
prop_disjoint =
  forAllItemOrder $ \(_ :: Proxy o) ->
    forAll arbitrary $ \(a :: ZDD o, b) ->
      ZDD.disjoint a b === ZDD.null (a `ZDD.intersection` b)

case_numNodes :: Assertion
case_numNodes = do
  let bdd, bdd1 :: ZDD ZDD.AscOrder
      bdd = ZDD.Branch 0 (ZDD.Branch 1 ZDD.empty bdd1) (ZDD.Branch 2 ZDD.empty bdd1)
      bdd1 = ZDD.Branch 3 ZDD.empty ZDD.base
  ZDD.numNodes bdd @?= 6

prop_flatten :: Property
prop_flatten =
  forAllItemOrder $ \(_ :: Proxy o) ->
    forAll arbitrary $ \(a :: ZDD o) ->
      ZDD.flatten a === IntSet.unions (ZDD.toListOfIntSets a)

prop_uniformM :: Property
prop_uniformM =
  forAllItemOrder $ \(_ :: Proxy o) ->
    forAll (arbitrary `suchThat` ((>= (2::Integer)) . ZDD.size)) $ \(a :: ZDD o) ->
      forAll arbitrary $ \(seed :: Vector Word32) ->
        let m :: Integer
            m = ZDD.size a
            n = 1000
            samples = runST $ do
              gen <- Rand.initialize seed
              replicateM n $ ZDD.uniformM a gen
            hist_actual = Map.fromListWith (+) [(s, 1) | s <- samples]
            hist_expected = [(s, fromIntegral n / fromIntegral m) | s <- ZDD.toListOfIntSets a]
            chi_sq = sum [(Map.findWithDefault 0 s hist_actual - cnt) ** 2 / cnt | (s, cnt) <- hist_expected]
            threshold = complQuantile (chiSquared (fromIntegral m - 1)) 0.0001
         in counterexample (show hist_actual ++ " /= " ++ show (Map.fromList hist_expected)) $
              and [xs `ZDD.member` a | xs <- Map.keys hist_actual] .&&.
              counterexample ("χ² = " ++ show chi_sq ++ " >= " ++ show threshold) (chi_sq < threshold)

prop_findMinSum :: Property
prop_findMinSum =
  forAllItemOrder $ \(_ :: Proxy o) ->
    forAll (arbitrary `suchThat` (not . ZDD.null)) $ \(a :: ZDD o) ->
      forAll arbitrary $ \(weight' :: Fun Int Integer) ->
        let weight = apply weight'
            setWeight s = sum [weight x | x <- IntSet.toList s]
            (obj0, s0) = ZDD.findMinSum weight a
         in counterexample (show [(x, weight x) | x <- IntSet.toList (ZDD.flatten a)]) $
            counterexample (show (obj0, s0)) $
              setWeight s0 === obj0
              .&&.
              conjoin [counterexample (show (s, setWeight s)) $ obj0 <= setWeight s | s <- ZDD.toListOfIntSets a]

prop_findMaxSum :: Property
prop_findMaxSum =
  forAllItemOrder $ \(_ :: Proxy o) ->
    forAll (arbitrary `suchThat` (not . ZDD.null)) $ \(a :: ZDD o) ->
      forAll arbitrary $ \(weight' :: Fun Int Integer) ->
        let weight = apply weight'
            setWeight s = sum [weight x | x <- IntSet.toList s]
            (obj0, s0) = ZDD.findMaxSum weight a
         in counterexample (show [(x, weight x) | x <- IntSet.toList (ZDD.flatten a)]) $
            counterexample (show (obj0, s0)) $
              setWeight s0 === obj0
              .&&.
              conjoin [counterexample (show (s, setWeight s)) $ obj0 >= setWeight s | s <- ZDD.toListOfIntSets a]

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

case_fold_laziness :: Assertion
case_fold_laziness = do
  let bdd :: ZDD ZDD.AscOrder
      bdd = ZDD.Branch 0 (ZDD.Branch 1 ZDD.Empty ZDD.Base) (ZDD.Branch 2 ZDD.Empty ZDD.Base)
      f x lo _hi =
        if x == 2 then
          error "unused value should not be evaluated"
        else
          lo
  seq (ZDD.fold f id bdd) $ return ()

case_fold'_strictness :: Assertion
case_fold'_strictness = do
  ref <- newIORef False
  let bdd :: ZDD ZDD.AscOrder
      bdd = ZDD.Branch 0 (ZDD.Branch 1 ZDD.Empty ZDD.Base) (ZDD.Branch 2 ZDD.Empty ZDD.Base)
      f x lo _hi = unsafePerformIO $ do
        when (x==2) $ writeIORef ref True
        return lo
  seq (ZDD.fold' f id bdd) $ do
    flag <- readIORef ref
    assertBool "unused value should be evaluated" flag

prop_fold_inSig :: Property
prop_fold_inSig =
  forAllItemOrder $ \(_ :: Proxy o) ->
    forAll arbitrary $ \(zdd :: ZDD o) ->
      ZDD.fold (\x lo hi -> ZDD.inSig (ZDD.SBranch x lo hi)) (ZDD.inSig . ZDD.SLeaf) zdd
      ===
      zdd

prop_fold'_inSig :: Property
prop_fold'_inSig =
  forAllItemOrder $ \(_ :: Proxy o) ->
    forAll arbitrary $ \(zdd :: ZDD o) ->
      ZDD.fold' (\x lo hi -> ZDD.inSig (ZDD.SBranch x lo hi)) (ZDD.inSig . ZDD.SLeaf) zdd
      ===
      zdd

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

prop_unfoldHashable_outSig :: Property
prop_unfoldHashable_outSig =
  forAllItemOrder $ \(_ :: Proxy o) ->
    forAll arbitrary $ \(zdd :: ZDD o) ->
      ZDD.unfoldHashable ZDD.outSig zdd === zdd

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

prop_toGraph_fromGraph :: Property
prop_toGraph_fromGraph = do
  forAllItemOrder $ \(_ :: Proxy o) ->
    forAll arbitrary $ \(a :: ZDD o) ->
      ZDD.fromGraph (ZDD.toGraph a) === a

-- ------------------------------------------------------------------------
-- Test cases from JDD library
-- https://bitbucket.org/vahidi/jdd/src/21e103689c697fa40022294a829cab04add8ff79/src/jdd/zdd/ZDD.java

case_jdd_test_1 :: Assertion
case_jdd_test_1 = do
  let x1 = 1
      x2 = 2

      a, b, c, d, e, f, g :: ZDD ZDD.DescOrder
      a = ZDD.empty          -- {}
      b = ZDD.base           -- {{}}
      c = ZDD.change x1 b    -- {{x1}}
      d = ZDD.change x2 b    -- {{x2}}
      e = ZDD.union c d      -- {{x1}, {x2}}
      f = ZDD.union b e      -- {{}, {x1}, {x2}}
      g = ZDD.difference f c -- {{}, {x2}}

  -- directly from minatos paper, figure 9
  -- [until we find a better way to test isomorphism...]
  a @?= ZDD.Empty
  b @?= ZDD.Base
  c @?= ZDD.Branch x1 ZDD.empty ZDD.base
  d @?= ZDD.Branch x2 ZDD.empty ZDD.base
  e @?= ZDD.Branch x2 c ZDD.base
  f @?= ZDD.Branch x2 (ZDD.Branch x1 ZDD.base ZDD.base) ZDD.base
  g @?= ZDD.Branch x2 ZDD.base ZDD.base

  -- intersect
  ZDD.intersection a b @?= a
  ZDD.intersection a ZDD.base @?= a
  ZDD.intersection b b @?= b
  ZDD.intersection c e @?= c
  ZDD.intersection e f @?= e
  ZDD.intersection e g @?= d

  -- union
  ZDD.union a a @?= a
  ZDD.union b b @?= b
  ZDD.union a b @?= b
  ZDD.union g e @?= f

  -- diff
  ZDD.difference a a @?= a
  ZDD.difference b b @?= a
  ZDD.difference d c @?= d
  ZDD.difference c d @?= c
  ZDD.difference e c @?= d
  ZDD.difference e d @?= c
  ZDD.difference g b @?= d

  ZDD.difference g d @?= b
  ZDD.difference f g @?= c
  ZDD.difference f e @?= b

  -- subset0
  ZDD.subset0 x1 b @?= b
  ZDD.subset0 x2 b @?= b
  ZDD.subset0 x2 d @?= a
  ZDD.subset0 x2 e @?= c

  -- subset1
  ZDD.subset1 x1 b @?= ZDD.empty
  ZDD.subset1 x2 b @?= ZDD.empty
  ZDD.subset1 x2 d @?= b
  ZDD.subset1 x2 g @?= b
  ZDD.subset1 x1 g @?= a
  ZDD.subset1 x2 e @?= b

case_jdd_test_2 :: Assertion
case_jdd_test_2 = do
  let [a, b, c, d] = [4, 3, 2, 1]

  -- this is the exact construction sequence of Fig.14 in "Zero-suppressed BDDs and their application" by Minato
  let fig14 :: ZDD ZDD.DescOrder
      fig14 = ZDD.union z00__ z0100
        where
          z0000 = ZDD.base
          z000_ = ZDD.union z0000 (ZDD.change d z0000)
          z00__ = ZDD.union z000_ (ZDD.change c z000_)
          z0100 = ZDD.change b z0000

  -- this is the exact construction sequence of Fig.13 in "Zero-suppressed BDDs and their application" by Minato
  let fig13 :: ZDD ZDD.DescOrder
      fig13 = ZDD.intersection z0___ tmp
        where
          z___0 = ZDD.subsets (IntSet.fromList [a, b, c])
          z__0_ = ZDD.subsets (IntSet.fromList [a, b, d])
          z_0__ = ZDD.subsets (IntSet.fromList [a, c, d])
          z0___ = ZDD.subsets (IntSet.fromList [b, c, d])
          z__00 = ZDD.intersection z___0 z__0_
          tmp = ZDD.union z_0__ z__00

  fig14 @?= fig13

case_jdd_test_3 :: Assertion
case_jdd_test_3 = do
  let tmp :: ZDD ZDD.DescOrder
      tmp  = ZDD.fromListOfIntSets $ map IntSet.fromList [[2], [0,1], [1]] -- "100 011 010"
      tmp2 = ZDD.union (ZDD.singleton (IntSet.fromList [0, 1])) ZDD.base   -- union("11", base)
  -- 1. INTERSECT
  ZDD.intersection tmp tmp2 @?= ZDD.singleton (IntSet.fromList [0, 1])
  -- 2. UNION
  ZDD.union tmp tmp2 @?= ZDD.union tmp ZDD.base
  -- 3. DIFF
  ZDD.difference tmp tmp2 @?= ZDD.fromListOfIntSets (map IntSet.fromList [[1], [2]])

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

subsetOf :: IntSet -> Gen IntSet
subsetOf = liftM IntSet.fromList . sublistOf . IntSet.toList

arbitrarySmallIntMap :: Arbitrary a => Gen (IntMap a)
arbitrarySmallIntMap = do
  n <- choose (0, 12)
  liftM IntMap.fromList $ replicateM n $ do
    k <- arbitrary
    v <- arbitrary
    return (k, v)

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

zddTestGroup :: TestTree
zddTestGroup = $(testGroupGenerator)


{-# LANGUAGE QuasiQuotes #-}
module Myelin.PyNN.PyNNSpec where

import Control.Lens
import Control.Monad.State.Lazy
import Control.Monad.Trans.Except
import Data.Aeson
import Data.ByteString.Lazy.Char8
import Data.Either
import qualified Data.Map.Strict as Map
import Data.String.Interpolate
import Numeric.LinearAlgebra
import Text.Regex as Regex

import Myelin.Model
import Myelin.Neuron
import Myelin.PyNN.PyNN
import Myelin.SNN

import Test.Hspec

main :: IO ()
main = hspec spec

eval :: PyNNState a -> Either String a
eval s = evalState (runExceptT s) emptyPyNNModel

exec :: PyNNState a -> PyNNModel
exec s = execState (runExceptT s) emptyPyNNModel

spec :: Spec
spec = do
  describe "PyNN backend" $ do
    it "can return the reference for a node" $ do
      let tpe = if_cond_exp
      let pop = Population 2 tpe "p0" 0
      let expected = "p0"
      eval (pyNNNode pop) `shouldBe` Right expected
    it "can translate an IF_exp_cond population to Python" $ do
      let tpe = if_cond_exp
      let dictWithType = unpack $ encode tpe
      let dict = Regex.subRegex typeRegex dictWithType ""
      let pop = Population 2 tpe "p0" 0
      let expectedPopulation = "p0 = pynn.Population(2, pynn.IF_cond_exp(**" ++ dict ++ "))"
      let actualModel = exec (pyNNNode pop)
      actualModel ^. populations `shouldBe` Map.singleton "p0" expectedPopulation
    it "can return the reference for a layer" $ do
      let effect = Static Excitatory (AllToAll (BiasGenerator (Constant 0)) (WeightGenerator (Constant 1.0)))
      let nodeIn = SpikeSourcePoisson 0.1 0 0
      let nodeOut = Population 2 if_cond_exp "p1" 1
      let edge = DenseProjection effect nodeIn nodeOut
      let expected = "layer0"
      eval (pyNNEdge edge) `shouldBe` Right expected
    it "can translate a dense edge to Python" $ do
      let effect = Static Excitatory (AllToAll (BiasGenerator (Constant 0)) (WeightGenerator (Constant 1.0)))
      let nodeIn = SpikeSourcePoisson 0.1 0 0
      let nodeOut = Population 2 if_cond_exp "p1" 1
      let edge = DenseProjection effect nodeIn nodeOut
      let expected = "layer0 = v.Dense(p0, p1, weights=1.0, biases=0.0)"
      exec (pyNNEdge edge) ^. layers `shouldBe` (Map.singleton "layer0" expected)
    it "can translate a replicate edge to Python" $ do
      let effect = Static Excitatory (AllToAll (BiasGenerator (Constant 0)) (WeightGenerator (Constant 1.0)))
      let nodeIn = SpikeSourcePoisson 0.1 0 0
      let nodeOut1 = Population 2 if_cond_exp "p1" 1
      let nodeOut2 = Population 2 if_cond_exp "p2" 2
      let edge = ReplicateProjection effect nodeIn (nodeOut1, nodeOut2)
      let expected = "layer0 = v.Replicate(p0, (p1, p2), weights=(1.0, 1.0), biases=0.0)"
      exec (pyNNEdge edge) ^. layers `shouldBe` (Map.singleton "layer0" expected)
    it "can translate a merge edge to Python" $ do
      let effect = Static Excitatory (AllToAll (BiasGenerator (Constant 0)) (WeightGenerator (Constant 1.0)))
      let nodeIn1 = Population 2 if_cond_exp "p0" 0
      let nodeIn2 = Population 2 if_cond_exp "p1" 1
      let nodeOut = Population 4 if_cond_exp "p2" 2
      let edge = MergeProjection effect (nodeIn1, nodeIn2) nodeOut
      let expected = "layer0 = v.Merge((p0, p1), p2)"
      exec (pyNNEdge edge) ^. layers `shouldBe` (Map.singleton "layer0" expected)
    it "can connect edges to learning nodes with constant weights" $ do
      let effect = Static Excitatory (AllToAll (BiasGenerator (Constant 4)) (WeightGenerator (Constant 12.0)))
      let nodeIn = SpikeSourcePoisson 0.1 0 0
      let nodeOut = Population 2 if_cond_exp "p1" 1
      let edge = DenseProjection effect nodeIn nodeOut
      let expected = "layer0 = v.Dense(p0, p1, weights=12.0, biases=4.0)"
      exec (pyNNEdge edge) ^. layers `shouldBe` Map.singleton "layer0" expected
    it "can correctly translate two nodes into a dictionary" $ do
      let node1 = Population 1 if_cond_exp "p0" 0
      let node2 = Population 2 if_cond_exp "p1" 1
      let network = Network 0 [node1] [node2] [] []
      let code1 = "p0 = " ++ (fromRight "" (pyNNPopulationString if_cond_exp 1 0))
      let code2 = "p1 = " ++ (fromRight "" (pyNNPopulationString if_cond_exp 2 1))
      let expected = Map.fromList [("p0", code1), ("p1", code2)]
      exec (translate' network) ^. populations `shouldBe` expected
    it "can create a model with two inputs and two outputs" $ do
      let effect1 = Static Excitatory (AllToAll (BiasGenerator (Constant 0)) (WeightGenerator (GaussianRandom 1 2)))
      let effect2 = Static Excitatory (AllToAll (BiasGenerator (Constant 0)) (WeightGenerator (Constant 1.0)))
      let nodeIn1 = Population 1 if_cond_exp "p0" 0
      let nodeIn2 = Population 2 if_cond_exp "p1" 1
      let nodeOut1 = Population 3 if_cond_exp "p2" 2
      let nodeOut2 = Population 4 if_cond_exp "p3" 3
      let edge1 = DenseProjection effect1 nodeIn1 nodeOut1
      let edge2 = DenseProjection effect2 nodeIn2 nodeOut2
      let expectedDeclarations = ["l_decode = v.Decode(p3)", "model = v.Model(layer0, layer1, l_decode)"]
      let network = Network 0 [nodeIn1, nodeIn2] [] [edge1, edge2] [nodeOut1, nodeOut2]
      let model = exec (translate' network)
      model ^. declarations `shouldBe` expectedDeclarations
    it "can translate SNN nodes and edges to a PyNN model" $ do
      let input = Population 2 if_cond_exp "p1" 1
      let hidden = Population 4 if_cond_exp "p2" 2
      let output = Population 3 if_cond_exp "p3" 3
      let effect = Static Excitatory (AllToAll (BiasGenerator (Constant 0)) (WeightGenerator (Constant 1)))
      let edges = [ DenseProjection effect input hidden, DenseProjection effect hidden output ]
      let network = Network 0 [input] [hidden] edges [output]
      let expectedDeclarations = ["l_decode = v.Decode(p3)", "model = v.Model(layer0, layer1, l_decode)"]
      let model = exec (translate' network)
      model ^. declarations `shouldBe` expectedDeclarations
      Map.size (model ^. populations) `shouldBe` 3
      Map.keys (model ^. populations) `shouldBe` ["p1", "p2", "p3"]
      Map.size (model ^. layers) `shouldBe` 2
    it "can translate SNN nodes and edges with given weights and biases to a PyNN model" $ do
      let input = Population 2 if_cond_exp "p1" 1
      let output = Population 3 if_cond_exp "p2" 2
      let biases = Biases [0..2]
      let weights=Weights ((3><2) [1..])
      let effect = Static Excitatory (AllToAll biases weights)
      let edges = [ DenseProjection effect input output ]
      let network = Network 0 [input] [] edges [output]
      let expectedDeclarations = ["l_decode = v.Decode(p2)", "model = v.Model(layer0, l_decode)"]
      let expectedLayer = "layer0 = v.Dense(p1, p2, weights=np.array([[1.0,2.0],[3.0,4.0],[5.0,6.0]]), biases=np.array([0.0,1.0,2.0]))"
      let model = exec (translate' network)
      model ^. declarations `shouldBe` expectedDeclarations
      Map.size (model ^. populations) `shouldBe` 2
      model ^. layers `shouldBe` Map.singleton "layer0" expectedLayer
    it "can fail to translate SNN nodes and edges with wrong weight/bias config" $ do
      let input = Population 2 if_cond_exp "p1" 1
      let output = Population 3 if_cond_exp "p2" 2
      let biases = Biases [0..2]
      let weights=Weights ((2><3) [1..])
      let effect = Static Excitatory (AllToAll biases weights)
      let edges = [ DenseProjection effect input output ]
      let network = Network 0 [input] [] edges [output]
      (isLeft $ eval (translate' network)) `shouldBe` True

    it "can translate an entire SNN model to a Python script" $ do
      let input = Population 2 if_cond_exp "p0" 0
      let hidden = Population 4 if_cond_exp "p1" 1
      let output = Population 3 if_cond_exp "p2" 2
      let dictWithType = unpack $ encode if_cond_exp
      let dict = Regex.subRegex typeRegex dictWithType ""
      let effect1 = Static Excitatory (AllToAll (BiasGenerator (Constant 0)) (WeightGenerator (Constant 1)))
      let effect2 = Static Excitatory (AllToAll (Biases [0.0, 1.1, 2.2]) (WeightGenerator (GaussianRandom 2 1)))
      let edges = [ DenseProjection effect1 input hidden, DenseProjection effect2 hidden output ]
      let network = Network 0 [input] [hidden] edges [output]
      let task = Task (Nest 0 0) network 50
      let preample = PyNNPreample "# some config"
      let code = [i|import numpy as np
import volrpynn.nest as v
import pyNN.nest as pynn

# some config

p0 = pynn.Population(2, pynn.IF_cond_exp(**#{dict}))
p1 = pynn.Population(4, pynn.IF_cond_exp(**#{dict}))
p2 = pynn.Population(3, pynn.IF_cond_exp(**#{dict}))
layer0 = v.Dense(p0, p1, weights=1.0, biases=0.0)
layer1 = v.Dense(p1, p2, weights=np.random.normal(2.0, 1.0, (4, 3)), biases=np.array([0.0,1.1,2.2]))
l_decode = v.Decode(p2)
model = v.Model(layer0, layer1, l_decode)

optimiser = v.GradientDescentOptimiser(0.1, simulation_time=50.0)
if __name__ == "__main__":
    v.Main(model).train(optimiser)
|]
      translate task preample `shouldBe` Right code
    it "can translate an parallel SNN model to a Python script" $ do
      let input = Population 100 if_cond_exp "p0" 0
      let h1 = Population 20 if_cond_exp "p1" 1
      let h2 = Population 20 if_cond_exp "p2" 2
      let h3 = Population 10 if_cond_exp "p3" 3
      let h4 = Population 20 if_cond_exp "p4" 4
      let h5 = Population 10 if_cond_exp "p5" 5
      let h6 = Population 20 if_cond_exp "p6" 6
      let output = Population 10 if_cond_exp "p7" 7
      let dictWithType = unpack $ encode if_cond_exp
      let dict = Regex.subRegex typeRegex dictWithType ""
      let effect = Static Excitatory (AllToAll (BiasGenerator (Constant 0)) (WeightGenerator (Constant 1)))
      let edges = [ DenseProjection effect input h1
		  , ReplicateProjection effect h1 (h2, h4)
		  , DenseProjection effect h2 h3
		  , DenseProjection effect h4 h5
		  , MergeProjection effect (h3, h5) h6
		  , DenseProjection effect h6 output
		  ]
      let network = Network 0 [input] [h1, h2, h3, h4, h5, h6] edges [output]
      let task = Task (Nest 0 0) network 50
      let preample = PyNNPreample "# some config"
      let code = [i|import numpy as np
import volrpynn.nest as v
import pyNN.nest as pynn

# some config

p0 = pynn.Population(100, pynn.IF_cond_exp(**#{dict}))
p1 = pynn.Population(20, pynn.IF_cond_exp(**#{dict}))
p2 = pynn.Population(20, pynn.IF_cond_exp(**#{dict}))
p3 = pynn.Population(10, pynn.IF_cond_exp(**#{dict}))
p4 = pynn.Population(20, pynn.IF_cond_exp(**#{dict}))
p5 = pynn.Population(10, pynn.IF_cond_exp(**#{dict}))
p6 = pynn.Population(20, pynn.IF_cond_exp(**#{dict}))
p7 = pynn.Population(10, pynn.IF_cond_exp(**#{dict}))
layer0 = v.Dense(p0, p1, weights=1.0, biases=0.0)
layer1 = v.Replicate(p1, (p2, p4), weights=(1.0, 1.0), biases=0.0)
layer2 = v.Dense(p2, p3, weights=1.0, biases=0.0)
layer3 = v.Dense(p4, p5, weights=1.0, biases=0.0)
layer4 = v.Merge((p3, p5), p6)
layer5 = v.Dense(p6, p7, weights=1.0, biases=0.0)
l_decode = v.Decode(p7)
model = v.Model(layer0, layer1, layer2, layer3, layer4, layer5, l_decode)

optimiser = v.GradientDescentOptimiser(0.1, simulation_time=50.0)
if __name__ == "__main__":
    v.Main(model).train(optimiser)
|]
      translate task preample `shouldBe` Right code


-- Copyright 2016 TensorFlow authors.
--
-- Licensed under the Apache License, Version 2.0 (the "License");
-- you may not use this file except in compliance with the License.
-- You may obtain a copy of the License at
--
--     http://www.apache.org/licenses/LICENSE-2.0
--
-- Unless required by applicable law or agreed to in writing, software
-- distributed under the License is distributed on an "AS IS" BASIS,
-- WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-- See the License for the specific language governing permissions and
-- limitations under the License.

{-# LANGUAGE ConstraintKinds #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE ViewPatterns #-}

module TensorFlow.Gradient
    ( GradientCompatible
    , gradients
    ) where

import Control.Monad (forM, zipWithM)
import Control.Monad.State.Strict (State, evalState, gets, modify)
import Data.ByteString (ByteString)
import Data.Complex (Complex)
import Data.Default (def)
import Data.Int (Int32, Int64)
import Data.Foldable (foldlM)
import Data.List (foldl', sortBy)
import Data.Map.Strict (Map)
import Data.Maybe (fromMaybe, maybeToList, mapMaybe)
import Data.Ord (comparing)
import Data.ProtoLens.TextFormat (showMessage)
import Data.Set (Set)
import Data.Text (Text)
import Data.Tuple (swap)
import Lens.Family2 (Lens', view, (&), (^.), (.~), (%~))
import Lens.Family2.State.Strict (uses)
import Lens.Family2.Stock (at, intAt)
import Lens.Family2.Unchecked (lens, iso)
import Prelude hiding (sum)
import Text.Printf (printf)
import qualified Data.Graph.Inductive.Basic as FGL
import qualified Data.Graph.Inductive.Graph as FGL
import qualified Data.Graph.Inductive.PatriciaTree as FGL
import qualified Data.Graph.Inductive.Query.DFS as FGL
import qualified Data.IntMap.Strict as IntMap
import qualified Data.Map.Strict as Map
import qualified Data.Set as Set
import qualified Data.Text as Text

import qualified TensorFlow.GenOps.Core as CoreOps
import TensorFlow.Build
    ( MonadBuild
    , Build
    , build
    , renderedNodeDefs
    , opDef
    , opAttr
    , opInputs
    )
import TensorFlow.BuildOp
import TensorFlow.Ops
    ( addN
    , broadcastGradientArgs
    , expandDims
    , fill
    , matMul
    , matMul'
    , reducedShape
    , reluGrad
    , reshape
    , scalar
    , shape
    , softmaxCrossEntropyWithLogits
    , sum
    , scalarize
    , vector
    , zerosLike
    )
import TensorFlow.Output
    ( NodeName(..)
    , Output(..)
    , OutputIx(..)
    , outputIndex
    )
import TensorFlow.Tensor
    ( Tensor(..)
    , Value
    , render
    , expr
    , Rendered
    , tensorNodeName
    , renderedOutput
    , renderValue
    , ToTensor(..)
    )
import TensorFlow.Types (Attribute, OneOf, TensorType, attrLens)
import Proto.Tensorflow.Core.Framework.NodeDef
    (NodeDef, attr, input, op, name)

type GradientCompatible a =
    -- TODO(fmayle): MaxPoolGrad doesn't support Double for some reason.
    (Num a, OneOf '[ Float, Complex Float, Complex Double ] a)

-- TODO(fmayle): Support control flow.
-- TODO(fmayle): Support gate_gradients-like option to avoid race conditions.
-- TODO(fmayle): Do we need to consider control inputs? See _PendingCount in
-- tensorflow/python/ops/gradients.py.
-- TODO(fmayle): Maybe store the gradient functions and numOutputs on the OpDef.

-- | Gradient of @y@ w.r.t. each element of @xs@.
gradients :: forall a v1 t m . ( MonadBuild m
                               , Rendered t
                               , ToTensor t
                               , GradientCompatible a
                               )
          => Tensor v1 a  -- ^ The output of the graph.
          -> [t a]        -- ^ Tensors for which gradients are computed.
          -> m [Tensor Value a]
gradients y xs = build $ do
    -- The gradients are computed using "reverse accumulation", similarly to
    -- what is described here:
    -- https://en.wikipedia.org/wiki/Automatic_differentiation#The_chain_rule.2C_forward_and_reverse_accumulation
    --
    -- The code is summarised as follows:
    --
    -- 1. Create an fgl graph of the relevant nodes (ops) and edges (tensors).
    -- 2. Initialize the gradient of y to 1 (∂y/∂y = 1) and the rest of tensor's
    --    gradients to nothing.
    -- 3. Process the nodes in reverse topological order (i.e. each node comes
    --    after all of its outputs so that the output gradients for a node have
    --    been completely calculated before it is processed):
    --      a. Record the gradient for each of the node's output tensors (∂y/∂w
    --         for each output tensor w).
    --      b. Calculate the gradient of y w.r.t. each of the node's input
    --         tensors using the gradients of the node's output tensors.
    --
    --         Written differently, for each output tensor w and input tensor v:
    --           ∂y/∂w = ...            (calculated in previous steps)
    --           ∂w/∂v = ...            (op specific)
    --           ∂y/∂v = ∂y/∂w * ∂w/∂v  (technically, if tensor v is an input
    --                                   to multiple nodes, then this is only
    --                                   part of ∂y/∂v)
    --
    -- 4. Lookup the recorded gradient for each x in xs.

    y' <- renderValue y
    let yName = tensorNodeName y'
    yOne <- render $ fill (shape y') (scalar 1)
    -- TODO(fmayle): Move this into Build.hs and call it unsafeNodeDefFromName?
    nodeDefLookup :: (NodeName -> NodeDef) <- uses renderedNodeDefs $
        (\f x -> fromMaybe (error $ "no NodeDef found for " ++ show x) (f x))
        . flip Map.lookup
    let (gr, nodeMap) = createGraph yName nodeDefLookup
    -- Set gradient of y to one.
    -- TODO: nicer
    let initPending :: Map.Map FGL.Node (PendingGradients a)
            = Map.empty & (at (nodeMap Map.! yName)
                                . nonEmpty
                                . outputIxAt (outputIndex $ renderedOutput y')
                                . nonEmpty
                                .~ [yOne]
                                )
    -- Calculate the gradients of y w.r.t. each node in the graph.
    gradientMap <- graphGrads gr initPending
    -- Lookup the gradients for each x.
    forM xs $ \x ->
        let Output i xName = renderedOutput x
        in maybe (render $ zerosLike $ toTensor x) return $ do
            n <- nodeMap ^. at xName
            gradientMap ^. at n . nonEmpty . outputIxAt i

outputIxAt :: OutputIx -> Lens' (IntMap.IntMap v) (Maybe v)
outputIxAt = intAt . unOutputIx

-- | Incomplete gradients of a node's outputs.
--
-- The lists represent partial sums. The key is an OutputIx sans newtype.
type PendingGradients a = IntMap.IntMap [Tensor Value a]

-- | Gradients of a node's outputs. The key is an OutputIx sans newtype.
-- TODO: precache the rendering?
type Gradients a = IntMap.IntMap (Tensor Value a)

-- | Graph of TensorFlow operations.
type Graph = FGL.Gr NodeDef EdgeLabel

-- | Data associated with an edge.
--
-- Pair of
--   1. Output index of a tensor from the source node.
--   2. Input index that the tensor connects to on the destination node.
type EdgeLabel = (OutputIx, OutputIx)

-- | State used for calculating gradients.
data GradientsState a = GradientsState
                      { _gradientsPending :: !(Map FGL.Node (PendingGradients a))
                      , _gradientsResult  :: !(Map FGL.Node (Gradients a))
                      }

gradientsPending :: Lens' (GradientsState a) (Map FGL.Node (PendingGradients a))
gradientsPending = lens _gradientsPending (\x y -> x { _gradientsPending = y })

gradientsResult :: Lens' (GradientsState a) (Map FGL.Node (Gradients a))
gradientsResult = lens _gradientsResult (\x y -> x { _gradientsResult = y })

-- TODO(fmayle): Use something like Data.List.Safe.
-- | Safe version of (!!).
safeIndex :: [a] -> Int -> Maybe a
_      `safeIndex` n | n < 0 = Nothing
[]     `safeIndex` _         = Nothing
(x:_)  `safeIndex` 0         = Just x
(_:xs) `safeIndex` n         = xs `safeIndex` (n-1)

-- Copy of http://hackage.haskell.org/package/lens-3.9.0.2/docs/Control-Lens-Iso.html#v%3anon
anon :: a -> (a -> Bool) -> Lens' (Maybe a) a
anon a p = iso (fromMaybe a) go where
  go b | p b       = Nothing
       | otherwise = Just b

non :: Eq a => a -> Lens' (Maybe a) a
non a = anon a (a==)

-- | Lens that defaults Nothing to mempty.
nonEmpty :: (Monoid (t v), Foldable t) => Lens' (Maybe (t v)) (t v)
nonEmpty = anon mempty null

-- TODO: strictness (e.g., foldlM')

-- | Calculate the gradients for every node in a graph.
graphGrads :: forall a. GradientCompatible a
           => Graph
           -> Map FGL.Node (PendingGradients a)
           -- ^ Initial gradients (usually just 1 for the node of interest).
           -> Build (Map FGL.Node (Gradients a))
graphGrads gr initPending = view gradientsResult <$> foldlM go initState nodeOrder
  where
    initState = GradientsState initPending Map.empty
    -- Reverse topological sort.
    -- TODO(fmayle): Filter out nodes that are not successors of any x in xs to
    -- avoid calculating gradients that won't be used.
    nodeOrder = FGL.topsort $ FGL.grev gr
    go :: GradientsState a -> Int -> Build (GradientsState a)
    go state node = do
        -- Aggregate the accumulated gradients for this node.
        outputGrads <-
                sumPendingGradient (state ^. gradientsPending . at node . nonEmpty)
        if null outputGrads
           then pure state
           else do
              let ctx = FGL.context gr node
              inputGrads <- calculateInputGrads ctx outputGrads gr
              -- Calculate the gradients for each of the node's inputs.
              let nextState = state & gradientsResult %~ Map.insert node outputGrads
              pure $ updatePendingGradients ctx inputGrads nextState

-- | Reduce accumulated gradients for each output to one Tensor.
sumPendingGradient :: GradientCompatible a
                   => PendingGradients a -> Build (Gradients a)
sumPendingGradient = sequence . IntMap.mapMaybe f
  where
    f [] = Nothing
    f [x] = Just (pure x)
    f xs = Just (render $ addN xs)

-- | Calculate the gradients of a node's input tensors.
--
-- This is mostly just a wrapper around opGrad.
calculateInputGrads :: forall a. GradientCompatible a
                    => FGL.Context NodeDef EdgeLabel
                    -> Gradients a  -- ^ Output gradients of the node.
                    -> Graph
                    -> Build [Maybe (Tensor Value a)]
calculateInputGrads (inputEdges, _, nodeDef, _) outputGrads gr = do
    fullOutGrads <- fullOutputGrads (numOutputs nodeDef) (nodeDefName nodeDef)
                        outputGrads
    traverse (traverse render) $ opGrad (nodeDef ^. op) nodeDef inputTensors fullOutGrads
  where
    -- Create a tensor from an edge (technically an Output, but it seems less
    -- confusing to refer to it as a tensor here).
    edgeToTensor :: (EdgeLabel, FGL.Node) -> Output
    edgeToTensor ((i, _), n) =
        case FGL.lab gr n of
            Just edgeNodeDef -> Output i (NodeName $ edgeNodeDef ^. name)
            Nothing -> error $ "calculateInputGrads: missing input node for "
                               ++ Text.unpack (nodeDef ^. name)
    -- Input tensors, sorted by input index.
    inputTensors = map edgeToTensor $ sortBy (comparing (snd . fst)) inputEdges

-- | Convert a Map of gradients to a list, with zeros for missing outputs.
fullOutputGrads :: (TensorType a, Num a)
                => OutputIx  -- ^ Number of outputs.
                -> NodeName
                -> Gradients a
                -> Build [Tensor Value a]
fullOutputGrads n o gs =
    mapM (\i -> maybe (render $ zero i) return (gs ^. outputIxAt i)) [0..n-1]
  where
    -- A tensor of zeros with the same shape as the i'th output.
    zero i = zerosLike $ toT (Output i o)

-- | Update the pending gradients of a node's inputs.
updatePendingGradients :: forall a. (TensorType a, Num a)
                       => FGL.Context NodeDef EdgeLabel
                       -> [Maybe (Tensor Value a)]
                       -- ^ Gradient of each input tensor.
                       -> GradientsState a
                       -> GradientsState a
updatePendingGradients (inputEdges, _, nodeDef, _) inputGrads initState =
    foldl' go initState inputEdges
  where
    go :: GradientsState a -> (EdgeLabel, FGL.Node) -> GradientsState a
    go state ((outIndex, OutputIx inIndex), node) =
        case maybeGradient of
            Nothing -> state
            Just g ->
                -- Add to the list of pending gradients for this tensor.
                state & gradientsPending
                      . at node
                      . nonEmpty
                      . outputIxAt outIndex
                      . nonEmpty
                      %~ (g:)
      where
        badSizeErr = error $ printf "updatePendingGradients: bad input index \
                                    \%d for inputGrads of length %d in %s"
                                    inIndex (length inputGrads)
                                    (show (nodeDef ^. name))
        maybeGradient = fromMaybe badSizeErr (safeIndex inputGrads inIndex)

-- | Create a graph that includes a node and its transitive dependencies.
createGraph :: NodeName -> (NodeName -> NodeDef)
            -> (Graph, Map NodeName FGL.Node)
createGraph nodeName nodeDefLookup = (FGL.nmap nodeDefLookup graph, nodeMap)
  where
    -- Parse a tensor name.
    parseTensorName :: Text -> Maybe (NodeName, OutputIx)
    parseTensorName n
        | Text.null n        = error "parseTensorName: empty name"
        | Text.head n == '^' = Nothing  -- Control edge
        | otherwise          =
            let (nm, indexStr) = Text.breakOn ":" n
                index | Text.null indexStr = 0
                      | otherwise = read $ Text.unpack $ Text.tail indexStr
            in Just (NodeName nm, OutputIx index)

    -- Build a map from node name to outward edges.
    --
    -- The state is the set of visited nodes.
    collect :: Maybe (NodeName, OutputIx, OutputIx)
            -> NodeName
            -> State (Set NodeName)
                     (Map NodeName [(NodeName, OutputIx, OutputIx)])
    collect outgoingEdge nm = do
        let nextLookup = Map.singleton nm (maybeToList outgoingEdge)
        seen <- gets (Set.member nm)
        modify (Set.insert nm)
        if seen
            then pure nextLookup
            else do
                let inputs = nodeDefLookup nm ^. input
                    recurse inIndex (parentName, outIndex) =
                        collect (Just (nm, outIndex, inIndex)) parentName
                subEdgeLookups <-
                    zipWithM recurse [0..] $ mapMaybe parseTensorName inputs
                pure $ Map.unionsWith (++) (nextLookup:subEdgeLookups)

    edgeLookup = evalState (collect Nothing nodeName) Set.empty
    -- Associate an ID with each node name.
    nodeMap = Map.fromList $ zip (Map.keys edgeLookup) [0..]
    -- Create the graph.
    graph = FGL.mkGraph (swap <$> Map.toList nodeMap)
                        [ (nodeMap Map.! n, nodeMap Map.! m, (i, j))
                        | (n, edges) <- Map.toList edgeLookup
                        , (m, i, j) <- edges
                        ]

-- | Function to compute the gradient of y w.r.t. each input.
--
-- Let y be an arbitrary tensor
-- and [w_0, ..., w_n] be the output tensors of a node
-- and [v_0, ..., v_n] be the input tensors of the same node.
--
-- Given [∂y/∂w_0, ..., ∂y/∂w_n] and [v_0, ..., v_n], a GradientFunc computes
-- [∂y/∂v_0, ..., ∂y/∂v_n] for a particular op type.
--
-- A Nothing gradient is equivalent to zero (but allows for short circuiting
-- computation when all the gradients for something are Nothing).
type GradientFunc a = NodeDef
                    -> [Output]
                    -- ^ Input tensors.
                    -> [Tensor Value a]
                    -- ^ Gradient of y w.r.t. each output tensor.
                    -> [Maybe (Tensor Build a)]
                    -- ^ Gradient of y w.r.t. each input tensor.

-- TODO(fmayle): Assert the type is correct.
-- | Create a Tensor from an Output.
toT :: Output -> Tensor Build a
toT = Tensor . pure

-- | Wrapper around `TensorFlow.GenOps.Core.slice` that builds vectors from scalars for
-- simple slicing operations.
flatSlice :: forall v1 t . TensorType t
         => Tensor v1 t    -- ^ __input__
         -> Int32          -- ^ __begin__: specifies the offset into the first dimension of
                           -- 'input' to slice from.
         -> Int32          -- ^ __size__: specifies the number of elements of the first dimension
                           -- of 'input' to slice. If size is -1, all remaining elements in the dimension
                           -- are included in the slice (i.e. this is equivalent to setting
                           -- size = input.dim_size(0) - begin).
         -> Tensor Build t -- ^ __output__
flatSlice t begin size = CoreOps.slice t (vector [begin]) (vector [size])

nodeDefName :: NodeDef -> NodeName
nodeDefName = NodeName . view name

-- | Gradient helper for binary component wise operations
-- See https://github.com/tensorflow/tensorflow/blob/e9de087fa7f59c39bbe12ac2c83c5547c83f746c/tensorflow/core/ops/math_grad.cc#L329
gradForBinaryCwise :: ( OneOf '[ Int32, Int64, Float, Double, Complex Float, Complex Double ] t
                      )
                   => (Tensor v1 t, Tensor v1 t)
                   -> (Tensor v1 t, Tensor v1 t)
                   -> [ Maybe (Tensor Build t) ]
gradForBinaryCwise (x, gx) (y, gy) =
    [ Just dx
    , Just dy ]
  where
    dx = reshape (sum gx rx) sx
    dy = reshape (sum gy ry) sy
    sx = shape x
    sy = shape y
    (rx, ry) = broadcastGradientArgs sx sy

-- | The gradient function for an op type.
--
-- These implementations should match their python counterparts in:
-- third_party/tensorflow/python/ops/*_grad.py
opGrad :: forall a . GradientCompatible a => Text -> GradientFunc a

opGrad "Abs" _ [toT -> x] [dz] = [Just $ expr dz * signum x]
opGrad "Neg" _ [_] [dz] = [Just $ negate $ expr dz]
opGrad "Relu" _ [toT -> x] [dz] = [Just $ reluGrad dz x]
opGrad "ReluGrad" _ [_, toT -> x ] [dz] = [Just $ reluGrad dz x, Just $ CoreOps.zerosLike x]

opGrad "Concat" _ _ix [dy]
    -- Concat concatenates input tensors
    --   x1 of shape s1 = [k1, ..., ki_1, ..., kn]
    --   x2 of shape s2 = [k1, ..., ki_2, ..., kn]
    --    .           .     .          .        .
    --    .           .     .          .        .
    --    .           .     .          .        .
    --   xm of shape sm = [k1, ..., ki_m, ..., kn]
    --  along dimension i to an output tensor
    --   y  of shape sy = [k1, ..., k, ..., kn]
    --  where k = sum ki = sum [ki_1,...,ki_m]
    --
    --  The incoming gradient dy from backpropagation is
    --   simply forwarded split across input tensors yielding dx.
    --   Forwarded gradients have shapes s = [s1, ..., sm].
    | m == 1    = Nothing : [Just $ expr dy]
    | otherwise = Nothing : map Just (dx `reshapeZip` s)
  where
    reshapeZip = zipWith reshape
    dx = CoreOps.splitV (fromIntegral m) dy ki _i
    s  :: [Tensor Build Int32]
    s  = map shape x
    x  :: [Tensor Build a]
    x  = map toT $ tail _ix
    -- i: concat dimension. Adjusted modulo n to handle negative indices.
    _i = toT (head _ix) `CoreOps.floorMod` n
    i  = reshape _i $ vector [1 :: Int32]
    -- sizes along concatenated dimension
    ki :: Tensor Build Int32
    ki = CoreOps.concat 0 $ map (\t -> CoreOps.slice t i $ vector [1 :: Int32]) s
    m  = length x
    n  = CoreOps.rank (head x)

opGrad "Square" _ [toT -> x] [dz] =
    -- TODO(fmayle): Handle complex numbers.
    -- TODO(fmayle): The python code makes dz a control dependency of the 2*x
    -- (for performance reasons?). Will need to put these functions in the Build
    -- monad to replicate that.
    [Just $ dz `CoreOps.mul` (2 * x)]

opGrad "Gather" _ [toT -> x, toT -> indices] [dz] =
    -- TODO(fmayle): The python version uses a better performance implementation
    -- when the shape is known without having to run the graph.
    -- TODO(fmayle): We shouldn't convert the result to a dense tensor. Sparse
    -- tensor support will require some thinking.
    [ Just $ CoreOps.unsortedSegmentSum values indices' numRows
    , Nothing
    ]
  where
    -- TODO(gnezdo): Use colocateWith but it requires Build monad.
    denseShape = shape (x :: Tensor Build a)
    numRows = scalarize $ flatSlice denseShape 0 1
    valuesShape = CoreOps.concat 0 [ allDimensions
                                   , flatSlice denseShape 1 (-1)
                                   ]
    values = reshape dz valuesShape
    -- TODO(fmayle): This could be either Int32 or Int64.
    indices' = reshape indices allDimensions :: Tensor Build Int32

opGrad "Max" _ [toT -> x, toT -> indices] [dz] =
    [Just $ indicators `CoreOps.div` numSelected * dz', Nothing]
  where
    sx = shape (x :: Tensor Build a)
    outputShapeKeptDims = reducedShape sx (indices :: Tensor Build Int32)
    y = CoreOps.max x indices
    y' = reshape y outputShapeKeptDims
    dz' = reshape dz outputShapeKeptDims
    indicators = CoreOps.cast $ CoreOps.equal y' x
    numSelected = reshape (sum indicators indices) outputShapeKeptDims

-- Min and Max have identical gradient implementations.
opGrad "Min" u v w = opGrad "Max" u v w

-- Element wise maximum gradient
-- See https://github.com/tensorflow/tensorflow/blob/e9de087fa7f59c39bbe12ac2c83c5547c83f746c/tensorflow/core/ops/math_grad.cc#L473
opGrad "Maximum" _ [toT -> x, toT -> y] [dz] =
    gradForBinaryCwise (x, gx) (y, gy)
  where
    xmask = CoreOps.greaterEqual x y
    gx = CoreOps.select xmask dz (CoreOps.zerosLike dz)
    gy = CoreOps.select (CoreOps.logicalNot xmask) dz (CoreOps.zerosLike dz)

opGrad "Sum" _ [toT -> x, toT -> indices] [dz] =
    [ Just $ CoreOps.tile grad tileScaling, Nothing ]
  where
    -- TODO(gnezdo): Implement the fast-path from math_grad._SumGrad.
    sx = shape (x :: Tensor Build a)
    outputShapeKeptDims = reducedShape sx (indices :: Tensor Build Int32)
    tileScaling = safeShapeDiv sx outputShapeKeptDims
    grad = reshape dz outputShapeKeptDims

opGrad "Mean" u v@[toT -> x, _] w =
    [Just $ dz `CoreOps.div` CoreOps.cast factor, Nothing]
  where
    [Just dz, Nothing] = opGrad "Sum" u v w
    inputShape = shape (x :: Tensor Build a)
    outputShape = shape (dz :: Tensor Build a)
    -- TODO(fmayle): Add fast path when shape is known.
    inputSize = CoreOps.prod inputShape $ rangeOfRank inputShape
    outputSize = CoreOps.prod outputShape $ rangeOfRank outputShape
    factor = safeShapeDiv inputSize outputSize

opGrad "Add" _ [toT -> x, toT -> y] [dz] =
    [ Just $ reshape (sum dz rx) sx
    , Just $ reshape (sum dz ry) sy ]
  where
    sx = shape (x :: Tensor Build a)
    sy = shape (y :: Tensor Build a)
    (rx, ry) = broadcastGradientArgs sx sy

-- Copies the gradients to all inputs
-- Not broadcasting
opGrad "AddN" _ inputs [dz] =
    map ((const . Just . expr) dz) inputs

opGrad "Sub" u v w =
    [Just x, Just (-y)]
  where
    [Just x, Just y] = opGrad "Add" u v w

opGrad "SoftmaxCrossEntropyWithLogits" _ [toT -> x, toT -> y] [dz, _] =
    [ Just $ expandDims dz (-1) * snd (softmaxCrossEntropyWithLogits x y)
    , Nothing ]

opGrad "Mul" _ [toT -> x, toT -> y] [dz] =
    -- TODO(fmayle): Handle complex numbers.
    [ Just $ reshape (sum (dz `CoreOps.mul` y) rx) sx
    , Just $ reshape (sum (x `CoreOps.mul` dz) ry) sy ]
  where
    sx = shape (x :: Tensor Build a)
    sy = shape (y :: Tensor Build a)
    (rx, ry) = broadcastGradientArgs sx sy

opGrad "Div" _ [toT -> x, toT -> y] [dz] =
    -- TODO(fmayle): Handle complex numbers.
    -- TODO(gnezdo): Provide Fractional instance and use '/' instead of div.
    [ Just $ reshape (sum (dz `CoreOps.div` y) rx) sx
    , Just $ reshape (sum (dz `CoreOps.mul` (negate x `CoreOps.div` (y * y)))
                         ry)
                sy
    ]
  where
    sx = shape (x :: Tensor Build a)
    sy = shape (y :: Tensor Build a)
    (rx, ry) = broadcastGradientArgs sx sy

opGrad "MatMul" nodeDef [toT -> x, toT -> y] [dz] =
    let transposeA = lookupAttr nodeDef "transpose_a"
        transposeB = lookupAttr nodeDef "transpose_b"
        transAttrs a b =
            (opAttr "transpose_a" .~ a) . (opAttr "transpose_b" .~ b)
    in case (transposeA, transposeB) of
       (False, False) ->
           [ Just $ matMul' (transAttrs False True) dz y
           , Just $ matMul' (transAttrs True False) x dz]
       (False, True) ->
           [ Just $ matMul dz y
           , Just $ matMul' (transAttrs True False) dz x]
       (True, False) ->
           [ Just $ matMul' (transAttrs False True) y dz
           , Just $ matMul x dz]
       (True, True) ->
           [ Just $ matMul' (transAttrs True True) y dz
           , Just $ matMul' (transAttrs True True) dz x]

opGrad "Transpose" _ [_, toT -> p] [dz] =
    [ Just $ CoreOps.transpose dz
            (CoreOps.invertPermutation p :: Tensor Build Int32)
    , Nothing
    ]

opGrad "Conv2D" nodeDef [toT -> x, toT -> y] [dz] =
    [ Just $ CoreOps.conv2DBackpropInput'
                ((opAttr "strides" .~ strides)
                    . (opAttr "padding" .~ padding)
                    . (opAttr "use_cudnn_on_gpu" .~ useCudnnOnGpu)
                    . (opAttr "data_format" .~ dataFormat))
                (shape x) y dz
    , Just $ CoreOps.conv2DBackpropFilter'
                ((opAttr "strides" .~ strides)
                    . (opAttr "padding" .~ padding)
                    . (opAttr "use_cudnn_on_gpu" .~ useCudnnOnGpu)
                    . (opAttr "data_format" .~ dataFormat))
                x (shape y) dz
    ]
  where
    strides = lookupAttr nodeDef "strides" :: [Int64]
    padding = lookupAttr nodeDef "padding" :: ByteString
    useCudnnOnGpu = lookupAttr nodeDef "use_cudnn_on_gpu" :: Bool
    dataFormat = lookupAttr nodeDef "data_format" :: ByteString

opGrad "Conv2DBackpropInput" nodeDef [_, toT -> x, toT -> y] [dz] =
    [ Nothing
    , Just $ CoreOps.conv2DBackpropFilter'
                ((opAttr "strides" .~ strides)
                    . (opAttr "padding" .~ padding)
                    . (opAttr "use_cudnn_on_gpu" .~ useCudnnOnGpu)
                    . (opAttr "data_format" .~ dataFormat))
                dz (shape x) y
    , Just $ CoreOps.conv2D'
                ((opAttr "strides" .~ strides)
                    . (opAttr "padding" .~ padding)
                    . (opAttr "use_cudnn_on_gpu" .~ useCudnnOnGpu)
                    . (opAttr "data_format" .~ dataFormat))
                dz x
    ]
  where
    strides = lookupAttr nodeDef "strides" :: [Int64]
    padding = lookupAttr nodeDef "padding" :: ByteString
    useCudnnOnGpu = lookupAttr nodeDef "use_cudnn_on_gpu" :: Bool
    dataFormat = lookupAttr nodeDef "data_format" :: ByteString

opGrad "MaxPool" nodeDef [toT -> x] [dz] =
    [ Just $ CoreOps.maxPoolGrad'
                ((opAttr "ksize" .~ ksize)
                    . (opAttr "strides" .~ strides)
                    . (opAttr "padding" .~ padding)
                    . (opAttr "data_format" .~ dataFormat))
                x output dz
    ]
  where
    output :: Tensor Build a
    output = toT $ Output 0 (nodeDefName nodeDef)
    ksize = lookupAttr nodeDef "ksize" :: [Int64]
    strides = lookupAttr nodeDef "strides" :: [Int64]
    padding = lookupAttr nodeDef "padding" :: ByteString
    dataFormat = lookupAttr nodeDef "data_format" :: ByteString

opGrad "Reshape" _ [toT -> x, _] [dz] =
    [Just $ reshape dz $ shape (x :: Tensor Build a), Nothing]

opGrad "OneHot" _ _ _ = [Nothing, Nothing, Nothing, Nothing]
opGrad "TruncatedNormal" _ _ _ = [Nothing]

opGrad "RefIdentity" _ _ [dz] = [Just $ expr dz]
opGrad "Cast" nodeDef _ [dz] = [Just reverseCast]
  where
    -- TODO(gnezdo): too permissive, python only allows float types as src_type.
    reverseCast =
        pureOp [] $ pure (opDef "Cast"
                 & opAttr "DstT" .~ (lookupAttr nodeDef "SrcT" :: ByteString)
                 & opAttr "SrcT" .~ (lookupAttr nodeDef "DstT" :: ByteString)
                 & opInputs .~ [renderedOutput dz])

opGrad "DynamicStitch" nodeDef inputs [dz] =
    replicate halfLen Nothing ++ valuesGrads
  where
    halfLen =
        let len = length inputs
            half = len `div` 2
        in if 2 * half == len
           then half
           else error ("Uneven input size " ++ show (len, showMessage nodeDef))
    valuesGrads = [ Just $ CoreOps.gather dz (toT idx :: Tensor Build Int32)
                  | idx <- take halfLen inputs
                  ]

opGrad "DynamicPartition" nodeDef [toT -> xs, toT -> indices] dz =
    [ Just reconstructed, Nothing ]
  where
    reconstructed = CoreOps.reshape stitched
                    (CoreOps.shape (xs :: Tensor Build a) :: Tensor Build Int32)
    stitched = CoreOps.dynamicStitch partitionedIndices dz
    partitionedIndices = CoreOps.dynamicPartition np originalIndices indices
    np = lookupAttr nodeDef "num_partitions" :: Int64
    originalIndices =
        CoreOps.reshape (CoreOps.range 0 (CoreOps.size indices) 1) prefixShape
    prefixShape = shapeInt32 indices
    shapeInt32 t = CoreOps.shape t :: Tensor Build Int32

opGrad "Select" _ [toT -> c, toT -> x, _] [dz] =
    [ Nothing
    , Just $ CoreOps.select c dz zeros
    , Just $ CoreOps.select c zeros dz
    ]
  where zeros = CoreOps.zerosLike x

-- TODO(gnezdo): Unlike Python, no control dependency on dz.
opGrad "Log" _ [toT -> x] [dz] = [ Just $ dz `CoreOps.mul` CoreOps.inv x ]
-- TODO(gnezdo): Reuse the output instead of doing another exp,
-- though, it is probably CSE'd away anyway.
opGrad "Exp" _ [toT -> x] [dz] = [ Just $ dz `CoreOps.mul` CoreOps.exp x ]
opGrad "SparseSegmentSum" _ [toT -> x, toT -> y, toT -> t] [dz] =
    [ Just $ CoreOps.unsortedSegmentSum
             (CoreOps.gather dz (t :: Tensor Build Int32))
             (y :: Tensor Build Int32) inputRows
    , Nothing
    , Nothing
    ]
  where inputRows = flatSlice (shape (x :: Tensor Build a)) 0 1

opGrad "LabelClasses" _ _ _ = [Nothing, Nothing]
opGrad "LabelWeights" _ _ _ = [Nothing]
opGrad "Size" _ _ _ = [Nothing]

-- TODO (jcberentsen): Python implementation uses set_shape for
-- static shape inference, which is unsupported.
-- TODO: implement support for static shape inference
opGrad "Tile" _ [toT -> x, toT -> multiples] [dz] =
    [Just inputGrad, Nothing]
  where
    inputGrad = sum reshapedDz axes
    inputShape = shape (x :: Tensor Build a)
    packed = CoreOps.pack [multiples, inputShape]
    perm = vector [1, 0 :: Int32]
    splitShape = CoreOps.reshape (CoreOps.transpose packed perm) allDimensions
    axes = CoreOps.range 0 (CoreOps.size splitShape) (2 :: Tensor Build Int32)
    reshapedDz = CoreOps.reshape dz splitShape

opGrad "ZerosLike" _ _ _ = [Nothing]
opGrad "Fill" _ _ [dz] = [Nothing, Just $ sum dz rx]
  where
    rx = rangeOfRank dz

-- Treat read ops as an identity function on the variable. This allows us to
-- take gradients w.r.t. to the variable handle instead of the result of a read
-- op. If a variable is read multiple times, the gradients will propagate back
-- through each read.
opGrad "ReadVariableOp" _ _ [dz] = [Just $ expr dz]

-- TODO(fmayle): These can go away if we properly prune the graph.
opGrad "Const" _ _ _ = [Nothing, Nothing]
opGrad "Placeholder" _ _ _ = []
opGrad "VarHandleOp" _ _ _ = []
opGrad "Variable" _ _ _ = []

opGrad n nodeDef ins grads =
    error $ "no gradient implemented for " ++
            show (n, length ins, length grads, showMessage nodeDef, ins)

-- | The number of outputs for an op type.
numOutputs :: NodeDef -> OutputIx
numOutputs o =
    case o ^. op of
        "Abs" -> 1
        "Add" -> 1
        "AddN" -> 1
        "Cast" -> 1
        "Const" -> 1
        "Concat" -> 1
        "Conv2D" -> 1
        "Conv2DBackpropInput" -> 1
        "Div" -> 1
        "DynamicStitch" -> 1
        "DynamicPartition" ->
            fromIntegral (lookupAttr o "num_partitions" :: Int64)
        "Exp" -> 1
        "Gather" -> 1
        "LabelClasses" -> 1
        "LabelWeights" -> 1
        "Log" -> 1
        "MatMul" -> 1
        "Max" -> 1
        "Maximum" -> 1
        "MaxPool" -> 1
        "Mean" -> 1
        "Min" -> 1
        "Mul" -> 1
        "Neg" -> 1
        "Placeholder" -> 1
        "OneHot" -> 1
        "ReadVariableOp" -> 1
        "RefIdentity" -> 1
        "Relu" -> 1
        "ReluGrad" -> 1
        "Reshape" -> 1
        "Select" -> 1
        "Size" -> 1
        "SoftmaxCrossEntropyWithLogits" -> 2
        "Square" -> 1
        "SparseSegmentSum" -> 1
        "Sub" -> 1
        "Sum" -> 1
        "Tile" -> 1
        "Transpose" -> 1
        "TruncatedNormal" -> 1
        "VarHandleOp" -> 1
        "Variable" -> 1
        "ZerosLike" -> 1
        "Fill" -> 1
        _ -> error $ "numOutputs not implemented for " ++ show (o ^. op)

-- Divides `x / y` assuming `x, y >= 0`, treating `0 / 0 = 0`
safeShapeDiv :: Tensor v1 Int32 -> Tensor v2 Int32 -> Tensor Build Int32
safeShapeDiv x y = x `CoreOps.div` (CoreOps.maximum y 1)

allDimensions :: Tensor Build Int32
allDimensions = vector [-1 :: Int32]

rangeOfRank :: forall v1 t. TensorType t => Tensor v1 t -> Tensor Build Int32
rangeOfRank x = CoreOps.range 0 (CoreOps.rank x) 1

lookupAttr ::  Attribute a1 => NodeDef -> Text -> a1
lookupAttr nodeDef attrName = nodeDef ^. attr . at attrName . non def . attrLens


{-# LANGUAGE KindSignatures #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE FunctionalDependencies #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE ViewPatterns #-}
{-# LANGUAGE DeriveGeneric #-}
{-# LANGUAGE DeriveAnyClass #-}
{-# LANGUAGE TypeApplications #-}

module NQS.Internal.BLAS
  ( AXPY(..)
  , COPY(..)
  , DOTC(..)
  , DOTU(..)
  , NRM2(..)
  , SCAL(..)
  , GEMV(..)
  , HERK(..)
  , fill
  , testGemv
  ) where

import Control.Monad.Primitive
import Control.Exception (assert)
import Data.Complex (Complex)
import Data.Proxy
import Foreign.Storable
import Foreign.Ptr
-- import Foreign.Marshal.Alloc

import Data.Semigroup((<>))
import Numerical.HBLAS.BLAS.FFI ( CBLAS_ORDERT(..)
                                , CBLAS_UPLOT(..)
                                , CBLAS_TRANSPOSET(..)
                                )
import Numerical.HBLAS.BLAS.FFI.Level1
import Numerical.HBLAS.BLAS.FFI.Level2
import Numerical.HBLAS.BLAS.FFI.Level3
import Numerical.HBLAS.Constants
import Numerical.HBLAS.UtilsFFI

import Data.Vector.Storable (Vector, MVector)
import qualified Data.Vector.Storable as V
import qualified Data.Vector.Storable.Mutable as MV

import Data.Singletons

import Lens.Micro((^.))

import NQS.Internal.Types
import NQS.Internal.FFI

instance ToNative Transpose CBLAS_TRANSPOSET where
  encode NoTranspose     = CBLAS_TransposeT 111
  encode Transpose       = CBLAS_TransposeT 112
  encode ConjTranspose   = CBLAS_TransposeT 113
  encode ConjNoTranspose = CBLAS_TransposeT 114

instance ToNative Orientation CBLAS_ORDERT where
  encode Row    = CBOInt 101
  encode Column = CBOInt 102

instance ToNative MatUpLo CBLAS_UPLOT where
  {-# INLINE encode #-}
  encode MatUpper = CBlasUPLO 121
  encode MatLower = CBlasUPLO 122

type ScalFun scale el m = scale -> MDenseVector 'Direct (PrimState m) el -> m ()

type AxpyFun el m = el -> MDenseVector 'Direct (PrimState m) el
                       -> MDenseVector 'Direct (PrimState m) el -> m ()

type DotFun el m res = MDenseVector 'Direct (PrimState m) el
                    -> MDenseVector 'Direct (PrimState m) el
                    -> m res

type CopyFun el m = MDenseVector 'Direct (PrimState m) el
                 -> MDenseVector 'Direct (PrimState m) el
                 -> m ()

type Nrm2Fun el m res = MDenseVector 'Direct (PrimState m) el -> m res

type GemvFun orient el m = Transpose
                        -> el -> MDenseMatrix orient (PrimState m) el
                        -> MDenseVector 'Direct (PrimState m) el
                        -> el -> MDenseVector 'Direct (PrimState m) el
                        -> m ()

type HerkFun scale el orient m
  = MatUpLo
 -> Transpose
 -> scale -> MDenseMatrix orient (PrimState m) el
 -> scale -> MDenseMatrix orient (PrimState m) el
 -> m ()

-- | Fills a vector with a given value.
fill ::
     (PrimMonad m, Storable el, COPY (MDenseVector 'Direct) el)
  => MDenseVector 'Direct (PrimState m) el -> el -> m ()
fill v x =
  V.unsafeThaw (V.singleton x) >>= \t -> copy (MDenseVector (v ^. dim)  0 t) v

{-
        SCAL
  ================
-}

class SCAL scale v el where
  scal :: PrimMonad m => scale -> v (PrimState m) el -> m ()

{-# SPECIALISE scal :: PrimMonad m => ScalFun Float Float m #-}
{-# SPECIALISE scal :: PrimMonad m => ScalFun Float (Complex Float) m #-}
{-# SPECIALISE scal :: PrimMonad m => ScalFun (Complex Float) (Complex Float) m #-}
{-# SPECIALISE scal :: PrimMonad m => ScalFun Double Double m #-}
{-# SPECIALISE scal :: PrimMonad m => ScalFun Double (Complex Double) m #-}
{-# SPECIALISE scal :: PrimMonad m => ScalFun (Complex Double) (Complex Double) m #-}

scalAbstraction ::
     (Storable el, PrimMonad m, Show el)
  => String
  -> ScalFunFFI scale el
  -> ScalFunFFI scale el
  -> (scaleplain -> (scale -> m ()) -> m ())
  -> ScalFun scaleplain el m
{-# NOINLINE scalAbstraction #-}
scalAbstraction scalName scalSafeFFI scalUnsafeFFI constHandler = doScal
  where
    shouldCallFast :: Int -> Bool
    shouldCallFast !n = flopsThreshold >= fromIntegral n
    doScal alpha (MDenseVector dim stride buffer)
      | not (isValidVector dim stride (MV.length buffer)) = error $!
          badVectorInfo scalName "X" dim stride (MV.length buffer)
      | otherwise =
        unsafeWithPrim buffer $ \xPtr ->
          constHandler alpha $ \alphaPtr -> do
            unsafePrimToPrim $!
              (if shouldCallFast dim
                 then scalUnsafeFFI
                 else scalSafeFFI)
                (fromIntegral dim)
                alphaPtr
                xPtr (fromIntegral stride)

instance SCAL Float (MDenseVector 'Direct) Float where
  scal = scalAbstraction "sscal" cblas_sscal_safe cblas_sscal_unsafe (\x f -> f x)

instance SCAL Float (MDenseVector 'Direct) (Complex Float) where
  scal = scalAbstraction "csscal" cblas_csscal_safe cblas_csscal_unsafe (\x f -> f x)

instance SCAL (Complex Float) (MDenseVector 'Direct) (Complex Float) where
  scal = scalAbstraction "cscal" cblas_cscal_safe cblas_cscal_unsafe withRStorable_

instance SCAL Double (MDenseVector 'Direct) Double where
  scal = scalAbstraction "dscal" cblas_dscal_safe cblas_dscal_unsafe (\x f -> f x)

instance SCAL Double (MDenseVector 'Direct) (Complex Double) where
  scal = scalAbstraction "zdscal" cblas_zdscal_safe cblas_zdscal_unsafe (\x f -> f x)

instance SCAL (Complex Double) (MDenseVector 'Direct) (Complex Double) where
  scal = scalAbstraction "zscal" cblas_zscal_safe cblas_zscal_unsafe withRStorable_

{-
        AXPY
  ================
-}

class AXPY v el where
  axpy :: PrimMonad m => el -> v (PrimState m) el -> v (PrimState m) el -> m ()

{-# SPECIALISE axpy :: PrimMonad m => AxpyFun Float m #-}
{-# SPECIALISE axpy :: PrimMonad m => AxpyFun (Complex Float) m #-}
{-# SPECIALISE axpy :: PrimMonad m => AxpyFun Double m #-}
{-# SPECIALISE axpy :: PrimMonad m => AxpyFun (Complex Double) m #-}

axpyAbstraction ::
     (PrimMonad m, Storable el, Show el)
  => String
  -> AxpyFunFFI scale el
  -> AxpyFunFFI scale el
  -> (el -> (scale -> m ()) -> m ())
  -> AxpyFun el m
{-# NOINLINE axpyAbstraction #-}
axpyAbstraction axpyName axpySafeFFI axpyUnsafeFFI constHandler = doAxpy
  where
    shouldCallFast :: Int -> Bool
    shouldCallFast !n = flopsThreshold >= 2 * (fromIntegral n) -- n for a*x, and n for +y
    doAxpy alpha (MDenseVector xDim xStride xBuff)
                 (MDenseVector yDim yStride yBuff)
      | not (isValidVector xDim xStride (MV.length xBuff)) = error $!
        badVectorInfo axpyName "X" xDim xStride (MV.length xBuff)
      | not (isValidVector yDim yStride (MV.length yBuff)) = error $!
        badVectorInfo axpyName "Y" yDim yStride (MV.length yBuff)
      | xDim /= yDim = error $! axpyName <> ": Inconsistent dimensions: " <>
                                show xDim <> " != " <> show yDim <> "."
      | otherwise =
        unsafeWithPrim xBuff $ \xPtr ->
          unsafeWithPrim yBuff $ \yPtr ->
            constHandler alpha $ \alphaPtr -> do
              unsafePrimToPrim $!
                (if shouldCallFast xDim
                   then axpyUnsafeFFI
                   else axpySafeFFI)
                  (fromIntegral xDim)
                  alphaPtr
                  xPtr (fromIntegral xStride)
                  yPtr (fromIntegral yStride)

instance AXPY (MDenseVector 'Direct) Float where
  axpy = axpyAbstraction "saxpy" cblas_saxpy_safe cblas_saxpy_unsafe (\x f -> f x)

instance AXPY (MDenseVector 'Direct) Double where
  axpy = axpyAbstraction "daxpy" cblas_daxpy_safe cblas_daxpy_unsafe (\x f -> f x)

instance AXPY (MDenseVector 'Direct) (Complex Float) where
  axpy = axpyAbstraction "caxpy" cblas_caxpy_safe cblas_caxpy_unsafe withRStorable_

instance AXPY (MDenseVector 'Direct) (Complex Double) where
  axpy = axpyAbstraction "zaxpy" cblas_zaxpy_safe cblas_zaxpy_unsafe withRStorable_

{-
        CDOT
  ================
-}

class DOTU v el res where
  dotu :: PrimMonad m => v (PrimState m) el -> v (PrimState m) el -> m res

{-# SPECIALISE dotu :: PrimMonad m => DotFun Float m Float #-}
{-# SPECIALISE dotu :: PrimMonad m => DotFun Float m Double #-}
{-# SPECIALISE dotu :: PrimMonad m => DotFun Double m Double #-}
{-# SPECIALISE dotu :: PrimMonad m => DotFun (Complex Float) m (Complex Float) #-}
{-# SPECIALISE dotu :: PrimMonad m => DotFun (Complex Double) m (Complex Double) #-}

class DOTC v el res where
  dotc :: PrimMonad m => v (PrimState m) el -> v (PrimState m) el -> m res

{-# SPECIALISE dotc :: PrimMonad m => DotFun Float m Float #-}
{-# SPECIALISE dotc :: PrimMonad m => DotFun Float m Double #-}
{-# SPECIALISE dotc :: PrimMonad m => DotFun Double m Double #-}
{-# SPECIALISE dotc :: PrimMonad m => DotFun (Complex Float) m (Complex Float) #-}
{-# SPECIALISE dotc :: PrimMonad m => DotFun (Complex Double) m (Complex Double) #-}

realDotAbstraction ::
     (PrimMonad m, Storable el, Num el, Show el)
  => String
  -> NoScalarDotFunFFI el res
  -> NoScalarDotFunFFI el res
  -> DotFun el m res
{-# NOINLINE realDotAbstraction #-}
realDotAbstraction dotName dotSafeFFI dotUnsafeFFI = doDot
  where
    shouldCallFast :: Int -> Bool
    shouldCallFast n = flopsThreshold >= fromIntegral n
    doDot (MDenseVector xDim xStride xBuff)
          (MDenseVector yDim yStride yBuff)
      | not (isValidVector xDim xStride (MV.length xBuff)) = error $!
        badVectorInfo dotName "X" xDim xStride (MV.length xBuff)
      | not (isValidVector yDim yStride (MV.length yBuff)) = error $!
        badVectorInfo dotName "Y" yDim yStride (MV.length yBuff)
      | xDim /= yDim = error $! dotName <> ": Inconsistent dimensions: " <>
                                show xDim <> " != " <> show yDim <> "."
      | otherwise =
        unsafeWithPrim xBuff $ \xPtr ->
          unsafeWithPrim yBuff $ \yPtr ->
            unsafePrimToPrim $!
              (if shouldCallFast xDim
                 then dotUnsafeFFI
                 else dotSafeFFI)
                (fromIntegral xDim)
                xPtr (fromIntegral xStride)
                yPtr (fromIntegral yStride)

complexDotAbstraction ::
     (PrimMonad m, Storable el, Num el, Show el)
  => String
  -> ComplexDotFunFFI el
  -> ComplexDotFunFFI el
  -> DotFun el m el
{-# NOINLINE complexDotAbstraction #-}
complexDotAbstraction dotName dotSafeFFI dotUnsafeFFI = doDot
  where
    shouldCallFast :: Int -> Bool
    shouldCallFast n = flopsThreshold >= fromIntegral n
    doDot (MDenseVector xDim xStride xBuff)
          (MDenseVector yDim yStride yBuff)
      | not (isValidVector xDim xStride (MV.length xBuff)) = error $!
        badVectorInfo dotName "X" xDim xStride (MV.length xBuff)
      | not (isValidVector yDim yStride (MV.length yBuff)) = error $!
        badVectorInfo dotName "Y" yDim yStride (MV.length yBuff)
      | xDim /= yDim = error $! dotName <> ": Inconsistent dimensions: " <>
                                show xDim <> " != " <> show yDim <> "."
      | otherwise =
        unsafeWithPrim xBuff $ \xPtr ->
          unsafeWithPrim yBuff $ \yPtr ->
            unsafePrimToPrim $!
              alloca $ \outPtr -> do
                poke outPtr 0
                (if shouldCallFast xDim
                   then dotUnsafeFFI
                   else dotSafeFFI)
                  (fromIntegral xDim)
                  xPtr (fromIntegral xStride)
                  yPtr (fromIntegral yStride)
                  outPtr
                peek outPtr

instance DOTC (MDenseVector 'Direct) Float Float where
  dotc = realDotAbstraction "sdot" cblas_sdot_safe cblas_sdot_unsafe

instance DOTC (MDenseVector 'Direct) Double Double where
  dotc = realDotAbstraction "ddot" cblas_ddot_safe cblas_ddot_unsafe

instance DOTC (MDenseVector 'Direct) Float Double where
  dotc = realDotAbstraction "dsdot" cblas_dsdot_safe cblas_dsdot_unsafe

instance DOTC (MDenseVector 'Direct) (Complex Float) (Complex Float) where
  dotc = complexDotAbstraction "cdotc" cblas_cdotc_safe cblas_cdotc_unsafe

instance DOTC (MDenseVector 'Direct) (Complex Double) (Complex Double) where
  dotc = complexDotAbstraction "zdotc" cblas_zdotc_safe cblas_zdotc_unsafe

instance DOTU (MDenseVector 'Direct) Float Float where
  dotu = realDotAbstraction "sdot" cblas_sdot_safe cblas_sdot_unsafe

instance DOTU (MDenseVector 'Direct) Double Double where
  dotu = realDotAbstraction "ddot" cblas_ddot_safe cblas_ddot_unsafe

instance DOTU (MDenseVector 'Direct) Float Double where
  dotu = realDotAbstraction "dsdot" cblas_dsdot_safe cblas_dsdot_unsafe

instance DOTU (MDenseVector 'Direct) (Complex Float) (Complex Float) where
  dotu = complexDotAbstraction "cdotu" cblas_cdotu_safe cblas_cdotu_unsafe

instance DOTU (MDenseVector 'Direct) (Complex Double) (Complex Double) where
  dotu = complexDotAbstraction "zdotu" cblas_zdotu_safe cblas_zdotu_unsafe

class COPY v el where
  copy :: PrimMonad m => v (PrimState m) el -> v (PrimState m) el -> m ()

{-# SPECIALISE copy :: PrimMonad m => CopyFun Float m #-}
{-# SPECIALISE copy :: PrimMonad m => CopyFun Double m #-}
{-# SPECIALISE copy :: PrimMonad m => CopyFun (Complex Float) m #-}
{-# SPECIALISE copy :: PrimMonad m => CopyFun (Complex Double) m #-}

copyAbstraction ::
     (PrimMonad m, Storable el, Num el, Show el)
  => String
  -> CopyFunFFI el
  -> CopyFunFFI el
  -> CopyFun el m
{-# NOINLINE copyAbstraction #-}
copyAbstraction copyName copySafeFFI copyUnsafeFFI = doCopy
  where
    shouldCallFast :: Int -> Bool
    shouldCallFast _ = True
    doCopy (MDenseVector xDim xStride xBuff)
           (MDenseVector yDim yStride yBuff)
      | not (isValidVector xDim xStride (MV.length xBuff)) = error $!
        badVectorInfo copyName "X" xDim xStride (MV.length xBuff)
      | not (isValidVector yDim yStride (MV.length yBuff)) = error $!
        badVectorInfo copyName "Y" yDim yStride (MV.length yBuff)
      | xDim /= yDim = error $! copyName <> ": Inconsistent dimensions: " <>
                                show xDim <> " != " <> show yDim <> "."
      | otherwise =
        unsafeWithPrim xBuff $ \xPtr ->
          unsafeWithPrim yBuff $ \yPtr ->
            unsafePrimToPrim $!
              (if shouldCallFast xDim
                 then copyUnsafeFFI
                 else copySafeFFI)
                (fromIntegral xDim)
                xPtr (fromIntegral xStride)
                yPtr (fromIntegral yStride)

instance COPY (MDenseVector 'Direct) Float where
  copy = copyAbstraction "scopy" cblas_scopy_safe cblas_scopy_unsafe

instance COPY (MDenseVector 'Direct) Double where
  copy = copyAbstraction "dcopy" cblas_dcopy_safe cblas_dcopy_unsafe

instance COPY (MDenseVector 'Direct) (Complex Float) where
  copy = copyAbstraction "ccopy" cblas_ccopy_safe cblas_ccopy_unsafe

instance COPY (MDenseVector 'Direct) (Complex Double) where
  copy = copyAbstraction "zcopy" cblas_zcopy_safe cblas_zcopy_unsafe

class NRM2 v el res | v el -> res where
  nrm2 :: PrimMonad m => v (PrimState m) el -> m res

{-# SPECIALISE nrm2 :: PrimMonad m => Nrm2Fun Float m Float #-}
{-# SPECIALISE nrm2 :: PrimMonad m => Nrm2Fun Double m Double #-}
{-# SPECIALISE nrm2 :: PrimMonad m => Nrm2Fun (Complex Float) m Float #-}
{-# SPECIALISE nrm2 :: PrimMonad m => Nrm2Fun (Complex Double) m Double #-}

nrm2Abstraction ::
     (PrimMonad m, Storable el, Num el, Show el)
  => String
  -> Nrm2FunFFI el res
  -> Nrm2FunFFI el res
  -> Nrm2Fun el m res
{-# NOINLINE nrm2Abstraction #-}
nrm2Abstraction nrm2Name nrm2SafeFFI nrm2UnsafeFFI = doNrm2
  where
    shouldCallFast :: Int -> Bool
    shouldCallFast n = fromIntegral n <= flopsThreshold
    doNrm2 (MDenseVector xDim xStride xBuff)
      | not (isValidVector xDim xStride (MV.length xBuff)) = error $!
        badVectorInfo nrm2Name "X" xDim xStride (MV.length xBuff)
      | otherwise =
        unsafeWithPrim xBuff $ \xPtr ->
            unsafePrimToPrim $!
              (if shouldCallFast xDim
                 then nrm2UnsafeFFI
                 else nrm2SafeFFI) (fromIntegral xDim) xPtr (fromIntegral xStride)

instance NRM2 (MDenseVector 'Direct) Float Float where
  nrm2 = nrm2Abstraction "snrm2" cblas_snrm2_safe cblas_snrm2_unsafe

instance NRM2 (MDenseVector 'Direct) Double Double where
  nrm2 = nrm2Abstraction "dnrm2" cblas_dnrm2_safe cblas_dnrm2_unsafe

instance NRM2 (MDenseVector 'Direct) (Complex Float) Float where
  nrm2 = nrm2Abstraction "scnrm2" cblas_scnrm2_safe cblas_scnrm2_unsafe

instance NRM2 (MDenseVector 'Direct) (Complex Double) Double where
  nrm2 = nrm2Abstraction "dznrm2" cblas_dznrm2_safe cblas_dznrm2_unsafe

class GEMV mat v el where
  gemv :: PrimMonad m
       => Transpose
       -> el -> mat (PrimState m) el -> v (PrimState m) el
       -> el -> v (PrimState m) el -> m ()

gemvAbstraction :: forall m el scale orient.
     (PrimMonad m, Storable el, Show el, SingI orient)
  => String
  -> GemvFunFFI scale el
  -> GemvFunFFI scale el
  -> (el -> (scale -> m ()) -> m ())
  -> GemvFun orient el m
{-# NOINLINE gemvAbstraction #-}
gemvAbstraction gemvName gemvSafeFFI gemvUnsafeFFI constHandler = doGemv
  where
    shouldCallFast :: Int -> Int -> Bool
    shouldCallFast !n !m = flopsThreshold >= (fromIntegral n) * (fromIntegral m)
    doGemv trans alpha a@(MDenseMatrix aRows aCols aStride aBuff)
                       x@(MDenseVector xDim xStride xBuff)
                  beta y@(MDenseVector yDim yStride yBuff) =
      assertValid gemvName "A" a $
      assertValid gemvName "X" x $
      assertValid gemvName "Y" y $
        unsafeWithPrim aBuff $ \aPtr ->
        unsafeWithPrim xBuff $ \xPtr ->
        unsafeWithPrim yBuff $ \yPtr ->
        constHandler alpha $ \alphaPtr ->
        constHandler beta $ \betaPtr ->
          do
            unsafePrimToPrim $!
              (if shouldCallFast xDim yDim then gemvUnsafeFFI else gemvSafeFFI)
                (encode (orientationOf a)) (encode trans)
                (fromIntegral aRows) (fromIntegral aCols)
                alphaPtr
                aPtr (fromIntegral aStride)
                xPtr (fromIntegral xStride)
                betaPtr
                yPtr (fromIntegral yStride)

instance SingI orient =>
  GEMV (Mutable (DenseMatrix orient)) (MDenseVector 'Direct) Float where
    gemv = gemvAbstraction "sgemv" cblas_sgemv_safe cblas_sgemv_unsafe (\x f -> f x)

instance SingI orient =>
  GEMV (Mutable (DenseMatrix orient)) (MDenseVector 'Direct) Double where
    gemv = gemvAbstraction "dgemv" cblas_dgemv_safe cblas_dgemv_unsafe (\x f -> f x)

instance SingI orient =>
  GEMV (Mutable (DenseMatrix orient)) (MDenseVector 'Direct) (Complex Float) where
    gemv = gemvAbstraction "cgemv" cblas_cgemv_safe cblas_cgemv_unsafe withRStorable_

instance SingI orient =>
  GEMV (Mutable (DenseMatrix orient)) (MDenseVector 'Direct) (Complex Double) where
    gemv = gemvAbstraction "zgemv" cblas_zgemv_safe cblas_zgemv_unsafe withRStorable_

testGemv :: IO ()
testGemv = do
  aBuff <- newVectorAligned 12 64 :: IO (MV.IOVector Double)
  V.copy aBuff (V.fromList [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12])
  let a = MDenseMatrix @'Row 4 3 3 aBuff
  print (orientationOf a)
  xBuff <- newVectorAligned 4 64 :: IO (MV.IOVector Double)
  V.copy xBuff (V.fromList [1, 2, 3, 4])
  let x = MDenseVector @'Direct 3 1 xBuff
  yBuff <- newVectorAligned 3 64 :: IO (MV.IOVector Double)
  -- V.copy yBuff (V.fromList [1, 2, 3])
  let y = MDenseVector @'Direct 3 1 yBuff
  gemv Transpose 1.0 a x 0.0 y
  print =<< V.unsafeFreeze yBuff

class HERK mat scale el where
  herk :: PrimMonad m
       => MatUpLo -> Transpose
       -> scale -> mat (PrimState m) el
       -> scale -> mat (PrimState m) el
       -> m ()

herkAbstraction
  :: forall m el scale scalePtr. (Storable el, PrimMonad m)
  => String
  -> HerkFunFFI scalePtr el
  -> HerkFunFFI scalePtr el
  -> (scale -> (scalePtr -> m ()) -> m ())
  -> (forall orient . SingI orient => HerkFun scale el orient m)
{-# NOINLINE herkAbstraction #-}
herkAbstraction herkName herkSafeFFI herkUnsafeFFI constHandler = doHerk
 where
  isBadHerk :: Transpose -> Int -> Int -> Int -> Int -> Bool
  isBadHerk NoTranspose   ax ay cx cy = (ax /= cx)
  isBadHerk ConjTranspose ax ay cx cy = (ay /= cx)
  isBadHerk trans _ _ _ _ =
    error $! herkName <> ": trans " <> show trans <> " is not supported."
  doHerk :: forall orient. (SingI orient) => HerkFun scale el orient m
  doHerk uplo trans alpha a@(MDenseMatrix ax ay astride abuff) beta c@(MDenseMatrix cx cy cstride cbuff)
    | cx /= cy = error $!
      herkName <> ": C must be a square matrix, but has dimensions " <> show (cx, cy) <> "."
    | isBadHerk trans ax ay cx cy = error $!
      herkName <> ": Incompatible dimensions: ax ay cx cy trans: " <> show [ax, ay, cx, cy]
               <> " " <> show trans
    | MV.overlaps abuff cbuff = error $! herkName <> ": Input and output buffers overlap."
    | otherwise =
      withMatrixFFI a $ \layout aRows aCols aPtr aStride ->
      withMatrixFFI c $ \_      cRows _     cPtr cStride ->
      constHandler alpha $ \alphaPtr ->
      constHandler beta $ \betaPtr ->
        let k = if (trans == NoTranspose) then aCols else aRows
        in  unsafePrimToPrim $!
          herkSafeFFI ((encode . fromSing) (sing :: Sing orient))
                      (encode uplo)
                      (encode trans)
                      cRows
                      k
                      alphaPtr
                      aPtr
                      aStride
                      betaPtr
                      cPtr
                      cStride

instance SingI orient => HERK (Mutable (DenseMatrix orient)) Float (Complex Float) where
  herk = herkAbstraction "cherk" cblas_cherk_safe cblas_cherk_unsafe (\x f -> f x)

instance SingI orient => HERK (Mutable (DenseMatrix orient)) Double (Complex Double) where
  herk = herkAbstraction "zherk" cblas_zherk_safe cblas_zherk_unsafe (\x f -> f x)


{-# LANGUAGE  ForeignFunctionInterface #-}
-----------------------------------------------------------------------------
-- |
-- Module     : Data.Elem.LAPACK.Zomplex
-- Copyright  : Copyright (c) , Patrick Perry <patperry@stanford.edu>
-- License    : BSD3
-- Maintainer : Patrick Perry <patperry@stanford.edu>
-- Stability  : experimental
--

module Data.Elem.LAPACK.Zomplex
    where

import Data.Complex( Complex )
import Foreign( Ptr )
import LAPACK.CTypes

type Zomplex = Complex Double

foreign import ccall unsafe "LAPACK.h lapack_zgeqrf"
    zgeqrf :: Int -> Int -> Ptr Zomplex -> Int -> Ptr Zomplex -> Ptr Zomplex -> Int -> IO Int

foreign import ccall unsafe "LAPACK.h lapack_zgelqf"
    zgelqf :: Int -> Int -> Ptr Zomplex -> Int -> Ptr Zomplex -> Ptr Zomplex -> Int -> IO Int

foreign import ccall unsafe "LAPACK.h lapack_zunmqr"
    zunmqr :: CBLASSide -> CBLASTrans -> Int -> Int -> Int -> Ptr Zomplex -> Int -> Ptr Zomplex
           -> Ptr Zomplex -> Int -> Ptr Zomplex -> Int -> IO Int

foreign import ccall unsafe "LAPACK.h lapack_zunmlq"
    zunmlq :: CBLASSide -> CBLASTrans -> Int -> Int -> Int -> Ptr Zomplex -> Int -> Ptr Zomplex
           -> Ptr Zomplex -> Int -> Ptr Zomplex -> Int -> IO Int

foreign import ccall unsafe "LAPACK.h lapack_zlarfg"
    zlarfg :: Int -> Ptr Zomplex -> Ptr Zomplex -> Int -> Ptr Zomplex -> IO ()