{-# LANGUAGE DeriveAnyClass, DeriveGeneric, TupleSections #-}

module Test.Circuit.Arithmetic where

import           Circuit.Affine
import           Circuit.Arithmetic
import           Data.Curve.Weierstrass.BN254 (Fr)
import           Data.Map                     (Map)
import qualified Data.Map                     as Map
import           Data.Pairing.BN254           (getRootOfUnity)
import           Fresh
import           Protolude
import           QAP
import           Test.Tasty.HUnit
import           Test.Tasty.QuickCheck

-------------------------------------------------------------------------------
-- Test values
-------------------------------------------------------------------------------

testEqualCircuit :: ArithCircuit Fr
testEqualCircuit = ArithCircuit [Equal (InputWire 0) (IntermediateWire 0) (OutputWire 0)]

testInputMap :: Fr -> Map Int Fr
testInputMap = Map.singleton 0

testSplitUnsplitCircuit :: Int -> ArithCircuit Fr
testSplitUnsplitCircuit nbits =
  ArithCircuit
    [ Split (InputWire 0) midWires,
      Mul (ConstGate 1) (unsplit midWires) (OutputWire 0)
    ]
  where
    midWires = fmap IntermediateWire [0 .. nbits - 1]

-------------------------------------------------------------------------------
-- Generators
-------------------------------------------------------------------------------

arbVars :: Int -> [Int] -> [Gen (AffineCircuit Wire f)]
arbVars numInps mids =
  varInps numInps ++ varMids mids
  where
    varInps size
      | size <= 0 = []
      | otherwise = [Var . InputWire <$> choose (0, numInps - 1)]
    varMids [] = []
    varMids ms@(_ : _) = [Var . IntermediateWire <$> elements ms]

arbAffineCircuitWithMids ::
  Arbitrary f =>
  Int ->
  [Int] ->
  Int ->
  Gen (AffineCircuit Wire f)
arbAffineCircuitWithMids numInps mids size
  | size <= 0 =
    oneof $ [ConstGate <$> arbitrary] ++ arbVars numInps mids
  | size > 0 =
    oneof
      [ ScalarMul <$> arbitrary <*> arbAffineCircuitWithMids numInps mids (size - 1),
        Add <$> arbAffineCircuitWithMids numInps mids (size - 1)
          <*> arbAffineCircuitWithMids numInps mids (size - 1)
      ]

arbInputVector :: (Arbitrary f) => Int -> Gen (Map Int f)
arbInputVector numVars = Map.fromList . zip [0 ..] <$> vector numVars

arbArithCircuit ::
  Arbitrary f =>
  -- | distribution of frequency of mul/equal/split
  -- gates, respectively
  (Int, Int, Int) ->
  Int ->
  Int ->
  Gen (ArithCircuit f)
arbArithCircuit (distMul, distEqual, distSplit) numInps size
  | size <= 0 =
    pure $ ArithCircuit []
  | size > 0 =
    do
      ArithCircuit gates <- arbArithCircuit (distMul, distEqual, distSplit) numInps (size - 1)
      let mids = [i | IntermediateWire i <- concatMap outputWires gates]
      frequency . catMaybes $
        [ (distMul,) <$> mulGate gates mids,
          (distEqual,) <$> equalGate gates mids,
          (distSplit,) <$> splitGate gates mids
        ]
  where
    mulGate gates mids =
      Just $ do
        lhs <- arbAffineCircuitWithMids numInps mids 1
        rhs <- arbAffineCircuitWithMids numInps mids 1
        let outWire = case mids of
              [] -> 0
              ms@(_ : _) -> maximum ms + 1
            gate = Mul lhs rhs (IntermediateWire outWire)
        pure . ArithCircuit $ gates ++ [gate]
    equalGate _ [] =
      Nothing
    equalGate gates mids@(_ : _) =
      Just $ do
        inp <- elements mids
        let outWire =
              case mids of
                [] -> 0
                ms@(_ : _) -> maximum ms + 1
            gate =
              Equal
                (IntermediateWire inp)
                (IntermediateWire outWire)
                (IntermediateWire $ outWire + 1)
        pure . ArithCircuit $ gates ++ [gate]
    splitGate _ [] =
      Nothing
    splitGate gates mids@(_ : _) =
      Just $ do
        inp <- IntermediateWire <$> elements mids
        let firstOutWire =
              case mids of
                [] -> 0
                ms@(_ : _) -> maximum ms + 1
            nbits = 256
            outWires = fmap IntermediateWire [firstOutWire .. firstOutWire + nbits - 1]
            gate = Split inp outWires
        pure . ArithCircuit $ gates ++ [gate]

-- | The input vector has to have the correct length, so we want to
-- generate the program and the test input simultaneously.
data ArithCircuitWithInputs f = ArithCircuitWithInputs (ArithCircuit f) [Map Int f]
  deriving (Show, Generic, NFData)

instance (Arbitrary f, Num f) => Arbitrary (ArithCircuitWithInputs f) where
  arbitrary = do
    numVars <- abs <$> arbitrary `suchThat` (> 0)
    program <- sized (arbArithCircuit (50, 10, 1) numVars)
    inputs <- vectorOf 5 $ arbInputVector numVars
    pure $ ArithCircuitWithInputs program inputs

data ArithCircuitWithInput f = ArithCircuitWithInput (ArithCircuit f) (Map Int f)
  deriving (Show, Generic, NFData)

instance (Arbitrary f, Num f) => Arbitrary (ArithCircuitWithInput f) where
  arbitrary = do
    numVars <- abs <$> arbitrary `suchThat` (> 0)
    program <- sized (arbArithCircuit (50, 10, 1) numVars)
    input <- arbInputVector numVars
    pure $ ArithCircuitWithInput program input

-------------------------------------------------------------------------------
-- Tests
-------------------------------------------------------------------------------

unit_eqGate ::
  Assertion
unit_eqGate =
  do
    testEvalWith 0 @?= Just 0
    testEvalWith 1 @?= Just 1
    testEvalWith 2 @?= Just 1
    testEvalWith 3 @?= Just 1
  where
    testEvalWith n =
      lookupAtWire (OutputWire 0) $
        evalArithCircuit
          lookupAtWire
          updateAtWire
          testEqualCircuit
          (initialQapSet $ testInputMap n)

unit_splitUnsplit :: Assertion
unit_splitUnsplit =
  mapM_ (\n -> testSplitUnsplit n @?= Just n) (fmap fromIntegral [0 .. 2 ^ nbits - 1])
  where
    nbits = 16
    testSplitUnsplit n =
      lookupAtWire (OutputWire 0) $
        evalArithCircuit
          lookupAtWire
          updateAtWire
          (testSplitUnsplitCircuit nbits)
          (initialQapSet $ testInputMap n)

prop_arithCircuitValid :: ArithCircuitWithInputs Fr -> Bool
prop_arithCircuitValid (ArithCircuitWithInputs program _) =
  validArithCircuit program

prop_arithCircuitToQAP_slow :: ArithCircuitWithInputs Fr -> Property
prop_arithCircuitToQAP_slow (ArithCircuitWithInputs program inputs) =
  withMaxSuccess 10 $
    case program of
      ArithCircuit [] -> True
      ArithCircuit (_ : _) -> all testInput inputs
  where
    roots = evalFresh $ generateRoots (fromIntegral . (+ 1) <$> fresh) program
    qap = arithCircuitToQAP roots program
    testInput input =
      verifyAssignment qap $ generateAssignment program input

prop_arithCircuitToQAP_fft :: ArithCircuitWithInputs Fr -> Property
prop_arithCircuitToQAP_fft (ArithCircuitWithInputs program inputs) =
  withMaxSuccess 10 $
    case program of
      ArithCircuit [] -> True
      ArithCircuit (_ : _) -> all testInput inputs
  where
    roots = evalFresh $ generateRoots (fromIntegral . (+ 1) <$> fresh) program
    qap = createPolynomialsFFT getRootOfUnity $ arithCircuitToGenQAP roots program
    testInput input = verifyAssignment qap $ generateAssignment program input