{-# LANGUAGE ScopedTypeVariables #-}

-- By Christian Oliveros and Minmin Chen
module PingPong.Player.Stig (stig)
where

import Control.Lens (view, (&), (.~), (^.))
import Data.Bool (bool)
import Data.Ext
import Data.Geometry hiding (head, Above, Below)
import Data.Maybe
import Debug.Trace
import GHC.Float
import qualified Numeric.LinearAlgebra as Numerical
import PingPong.Model
import PingPong.Player.Stig.General
import PingPong.Player.Stig.GeometryHelpers
import PingPong.Player.Stig.Kinematics
import PingPong.Player.Stig.Fabrik

import Control.Monad

-- | Simulation's gravity
simulationGravity :: Float
simulationGravity = 2

-- | Simulation's max speed
simulationMaxSpeed :: Float
simulationMaxSpeed = 2

-- | Simulation's Table Height
simulationTableHeight :: Float
simulationTableHeight = 0.5

-- | Simulation's Table Center X Position
simulationTableCenterX :: Float
simulationTableCenterX = 0

-- | Simulation's Table Max X Position
simulationTableMaxX :: Float
simulationTableMaxX = 1

-- | Normalized direction of the table
simulationTableDir :: Vector 2 Float
simulationTableDir = Vector2 (-1) 0

-- | Bat's Length
simulationBatLength :: Float
simulationBatLength = 0.1

-- | Simulation's Table Center X Position
simulationTableOpponentCenterX :: Float
simulationTableOpponentCenterX = -0.5

-- | Simulation's Table Max X Position
simulationTableOpponentMinX :: Float
simulationTableOpponentMinX = -0.8

-- | If the x value is inside the opponent's table range
insideOpponentTableX :: Float -> DistanceToRange Float
insideOpponentTableX = signedDistanceToRange simulationTableOpponentMinX simulationTableOpponentCenterX

-- | Predict the time (relative to current time) the gravity parabole is going to intersect a height
predictFreefallHeightInter :: Point 2 Float -> Vector 2 Float -> Float -> Maybe Float
predictFreefallHeightInter p v tH = bool res Nothing (isNothing possible || null tsRaw)
  where
    a2 = -simulationGravity / 2
    a1 = view yComponent v
    a0 = view yCoord p - tH
    possible = solveQuadratic a2 a1 a0
    tsRaw = fromJust possible
    -- If solveQuadratic is [] change for 0, because all t are valid
    -- Also, make sure the times are valid between 0 and 1 (threshold is used for approximations)
    ts = filter (0 <=) $ map (globalThreshold 0) $ bool tsRaw [0] (null tsRaw)
    res = case possible of
        Nothing -> Nothing
        Just _ -> case ts of
                    [] -> Nothing -- No solution
                    _ -> Just $ minimum ts -- Minimum Valid Time

-- | Evaluate the free fall formulas for a time
freeFallEvaluateTime :: Point 2 Float -> Vector 2 Float -> Float -> (Point 2 Float, Vector 2 Float)
freeFallEvaluateTime p v t = (Point2 x y, Vector2 vx vy)
  where
    vx = view xComponent v
    x = vx * t + view xCoord p

    vy0 = view yComponent v
    vy = -simulationGravity * t + vy0
    y = -simulationGravity * t * t / 2 + vy0 * t + view yCoord p

-- | Reflect a velocity vector againts the table like in a collision
reflectVelocityTable :: Vector 2 Float -> Vector 2 Float
reflectVelocityTable v = reflect v simulationTableDir

-- | Reflect Ball Velocity against a bat
reflectVelocityBat :: LineSegment 2 () Float -> Vector 2 Float -> Vector 2 Float
reflectVelocityBat bat v = reflect v (segmentDir bat)

-- | Get the normal of the free fall velocity
freefallNormal :: Vector 2 Float -> Vector 2 Float
freefallNormal v = n
  where
    nUnorm = Vector2 (-view yComponent v) (view xComponent v)
    (n, _) = normalizeVector nUnorm

-- | Get the (relative) time of maximum height of a freefall. Note, it might be negative, meaning it already peaked
freeFallMaxHeightTime :: Vector 2 Float -> Float
freeFallMaxHeightTime v = view yComponent v / simulationGravity

-- | Calculates the (relative) time to reach a x position in a freefall. Note, it might be negative, meaning it already passed
freeFallTimeToXPos :: Point 2 Float -> Vector 2 Float -> Float -> Float
freeFallTimeToXPos p v x = (x - view xCoord p) / view xComponent v

-- | Predict the info of the ball bouncing on our side of the table
predictTableBounce :: Point 2 Float -> Vector 2 Float -> Maybe (Point 2 Float, Vector 2 Float, Float)
predictTableBounce p v = res
  where
    tPossible = predictFreefallHeightInter p v simulationTableHeight
    (pB, vB) = freeFallEvaluateTime p v (fromJust tPossible)
    xPos = view xCoord pB
    xPosValid =  simulationTableCenterX <= xPos && xPos <= simulationTableMaxX
    res = case tPossible of
            Nothing -> Nothing
            Just t -> if xPosValid then return (pB, reflectVelocityTable vB, t) else Nothing

-- | Checks if a position is valid for our side of the table
checkValidPos :: Point 2 Float -> Bool
checkValidPos p0 = simulationTableCenterX <= xPos && xPos <= simulationTableMaxX && simulationTableHeight <= yPos
  where
    xPos = view xCoord p0
    yPos = view yCoord p0

-- | Try to predict the best interception time
predictBestInterceptionTime :: Point 2 Float -> Vector 2 Float -> Maybe Float
predictBestInterceptionTime p v = res
  where
    ts = [freeFallMaxHeightTime v, freeFallTimeToXPos p v simulationTableMaxX, freeFallTimeToXPos p v (simulationTableMaxX / 2), freeFallTimeToXPos p v (simulationTableMaxX / 4)]
    tF = filter tFilter ts

    res =  if null tF then Nothing else Just $ traceShowId $ minimum tF

    tFilter :: Float -> Bool
    tFilter t = 0 <= t && checkValidPos p0
      where
        (p0, _) = freeFallEvaluateTime p v t

-- | Place the bat normal to a ball point and velocity
placeBatInBounceCurve :: Point 2 Float -> Vector 2 Float -> Float -> LineSegment 2 () Float
placeBatInBounceCurve p v q = line
  where

    -- Normal to Velocity
    n0 = freefallNormal v
    n = rotateVector2D q n0
    --(n, _) = normalizeVector v
    p0 = p .+^ (n ^* (simulationBatLength / 2))
    p1 = p .-^ (n ^* (simulationBatLength / 2))
    {- p0 = p
    p1 = p .+^ (n ^* simulationBatLength) -}
    p0Dist = distToSpecialBase p0
    p1Dist = distToSpecialBase p1

    line = bool (ClosedLineSegment (p1 :+ ()) (p0 :+ ())) (ClosedLineSegment (p0 :+ ()) (p1 :+ ()))  (p0Dist < p1Dist)

    -- | Distance to Stig base at table height
    distToSpecialBase :: Point 2 Float -> Float
    distToSpecialBase po = norm $ po .-. Point2 stigFoot simulationTableHeight

-- | Maximum ITerations to guess
binaryGuessMaxIter :: Int
binaryGuessMaxIter = 10

-- | Place the ball at the center of the opponents table
rotateBatToCenter :: Point 2 Float -> Vector 2 Float -> DistanceToRange (Float, LineSegment 2 () Float)
rotateBatToCenter p v = fromMaybe (trace "Never interception with table?!" Below (read "Infinity" :: Float, placeBatInBounceCurve p v 0)) finalGuess
  where
    -- If we are going up
    goingUp = view yComponent v > 0

    finalGuess = -- binaryGuess 0 (-binaryGuessLimit) binaryGuessLimit 0
      bool
        (binaryGuess 0 0 (pi/2) 0) -- Case going down
        (binaryGuess 0 (-pi/4) 0 (-pi/4)) -- Case going up
        goingUp

    -- | Guess a possible bat possition
    binaryGuess iter qmin qmax qcurr
      | noTPossible = trace "No interception with table?!" Nothing
      {- | iter >= binaryGuessMaxIter = trace ("Max Iter Selected q=" ++ show qcurr) Just (useInsideInfo insideInfo bat)-}
      | iter >= binaryGuessMaxIter = Just (useInsideInfo insideInfo bat)
      | otherwise = guess
      where
        -- Generate bat
        bat = placeBatInBounceCurve p v qcurr
        nV = reflectVelocityBat bat v
        tPossible = predictFreefallHeightInter p nV simulationTableHeight
        noTPossible = isNothing tPossible
        t = fromJust tPossible
        (pI, _) = freeFallEvaluateTime p nV t
        pIX = view xCoord pI

        -- If Bounce is going up
        bounceGoingUp = view yComponent nV > 0

        insideInfo = insideOpponentTableX pIX

        useInsideInfo (Inside _) b = Inside (pIX, b)
        useInsideInfo (Above _) b = Below (pIX, b)
        useInsideInfo (Below _) b = Above (pIX, b)

        -- Guess Selection to improve
        -- Ball Going Up
        guessTry x True = case insideOpponentTableX x of
                  {- Above _ -> trace ("X=" ++ show x ++ " Up:Below " ++ show bounceGoingUp ++ " qs=" ++ show (qmin, qmax, qcurr)) $ binaryGuess (iter + 1) qmin qcurr ((qmin + qcurr) / 2)
                  Below _ -> trace ("X=" ++ show x ++ " Up:Above " ++ show bounceGoingUp ++ " qs=" ++ show (qmin, qmax, qcurr)) $ binaryGuess (iter + 1) qcurr qmax ((qmax + qcurr) / 2)
                  Inside _ -> trace ("X=" ++ show x ++ " Up:Inside Selected q=" ++ show qcurr) Just (Inside (pIX, bat)) -}
                  Above _ ->  binaryGuess (iter + 1) qmin qcurr ((qmin + qcurr) / 2)
                  Below _ ->  binaryGuess (iter + 1) qcurr qmax ((qmax + qcurr) / 2)
                  Inside _ -> Just (Inside (pIX, bat))

        -- Ball Going Down
        guessTry x False = case insideOpponentTableX x of
          {- Above _ -> trace ("X=" ++ show x ++ " Down:Below " ++ show bounceGoingUp ++ " qs=" ++ show (qmin, qmax, qcurr)) $ binaryGuess (iter + 1) qmin qcurr ((qmin + qcurr) / 2)
          Below _ -> trace ("X=" ++ show x ++ " Down:Above " ++ show bounceGoingUp ++ " qs=" ++ show (qmin, qmax, qcurr)) $ binaryGuess (iter + 1) qcurr qmax ((qmax + qcurr) / 2)
          Inside _ -> trace ("X=" ++ show x ++ " Down:Inside Selected q=" ++ show qcurr) Just (Inside (pIX, bat)) -}
          Above _ -> binaryGuess (iter + 1) qmin qcurr ((qmin + qcurr) / 2)
          Below _ -> binaryGuess (iter + 1) qcurr qmax ((qmax + qcurr) / 2)
          Inside _ -> Just (Inside (pIX, bat))

        --guess = traceShow (p, nV, t, pI) guessTry pIX goingUp
        guess = guessTry pIX goingUp

-- | Move bat to center low limit
moveBatToCenterLowLim :: Float
moveBatToCenterLowLim = 0.6

-- | Move bat to center high limit
moveBatToCenterHighLim :: Float
moveBatToCenterHighLim = 1.3

-- | Move bat to center number of steps
moveBatToCenterSteps :: Float
moveBatToCenterSteps = 40

-- | Move bat to center step size
moveBatToCenterStepSize :: Float
moveBatToCenterStepSize = (moveBatToCenterHighLim - moveBatToCenterLowLim) / moveBatToCenterSteps

-- | Maximum Time apply Motion
maxTimeToMotion :: Motion-> Float
maxTimeToMotion m = maximum $ map (\v -> abs v / simulationMaxSpeed) m

bestMotion :: Arm -> Point 2 Float -> Vector 2 Float -> Motion
bestMotion arm p v = finalMotion
  where
    (finalMotion, _) = move 0 moveBatToCenterHighLim
    -- Approach to good bat center
    move iter x
      | iter >= moveBatToCenterSteps || moveBatToCenterLowLim > x = (bM, bMV) -- Limit
      | isInside possible = bool (bMN, mVN) (bM, bMV) (bMV <= mVN) -- We are inside check if we are better
      | otherwise = (bMN, bMV) -- We are not inside, continue checking
      where
        tX = freeFallTimeToXPos p v x
        (pN, vN) = freeFallEvaluateTime p v tX
        possible = rotateBatToCenter pN vN
        (_, bat) = getDistanceToRangeContent possible

        batInv = segmentInvert bat
        (mIntercept, _, _) = fabrikToSegment stigFoot arm bat
        (mInterceptInv, _, _) = fabrikToSegment stigFoot arm batInv
        m = armToMotion arm mIntercept
        mV = maxTimeToMotion m
        mInv = armToMotion arm mInterceptInv
        mVInv = maxTimeToMotion mInv

        -- Current Best Motion
        bM = bool mInv m (mV <= mVInv)
        bMV = maxTimeToMotion bM

        -- Possible Next Best Motion
        (bMN, mVN) = move (iter + 1) (x - moveBatToCenterStepSize)

-- | Try to intercept Ball
tryInterceptBall :: Arm -> Point 2 Float -> Vector 2 Float -> Float -> IO Motion
tryInterceptBall arm p v tColl =
   do
     let bM = bestMotion arm p v
     --return $ trace ("Opponent did a proper hit we can catch at " ++ show tColl ++ "\n" ++ show bM) bM
     return bM

-- | Stig's player
stig :: Player
stig =
  Player
    {
      name = "Stig",
      arm = stigArm,
      initArm = stigArm,
      foot = stigFoot,
      prepare = return (),
      terminate = return (),
      action = stigAction,
      collide = stigCollide,
      planPnt = stigPlanPnt,
      planSeg = stigPlanSeg,
      stretch = \_ arm -> return $ armToStigRestMotion arm,
      dance   = \_ arm -> return $ armToStigRestMotion arm
    }

-- Internal Stig Arm
stigInternalArm :: Arm
stigInternalArm = checkArm
    [ Link paleBlue 0.5,
      Joint red (0.5112798), -- (0.1)
      Link paleBlue 0.3,
      Joint red 1.247675, -- (0.1)
      Link paleBlue 0.25,
      Joint red 0.4820231, -- (-0.1)
      Link paleBlue 0.1,
      Joint red (-2.2409778), -- (-0.1)
      Link hotPink 0.1 -- Bat
    ]

-- | Arm to use
stigArm :: Arm
stigArm = mapMotion stigInternalArm stigRest

-- | Stig Arm Length
stigArmLength :: Float
stigArmLength = armLength stigArm

-- | Separation from the center of the table
stigFoot :: Float
stigFoot = 1.3

-- | Stig rest postion
stigRest :: Motion
stigRest = m
  where
    (m, _, _) = fabrikToSegment stigFoot stigInternalArm (ClosedLineSegment (Point2 1.1 0.65 :+()) (Point2 1.1 0.75 :+()))

-- | Stig rest 2 postion
stigRest2 :: Motion
stigRest2 = m
  where
    (m, _, _) = fabrikToSegment stigFoot stigInternalArm (ClosedLineSegment (Point2 0.9 0.83 :+()) (Point2 0.9 0.93 :+()))

-- | Get the a zeroed Motion list for Stig's arm
stigNoMotion :: Motion
stigNoMotion = map f stigRest
  where
    f = const 0

-- | Calculate Motion Velocity to Rest Motion. !Warning: no limits are applied
armToStigRestMotion :: Arm -> Motion
armToStigRestMotion ar = armToMotion ar stigRest

-- | Calculate Motion Velocity to Rest Motion2. !Warning: no limits are applied
armToStigRestMotion2 :: Arm -> Motion
armToStigRestMotion2 ar = armToMotion ar stigRest2

-- | Check collision of moving line and point
stigCollide ::
  forall r.
  (Num r, Floating r, Ord r, Eq r, Show r) =>
  (r, Point 2 r, LineSegment 2 () r) ->
  (r, Point 2 r, LineSegment 2 () r) ->
  IO (Point 2 r)
stigCollide t1 t2 = bool (error "Stig Collide Failed a Test Case") (return (f t1 t2)) completeCheck
  where
    f = movingBallMovingLineCollide
    generateTestState :: (Num r, Floating r, Ord r) => r -> (r, r) -> (r, r) -> (r, r) -> (r, Point 2 r, LineSegment 2 () r)
    generateTestState t (px, py) (pl0x, pl0y) (pl1x, pl1y) = (t, Point2 px py, ClosedLineSegment (Point2 pl0x pl0y :+ ()) (Point2 pl1x pl1y :+ ()))

    testCases =
      [ (Point2 0 0, generateTestState 0 (0, 0) (1, -1) (1, 1), generateTestState 1 (0, 0) (1, -1) (1, 1)),
        (Point2 0 2, generateTestState 0 (0, 0) (0, 0) (1, 0), generateTestState 1 (0, 0) (0, 1) (1, 1)),
        (Point2 0 0, generateTestState 0 (0, 0) (1, -1) (1, 1), generateTestState 1 (2, 0) (1, -1) (1, 1)),
        (Point2 (-1) 0, generateTestState 0 (0, 0) (1, -1) (1, 1), generateTestState 1 (1, 0) (0, -1) (0, 1)),
        (Point2 1 1, generateTestState 0 (0, 0) (0, -1) (2, 1), generateTestState 1 (2, 0) (0, -1) (2, 1)),
        (Point2 (-1) 1, generateTestState 0 (0, 0) (0, -1) (2, 1), generateTestState 1 (0, 0) (-2, -1) (0, 1)),
        (Point2 1.2 0.8, generateTestState 0 (0.3, 1.0) (0.1, 2.1) (-0.5, 0.9), generateTestState 1 (1.2, 0.8) (-0.2, 2.2) (-0.3, 1.1))
        --(Point2 (-6) (-4), generateTestState 0 (-5, -5) (0, 1) (1, 0), generateTestState 1 (5, 7) (0, 1) (1, 0))
      ]

    checkCollision :: (Num r, Floating r, Ord r, Show r) => (Point 2 r, (r, Point 2 r, LineSegment 2 () r), (r, Point 2 r, LineSegment 2 () r)) -> (Bool, Diff (Point 2) r, Point 2 r)
    checkCollision (ans, s1, s2) = (diffX == 0 && diffY == 0, diff, c)
      where
        c = f s1 s2
        diff = c .-. ans
        diffX = globalThreshold 0 $ abs $ view xComponent diff
        diffY = globalThreshold 0 $ abs $ view yComponent diff

    performTest :: (Num r, Floating r, Ord r, Show r) => (Point 2 r, (r, Point 2 r, LineSegment 2 () r), (r, Point 2 r, LineSegment 2 () r)) -> Bool
    performTest testCase@(ans, s1, s2) = bool (error showError) True correct
      where
        (correct, diff, p) = checkCollision testCase
        showError =
          "Expected " ++ show ans ++ " but got "
            ++ show p
            ++ " with Diff "
            ++ show diff
            ++ " \n\tFor case:\n\t"
            ++ show s1
            ++ "\n\t->\n\t"
            ++ show s2

    completeCheck = all performTest testCases

--test = stigCollide (0, Point2 0 0, ClosedLineSegment (Point2 0 0 :+ ()) (Point2 1 0 :+ ())) (1, Point2 0 0, ClosedLineSegment (Point2 1 0 :+ ()) (Point2 1 0 :+ ()))
--test = stigCollide (0, Point2 (-1) 1, ClosedLineSegment (Point2 0 0 :+ ()) (Point2 1 1 :+ ())) (1, Point2 0 0, ClosedLineSegment (Point2 0 0 :+ ()) (Point2 (-1) 1 :+ ()))
--test = stigCollide (0, Point2 (-1) 1, ClosedLineSegment (Point2 0 (-1) :+ ()) (Point2 1 1 :+ ())) (1, Point2 0 0, ClosedLineSegment (Point2 0 (-1) :+ ()) (Point2 (-1) 1 :+ ()))

stigAction :: Float -> (Float, Item) -> BallState -> Arm -> IO Motion
stigAction _ (tColl, Net) _ arm =
  return $
    -- Ball hit the net, this means someone scored
    -- Go to rest
    let toBase = armToMotion arm stigNoMotion
     --in trace ("Someone Scored at " ++ show tColl) applyMotionLimits toBase -- Velocity limits
     in applyMotionLimits toBase
stigAction _ (tColl, Other _) _ arm =
  return $
    -- Ball hit something out of the game, this means someone scored
    -- Go to rest
    let toBase = armToMotion arm stigNoMotion
     --in trace ("Someone Scored at " ++ show tColl) applyMotionLimits toBase -- Velocity limits
     in applyMotionLimits toBase
stigAction _ (tColl, Bat Self) _ arm =
  return $
    -- We hit the ball, go to rest motion
    let toRest2 = armToStigRestMotion2 arm
     --in trace ("We just hit the ball at " ++ show tColl) toRest2 -- Velocity limits
     in toRest2
stigAction _ (tColl, Table Opponent) _ arm =
  return $
    -- Our hit was correct and we reached the other player's side
    -- So rest
    let toRest2 = armToStigRestMotion2 arm
     -- in trace ("We did a proper hit at " ++ show tColl) toRest2
     in toRest2
stigAction t (tColl, Air) bs arm =
  return $
    -- Invalid State
    let toBase = armToMotion arm stigNoMotion
     -- in trace ("Impossible State " ++ show tColl) applyMotionLimits toBase -- Velocity limits
     in applyMotionLimits toBase
stigAction t (tColl, Table Self) bs arm =
    do
      -- Other player did a proper hit we have to respond to
      -- Distance to max point for now
      let p = loc bs
      let v = dir bs
      tryInterceptBall arm p v tColl
stigAction t (tColl, Bat Opponent) bs arm =
  do

      -- Other player did a hit we have to respond to
      let p = loc bs
      let v = dir bs

      let mayBounce = predictTableBounce p v

      case mayBounce of
        {- Nothing -> return $ trace ("Opponent did a wrong hit at " ++ show tColl) (armToStigRestMotion arm) -- Velocity limits
        Just (pT, vT, _) -> trace ("Opponent did a proper hit at " ++ show tColl) tryInterceptBall arm pT vT tColl -}
        Nothing -> return $ armToStigRestMotion arm
        Just (pT, vT, _) -> tryInterceptBall arm pT vT tColl
stigAction _ (tColl, _) _ arm =
  return $
    -- Ball: Don't know what happened
    -- Go to rest
    let toBase = armToMotion arm stigNoMotion
     --in trace ("Someone Scored at " ++ show tColl) applyMotionLimits toBase -- Velocity limits
     in trace ("Don't know what happened at: " ++ show tColl) applyMotionLimits toBase

-- | Stig Plan Threshold
stigPlanThreshold :: (Num r, Ord r, Fractional r) => r -> r -> r
stigPlanThreshold = threshold 0.01

-- | Calculates the possible motion values to achieve a point. If it fails it returns []
stigPlanPnt :: Float -> Arm -> Point 2 Float -> IO Motion
stigPlanPnt foot arm p
  | isTooFar = trace "Too Far" return []
  | stigPlanThreshold 0 eB == 0 = return $ map normalizeAngle qB
  | otherwise = trace ("Error: " ++ show eB) return []
  where
    isTooFar = norm (p .-. Point2 foot 0) > 1.2 * armLength arm
    (qB, eB) = fabrikToPoint foot arm p

-- | Calculates the possible motion values to achieve a line segment. If it fails it returns []
stigPlanSeg :: Float -> Arm -> LineSegment 2 () Float -> IO Motion
stigPlanSeg foot arm s
  | stigPlanThreshold 0 eB == 0 && stigPlanThreshold 0 eBat == 0 = return m
  | otherwise = return []
  where
    (m, eB, eBat) = fabrikToSegment foot arm s

-- Testing Starts Here
-- | Create a Test Case for stigPlanPnt
createPlanPntCase :: Float -> Arm -> (Float, Float) -> Motion -> (Float, Arm, Point 2 Float, Motion)
createPlanPntCase f a (xT, yT) m = (f, a, Point2 xT yT, m)

-- | Test stigPlanPnt
testPlanPnt :: (Float, Arm, Point 2 Float, Motion) -> IO Bool
testPlanPnt (f, arm, pT, mT)
  = do
      -- Expected Position
      let qT = motionToJointVector mT
      let xTargetGlobal = pointToHomogenousPoint pT

      -- Arm
      let (a, _) = getArmKinematicAndMotion f arm

      -- Calculate Answer
      mB <- stigPlanPnt f arm pT
      if null mB then
        return $ bool (trace "Wrong Null" False) (null mT) (null mT)
      else
        do
          let qB = motionToJointVector mB

          -- Forward Transforms
          let fwdTT = applyForwardKinematicTrans a qT
          let fwdTB = applyForwardKinematicTrans a qB

          -- Bat Global Position
          let batGlobalB = applyForwardKinematicMatrixTrans fwdTB Numerical.#> homogeneousZero
          let batGlobalT = applyForwardKinematicMatrixTrans fwdTT Numerical.#> homogeneousZero

          -- Error
          let eB = xTargetGlobal - batGlobalB
          let eNormB = Numerical.norm_2 eB

          let eT = trace ("Best " ++ show (batGlobalB, eNormB, mB)) $ xTargetGlobal - batGlobalT
          let eNormT = Numerical.norm_2 eT

          let result = trace ("Target " ++ show (batGlobalT, eNormT, mT) ++ "\n") $ (globalThreshold 0 eNormB == 0) && (eNormB <= eNormT)
          return $ result && (length mB == length mT)

-- | Test Cases for testPlanPnt
planPntTestCases :: [(Float, Arm, Point 2 Float, Motion)]
planPntTestCases =
  [ createPlanPntCase
      0
      [
        Joint red 0.0,
        Link red 0.1
      ]
      (0, 0.1)
      [0],
    createPlanPntCase
      0
      [
        Joint red 0.0,
        Link red 0.1
      ]
      (0.1, 0)
      [pi/2],
    createPlanPntCase
      1.5
      [ Link red 0.2,
        Joint red 0.0,
        Link red 0.2,
        Joint red 0.0,
        Link red 0.2,
        Joint red 0.0,
        Link red 0.2,
        Joint red 0.0,
        Link red 0.1
      ]
      (1.22385, 0.80917)
      [0.1, 0.2, 0.3, 0.4],
    createPlanPntCase
      1.5
      [ Link red 0.2,
        Joint red 0.0,
        Link red 0.2,
        Joint red 0.0,
        Link red 0.2,
        Joint red 0.0,
        Link red 0.2,
        Joint red 0.0,
        Link red 0.1
      ]
      (1.77615, 0.80917)
      [-0.5, 0.0, 0.4, 0.6],
    createPlanPntCase
      1.5
      [ Link red 0.2,
        Joint red 0.0,
        Link red 0.2,
        Joint red 0.0,
        Link red 0.2,
        Joint red 0.0,
        Link red 0.2,
        Joint red 0.0,
        Link red 0.1
      ]
      (1.48013, 0.89202)
      [0.1, -0.2, 0.3, -0.4],
      createPlanPntCase
      1.5
      [ Link red 0.1,
        Joint red 0.0,
        Link red 0.1,
        Joint red 0.0,
        Link red 0.1
      ]
      (1.5, 0)
      [2.094,2.094],
      createPlanPntCase
      1.5
      [ Link red 0.1,
        Joint red 0.0,
        Link red 0.1,
        Joint red 0.0,
        Link red 0.1
      ]
      (1.5, 0)
      [2.094,2.094]
  ]

-- | Executes testPlanPnt
executePlanPntTestCases :: IO Bool
executePlanPntTestCases = foldM f True planPntTestCases
  where
    f True params = testPlanPnt params
    f False _ = return False

-- | Create a Test Case for stigPlanSeg
createPlanSegCase :: Float -> Arm -> (Float, Float) -> (Float, Float) -> Motion -> (Float, Arm, LineSegment 2 () Float, Motion)
createPlanSegCase f a (xP0, yP0) (xP1, yP1) m = (f, a, ClosedLineSegment (Point2 xP0 yP0 :+ ()) (Point2 xP1 yP1 :+ ()), m)

-- | Test stigPlanSeg
testPlanSeg :: (Float, Arm, LineSegment 2 () Float, Motion) -> IO Bool
testPlanSeg (f, arm, sT, mT)
  =
    do
      -- Expected Position
      let pT = sT ^. (end . core)
      let qT = motionToJointVector mT
      let xTargetGlobal = pointToHomogenousPoint pT

      -- Arm
      let (a, _) = getArmKinematicAndMotion f arm

      -- Calculate Answer
      mB <- stigPlanSeg f arm sT
      if null mB then
        return $ bool (trace "Wrong Null" False) (null mT) (null mT)
      else
        do
          let qB = motionToJointVector mB

          -- Forward Transforms
          let fwdTT = applyForwardKinematicTrans a qT
          let fwdTB = applyForwardKinematicTrans a qB

          -- Bat Global Position
          let batGlobalB = applyForwardKinematicMatrixTrans fwdTB Numerical.#> homogeneousZero
          let batGlobalT = applyForwardKinematicMatrixTrans fwdTT Numerical.#> homogeneousZero

          -- Error
          let eB = xTargetGlobal - batGlobalB
          let eNormB = Numerical.norm_2 eB

          let eT = trace ("Best " ++ show (batGlobalB, eNormB, mB)) $ xTargetGlobal - batGlobalT
          let eNormT = Numerical.norm_2 eT

          let result = trace ("Target " ++ show (batGlobalT, eNormT, mT)) $
                (globalThreshold 0 eNormB == 0)
                  && ((eNormB <= eNormT) || (eNormT > 0 && eNormB / eNormT <= 1.1) || (eNormT == 0 && globalThreshold 0 eNormB == 0))
          return $ result && bool (trace "Different Motion Sizes" False) True (length mB == length mT)

-- | Test Cases for testPlanPnt
planSegTestCases :: [(Float, Arm, LineSegment 2 () Float, Motion)]
planSegTestCases =
  [ createPlanSegCase
      1.5
      [ Link red 0.2,
        Joint red 0.0,
        Link red 0.2,
        Joint red 0.0,
        Link red 0.2,
        Joint red 0.0,
        Link red 0.2,
        Joint red 0.0,
        Link red 0.1
      ]
      (1.48023, 0.79501)
      (1.48023, 0.89501)
      [0.2, -0.2, -0.1, 0.1],
    createPlanSegCase
      1.5
      [ Link red 0.2,
        Joint red 0.0,
        Link red 0.2,
        Joint red 0.0,
        Link red 0.2,
        Joint red 0.0,
        Link red 0.2,
        Joint red 0.0,
        Link red 0.1
      ]
      (1.71174, 0.75003)
      (1.66379, 0.83779)
      [-0.5, 0.0, 0.4, 0.6],
    createPlanSegCase
      1.5
      [ Link red 0.1,
        Joint red 0.0,
        Link red 0.1,
        Joint red 0.0,
        Link red 0.1
      ]
      (1.4134, 5.0e-2)
      (1.5, 0)
      [2.094,2.094]
  ]

-- | Executes testPlanSeg
executePlanSegTestCases :: IO Bool
executePlanSegTestCases = foldM f True planSegTestCases
  where
    f True params = testPlanSeg params
    f False _ = return False


{-# LANGUAGE RankNTypes #-}

module Observable.MCMC.AffineInvariantEnsemble where

import Control.Monad.Primitive
import Control.Monad.State.Strict
import Data.Vector.Unboxed (Vector)
import qualified Data.Vector.Unboxed as V
import Observable.Core
import Statistics.Distribution
import Statistics.Distribution.Normal

type Ensemble = Vector [Double]

-- | Generate a random value from a distribution having the property that
--   g(1/z) = z g(z).
symmetricVariate :: PrimMonad m => Observable m Double
symmetricVariate = do
  z <- unit
  return $ 0.5 * (z + 1) ^ 2

perturbParticle w0 w1 = do
  zs <- replicateM (length w0) symmetricVariate

V.zipWith (+) (V.map (* z) w0) (V.map (* (1 - z)) w1)

-- acceptanceRatio
--   :: [Double] -> [Double]
--   -> Double -> Double
--   -> ([Double] -> Double)
--   -> ([Double], Int)
-- acceptanceRatio target w0 w1 z zc =
--   let val      = target proposal - target w0 + (fromIntegral (length w0) - 1) * log z
--       proposal = zipWith (+) (map (*z) w0) (map (*(1-z)) w1)
--   in  if zc <= min 1 (exp val) then (proposal, 1) else (w0, 0)
--
-- -- | Calculate the acceptance ratio for a proposed move.
-- acceptRejectRatio
--   :: Target Double -> Double -> Vector Double -> Vector Double -> Double
-- acceptRejectRatio target e current proposed = exp . min 0 $
--     logObjective target proposed + log (isoGauss current proposed e)
--   - logObjective target current  - log (isoGauss proposed current e)


{-# LANGUAGE RankNTypes #-}

module Sim
  ( State (..),
    Plant (..),
    Measurement (..),
    simulate,
  )
where

import Numeric.LinearAlgebra

type Measurement = Vector Double

type State = Vector Double

data Plant = Plant {fA, fC, fB, fD :: forall a. Floating a => ([a] -> [a])}

simulate :: Plant -> State -> Int -> [(State, Measurement)]
simulate (Plant fA fB fC fD) x n = (x', y') : simulate (Plant fA fB fC fD) xh (n + 1)
  where
    x' = vector $ fA (toList x)
    y = vector $ fC (toList x')
    xh = x' + scale 0.0001 (randomVector (n + 1000) Gaussian (size x))
    y' = y + scale 0.5 (randomVector n Gaussian (size y))


{-# LANGUAGE DataKinds #-}
{-# LANGUAGE DeriveFoldable #-}
{-# LANGUAGE DeriveFunctor #-}
{-# LANGUAGE DeriveTraversable #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE KindSignatures #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE UndecidableInstances #-}

module PLC1 (PLC1(..)) where

import Control.Applicative
import Control.Comonad
import Control.Comonad.Store
import Control.Exception (assert)
-- import Control.Monad
-- import Data.Distributive
import Data.Proxy
import Data.Unfoldable
import Data.Validity
import qualified Data.Vector as V
import Data.VectorSpace
import GHC.TypeLits
-- import Linear.Metric
-- import Linear.Vector
import qualified Numeric.LinearAlgebra as L
import qualified Test.QuickCheck as QC
import qualified Test.QuickCheck.Function as QC

import ComonadicVector
import Function
import Integration

-- https://hackage.haskell.org/package/polynomial-0.7.3

newtype PLC1 (n' :: Nat) c s = PLC1
  { getPLC1 :: ComonadicVector V.Vector s
  } deriving (Eq, Ord, Read, Show, Foldable, Functor, Traversable)

instance (KnownNat n', Validity s) => Validity (PLC1 n' c s) where
    validate (PLC1 xs) =
        mconcat
        [ n > 0 <?@> "length is positive"
        , length xs == n <?@> "vector has correct length "
        , xs <?!> "vector"
        ]
        where n = fromInteger (natVal (Proxy :: Proxy n'))
    isValid = isValidByValidating

instance KnownNat n' => Unfoldable (PLC1 n' c) where
    -- unfold xs = PLC1 <$> unfold xs
    -- unfold xs = (PLC1 . ComonadicVector 0 . V.fromListN n) <$> mklist n
    --     where n = fromInteger (natVal (Proxy :: Proxy n')) :: Int
    --           mklist i | i == 0 = pure []
    --                    | otherwise = (:) <$> xs <*> mklist (i-1)
    unfold xs = PLC1 . goodlength <$> unfold xs
        where n = fromInteger (natVal (Proxy :: Proxy n')) :: Int
              goodlength ys = assert (n > 0 && length ys == n) ys

instance (Applicative (ComonadicVector V.Vector), KnownNat n') =>
         Applicative (PLC1 n' c) where
  -- pure x = PLC1 (pure x)
  -- Note: 'pure' should use the neutral element of the monoid used
  -- for 'pos_' in '(<*>)'
  pure x = PLC1 (ComonadicVector (n - 1) (V.replicate n x))
      where n = fromInteger (natVal (Proxy :: Proxy n'))
  PLC1 fs <*> PLC1 xs = PLC1 (fs <*> xs)

-- instance (Alternative (ComonadicVector V.Vector), KnownNat n') =>
--          Alternative (PLC1 n' c) where
--   empty = PLC1 empty
--   PLC1 xs <|> PLC1 ys = PLC1 (xs <|> ys)

-- instance Monad (ComonadicVector V.Vector) => Monad (PLC1 n' c) where
--     -- (>>=) :: PLC1 n' c s -> (s -> PLC1 n' c t) -> PLC1 n' c t
--     PLC1 x >>= f = PLC1 (x >>= getPLC1 . f)
--
-- instance MonadPlus (ComonadicVector V.Vector) => MonadPlus (PLC1 n' c) where
--     mzero = PLC1 mzero
--     PLC1 xs `mplus` PLC1 ys = PLC1 (xs `mplus` ys)

-- instance Distributive (ComonadicVector V.Vector) =>
--          Distributive (PLC1 n' c) where
--   collect f arr = PLC1 (collect (getPLC1 . f) arr)

instance Comonad (ComonadicVector V.Vector) => Comonad (PLC1 n' c) where
  extract (PLC1 arr) = extract arr
  extend f (PLC1 arr) = PLC1 (extend (f . PLC1) arr)

instance ComonadStore Int (ComonadicVector V.Vector) =>
         ComonadStore Int (PLC1 n' c) where
  pos (PLC1 arr) = pos arr
  peek i (PLC1 arr) = peek i arr

instance (KnownNat n', QC.Arbitrary s) => QC.Arbitrary (PLC1 n' c s) where
  -- arbitrary = do
  --   arr <- QC.arbitrary
  --   return (PLC1 arr)
  -- shrink (PLC1 arr) = [PLC1 arr' | arr' <- QC.shrink arr]
  -- arbitrary = arbitraryDefault
  arbitrary = do
      xs <- V.generateM n (const QC.arbitrary)
      i <- QC.choose (0, V.length xs - 1)
      return (PLC1 (ComonadicVector i xs))
      where n = fromInteger (natVal (Proxy :: Proxy n'))
  shrink arr = [arr' | pos arr > 0] ++ traverse QC.shrink arr'
      where arr' = seek 0 arr

instance QC.CoArbitrary s => QC.CoArbitrary (PLC1 n' c s) where
    coarbitrary = QC.coarbitrary . getPLC1

instance QC.Function s => QC.Function (PLC1 n' c s) where
    function = QC.functionMap getPLC1 PLC1

-- instance Additive (ComonadicVector V.Vector) => Additive (PLC1 n' c) where
--     zero = PLC1 zero
--
-- instance (Additive (PLC1 n' c), Bounded c, RealFrac c) => Metric (PLC1 n' c) where
--     -- dot :: Num a => f a -> f a -> a
--     dot xs ys =
--         assert (length ys == n) $
--         case n of
--           0 -> 0
--           _ -> sum [w i * peek i xs * peek i ys | i <- [0 .. n-1]]
--         where n = length xs
--               xmin = minBound :: c
--               xmax = maxBound :: c
--               w i = realToFrac (weight n xmin xmax i)

instance ( KnownNat n'
         , AdditiveGroup (ComonadicVector V.Vector s)
         , AdditiveGroup s
         ) =>
         AdditiveGroup (PLC1 n' c s) where
  zeroV = pure zeroV
  negateV = fmap negateV
  (^+^) = liftA2 (^+^)

instance ( KnownNat n'
         , VectorSpace (ComonadicVector V.Vector s)
         , VectorSpace s
         ) =>
         VectorSpace (PLC1 n' c s) where
  type Scalar (PLC1 n' c s) = Scalar s
  (*^) x = fmap (x *^)

-- instance ( KnownNat n'
--          , Bounded c
--          , RealFrac c
--          , InnerSpace s
--          , VectorSpace (Scalar s)
--          , Fractional (Scalar (Scalar s))
--          ) =>
--          InnerSpace (PLC1 n' c s) where
--   -- xs <.> ys =
--   --   assert (length ys == length xs) $ sumV (weighted ((<.>) <$> xs <*> ys))

instance ( KnownNat n'
         , Integrable (PLC1 n')
         , IntegrableOk (PLC1 n') c s
         , Bounded c
         , RealFrac c
         , InnerSpace s
         , Fractional (Scalar s)
         ) => InnerSpace (PLC1 n' c s) where
  xs <.> ys = getPLC1 (extend weighted xs) <.> getPLC1 ys
    where
      weighted zs = weight n (xmin, xmax) (pos zs) *^ extract zs
      -- n = length xs
      n = fromInteger (natVal (Proxy :: Proxy n'))
      (xmin, xmax) = bounds xs

instance Function (PLC1 n') where
  type FunctionOk (PLC1 n') a b = ( KnownNat n'
                                  , Bounded a
                                  , RealFrac a
                                  , InnerSpace b
                                  , Floating (Scalar b)
                                  , Ord (Scalar b)
                                  , b ~ Double) -- TODO
  eval (PLC1 arr) x =
    case n of
      0 -> zeroV
      1 -> extract arr
      _ ->
        let y = ((x - cx) / rx + 1) / 2 * fromIntegral (n - 1)
            i = max 0 $ min (n - 2) $ floor y
            f0 = 1 - f1
            f1 = realToFrac (y - fromIntegral i)
        in f0 *^ peek i arr ^+^ f1 *^ peek (i + 1) arr
    where
      -- n = length arr
      n = fromInteger (natVal (Proxy :: Proxy n'))
      xmin = minBound `asTypeOf` x
      xmax = maxBound `asTypeOf` x
      cx = (xmax + xmin) / 2
      rx = (xmax - xmin) / 2

instance KnownNat n' => Discretization (PLC1 n') where
  discretized f = PLC1 (ComonadicVector 0 res)
    where
      n = fromInteger (natVal (Proxy :: Proxy n'))
      -- res = V.generate n go
      -- TODO: use (symmetric) triangular solver
      rhs = L.vector [go i | i <- [0..n-1]]
      mat = L.matrix n [ overlap n (xmin, xmax) i j
                       | i <- [0..n-1], j <- [0..n-1]]
      res = V.fromListN n $ L.toList $ mat L.<\> rhs
      (xmin, xmax) = (minBound, maxBound)
      _ = (f xmin, f xmax)      -- constrain types
      go i = integrate bf lo hi
        where
          bf x = b x *^ f x
          b x = basis n (xmin, xmax) i x
          (lo, hi) = support n (xmin, xmax) i

instance KnownNat n' => Integrable (PLC1 n') where
  type IntegrableOk (PLC1 n') a b = ()
  integral xs = sumV (extend weighted xs)
    where
      weighted ys = weight n (xmin, xmax) (pos ys) *^ extract ys
      -- n = length xs
      n = fromInteger (natVal (Proxy :: Proxy n'))
      (xmin, xmax) = bounds xs

instance Differentiable (PLC1 n') where
  type DifferentiableOk (PLC1 n') a b = ()
  type Dir (PLC1 n') = ()
  boundary () xs = extend bndry xs
    where
      bndry ys
        | n == 1     = zeroV *^ extract ys -- or zeroV
        | i == 0     = negateV (extract ys ^/ realToFrac dx)
        | i == n - 1 = extract ys ^/ realToFrac dx
        | otherwise  = zeroV *^ extract ys -- or zeroV
        where
          i = pos ys
      -- n = length xs
      n = fromInteger (natVal (Proxy :: Proxy n'))
      (xmin, xmax) = bounds xs
      rx = (xmax - xmin) / 2
      hdx = rx / fromIntegral (n - 1)
      dx = 2 * hdx
  derivative () xs = extend deriv xs
    where
      deriv ys
        | n == 1     = zeroV *^ extract ys -- or zeroV
        | i == 0     = (peek (i + 1) ys ^-^ extract ys) ^/ realToFrac dx
        | i == n - 1 = (extract ys ^-^ peek (i - 1) ys) ^/ realToFrac dx
        | otherwise  =
            (peek (i + 1) ys ^-^ peek (i - 1) ys) ^/ realToFrac (2 * dx)
        where
          i = pos ys
      -- n = length xs
      n = fromInteger (natVal (Proxy :: Proxy n'))
      (xmin, xmax) = bounds xs
      rx = (xmax - xmin) / 2
      hdx = rx / fromIntegral (n - 1)
      dx = 2 * hdx

triangle :: (Num a, Ord a) => a -> a
triangle x = max 0 (1 - abs x)

typed :: f a b -> a -> a
typed _ = id

bounds :: Bounded a => f a b -> (a, a)
bounds xs = (typed xs minBound, typed xs maxBound)

-- The basis functions are defined to have a maximum of '1'
basis :: (RealFrac a, Fractional b) => Int -> (a, a) -> Int -> a -> b
basis n (xmin, xmax) i x | n == 1    = 1
                         | otherwise = realToFrac (triangle y)
  where
    cx  = (xmin + xmax) / 2
    rx  = (xmax - xmin) / 2
    hdx = rx / fromIntegral (n - 1)
    xi  = cx + hdx * fromIntegral (2 * i + 1 - n)
    dx  = 2 * hdx
    y   = (x - xi) / dx

support :: (Fractional a, Ord a) => Int -> (a, a) -> Int -> (a, a)
support n (xmin, xmax) i | n == 1    = (xmin, xmax)
                         | otherwise = (max xmin (xi - dx), min xmax (xi + dx))
  where
    cx  = (xmin + xmax) / 2
    rx  = (xmax - xmin) / 2
    hdx = rx / fromIntegral (n - 1)
    xi  = cx + hdx * fromIntegral (2 * i + 1 - n)
    dx  = 2 * hdx

-- Volume
weight :: (RealFrac a, Fractional b) => Int -> (a, a) -> Int -> b
weight n (xmin, xmax) i | n == 1     = realToFrac (2 * rx)
                        | i == 0     = realToFrac hdx
                        | i == n - 1 = realToFrac hdx
                        | otherwise  = realToFrac (2 * hdx)
  where
    rx  = (xmax - xmin) / 2
    hdx = rx / fromIntegral (n - 1)

overlap :: (RealFrac a, Fractional b) => Int -> (a, a) -> Int -> Int -> b
overlap n (xmin, xmax) i j
    | n == 1                           = realToFrac (2 * rx)
    | i == j && (i == 0 || i == n - 1) = realToFrac (2 / 3 * hdx)
    | i == j                           = realToFrac (4 / 3 * hdx)
    | abs (i - j) == 1                 = realToFrac (1 / 3 * hdx)
    | otherwise                        = 0
  where
    rx  = (xmax - xmin) / 2
    hdx = rx / fromIntegral (n - 1)


{-# LANGUAGE BangPatterns, FlexibleContexts, UnboxedTuples #-}
-- |
-- Module    : Statistics.Sample.KernelDensity
-- Copyright : (c) 2011 Bryan O'Sullivan
-- License   : BSD3
--
-- Maintainer  : bos@serpentine.com
-- Stability   : experimental
-- Portability : portable
--
-- Kernel density estimation.  This module provides a fast, robust,
-- non-parametric way to estimate the probability density function of
-- a sample.
--
-- This estimator does not use the commonly employed \"Gaussian rule
-- of thumb\".  As a result, it outperforms many plug-in methods on
-- multimodal samples with widely separated modes.

module Statistics.Sample.KernelDensity
    (
    -- * Estimation functions
      kde
    -- , kde_
    -- * References
    -- $references
    ) where

import Numeric.MathFunctions.Constants (m_sqrt_2_pi)
import Prelude hiding (const, min, max, sum)
import Statistics.Function (minMax, nextHighestPowerOfTwo)
import Statistics.Math.RootFinding (fromRoot, ridders)
import Statistics.Sample.Histogram (histogram_)
import Statistics.Sample.Internal (sum)
import Statistics.Transform (CD, dct, idct)
import qualified Data.Vector.Generic  as G
import qualified Data.Vector.Unboxed  as U
import qualified Data.Vector          as V

-- | Gaussian kernel density estimator for one-dimensional data, using
-- the method of Botev et al.
--
-- The result is a pair of vectors, containing:
--
-- * The coordinates of each mesh point.  The mesh interval is chosen
--   to be 20% larger than the range of the sample.  (To specify the
--   mesh interval, use 'kde_'.)
--
-- * Density estimates at each mesh point.
kde :: (G.Vector v CD, G.Vector v Double, G.Vector v Int)
    => Int
    -- ^ The number of mesh points to use in the uniform discretization
    -- of the interval @(min,max)@.  If this value is not a power of
    -- two, then it is rounded up to the next power of two.
    -> v Double -> (v Double, v Double)
kde n0 xs = kde_ n0 (lo - range / 10) (hi + range / 10) xs
  where
    (lo,hi) = minMax xs
    range   | G.length xs <= 1 = 1       -- Unreasonable guess
            | lo == hi         = 1       -- All elements are equal
            | otherwise        = hi - lo
{-# INLINABLE  kde #-}
{-# SPECIAlIZE kde :: Int -> U.Vector Double -> (U.Vector Double, U.Vector Double) #-}
{-# SPECIAlIZE kde :: Int -> V.Vector Double -> (V.Vector Double, V.Vector Double) #-}

-- | Gaussian kernel density estimator for one-dimensional data, using
-- the method of Botev et al.
--
-- The result is a pair of vectors, containing:
--
-- * The coordinates of each mesh point.
--
-- * Density estimates at each mesh point.
kde_ :: (G.Vector v CD, G.Vector v Double, G.Vector v Int)
     => Int
     -- ^ The number of mesh points to use in the uniform discretization
     -- of the interval @(min,max)@.  If this value is not a power of
     -- two, then it is rounded up to the next power of two.
     -> Double
     -- ^ Lower bound (@min@) of the mesh range.
     -> Double
     -- ^ Upper bound (@max@) of the mesh range.
     -> v Double
     -> (v Double, v Double)
kde_ n0 min max xs
  | G.null xs = error "Statistics.KernelDensity.kde: empty sample"
  | n0 <= 1   = error "Statistics.KernelDensity.kde: invalid number of points"
  | otherwise = (mesh, density)
  where
    mesh = G.generate ni $ \z -> min + (d * fromIntegral z)
        where d = r / (n-1)
    density = G.map (/(2 * r)) . idct $ G.zipWith f a (G.enumFromTo 0 (n-1))
      where f b z = b * exp (sqr z * sqr pi * t_star * (-0.5))
    !n  = fromIntegral ni
    !ni = nextHighestPowerOfTwo n0
    !r  = max - min
    a   = dct . G.map (/ sum h) $ h
        where h = G.map (/ len) $ histogram_ ni min max xs
    !len    = fromIntegral (G.length xs)
    !t_star = fromRoot (0.28 * len ** (-0.4)) . ridders 1e-14 (0,0.1) $ \x ->
              x - (len * (2 * sqrt pi) * go 6 (f 7 x)) ** (-0.4)
      where
        f q t = 2 * pi ** (q*2) * sum (G.zipWith g iv a2v)
          where g i a2 = i ** q * a2 * exp ((-i) * sqr pi * t)
                a2v = G.map (sqr . (*0.5)) $ G.tail a
                iv = G.map sqr $ G.enumFromTo 1 (n-1)
        go s !h | s == 1    = h
                | otherwise = go (s-1) (f s time)
          where time  = (2 * const * k0 / len / h) ** (2 / (3 + 2 * s))
                const = (1 + 0.5 ** (s+0.5)) / 3
                k0    = U.product (G.enumFromThenTo 1 3 (2*s-1)) / m_sqrt_2_pi
    sqr x = x * x
{-# INLINABLE  kde_ #-}
{-# SPECIAlIZE kde_ :: Int -> Double -> Double -> U.Vector Double -> (U.Vector Double, U.Vector Double) #-}
{-# SPECIAlIZE kde_ :: Int -> Double -> Double -> V.Vector Double -> (V.Vector Double, V.Vector Double) #-}

-- $references
--
-- Botev. Z.I., Grotowski J.F., Kroese D.P. (2010). Kernel density
-- estimation via diffusion. /Annals of Statistics/
-- 38(5):2916–2957. <http://arxiv.org/pdf/1011.2602>


{-# LANGUAGE MultiWayIf #-}
module Test.BigOh.Fit.Hakaru
  ( -- * Complexity defs
    Coefficient
  , Order(..)
  , Fit(..)
  , knownOrders
  , exponential
  , constant
  , linear
  , quadratic
  , cubic
  , nlogn
  , poly1d

  -- * Fitting curves
  , bestFit
  , fit
  ) where

import           Control.Arrow
import           Control.Monad
import           Data.Dynamic
import           Data.List
import           Data.Monoid
import           Data.Ord
import qualified Data.Vector                  as V
import           Language.Hakaru.Distribution
import           Language.Hakaru.Metropolis
import           Language.Hakaru.Types
import qualified Statistics.Sample            as S

import           Test.BigOh.Fit.Base

data Fit
  = Fit { xVals :: [Double]
        , yVals :: [Double]
        , burnN :: Int
        , takeN :: Int
        }

-- | Given a complexity order and some data,
--   determine if the order best describes the data.
fit :: (Order -> Bool) -> [Order] -> [Double] -> [Double] -> IO Bool
fit predi orders xs ys
  = go (50 :: Int) burnStart takeStart
  where
   burnStart = 10000
   takeStart = 1000
   go try b t
    = do ranked    <- bestFit' orders (Fit xs ys b t)
         putStrLn ("ranked: " <> show (fmap (first name) ranked))
         let order' = fst $ head ranked
         if | predi order'              -> return True
            | try        == 0           -> return False
            | otherwise                 -> go (try - 1) (b * 2) (t * 2)

-- | Find the complexity order that best fits the data.
--
bestFit :: [Order] -> Fit ->  IO Order
bestFit orders f
  = fst . head <$> bestFit' orders f

bestFit' :: [Order] -> Fit ->  IO [(Order, Double)]
bestFit' orders f@(Fit xs ys _ _)
  = do fits   <- mapM (`curveFit` f) orders
       let fs  = fmap (<$> xs) fits
           rs  = fmap (rSquared ys) fs
           rs' = zip orders rs
           rs''= sortBy (comparing snd) rs'
       return rs''

-- | Given a complexity order and some data, generate a curve of
--   that order that fit the data.
--
curveFit :: Order -> Fit -> IO (Double -> Double)
curveFit thing (Fit xs ys dropn taken)
 = do l <- mcmc (measureForOrder thing xs ys)
                (map (Just . toDyn . Lebesgue) ys)
      let means = expectations $ take taken $ drop dropn l
      return $ mkCurve thing means

-- | Create a sampler for a class of curves with some x values.
--   e.g. sample @y = a*x^2 + b*x +c@
--
measureForOrder :: Order -> [Double] -> [Double] -> Measure [Double]
measureForOrder order xs ys
  = measureForOrder' 0 (maximum $ fmap abs ys) (sd ys) order xs

measureForOrder' :: Double -> Double -> Double -> Order -> [Double] -> Measure [Double]
measureForOrder' mean range sdev (Order _ n func) xs
  = do w <- replicateM n $ unconditioned (normal mean range)
       y <- mapM (conditioned . withinNormal w) xs
       return w
  where
   withinNormal w x
     = normal (func w x) sdev

-- | Given a bunch of possible coefficient sets, return the
--   expected value of each coeffient.
--   e.g. for @a*x^2 + b^x + c@,some possible coffients might be:
--        @[[a=0,b=1,c=4], [a=3,b=4,c=2]]@, @expectations@ returns
--        the expected values for @a, b, c@.
--
expectations :: [[Double]] -> [Double]
expectations l = map (S.mean . V.fromList) (transpose l)

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

type Coefficient = Double

-- | A complexity order, e.g. exponential, quadratic.
data Order
  = Order
    { name      :: String
    , numCoeffs :: Int
    , mkCurve   :: [Coefficient] -> Double -> Double
    }

knownOrders :: [Order]
knownOrders = [exponential, constant, linear, quadratic, cubic, nlogn]

exponential :: Order
exponential
  = Order "exp" 4 $ \[a, b, c, d] x -> a * (2 ** (b * x + c)) + d

constant :: Order
constant
  = Order "constant" 1 poly1d

linear :: Order
linear
  = Order "linear" 2 poly1d

quadratic :: Order
quadratic
  = Order "quadratic" 3 poly1d

cubic :: Order
cubic
  = Order "cubic" 4 poly1d

nlogn :: Order
nlogn
  = Order "nlogn" 2
  $ \[a, b] n -> a * n * log  n + b

poly1d :: [Double] -> Double -> Double
poly1d weights a = poly weights a 1
   where
     poly [] _ _ = 0
     poly (w:ws) x acc = w*acc + (poly ws x acc*x)

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

rSquared
  :: [Double] -- ^ data set y1..yn
  -> [Double] -- ^ model f1..fn
  -> Double   -- ^ r squared
rSquared ys fs
 = let yBar  = sum ys / fromIntegral (length ys)
       ssTot = sum (fmap (square . subtract yBar) ys)
       ssRes = sum (fmap square (zipWith (-) ys fs))
   in  (ssTot - ssRes) / ssTot


{-# LANGUAGE PartialTypeSignatures #-}
module PMMHTest where

import           Control.Monad.Bayes.Class
import           Control.Monad.Bayes.Sampler
import           Control.Monad
import           Control.Monad.Bayes.Inference.PMMH
                                               as PMMH
import           Control.Monad.Bayes.Weighted
import           Control.Monad.State
import           Numeric.Log
import           DataParser
import           Statistics
import           Utils

generateData
  :: MonadSample m
  =>
  -- | T
     Int
  ->
  -- | list of latent and observable states from t=1
     m [(Double, Double)]
generateData t = do
  (sigmaX, sigmaY) <- param
  let sq x = x * x
      simulate 0 _ acc = return acc
      simulate k x acc = do
        let n = length acc
        x' <- normal (mean x n) sigmaX
        y' <- normal (sq x' / 20) sigmaY
        simulate (k - 1) x' ((x', y') : acc)
  x0  <- normal 0 (sqrt 5)
  xys <- simulate t x0 []
  return $ reverse xys

param :: MonadSample m => m (Double, Double)
param = do
  let a = 0.01
  let b = 0.01
  precX <- gamma a b
  let sigmaX = 1 / sqrt precX
  precY <- gamma a b
  let sigmaY = 1 / sqrt precY
  return (sigmaX, sigmaY)

mean :: Double -> Int -> Double
mean x n =
  let sq x = x * x
  in  0.5 * x + 25 * x / (1 + sq x) + 8 * cos (1.2 * fromIntegral n)

model
  :: (MonadInfer m)
  =>
  -- | observed data
     [Double]
  ->
  -- | prior on the parameters
     (Double, Double)
  ->
  -- | list of latent states from t=1
     m [Double]
model obs (sigmaX, sigmaY) = do
  let sq x = x * x
      simulate []       _ acc = return acc
      simulate (y : ys) x acc = do
        let n = length acc
        x' <- normal (mean x n) sigmaX
        factor $ normalPdf (sq x' / 20) sigmaY y
        simulate ys x' (x' : acc)
  x0 <- normal 0 (sqrt 5)
  xs <- simulate obs x0 []
  return $ reverse xs

test :: IO [[([Double], Numeric.Log.Log Double)]]
test = sampleIO $ do
  let t = 5
  dat <- generateData t
  let ys = map snd dat
  pmmhRes <- prior $ pmmh 2 t 3 param (model ys)
  liftIO $ print pmmhRes
  return pmmhRes


{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE RankNTypes  #-}

module Main where

import qualified Data.HashTable.Class                 as C
import           Data.HashTable.IO
import           Data.HashTable.Test.Common
import qualified Data.Vector.Unboxed as V
import qualified Data.Vector.Unboxed.Mutable as VM
import           Statistics.Quantile (continuousBy, cadpw)
import           Statistics.Sample
import           System.Environment
import           System.Random.MWC

overhead :: C.HashTable h =>
            FixedTableType h ->
            GenIO ->
            IO Double
overhead dummy rng = do
    size <- uniformR (1000,50000) rng
    !v <- replicateM' size $ uniform rng
    let _ = v :: [(Int,Int)]

    !ht <- fromList v
    forceType dummy ht

    x <- computeOverhead ht
    return x

  where
    replicateM' :: Int -> IO a -> IO [a]
    replicateM' !sz f = go sz []
      where
        go !i !l | i == 0 = return l
                 | otherwise = do
                     !x <- f
                     go (i-1) (x:l)

-- Returns mean / stddev
runTrials :: C.HashTable h =>
             FixedTableType h
          -> GenIO
          -> Int
          -> IO (Double, Double, Double, Double)
runTrials dummy rng ntrials = do
    sample <- rep ntrials $ overhead dummy rng

    let (m, v) = meanVarianceUnb sample
    return (m, sqrt v, p95 sample, pMax sample)
  where
    p95 sample = continuousBy cadpw 19 20 sample

    pMax sample = V.foldl' max (-1) sample

    rep !n !f = do
        mv <- VM.new n
        go mv

      where
        go !mv = go' 0
          where
            go' !i | i >= n = V.unsafeFreeze mv
                   | otherwise = do
                !d <- f
                VM.unsafeWrite mv i d
                go' $ i+1

main :: IO ()
main = do
    rng <- do
        args <- getArgs
        if null args
          then withSystemRandom (\x -> (return x) :: IO GenIO)
          else initialize $ V.fromList [read $ head args]

    runTrials dummyLinearTable rng nTrials >>= report "linear hash table"
    runTrials dummyBasicTable rng nTrials >>= report "basic hash table"
    runTrials dummyCuckooTable rng nTrials >>= report "cuckoo hash table"

  where
    nTrials = 200

    report name md = putStrLn msg
      where msg = concat [ "\n(Mean,StdDev,95%,Max) for overhead of "
                         , name
                         , " ("
                         , show nTrials
                         , " trials): "
                         , show md
                         , "\n" ]

    dummyBasicTable = dummyTable
                      :: forall k v . BasicHashTable k v

    dummyLinearTable = dummyTable
                       :: forall k v . LinearHashTable k v

    dummyCuckooTable = dummyTable
                       :: forall k v . CuckooHashTable k v


{-# OPTIONS_GHC -Wall -Wno-orphans #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TemplateHaskell #-}
module TestBDD (bddTestGroup) where

import Control.Monad
import Control.Monad.ST
import Data.IntMap.Lazy (IntMap)
import qualified Data.IntMap.Lazy as IntMap
import Data.IntSet (IntSet)
import qualified Data.IntSet as IntSet
import Data.IORef
import Data.List
import qualified Data.Map.Lazy as Map
import Data.Proxy
import qualified Data.Set as Set
import Data.Vector (Vector)
import Data.Word
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 (BDD (..), ItemOrder (..))
import qualified Data.DecisionDiagram.BDD as BDD

import Utils

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

instance BDD.ItemOrder a => Arbitrary (BDD a) where
  arbitrary = arbitraryBDDOver =<< arbitrary

  shrink (BDD.Leaf _) = []
  shrink (BDD.Branch x p0 p1) =
    [p0, p1]
    ++
    [ BDD.Branch x p0' p1'
    | (p0', p1') <- shrink (p0, p1), p0' /= p1'
    ]

arbitraryBDDOver :: forall a. BDD.ItemOrder a => IntSet -> Gen (BDD a)
arbitraryBDDOver xs = do
  let f vs n = oneof $
        [ return BDD.true
        , return BDD.false
        ]
        ++
        [ do v <- elements vs
             let vs' = dropWhile (\v' -> compareItem (Proxy :: Proxy a) v' v  /= GT) vs
             lo <- f vs' (n `div` 2)
             hi <- f vs' (n `div` 2) `suchThat` (/= lo)
             return (BDD.Branch v lo hi)
        | n > 0, not (null vs)
        ]
  sized $ f (sortBy (BDD.compareItem (Proxy :: Proxy a)) $ IntSet.toList xs)

arbitrarySatisfyingAssignment :: forall a. BDD.ItemOrder a => BDD a -> IntSet -> Gen (IntMap Bool)
arbitrarySatisfyingAssignment bdd xs = do
  m1 <- arbitrarySatisfyingPartialAssignment bdd
  let ys = xs `IntSet.difference` IntMap.keysSet m1
  m2 <- liftM(IntMap.fromAscList) $ forM (IntSet.toAscList ys) $ \y -> do
    v <- arbitrary
    return (y,v)
  return $ m1 `IntMap.union` m2

arbitrarySatisfyingPartialAssignment :: forall a. BDD.ItemOrder a => BDD a -> Gen (IntMap Bool)
arbitrarySatisfyingPartialAssignment = f
  where
    f (BDD.Leaf True) = return IntMap.empty
    f (BDD.Leaf False) = undefined
    f (BDD.Branch x lo hi) = oneof $
      [liftM (IntMap.insert x False) (f lo) | lo /= BDD.Leaf False]
      ++
      [liftM (IntMap.insert x True) (f hi) | hi /= BDD.Leaf False]

-- ------------------------------------------------------------------------
-- conjunction
-- ------------------------------------------------------------------------

prop_and_unitL :: Property
prop_and_unitL =
  forAllItemOrder $ \(_ :: Proxy o) ->
    forAll arbitrary $ \(a :: BDD o) ->
      (BDD.true BDD..&&. a) === a

prop_and_unitR :: Property
prop_and_unitR =
  forAllItemOrder $ \(_ :: Proxy o) ->
    forAll arbitrary $ \(a :: BDD o) ->
      (a BDD..&&. BDD.true) === a

prop_and_falseL :: Property
prop_and_falseL =
  forAllItemOrder $ \(_ :: Proxy o) ->
    forAll arbitrary $ \(a :: BDD o) ->
      (BDD.false BDD..&&. a) === BDD.false

prop_and_falseR :: Property
prop_and_falseR =
  forAllItemOrder $ \(_ :: Proxy o) ->
    forAll arbitrary $ \(a :: BDD o) ->
      (a BDD..&&. BDD.false) === BDD.false

prop_and_comm :: Property
prop_and_comm =
  forAllItemOrder $ \(_ :: Proxy o) ->
    forAll arbitrary $ \(a :: BDD o, b) ->
      (a BDD..&&. b) === (b BDD..&&. a)

prop_and_assoc :: Property
prop_and_assoc =
  forAllItemOrder $ \(_ :: Proxy o) ->
    forAll arbitrary $ \(a :: BDD o, b, c) ->
      (a BDD..&&. (b BDD..&&. c)) === ((a BDD..&&. b) BDD..&&. c)

prop_and_idempotent :: Property
prop_and_idempotent =
  forAllItemOrder $ \(_ :: Proxy o) ->
    forAll arbitrary $ \(a :: BDD o) ->
      (a BDD..&&. a) === a

-- ------------------------------------------------------------------------
-- disjunction
-- ------------------------------------------------------------------------

prop_or_unitL :: Property
prop_or_unitL =
  forAllItemOrder $ \(_ :: Proxy o) ->
    forAll arbitrary $ \(a :: BDD o) ->
      (BDD.false BDD..||. a) === a

prop_or_unitR :: Property
prop_or_unitR =
  forAllItemOrder $ \(_ :: Proxy o) ->
    forAll arbitrary $ \(a :: BDD o) ->
      (a BDD..||. BDD.false) === a

prop_or_trueL :: Property
prop_or_trueL =
  forAllItemOrder $ \(_ :: Proxy o) ->
    forAll arbitrary $ \(a :: BDD o) ->
      (BDD.true BDD..||. a) === BDD.true

prop_or_trueR :: Property
prop_or_trueR =
  forAllItemOrder $ \(_ :: Proxy o) ->
    forAll arbitrary $ \(a :: BDD o) ->
      (a BDD..||. BDD.true) === BDD.true

prop_or_comm :: Property
prop_or_comm =
  forAllItemOrder $ \(_ :: Proxy o) ->
    forAll arbitrary $ \(a :: BDD o, b) ->
      (a BDD..||. b) === (b BDD..||. a)

prop_or_assoc :: Property
prop_or_assoc =
  forAllItemOrder $ \(_ :: Proxy o) ->
    forAll arbitrary $ \(a :: BDD o, b, c) ->
      (a BDD..||. (b BDD..||. c)) === ((a BDD..||. b) BDD..||. c)

prop_or_idempotent :: Property
prop_or_idempotent =
  forAllItemOrder $ \(_ :: Proxy o) ->
    forAll arbitrary $ \(a :: BDD o) ->
      (a BDD..||. a) === a

-- ------------------------------------------------------------------------
-- xor
-- ------------------------------------------------------------------------

prop_xor_unitL :: Property
prop_xor_unitL =
  forAllItemOrder $ \(_ :: Proxy o) ->
    forAll arbitrary $ \(a :: BDD o) ->
      (BDD.false `BDD.xor` a) === a

prop_xor_unitR :: Property
prop_xor_unitR =
  forAllItemOrder $ \(_ :: Proxy o) ->
    forAll arbitrary $ \(a :: BDD o) ->
      (a `BDD.xor` BDD.false) === a

prop_xor_comm :: Property
prop_xor_comm =
  forAllItemOrder $ \(_ :: Proxy o) ->
    forAll arbitrary $ \(a :: BDD o, b) ->
      (a `BDD.xor` b) === (b `BDD.xor` a)

prop_xor_assoc :: Property
prop_xor_assoc =
  forAllItemOrder $ \(_ :: Proxy o) ->
    forAll arbitrary $ \(a :: BDD o, b, c) ->
      (a `BDD.xor` (b `BDD.xor` c)) === ((a `BDD.xor` b) `BDD.xor` c)

prop_xor_involution :: Property
prop_xor_involution =
  forAllItemOrder $ \(_ :: Proxy o) ->
    forAll arbitrary $ \(a :: BDD o) ->
      (a `BDD.xor` a) === BDD.false

prop_xor_dist :: Property
prop_xor_dist =
  forAllItemOrder $ \(_ :: Proxy o) ->
    forAll arbitrary $ \(a :: BDD o, b, c) ->
      (a BDD..&&. (b `BDD.xor` c)) === ((a BDD..&&. b) `BDD.xor` (a BDD..&&. c))

-- ------------------------------------------------------------------------
-- distributivity
-- ------------------------------------------------------------------------

prop_dist_1 :: Property
prop_dist_1 =
  forAllItemOrder $ \(_ :: Proxy o) ->
    forAll arbitrary $ \(a :: BDD o, b, c) ->
      (a BDD..&&. (b BDD..||. c)) === ((a BDD..&&. b) BDD..||. (a BDD..&&. c))

prop_dist_2 :: Property
prop_dist_2 =
  forAllItemOrder $ \(_ :: Proxy o) ->
    forAll arbitrary $ \(a :: BDD o, b, c) ->
      (a BDD..||. (b BDD..&&. c)) === ((a BDD..||. b) BDD..&&. (a BDD..||. c))

prop_absorption_1 :: Property
prop_absorption_1 =
  forAllItemOrder $ \(_ :: Proxy o) ->
    forAll arbitrary $ \(a :: BDD o, b) ->
      (a BDD..&&. (a BDD..||. b)) === a

prop_absorption_2 :: Property
prop_absorption_2 =
  forAllItemOrder $ \(_ :: Proxy o) ->
    forAll arbitrary $ \(a :: BDD o, b) ->
      (a BDD..||. (a BDD..&&. b)) === a

-- ------------------------------------------------------------------------
-- negation
-- ------------------------------------------------------------------------

prop_double_negation :: Property
prop_double_negation =
  forAllItemOrder $ \(_ :: Proxy o) ->
    forAll arbitrary $ \(a :: BDD o) ->
      BDD.notB (BDD.notB a) === a

prop_and_complement :: Property
prop_and_complement =
  forAllItemOrder $ \(_ :: Proxy o) ->
    forAll arbitrary $ \(a :: BDD o) ->
      (a BDD..&&. BDD.notB a) === BDD.false

prop_or_complement :: Property
prop_or_complement =
  forAllItemOrder $ \(_ :: Proxy o) ->
    forAll arbitrary $ \(a :: BDD o) ->
      (a BDD..||. BDD.notB a) === BDD.true

prop_de_morgan_1 :: Property
prop_de_morgan_1 =
  forAllItemOrder $ \(_ :: Proxy o) ->
    forAll arbitrary $ \(a :: BDD o, b) ->
      BDD.notB (a BDD..||. b) === (BDD.notB a BDD..&&. BDD.notB b)

prop_de_morgan_2 :: Property
prop_de_morgan_2 =
  forAllItemOrder $ \(_ :: Proxy o) ->
    forAll arbitrary $ \(a :: BDD o, b) ->
      BDD.notB (a BDD..&&. b) === (BDD.notB a BDD..||. BDD.notB b)

-- ------------------------------------------------------------------------
-- Implication
-- ------------------------------------------------------------------------

prop_imply :: Property
prop_imply =
  forAllItemOrder $ \(_ :: Proxy o) ->
    forAll arbitrary $ \(a :: BDD o, b) ->
      (a BDD..=>. b) === (BDD.notB a BDD..||. b)

prop_imply_currying :: Property
prop_imply_currying =
  forAllItemOrder $ \(_ :: Proxy o) ->
    forAll arbitrary $ \(a :: BDD o, b, c) ->
      ((a BDD..&&. b) BDD..=>. c) === (a BDD..=>. (b BDD..=>. c))

-- ------------------------------------------------------------------------
-- Equivalence
-- ------------------------------------------------------------------------

prop_equiv :: Property
prop_equiv =
  forAllItemOrder $ \(_ :: Proxy o) ->
    forAll arbitrary $ \(a :: BDD o, b) ->
      (a BDD..<=>. b) === ((a BDD..=>. b) BDD..&&. (b BDD..=>. a))

-- ------------------------------------------------------------------------
-- If-then-else
-- ------------------------------------------------------------------------

prop_ite :: Property
prop_ite =
  forAllItemOrder $ \(_ :: Proxy o) ->
    forAll arbitrary $ \(c :: BDD o, t, e) ->
      BDD.ite c t e === ((c BDD..&&. t) BDD..||. (BDD.notB c BDD..&&. e))

prop_ite_swap_branch :: Property
prop_ite_swap_branch =
  forAllItemOrder $ \(_ :: Proxy o) ->
    forAll arbitrary $ \(c :: BDD o, t, e) ->
      BDD.ite c t e === BDD.ite (BDD.notB c) e t

prop_ite_dist_not :: Property
prop_ite_dist_not =
  forAllItemOrder $ \(_ :: Proxy o) ->
    forAll arbitrary $ \(c :: BDD o, t, e) ->
      BDD.notB (BDD.ite c t e) === BDD.ite c (BDD.notB t) (BDD.notB e)

prop_ite_dist_and :: Property
prop_ite_dist_and =
  forAllItemOrder $ \(_ :: Proxy o) ->
    forAll arbitrary $ \(c :: BDD o, t, e, d) ->
      (d BDD..&&. BDD.ite c t e) === BDD.ite c (d BDD..&&. t) (d BDD..&&. e)

prop_ite_dist_or :: Property
prop_ite_dist_or =
  forAllItemOrder $ \(_ :: Proxy o) ->
    forAll arbitrary $ \(c :: BDD o, t, e, d) ->
      (d BDD..||. BDD.ite c t e) === BDD.ite c (d BDD..||. t) (d BDD..||. e)

-- ------------------------------------------------------------------------
-- Pseudo-Boolean
-- ------------------------------------------------------------------------

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

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

prop_pbExactly :: Property
prop_pbExactly =
  forAllItemOrder $ \(_ :: Proxy o) ->
    forAll arbitrarySmallIntMap $ \(xs :: IntMap Integer) ->
      forAll arbitrary $ \k ->
        let a :: BDD o
            a = BDD.pbExactly xs k
         in counterexample (show a) $
              if a == BDD.Leaf False then
                property True
              else
                forAll (arbitrarySatisfyingAssignment a (IntMap.keysSet xs)) $ \ys ->
                  (IntMap.keysSet ys `IntSet.isSubsetOf` IntMap.keysSet xs)
                  .&&.
                  sum [xs IntMap.! y | (y,b) <- IntMap.toList ys, b] === k

prop_pbExactly_2 :: Property
prop_pbExactly_2 =
  forAllItemOrder $ \(_ :: Proxy o) ->
    forAll arbitrarySmallIntMap $ \(xs :: IntMap Integer) ->
      forAll (gen xs) $ \(m, k) ->
        let a :: BDD o
            a = BDD.pbExactly xs k
         in counterexample (show a) $ BDD.evaluate (m IntMap.!) a
  where
    gen :: IntMap Integer -> Gen (IntMap Bool, Integer)
    gen xs = do
      ys <- sublistOf (IntMap.toList xs)
      let ys' = IntSet.fromList [y | (y,_) <- ys]
      return
        ( IntMap.mapWithKey (\x _ -> x `IntSet.member` ys') xs
        , sum [w | (_,w) <- ys]
        )

prop_pbExactlyIntegral :: Property
prop_pbExactlyIntegral =
  forAllItemOrder $ \(_ :: Proxy o) ->
    forAll arbitrarySmallIntMap $ \(xs :: IntMap Integer) ->
      forAll arbitrary $ \k ->
        (BDD.pbExactlyIntegral xs k :: BDD o) === BDD.pbExactly xs k

-- ------------------------------------------------------------------------
-- Quantification
-- ------------------------------------------------------------------------

prop_forAll :: Property
prop_forAll =
  forAllItemOrder $ \(_ :: Proxy o) ->
    forAll arbitrary $ \(x, a :: BDD o) ->
      BDD.forAll x a === (BDD.restrict x True a BDD..&&. BDD.restrict x False a)

prop_exists :: Property
prop_exists =
  forAllItemOrder $ \(_ :: Proxy o) ->
    forAll arbitrary $ \(x, a :: BDD o) ->
      BDD.exists x a === (BDD.restrict x True a BDD..||. BDD.restrict x False a)

prop_existsUnique :: Property
prop_existsUnique =
  forAllItemOrder $ \(_ :: Proxy o) ->
    forAll arbitrary $ \(x, a :: BDD o) ->
      BDD.existsUnique x a === (BDD.restrict x True a `BDD.xor` BDD.restrict x False a)

prop_forAll_support :: Property
prop_forAll_support =
  forAllItemOrder $ \(_ :: Proxy o) ->
    forAll arbitrary $ \(x, a :: BDD o) ->
      let b = BDD.forAll x a
          xs = BDD.support b
       in counterexample (show (b, xs)) $
            x `IntSet.notMember` xs

prop_exists_support :: Property
prop_exists_support =
  forAllItemOrder $ \(_ :: Proxy o) ->
    forAll arbitrary $ \(x, a :: BDD o) ->
      let b = BDD.exists x a
          xs = BDD.support b
       in counterexample (show (b, xs)) $
            x `IntSet.notMember` xs

prop_existsUnique_support :: Property
prop_existsUnique_support =
  forAllItemOrder $ \(_ :: Proxy o) ->
    forAll arbitrary $ \(x, a :: BDD o) ->
      let b = BDD.existsUnique x a
          xs = BDD.support b
       in counterexample (show (b, xs)) $
            x `IntSet.notMember` xs

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

prop_forAllSet_empty :: Property
prop_forAllSet_empty =
  forAllItemOrder $ \(_ :: Proxy o) ->
    forAll arbitrary $ \(a :: BDD o) ->
      BDD.forAllSet IntSet.empty a === a

prop_existsSet_empty :: Property
prop_existsSet_empty =
  forAllItemOrder $ \(_ :: Proxy o) ->
    forAll arbitrary $ \(a :: BDD o) ->
      BDD.existsSet IntSet.empty a === a

prop_existsUniqueSet_empty :: Property
prop_existsUniqueSet_empty =
  forAllItemOrder $ \(_ :: Proxy o) ->
    forAll arbitrary $ \(a :: BDD o) ->
      BDD.existsUniqueSet IntSet.empty a === a

prop_forAllSet_singleton :: Property
prop_forAllSet_singleton =
  forAllItemOrder $ \(_ :: Proxy o) ->
    forAll arbitrary $ \(x, a :: BDD o) ->
      BDD.forAllSet (IntSet.singleton x) a === BDD.forAll x a

prop_existsSet_singleton :: Property
prop_existsSet_singleton =
  forAllItemOrder $ \(_ :: Proxy o) ->
    forAll arbitrary $ \(x, a :: BDD o) ->
      BDD.existsSet (IntSet.singleton x) a === BDD.exists x a

prop_existsUniqueSet_singleton :: Property
prop_existsUniqueSet_singleton =
  forAllItemOrder $ \(_ :: Proxy o) ->
    forAll arbitrary $ \(x, a :: BDD o) ->
      BDD.existsUniqueSet (IntSet.singleton x) a === BDD.existsUnique x a

prop_forAllSet_union :: Property
prop_forAllSet_union =
  forAllItemOrder $ \(_ :: Proxy o) ->
    forAll arbitrary $ \(xs1, xs2, a :: BDD o) ->
      BDD.forAllSet (xs1 `IntSet.union` xs2) a === BDD.forAllSet xs2 (BDD.forAllSet xs1 a)

prop_existsSet_union :: Property
prop_existsSet_union =
  forAllItemOrder $ \(_ :: Proxy o) ->
    forAll arbitrary $ \(xs1, xs2, a :: BDD o) ->
      BDD.existsSet (xs1 `IntSet.union` xs2) a === BDD.existsSet xs2 (BDD.existsSet xs1 a)

prop_existsUniqueSet_union :: Property
prop_existsUniqueSet_union =
  forAllItemOrder $ \(_ :: Proxy o) ->
    forAll arbitraryDisjointSets $ \(xs1, xs2) ->
    forAll arbitrary $ \(a :: BDD o) ->
      BDD.existsUniqueSet (xs1 `IntSet.union` xs2) a === BDD.existsUniqueSet xs2 (BDD.existsUniqueSet xs1 a)
  where
    arbitraryDisjointSets = do
      (u, v) <- arbitrary
      return (u `IntSet.intersection` v, u IntSet.\\ v)

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

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

case_fold'_strictness :: Assertion
case_fold'_strictness = do
  ref <- newIORef False
  let bdd :: BDD BDD.AscOrder
      bdd = BDD.Branch 0 (BDD.Branch 1 (BDD.Leaf False) (BDD.Leaf True)) (BDD.Branch 2 (BDD.Leaf False) (BDD.Leaf True))
      f x lo _hi = unsafePerformIO $ do
        when (x==2) $ writeIORef ref True
        return lo
  seq (BDD.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 $ \(bdd :: BDD o) ->
      BDD.fold (\x lo hi -> BDD.inSig (BDD.SBranch x lo hi)) (BDD.inSig . BDD.SLeaf) bdd
      ===
      bdd

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

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

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

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

case_support_false :: Assertion
case_support_false = BDD.support BDD.false @?= IntSet.empty

case_support_true :: Assertion
case_support_true = BDD.support BDD.true @?= IntSet.empty

prop_support_var :: Property
prop_support_var =
  forAll arbitrary $ \x ->
    BDD.support (BDD.var x) === IntSet.singleton x

prop_support_not :: Property
prop_support_not =
  forAllItemOrder $ \(_ :: Proxy o) ->
    forAll arbitrary $ \(a :: BDD o) ->
      let lhs = BDD.support a
          rhs = BDD.support (BDD.notB a)
       in lhs === rhs

prop_support_and :: Property
prop_support_and =
  forAllItemOrder $ \(_ :: Proxy o) ->
    forAll arbitrary $ \(a :: BDD o, b) ->
      let lhs = BDD.support (a BDD..&&. b)
          rhs = BDD.support a `IntSet.union` BDD.support b
       in counterexample (show (lhs, rhs)) $ lhs `IntSet.isSubsetOf` rhs

prop_support_or :: Property
prop_support_or =
  forAllItemOrder $ \(_ :: Proxy o) ->
    forAll arbitrary $ \(a :: BDD o, b) ->
      let lhs = BDD.support (a BDD..||. b)
          rhs = BDD.support a `IntSet.union` BDD.support b
       in counterexample (show (lhs, rhs)) $ lhs `IntSet.isSubsetOf` rhs

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

prop_evaluate_var :: Property
prop_evaluate_var =
 forAll arbitrary $ \(f' :: Fun Int Bool) ->
   let f = apply f'
    in forAllItemOrder $ \(_ :: Proxy o) ->
         forAll arbitrary $ \x ->
           BDD.evaluate f (BDD.var x :: BDD o) === f x

prop_evaluate_not :: Property
prop_evaluate_not =
 forAll arbitrary $ \(f' :: Fun Int Bool) ->
   let f = apply f'
    in forAllItemOrder $ \(_ :: Proxy o) ->
         forAll arbitrary $ \(a :: BDD o) ->
           BDD.evaluate f (BDD.notB a) === not (BDD.evaluate f a)

prop_evaluate_and :: Property
prop_evaluate_and =
 forAll arbitrary $ \(f' :: Fun Int Bool) ->
   let f = apply f'
    in forAllItemOrder $ \(_ :: Proxy o) ->
         forAll arbitrary $ \(a :: BDD o, b) ->
           BDD.evaluate f (a BDD..&&. b) === (BDD.evaluate f a && BDD.evaluate f b)

prop_evaluate_or :: Property
prop_evaluate_or =
 forAll arbitrary $ \(f' :: Fun Int Bool) ->
   let f = apply f'
    in forAllItemOrder $ \(_ :: Proxy o) ->
         forAll arbitrary $ \(a :: BDD o, b) ->
           BDD.evaluate f (a BDD..||. b) === (BDD.evaluate f a || BDD.evaluate f b)

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

case_numNodes :: Assertion
case_numNodes = do
  let bdd, bdd1 :: BDD BDD.AscOrder
      bdd = BDD.Branch 0 (BDD.Branch 1 BDD.false bdd1) (BDD.Branch 2 BDD.false bdd1)
      bdd1 = BDD.Branch 3 BDD.true BDD.false
  BDD.numNodes bdd @?= 6

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

prop_restrict :: Property
prop_restrict =
  forAllItemOrder $ \(_ :: Proxy o) ->
    forAll arbitrary $ \(a :: BDD o) ->
      forAll arbitrary $ \x ->
        let b = (BDD.var x BDD..&&. BDD.restrict x True a) BDD..||.
                (BDD.notB (BDD.var x) BDD..&&. BDD.restrict x False a)
         in a === b

prop_restrict_idempotent :: Property
prop_restrict_idempotent =
  forAllItemOrder $ \(_ :: Proxy o) ->
    forAll arbitrary $ \(a :: BDD o) ->
      forAll arbitrary $ \(x, val) ->
        let b = BDD.restrict x val a
            c = BDD.restrict x val b
         in counterexample (show (b, c)) $ b === c

prop_restrict_not :: Property
prop_restrict_not =
  forAllItemOrder $ \(_ :: Proxy o) ->
    forAll arbitrary $ \(a :: BDD o) ->
      forAll arbitrary $ \(x, val) ->
        BDD.restrict x val (BDD.notB a) === BDD.notB (BDD.restrict x val a)

prop_restrict_and :: Property
prop_restrict_and =
  forAllItemOrder $ \(_ :: Proxy o) ->
    forAll arbitrary $ \(a :: BDD o, b) ->
      forAll arbitrary $ \(x, val) ->
        BDD.restrict x val (a BDD..&&. b) === (BDD.restrict x val a BDD..&&. BDD.restrict x val b)

prop_restrict_or :: Property
prop_restrict_or =
  forAllItemOrder $ \(_ :: Proxy o) ->
    forAll arbitrary $ \(a :: BDD o, b) ->
      forAll arbitrary $ \(x, val) ->
        BDD.restrict x val (a BDD..||. b) === (BDD.restrict x val a BDD..||. BDD.restrict x val b)

prop_restrict_var :: Property
prop_restrict_var =
  forAllItemOrder $ \(_ :: Proxy o) ->
    forAll arbitrary $ \x ->
      let a :: BDD o
          a = BDD.var x
       in (BDD.restrict x True a === BDD.true) .&&.
          (BDD.restrict x False a === BDD.false)

prop_restrict_support :: Property
prop_restrict_support =
  forAllItemOrder $ \(_ :: Proxy o) ->
    forAll arbitrary $ \(a :: BDD o) ->
      forAll arbitrary $ \(x, val) ->
        let b = BDD.restrict x val a
            xs = BDD.support b
         in counterexample (show b) $
            counterexample (show xs) $
              x `IntSet.notMember` xs

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

prop_restrictSet_empty :: Property
prop_restrictSet_empty =
  forAllItemOrder $ \(_ :: Proxy o) ->
    forAll arbitrary $ \(a :: BDD o) ->
      BDD.restrictSet IntMap.empty a === a

prop_restrictSet_singleton :: Property
prop_restrictSet_singleton =
  forAllItemOrder $ \(_ :: Proxy o) ->
    forAll arbitrary $ \(a :: BDD o) ->
      forAll arbitrary $ \(x, val) ->
        BDD.restrict x val a === BDD.restrictSet (IntMap.singleton x val) a

prop_restrictSet_union :: Property
prop_restrictSet_union =
  forAllItemOrder $ \(_ :: Proxy o) ->
    forAll arbitrary $ \(a :: BDD o) ->
      forAll arbitrary $ \(val1, val2) ->
        and (IntMap.intersectionWith (==) val1 val2)
        ==>
        (BDD.restrictSet val2 (BDD.restrictSet val1 a) === BDD.restrictSet (val1 `IntMap.union` val2) a)

prop_restrictSet_idempotent :: Property
prop_restrictSet_idempotent =
  forAllItemOrder $ \(_ :: Proxy o) ->
    forAll arbitrary $ \(a :: BDD o) ->
      forAll arbitrary $ \val ->
        let b = BDD.restrictSet val a
            c = BDD.restrictSet val b
         in counterexample (show (b, c)) $ b === c

prop_restrictSet_not :: Property
prop_restrictSet_not =
  forAllItemOrder $ \(_ :: Proxy o) ->
    forAll arbitrary $ \(a :: BDD o) ->
      forAll arbitrary $ \val ->
        BDD.restrictSet val (BDD.notB a) === BDD.notB (BDD.restrictSet val a)

prop_restrictSet_and :: Property
prop_restrictSet_and =
  forAllItemOrder $ \(_ :: Proxy o) ->
    forAll arbitrary $ \(a :: BDD o, b) ->
      forAll arbitrary $ \val ->
        BDD.restrictSet val (a BDD..&&. b) === (BDD.restrictSet val a BDD..&&. BDD.restrictSet val b)

prop_restrictSet_or :: Property
prop_restrictSet_or =
  forAllItemOrder $ \(_ :: Proxy o) ->
    forAll arbitrary $ \(a :: BDD o, b) ->
      forAll arbitrary $ \val ->
        BDD.restrictSet val (a BDD..||. b) === (BDD.restrictSet val a BDD..||. BDD.restrictSet val b)

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

prop_restrictLaw :: Property
prop_restrictLaw =
  forAllItemOrder $ \(_ :: Proxy o) ->
    forAll arbitrary $ \(a :: BDD o) ->
      forAll arbitrary $ \law ->
        (law BDD..&&. BDD.restrictLaw law a) === (law BDD..&&. a)

case_restrictLaw_case_0 :: Assertion
case_restrictLaw_case_0 = (law BDD..&&. BDD.restrictLaw law a) @?= (law BDD..&&. a)
  where
    a, law :: BDD BDD.AscOrder
    a = BDD.Branch 2 (BDD.Leaf False) (BDD.Leaf True)
    law = BDD.Branch 1 (BDD.Branch 2 (BDD.Leaf True) (BDD.Leaf False)) (BDD.Branch 2 (BDD.Leaf False) (BDD.Leaf True))

prop_restrictLaw_true :: Property
prop_restrictLaw_true =
  forAllItemOrder $ \(_ :: Proxy o) ->
    forAll arbitrary $ \(a :: BDD o) ->
      BDD.restrictLaw BDD.true a === a

prop_restrictLaw_self :: Property
prop_restrictLaw_self =
  forAllItemOrder $ \(_ :: Proxy o) ->
    forAll arbitrary $ \(a :: BDD o) ->
      (a /= BDD.false) ==> BDD.restrictLaw a a === BDD.true

prop_restrictLaw_not_self :: Property
prop_restrictLaw_not_self =
  forAllItemOrder $ \(_ :: Proxy o) ->
    forAll arbitrary $ \(a :: BDD o) ->
      (a /= BDD.true) ==> BDD.restrictLaw (BDD.notB a) a === BDD.false

prop_restrictLaw_restrictSet :: Property
prop_restrictLaw_restrictSet =
  forAllItemOrder $ \(_ :: Proxy o) ->
    forAll arbitrary $ \(a :: BDD o) ->
      forAll arbitrary $ \val ->
        let b = BDD.andB [if v then BDD.var x else BDD.notB (BDD.var x) | (x,v) <- IntMap.toList val]
         in BDD.restrictLaw b a === BDD.restrictSet val a

-- prop_restrictLaw_and_condition :: Property
-- prop_restrictLaw_and_condition =
--   forAllItemOrder $ \(_ :: Proxy o) ->
--     forAll arbitrary $ \(a :: BDD o) ->
--       forAll arbitrary $ \(val1, val2) ->
--         let val = val1 BDD..&&. val2
--          in counterexample (show val) $
--               (val /= BDD.false)
--               ==>
--               (BDD.restrictLaw val a === BDD.restrictLaw val2 (BDD.restrictLaw val1 a))

-- counterexample to the above prop_restrictLaw_and_condition
case_restrictLaw_case_1 :: Assertion
case_restrictLaw_case_1 = do
  -- BDD.restrictLaw val a @?= BDD.restrictLaw val2 (BDD.restrictLaw val1 a)
  BDD.restrictLaw val a @?= BDD.Branch 2 (BDD.Leaf False) (BDD.Leaf True)
  BDD.restrictLaw val2 (BDD.restrictLaw val1 a) @?= BDD.Branch 1 (BDD.Leaf True) (Branch 2 (BDD.Leaf False) (BDD.Leaf True))
  where
    a :: BDD BDD.AscOrder
    a = Branch 2 (BDD.Leaf False) (BDD.Leaf True) -- x2
    val1 = BDD.Branch 1 (BDD.Leaf False) (BDD.Leaf True) -- x1
    val2 = BDD.Branch 1 (BDD.Branch 2 (BDD.Leaf False) (BDD.Leaf True)) (BDD.Leaf True) -- x1 ∨ x2
    val = val1 BDD..&&. val2 -- x1

prop_restrictLaw_or_condition :: Property
prop_restrictLaw_or_condition =
  forAllItemOrder $ \(_ :: Proxy o) ->
    forAll arbitrary $ \(a :: BDD o) ->
      forAll arbitrary $ \(val1, val2) ->
        let val = val1 BDD..||. val2
         in counterexample (show val) $
              (val BDD..&&. BDD.restrictLaw val a) === (val1 BDD..&&. BDD.restrictLaw val1 a BDD..||. val2 BDD..&&. BDD.restrictLaw val2 a)

prop_restrictLaw_idempotent :: Property
prop_restrictLaw_idempotent =
  forAllItemOrder $ \(_ :: Proxy o) ->
    forAll arbitrary $ \(a :: BDD o) ->
      forAll arbitrary $ \val ->
        let b = BDD.restrictLaw val a
            c = BDD.restrictLaw val b
         in counterexample (show (b, c)) $ b === c

prop_restrictLaw_not :: Property
prop_restrictLaw_not =
  forAllItemOrder $ \(_ :: Proxy o) ->
    forAll arbitrary $ \(a :: BDD o) ->
      forAll (arbitrary `suchThat` (/= BDD.false)) $ \val ->
        BDD.restrictLaw val (BDD.notB a) === BDD.notB (BDD.restrictLaw val a)

prop_restrictLaw_and :: Property
prop_restrictLaw_and =
  forAllItemOrder $ \(_ :: Proxy o) ->
    forAll arbitrary $ \(a :: BDD o, b) ->
      forAll (arbitrary `suchThat` (/= BDD.false)) $ \val ->
        BDD.restrictLaw val (a BDD..&&. b) === (BDD.restrictLaw val a BDD..&&. BDD.restrictLaw val b)

prop_restrictLaw_or :: Property
prop_restrictLaw_or =
  forAllItemOrder $ \(_ :: Proxy o) ->
    forAll arbitrary $ \(a :: BDD o, b) ->
      forAll (arbitrary `suchThat` (/= BDD.false)) $ \val ->
        BDD.restrictLaw val (a BDD..||. b) === (BDD.restrictLaw val a BDD..||. BDD.restrictLaw val b)

-- prop_restrictLaw_minimality :: Property
-- prop_restrictLaw_minimality =
--   forAllItemOrder $ \(_ :: Proxy o) ->
--     forAll arbitrary $ \(a :: BDD o) ->
--       forAll arbitrary $ \law ->
--         let b = BDD.restrictLaw law a
--          in counterexample (show b) $
--               ((law BDD..&&. b) === (law BDD..&&. a))
--               .&&.
--               conjoin [counterexample (show b') $ (law BDD..&&. b') =/= (law BDD..&&. a) | b' <- shrink b]

case_restrictLaw_non_minimal_1 :: Assertion
case_restrictLaw_non_minimal_1 = do
  (law BDD..&&. BDD.restrictLaw law a) @?= (law BDD..&&. a)
  BDD.restrictLaw law a @?= b -- should be 'a'?
  where
    law, a :: BDD BDD.AscOrder
    law = BDD.Branch 1 (BDD.Branch 2 (BDD.Leaf False) (BDD.Leaf True)) (BDD.Leaf True) -- x1 ∨ x2
    a = BDD.Branch 2 (BDD.Leaf True) (BDD.Leaf False) -- ¬x2
    b = BDD.Branch 1 (BDD.Leaf False) (BDD.Branch 2 (BDD.Leaf True) (BDD.Leaf False)) -- x1 ∧ ¬x2

case_restrictLaw_non_minimal_2 :: Assertion
case_restrictLaw_non_minimal_2 = do
  (law BDD..&&. BDD.restrictLaw law a) @?= (law BDD..&&. a)
  BDD.restrictLaw law a @?= b -- should be 'a'?
  where
    law, a, b :: BDD BDD.AscOrder
    law = BDD.Branch 1 (BDD.Leaf True) (BDD.Branch 2 (BDD.Leaf False) (BDD.Leaf True)) -- ¬x1 ∨ x2
    a = BDD.Branch 2 (BDD.Leaf False) (BDD.Leaf True) -- x2
    b = BDD.Branch 1 (BDD.Branch 2 (BDD.Leaf False) (BDD.Leaf True)) (BDD.Leaf True) -- x1 ∨ x2

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

prop_subst_restrict_constant :: Property
prop_subst_restrict_constant =
  forAllItemOrder $ \(_ :: Proxy o) ->
    forAll arbitrary $ \(m :: BDD o) ->
    forAll arbitrary $ \x ->
    forAll arbitrary $ \val ->
      BDD.subst x (if val then BDD.true else BDD.false) m === BDD.restrict x val m

prop_subst_restrict :: Property
prop_subst_restrict =
  forAllItemOrder $ \(_ :: Proxy o) ->
    forAll arbitrary $ \(m :: BDD o) ->
    forAll arbitrary $ \x ->
    forAll arbitrary $ \(n :: BDD o) ->
      BDD.subst x n m === ((n BDD..&&. BDD.restrict x True m) BDD..||. (BDD.notB n BDD..&&. BDD.restrict x False m))

prop_subst_same_var :: Property
prop_subst_same_var =
  forAllItemOrder $ \(_ :: Proxy o) ->
    forAll arbitrary $ \(m :: BDD o) ->
    forAll arbitrary $ \x ->
      BDD.subst x (BDD.var x) m === m

prop_subst_not_occured :: Property
prop_subst_not_occured =
  forAllItemOrder $ \(_ :: Proxy o) ->
    forAll arbitrary $ \(m :: BDD o) ->
    forAll (arbitrary `suchThat` (\x -> x `IntSet.notMember` (BDD.support m))) $ \x ->
    forAll arbitrary $ \(n :: BDD o) ->
      BDD.subst x n m === m

-- If x1≠x2 and x1∉FV(M2) then M[x1 ↦ M1][x2 ↦ M2] = M[x2 ↦ M2][x1 ↦ M1[x2 ↦ M2]].
prop_subst_dist :: Property
prop_subst_dist =
  forAllItemOrder $ \(_ :: Proxy o) ->
    forAll arbitrary $ \(x1, m1) ->
    forAll ((,) <$> (arbitrary `suchThat` (/= x1)) <*> (arbitrary `suchThat` (\m2 -> x1 `IntSet.notMember` BDD.support m2))) $ \(x2, m2) ->
    forAll arbitrary $ \(m :: BDD o) ->
      BDD.subst x2 m2 (BDD.subst x1 m1 m) === BDD.subst x1 (BDD.subst x2 m2 m1) (BDD.subst x2 m2 m)

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

prop_substSet_empty :: Property
prop_substSet_empty =
  forAllItemOrder $ \(_ :: Proxy o) ->
    forAll arbitrary $ \(m :: BDD o) ->
      BDD.substSet IntMap.empty m === m

prop_substSet_singleton :: Property
prop_substSet_singleton =
  forAllItemOrder $ \(_ :: Proxy o) ->
    forAll arbitrary $ \(m :: BDD o) ->
    forAll arbitrary $ \x ->
    forAll arbitrary $ \m1 ->
      BDD.substSet (IntMap.singleton x m1) m === BDD.subst x m1 m

case_substSet_case_1 :: Assertion
case_substSet_case_1 = do
  BDD.substSet (IntMap.singleton x m1) m @?= BDD.subst x m1 m
  where
    m :: BDD BDD.AscOrder
    m = BDD.Branch 1 (BDD.Branch 2 (BDD.Leaf True) (BDD.Leaf False)) (BDD.Branch 2 (BDD.Leaf False) (BDD.Leaf True))
    x = 1
    m1 = BDD.Branch 1 (BDD.Leaf True) (BDD.Leaf False)

prop_substSet_same_vars :: Property
prop_substSet_same_vars =
  forAllItemOrder $ \(_ :: Proxy o) ->
    forAll arbitrary $ \(m :: BDD o) ->
    forAll arbitrary $ \xs ->
      BDD.substSet (IntMap.fromAscList [(x, BDD.var x) | x <- IntSet.toAscList xs]) m === m

prop_substSet_not_occured :: Property
prop_substSet_not_occured =
  forAllItemOrder $ \(_ :: Proxy o) ->
    forAll arbitrary $ \(m :: BDD o) ->
    forAll (f (BDD.support m)) $ \s ->
      BDD.substSet s m === m
  where
    f xs = liftM IntMap.fromList $ listOf $ do
      y <- arbitrary `suchThat` (`IntSet.notMember` xs)
      m <- arbitrary
      return (y, m)

prop_substSet_compose :: Property
prop_substSet_compose =
  forAllItemOrder $ \(_ :: Proxy o) ->
    forAll arbitrary $ \xs ->
    forAll arbitrary $ \ys ->
    forAll (liftM IntMap.fromList $ mapM (\x -> (,) <$> pure x <*> arbitraryBDDOver ys) (IntSet.toList xs)) $ \s1 ->
    forAll (liftM IntMap.fromList $ mapM (\y -> (,) <$> pure y <*> arbitrary) (IntSet.toList ys)) $ \s2 ->
    forAll (arbitraryBDDOver xs) $ \(m :: BDD o) ->
      BDD.substSet s2 (BDD.substSet s1 m) === BDD.substSet (IntMap.map (BDD.substSet s2) s1) m

case_substSet_case_2 :: Assertion
case_substSet_case_2 = do
  let m :: BDD BDD.AscOrder
      m = BDD.var 1 BDD..&&. BDD.var 2
  BDD.substSet (IntMap.fromList [(1, BDD.var 2), (2, BDD.var 3)]) m @?= BDD.var 2 BDD..&&. BDD.var 3
  BDD.substSet (IntMap.fromList [(1, BDD.var 3), (2, BDD.var 1)]) m @?= BDD.var 3 BDD..&&. BDD.var 1

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

data MonotoneExpr a
  = MVar a
  | MAnd (MonotoneExpr a) (MonotoneExpr a)
  | MOr (MonotoneExpr a) (MonotoneExpr a)
  | MConst Bool
  deriving (Show)

arbitraryMonotoneExpr :: forall a. Gen a -> Gen (MonotoneExpr a)
arbitraryMonotoneExpr gen = sized f
  where
    f :: Int -> Gen (MonotoneExpr a)
    f n = oneof $
      [ liftM MConst arbitrary
      , liftM MVar gen
      ]
      ++
      concat
      [ [liftM2 MAnd sub sub, liftM2 MOr sub sub]
      | n > 0, let sub = f (n `div` 2)
      ]

evalMonotoneExpr :: ItemOrder a => (b -> BDD a) -> MonotoneExpr b -> BDD a
evalMonotoneExpr f = g
  where
    g (MVar a) = f a
    g (MConst v) = BDD.Leaf v
    g (MOr a b) = g a BDD..||. g b
    g (MAnd a b) = g a BDD..&&. g b

forAllMonotonicFunction :: forall o prop. (ItemOrder o, Testable prop) => IntSet -> ((BDD o -> BDD o) -> prop) -> Property
forAllMonotonicFunction xs k =
  forAll (arbitraryMonotoneExpr (elements (Nothing : map Just (IntSet.toList xs)))) $ \e -> do
    let f :: BDD o -> BDD o
        f x = evalMonotoneExpr g e
          where
            g Nothing = x
            g (Just v) = BDD.var v
     in k f

prop_lfp_is_fixed_point :: Property
prop_lfp_is_fixed_point =
 forAll arbitrary $ \(xs :: IntSet) ->
   forAllItemOrder $ \(_ :: Proxy o) ->
     forAllMonotonicFunction xs $ \(f :: BDD o -> BDD o) -> do
       let a = BDD.lfp f
        in counterexample (show a) $ f a === a

prop_gfp_is_fixed_point :: Property
prop_gfp_is_fixed_point =
 forAll arbitrary $ \(xs :: IntSet) ->
   forAllItemOrder $ \(_ :: Proxy o) ->
     forAllMonotonicFunction xs $ \(f :: BDD o -> BDD o) -> do
       let a = BDD.gfp f
        in counterexample (show a) $ f a === a

prop_lfp_imply_gfp :: Property
prop_lfp_imply_gfp =
 forAll arbitrary $ \(xs :: IntSet) ->
   forAllItemOrder $ \(_ :: Proxy o) ->
     forAllMonotonicFunction xs $ \(f :: BDD o -> BDD o) -> do
       let a = BDD.lfp f
           b = BDD.gfp f
        in counterexample (show (a, b)) $ (a BDD..=>. b) === BDD.true

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

prop_anySat :: Property
prop_anySat =
  forAllItemOrder $ \(_ :: Proxy o) ->
    forAll arbitrary $ \(bdd :: BDD o) ->
      case BDD.anySat bdd of
        Just p -> counterexample (show p) $ BDD.evaluate (p IntMap.!) bdd
        Nothing -> bdd === BDD.Leaf False

prop_allSat :: Property
prop_allSat =
  forAllItemOrder $ \(_ :: Proxy o) ->
    forAll arbitrarySmallIntSet $ \xs ->
      forAll (arbitraryBDDOver xs) $ \(bdd :: BDD o) ->
         let ps = BDD.allSat bdd
         in null ps === (bdd == BDD.Leaf False)
            .&&.
            conjoin [counterexample (show p) $ BDD.evaluate (p IntMap.!) bdd |  p <- ps]

prop_anySatComplete :: Property
prop_anySatComplete =
  forAllItemOrder $ \(_ :: Proxy o) ->
    forAll arbitrary $ \xs ->
      forAll (arbitraryBDDOver xs) $ \(bdd :: BDD o) ->
        case BDD.anySatComplete xs bdd of
          Just p -> counterexample (show p) $
            IntMap.keysSet p === xs
            .&&.
            BDD.evaluate (p IntMap.!) bdd
          Nothing -> bdd === BDD.Leaf False

prop_allSatComplete :: Property
prop_allSatComplete =
  forAllItemOrder $ \(_ :: Proxy o) ->
    forAll arbitrarySmallIntSet $ \xs ->
      forAll (arbitraryBDDOver xs) $ \(bdd :: BDD o) ->
        let ps = BDD.allSatComplete xs bdd
            qs = [q | q <- foldM (\m x -> [IntMap.insert x v m | v <- [False, True]]) IntMap.empty (IntSet.toList xs)
                    , BDD.evaluate (q IntMap.!) bdd]
         in conjoin [counterexample (show p) (IntMap.keysSet p === xs) | p <- ps]
            .&&.
            Set.fromList ps === Set.fromList qs

prop_countSat_allSatComplete :: Property
prop_countSat_allSatComplete =
  forAllItemOrder $ \(_ :: Proxy o) ->
    forAll arbitrarySmallIntSet $ \xs ->
      forAll (arbitraryBDDOver xs) $ \(bdd :: BDD o) ->
        let ps = BDD.allSatComplete xs bdd
            n = BDD.countSat xs bdd
         in counterexample (show n) $
              if bdd == BDD.Leaf False then
                n === 0
              else
                -- Note that the number of partial assignments is smaller than the number of total assignments
                (n > 0) .&&. n === length ps

prop_uniformSatM :: Property
prop_uniformSatM =
  forAllItemOrder $ \(_ :: Proxy o) ->
    forAll (arbitrarySmallIntSet `suchThat` ((>= 2) . IntSet.size)) $ \xs ->
      forAll (arbitraryBDDOver xs `suchThat` ((>= (2::Integer)) . BDD.countSat xs)) $ \(bdd :: BDD o) ->
        forAll arbitrary $ \(seed :: Vector Word32) ->
          let m :: Integer
              m = BDD.countSat xs bdd
              n = 1000
              samples = runST $ do
                gen <- Rand.initialize seed
                replicateM n $ BDD.uniformSatM xs bdd gen
              hist_actual = Map.fromListWith (+) [(s, 1 :: Double) | s <- samples]
              hist_expected = [(s, fromIntegral n / fromIntegral m :: Double) | s <- BDD.allSatComplete xs bdd]
              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 [BDD.evaluate (a IntMap.!) bdd | a <- Map.keys hist_actual]
                .&&.
                counterexample ("χ² = " ++ show chi_sq ++ " >= " ++ show threshold) (chi_sq < threshold)

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

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

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

arbitrarySmallIntSet :: Gen IntSet
arbitrarySmallIntSet = do
  n <- choose (0, 12)
  liftM IntSet.fromList $ replicateM n arbitrary

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)

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

bddTestGroup :: TestTree
bddTestGroup = $(testGroupGenerator)


{-# LANGUAGE FlexibleContexts #-}

module Main (main) where

import Control.Monad (mzero)
import Control.Monad.IO.Class (liftIO)
import Data.Complex (Complex)
import Data.Foldable (toList)
import Data.Typeable (Typeable)
import System.Console.GetOpt
import Text.PrettyPrint.Mainland
import Text.PrettyPrint.Mainland.Class

import Spiral
import Spiral.Backend.C
import Spiral.Config
import Spiral.Driver
import Spiral.Exp
import Spiral.FFT.CooleyTukey
import Spiral.Monad
import Spiral.OpCount
import Spiral.Program
import Spiral.SPL
import Spiral.SPL.Run
import Spiral.Search
import Spiral.Search.FFTBreakdowns
import Spiral.Util.Uniq

main :: IO ()
main = defaultMainWith' options mempty $ \fs args -> do
    useComplexType <- asksConfig $ testDynFlag UseComplex
    n <- case args of
           [s] -> return (read s)
           _   -> return 4
    f <- formula fs n
    pprint f
    if useComplexType
      then toProgram ("hspiral_dft_" ++ show n) f >>= go
      else toProgram ("hspiral_dft_" ++ show n) (Re f) >>= go
  where
    go :: (Typed a, Num (Exp a)) => Program a -> Spiral ()
    go prog = do
      pprint prog
      ops <- countProgramOps prog
      resetUnique
      defs <- evalCg $ cgProgram prog
      outp <- asksConfig output
      case outp of
        Nothing -> return ()
        Just{}  -> writeOutput (toList defs)
      liftIO $ putDocLn $
          text "Multiplications:" <+> ppr (mulOps ops) </>
          text "      Additions:" <+> ppr (addOps ops) </>
          text "          Total:" <+> ppr (allOps ops)

-- The SPL formula for which we generate code and count operations.
formula :: MonadSpiral m => [Flag] -> Int -> m (SPL (Exp (Complex Double)))
formula fs n =
  case fs of
    [Dif]                -> return $ dif n
    [Dit]                -> return $ dit n
    [SplitRadix]         -> runSearch () splitRadixSearch (DFT n)
    [ConjPairSplitRadix] -> runSearch () conjSplitRadixSearch (DFT n)
    [ImpSplitRadix]      -> runSearch () impSplitRadixSearch (DFT n)
    _                    -> fail "Must specify exactly on of --dif, --dit, --split-radix, conj-split-radix, or --improved-split-radix"
  where
    splitRadixSearch :: (Typeable a, Typed a, MonadSpiral m)
                     => SPL (Exp a)
                     -> S s m (SPL (Exp a))
    splitRadixSearch (F n w) = splitRadixBreakdown n w
    splitRadixSearch _       = mzero

    conjSplitRadixSearch :: (Typeable a, Typed a, MonadSpiral m)
                         => SPL (Exp a)
                         -> S s m (SPL (Exp a))
    conjSplitRadixSearch (F n w) = conjPairSplitRadixBreakdown n w
    conjSplitRadixSearch _       = mzero

    impSplitRadixSearch :: (Typeable a, Typed a, MonadSpiral m)
                        => SPL (Exp a)
                        -> S s m (SPL (Exp a))
    impSplitRadixSearch (F n w) = improvedSplitRadixBreakdown n w
    impSplitRadixSearch _       = mzero

data Flag = Dif
          | Dit
          | SplitRadix
          | ConjPairSplitRadix
          | ImpSplitRadix
  deriving (Eq, Ord, Show)

options :: [OptDescr Flag]
options =
    [ Option [] ["dif"] (NoArg Dif)                             "Use DIF"
    , Option [] ["dit"] (NoArg Dit)                             "Use DIT"
    , Option [] ["split-radix"] (NoArg SplitRadix)              "Use split radix"
    , Option [] ["conj-split-radix"] (NoArg ConjPairSplitRadix) "Use conjugate pair split radix"
    , Option [] ["imp-split-radix"] (NoArg ImpSplitRadix)       "Use improved split radix"
    ]