{-# LANGUAGE GADTs #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE DefaultSignatures #-}
{-# LANGUAGE UndecidableInstances #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE DeriveFunctor #-}
{-# LANGUAGE DeriveGeneric #-}
{-# LANGUAGE InstanceSigs #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE PolyKinds #-}
{-# LANGUAGE KindSignatures #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE FunctionalDependencies #-}

{-# LANGUAGE AllowAmbiguousTypes #-}

module Expr where

import           Data.Void
import           Data.Tagged
import           Data.Proxy
import           Data.Constraint (Constraint)

import           Data.Bifunctor

import           GHC.Generics
import           Language.Haskell.TH (Name)

import           Data.Complex

import           GHC.Types

data Tagged2 (h :: k) (t :: *) (a :: *) = Tagged2 { unTagged2 :: a }

-- Idea: Represent a pattern match of a particular type as
--    KnownSymbol name => Proxy name -> MatchType name
-- with 'MatchType' being a type family that gives the type of the match
-- associated to matching on the constructor named 'name'. Maybe these
-- could be stored in a heterogenous list.

-- | Idea: Deep embedding for a pattern match (possibly using a non-regular
-- type?). Convert types to a "canonical form" to use this (like 'Either's
-- and '(,)'s)

-- TODO: Should we have a phantom type argument with the original type to preserve (possibly) more type safety?

-- TODO: Should this just be a function type? Maybe it could be a function
-- from s -> t, where is a nested pair type, and we use projection
-- combinators to extract the values inside the function.
data ProdMatch s t where
  ProdMatch ::
    (GPURep k1 a, GPURep k2 b)
       => (GPUExp (GetType a) -> ProdMatch b r) -> ProdMatch (Tagged2 (TypeHead t) t (a, b)) r

  OneProdMatch :: (GPURep k a) => (GPUExp (GetType a) -> GPUExp r) -> ProdMatch a r
  NullaryMatch :: GPURep k r => GPUExp r -> ProdMatch a r

data SumMatch s t where
  SumMatch ::
    (GPURep k1 a, GPURep k2 b, GPURepTy b ~ b)
        => ProdMatch a r -> SumMatch b r -> SumMatch (Tagged2 (TypeHead t) t (Either a b)) r

  EmptyMatch :: SumMatch Void r

  OneSumMatch :: (GPURep k1 a, GPURep k2 b, GPURepTy a ~ a) => ProdMatch a b -> SumMatch a b

type family TagMatches t x :: Constraint where
  -- TagMatches t (Tagged2 h t x) = ()
  TagMatches t (Tagged2 h t x) = TypeHead t ~ h

data SafeSumMatch origType s t where
  SafeSumMatch :: TagMatches origType s => Proxy (TypeHead s) -> SumMatch s t -> SafeSumMatch origType s t

  SafeEmptyMatch :: SafeSumMatch origType Void r
  SafeOneMatch :: (GPURep k1 a, GPURep k2 b, GPURepTy a ~ a) => ProdMatch a b -> SafeSumMatch origType a b

type family TagWith t x where
  TagWith t (Tagged2 h t' x) = Tagged2 (TypeHead t) t x
  TagWith t x                = Tagged2 (TypeHead t) t x

type family GetType x where
  GetType (Tagged2 h t x) = t
  GetType t            = t

type family Untag x where
  Untag (Tagged2 h t x) = x
  Untag x            = x

-- Done (case ... of A -> x; B -> y)  ==>  (case ... of A -> Done x; B -> Done y)
data Iter a b
  = Step b
  | Done a
  deriving (Functor, Generic)

data GPUExp t where
  CaseExp :: (GPURep k t) => GPUExp t -> SafeSumMatch t (GPURepTy t) r -> GPUExp r

  FalseExp :: GPUExp Bool
  TrueExp :: GPUExp Bool

  Repped :: GPURep k a => GPURepTy a -> GPUExp a

  Lam :: GPURep k a => Name -> GPUExp b -> GPUExp (a -> b)
  Var :: Proxy a -> Name -> GPUExp a -- Constructed internally

  -- Lam :: GPURep a =>  -> GPUExp (a -> b)

  Apply :: GPUExp (a -> b) -> GPUExp a -> GPUExp b

  Lit :: Num a => a -> GPUExp a

  Add :: Num a => GPUExp a -> GPUExp a -> GPUExp a
  Sub :: Num a => GPUExp a -> GPUExp a -> GPUExp a
  Mul :: Num a => GPUExp a -> GPUExp a -> GPUExp a

  Twice :: GPUExp (a -> a) -> GPUExp a -> GPUExp a

  FromEnum :: Enum a => GPUExp a -> GPUExp Int
  FromIntegral :: (Integral a, Num b) => GPUExp a -> GPUExp b

  Sqrt :: Floating a => GPUExp a -> GPUExp a

  Equal :: Eq a => GPUExp a -> GPUExp a -> GPUExp Bool
  Lte :: Ord a => GPUExp a -> GPUExp a -> GPUExp Bool
  Gt :: Ord a => GPUExp a -> GPUExp a -> GPUExp Bool

  Not :: GPUExp Bool -> GPUExp Bool

  LeftExp :: GPUExp a -> GPUExp (Either a b)
  RightExp :: GPUExp b -> GPUExp (Either a b)

  PairExp :: GPUExp a -> GPUExp b -> GPUExp (a, b)

  StepExp :: GPUExp b -> GPUExp (Iter a b)
  DoneExp :: GPUExp a -> GPUExp (Iter a b)

  IfThenElse :: GPUExp Bool -> GPUExp a -> GPUExp a -> GPUExp a

  TailRec :: (GPURep k1 a, GPURep k2 b) => (GPUExp b -> GPUExp (Iter a b)) -> GPUExp (b -> a)

  Construct :: a -> GPUExp a
  ConstructAp :: (GPURep k a) => GPUExp (a -> b) -> GPUExp a -> GPUExp b

class GPURep k t where
    -- Should we wrap these types in a 'Tagged t' in order to preserve more
    -- type safety?
  type GPURepTy t
  type GPURepTy t = TagWith t (GPURepTy (Rep t Void))

  -- | This should be unapplied type (without type arguments)
  type TypeHead t :: k
  -- type TypeHead t = t

  rep :: t -> GPUExp t
  rep = Repped . rep'

  rep' :: t -> GPURepTy t

  default rep' :: (Generic t, GenericRep (Rep t Void), GPURepTy t ~ Tagged2 (TypeHead t) t (GenericRepTy (Rep t Void)))
    => t -> GPURepTy t
  rep' = Tagged2 . genericRep' . (from :: t -> Rep t Void)

  unrep' :: GPURepTy t -> t

  default unrep' :: (Generic t, GenericRep (Rep t Void), GPURepTy t ~ Tagged2 (TypeHead t) t (GenericRepTy (Rep t Void)))
    => GPURepTy t -> t
  unrep' = (to :: Rep t Void -> t) . genericUnrep' . unTagged2

  repGetType :: t -> GPUExp (GetType t)

  default repGetType :: GetType t ~ t => t -> GPUExp (GetType t)
  repGetType = rep

instance (GPURep k t, h ~ TypeHead t, GPURepTy t ~ Tagged2 h t x) => GPURep k (Tagged2 (h :: k) t x) where
  type GPURepTy (Tagged2 h t x) = Tagged2 h t x
  type TypeHead (Tagged2 h t x) = h -- This instance is special

  rep' :: Tagged2 h t x -> GPURepTy (Tagged2 h t x)
  rep' = id

  unrep' :: GPURepTy (Tagged2 h t x) -> Tagged2 h t x
  unrep' = id

  repGetType :: Tagged2 h t x -> GPUExp (GetType (Tagged2 h t x))
  repGetType = rep . unrep'

instance GPURep * Int where
  type GPURepTy Int = Int
  type TypeHead Int = Int
  rep = Lit
  rep' = id
  unrep' = id
instance GPURep * Integer where
  type GPURepTy Integer = Integer
  type TypeHead Integer = Integer
  rep = Lit
  rep' = id
  unrep' = id
instance GPURep * Float where
  type GPURepTy Float = Float
  type TypeHead Float = Float
  rep = Lit
  rep' = id
  unrep' = id
instance GPURep * Double where
  type GPURepTy Double = Double
  type TypeHead Double = Double
  rep = Lit
  rep' = id
  unrep' = id

instance GPURep * Bool where
  type GPURepTy Bool = Bool
  type TypeHead Bool = Bool
  rep False = FalseExp
  rep True  = TrueExp
  rep' = id
  unrep' = id

instance GPURep * a => GPURep (* -> *) (Complex a) where
  type TypeHead (Complex a) = Complex

instance GPURep * a => GPURep (* -> *) (Maybe a) where
  type TypeHead (Maybe a) = Maybe

instance (GPURep * a, GPURep * b) => GPURep (* -> * -> *) (Either a b) where
  type GPURepTy (Either a b) = Tagged2 (Either) (Either a b) (Either (GPURepTy a) (GPURepTy b))
  type TypeHead (Either a b) = Either

  rep (Left x) = LeftExp (rep x)
  rep (Right y) = RightExp (rep y)

  rep' (Left x) = Tagged2 . Left $ rep' x
  rep' (Right y) = Tagged2 . Right $ rep' y

  unrep' (Tagged2 (Left x)) = Left $ unrep' x
  unrep' (Tagged2 (Right y)) = Right $ unrep' y

instance (GPURep * a, GPURep * b) => GPURep (* -> * -> *) (a, b) where
  type GPURepTy (a, b) = Tagged2 (,) (a, b) (GPURepTy a, GPURepTy b)
  type TypeHead (a, b) = (,)

  rep (x, y) = PairExp (rep x) (rep y)
  rep' (x, y) = Tagged2 (rep' x, rep' y)
  unrep' (Tagged2 (x, y)) = (unrep' x, unrep' y)

instance (GPURep * a, GPURep * b, GPURep * c) => GPURep (* -> * -> * -> *) (a, b, c) where
  type TypeHead (a, b, c) = (,,)
instance (GPURep * a, GPURep * b, GPURep * c, GPURep * d) =>
    GPURep (* -> * -> * -> * -> *) (a, b, c, d) where
  type TypeHead (a, b, c, d) = (,,,)

-- XXX: Should this instance exist?
instance (GPURep * a, GPURep * b) => GPURep (* -> * -> *) (Iter a b) where
  type TypeHead (Iter a b) = Iter

-- Generics instances
-- i (c :: Meta) (f :: k -> *) (p :: k)
instance GPURep * (f p) => GPURep (* -> Meta -> (k -> *) -> k -> *) (M1 i c f p) where
  type GPURepTy (M1 i c f p) = GPURepTy (f p)
  type TypeHead (M1 i c f p) = M1

  rep = Repped . rep'
  rep' (M1 x) = rep' x
  unrep' = M1 . unrep'

instance (GPURep k (p x), GPURep k (q x)) => GPURep ((k -> *) -> (k -> *) -> k -> *) ((p :+: q) x) where
  type GPURepTy ((p :+: q) x) = Tagged2 (:+:) (Either (p x) (q x)) (Either (GPURepTy (p x)) (GPURepTy (q x)))
  type TypeHead ((p :+: q) x) = (:+:)

  rep = Repped . rep'

  rep' (L1 x) = Tagged2 $ Left (rep' x)
  rep' (R1 y) = Tagged2 $ Right (rep' y)

  unrep' (Tagged2 (Left x)) = L1 (unrep' x)
  unrep' (Tagged2 (Right y)) = R1 (unrep' y)

instance (GPURep k (p x), GPURep k (q x)) => GPURep ((k -> *) -> (k -> *) -> k -> *) ((p :*: q) x) where
  type GPURepTy ((p :*: q) x) = Tagged2 (:*:) (p x, q x) (GPURepTy (p x), GPURepTy (q x))
  type TypeHead ((p :*: q) x) = (:*:)

  rep' (x :*: y) = Tagged2 (rep' x, rep' y)
  unrep' (Tagged2 (x, y)) = (unrep' x :*: unrep' y)

instance GPURep * c => GPURep (* -> * -> k -> *) (K1 i c p) where
  type GPURepTy (K1 i c p) = GPURepTy c
  type TypeHead (K1 i c p) = K1

  rep = Repped . rep'

  rep' (K1 x) = rep' x
  unrep' = K1 . unrep'

instance GPURep (k -> *) (U1 p) where
  type GPURepTy (U1 p) = ()
  type TypeHead (U1 p) = U1

  rep' U1 = ()
  unrep' () = U1

class GenericRep repr where
    type GenericRepTy repr

    genericRep' :: repr -> GenericRepTy repr
    genericUnrep' :: GenericRepTy repr -> repr

instance forall k p q. (GPURep k (p Void), GPURep k (q Void)) =>
  GenericRep ((p :+: q) Void) where

    type GenericRepTy ((p :+: q) Void) = Either (GPURepTy (p Void)) (GPURepTy (q Void))

    genericRep' = bimap rep' rep' . toCanonical
    genericUnrep' = fromCanonical . bimap unrep' unrep'

instance (GPURep k (p Void), GPURep k (q Void)) =>
  GenericRep ((p :*: q) Void) where

    type GenericRepTy ((p :*: q) Void) = (GPURepTy (p Void), GPURepTy (q Void))

    genericRep' = bimap rep' rep' . toCanonical
    genericUnrep' = fromCanonical . bimap unrep' unrep'

instance (GenericRep (f Void)) =>
  GenericRep (M1 i c f Void) where
    type GenericRepTy (M1 i c f Void) = GenericRepTy (f Void)

    genericRep' = genericRep' . unM1
    genericUnrep' = M1 . genericUnrep'

the :: a -> a
the = id

the_repr :: GPUExp a -> GPUExp a
the_repr = id

type family LiftedFn t where
  LiftedFn (a -> b) = GPUExp a -> LiftedFn b
  LiftedFn b = GPUExp b

construct :: (Construct t)
  => t -> LiftedFn t
construct = construct' . Construct

class Construct t where
    construct' :: GPUExp t -> LiftedFn t

instance (GPURep k a, Construct b) => Construct (a -> b) where
    construct' :: GPUExp (a -> b) -> GPUExp a -> LiftedFn b
    construct' c = construct' . ConstructAp c

instance {-# INCOHERENT #-} (LiftedFn b ~ GPUExp b) => Construct b where
    construct' :: GPUExp b -> GPUExp b
    construct' = id

-- Should this just produce an error?
sumMatchAbs :: GPURep k s => SafeSumMatch s (GPURepTy s) t -> s -> t
sumMatchAbs (SafeSumMatch Proxy s) = go s
  where
    go :: GPURep k s => SumMatch (GPURepTy s) t -> s -> t
    go (SumMatch p q) =
      \x0 ->
        let Tagged2 x = rep' x0
        in
        case x of
          Left  a -> prodMatchAbs p a
          Right b -> go q b
    go EmptyMatch = \_ -> error "Non-exhaustive match"
    go (OneSumMatch f) = prodMatchAbs f . unrep' . rep' -- TODO: Is this reasonable?

prodMatchAbs :: GPURep k s => ProdMatch s t -> s -> t
prodMatchAbs (ProdMatch f) =
  \(Tagged2 pair) ->
    case pair of
      (x, y) -> prodMatchAbs (f (repGetType x)) y

prodMatchAbs (OneProdMatch f) = \x -> gpuAbs (f (repGetType x))
prodMatchAbs (NullaryMatch x) = \_ -> gpuAbs x

gpuAbs :: GPUExp t -> t
gpuAbs (CaseExp x f) = sumMatchAbs f (gpuAbs x)
gpuAbs FalseExp = False
gpuAbs TrueExp  = True
gpuAbs (Repped x) = unrep' x
-- gpuAbs (Lam f) = gpuAbs . f . rep
-- gpuAbs (Apply f x) = gpuAbs f (gpuAbs x)
gpuAbs (Lit x)  = x
gpuAbs (Add x y) = gpuAbs x + gpuAbs y
gpuAbs (Sub x y) = gpuAbs x - gpuAbs y
gpuAbs (Mul x y) = gpuAbs x * gpuAbs y
gpuAbs (FromEnum x) = fromEnum (gpuAbs x)
gpuAbs (FromIntegral x) = fromIntegral (gpuAbs x)
gpuAbs (Sqrt x) = sqrt (gpuAbs x)
gpuAbs (Equal x y) = gpuAbs x == gpuAbs y
gpuAbs (Lte x y) = gpuAbs x <= gpuAbs y
gpuAbs (Gt x y) = gpuAbs x > gpuAbs y
gpuAbs (Not x) = not (gpuAbs x)
gpuAbs (LeftExp x) = Left (gpuAbs x)
gpuAbs (RightExp y) = Right (gpuAbs y)
gpuAbs (PairExp x y) = (gpuAbs x, gpuAbs y)
gpuAbs (StepExp x) = Step $ gpuAbs x
gpuAbs (DoneExp y) = Done $ gpuAbs y
gpuAbs (TailRec f) = \x ->
  case gpuAbs (f (rep x)) of
    Step x' -> gpuAbs (TailRec f) x'
    Done r  -> r
gpuAbs (IfThenElse cond t f)
  | gpuAbs cond = gpuAbs t
  | otherwise = gpuAbs f
gpuAbs (Construct x) = x
gpuAbs (ConstructAp f x) = gpuAbs f (gpuAbs x)

class Canonical t where
  type GenericOp t :: (* -> *) -> (* -> *) -> * -> *

  toCanonical :: GenericOp t p q x -> t (p x) (q x)
  fromCanonical :: t (p x) (q x) -> GenericOp t p q x

instance Canonical Either where
  type GenericOp Either = (:+:)

  toCanonical (L1 x) = Left x
  toCanonical (R1 y) = Right y

  fromCanonical (Left x) = L1 x
  fromCanonical (Right y) = R1 y

instance Canonical (,) where
  type GenericOp (,) = (:*:)

  toCanonical (x :*: y) = (x, y)
  fromCanonical (x, y) = x :*: y


import Test.Tasty
import Test.Tasty.HUnit

import Control.Lens
import Control.Monad

import qualified Data.ByteString as B
import Data.Default
import Data.Function
import Data.List
import Data.Ord
import Data.Maybe

import Linear

import Linear.Vector
import qualified Numeric.LinearAlgebra.HMatrix as HM
import qualified FrameGrabber
import StartFiducial as SF
import StartFiducialDebug
import TileMatcher
import TrackTracker as TT
import qualified TileMatcherDebug
import qualified Loop
import qualified Track
import Transform
import qualified OpenCV as CV
import TrackDebug
import OpenCV.Core.Types.Mat
import OpenCV.VideoIO.Types
import qualified Data.Vector as V

import Image
import qualified Video
import MaskDebug
import qualified SegmentPositionerTests
import FrameWriter

import Text.Printf

import System.IO.Unsafe ( unsafePerformIO )

main :: IO ()
main = defaultMain unitTests

unitTests = testGroup "Unit tests"
  [ Video.tests
  , SegmentPositionerTests.tests
  , startFiducialTests
  , tileMatcherTests
  , trackDebugTests
  , loopTests
  , trackTests
  , trackTrackingTests
  ]

startFiducialTests :: TestTree
startFiducialTests = testGroup "Start fiducial tests"
  [ testCase "Start fiducial position" testStartFiducialPosition
  , testCase "Start fiducial transform" testStartFiducialTransform
  , testCase "Start fiducial consistency" testStartFiducialConsistency
  , testStartFiducialIsRectangle
  ]

testStartFiducialPosition :: Assertion
testStartFiducialPosition = do
  points@(V2 center tip) <- fromJust <$> SF.findCenter idleNoCars
  renderImage "/tmp/drawCenter.png" $ drawArrow idleNoCars points

  (round <$> center) @?= V2 485 151
  (round <$> tip) @?= V2 383 430

testStartFiducialTransform :: Assertion
testStartFiducialTransform = do
  points <- fromJust <$> SF.findCenter idleNoCars
  roundTransform (transformFromVector points) @?= Track.transform idleNoCarsStart

video :: FilePath
video = "test/video/idle-no-cars-0-3-frames.mp4"

testStartFiducialConsistency :: Assertion
testStartFiducialConsistency = do
  (frames :: [FrameGrabber.TestMat]) <- FrameGrabber.getFrames video
  (points :: [Maybe (V2 (V2 Double))]) <- mapM SF.findCenter frames

  let (debugs :: [SF.FrameMat]) = zipWith drawArrow frames (map fromJust points)

  let renderFrame n = renderImage ("/tmp/testStartFiducial_" ++ show n ++ ".png")
  zipWithM_ renderFrame [0..] debugs

  let centers = fmap (^._x) (catMaybes points)
  let mean = sumV centers ^/ 3
  let deltas = fmap (mean ^-^) centers
  (< V2 1.0 1.0) <$> deltas @?= replicate 3 True

testStartFiducialIsRectangle :: TestTree
testStartFiducialIsRectangle = testGroup "isRectangle"
  [ testCase "perfect" $ isIt True [V2 5 0, V2 5 10, V2 0 10, V2 0 0]
  , testCase "close" $ isIt True [V2 4 0, V2 5 10, V2 0 13, V2 0 0]
  , testCase "way off" $ isIt False [V2 2 0, V2 5 10, V2 0 12, V2 0 0]
  ]
  where isIt b = (@?= b) . SF.isRectangle . V.fromList

tileMatcherTests :: TestTree
tileMatcherTests = testGroup "Tile matcher tests"
  [ testCase "Straight is a straight" tileMatcherStraight
  , testCase "Left is a left" tileMatcherLeft
  , testCase "Find track" tileMatcherFindTrack
  ]

tileMatcherStraight :: Assertion
tileMatcherStraight = do
  let straight = idleNoCarsTrack Loop.!! 1
  renderImage "/tmp/tileMatcherStraight.png" $ TileMatcherDebug.drawTileMasks idleNoCars straight

tileMatcherLeft :: Assertion
tileMatcherLeft = do
  let left = idleNoCarsTrack Loop.!! 2
  renderImage "/tmp/tileMatcherLeft.png" $ TileMatcherDebug.drawTileMasks idleNoCars left

tileMatcherFindTrack :: Assertion
tileMatcherFindTrack = do
  let foundTrack = TileMatcher.findTrack idleNoCars idleNoCarsStart
  renderImage "/tmp/tileMatcherTrackMask.png" $ drawTrackMask idleNoCars foundTrack
  foundTrack @?= idleNoCarsTrack

trackDebugTests :: TestTree
trackDebugTests = testGroup "TrackDebug tests"
  [ testCase "drawTrack" trackDebugArrows
  , testCase "outline" trackDebugOutline
  , testCase "Draw track mask" trackDrawTrackMask
  ]

trackDebugArrows :: Assertion
trackDebugArrows = renderImage "/tmp/trackArrows.png" $ drawTrackArrows idleNoCars idleNoCarsTrack

trackDebugOutline :: Assertion
trackDebugOutline = renderImage "/tmp/trackOutline.png" $ drawTrackOutline idleNoCars idleNoCarsTrack

trackDrawTrackMask :: Assertion
trackDrawTrackMask =  renderImage "/tmp/trackMask.png" $ drawTrackMask idleNoCars idleNoCarsTrack

loopTests :: TestTree
loopTests = testGroup "Loop tests"
  [ testCase "mkLoop single" mkLoopSingleTest
  , testCase "mkLoop multiple" mkLoopMultipleTest
  , testCase "mkLoop multiple second" mkLoopMultipleSecondTest
  , testCase "mkLoop multiple last" mkLoopMultipleLastTest
  , testCase "loop unfold" loopUnfold
  , testCase "loop unfold from middle" loopUnfoldMiddle
  ]

mkLoopSingleTest :: Assertion
mkLoopSingleTest = let
  loop = Loop.mkLoop [1]
  Loop.Loop end (Loop.Start one) next = loop
  in one @?= 1

mkLoopMultipleTest :: Assertion
mkLoopMultipleTest = let
  loop = Loop.mkLoop [1,2,3]
  Loop.Loop end (Loop.Start one) next = loop
  Loop.Loop _ (Loop.Node two) _ = next
  Loop.Loop _ (Loop.Node three) _ = end
  in [one, two, three] @?= [1,2,3]

mkLoopMultipleSecondTest :: Assertion
mkLoopMultipleSecondTest = let
  loop = Loop.mkLoop [1,2,3]
  Loop.Loop _ _ (Loop.Loop left x right) = loop
  Loop.Loop _ (Loop.Start one) _ = left
  Loop.Node two = x
  Loop.Loop _ (Loop.Node three) _ = right
  in [one, two, three] @?= [1,2,3]

mkLoopMultipleLastTest :: Assertion
mkLoopMultipleLastTest = let
  loop = Loop.mkLoop [1,2,3]
  Loop.Loop (Loop.Loop left x right) _ _ = loop
  Loop.Loop _ (Loop.Node two) _ = left
  Loop.Node three = x
  Loop.Loop _ (Loop.Start one) _ = right
  in [two, three, one] @?= [2,3, 1]

loopUnfold :: Assertion
loopUnfold = let
  list = [1, 2]
  loop = Loop.mkLoop list
  in Loop.unfold loop @?= list

loopUnfoldMiddle :: Assertion
loopUnfoldMiddle = let
  list = [1..6]
  loop = Loop.mkLoop list
  Loop.Loop _ _ second = loop
  in Loop.unfold second @?= list

trackTests :: TestTree
trackTests = testGroup "Track tests"
  [testCase "parseTrack" parseTestTrack
  , testCase "parseBadTrack" parseBadTrack
  , testCase "start" trackStart
  , testCase "nextSegment" trackNextSegment
  , testCase "parse to length" trackLength
  , testCase "scanl" trackScanl
  , testCase "shows" trackShow
  , testCase "moves" trackMoves
  , testCase "loops" trackLoops
  , testCase "transform" trackTransform
  ]

testTrack = fromJust $ Track.parseTrack Track.start "srrsrr"

parseTestTrack :: Assertion
parseTestTrack = let
  tiles = Track.tile <$> testTrack
  in tiles @?= Loop.mkLoop [Track.Straight, Track.Right, Track.Right, Track.Straight, Track.Right, Track.Right]

parseBadTrack :: Assertion
parseBadTrack = Track.parseTrack Track.start "fail" @?= Nothing

trackNextSegment :: Assertion
trackNextSegment = let
  (Track.Segment tile p t) = Track.nextSegment Track.start Track.Right
  in do
    tile @?= Track.Right
    p @?= V2 Track.trackLength 0
    getMat t @?= V2 Track.trackUnitVector (V2 0 1)

trackScanl :: Assertion
trackScanl = let
  [x, y, z] = scanl Track.nextSegment Track.start [Track.Right, Track.Right]
  (Track.Segment x_tile x_p x_t) = x
  (Track.Segment y_tile y_p y_t) = y
  (Track.Segment z_tile z_p z_t) = z
  in do
    x_tile @?= Track.Straight
    x_p @?= zero
    getMat x_t @?= V2 Track.trackUnitVector (V2 0 1)

    y_tile @?= Track.Right
    y_p @?= V2 Track.trackLength 0
    getMat y_t @?= V2 Track.trackUnitVector (V2 0 1)

    z_tile @?= Track.Right
    z_p @?= V2 (Track.trackLength + Track.cornerWidth) Track.cornerWidth
    getMat z_t @?= V2 (V2 0 1) (V2 (-1) 0)

trackStart :: Assertion
trackStart = let
  Loop.Loop _ (Loop.Start start) _  = testTrack
  Track.Segment startTile _ _ = start
  in startTile @?= Track.Straight

trackLength :: Assertion
trackLength = Loop.length testTrack  @?= 6

trackShow :: Assertion
trackShow = let string = concatMap (show . Track.tile) . Loop.unfold $ testTrack
            in string @?= "srrsrr"

trackMoves :: Assertion
trackMoves = let
  lengthNub = length . nubBy ((==) `on` Track.position) . Loop.unfold
  in lengthNub testTrack @?= Loop.length testTrack

trackLoops :: Assertion
trackLoops = let
  Loop.Loop _ (Loop.Node end) _ = Loop.prev testTrack
  Track.Segment tile p t = end
  in show (Track.nextSegment end Track.Straight) @?= show Track.start

trackTransform :: Assertion
trackTransform = do
  transformFromAngle eye 0 @?= mkTransform (V2 (V2 1 0) (V2 0 1))
  (round <$>) <$> getMat (transformFromAngle eye (pi / 2) ) @?= V2 (V2 0 1) (V2 (-1) 0)
  (round <$>) <$> getMat (transformFromAngle eye pi ) @?= V2 (V2 (-1) 0) (V2 0 (-1))
  (round <$>) <$> getMat (transformFromAngle eye (3* (pi/2)) ) @?= V2 (V2 0 (-1)) (V2 1 0)

  Transform.angleFromTransform (mkTransform (V2 (V2 1 0) (V2 0 1))) @?= 0
  Transform.angleFromTransform (mkTransform (V2 (V2 0 1) (V2 (-1) 0))) @?= (pi/2)
  Transform.angleFromTransform (transformFromAngle eye 0) @?= 0
  Transform.angleFromTransform (transformFromAngle eye (pi/2)) @?= (pi/2)
  Transform.angleFromTransform (transformFromAngle eye 2) @?= 2

trackTrackingTests :: TestTree
trackTrackingTests = testGroup "Track tracking"
  [ testCase "track" trackTrackingTrack
  , testCase "rotated" trackTracking
  -- , testCase "pertubations" trackTrackingPertubations
  ]

trackTrackingTrack :: Assertion
trackTrackingTrack = do
  track <- fromJust <$> TT.track idleNoCars
  renderImage "/tmp/trackTrackerOutline.png" $ drawTrackOutline idleNoCars track

trackTracking :: Assertion
trackTracking = do
  (frames :: [FrameGrabber.TestMat]) <- FrameGrabber.getFrames "test/video/idle-no-cars-0-3-frames.mp4"
  outlines <- mapM maybeOutline frames
  FrameWriter.writeFrames "/tmp/trackTracking.mov" outlines

trackTrackingPertubations :: Assertion
trackTrackingPertubations = do
  (frames :: [FrameGrabber.TestMat]) <- FrameGrabber.getFrames "test/video/track/pertubations.mov"
  outlines <- mapM maybeOutline frames
  FrameWriter.writeFrames "/tmp/trackTrackingPertubations.mov" outlines

maybeOutline :: FrameGrabber.TestMat -> IO FrameGrabber.TestMat
maybeOutline frame = do
  maybeTrack <- TT.track frame
  case maybeTrack of
    Nothing -> return frame
    Just track -> return $ drawTrackOutline frame track


{-# LANGUAGE ViewPatterns, MagicHash, BangPatterns #-}

module CVUtils where

import Codec.Image.DevIL
import Data.Array.Unboxed
import Math.Probably.MCMC
import Math.Probably.Sampler
import Math.Probably.FoldingStats
import Control.Monad.State.Strict
import System.Cmd
import System.Environment
import Data.Array.IO
import System.IO

import Data.Array.Unboxed
import Numeric.LinearAlgebra hiding (find)
import qualified Math.Probably.PDF as PDF
import Data.List
import Data.Maybe
import Control.Applicative
import Data.Ord

import Data.IORef
import System.IO.Unsafe
import Data.BitArray

import GHC.Exts
import GHC.Word

type R = Double

type Pos = (R,R)
type Image = UArray (Int, Int, Int) Word8
type MImage = IOUArray (Int, Int, Int) Word8
data BitImage = BitImage Int BitArray

mkBitImage :: Int -> Int -> (Int -> Int -> Bool) -> BitImage
mkBitImage hix hiy f = BitImage hiy $ bitArray (0,(hix)*(hiy)-1) [ (x*hiy+y, f x y) | x<- [0..hix-1], y<- [0..hiy-1]]

readBitImage :: BitImage -> Int -> Int -> Bool
readBitImage (BitImage hiy im) x y = unsafeLookupBit im (x*hiy+y)

{-track1 :: Image -> [Obj] -> Image -> Sampler Obj
track1 bgIm objs im =
   samplingImportanceResampling $ particleLike bgIm im objs -}

type Square = (Pos,Pos)

within ((lox,loy), (hix,hiy)) ox oy
       =    ox > min lox hix
         && ox < max lox hix
         && oy > min loy hiy
         && oy < max loy hiy

data Obj =
  Obj {vellen :: !R,
       sideDisp :: !R,
       objlen:: !R,
       posx :: !R,
       posy :: !R,
       rot :: !R  } deriving Show

fileroot = reverse . takeWhile (/='/') . reverse . takeWhile (/='.')

print2 x y = putStrLn $ x ++show y

smallest ws = foldl1' (min) $ map snd ws
largest ws = map snd $ sortBy (comparing snd) ws
sumWeights ws = sum $ map ( (subtr (smallest ws)) . snd) ws
cummWeightedSamples ws=
   let least = smallest ws
   in scanl (\(_,csum) (x,w) -> (x,csum+ (w-least))) (fst $ head ws ,0) ws

dropLosers ws =
  let topW = foldl1' (max) $ map snd ws
      p (x,w) = w>topW-100
  in sortBy (comparing snd) $ filter p ws

subtr x y = y - x

maxOn f xs = ceiling $ foldl1' max $ map f xs
minOn f xs = floor $ foldl1' min $ map f xs

data StaticParams = SP {noise :: !R,
                        wlength :: !R,
                        eccentric :: !R } deriving Show

dist :: R -> R -> Int ->Int -> R
dist  !cx !cy !x !y  = let !dx = cx-(unround x+0.5)
                           !dy = cy-(unround y+0.5)
                       in sqrt(dx*dx + dy*dy)

dist' :: R -> R -> Int ->Int -> R
dist' (D# cx) (D# cy) (I# x) (I# y)
       = let !dx = cx -## (int2Double# x+##half)
             !dy = cy -## (int2Double# y+##half)
             !(D# half) = 0.5
         in  D# (sqrtDouble# (dx*##dx +## dy*##dy))

gaussW8 :: R -> Word8 -> Word8 -> R
gaussW8 tau muw8 = lpdf . word8ToR
   where lpdf x = log (sqrt (tau/2.0*pi)) + (0.0-((x-mu)*(x-mu)*tau))
         mu = word8ToR  muw8
gaussR :: R-> R -> Word8 -> R
gaussR tau mu = lpdf . word8ToR
   where lpdf :: R -> R
         lpdf x = log (sqrt (tau/2.0*pi)) + (0.0-((x-mu)*(x-mu)*tau))

--no-noise versions
gaussW8nn :: R -> Word8 -> Word8 -> R
gaussW8nn tau muw8 = lpdf . word8ToR
   where lpdf x = let y = x - mu in negate $ y*y
         mu = word8ToR  muw8
gaussRnn :: R-> R -> Word8 -> R
gaussRnn tau mu = lpdf . word8ToR
   where lpdf :: R -> R
         lpdf x = let y = x - mu in negate $ y*y

word8ToR :: Word8 -> R
--word8ToR (W8# whash) = (/256) $ F# $ int2Float# $ word2Int# whash
word8ToR (W8# whash) = (/256) $ D# $ int2Double# $ word2Int# whash
realToW8 :: R -> Word8
realToW8 =  round . (*256)

unround :: Int -> R
--word8ToR (W8# whash) = (/256) $ F# $ int2Float# $ word2Int# whash
unround (I# whash) = D# $ int2Double# whash

markObjsOnImage :: [Obj] -> Image -> IO (Image)
markObjsOnImage objs im = do
    mutIm <- thaw im
    forM_ objs $ \(Obj _ _ _ cx cy _) -> do
          writeArray mutIm (round cy, round cx, 2) 255
          writeArray mutIm (round cy, round cx, 0) 0
          writeArray mutIm (round cy, round cx, 1) 0

    freeze (mutIm::MImage)

--updateBgIm :: Image ->

mkRed arr cx cy = do
          writeArray arr ( cy,
                             cx, 0) 255
--          writeArray arr ( cy,
--                             cx, 1) 0
--          writeArray arr ( cy,
--                             cx, 2) 0

markEllipse :: StaticParams -> Obj -> Image -> IO (Image)
markEllipse sp@(SP noise _ ecc) (Obj _ _ len px py rot) im = do
    mutIm <- thaw im

    let cx = round px
    let cy = round py
        f1x = px+(len*ecc)*cos rot
        f1y = py+(len*ecc)*sin rot
        f2x = px-(len*ecc)*cos rot
        f2y = py-(len*ecc)*sin rot
    let f x y = do -- print2 "f at" (x,y,dist (f1x,f1y) (x, y) + dist (f2x,f2y) (x, y) )
                   when (dist  f1x f1y   x  y  + dist  f2x f2y   x  y  < 2 * len) $ do
                      --print2 "red at " (x,y)
                      mkRed mutIm x y
                      return ()
    sequence_ [ f x y  |
                   x <- [ (cx::Int) -50.. cx +50],
                   y <- [ cy -50.. cy +50]]
    freeze (mutIm::MImage)

uniformLogPdf :: R -> R-> R-> R
uniformLogPdf from to = \x-> if x>=from && x <=to
                               then log $ realToFrac $ 1/(to-from )
                               else -1e20