{-# language DeriveGeneric #-}
{-# options_ghc -Wno-unused-imports #-}
module Data.VPTree.TestData where

-- import Data.Foldable (toList)
import GHC.Generics (Generic(..))
import Text.Printf (printf)

-- deepseq
import Control.DeepSeq (NFData())
-- mwc-probability
import qualified System.Random.MWC.Probability as P (Gen, Prob, withSystemRandom, asGenIO, GenIO, create, initialize, sample, samples, normal, bernoulli, uniformR)
-- primitive
import Control.Monad.Primitive (PrimMonad(..), PrimState)
-- vector
import qualified Data.Vector as V (Vector, map, filter, length, toList, replicate, partition, zipWith, head, tail, fromList, thaw, freeze, (!), foldl)


import Data.VPTree.Build (build)
-- import Data.VPTree.Draw (draw)
import Data.VPTree.Internal (VT, VPTree, withST, withST_, withIO)
-- import Data.VPTree.Query (range, distances)


-- test data

data P = P !Double !Double deriving (P -> P -> Bool
(P -> P -> Bool) -> (P -> P -> Bool) -> Eq P
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: P -> P -> Bool
$c/= :: P -> P -> Bool
== :: P -> P -> Bool
$c== :: P -> P -> Bool
Eq, (forall x. P -> Rep P x) -> (forall x. Rep P x -> P) -> Generic P
forall x. Rep P x -> P
forall x. P -> Rep P x
forall a.
(forall x. a -> Rep a x) -> (forall x. Rep a x -> a) -> Generic a
$cto :: forall x. Rep P x -> P
$cfrom :: forall x. P -> Rep P x
Generic)
instance NFData P
instance Show P where
  show :: P -> String
show (P Double
x Double
y) = String -> Double -> Double -> String
forall r. PrintfType r => String -> r
printf String
"(%2.2f, %2.2f)" Double
x Double
y --show (x,y)

(.+.) :: P -> P -> P
P Double
x1 Double
y1 .+. :: P -> P -> P
.+. P Double
x2 Double
y2 = Double -> Double -> P
P (Double
x1 Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
x2) (Double
y1 Double -> Double -> Double
forall a. Num a => a -> a -> a
+ Double
y2)

distp :: P -> P -> Double
distp :: P -> P -> Double
distp (P Double
x1 Double
y1) (P Double
x2 Double
y2) = Double -> Double
forall a. Floating a => a -> a
sqrt (Double -> Double) -> Double -> Double
forall a b. (a -> b) -> a -> b
$ (Double
x1 Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
x2)Double -> Double -> Double
forall a. Floating a => a -> a -> a
**Double
2 Double -> Double -> Double
forall a. Num a => a -> a -> a
+ (Double
y1 Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
y2)Double -> Double -> Double
forall a. Floating a => a -> a -> a
**Double
2



-- t2, t2', t3 :: VPTree Double P
-- t3 = buildP $ genN3 12
-- t2 = buildP $ genN2 12
-- t2' = buildP $ genN2 10000

genN1, gaussMixSamples, binDiskSamples :: Int -> V.Vector P
gaussMixSamples :: Int -> Vector P
gaussMixSamples Int
n = [P] -> Vector P
forall a. [a] -> Vector a
V.fromList ([P] -> Vector P) -> [P] -> Vector P
forall a b. (a -> b) -> a -> b
$ (forall s. Gen s -> ST s [P]) -> [P]
forall a. (forall s. Gen s -> ST s a) -> a
withST_ (Int -> Prob (ST s) P -> Gen (PrimState (ST s)) -> ST s [P]
forall (m :: * -> *) a.
PrimMonad m =>
Int -> Prob m a -> Gen (PrimState m) -> m [a]
P.samples Int
n (Double -> Double -> Double -> Double -> Prob (ST s) P
forall (m :: * -> *).
PrimMonad m =>
Double -> Double -> Double -> Double -> Prob m P
gaussMix Double
0 Double
25 Double
1 Double
1))

genN1 :: Int -> Vector P
genN1 Int
n = [P] -> Vector P
forall a. [a] -> Vector a
V.fromList ([P] -> Vector P) -> [P] -> Vector P
forall a b. (a -> b) -> a -> b
$ (forall s. Gen s -> ST s [P]) -> [P]
forall a. (forall s. Gen s -> ST s a) -> a
withST_ (Int -> Prob (ST s) P -> Gen (PrimState (ST s)) -> ST s [P]
forall (m :: * -> *) a.
PrimMonad m =>
Int -> Prob m a -> Gen (PrimState m) -> m [a]
P.samples Int
n (Double -> Double -> Prob (ST s) P
forall (m :: * -> *). PrimMonad m => Double -> Double -> Prob m P
isoNormal2d Double
0 Double
1))

binDiskSamples :: Int -> Vector P
binDiskSamples Int
n = [P] -> Vector P
forall a. [a] -> Vector a
V.fromList ([P] -> Vector P) -> [P] -> Vector P
forall a b. (a -> b) -> a -> b
$ (forall s. Gen s -> ST s [P]) -> [P]
forall a. (forall s. Gen s -> ST s a) -> a
withST_ ((forall s. Gen s -> ST s [P]) -> [P])
-> (forall s. Gen s -> ST s [P]) -> [P]
forall a b. (a -> b) -> a -> b
$ Int -> Prob (ST s) P -> Gen (PrimState (ST s)) -> ST s [P]
forall (m :: * -> *) a.
PrimMonad m =>
Int -> Prob m a -> Gen (PrimState m) -> m [a]
P.samples Int
n (Double -> Double -> P -> P -> Prob (ST s) P
forall (m :: * -> *).
PrimMonad m =>
Double -> Double -> P -> P -> Prob m P
binDisk Double
1 Double
1 P
z P
fv)
  where
    z :: P
z = Double -> Double -> P
P Double
0 Double
0
    fv :: P
fv = Double -> Double -> P
P Double
5 Double
5


-- | binary mixture of isotropic 2d normal distribs
gaussMix :: PrimMonad m =>
          Double -> Double -> Double -> Double -> P.Prob m P
gaussMix :: Double -> Double -> Double -> Double -> Prob m P
gaussMix Double
mu1 Double
mu2 Double
sig1 Double
sig2 = do
  Bool
b <- Prob m Bool
forall (m :: * -> *). PrimMonad m => Prob m Bool
coin
  if Bool
b
    then Double -> Double -> Prob m P
forall (m :: * -> *). PrimMonad m => Double -> Double -> Prob m P
isoNormal2d Double
mu1 Double
sig1
    else Double -> Double -> Prob m P
forall (m :: * -> *). PrimMonad m => Double -> Double -> Prob m P
isoNormal2d Double
mu2 Double
sig2

coin :: PrimMonad m => P.Prob m Bool
coin :: Prob m Bool
coin = Double -> Prob m Bool
forall (m :: * -> *). PrimMonad m => Double -> Prob m Bool
P.bernoulli Double
0.5

isoNormal2d :: PrimMonad m => Double -> Double -> P.Prob m P
isoNormal2d :: Double -> Double -> Prob m P
isoNormal2d Double
mu Double
sig = Double -> Double -> P
P (Double -> Double -> P) -> Prob m Double -> Prob m (Double -> P)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Double -> Double -> Prob m Double
forall (m :: * -> *).
PrimMonad m =>
Double -> Double -> Prob m Double
P.normal Double
mu Double
sig Prob m (Double -> P) -> Prob m Double -> Prob m P
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Double -> Double -> Prob m Double
forall (m :: * -> *).
PrimMonad m =>
Double -> Double -> Prob m Double
P.normal Double
mu Double
sig

binDisk :: PrimMonad m => Double -> Double -> P -> P -> P.Prob m P
binDisk :: Double -> Double -> P -> P -> Prob m P
binDisk Double
r0 Double
r1 P
p0 P
p1 = do
  Bool
b <- Prob m Bool
forall (m :: * -> *). PrimMonad m => Prob m Bool
coin
  if Bool
b
    then Double -> P -> Prob m P
forall (m :: * -> *). PrimMonad m => Double -> P -> Prob m P
uniformDisk Double
r0 P
p0
    else Double -> P -> Prob m P
forall (m :: * -> *). PrimMonad m => Double -> P -> Prob m P
uniformDisk Double
r1 P
p1

-- point in a disk of radius r and centered at P
uniformDisk :: PrimMonad m => Double -> P -> P.Prob m P
uniformDisk :: Double -> P -> Prob m P
uniformDisk Double
rmax P
p = do
  Double
r <- (Double, Double) -> Prob m Double
forall (m :: * -> *) a.
(PrimMonad m, Variate a) =>
(a, a) -> Prob m a
P.uniformR (Double
0, Double
rmax)
  Double
aa <- (Double, Double) -> Prob m Double
forall (m :: * -> *) a.
(PrimMonad m, Variate a) =>
(a, a) -> Prob m a
P.uniformR (Double
0, Double
2 Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
forall a. Floating a => a
pi)
  let
    x :: Double
x = Double
r Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double -> Double
forall a. Floating a => a -> a
cos Double
aa
    y :: Double
y = Double
r Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double -> Double
forall a. Floating a => a -> a
sin Double
aa
    p0 :: P
p0 = Double -> Double -> P
P Double
x Double
y
  P -> Prob m P
forall (f :: * -> *) a. Applicative f => a -> f a
pure (P -> Prob m P) -> P -> Prob m P
forall a b. (a -> b) -> a -> b
$ P
p0 P -> P -> P
.+. P
p

buildP :: V.Vector P -> VPTree Double P
buildP :: Vector P -> VPTree Double P
buildP = (P -> P -> Double) -> Double -> Vector P -> VPTree Double P
forall p d a.
(RealFrac p, Floating d, Ord d, Eq a) =>
(a -> a -> d) -> p -> Vector a -> VPTree d a
build P -> P -> Double
distp (Double
1.0 :: Double)