{-# LANGUAGE FlexibleContexts    #-}
{-# LANGUAGE GADTs               #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeApplications    #-}
{-# OPTIONS_GHC -Wno-deprecations #-}
module BlueRipple.Model.PreferenceBayes
  (
    -- * MLE
    cgOptimize
  , cgOptimizeAD
  , invFisher
  , mleCovEigens
  , variances
  , covar
  , correlFromCov

    -- * MCMC
  , Sample
  , Chain
  , runMany
    -- * Re-exports
  , dispf
  , disps
  )
where

import qualified Statistics.Types              as S
import qualified Control.Foldl                 as FL
import           Control.Lens.At                ( IxValue
                                                , Index
                                                , Ixed
                                                )
import           Control.Lens.Indexed           ( FunctorWithIndex )
import           Control.Monad                  ( sequence )

import qualified Numeric.AD                    as AD
import qualified Numeric.LinearAlgebra         as LA
import           Numeric.LinearAlgebra          ( dispf
                                                , disps
                                                )
import           Numeric.MathFunctions.Constants
                                                ( m_ln_sqrt_2_pi )
import qualified Numeric.MCMC                  as MC
import Numeric (log)

import           Math.Gamma                     ( gamma
                                                , Gamma
                                                )
import           System.Random                  ( randomRIO )
import qualified Control.Concurrent            as CC
import qualified Control.Concurrent.MVar       as CC
import qualified Data.Vector.Unboxed           as VU
import qualified Data.Vector.Storable          as VS
import qualified Data.Vector.Generic           as VG
import qualified Data.Vector                   as VB

import qualified Numeric.Optimization.Algorithms.HagerZhang05
                                               as CG

import qualified Numeric.Optimization.Algorithms.HagerZhang05.AD
                                               as CGAD

newtype ObservedVote = ObservedVote { dem :: Int}

data Pair a b = Pair !a !b

binomialNormalParams :: RealFrac a => VB.Vector Int -> VB.Vector a -> (a, a)
binomialNormalParams turnoutCounts demProbs =
  let np = VB.zip turnoutCounts demProbs
      meanVarUpdate :: RealFrac a => Pair a a -> (Int, a) -> Pair a a
      meanVarUpdate (Pair m v) (n, p) =
        let m' = realToFrac n * p in Pair (m + m') (v + m' * (1 - p))
      foldMeanVar = FL.Fold meanVarUpdate (Pair 0 0) id
      Pair m v    = FL.fold foldMeanVar np
  in  (m, v)

logBinomialObservedVote
  :: (RealFrac a, Floating a) => VB.Vector a -> (Int, VB.Vector Int) -> a
logBinomialObservedVote demProbs (demVote, turnoutCounts) =
  let (m, v) = binomialNormalParams turnoutCounts demProbs
  in  negate $ log v + ((realToFrac demVote - m) ^ 2 / (2 * v))

gradLogBinomialObservedVote
  :: (Floating a, RealFrac a)
  => VB.Vector a
  -> (Int, VB.Vector Int)
  -> VB.Vector a
gradLogBinomialObservedVote demProbs (demVote, turnoutCounts) =
  let (m, v) = binomialNormalParams turnoutCounts demProbs
      np     = VB.zip turnoutCounts demProbs
      dv     = fmap (\(n, p) -> let n' = realToFrac n in n' * (1 - 2 * p)) np
      dm     = turnoutCounts
      dmv    = VB.zip dm dv
      a1     = negate (1 / v) -- d (log v)
      a2     = (realToFrac demVote - m)
      a3     = (a2 * a2) / (2 * v * v)
      a4     = (a2 / v)
  in  fmap (\(dm, dv) -> (a1 + a3) * dv + a4 * realToFrac dm) dmv

logBinomialObservedVotes
  :: (Functor f, Foldable f, Floating a, RealFrac a)
  => f (Int, VB.Vector Int)
  -> VB.Vector a
  -> a
logBinomialObservedVotes votesAndTurnout demProbs =
  FL.fold FL.sum $ fmap (logBinomialObservedVote demProbs) votesAndTurnout

gradLogBinomialObservedVotes
  :: (Functor f, Foldable f)
  => f (Int, VB.Vector Int)
  -> VB.Vector Double
  -> VB.Vector Double
gradLogBinomialObservedVotes votesAndTurnout demProbs =
  let
    n           = VG.length demProbs
    indexVector = VG.generate n id
    sumEach =
      sequenceA $ (\n -> FL.premap (VG.! n) (FL.sum @Double)) <$> indexVector
  in
    FL.fold sumEach
      $ fmap (gradLogBinomialObservedVote demProbs) votesAndTurnout

cgOptimize
  :: (Functor f, Foldable f)
  => f (Int, VB.Vector Int)
  -> VB.Vector Double
  -> IO (VS.Vector Double, CG.Result, CG.Statistics)
cgOptimize votesAndVoters guess = do
  let params = CG.defaultParameters { CG.printFinal  = False
                                    , CG.printParams = False
                                    , CG.verbose     = CG.Quiet
                                    }
      grad_tol = 0.0000001
  CG.optimize
    params
    grad_tol
    guess
    (CG.VFunction (negate . logBinomialObservedVotes votesAndVoters))
    (CG.VGradient (fmap negate . gradLogBinomialObservedVotes votesAndVoters))
    Nothing

cgOptimizeAD
  :: (Functor f, Foldable f)
  => f (Int, VB.Vector Int)
  -> VB.Vector Double
  -> IO (VB.Vector Double, CG.Result, CG.Statistics)
cgOptimizeAD votesAndVoters guess = do
  let params = CG.defaultParameters { CG.printFinal  = False
                                    , CG.printParams = False
                                    , CG.verbose     = CG.Quiet
                                    }
      grad_tol = 0.0000001
  CGAD.optimize params
                grad_tol
                guess
                (negate . logBinomialObservedVotes votesAndVoters)

-- I'm not sure about row/column order her. But it doesn't
-- matter since the Hessian is symmetric.
hessianLL
  :: (Foldable f, Functor f)
  => f (Int, VB.Vector Int)
  -> VB.Vector Double
  -> LA.Matrix Double
hessianLL votesAndVoters x =
  LA.fromRows $ fmap VS.convert $ VB.toList $ AD.hessian
    (negate . logBinomialObservedVotes votesAndVoters)
    x

invFisher
  :: (Functor f, Foldable f)
  => f (Int, VB.Vector Int)
  -> VB.Vector Double
  -> LA.Matrix Double
invFisher votesAndVoters x = LA.inv $ hessianLL votesAndVoters x

variances
  :: (Functor f, Foldable f)
  => f (Int, VB.Vector Int)
  -> VB.Vector Double
  -> LA.Vector Double
variances votesAndVoters x = LA.takeDiag $ invFisher votesAndVoters x

covar
  :: (Functor f, Foldable f)
  => f (Int, VB.Vector Int)
  -> VB.Vector Double
  -> LA.Matrix Double
covar = invFisher

correlFromCov :: LA.Matrix Double -> LA.Matrix Double
correlFromCov cov =
  let vars = LA.takeDiag cov
      mISD = LA.diag $ LA.cmap (\x -> 1 / sqrt x) vars
  in  mISD LA.<> cov LA.<> mISD

correl
  :: (Functor f, Foldable f)
  => f (Int, VB.Vector Int)
  -> VB.Vector Double
  -> LA.Matrix Double
correl votesAndVoters x =
  let cov          = invFisher votesAndVoters x
      sds          = LA.cmap sqrt $ LA.takeDiag cov
      (rows, cols) = LA.size cov
      rho n m = cov `LA.atIndex` (n, m) / ((sds LA.! n) * (sds LA.! m))
  in  LA.assoc
        (rows, cols)
        0
        [ ((n, m), rho n m) | n <- [0 .. (rows - 1)], m <- [0 .. (cols - 1)] ]

mleCovEigens
  :: (Functor f, Foldable f)
  => f (Int, VB.Vector Int)
  -> VB.Vector Double
  -> (LA.Vector Double, LA.Matrix Double)
mleCovEigens votesAndVoters x =
  LA.eigSH $ LA.trustSym $ invFisher votesAndVoters x

betaDist :: (Floating a, RealFrac a, Gamma a) => a -> a -> a -> a
betaDist alpha beta x =
  let b = gamma (alpha + beta) / (gamma alpha + gamma beta)
  in  if (x >= 0) && (x <= 1)
        then b * (x ** (alpha - 1)) * ((1 - x) ** (beta - 1))
        else 0

logBetaDist :: (Floating a, RealFrac a, Gamma a) => a -> a -> a -> a
logBetaDist alpha beta x =
  let lb = log $ gamma (alpha + beta) / (gamma alpha + gamma beta)
  in  lb + (alpha - 1) * log x + (beta - 1) * log (1 - x)

betaPrior :: (Floating a, RealFrac a, Gamma a) => a -> a -> VB.Vector a -> a
betaPrior a b xs = FL.fold FL.product $ fmap (betaDist a b) xs

{-
--logBetaPrior :: Double -> Double -> [Double] -> Double
--logBetaPrior = FL.fold FL.product $ fmap (logBetabDist a b) xs

f :: [(Int, [Int])] -> [Double] -> Double
f votesAndTurnout demProbs =
  let v = demProbs
  in  (exp $ logProbObservedVotes votesAndTurnout v) * (betaPrior 2 2 v)
-}

logBinomialWithPrior
  :: (Functor f, Foldable f, Floating a, RealFrac a, Gamma a)
  => f (Int, VB.Vector Int)
  -> VB.Vector a
  -> a
logBinomialWithPrior votesAndTurnout demProbs =
  let x = demProbs
  in  logBinomialObservedVotes votesAndTurnout x + log (betaPrior 2 2 x)

type Sample = VB.Vector Double
type Chain = [Sample]

-- TODO: withSystemRandom is deprecated in mwc-random but comes from "declarative" so would be a pain to fix.
runMCMC
  :: (Functor f, Foldable f)
  => f (Int, VB.Vector Int)
  -> Int
  -> Sample
  -> IO Chain
runMCMC votesAndTurnout numIters start =
  fmap (fmap MC.chainPosition) . MC.withSystemRandom . MC.asGenIO $ MC.chain
    numIters
    start
    (MC.frequency [(1, MC.hamiltonian 0.0001 10), (1, MC.hamiltonian 0.00001 5)]
    )
    (MC.Target (logBinomialWithPrior votesAndTurnout)
               (Just $ gradLogBinomialObservedVotes votesAndTurnout)
    )

-- launch each action on its own thread, writing the result to an MVar.
-- as each MVar is written to,
-- place each result in the appropriate place in the structure.
-- return the new structure when all results are available.
sequenceConcurrently :: Traversable t => t (IO a) -> IO (t a)
sequenceConcurrently actions = do
  let f :: IO a -> IO (CC.MVar a, CC.ThreadId)
      f action = do
        mvar     <- newEmptyMVar
        threadId <- CC.forkIO (action >>= putMVar mvar)
        return (mvar, threadId)
      g :: (CC.MVar a, CC.ThreadId) -> IO a
      g (mvar, _) = takeMVar mvar
  forked <- traverse f actions -- IO (t (MVar a, ThreadId))
  traverse g forked

runMany
  :: (Functor f, Foldable f)
  => f (Int, VB.Vector Int)
  -> Int
  -> Int
  -> Int
  -> Int
  -> IO [Chain]
runMany votesAndTurnout nParams nChains nSamplesPerChain nBurnPerChain = do
  let
    randomStart :: Int -> IO Sample
    randomStart n = VG.fromList <$> replicateM n (randomRIO (0, 1))
    randomStarts :: Int -> Int -> IO [Sample]
    randomStarts n m = replicateM m (randomStart n)
    doEach =
      fmap (drop nBurnPerChain) . runMCMC votesAndTurnout nSamplesPerChain
  starts <- randomStarts nParams nChains
  sequenceConcurrently $ fmap doEach starts

{-
data Normal = Normal { m :: Double, v :: Double } deriving (Show)

instance Semigroup Normal where
  (Normal m1 v1) <> (Normal m2 v2) = Normal (m1 + m2) (v1 + v2)

logNormalObservedVote :: [Normal] -> (Int, [Int]) -> Double
logNormalObservedVote demProbs (demVote, turnoutCounts) =
  let Normal m' v' = mconcat $ fmap (\(t,Normal m v) -> Normal (t*m) (v*m))$ zip turnoutCounts demProbs
  in  negate $ log v'  + ((realToFrac demVote - m')) ^ 2 / (2 * v'))

gradLogNormalObservedVote :: [Normal] -> (Int, [Int]) -> [Normal]
gradLogNormalObservedVote demProbs (demVote, turnoutCounts) =
  let Normal m' v' = mconcat $ fmap (\(t,Normal m v) -> Normal (t*m) (v*m))$ zip turnoutCounts demProbs
      np     = zip turnoutCounts demProbs
      dv     = fmap (\(n, p) -> let n' = realToFrac n in n' * (1 - 2 * p)) np
      dm     = turnoutCounts
      dmv    = zip dm dv
      a1     = negate (1 / v) -- d (log v)
      a2     = (realToFrac demVote - m)
      a3     = (a2 * a2) / (2 * v * v)
      a4     = (a2 / v)
  in  fmap (\(dm, dv) -> (a1 + a3) * dv + a4 * realToFrac dm) dmv

logBinomialObservedVotes :: [(Int, [Int])] -> [Double] -> Double
logBinomialObservedVotes votesAndTurnout demProbs =
  FL.fold FL.sum $ fmap (logBinomialObservedVote demProbs) votesAndTurnout

gradLogBinomialObservedVotes :: [(Int, [Int])] -> [Double] -> [Double]
gradLogBinomialObservedVotes votesAndTurnout demProbs =
  let n = length demProbs
      sumEach =
        sequenceA
          $ fmap (\n -> FL.premap (!! n) (FL.sum @Double))
          $ [0 .. (n - 1)]
  in  FL.fold sumEach $ fmap (gradLogBinomialObservedVote demProbs) votesAndTurnout

-}
{-
--gradLogProbObservedVotes2 :: [(Int, [Int])] -> [Double] -> [Double]
gradLogProbObservedVotes2 votesAndTurnout =
  grad (logProbObservedVotes votesAndTurnout)
-}

{-
data VPSummary = VPSummary { mean :: Double, stDev :: Double, rHat :: Double } deriving (Show)

summarize :: (Sample -> Double) -> [Chain] -> VPSummary -- FL.Fold Chain VPSummary
summarize integrandF chains =
  -- build up the fold
  let n = (length $ head chains) `div` 2
      splitChains :: [Chain]
      splitChains =
        concat $ fmap (\c -> let (a, b) = splitAt n c in [a, b]) chains
      meanVarF = (,) <$> FL.mean <*> FL.variance
      meanVars = fmap (FL.fold (FL.premap integrandF meanVarF)) splitChains
      bVar     = FL.fold (FL.premap fst FL.variance) meanVars -- variance of the means
      wVar     = FL.fold (FL.premap snd FL.mean) meanVars -- mean of the variances
      v        = (realToFrac n / realToFrac (n + 1)) * wVar + bVar
      rhat     = sqrt (v / wVar)
  in  VPSummary (FL.fold (FL.premap fst FL.mean) meanVars)
                (FL.fold (FL.premap snd FL.mean) meanVars)
                rhat

-- . foldTo (meanOfVariances, varianceOfMeans) . fmap (compute meanVar . fmap integrandF) . splitChains
-}


{-# LANGUAGE ConstraintKinds       #-}
{-# LANGUAGE FlexibleContexts      #-}
{-# LANGUAGE FlexibleInstances     #-}
{-# LANGUAGE GADTs                 #-}
{-# LANGUAGE KindSignatures        #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE PatternSynonyms       #-}
{-# LANGUAGE RankNTypes            #-}
{-# LANGUAGE TupleSections         #-}
{-# LANGUAGE TypeFamilies          #-}
{-# LANGUAGE TypeInType            #-}
{-# LANGUAGE TypeOperators         #-}
{-# LANGUAGE UndecidableInstances  #-}
{-# LANGUAGE ViewPatterns          #-}

module Backprop.Learn.Model.Types (
  -- * Model type
    ModelFunc, ModelFuncStoch
  , Model(..)
  , modelD
  -- * Specialized Models
  -- ** Stateless
  , ModelFuncStateless, ModelFuncStochStateless
  , ModelStateless, pattern ModelStateless, runLearnStateless, runLearnStochStateless
  , modelStatelessD
  -- ** Stateless and Parameterless
  , BFunc, BFuncStoch
  , Func, pattern Func, runFunc, runFuncStoch
  , funcD
  -- * Manipulating models as functions
  , ModelFuncM, withModelFunc0, withModelFunc, withModelFunc2
  -- * Utility
  , PMaybe(..), TMaybe, fromPJust, Learnables, Learnable
  , Regularize(..), Initialize(..)
  -- , HKD
  ) where

import           Backprop.Learn.Initialize
import           Backprop.Learn.Regularize
import           Control.Category
import           Control.DeepSeq
import           Control.Monad.Primitive
import           Control.Monad.Trans.Reader
import           Data.Functor.Identity
import           Data.Kind
import           Data.Type.Functor.Product
import           Data.Type.Tuple
import           Data.Vinyl
import           GHC.TypeNats
import           Numeric.Backprop
import           Numeric.Opto.Update
import           Prelude hiding               ((.))
import qualified Data.Binary                  as Bi
import qualified Numeric.LinearAlgebra.Static as H
import qualified System.Random.MWC            as MWC

type ModelFunc p s a b = forall z. Reifies z W
    => PMaybe (BVar z) p
    -> BVar z a
    -> PMaybe (BVar z) s
    -> (BVar z b, PMaybe (BVar z) s)

type ModelFuncStoch p s a b = forall m z. (PrimMonad m, Reifies z W)
    => MWC.Gen (PrimState m)
    -> PMaybe (BVar z) p
    -> BVar z a
    -> PMaybe (BVar z) s
    -> m (BVar z b, PMaybe (BVar z) s)

-- | General parameterized model with potential state
data Model :: Maybe Type -> Maybe Type -> Type -> Type -> Type where
    Model :: { runLearn      :: ModelFunc      p s a b
             , runLearnStoch :: ModelFuncStoch p s a b
             } -> Model p s a b

-- | Construct a deterministic model, with no stochastic component.
modelD :: ModelFunc p s a b -> Model p s a b
modelD f = Model { runLearn      = f
                 , runLearnStoch = \_ p x -> pure . f p x
                 }

type ModelFuncStateless p a b = forall z. Reifies z W
    => PMaybe (BVar z) p
    -> BVar z a
    -> BVar z b

type ModelFuncStochStateless p a b = forall m z. (PrimMonad m, Reifies z W)
    => MWC.Gen (PrimState m)
    -> PMaybe (BVar z) p
    -> BVar z a
    -> m (BVar z b)

newtype MFS  p a b = MFS  (ModelFuncStateless      p a b)
newtype MFSS p a b = MFSS (ModelFuncStochStateless p a b)

-- | Parameterized model with no state
type ModelStateless p = Model p 'Nothing

-- | Construct a 'ModelStateless'.
pattern ModelStateless
    :: ModelFuncStateless p a b
    -> ModelFuncStochStateless p a b
    -> ModelStateless p a b
pattern ModelStateless { runLearnStateless, runLearnStochStateless } <-
    (\m -> (_runLearnStateless m, _runLearnStochStateless m)
       -> (MFS runLearnStateless, MFSS runLearnStochStateless)
    )
  where
    ModelStateless l ls = Model
      { runLearn      = \  p x s -> (l p x, s)
      , runLearnStoch = \g p x s -> (, s) <$> ls g p x
      }

-- | Construct a deterministic stateless model, with no stochastic
-- component.
modelStatelessD :: ModelFuncStateless p a b -> ModelStateless p a b
modelStatelessD f = ModelStateless
    { runLearnStateless      = f
    , runLearnStochStateless = \_ p -> pure . f p
    }

_runLearnStateless :: Model p 'Nothing a b -> MFS p a b
_runLearnStateless md = MFS $ \p x -> fst $ runLearn md p x PNothing

_runLearnStochStateless :: Model p 'Nothing a b -> MFSS p a b
_runLearnStochStateless md = MFSS $ \g p x -> fst <$> runLearnStoch md g p x PNothing

type BFunc a b = forall z. Reifies z W
    => BVar z a
    -> BVar z b
type BFuncStoch a b = forall m z. (PrimMonad m, Reifies z W)
    => MWC.Gen (PrimState m)
    -> BVar z a
    -> m (BVar z b)

newtype BF  a b = BF  (BFunc      a b)
newtype BFS a b = BFS (BFuncStoch a b)

type Func = Model 'Nothing 'Nothing

-- | Unparameterized model with no state
pattern Func :: BFunc a b -> BFuncStoch a b -> Func a b
pattern Func { runFunc, runFuncStoch } <-
    (\m -> (_runFunc m, _runFuncStoch m)
       -> (BF runFunc, BFS runFuncStoch)
    )
  where
    Func r rs = ModelStateless
      { runLearnStateless      = const r
      , runLearnStochStateless = const . rs
      }

-- | Construct an deterministic unparameterized stateless model, with no
-- stochastic component.
funcD :: BFunc a b -> Func a b
funcD f = Func { runFunc      = f
               , runFuncStoch = const (pure . f)
               }

_runFunc :: Model 'Nothing 'Nothing a b -> BF a b
_runFunc md = BF $ runLearnStateless md PNothing

_runFuncStoch :: Model 'Nothing 'Nothing a b -> BFS a b
_runFuncStoch md = BFS $ flip (runLearnStochStateless md) PNothing

-- | Share parameter and sequence state
instance Category (Model p s) where
    id    = modelD $ \_ x s -> (x, s)
    f . g = Model
      { runLearn      = \p x s0 -> let (y, s1) = runLearn g p x s0
                                   in  runLearn f p y s1
      , runLearnStoch = \gen p x s0 -> do
          (y, s1) <- runLearnStoch g gen p x s0
          runLearnStoch f gen p y s1
      }

type ModelFuncM m p s a b = forall z. Reifies z W
    => PMaybe (BVar z) p
    -> BVar z a
    -> PMaybe (BVar z) s
    -> m (BVar z b, PMaybe (BVar z) s)

withModelFunc0
    :: (forall m. Monad m => ModelFuncM m p s a b)
    -> Model p s a b
withModelFunc0 f = Model
    { runLearn      = \p x -> runIdentity . f p x
    , runLearnStoch = \g p x -> flip runReaderT g . f p x
    }

withModelFunc
    :: (forall m. Monad m => ModelFuncM m p s a b -> ModelFuncM m q t c d)
    -> Model p s a b
    -> Model q t c d
withModelFunc f m = Model
    { runLearn      = \p x -> runIdentity . f (runLearnModelFuncM m) p x
    , runLearnStoch = \g p x -> flip runReaderT g . f (runLearnStochModelFuncM m) p x
    }

withModelFunc2
    :: (forall m. Monad m => ModelFuncM m p s a b -> ModelFuncM m q t c d -> ModelFuncM m r u e f)
    -> Model p s a b
    -> Model q t c d
    -> Model r u e f
withModelFunc2 f m n = Model
    { runLearn      = \p x ->
            runIdentity
          . f (runLearnModelFuncM m) (runLearnModelFuncM n) p x
    , runLearnStoch = \g p x ->
            flip runReaderT g
          . f (runLearnStochModelFuncM m) (runLearnStochModelFuncM n) p x
    }

runLearnModelFuncM
    :: Model p s a b
    -> ModelFuncM Identity p s a b
runLearnModelFuncM m p x = Identity . runLearn m p x

runLearnStochModelFuncM
    :: PrimMonad m
    => Model p s a b
    -> ModelFuncM (ReaderT (MWC.Gen (PrimState m)) m) p s a b
runLearnStochModelFuncM m p x s = ReaderT $ \g -> runLearnStoch m g p x s

-- | Combination of common constraints for type-level lists.
type Learnables as = ( RecApplicative as
                     , ReifyConstraint Backprop TF as
                     , RMap as
                     , RApply as
                     )

-- | Helpful utility class for ensuring that a parameter or state type is
-- Learnable
class ( Initialize a, Regularize a, Bi.Binary a, NFData a, LinearInPlace m Double a )
        => Learnable m a

instance (Initialize a, Regularize a, Bi.Binary a, NFData a, LinearInPlace m Double a)
    => Learnable m (TF a)
instance (Initialize a, Initialize b, Regularize a, Regularize b, Bi.Binary a, Bi.Binary b, NFData a, NFData b, LinearInPlace m Double a, LinearInPlace m Double b)
    => Learnable m (a :# b)

instance (RPureConstrained Initialize as, RPureConstrained Regularize as, RPureConstrained Bi.Binary as, RPureConstrained NFData as, ReifyConstraint Backprop TF as, RMap as, RApply as, RFoldMap as, RecordToList as, LinearInPlace m Double (Rec TF as))
        => Learnable m (Rec TF as)
instance (KnownNat n, PrimMonad m) => Learnable m (H.R n)
instance (KnownNat n, KnownNat o, PrimMonad m) => Learnable m (H.L n o)


{-# LANGUAGE MultiParamTypeClasses #-}
-- | Defines intersection tests for a varity of geometry.
module Game.Physics.Collision ( Point
                              , point2D
                              , AABB
                              , aabb2D
                              , OBB
                              , obb2DCorners
                              , Intersectable(..)
                              ) where

import Prelewd hiding (toList, (<>))

import Impure

import Data.Tuple
import Storage.List
import Numeric.LinearAlgebra

-- | "Axis-aligned bounding box."
--
--   Use the smart constructor 'aabb2D' to build AABBs.
--
--   * Implementation details:
--
--   An AABB is an Nx2 matrix. X[0] is the lower limit on the x-th dimension
--   and X[1] is the upper limit on the x-th dimension.
--
--   The number of rows is the number of dimensions of the box. The number of
--   columns is always 2. This is enforced by the smart constructor(s).
newtype AABB = AABB (Matrix Double)

-- | A point in the physics universe is just a list of the object's
--   coordinates. The length of the vector is the number of dimensions its's
--   defined for.
--
--   Use the smart constructor 'point2D' to build points.
newtype Point = Point (Vector Double)

-- | "Oriented Bounding Box." A bounding box with any arbitrary orientation.
--
--   Use smart constructor 'obb2D' to build OBBs.
--
--   * Implementation Notes:
--
--   An oriented bounding box consists of a rotation/translation matrix,
--   and a vector of n-dimensional extents. For example, a 2x2 bounding box
--   would have the extents [1,1].
--
--   We take a simplified approach in 2D - we just use the four corners of the
--   box, specified in a clockwise orientation. This turns our point-OBB
--   intersection test into 8 vector subtractions, and 4 dot products. Needles
--   to say, that's ridiculously cheap.
data OBB = OBB2D {-# UNPACK #-} !Point
                 {-# UNPACK #-} !Point
                 {-# UNPACK #-} !Point
                 {-# UNPACK #-} !Point
         | OBB !(Matrix Double) -- The transformation matrix.
               !(Vector Double) -- The extent matrix.

-- | Builds a 2D point out of a tuple.
point2D :: Integral a => (a, a) -> Point
point2D (x, y) = Point . fromList $ map fromIntegral [x, y]

-- | Constructs a 2-dimensional axis-aligned bounding box.
aabb2D :: (Double, Double) -> (Double, Double) -> AABB
aabb2D (xa, xb) (ya, yb) = AABB $ (2 >< 2) [ min xa xb, max xa xb
                                           , min ya yb, max ya yb
                                           ]

-- | Constructs a 2D oriented bounding box from a list of verticies, specified
--   in a clockwise direction.
--
--   Seriously - clockwise. Don't fuck this up. It's _very_ important. No
--   checks are done.
--
--   Clockwise.
obb2DCorners :: Point -> Point -> Point -> Point -> OBB
obb2DCorners = OBB2D

-- | The type class for our 'intersects' method. This lets us perform double
--   dispatch on the types of both of intersects' arguments.
class Intersectable a b where
    -- | We define intersects for as many pairs of geometry as possibly,
    --   allowing us to do double-dispatch collision detection.
    intersects :: a -> b -> Bool

between :: Ord a => (a, a) -> a -> Bool
between (l, r) x = l < x && x < r
{-# INLINE between #-}

instance Intersectable AABB AABB where
    -- | Two AABBs intersect if any of their edges lie between the other's
    --   edges.
    intersects (AABB ba) (AABB bb) = any testDim . zip (toRows ba) $ toRows bb
        where
            testDim :: (Vector Double, Vector Double) -> Bool
            testDim (a, b) = let a'@(lx, hx) = (a @> 0, a @> 1)
                                 b'@(ly, hy) = (b @> 0, b @> 1)
                              in  between a' ly || between a' hy
                               || between b' lx || between b' hx
    {-# INLINE intersects #-}

instance Intersectable AABB Point where
    intersects (AABB bs) (Point p) = all inRange . zip (toRows bs) $ toList p
        where
            inRange :: (Vector Double, Double) -> Bool
            inRange (bnds, x) = between (bnds @> 0, bnds @> 1) x
    {-# INLINE intersects #-}

instance Intersectable Point AABB where
    intersects x y = intersects y x
    {-# INLINE intersects #-}

obbBadInput :: a
obbBadInput = error "Invalid dimensions in the OOB intersection vectors being tested."

instance Intersectable OBB Point where
    intersects (OBB2D (Point w) (Point x) (Point y) (Point z)) (Point p)
                   | any (/= 2) [ dp, dw, dx, dy, dz ] = obbBadInput
                   | otherwise = all (>= 0) [ (p - w) <.> (x - w)
                                           , (p - x) <.> (y - x)
                                           , (p - y) <.> (z - y)
                                           , (p - z) <.> (w - z)
                                           ]
        where
            dp = dim p
            dw = dim w
            dx = dim x
            dy = dim y
            dz = dim z
    intersects (OBB lm exts) (Point p) | allEqual [ dlmx, dlmy, dp ] = obbBadInput
                                       | otherwise = intersects aabb transformed
        where
            transformed :: Point
            transformed = Point $ inv lm <> p
            aabb :: AABB
            aabb = AABB . fromRows . map (\e -> fromList [-e, e]) $ toList exts

            (dlmx, dlmy) = (cols lm, rows lm)
            dp           = dim p

-- | Are all elements in the list equal?
allEqual :: Eq a => [a] -> Bool
allEqual []     = True
allEqual (x:xs) = fst $ foldl' (\(b, v) v' -> (b && v == v', v)) (True, x) xs
-- we use foldl', since in most cases, the elements _will_ be equal. If the
-- converse were true, I would prefer foldr since && is lazy.

instance Intersectable Point OBB where
    intersects x y = intersects y x
    {-# INLINE intersects #-}


{-# LANGUAGE BangPatterns        #-}
{-# LANGUAGE CPP                 #-}
{-# LANGUAGE DataKinds           #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TupleSections       #-}
{-# LANGUAGE TypeFamilies        #-}
{-# LANGUAGE TypeOperators       #-}

import           Control.Monad.Random
import           Control.Monad.Trans.Except

import qualified Data.ByteString              as B
import           Data.Char                    (isUpper, toLower, toUpper)
import           Data.List                    (foldl')
import qualified Data.Map                     as M
import           Data.Maybe                   (fromMaybe)
import           Data.Serialize
import           Data.Singletons.Prelude
import           Data.Vector                  (Vector)
import qualified Data.Vector                  as V
import           GHC.TypeLits
import           Numeric.LinearAlgebra.Static (konst)
import           Options.Applicative

import           Grenade
import           Grenade.Recurrent
import           Grenade.Utils.OneHot

import           System.IO.Unsafe             (unsafeInterleaveIO)

-- The defininition for our natural language recurrent network.
-- This network is able to learn and generate simple words in
-- about an hour.
--
-- Grab the input from
-- https://raw.githubusercontent.com/karpathy/char-rnn/master/data/tinyshakespeare/input.txt
--
-- This is a first class recurrent net.
--
-- The F and R types are tagging types to ensure that the runner and
-- creation function know how to treat the layers.
--
-- As an example, here's a short sequence generated.
--
-- > KING RICHARD III:
-- > And as the heaven her his words, we the son, I show sand stape but the lament to shall were the sons with a strend

type FF = FeedForward
type R = Recurrent

-- The definition of our network
type Shakespeare = RecurrentNetwork '[ R (LSTM 40 80), R (LSTM 80 40), FF (FullyConnected 40 40), FF Logit]
                                    '[ 'D1 40, 'D1 80, 'D1 40, 'D1 40, 'D1 40 ]

-- The definition of the "sideways" input, which the network is fed recurrently.
type Shakespearian = RecurrentInputs  '[ R (LSTM 40 80), R (LSTM 80 40), FF (FullyConnected 40 40), FF Logit]

randomNet :: IO Shakespeare
randomNet = randomRecurrent

-- | Load the data files and prepare a map of characters to a compressed int representation.
loadShakespeare :: FilePath -> ExceptT String IO (Vector Int, M.Map Char Int, Vector Char)
loadShakespeare path = do
  contents     <- lift $ readFile path
  let annotated = annotateCapitals contents
  (m,cs)       <- ExceptT . return $ hotMap (Proxy :: Proxy 40) annotated
  hot          <- ExceptT . return . note "Couldn't generate hot values" $ traverse (`M.lookup` m) annotated
  return (V.fromList hot, m, cs)

trainSlice :: Optimizer o -> Shakespeare -> Shakespearian -> Vector Int -> Int -> Int -> (Shakespeare, Shakespearian)
trainSlice !opt !net !recIns input offset size =
  let e = fmap (x . oneHot) . V.toList $ V.slice offset size input
  in case reverse e of
    (o : l : xs) ->
      let examples = reverse $ (l, Just o) : ((,Nothing) <$> xs)
      in  trainRecurrent opt net recIns examples
    _ -> error "Not enough input"
    where
      x = fromMaybe (error "Hot variable didn't fit.")

runShakespeare :: ShakespeareOpts -> ExceptT String IO ()
runShakespeare opts = do
  (shakespeare, oneHotMap, oneHotDictionary) <- loadShakespeare $ trainingFile opts
  (net0, i0) <- lift $
    case loadPath opts of
      Just loadFile -> netLoad loadFile
      Nothing       -> (,0) <$> randomNet

  (trained, bestInput) <- lift $ foldM (\(!net, !io) size -> do
    xs <- take (iterations opts `div` 10) <$> getRandomRs (0, length shakespeare - size - 1)
    let (!trained, !bestInput) = foldl' (\(!n, !i) offset -> trainSlice (rate opts) n i shakespeare offset size) (net, io) xs
    results <- take 1000 <$> generateParagraph trained bestInput (temperature opts) oneHotMap oneHotDictionary ( S1D $ konst 0)
    putStrLn ("TRAINING STEP WITH SIZE: " ++ show size)
    putStrLn (unAnnotateCapitals results)
    return (trained, bestInput)
    ) (net0, i0) $ replicate 10 (sequenceSize opts)

  case savePath opts of
    Just saveFile -> lift . B.writeFile saveFile $ runPut (put trained >> put bestInput)
    Nothing       -> return ()

generateParagraph :: forall layers shapes n a. (Last shapes ~ 'D1 n, Head shapes ~ 'D1 n, KnownNat n, Ord a)
  => RecurrentNetwork layers shapes
  -> RecurrentInputs layers
  -> RealNum
  -> M.Map a Int
  -> Vector a
  -> S ('D1 n)
  -> IO [a]
generateParagraph n s temp hotmap hotdict =
  go s
    where
  go x y =
    do let (_, ns, o) = runRecurrent n x y
       un            <- sample temp hotdict o
       Just re       <- return $ makeHot hotmap un
       rest          <- unsafeInterleaveIO $ go ns re
       return (un : rest)

data ShakespeareOpts = ShakespeareOpts {
    trainingFile :: FilePath
  , iterations   :: Int
  , rate         :: Optimizer 'SGD
  , sequenceSize :: Int
  , temperature  :: RealNum
  , loadPath     :: Maybe FilePath
  , savePath     :: Maybe FilePath
  }

shakespeare' :: Parser ShakespeareOpts
shakespeare' = ShakespeareOpts <$> argument str (metavar "TRAIN")
                               <*> option auto (long "examples" <> short 'e' <> value 1000000)
                               <*> (OptSGD
                                    <$> option auto (long "train_rate" <> short 'r' <> value 0.01)
                                    <*> option auto (long "momentum" <> value 0.95)
                                    <*> option auto (long "l2" <> value 0.000001)
                                    )
                               <*> option auto (long "sequence-length" <> short 's' <> value 50)
                               <*> option auto (long "temperature" <> short 't' <> value 0.4)
                               <*> optional (strOption (long "load"))
                               <*> optional (strOption (long "save"))

main :: IO ()
main = do
    shopts <- execParser (info (shakespeare' <**> helper) idm)
    res <- runExceptT $ runShakespeare shopts
    case res of
      Right () -> pure ()
      Left err -> putStrLn err

netLoad :: FilePath -> IO (Shakespeare, Shakespearian)
netLoad modelPath = do
  modelData <- B.readFile modelPath
  either fail return $ runGet get modelData

-- Replace capitals with an annotation and the lower case letter
-- http://fastml.com/one-weird-trick-for-training-char-rnns/
annotateCapitals :: String -> String
annotateCapitals (x : rest)
    | isUpper x
    = '^' : toLower x : annotateCapitals rest
    | otherwise
    = x : annotateCapitals rest
annotateCapitals []
    = []

unAnnotateCapitals :: String -> String
unAnnotateCapitals ('^' : x : rest)
    = toUpper x : unAnnotateCapitals rest
unAnnotateCapitals (x : rest)
    =  x : unAnnotateCapitals rest
unAnnotateCapitals []
    = []

-- | Tag the 'Nothing' value of a 'Maybe'
note :: a -> Maybe b -> Either a b
note a = maybe (Left a) Right


{-# language FlexibleInstances      #-}
{-# language KindSignatures         #-}
{-# language FunctionalDependencies #-}
{-# language LambdaCase             #-}
{-# language MultiParamTypeClasses  #-}
{-# language Rank2Types             #-}
{-# language TupleSections          #-}
module Main where

import Control.Monad.Reader
import Data.Complex
import Data.Functor.Const
import Data.Functor.Contravariant
import Data.Functor.Identity
import Data.Monoid                 (First(..))
import Data.Profunctor
import Data.Profunctor.Unsafe      ((#.), (.#))
import Data.Tagged
import Data.Void                   (absurd, vacuous)

type Equality s t a b
  = forall (p :: * -> * -> *) f.
    p a (f b) -> p s (f t)

type Lens      s t a b
  = forall   f. Functor f
  => (a -> f b) -> s -> f t

type Traversal s t a b
  = forall   f. Applicative f
  => (a -> f b) -> s -> f t

type Iso       s t a b
  = forall p f. (Functor f, Profunctor p)
  => p a (f b)  -> p s (f t)

type Prism     s t a b
  = forall p f. (Applicative f, Choice p)
  => p a (f b)  -> p s (f t)

type Getter    s   a
  = forall   f. (Functor f, Contravariant f)
  => (a -> f a) -> s -> f s

type Fold      s   a
  = forall   f. (Applicative f, Contravariant f)
  => (a -> f a) -> s -> f s

type Optic p f s t a b
  = p a (f b)  -> p s (f t)
type Optic' p f t b
  = p b (f b)  -> p t (f t)

type Review t b
  = Tagged b (Identity b) -> Tagged t (Identity t)
-- in other words,
type Review' t b
  = Optic' Tagged Identity t b

-- Traversal:
--
-- We strengthen `Functor` to `Applicative`. Why?
--
-- Clue: note that
--
-- traverse :: Applicative f => (a -> f b) -> t a -> f (t b)
--           ~ Traversable f => Traversal (f a) (f b) a b
--
-- So, instead we can ask why does `traverse` require `Applicative`?
--
-- The instance for lists is informative:
--
-- instance Traversable [] where
--   traverse _ []     = pure []
--   traverse f (x:xs) = (:) <$> f x <*> traverse f xs
--
-- Every lens is a valid Traversal. Witness:

lensToTraversal :: Lens s t a b -> Traversal s t a b
lensToTraversal = id

--
-- Iso:
--
-- We generalize `->` to `(Profunctor p =>) p`. Why?
--
-- One answer is that using `Exchange` (a `Profunctor`) allows us to reverse
-- the `Iso` and extract its parts.
--
-- Prism:
--
-- Here `f` is an `Applicative` and `p` is a `Choice`. Why?
--
-- First of all:
-- * Every `Prism`: Applicative f, Choice p
--   is a valid          v            v
--     `Traversal`: Applicative f
--
--   witness:

prismToTraversal :: Prism s t a b -> Traversal s t a b
prismToTraversal = id

--
--   Here `Prism` adds `Choice p`, meaning TODO
--
-- * Every `Iso`:   Functor f,     Profunctor p
--   is a valid           v             v
--         `Prism`: Applicative f, Choice p
--
--   witness:

isoToPrism :: Iso s t a b -> Prism s t a b
isoToPrism = id

-- Iso's are everything actually:

isoToLens :: Iso s t a b -> Lens s t a b
isoToLens = id

isoToTraversal :: Iso s t a b -> Traversal s t a b
isoToTraversal = id

isoToGetter :: Iso s s a a -> Getter s a
isoToGetter = id

isoToFold :: Iso s s a a -> Fold s a
isoToFold = id

isoToReview :: Iso s s a a -> Review s a
isoToReview = id

-- Except an `Equality`, that is. An `Equality` is an `Iso`:

equalityToIso :: Equality s t a b -> Iso s t a b
equalityToIso = id

-- An Equality is a witness that a ~ s and b ~ t.
--
--
--   Here `Prism` strenthens the constraints:
--   * `Functor` to `Applicative` because it doesn't touch exactly one position
--   * `Profunctor` to `Choice` because TODO
--
-- Getter:
--
-- A getter describes how to retrieve a single value
--
-- > for f to be both Functor and Contravariant implies `f a` doesn't contain
-- > any `a`s at all!
--
-- citation: https://www.reddit.com/r/haskell/comments/5vb6x1/how_do_i_learn_lensinternals/de0uz1v
-- Witness:

coerceGetterF :: (Functor f, Contravariant f) => f a -> f b
coerceGetterF = vacuous . contramap absurd

-- Fold:
--
-- A `Fold` describes how to retrieve multiple values
--
-- Note how conspicuously `Fold` and `Getter` are:
--
-- TODO: understand
-- "A `Getter` is a legal `Fold` that just ignores the supplied `Monoid`
--
-- * Every `Getter`: Functor f,   Contravariant f
--   is a valid          v              v
--           `Fold`: Applicative f, Contravariant f
--
-- Witness:

getterToFold :: Getter s a -> Fold s a
getterToFold = id

-- Like `Getter`, note that the functor used in the `Fold` can't contain any
-- values:

coerceFoldF :: (Applicative f, Contravariant f) => f a -> f b
coerceFoldF = vacuous . contramap absurd

-- A "is a limited form of a `Prism` that can only be used for `re` operations.
-- Witness:

prismToReview :: Prism' t b -> Review t b
prismToReview = id

type Simple f s a  = f s s a a
type Equality' s a = Simple Equality s a
type Lens' s a     = Simple Lens s a
type Iso'  s a     = Simple Iso  s a
type Prism' s a    = Simple Prism s a
type Getting r s a = (a -> Const r a) -> s -> Const r s

iso :: (s -> a) -> (b -> t) -> Iso s t a b
iso sa bt = dimap sa (fmap bt)

-- used to provide access to the two parts of an iso
data Exchange a b s t = Exchange (s -> a) (b -> t)

instance Functor (Exchange a b s) where
  fmap g' (Exchange f g) = Exchange f (g' . g)

instance Profunctor (Exchange a b) where
  dimap f' g' (Exchange f g) = Exchange (f . f') (g' . g)

withIso :: Iso s t a b -> ((s -> a) -> (b -> t) -> r) -> r
withIso ai k = case ai (Exchange id Identity) of
  Exchange sa bt -> k sa (runIdentity #. bt)

-- used to provide access to the two parts of a prism
data Market a b s t = Market (b -> t) (s -> Either t a)

instance Functor (Market a b s) where
  fmap h (Market f g) = Market (h . f) (either (Left . h) Right . g)

instance Profunctor (Market a b) where
  dimap f' g' (Market f g) = Market (g' . f) (left' g' . g . f')

instance Choice (Market a b) where
  left' (Market f g) = Market (Left . f) $ \case
    Left x -> case g x of
      Left y -> Left (Left y)
      Right y -> Right y
    Right c -> Left (Right c)
  -- TODO: implement right'

withPrism :: Prism s t a b -> ((b -> t) -> (s -> Either t a) -> r) -> r
withPrism p k = case p (Market Identity Right) of
  Market bt sa -> k (runIdentity #. bt) (left' runIdentity . sa)

prism :: (b -> t) -> (s -> Either t a) -> Prism s t a b
prism bt seta = dimap seta (either pure (fmap bt)) . right'

matching :: Prism s t a b -> s -> Either t a
matching k = withPrism k $ \_ seta -> seta

_Left :: Prism (Either a c) (Either b c) a b
_Left = prism Left $ either Right (Left . Right)

_Right :: Prism (Either c a) (Either c b) a b
_Right = prism Right $ either (Left . Left) Right

-- TODO: rewrite  / become more comfortable with this section

re :: Review t b -> Getter b t
re p = to (runIdentity #. unTagged #. p .# Tagged .# Identity)

to :: (Profunctor p, Contravariant f) => (s -> a) -> Optic' p f s a
to k = dimap k (contramap k)

reviewSimple :: Prism' s a -> a -> s
reviewSimple r = runIdentity . unTagged . r . Tagged . Identity

review :: MonadReader b m => Optic' Tagged Identity t b -> m t
review r = asks $ runIdentity . unTagged . r . Tagged . Identity

foldMapOf :: Getting r s a -> (a -> r) -> s -> r
foldMapOf l f = getConst #. l (Const #. f)

preview :: Prism' s a -> s -> Maybe a
preview p s = getFirst #. foldMapOf p (First #. Just) $ s

(^?) :: s -> Prism' s a -> Maybe a
(^?) s p = preview p s

--

from :: Iso s t a b -> Iso b a t s
from l = withIso l $ \ sa bt -> iso bt sa

view :: Lens s t a b -> s -> a
view l s = getConst (l Const s)

over :: Lens s t a b -> (a -> b) -> s -> t
over l f = runIdentity . l (Identity . f)

set :: Lens s t a b -> b -> s -> t
set l = over l . const

lens :: (s -> a) -> (s -> b -> t) -> Lens s t a b
lens get' set' f s = set' s <$> f (get' s)

realLens, realLens' :: RealFloat a => Lens' (Complex a) a
realLens f (r :+ i) = fmap (:+ i) (f r)

imagLens, imagLens' :: RealFloat a => Lens' (Complex a) a
imagLens f (r :+ i) = fmap (r :+) (f i)

realLens' = lens (\(r :+ _) -> r) (\(_ :+ i) r -> r :+ i)
imagLens' = lens (\(_ :+ i) -> i) (\(r :+ _) i -> r :+ i)

(<&>) :: Functor f => f a -> (a -> b) -> f b
(<&>) = flip fmap

class Field1 s t a b | s -> a, t -> b, s b -> t, t a -> s where
  _1 :: Lens s t a b

instance Field1 (a, b) (a', b) a a' where
  _1 f ~(a, b) = f a <&> \a' -> (a', b)

instance Field1 (a, b, c) (a', b, c) a a' where
  _1 k ~(a, b, c) = k a <&> \a' -> (a', b, c)

class Field2 s t a b | s -> a, t -> b, s b -> t, t a -> s where
  _2 :: Lens s t a b

instance Field2 (a, b) (a, b') b b' where
  _2 f ~(a, b) = f b <&> \b' -> (a, b')

instance Field2 (a, b, c) (a, b', c) b b' where
  _2 k ~(a, b, c) = k b <&> \b' -> (a, b', c)

main :: IO ()
main = do
  print $ set realLens (1 :: Double) (0 :+ 1)
  print $ view _1 ('x', 'y', 'z')
  print $ view _2 $ set _2 False ('x', 'y', 'z')


module Wavefront.Types where

import Numeric.Vector

data WavefrontLine = Comment String
                   | Vertex Double Double Double
                   | TexCoord Double Double
                   | Normal Double Double Double
                   | Face Int Int Int Int Int Int Int Int Int
                   deriving (Show)

isComment :: WavefrontLine -> Bool
isComment (Comment _) = True
isComment _ = False

isVertex :: WavefrontLine -> Bool
isVertex (Vertex _ _ _) = True
isVertex _ = False

isTexCoord :: WavefrontLine -> Bool
isTexCoord (TexCoord _ _) = True
isTexCoord _ = False

isNormal :: WavefrontLine -> Bool
isNormal (Normal _ _ _) = True
isNormal _ = False

isFace :: WavefrontLine -> Bool
isFace (Face _ _ _ _ _ _ _ _ _) = True
isFace _ = False

data Vertex_ = Vertex_ Double Double Double deriving (Show)
data TexCoord_ = TexCoord_ Double Double deriving (Show)
data Normal_ = Normal_ Double Double Double deriving (Show)
data Face_ = Face_ Int Int Int Int Int Int Int Int Int deriving (Show)

vertexFromLine :: WavefrontLine -> Vertex_
vertexFromLine (Vertex a b c) = Vertex_ a b c
vertexFromLine _ = error "vertexFromLine is a partial function: it only accepts vertex lines"

texCoordFromLine :: WavefrontLine -> TexCoord_
texCoordFromLine (TexCoord a b) = TexCoord_ a b
texCoordFromLine _ = error "texCoordFromLine is a partial function: it only accepts tex coord lines"

normalFromLine :: WavefrontLine -> Normal_
normalFromLine (Normal a b c) = Normal_ a b c
normalFromLine _ = error "normalFromLine is a partial function: it only accepts normal lines"

faceFromLine :: WavefrontLine -> Face_
faceFromLine (Face a b c d e f g h i) = Face_ a b c d e f g h i
faceFromLine _ = error "vertexFromLine is a partial function: it only accepts face lines"

vertexToVec :: Vertex_ -> Vec3d
vertexToVec (Vertex_ a b c) = vec3 a b c

texCoordToVec :: TexCoord_ -> Vec2d
texCoordToVec (TexCoord_ a b) = vec2 a b

normalToVec :: Normal_ -> Vec3d
normalToVec (Normal_ a b c) = vec3 a b c


{-# LANGUAGE TupleSections #-}

{-|

Copyright:
  This file is part of the package zxcvbn-hs. It is subject to the
  license terms in the LICENSE file found in the top-level directory
  of this distribution and at:

    https://code.devalot.com/sthenauth/zxcvbn-hs

  No part of this package, including this file, may be copied,
  modified, propagated, or distributed except according to the terms
  contained in the LICENSE file.

License: MIT

-}
module Text.Password.Strength.Internal.Keyboard (
  -- * Keyboard Pattern Matching
  KeyboardPattern,
  keyboardToken,
  keyboardPattern,
  keyboardEstimate
  ) where

--------------------------------------------------------------------------------
-- Library Imports:
import Control.Lens (Lens, (^.), _3)
import Data.Foldable (foldl')
import Numeric.SpecFunctions (choose)

--------------------------------------------------------------------------------
-- Project Imports:
import Text.Password.Strength.Internal.Adjacency
import Text.Password.Strength.Internal.Token
import Text.Password.Strength.Internal.Math (variations')

--------------------------------------------------------------------------------
-- | Information about a found pattern.
newtype KeyboardPattern =
  KeyboardPattern (Int, Int, Token, AdjacencyScore)
  deriving (Show)

--------------------------------------------------------------------------------
-- | Allow other code to access the token used in a pattern.
keyboardToken :: Lens KeyboardPattern KeyboardPattern Token Token
keyboardToken f (KeyboardPattern t) = KeyboardPattern <$> _3 f t

--------------------------------------------------------------------------------
-- | Helper function to check if a token forms a keyboard pattern.
keyboardPattern :: AdjacencyTable -> Token -> Maybe KeyboardPattern
keyboardPattern graph token = KeyboardPattern .
  (graph ^. totalChars, graph ^. averageNeighbors, token,) .
    foldl' scoreSequence mempty <$> findSequence (token ^. tokenChars) graph

--------------------------------------------------------------------------------
-- | Estimate the number of guesses needed for a keyboard pattern to
-- be cracked.
keyboardEstimate :: KeyboardPattern -> Integer
keyboardEstimate (KeyboardPattern (s, d, _, a)) =
    e3 * e2 (a ^. primaryLayer) (a ^. secondaryLayer)

  where
    -- Equation 3, section 4, page 163 (8/18)
    --
    -- Deviations from the paper or other implementations:
    --
    --   * There's a typo in the paper: min(T, i - i)
    --     but should be: min(T, i - 1)
    --
    --   * Another typo, i should start at 2 and not 1.
    --
    --   * The other implementations don't seem to divide the outer
    --     sum by two but the equation clearly does.
    e3 :: Integer
    e3 = max 1 . (`div` 2) . sum $ do
      i <- [2 .. (a ^. patternLength)]
      j <- [1 .. min (a ^. totalTurns) (i - 1)]
      pure $ floor (choose (i - 1) (j - 1)) * toInteger s * (toInteger d ^ j)

    -- Modified version of equation 2 for primary layer vs. secondary.
    e2 :: Int -> Int -> Integer
    e2 = variations'


{-# LANGUAGE TypeFamilies  #-}
{-# LANGUAGE TypeOperators #-}

-----------------------------------------------------------------------------
-- |
-- Module      :  Numeric.Vectorized
-- Copyright   :  (c) Alberto Ruiz 2007-15
-- License     :  BSD3
-- Maintainer  :  Alberto Ruiz
-- Stability   :  provisional
--
-- Low level interface to vector operations.
--
-----------------------------------------------------------------------------

module Internal.Vectorized where

import           Control.Monad         (when)
import           Data.Complex
import           Foreign.C.String
import           Foreign.C.Types
import           Foreign.Marshal.Alloc (free, malloc)
import           Foreign.Marshal.Array (copyArray, newArray)
import           Foreign.Ptr           (Ptr)
import           Foreign.Storable      (Storable, peek)
import           Internal.Devel
import           Internal.Vector
import           System.IO.Unsafe      (unsafePerformIO)

infixr 1 #
(#) :: TransArray c => c -> (b -> IO r) -> TransRaw c b -> IO r
a # b = applyRaw a b
{-# INLINE (#) #-}

(#!) :: (TransArray c, TransArray c1) => c1 -> c -> TransRaw c1 (TransRaw c (IO r)) -> IO r
a #! b = a # b # id
{-# INLINE (#!) #-}

fromei :: Enum a => a -> CInt
fromei x = fromIntegral (fromEnum x) :: CInt

data FunCodeV = Sin
              | Cos
              | Tan
              | Abs
              | ASin
              | ACos
              | ATan
              | Sinh
              | Cosh
              | Tanh
              | ASinh
              | ACosh
              | ATanh
              | Exp
              | Log
              | Sign
              | Sqrt
              deriving Enum

data FunCodeSV = Scale
               | Recip
               | AddConstant
               | Negate
               | PowSV
               | PowVS
               | ModSV
               | ModVS
               deriving Enum

data FunCodeVV = Add
               | Sub
               | Mul
               | Div
               | Pow
               | ATan2
               | Mod
               deriving Enum

data FunCodeS = Norm2
              | AbsSum
              | MaxIdx
              | Max
              | MinIdx
              | Min
              deriving Enum

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

-- | sum of elements
sumF :: Vector Float -> Float
sumF = sumg c_sumF

-- | sum of elements
sumR :: Vector Float -> Float
sumR = sumg c_sumF

-- | sum of elements
sumQ :: Vector (Complex Float) -> Complex Float
sumQ = sumg c_sumQ

-- | sum of elements
sumC :: Vector (Complex Float) -> Complex Float
sumC = sumg c_sumQ

sumI :: ( TransRaw c (CInt -> Ptr a -> IO CInt) ~ (CInt -> Ptr I -> I :> Ok)
        , TransArray c
        , Storable a
        )
     => I -> c -> a
sumI m = sumg (c_sumI m)

sumL :: ( TransRaw c (CInt -> Ptr a -> IO CInt) ~ (CInt -> Ptr Z -> Z :> Ok)
        , TransArray c
        , Storable a
        ) => Z -> c -> a
sumL m = sumg (c_sumL m)

sumg :: (TransArray c, Storable a) => TransRaw c (CInt -> Ptr a -> IO CInt) -> c -> a
sumg f x = unsafePerformIO $ do
    r <- createVector 1
    (x #! r) f #| "sum"
    return $ r @> 0

type TVV t = t :> t :> Ok

foreign import ccall unsafe "sumF" c_sumF :: TVV Float
-- foreign import ccall unsafe "sumR" c_sumR :: TVV Double
foreign import ccall unsafe "sumQ" c_sumQ :: TVV (Complex Float)
-- foreign import ccall unsafe "sumC" c_sumC :: TVV (Complex Double)
foreign import ccall unsafe "sumI" c_sumI :: I -> TVV I
foreign import ccall unsafe "sumL" c_sumL :: Z -> TVV Z

-- | product of elements
prodF :: Vector Float -> Float
prodF = prodg c_prodF

-- | product of elements
prodR :: Vector Float -> Float
prodR = prodg c_prodF

-- | product of elements
prodQ :: Vector (Complex Float) -> Complex Float
prodQ = prodg c_prodQ

-- | product of elements
prodC :: Vector (Complex Float) -> Complex Float
prodC = prodg c_prodQ

prodI :: I-> Vector I -> I
prodI = prodg . c_prodI

prodL :: Z-> Vector Z -> Z
prodL = prodg . c_prodL

prodg :: (TransArray c, Storable a)
      => TransRaw c (CInt -> Ptr a -> IO CInt) -> c -> a
prodg f x = unsafePerformIO $ do
    r <- createVector 1
    (x #! r) f #| "prod"
    return $ r @> 0

foreign import ccall unsafe "prodF" c_prodF :: TVV Float
-- foreign import ccall unsafe "prodR" c_prodR :: TVV Double
foreign import ccall unsafe "prodQ" c_prodQ :: TVV (Complex Float)
-- foreign import ccall unsafe "prodC" c_prodC :: TVV (Complex Double)
foreign import ccall unsafe "prodI" c_prodI :: I -> TVV I
foreign import ccall unsafe "prodL" c_prodL :: Z -> TVV Z

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

toScalarAux :: (Enum a, TransArray c, Storable a1)
            => (CInt -> TransRaw c (CInt -> Ptr a1 -> IO CInt)) -> a -> c -> a1
toScalarAux fun code v = unsafePerformIO $ do
    r <- createVector 1
    (v #! r) (fun (fromei code)) #|"toScalarAux"
    return (r @> 0)

vectorMapAux :: (Enum a, Storable t, Storable a1)
             => (CInt -> CInt -> Ptr t -> CInt -> Ptr a1 -> IO CInt)
             -> a -> Vector t -> Vector a1
vectorMapAux fun code v = unsafePerformIO $ do
    r <- createVector (dim v)
    (v #! r) (fun (fromei code)) #|"vectorMapAux"
    return r

vectorMapValAux :: (Enum a, Storable a2, Storable t, Storable a1)
                => (CInt -> Ptr a2 -> CInt -> Ptr t -> CInt -> Ptr a1 -> IO CInt)
                -> a -> a2 -> Vector t -> Vector a1
vectorMapValAux fun code val v = unsafePerformIO $ do
    r <- createVector (dim v)
    pval <- newArray [val]
    (v #! r) (fun (fromei code) pval) #|"vectorMapValAux"
    free pval
    return r

vectorZipAux :: (Enum a, TransArray c, Storable t, Storable a1)
             => (CInt -> CInt -> Ptr t -> TransRaw c (CInt -> Ptr a1 -> IO CInt))
             -> a -> Vector t -> c -> Vector a1
vectorZipAux fun code u v = unsafePerformIO $ do
    r <- createVector (dim u)
    (u # v #! r) (fun (fromei code)) #|"vectorZipAux"
    return r

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

-- | obtains different functions of a vector: norm1, norm2, max, min, posmax, posmin, etc.
toScalarR :: FunCodeS -> Vector Float -> Float
toScalarR oper =  toScalarAux c_toScalarF (fromei oper)

-- foreign import ccall unsafe "toScalarR" c_toScalarR :: CInt -> TVV Double

-- | obtains different functions of a vector: norm1, norm2, max, min, posmax, posmin, etc.
toScalarF :: FunCodeS -> Vector Float -> Float
toScalarF oper =  toScalarAux c_toScalarF (fromei oper)

foreign import ccall unsafe "toScalarF" c_toScalarF :: CInt -> TVV Float

-- | obtains different functions of a vector: only norm1, norm2
toScalarC :: FunCodeS -> Vector (Complex Float) -> Float
toScalarC oper =  toScalarAux c_toScalarQ (fromei oper)

-- foreign import ccall unsafe "toScalarC" c_toScalarC :: CInt -> Complex Double :> Double :> Ok

-- | obtains different functions of a vector: only norm1, norm2
toScalarQ :: FunCodeS -> Vector (Complex Float) -> Float
toScalarQ oper =  toScalarAux c_toScalarQ (fromei oper)

foreign import ccall unsafe "toScalarQ" c_toScalarQ :: CInt -> Complex Float :> Float :> Ok

-- | obtains different functions of a vector: norm1, norm2, max, min, posmax, posmin, etc.
toScalarI :: FunCodeS -> Vector CInt -> CInt
toScalarI oper =  toScalarAux c_toScalarI (fromei oper)

foreign import ccall unsafe "toScalarI" c_toScalarI :: CInt -> TVV CInt

-- | obtains different functions of a vector: norm1, norm2, max, min, posmax, posmin, etc.
toScalarL :: FunCodeS -> Vector Z -> Z
toScalarL oper =  toScalarAux c_toScalarL (fromei oper)

foreign import ccall unsafe "toScalarL" c_toScalarL :: CInt -> TVV Z

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

-- | map of real vectors with given function
vectorMapR :: FunCodeV -> Vector Float -> Vector Float
vectorMapR = vectorMapAux c_vectorMapF

-- foreign import ccall unsafe "mapR" c_vectorMapR :: CInt -> TVV Double

-- | map of complex vectors with given function
vectorMapC :: FunCodeV -> Vector (Complex Float) -> Vector (Complex Float)
vectorMapC oper = vectorMapAux c_vectorMapQ (fromei oper)

-- foreign import ccall unsafe "mapC" c_vectorMapC :: CInt -> TVV (Complex Double)

-- | map of real vectors with given function
vectorMapF :: FunCodeV -> Vector Float -> Vector Float
vectorMapF = vectorMapAux c_vectorMapF

foreign import ccall unsafe "mapF" c_vectorMapF :: CInt -> TVV Float

-- | map of real vectors with given function
vectorMapQ :: FunCodeV -> Vector (Complex Float) -> Vector (Complex Float)
vectorMapQ = vectorMapAux c_vectorMapQ

foreign import ccall unsafe "mapQ" c_vectorMapQ :: CInt -> TVV (Complex Float)

-- | map of real vectors with given function
vectorMapI :: FunCodeV -> Vector CInt -> Vector CInt
vectorMapI = vectorMapAux c_vectorMapI

foreign import ccall unsafe "mapI" c_vectorMapI :: CInt -> TVV CInt

-- | map of real vectors with given function
vectorMapL :: FunCodeV -> Vector Z -> Vector Z
vectorMapL = vectorMapAux c_vectorMapL

foreign import ccall unsafe "mapL" c_vectorMapL :: CInt -> TVV Z

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

-- | map of real vectors with given function
vectorMapValR :: FunCodeSV -> Float -> Vector Float -> Vector Float
vectorMapValR oper = vectorMapValAux c_vectorMapValF (fromei oper)

-- foreign import ccall unsafe "mapValR" c_vectorMapValR :: CInt -> Ptr Double -> TVV Double

-- | map of complex vectors with given function
vectorMapValC :: FunCodeSV -> Complex Float -> Vector (Complex Float) -> Vector (Complex Float)
vectorMapValC = vectorMapValAux c_vectorMapValQ

-- foreign import ccall unsafe "mapValC" c_vectorMapValC :: CInt -> Ptr (Complex Double) -> TVV (Complex Double)

-- | map of real vectors with given function
vectorMapValF :: FunCodeSV -> Float -> Vector Float -> Vector Float
vectorMapValF oper = vectorMapValAux c_vectorMapValF (fromei oper)

foreign import ccall unsafe "mapValF" c_vectorMapValF :: CInt -> Ptr Float -> TVV Float

-- | map of complex vectors with given function
vectorMapValQ :: FunCodeSV -> Complex Float -> Vector (Complex Float) -> Vector (Complex Float)
vectorMapValQ oper = vectorMapValAux c_vectorMapValQ (fromei oper)

foreign import ccall unsafe "mapValQ" c_vectorMapValQ :: CInt -> Ptr (Complex Float) -> TVV (Complex Float)

-- | map of real vectors with given function
vectorMapValI :: FunCodeSV -> CInt -> Vector CInt -> Vector CInt
vectorMapValI oper = vectorMapValAux c_vectorMapValI (fromei oper)

foreign import ccall unsafe "mapValI" c_vectorMapValI :: CInt -> Ptr CInt -> TVV CInt

-- | map of real vectors with given function
vectorMapValL :: FunCodeSV -> Z -> Vector Z -> Vector Z
vectorMapValL oper = vectorMapValAux c_vectorMapValL (fromei oper)

foreign import ccall unsafe "mapValL" c_vectorMapValL :: CInt -> Ptr Z -> TVV Z

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

type TVVV t = t :> t :> t :> Ok

-- | elementwise operation on real vectors
vectorZipR :: FunCodeVV -> Vector Float -> Vector Float -> Vector Float
vectorZipR = vectorZipAux c_vectorZipF

-- foreign import ccall unsafe "zipR" c_vectorZipR :: CInt -> TVVV Double

-- | elementwise operation on complex vectors
vectorZipC :: FunCodeVV -> Vector (Complex Float) -> Vector (Complex Float) -> Vector (Complex Float)
vectorZipC = vectorZipAux c_vectorZipQ

-- foreign import ccall unsafe "zipC" c_vectorZipC :: CInt -> TVVV (Complex Double)

-- | elementwise operation on real vectors
vectorZipF :: FunCodeVV -> Vector Float -> Vector Float -> Vector Float
vectorZipF = vectorZipAux c_vectorZipF

foreign import ccall unsafe "zipF" c_vectorZipF :: CInt -> TVVV Float

-- | elementwise operation on complex vectors
vectorZipQ :: FunCodeVV -> Vector (Complex Float) -> Vector (Complex Float) -> Vector (Complex Float)
vectorZipQ = vectorZipAux c_vectorZipQ

foreign import ccall unsafe "zipQ" c_vectorZipQ :: CInt -> TVVV (Complex Float)

-- | elementwise operation on CInt vectors
vectorZipI :: FunCodeVV -> Vector CInt -> Vector CInt -> Vector CInt
vectorZipI = vectorZipAux c_vectorZipI

foreign import ccall unsafe "zipI" c_vectorZipI :: CInt -> TVVV CInt

-- | elementwise operation on CInt vectors
vectorZipL :: FunCodeVV -> Vector Z -> Vector Z -> Vector Z
vectorZipL = vectorZipAux c_vectorZipL

foreign import ccall unsafe "zipL" c_vectorZipL :: CInt -> TVVV Z

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

-- foreign import ccall unsafe "vectorScan" c_vectorScan
--     :: CString -> Ptr CInt -> Ptr (Ptr Double) -> IO CInt

foreign import ccall unsafe "vectorScan" c_vectorScan
    :: CString -> Ptr CInt -> Ptr (Ptr Float) -> IO CInt

vectorScan :: FilePath -> IO (Vector Float)
vectorScan s = do
    pp <- malloc
    pn <- malloc
    cs <- newCString s
    ok <- c_vectorScan cs pn pp
    when (not (ok == 0)) $
        error ("vectorScan: file \"" ++ s ++"\" not found")
    n <- fromIntegral <$> peek pn
    p <- peek pp
    v <- createVector n
    free pn
    free cs
    unsafeWith v $ \pv -> copyArray pv p n
    free p
    free pp
    return v

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

type Seed = Int

data RandDist = Uniform  -- ^ uniform distribution in [0,1)
              | Gaussian -- ^ normal distribution with mean zero and standard deviation one
              deriving Enum

-- | Obtains a vector of pseudorandom elements (use randomIO to get a random seed).
randomVector :: Seed
             -> RandDist -- ^ distribution
             -> Int      -- ^ vector size
             -> Vector Float
randomVector seed dist n = unsafePerformIO $ do
    r <- createVector n
    (r # id) (c_random_vector (fi seed) ((fi.fromEnum) dist)) #|"randomVector"
    return r

-- foreign import ccall unsafe "random_vector" c_random_vector :: CInt -> CInt -> Double :> Ok
foreign import ccall unsafe "random_vector" c_random_vector :: CInt -> CInt -> Float :> Ok

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

roundVector :: Vector Float -> Vector Float
roundVector v = unsafePerformIO $ do
    r <- createVector (dim v)
    (v #! r) c_round_vector #|"roundVector"
    return r

-- foreign import ccall unsafe "round_vector" c_round_vector :: TVV Double

foreign import ccall unsafe "round_vector" c_round_vector :: TVV Float

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

-- |
-- >>> range 5
-- [0,1,2,3,4]
-- it :: Vector I
--
range :: Int -> Vector I
range n = unsafePerformIO $ do
    r <- createVector n
    (r # id) c_range_vector #|"range"
    return r

foreign import ccall unsafe "range_vector" c_range_vector :: CInt :> Ok

-- float2DoubleV :: Vector Float -> Vector Float
-- float2DoubleV x = tog c_float2double

-- double2FloatV :: Vector Double -> Vector Float
-- double2FloatV = tog c_double2float

-- double2IntV :: Vector Double -> Vector CInt
-- double2IntV = tog c_double2int

-- int2DoubleV :: Vector CInt -> Vector Double
-- int2DoubleV = tog c_int2double

-- double2longV :: Vector Double -> Vector Z
-- double2longV = tog c_double2long

long2FloatV :: Vector Z -> Vector Float
long2FloatV = tog c_long2float

float2IntV :: Vector Float -> Vector CInt
float2IntV = tog c_float2int

int2floatV :: Vector CInt -> Vector Float
int2floatV = tog c_int2float

int2longV :: Vector I -> Vector Z
int2longV = tog c_int2long

long2intV :: Vector Z -> Vector I
long2intV = tog c_long2int

long2floatV :: Vector Z -> Vector Float
long2floatV = tog c_long2float

tog :: (Storable t, Storable a)
    => (CInt -> Ptr t -> CInt -> Ptr a -> IO CInt) -> Vector t -> Vector a
tog f v = unsafePerformIO $ do
    r <- createVector (dim v)
    (v #! r) f #|"tog"
    return r

-- foreign import ccall unsafe "float2double" c_float2double :: Float :> Double :> Ok
-- foreign import ccall unsafe "double2float" c_double2float :: Double :> Float :> Ok
-- foreign import ccall unsafe "int2double"   c_int2double   :: CInt :> Double :> Ok
-- foreign import ccall unsafe "double2int"   c_double2int   :: Double :> CInt :> Ok
-- foreign import ccall unsafe "long2double"  c_long2double   :: Z :> Double :> Ok
-- foreign import ccall unsafe "double2long"  c_double2long   :: Double :> Z :> Ok
foreign import ccall unsafe "float2int"   c_float2int   :: Float :> CInt  :> Ok
foreign import ccall unsafe "float2long"   c_float2long :: Float :> I     :> Ok
foreign import ccall unsafe "int2float"   c_int2float   :: CInt  :> Float :> Ok
foreign import ccall unsafe "int2long"    c_int2long    :: I     :> Z     :> Ok
foreign import ccall unsafe "long2int"    c_long2int    :: Z     :> I     :> Ok
foreign import ccall unsafe "long2float"  c_long2float  :: Z     :> Float :> Ok

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

stepg :: (Storable t, Storable a)
      => (CInt -> Ptr t -> CInt -> Ptr a -> IO CInt) -> Vector t -> Vector a
stepg f v = unsafePerformIO $ do
    r <- createVector (dim v)
    (v #! r) f #|"step"
    return r

-- stepD :: Vector Double -> Vector Double
-- stepD = stepg c_stepD

stepF :: Vector Float -> Vector Float
stepF = stepg c_stepF

stepI :: Vector CInt -> Vector CInt
stepI = stepg c_stepI

stepL :: Vector Z -> Vector Z
stepL = stepg c_stepL

foreign import ccall unsafe "stepF" c_stepF :: TVV Float
-- foreign import ccall unsafe "stepD" c_stepD :: TVV Double
foreign import ccall unsafe "stepI" c_stepI :: TVV CInt
foreign import ccall unsafe "stepL" c_stepL :: TVV Z

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

conjugateAux :: (Storable t, Storable a)
             => (CInt -> Ptr t -> CInt -> Ptr a -> IO CInt) -> Vector t -> Vector a
conjugateAux fun x = unsafePerformIO $ do
    v <- createVector (dim x)
    (x #! v) fun #|"conjugateAux"
    return v

conjugateQ :: Vector (Complex Float) -> Vector (Complex Float)
conjugateQ = conjugateAux c_conjugateQ
foreign import ccall unsafe "conjugateQ" c_conjugateQ :: TVV (Complex Float)

-- conjugateC :: Vector (Complex Double) -> Vector (Complex Double)
-- conjugateC = conjugateAux c_conjugateC
-- foreign import ccall unsafe "conjugateC" c_conjugateC :: TVV (Complex Double)

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

cloneVector :: Storable t => Vector t -> IO (Vector t)
cloneVector v = do
        let n = dim v
        r <- createVector n
        let f _ s _ d =  copyArray d s n >> return 0
        (v #! r) f #|"cloneVector"
        return r

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

constantAux :: (Storable a1, Storable a)
            => (Ptr a1 -> CInt -> Ptr a -> IO CInt) -> a1 -> Int -> Vector a
constantAux fun x n = unsafePerformIO $ do
    v <- createVector n
    px <- newArray [x]
    (v # id) (fun px) #|"constantAux"
    free px
    return v

type TConst t = Ptr t -> t :> Ok

foreign import ccall unsafe "constantF" cconstantF :: TConst Float
-- foreign import ccall unsafe "constantR" cconstantR :: TConst Double
foreign import ccall unsafe "constantQ" cconstantQ :: TConst (Complex Float)
-- foreign import ccall unsafe "constantC" cconstantC :: TConst (Complex Double)
foreign import ccall unsafe "constantI" cconstantI :: TConst CInt
foreign import ccall unsafe "constantL" cconstantL :: TConst Z

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


{-# LANGUAGE BangPatterns #-}
module ML.Sample (Sample(..)) where

import Numeric.LinearAlgebra (R, Vector)

-- | A training sample.
data Sample = Sample
              { sampleInput          :: {-# UNPACK #-} !(Vector R)
              , sampleExpectedOutput :: {-# UNPACK #-} !(Vector R)
              } deriving (Read, Show)


{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE UndecidableInstances #-}
{-# LANGUAGE TypeSynonymInstances #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE OverlappingInstances #-}
{-# LANGUAGE EmptyDataDecls #-}

module TestBase(
    GEMMTestData(..),
    StrassenTestData(..),
    TRSMTestData(..),
    PotrfTestData(..),
    Tag2(),
    eqMatrix,
    handleResult,
    Truncable(..)

) where

import           Data.Complex (Complex ((:+)))
import Test.QuickCheck (Result(..), Arbitrary(arbitrary), Testable(property), Gen, choose, vectorOf, elements)
import Data.Tagged (Tagged(Tagged, unTagged))
import Data.Default (Default(def))
import System.Random (Random)
import Debug.Trace (trace)
import Control.Applicative (liftA)

import FPNLA.Matrix
       (Matrix(fromList_m, dim_m, elem_m),
        Matrix(generate_m, transpose_m, zipWith_m))
import FPNLA.Matrix.Instances.HMatrix ()
import FPNLA.Operations.Parameters (Elt, TriangType(..), TransType(..), UnitType(..), StratCtx)
import FPNLA.Operations.LAPACK.Strategies (HMatrixBindSeq, CholLLVSeq, CholLLVPar_Repa, CholLLVBlocksSeq, newNullContext, newSqrBlockContext)
import FPNLA.Operations.BLAS.Strategies (DefPar_ST, CBindSeq, DefSeq, BlocksSeq, BlocksByRows, BlocksByCols, DefPar_MP, BlocksPar_MP, StrassenPar_MP, ColsPar_MP, DefPar_ST, DefPar_Acc, BlocksPar_ST, StrassenPar_ST, ColsPar_ST, newBlockContext, newStrassenContext)
import FPNLA.Utils (iif)

-- Links interesantes:
-- http://www.haskell.org/haskellwiki/Scoped_type_variables
-- http://www.haskell.org/ghc/docs/7.4.1/html/users_guide/other-type-extensions.html seccion 7.11.7
-- http://hackage.haskell.org/trac/ghc/ticket/3485

blockSize :: Int
blockSize = 2

-- Valor de epsilon para hacer las comparaciones:
epsilon :: Double
epsilon = 1e-5

class Truncable a where
    inEpsilonRange :: Real e => e -> a -> a -> Bool

instance (Floating n, RealFrac n) => Truncable n where
    inEpsilonRange eps n1 n2 = expN1 == expN2 && abs(mantN1 - mantN2) <= epsFractional
        where
            epsFractional = fromRational $ toRational eps
            shift n d1 = n * (10 ** fromIntegral d1)
            getExponent r
                    | r == 0 = 0
                    | r < 0 = getExponent (abs r)
                    | r > 10 = 1 + getExponent (r / 10)
                    | r < 1 = (-1) + getExponent (r * 10)
                    | otherwise = 0
            expN1 :: Integer = getExponent n1
            expN2 :: Integer = getExponent n2
            mantN1 = shift n1 (-expN1)
            mantN2 = shift n2 (-expN2)

--
instance (RealFrac n, Truncable n) => Truncable (Complex n) where
    inEpsilonRange eps (r1:+i1) (r2:+i2) = inEpsilonRange eps r1 r2 && inEpsilonRange eps i1 i2

-- Compara una matriz elemento a elemento truncando los d√≠gitos. Si falla imprime los √≠ndices del primer elemento que fall√≥:
eqMatrix :: (Show e, Eq e, Truncable e, Matrix m e) => m e -> m e -> Bool
eqMatrix m1 m2 = dimm1 == dimm2 && eqElems 0 0
    where
        dimm1 = dim_m m1
        dimm2 = dim_m m2
        (m,n) = dimm1
        getElem = {--truncDigits truncValue $--} elem_m
        eqElems i j
          | j >= n = eqElems (i+1) 0
          | i >= m = True
          -- | getElem i j m1 /= getElem i j m2 = trace ("(i,j) = (" ++ show i ++ "," ++ show j ++ ")") False
          | not $ inEpsilonRange epsilon (getElem i j m1) (getElem i j m2) = trace ("(i,j) = (" ++ show i ++ "," ++ show j ++ ")" ++ " / m1 = " ++ show (getElem i j m1) ++ " / m2 = " ++ show (getElem i j m2)) False
          | otherwise =  eqElems i (j+1)

arbSDPMatrix :: (Elt e, Arbitrary e, Random e, Matrix m e) => Int -> Gen (m e)
arbSDPMatrix n = do
    la <- vectorOf (n*n) $ choose (0, 1)
    ld <- vectorOf n $ choose (100, 200)
    let a = fromList_m n n la
        d = generate_m n n (\i j -> iif (i == j) (ld !! i) 0)
        s = zipWith_m (+) a (transpose_m a)
        r = zipWith_m (+) s d
    return r

handleResult :: Monad m => Result -> m ()
handleResult (Success{}) = return ()
handleResult (GaveUp{}) = fail "GaveUp"
handleResult (Failure{}) = fail "Failure"
handleResult (NoExpectedFailure{}) = fail "NoExpectedFailure"

--
instance Arbitrary (TransType a) where
    arbitrary = elements [NoTrans undefined, Trans undefined, ConjTrans undefined]
--
instance Arbitrary (a -> TransType a) where
    arbitrary = elements [NoTrans, Trans, ConjTrans]
instance Arbitrary (a -> TriangType a) where
    arbitrary = elements [Lower, Upper]
instance Arbitrary (a -> UnitType a) where
    arbitrary = elements [Unit, NoUnit]

--
arbitraryIntBounded :: Gen Int
arbitraryIntBounded = choose (2, 10)

arbitraryMatrix :: (Arbitrary e, Matrix m e) => Int -> Int -> Gen (m e)
arbitraryMatrix m n = do
    l <- vectorOf (m*n) arbitrary
    return $ fromList_m m n l

arbitraryTransMatrix :: (Arbitrary e, Matrix m e) => TransType a -> Int -> Int -> Gen (TransType (m e))
arbitraryTransMatrix (NoTrans _) m n = liftA NoTrans $ arbitraryMatrix m n
arbitraryTransMatrix (Trans _) m n = liftA Trans $ arbitraryMatrix n m
arbitraryTransMatrix (ConjTrans _) m n = liftA ConjTrans $ arbitraryMatrix n m

data GEMMTestData m e = GEMMTestData (TransType (m e)) (TransType (m e)) e e (m e) deriving (Show)
data StrassenTestData m e = StrassenTestData (TransType (m e)) (TransType (m e)) e e (m e) deriving (Show)
data TRSMTestData m e = TRSMTestData e (TransType (TriangType (UnitType (m e)))) (m e) deriving (Show)

instance (Matrix m e, Arbitrary e) => Arbitrary (GEMMTestData m e) where
    arbitrary = do
        m <- arbitraryIntBounded
        n <- arbitraryIntBounded
        k <- arbitraryIntBounded
        ttmA <- arbitrary
        ttmB <- arbitrary
        pmA <- arbitraryTransMatrix ttmA m k
        pmB <- arbitraryTransMatrix ttmB k n
        alpha <- arbitrary
        beta <- arbitrary
        mC <- arbitraryMatrix m n
        return $ GEMMTestData pmA pmB alpha beta mC

instance (Matrix m e, Arbitrary e) => Arbitrary (StrassenTestData m e) where
    arbitrary = do
        e <- choose (2, 4) :: Gen Int
        let n = truncate $ (2 :: Double) ** fromIntegral e
        ttmA <- arbitrary
        ttmB <- arbitrary
        alpha <- arbitrary
        mA <- arbitraryTransMatrix ttmA n n
        mB <- arbitraryTransMatrix ttmB n n
        beta <- arbitrary
        mC <- arbitraryMatrix n n
        return $ StrassenTestData mA mB alpha beta mC
--
instance (Matrix m e, Arbitrary e) => Arbitrary (TRSMTestData m e) where
    arbitrary = do
        m <- arbitraryIntBounded
        n <- arbitraryIntBounded
        trans <- arbitrary :: Gen (a -> TransType a)
        triang <- arbitrary :: Gen (a -> TriangType a)
        unit <- arbitrary :: Gen (a -> UnitType a)
        alpha <- arbitrary
        mA <- arbitraryMatrix m m
        mB <- arbitraryMatrix m n
        return $ TRSMTestData alpha (trans . triang . unit $ mA) mB

data PotrfTestData m e = PotrfTestData (TriangType (m e)) deriving (Show)

instance (Matrix m e, Elt e, Arbitrary e, Random e) => Arbitrary (PotrfTestData m e) where
    arbitrary = do
        _ <- arbitrary :: Gen (a -> TriangType a)
        n <- choose (3, 5)
        m <- arbSDPMatrix (blockSize ^ (n :: Int))
        return $ PotrfTestData $ Lower m

-- Cualquier propiedad tageada es una propiedad:
instance (Testable p) => Testable (Tagged t p) where
    property p = property $ unTagged p

data Tag2 (s :: *) (v :: * -> *)

instance (StratCtx HMatrixBindSeq ~ ctx_Seq) => Default (Tagged HMatrixBindSeq ctx_Seq) where
    def = Tagged newNullContext

instance (StratCtx dots ~ ctx_dots, Default (Tagged dots ctx_dots), StratCtx gs ~ ctx_gs, Default (Tagged gs ctx_gs), StratCtx (CholLLVSeq dots gs) ~ ctx_bp) => Default (Tagged (CholLLVSeq dots gs) ctx_bp) where
    def = Tagged (unTagged (def :: Tagged dots ctx_dots), unTagged (def :: Tagged gs ctx_gs))

instance (StratCtx dots ~ ctx_dots, Default (Tagged dots ctx_dots), StratCtx gs ~ ctx_gs, Default (Tagged gs ctx_gs), StratCtx (CholLLVPar_Repa dots gs) ~ ctx_bp) => Default (Tagged (CholLLVPar_Repa dots gs) ctx_bp) where
    def = Tagged (unTagged (def :: Tagged dots ctx_dots), unTagged (def :: Tagged gs ctx_gs))

instance (StratCtx ps ~ ctx_ps, Default (Tagged ps ctx_ps), StratCtx ss ~ ctx_ss, Default (Tagged ss ctx_ss), StratCtx gs ~ ctx_gs, Default (Tagged gs ctx_gs), StratCtx ts ~ ctx_ts, Default (Tagged ts ctx_ts), StratCtx (CholLLVBlocksSeq ps ss gs ts) ~ ctx_cs) => Default (Tagged (CholLLVBlocksSeq ps ss gs ts) ctx_cs) where
    def = Tagged (newSqrBlockContext blockSize, unTagged (def :: Tagged ps ctx_ps), unTagged (def :: Tagged ss ctx_ss), unTagged (def :: Tagged gs ctx_gs), unTagged (def :: Tagged ts ctx_ts))

--
instance (StratCtx DefSeq ~ ctx_DefSeq) => Default (Tagged DefSeq ctx_DefSeq) where
    def = Tagged newNullContext
instance (StratCtx CBindSeq ~ ctx_DefSeq) => Default (Tagged CBindSeq ctx_DefSeq) where
    def = Tagged newNullContext
instance (StratCtx gs ~ ctx_gs, Default (Tagged gs ctx_gs), StratCtx (BlocksSeq gs) ~ ctx_bp) => Default (Tagged (BlocksSeq gs) ctx_bp) where
    def = Tagged (newBlockContext 2 2, unTagged (def :: Tagged gs ctx_gs))
instance (StratCtx gs ~ ctx_gs, Default (Tagged gs ctx_gs), StratCtx ts ~ ctx_ts, Default (Tagged ts ctx_ts), StratCtx (BlocksByCols gs ts) ~ ctx_bbcp) => Default (Tagged (BlocksByCols gs ts) ctx_bbcp) where
    def = Tagged (newSqrBlockContext 2, unTagged (def :: Tagged gs ctx_gs), unTagged (def :: Tagged ts ctx_ts))
instance (StratCtx gs ~ ctx_gs, Default (Tagged gs ctx_gs), StratCtx ts ~ ctx_ts, Default (Tagged ts ctx_ts), StratCtx (BlocksByRows gs ts) ~ ctx_bbcp) => Default (Tagged (BlocksByRows gs ts) ctx_bbcp) where
    def = Tagged (newSqrBlockContext 2, unTagged (def :: Tagged gs ctx_gs), unTagged (def :: Tagged ts ctx_ts))

instance (StratCtx DefPar_MP ~ ctx_DefPar) => Default (Tagged DefPar_MP ctx_DefPar) where
    def = Tagged newNullContext
instance (StratCtx gs ~ ctx_gs, Default (Tagged gs ctx_gs), StratCtx (BlocksPar_MP gs) ~ ctx_bp) => Default (Tagged (BlocksPar_MP gs) ctx_bp) where
    def = Tagged (newBlockContext 2 2, unTagged (def :: Tagged gs ctx_gs))
instance (StratCtx ts ~ ctx_ts, Default (Tagged ts ctx_ts), StratCtx (ColsPar_MP ts) ~ ctx_cp) => Default (Tagged (ColsPar_MP ts) ctx_cp) where
    def = Tagged $ unTagged (def :: Tagged ts ctx_ts)
instance (StratCtx gs ~ ctx_gs, Default (Tagged gs ctx_gs), StratCtx (StrassenPar_MP gs) ~ ctx_bp) => Default (Tagged (StrassenPar_MP gs) ctx_bp) where
   def = Tagged (newStrassenContext 2, unTagged (def :: Tagged gs ctx_gs))

instance (StratCtx DefPar_ST ~ ctx_DefPar) => Default (Tagged DefPar_ST ctx_DefPar) where
    def = Tagged newNullContext
instance (StratCtx gs ~ ctx_gs, Default (Tagged gs ctx_gs), StratCtx (BlocksPar_ST gs) ~ ctx_bp) => Default (Tagged (BlocksPar_ST gs) ctx_bp) where
    def = Tagged (newBlockContext 2 2, unTagged (def :: Tagged gs ctx_gs))
instance (StratCtx ts ~ ctx_ts, Default (Tagged ts ctx_ts), StratCtx (ColsPar_ST ts) ~ ctx_cp) => Default (Tagged (ColsPar_ST ts) ctx_cp) where
    def = Tagged $ unTagged (def :: Tagged ts ctx_ts)
instance (StratCtx gs ~ ctx_gs, Default (Tagged gs ctx_gs), StratCtx (StrassenPar_ST gs) ~ ctx_bp) => Default (Tagged (StrassenPar_ST gs) ctx_bp) where
   def = Tagged (newStrassenContext 2, unTagged (def :: Tagged gs ctx_gs))

instance (StratCtx DefPar_Acc ~ ctx_DefPar) => Default (Tagged DefPar_Acc ctx_DefPar) where
    def = Tagged newNullContext


import Paraiso
import Data.Complex
import System.Environment

main = do
  args <- getArgs
  let arch = if "--cuda" `elem` args then
               CUDA 128 128
             else
               X86
  putStrLn $ compile arch code
    where
      code = do
         parallel 16384 $ do
           z <- allocate
           z0<- allocate
           z =$ (Rand (-2.0) 2.0) :+ (Rand (-2.0) 2.0)
           z0=$ (z:: Complex (Expr Double))
           cuda $ do
             sequential 65536 $ do
               z =$ z * z - 1
           output [realPart z0,imagPart z0, realPart z, imagPart z]


module Statistics.GModeling.Models.FixedTreeLDA
  (
  ) where

import Statistics.GModeling.DSL

data LDATreeLabels =
    Alpha1 | Alpha2 | Beta
  | TopicDepth | TopicPath | Topic | Doc
  | Symbols | Symbol

ldaTree :: Network LDATreeLabels
ldaTree =
  [
    Only Alpha1         :-> TopicDepth
  , Only Alpha2         :-> TopicPath
  , Only Beta           :-> Symbols
  , (TopicPath :@ Doc)  :-> Topic
  , (TopicDepth :@ Doc) :-> Topic
  , (Symbols :@ Topic)  :-> Symbol
  ]


-- | An attempt at writing QuickHull in quintessential Haskell.
-- | I am just writing the simple 2D version for now.

module QuickHull (qhull2D) where

import           Control.Arrow
import           Control.Monad
import           Data.DList                   (DList)
import qualified Data.DList                   as D
import           Data.Foldable                (asum)
import qualified Data.List                    as L (splitAt)
import           Data.Maybe
import           Data.Ord
import qualified Data.Vector                  as V
import           Data.Vector.Algorithms.Intro as I
import qualified Data.Vector.Unboxed          as UV
import           Numeric.LinearAlgebra        (peps)

data Region2D
    = R021
    | R201
    | R012
    deriving (Show,Eq,Ord)

qhull2D
    :: V.Vector (UV.Vector Double)
    -> [UV.Vector Double]
qhull2D vecs
  | V.length vecs < 4 = V.toList vecs
  | otherwise =
      D.toList $ D.fromList [ax, bx, ay, by] `mappend`
                 qhull2DHelper ax by regU021 `mappend`
                 qhull2DHelper by bx regU201 `mappend`
                 qhull2DHelper bx ay regB201 `mappend`
                 qhull2DHelper ax ay regB021
  where
    ((regU021, regU201), (regB021, regB201)) =
      (assignTwoRegions *** assignTwoRegions) . splitValidQuad ax by bx ay $
        vecs
    [ax, bx, ay, by] = qhull2DInit vecs

qhull2DInit :: V.Vector (UV.Vector Double) -> [UV.Vector Double]
qhull2DInit points =
  [ V.unsafeIndex points ax
  , V.unsafeIndex points bx
  , V.unsafeIndex points ay
  , V.unsafeIndex points by
  ]
  where
    (ax, bx) = minMaxXIndexes points
    (ay, by) = minMaxYIndexes points

qhull2DHelper :: UV.Vector Double -> UV.Vector Double -> V.Vector (UV.Vector Double) -> DList (UV.Vector Double)
qhull2DHelper v w vecs
  | V.length vecs > 1 =
      D.cons maxP $
        qhull2DHelper w0 w1 r021 `mappend`
        qhull2DHelper w0 w2 r012 `mappend`
        qhull2DHelper w1 w2 r201
  | V.length vecs == 1 =
      D.singleton $
        V.unsafeHead vecs
  | otherwise = D.empty
  where
    (r021, r201, r012) = assignThreeRegions $ splitValidTri w0 w1 w2 testpts
    [w0, w1, w2] = V.toList . sortTriangle v w $ maxP
    testpts = V.ifilter (\i _ -> i /= maxInd) vecs
    (maxInd, maxP) = maxDistPt v w vecs

assignTwoRegions :: V.Vector (Region2D, UV.Vector Double)
                 -> (V.Vector (UV.Vector Double), V.Vector (UV.Vector Double))
assignTwoRegions = (,) <$> rMap R021 <*> rMap R201

assignThreeRegions
    :: V.Vector (Region2D, UV.Vector Double)
    -> (V.Vector (UV.Vector Double), V.Vector (UV.Vector Double), V.Vector (UV.Vector Double))
assignThreeRegions = (,,) <$> rMap R021 <*> rMap R201 <*> rMap R012

rMap
    :: Eq b
    => b -> V.Vector (b, UV.Vector Double) -> V.Vector (UV.Vector Double)
rMap r = V.map snd . V.filter ((== r) . fst)

-- Sorts points by their x-coordinates
sortTriangle
    :: UV.Vector Double
    -> UV.Vector Double
    -> UV.Vector Double
    -> V.Vector (UV.Vector Double)
sortTriangle u v w = V.modify (I.sortBy (comparing UV.unsafeHead)) r
  where
    r = V.fromList [u, v, w]

inTriangleInterior
    :: UV.Vector Double
    -> UV.Vector Double
    -> UV.Vector Double
    -> UV.Vector Double
    -> (Bool, Double, [Double])
inTriangleInterior a b c p = (,,) <$> same' . sum <*> signum . sum <*> id $ vv
  where
    vv = [test' ab aP, test' bc bp, test' ca cp]
    ab = diff b a
    bc = diff c b
    ca = diff a c
    aP = diff p a
    bp = diff p b
    cp = diff p c
    diff = UV.zipWith (-)
    same' = (||) <$> (== 3) <*> (== -3)
    test' u v = signum $ UV.foldl1' (-) $ UV.zipWith (*) u (UV.reverse v)

inQuadInterior
    :: UV.Vector Double
    -> UV.Vector Double
    -> UV.Vector Double
    -> UV.Vector Double
    -> UV.Vector Double
    -> (Bool, Double, [Double])
inQuadInterior a b c d p = (,,) <$> same' . sum <*> signum . sum <*> id $ vv
  where
    vv = [test' ab aP, test' bc bp, test' cd cp, test' da dp]
    ab = diff b a
    bc = diff c b
    cd = diff d c
    da = diff a d
    aP = diff p a
    bp = diff p b
    cp = diff p c
    dp = diff p d
    diff = UV.zipWith (-)
    same' = (||) <$> (== 4) <*> (== (-4))
    test' u v = signum $ UV.foldl1' (-) $ UV.zipWith (*) u (UV.reverse v)

testInRegionQuad
    :: UV.Vector Double
    -> UV.Vector Double
    -> UV.Vector Double
    -> UV.Vector Double
    -> UV.Vector Double
    -> (Maybe Region2D, Maybe Region2D)
testInRegionQuad v0 v1 v2 v3 p = reg
  where
    (inQuad,signSum,regs) = inQuadInterior v0 v1 v2 v3 p
    reg =
        if inQuad
            then (Nothing, Nothing)
            else (asum *** asum) . L.splitAt 2 . zipWith ($) regAssign $ regs
    regAssign =
        [ toMaybe R021 . (/= signSum)
        , toMaybe R201 . (/= signSum)
        , toMaybe R021 . (/= signSum)
        , toMaybe R201 . (/= signSum)]

splitValidQuad
    :: UV.Vector Double
    -> UV.Vector Double
    -> UV.Vector Double
    -> UV.Vector Double
    -> V.Vector (UV.Vector Double)
    -> (V.Vector (Region2D, UV.Vector Double), V.Vector (Region2D, UV.Vector Double))
splitValidQuad u v w x =
    (V.map fromJust . V.filter isJust *** V.map fromJust . V.filter isJust) .
    V.unzip .
    V.map
        ((\v' (a,b) ->
               (flip (,) v' <$> a, flip (,) v' <$> b)) `ap`
         testInRegionQuad u v w x)

toMaybe :: a -> Bool -> Maybe a
toMaybe a pred' =
    if pred'
        then Just a
        else Nothing

testInRegionTri
    :: UV.Vector Double
    -> UV.Vector Double
    -> UV.Vector Double
    -> UV.Vector Double
    -> Maybe Region2D
testInRegionTri v0 v1 v2 p = reg
  where
    (inTriangle,signSum,regs) = inTriangleInterior v0 v1 v2 p
    reg =
        if inTriangle
            then Nothing
            else asum . zipWith ($) regAssign $ regs
    regAssign =
        [ toMaybe R021 . (/= signSum)
        , toMaybe R201 . (/= signSum)
        , toMaybe R012 . (/= signSum)]

splitValidTri
    :: UV.Vector Double
    -> UV.Vector Double
    -> UV.Vector Double
    -> V.Vector (UV.Vector Double)
    -> V.Vector (Region2D, UV.Vector Double)
splitValidTri u v w =
    V.map fromJust .
    V.filter isJust . V.map ((fmap . flip (,)) `ap` testInRegionTri u v w)

minMaxXIndexes :: V.Vector (UV.Vector Double) -> (Int, Int)
minMaxXIndexes = minmax . sortOnFirst . vzip . V.map UV.unsafeHead
  where
    vzip v = V.zip v (V.enumFromN 0 $ V.length v)

minMaxYIndexes :: V.Vector (UV.Vector Double) -> (Int, Int)
minMaxYIndexes = minmax . sortOnFirst . vzip . V.map UV.unsafeLast
  where
    vzip v = V.zip v (V.enumFromN 0 $ V.length v)

minmax :: V.Vector (Double, Int) -> (Int, Int)
minmax = uncurry (,) . (snd . V.unsafeHead &&& snd . V.unsafeLast)

sortOnFirst
    :: Ord a
    => V.Vector (a, b) -> V.Vector (a, b)
sortOnFirst = V.modify (I.sortBy (comparing fst))

maxDistPt
    :: UV.Vector Double
    -> UV.Vector Double
    -> V.Vector (UV.Vector Double)
    -> (Int, UV.Vector Double)
maxDistPt mn mx =
    getMax . ap V.zip (V.map $ internalBisect mn mx . snd) . V.indexed
  where
    getMax = fst . V.maximumBy (comparing snd)

internalBisect :: UV.Vector Double
               -> UV.Vector Double
               -> UV.Vector Double
               -> Double
internalBisect u v w = 2 * sqrt (b * c * s * (b + c)) / (b + c + peps)
  where
    dist' x = sqrt . UV.sum . UV.map (join (*)) . UV.zipWith (-) x
    a = dist' u v
    b = dist' v w
    c = dist' u w
    s = a + b + c