{-# LANGUAGE FlexibleContexts #-}
module Image.Transform where

import           Data.Array.Repa              as R
import           Data.Array.Repa.Stencil      as R
import           Data.Array.Repa.Stencil.Dim2 as R
import           Data.DList                   as DL
import           Data.List                    as L
import           Data.Vector.Unboxed          as VU
import           Numeric.LinearAlgebra        as NL

-- factor = 2^n, n = 0,1,..
-- the first factor in the list corresponds to the inner-most (right-most) dimension.
{-# INLINE downsample #-}
downsample ::
     (Source s e, Shape sh) => [Int] -> R.Array s sh e -> R.Array D sh e
downsample factorList arr
  | L.all (== 1) factorList = delay arr
  | L.any (< 1) newDList = error "Downsample factors are too large."
  | otherwise =
    R.backpermute
      (shapeOfList newDList)
      (shapeOfList . L.zipWith (*) factorList . listOfShape)
      arr
  where
    dList = listOfShape . extent $ arr
    newDList = L.zipWith div dList factorList

{-# INLINE downsampleUnsafe #-}
downsampleUnsafe ::
     (Source s e, Shape sh) => [Int] -> R.Array s sh e -> R.Array D sh e
downsampleUnsafe factorList arr =
  R.backpermute newSh (shapeOfList . L.zipWith (*) factorList . listOfShape) arr
  where
    dList = listOfShape $ extent arr
    newSh = shapeOfList $ L.zipWith div dList factorList

{-# INLINE upsample #-}
upsample ::
     (Source s e, Shape sh, Num e) => [Int] -> R.Array s sh e -> R.Array D sh e
upsample factorList arr
  | L.all (== 1) factorList = delay arr
  | L.any (< 1) factorList =
    error $ "upsample: factor must be >= 1.\n" L.++ show factorList
  | otherwise =
    R.backpermuteDft
      (fromFunction
         (shapeOfList $ L.zipWith (*) (listOfShape . extent $ arr) factorList)
         (const 0))
      (\shape ->
         if L.all (== 0) . L.zipWith (flip mod) factorList . listOfShape $ shape
           then Just .
                shapeOfList . L.zipWith (flip div) factorList . listOfShape $
                shape
           else Nothing)
      arr

{-# INLINE crop #-}
crop ::
     (Source s e, Shape sh)
  => [Int]
  -> [Int]
  -> R.Array s sh e
  -> R.Array D sh e
crop start len arr
  | L.any (< 0) start ||
      L.or (L.zipWith3 (\x y z -> x > (z - y)) start len dList) =
    error $
    "Crop out of boundary!\n" L.++ show start L.++ "\n" L.++ show len L.++ "\n" L.++
    show dList
  | L.length start /= L.length len || L.length start /= L.length dList =
    error $
    "crop: dimension error. \n start: " L.++ show (L.length start) L.++ " len:" L.++
    show (L.length len) L.++
    " arr:" L.++
    show (L.length dList)
  | otherwise =
    R.backpermute
      (shapeOfList len)
      (shapeOfList . L.zipWith (+) start . listOfShape)
      arr
  where
    dList = listOfShape $ extent arr

{-# INLINE cropUnsafe #-}
cropUnsafe ::
     (Source s e, Shape sh)
  => [Int]
  -> [Int]
  -> R.Array s sh e
  -> R.Array D sh e
cropUnsafe start len =
  R.backpermute
    (shapeOfList len)
    (shapeOfList . L.zipWith (+) start . listOfShape)

{-# INLINE pad #-}
pad :: (Source s e, Shape sh) => [Int] -> e -> R.Array s sh e -> R.Array D sh e
pad newDims padVal arr
  | L.all (== 0) diff = delay arr
  | otherwise =
    backpermuteDft
      (fromFunction (shapeOfList dimList) (const padVal))
      (\sh' ->
         let idx = L.zipWith (-) (listOfShape sh') diff
          in if L.or (L.zipWith (\i j -> i < 0 || (i >= j)) idx oldDimList)
               then Nothing
               else Just $ shapeOfList idx)
      arr
  where
    oldDimList = listOfShape . extent $ arr
    dimList = L.zipWith max newDims oldDimList
    diff =
      L.zipWith
        (\a b ->
           if a - b <= 0
             then 0
             else if a - b == 1
                    then 1
                    else div (a - b) 2)
        newDims
        oldDimList

{-# INLINE normalizeValueRange #-}
normalizeValueRange ::
     (Source r e, Shape sh, Num e, Fractional e, Ord e, Unbox e)
  => (e, e)
  -> R.Array r sh e
  -> R.Array D sh e
normalizeValueRange (minVal, maxVal) arr
  | maxV == minV = delay arr
  | otherwise =
    R.map (\x -> (x - minV) / (maxV - minV) * (maxVal - minVal) + minVal) $ arr
  where
    vec = computeUnboxedS . delay $ arr
    minV = VU.minimum . toUnboxed $ vec
    maxV = VU.maximum . toUnboxed $ vec

{-# INLINE computeDerivativeS #-}
computeDerivativeS ::
     (R.Source r e, Num e, Fractional e) => R.Array r DIM2 e -> R.Array D DIM3 e
computeDerivativeS arr =
  L.foldl' R.append (R.extend (Z :. R.All :. R.All :. (1 :: Int)) arr) ds
  where
    xStencil =
      makeStencil2 3 3 $ \ix ->
        case ix of
          Z :. 0 :. (-1) -> Just (-1)
          Z :. 0 :. 1 -> Just 1
          _ -> Nothing
    yStencil =
      makeStencil2 3 3 $ \ix ->
        case ix of
          Z :. -1 :. 0 -> Just (-1)
          Z :. 1 :. 0 -> Just 1
          _ -> Nothing
    xyStencil =
      makeStencil2 3 3 $ \ix ->
        case ix of
          Z :. -1 :. -1 -> Just 1
          Z :. -1 :. 1 -> Just (-1)
          Z :. 1 :. -1 -> Just (-1)
          Z :. 1 :. 1 -> Just 1
          _ -> Nothing
    ds =
      L.map
        (\s ->
           extend (Z :. R.All :. R.All :. (1 :: Int)) . R.map (/ 2) $
           mapStencil2 BoundClamp s arr)
        [xStencil, yStencil, xyStencil]

{-# INLINE bicubicInterpolation #-}
bicubicInterpolation ::
     (R.Source r Double)
  => ((Int, Int) -> (Double, Double))
  -> (Int, Int)
  -> R.Array r DIM2 Double
  -> R.Array U DIM2 Double
bicubicInterpolation idxFunc (newCols, newRows) arr =
  let (Z :. cols :. rows) = extent arr
      derivative = computeDerivativeS arr
      derivative16 =
        R.traverse
          derivative
          (const (Z :. newCols :. newRows :. 4 :. (4 :: Int)))
          (\f (Z :. a :. b :. derivativeIdx :. pairIdx) ->
             let (x, y) = idxFunc (a, b)
                 (x', y') =
                   case pairIdx of
                     0 -> (floor x, floor y)
                     1 -> (floor x, ceiling y)
                     2 -> (ceiling x, floor y)
                     3 -> (ceiling x, ceiling y)
              in if (x' < 0) ||
                    (x' > (fromIntegral cols - 1)) ||
                    (y' < 0) || (y' > (fromIntegral rows - 1))
                   then 0
                   else f (Z :. x' :. y' :. derivativeIdx))
      mat = ((newRows * newCols) >< 16) . R.toList $ derivative16
      multiplicationResultArray =
        fromListUnboxed (Z :. newCols :. newRows :. 4 :. 4) .
        NL.toList . flatten $
        mat NL.<> maxtrixA
   in R.sumS . R.sumS $
      R.traverse
        multiplicationResultArray
        id
        (\f idx@(Z :. a :. b :. i :. j) ->
           let (x, y) = idxFunc (a, b)
               x' = x - fromIntegral (floor x :: Int)
               y' = y - fromIntegral (floor y :: Int)
            in f idx * (x' ^ i) * (y' ^ j))
  where
    maxtrixA =
      NL.fromColumns . L.map NL.fromList $
      [ [1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]
      , [0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]
      , [-3, 3, 0, 0, -2, -1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]
      , [2, -2, 0, 0, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]
      , [0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0]
      , [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0]
      , [0, 0, 0, 0, 0, 0, 0, 0, -3, 3, 0, 0, -2, -1, 0, 0]
      , [0, 0, 0, 0, 0, 0, 0, 0, 2, -2, 0, 0, 1, 1, 0, 0]
      , [-3, 0, 3, 0, 0, 0, 0, 0, -2, 0, -1, 0, 0, 0, 0, 0]
      , [0, 0, 0, 0, -3, 0, 3, 0, 0, 0, 0, 0, -2, 0, -1, 0]
      , [9, -9, -9, 9, 6, 3, -6, -3, 6, -6, 3, -3, 4, 2, 2, 1]
      , [-6, 6, 6, -6, -3, -3, 3, 3, -4, 4, -2, 2, -2, -2, -1, -1]
      , [2, 0, -2, 0, 0, 0, 0, 0, 1, 0, 1, 0, 0, 0, 0, 0]
      , [0, 0, 0, 0, 2, 0, -2, 0, 0, 0, 0, 0, 1, 0, 1, 0]
      , [-6, 6, 6, -6, -4, -2, 4, 2, -3, 3, -3, 3, -2, -1, -2, -1]
      , [4, -4, -4, 4, 2, 2, -2, -2, 2, -2, 2, -2, 1, 1, 1, 1]
      ]

{-# INLINE bilinearInterpolation #-}
bilinearInterpolation ::
     (R.Source r Double)
  => ((Int, Int) -> (Double, Double))
  -> (Int, Int)
  -> R.Array r DIM2 Double
  -> R.Array U DIM2 Double
bilinearInterpolation idxFunc (newCols, newRows) arr =
  let (Z :. cols :. rows) = extent arr
  in computeS . R.traverse arr id $ \f (Z :. a :. b) ->
       let (x, y) = idxFunc (a, b)
           x1 = floor x
           x2 = ceiling x
           y1 = floor y
           y2 = ceiling y
           q11 = f (Z :. x1 :. y1)
           q12 = f (Z :. x1 :. y2)
           q21 = f (Z :. x2 :. y1)
           q22 = f (Z :. x2 :. y2)
           r1 =
             ((fromIntegral x2 - x) / fromIntegral (x2 - x1)) * q11 +
             ((x - fromIntegral x1) / fromIntegral (x2 - x1)) * q21
           r2 =
             ((fromIntegral x2 - x) / fromIntegral (x2 - x1)) * q12 +
             ((x - fromIntegral x1) / fromIntegral (x2 - x1)) * q22
           p =
             ((fromIntegral y2 - y) / fromIntegral (y2 - y1)) * r1 +
             ((y - fromIntegral y1) / fromIntegral (y2 - y1)) * r2
       in if (x < 0) ||
             (x > (fromIntegral cols - 1)) ||
             (y < 0) || (y > (fromIntegral rows - 1))
            then 0
            else if x1 == x2
                   then if y1 == y2
                          then q11
                          else ((fromIntegral y2 - y) / fromIntegral (y2 - y1)) *
                               q11 +
                               ((y - fromIntegral y1) / fromIntegral (y2 - y1)) *
                               q12
                   else if y1 == y2
                          then r1
                          else p

{-# INLINE resize2D #-}
resize2D ::
     (Source s Double)
  => (Int, Int)
  -> (Double, Double)
  -> R.Array s DIM2 Double
  -> R.Array D DIM2 Double
resize2D newSize@(newNx, newNy) bound arr =
  normalizeValueRange bound $
  bicubicInterpolation
    (\(i, j) -> (ratioX * fromIntegral i, ratioY * fromIntegral j))
    newSize
    arr
  where
    (Z :. nx' :. ny') = extent arr
    ratioX = fromIntegral nx' / fromIntegral newNx
    ratioY = fromIntegral ny' / fromIntegral newNy

{-# INLINE resize25D #-}
resize25D ::
     (Source s Double)
  => (Int, Int)
  -> (Double, Double)
  -> R.Array s DIM3 Double
  -> R.Array D DIM3 Double
resize25D newSize@(newNx, newNy) bound arr =
  normalizeValueRange bound .
  fromUnboxed (Z :. nf' :. newNx :. newNy) .
  VU.concat .
  L.map
    (\i ->
       toUnboxed .
       bicubicInterpolation
         (\(i, j) -> (ratioX * fromIntegral i, ratioY * fromIntegral j))
         newSize .
       R.slice arr $
       (Z :. i :. R.All :. R.All)) $
  [0 .. nf' - 1]
  where
    (Z :. nf' :. nx' :. ny') = extent arr
    ratioX = fromIntegral nx' / fromIntegral newNx
    ratioY = fromIntegral ny' / fromIntegral newNy

{-# INLINE resize2DFixedRatio #-}
resize2DFixedRatio ::
     (Source s Double)
  => Int
  -> (Double, Double)
  -> R.Array s DIM2 Double
  -> R.Array D DIM2 Double
resize2DFixedRatio n bound arr =
  normalizeValueRange bound $
  bicubicInterpolation
    (\(i, j) -> (ratio * fromIntegral i, ratio * fromIntegral j))
    newSize
    arr
  where
    (Z :. nx' :. ny') = extent arr
    newSize =
      if nx' > ny'
        then (n, round $ fromIntegral ny' * ratio)
        else (round $ fromIntegral nx' * ratio, n)
    ratio = fromIntegral (max nx' ny') / fromIntegral n

{-# INLINE resize25DFixedRatio #-}
resize25DFixedRatio ::
     (Source s Double)
  => Int
  -> (Double, Double)
  -> R.Array s DIM3 Double
  -> R.Array D DIM3 Double
resize25DFixedRatio n bound arr =
  normalizeValueRange bound .
  fromUnboxed (Z :. nf' :. newNx :. newNy) .
  VU.concat .
  L.map
    (\i ->
       toUnboxed .
       bicubicInterpolation
         (\(i, j) -> (ratio * fromIntegral i, ratio * fromIntegral j))
         newSize .
       R.slice arr $
       (Z :. i :. R.All :. R.All)) $
  [0 .. nf' - 1]
  where
    (Z :. nf' :. nx' :. ny') = extent arr
    newSize@(newNx, newNy) =
      if nx' > ny'
        then (n, round $ fromIntegral ny' * ratio)
        else (round $ fromIntegral nx' * ratio, n)
    ratio = fromIntegral (max nx' ny') / fromIntegral n

{-# INLINE rescale2D #-}
rescale2D ::
     (Source s Double)
  => Double
  -> (Double, Double)
  -> R.Array s DIM2 Double
  -> R.Array D DIM2 Double
rescale2D scale bound arr =
  normalizeValueRange bound $
  bicubicInterpolation
    (\(i, j) -> (fromIntegral i / scale, fromIntegral j / scale))
    (round $ scale * fromIntegral nx', round $ scale * fromIntegral ny')
    arr
  where
    (Z :. nx' :. ny') = extent arr

{-# INLINE rescale25D #-}
rescale25D ::
     (Source s Double)
  => Double
  -> (Double, Double)
  -> R.Array s DIM3 Double
  -> R.Array D DIM3 Double
rescale25D scale bound arr =
  normalizeValueRange bound .
  fromUnboxed (Z :. nf' :. newNx :. newNy) .
  VU.concat .
  L.map
    (\i ->
       toUnboxed .
       bicubicInterpolation
         (\(i, j) -> (fromIntegral i / scale, fromIntegral j / scale))
         (newNx, newNy) .
       R.slice arr $
       (Z :. i :. R.All :. R.All)) $
  [0 .. nf' - 1]
  where
    (Z :. nf' :. nx' :. ny') = extent arr
    newNx = round $ scale * fromIntegral nx'
    newNy = round $ scale * fromIntegral ny'

{-# INLINE rotate2D #-}
rotate2D ::
     (Source s Double)
  => Double
  -> (Double, Double)
  -> R.Array s DIM2 Double
  -> R.Array U DIM2 Double
rotate2D theta (centerX, centerY) arr =
  bicubicInterpolation
    (\(i, j) ->
       let i' = fromIntegral i - centerX
           j' = fromIntegral j - centerY
        in ( i' * cos theta - j' * sin theta + centerX
           , i' * sin theta + j' * cos theta + centerY))
    (nx', ny')
    arr
  where
    (Z :. nx' :. ny') = extent arr

{-# INLINE rotate25D #-}
rotate25D ::
     (Source s Double)
  => Double
  -> (Double, Double)
  -> R.Array s DIM3 Double
  -> R.Array U DIM3 Double
rotate25D theta (centerX, centerY) arr =
  fromUnboxed (Z :. nf' :. nx' :. ny') .
  VU.concat .
  L.map
    (\i ->
       toUnboxed .
       -- bicubicInterpolation
       bilinearInterpolation
         (\(i, j) ->
            let i' = fromIntegral i - centerX
                j' = fromIntegral j - centerY
            in ( i' * cos theta - j' * sin theta + centerX
               , i' * sin theta + j' * cos theta + centerY))
         (nx', ny') .
       R.slice arr $
       (Z :. i :. R.All :. R.All)) $
  [0 .. nf' - 1]
  where
    (Z :. nf' :. nx' :. ny') = extent arr


{-|
Module: Class.Mappable
Description: Map functions over data as polymorphically as possible.
License: 0BSD
Maintainer: wygulmage@users.noreply.github.com
Stability: Experimental
Portability: GHC

'Mappable' lets you 'map' over various structures. It's intended to unify all the functions named @map@ ('Data.List.map', 'Data.Set.map', 'Data.Text.map', etc.). To do this it uses type families.
'Data.Functor.Functor' and 'Data.MonoTraversable.MonoFunctor' provide simpler interfaces. But, 'Functor' 'fmap' requires a fully polymorphic container, while 'Data.MonoTraversable.MonoFunctor' 'Data.MonoTraversable.omap' can only map with endofunctions (and neither provides 'mapCoerce' yet).
-}

{-# LANGUAGE NoImplicitPrelude
           , ScopedTypeVariables
           , TypeFamilies
           , UndecidableInstances
           , FlexibleContexts
           , FlexibleInstances
           , DefaultSignatures
           , RankNTypes
           , MultiParamTypeClasses
           , GeneralizedNewtypeDeriving
   #-}

module Class.Mappable (
Mappable (..), omap,
ValidMorph,
WrapMappable (..),
WrapFunctor (..),
) where

import Control.Arrow (Arrow)
import Control.Applicative
import Control.Monad (Monad)
import Control.Monad.ST (ST)
import qualified Control.Exception as Exception
import Data.Ord
import Data.Semigroup
import qualified Data.Semigroup as Semigroup
import qualified Data.Monoid as Monoid
import qualified Data.Functor as Data
import Data.Bool (Bool)
import Data.Char
import Data.Int (Int)
import Data.Word (Word8)
import qualified Data.ByteString as Bytes
import qualified Data.ByteString.Lazy as Lazy.Bytes
import qualified Data.Text as Text
import qualified Data.Text.Lazy as Lazy.Text
import qualified Data.Sequence as Seq
import qualified Data.Tree as Rose
import qualified Data.Set as Set
import qualified Data.IntSet as IntSet
import qualified Data.Map as Map
import qualified Data.IntMap as IntMap
import Data.Complex
import Data.Tagged
import Data.Maybe (Maybe)
import Data.Either
import Data.Proxy
import Data.Functor.Identity
import qualified Data.List.NonEmpty as NonEmpty
import Data.Coerce (Coercible, coerce)
import Data.Kind

import GHC.Arr (Array)
import GHC.Conc (STM)

import Prelude (IO)

import Text.ParserCombinators.ReadP (ReadP)
import Text.ParserCombinators.ReadPrec (ReadPrec)

newtype WrapFunctor m a = WrapFunctor (m a)
   deriving (Data.Functor, Applicative, Alternative, Monad)

newtype WrapMappable t = WrapMappable t

type family ObjectDefault (ma :: Type) where ObjectDefault (m (a :: Type)) = a
{- ^ the rightmost parameter of a type -}

type family MorphDefault (x :: Type) (y :: Type) where
   MorphDefault (polymorphic a) b = polymorphic b
   MorphDefault monomorphic b     = monomorphic
{- ^ Try to replace the rightmost parameter of a type. If the type is polymorphic, the rightmost parameter is replaced; otherwise the original type is returned. -}

class    (Object (Morph mb a) ~ a)=> ValidMorph (mb :: Type) (a :: Type)
instance (Object (Morph mb a) ~ a)=> ValidMorph (mb :: Type) (a :: Type)

class (Morph mb (Object mb) ~ mb)=> Mappable (mb :: Type) where
{-^ @Mappable@ extends the power of 'Functor' to monomorphic containers and containers that constrain their contents.
To create an instance of @Mappable@ for a 'Functor', write e.g. @instance Mappable [b]@.
To create an instance of @Mappable@ for a monomorphic type, write e.g.
@instance Mappable 'Text.Text' where
   type Object 'Text.Text' = 'Char'
   'map' = 'Text.map'
@
To create an instance for a type that constraints its objects, you need to include a definition of 'mapCoerce'.
@instance (Ord b)=> Mappable ('Set.Set' b) where
   'map' = 'Set.map'
   'mapCoerce' = 'map' 'coerce'
@
To define an instance where the 'Object' is phantom, use a kind annotation: @instance Mappable ('Proxy' (o :: 'Type'))@.

Laws:
   1. for all f g. @map (g '$!') . map f@ === @map ((g '$!') . f)@, as long as equal ('==') objects produced by @f@ are indistinguishable to 'g'.
   2. @map 'id'@ === 'id'
   3. 'mapCoerce' === @'map' 'coerce'@
The first law forces @g@'s argument to accommodate both strict and lazy containers, and has the condition on equal objects to accommodate both 'Arg' and 'Set.Set'.
-}
   type Morph mb (a :: Type) :: Type
   {-^ Try to change the type of an object in a container. @Morph [a] b = [b]@; @Morph 'Text.Text' b = 'Text.Text'. -}
   type Morph mb a = MorphDefault mb a
   type Object mb :: Type
   {-^ The type of the objects in the container. By default it's the type's rightmost parameter. -}
   type Object mb = ObjectDefault mb

   map ::
      forall a. (ValidMorph mb a)=>
      (a -> Object mb) -> Morph mb a -> mb
   {-^ Map a function over all the 'Object's in a container. -}

   mapCoerce ::
      forall a. (ValidMorph mb a, Coercible a (Object mb))=>
      Morph mb a -> mb
   {-^ Coerce all the 'Object's in a container to a new type. -}
   -- mapCoerce = map coerce

   default map ::
      (ValidMorph mb a, Morph mb a ~ m a, m (Object mb) ~ mb, Data.Functor m)=>
      (a -> Object mb) -> Morph mb a -> mb
   map = Data.fmap
   {-# INLINE map #-}

   default mapCoerce ::
      (ValidMorph mb a, Coercible (Morph mb a) mb)=> Morph mb a -> mb
   mapCoerce = coerce
   {-# INLINE mapCoerce #-}

omap :: (Mappable t)=> (Object t -> Object t) -> t -> t
omap = map
{-# INLINE omap #-}

------ Instances ------

-- TODO: Add instances for lifted `Functor`s (e.g. `Ap m b`). Should these rely on the `Functor` instance or the `Mappable` instance? In favor of `Mappble` is being able to use constrained containers like `Set`.

--- State Functors ---
instance Mappable (IO b)
instance Mappable (ST s b)
instance Mappable (STM b)

--- Empty Containers ---
instance Mappable (Proxy (b :: Type))

--- Identity Containers ---
instance Mappable (Identity b)
instance Mappable (Down b)
instance Mappable (Dual b)
instance Mappable (Semigroup.First b)
instance Mappable (Semigroup.Last b)
instance Mappable (Max b)
instance Mappable (Min b)
instance Mappable (Product b)
instance Mappable (Sum b)

instance Mappable Semigroup.All where
   type Object Semigroup.All = Bool
   map = coerce

instance Mappable Semigroup.Any where
   type Object Semigroup.Any = Bool
   map = coerce

--- Affine Containers ---
instance Mappable (Maybe b)
instance Mappable (Monoid.First b)
instance Mappable (Monoid.Last b)

--- Fixed-size Vectors ---
instance Mappable (Complex b)

--- Sequences ---
instance Mappable [b]
instance Mappable (ZipList b)

instance Mappable (NonEmpty.NonEmpty b)

instance Mappable (Seq.Seq b)
instance Mappable (Seq.ViewL b)
instance Mappable (Seq.ViewR b)

instance Mappable (Rose.Tree b)

instance Mappable (Array i b)

instance Mappable Bytes.ByteString where
   type Object Bytes.ByteString = Word8
   map = Bytes.map

instance Mappable Lazy.Bytes.ByteString where
   type Object Lazy.Bytes.ByteString = Word8
   map = Lazy.Bytes.map

instance Mappable Text.Text where
   type Object Text.Text = Char
   map = Text.map

instance Mappable Lazy.Text.Text where
   type Object Lazy.Text.Text = Char
   map = Lazy.Text.map

--- Sets ---
instance (Ord b)=> Mappable (Set.Set b) where
   map = Set.map
   mapCoerce = map coerce
{-^ WARNING: @map (g $!) . map f@ is not equivalent to @map ((g $!) . f)@ when @g@ can distinguish between different but equal ('==') values produced by @f@. For example, @map (\ ('Arg' () n) -> n) . map ('Arg' ())@ is likely to give a very different result than @map ((\ ('Arg' () n) -> n) . 'Arg' ())@. -}

instance Mappable IntSet.IntSet where
   type Object IntSet.IntSet = Int
   map = IntSet.map

--- Bifunctors ---
instance Mappable (Const c (b :: Type))

instance Mappable (Tagged c b)

instance Mappable (Either c b)

instance Mappable (c, b)
instance Mappable (Arg c b)

--- Functions ---
instance Mappable (a -> b)

instance Mappable (Exception.Handler b)

instance Mappable (ReadP b)
instance Mappable (ReadPrec b)

--- Maps ---
instance Mappable (Map.Map i b)
instance Mappable (IntMap.IntMap b)

--- Wrapped Instances ---
instance (Arrow p)=> Mappable (WrappedArrow p c b) where
   mapCoerce = map coerce

instance (Data.Functor m)=> Mappable (WrapFunctor m a) where
   mapCoerce = map coerce

instance (Mappable t)=> Mappable (WrapMappable t) where
   type Object (WrapMappable t) = Object t
   type Morph (WrapMappable t) a = WrapMappable (Morph t a)
   map f (WrapMappable xs) = WrapMappable (map f xs)
   mapCoerce (WrapMappable xs) = WrapMappable (mapCoerce xs)


-- |
-- Module      : Southpaw.Picasso.Render
-- Description : Vector rendering in 2D with Cairo
-- Copyright   : (c) Jonatan H Sundqvist, 2015
-- License     : MIT
-- Maintainer  : Jonatan H Sundqvist
-- Stability   : experimental|stable
-- Portability : POSIX (not sure)
--

-- Created September 1 2015

-- TODO | - Rename (eg. Draw, Shapes) (?)
--        - Text, typography (find good library)
--        - Debug versions (cf. Occlusion.Render)
--        - Use BoundingBox (?)

-- SPEC | -
--        -

--------------------------------------------------------------------------------------------------------------------------------------------
-- GHC directives
--------------------------------------------------------------------------------------------------------------------------------------------
{-# LANGUAGE TupleSections #-}

--------------------------------------------------------------------------------------------------------------------------------------------
-- API
--------------------------------------------------------------------------------------------------------------------------------------------
module Southpaw.Picasso.Render where

--------------------------------------------------------------------------------------------------------------------------------------------
-- We'll need these
--------------------------------------------------------------------------------------------------------------------------------------------
import Data.Complex
import Control.Monad (forM, forM_, liftM, liftM2, void)

import Control.Monad.IO.Class

import qualified Graphics.Rendering.Cairo                         as Cairo
-- import qualified Graphics.Rendering.Cairo.Internal.Surfaces.Image as Image

import qualified Southpaw.Picasso.Palette     as Palette
import qualified Southpaw.Picasso.Shapes      as Shapes
import           Southpaw.Picasso.RenderUtils
import           Southpaw.Math.Constants
import           Southpaw.Cartesian.Plane.Utilities

--------------------------------------------------------------------------------------------------------------------------------------------
-- Functions
--------------------------------------------------------------------------------------------------------------------------------------------
-- General ---------------------------------------------------------------------------------------------------------------------------------
-- |
-- TODO: Rename (eg. plot)
-- TODO: Generalise (not just circles)
trail :: Palette.Colour Double -> [Complex Double] -> Cairo.Render ()
trail fill trail = forM_ trail $ \dot -> do
    choose fill
    circle dot 3

-- |
grid :: Int -> Int -> Double -> Cairo.Render ()
grid cols rows size = do
    -- TODO: Figure out how to use fill AND stroke
    Cairo.setLineWidth 4
    gridM_ cols rows $ \ cl rw -> tilePath cl rw >> Cairo.fill   --
    gridM_ cols rows $ \ cl rw -> tilePath cl rw >> Cairo.stroke --
    where
      chooseColour cl rw = if (cl `mod` 2) == (rw `mod` 2) then 0.3 else 0.75 -- TODO: This should be a utility function
      tilePath cl rw     = Cairo.rectangle (fromIntegral cl*size) (fromIntegral rw*size) size size >> Cairo.setSourceRGBA 0.22 0.81 (chooseColour cl rw) 0.32

-- Primitives ------------------------------------------------------------------------------------------------------------------------------

-- |
line :: Complex Double -> Complex Double -> Cairo.Render ()
line (fr:+om) (t:+o) = Cairo.moveTo fr om >> Cairo.lineTo t o

-- | Renders a path of connected lines
-- TODO: Options for colour, width, closed/open, etcairo.
linepath :: [Complex Double] -> Cairo.Render ()
linepath []      = return ()
linepath (e:dge) = void $ vectorise Cairo.moveTo e >> forM dge (vectorise Cairo.lineTo)

-- |
-- TODO: Support asymmetrical crosshairs (?)
crosshairs :: Complex Double -> Complex Double -> Cairo.Render ()
crosshairs centre size = do
  line (centre - (hdx:+0)) (centre + (hdx:+0))
  line (centre - (0:+hdy)) (centre + (0:+hdy))
  where
    (hdx:+hdy) = 0.5*size

-- Shapes ----------------------------------------------------------------------------------------------------------------------------------

-- |
-- TODO: Extract argument conversion logic (centre/size vectors to unpacked left-top/dx/dy)
rectangle :: Complex Double -> Complex Double -> Cairo.Render ()
rectangle (cx:+cy) (dx:+dy) = Cairo.rectangle (cx-dx/2) (cy-dy/2) dx dy

-- | A rectangle with an 'anchor' (in normalised coordinates) that is relative to the centre
-- TODO: Should anchor point be relative to centre or topleft corner, use normalised or absolute coords (?)
-- TODO: Less confusing terminology...
-- TODO: Refactor
-- TODO: Test
anchoredRectangle :: Complex Double -> Complex Double -> Complex Double -> Cairo.Render ()
anchoredRectangle p size anchor = rectangle (p-dotwise (*) (anchor+(0.5:+0.5)) size) size

-- |
-- TODO: Add arguments for colour, stroke, etcairo.
-- TODO: Maybe it'd be better if we stuck to the normal pattern of path-config-action that Cairo follows
-- TODO: Make polymorphic
polygon :: Integral int => int -> Double -> Complex Double -> (Double, Double, Double, Double) -> Bool -> Cairo.Render ()
polygon sides radius origin (r,g,b,a) filled = do
    -- TODO: Refine 'wrap-around logic'
    Cairo.moveTo fx fy
    forM_ rest $ \(x:+y) -> Cairo.lineTo x y

    Cairo.setSourceRGBA r g b a
    Cairo.setLineWidth 12
    if filled
        then Cairo.fill
        else Cairo.stroke
    where
      ((fx:+fy):rest) = Shapes.polygon sides radius origin ++ [fx:+fy]

-- |
-- TODO: Options for fill/stroke, colour, width, etcairo.
circle :: Complex Double -> Double -> Cairo.Render ()
circle (cx:+cy) radius = do
    Cairo.arc cx cy radius 0 τ
    -- Cairo.fill

-- Composite -------------------------------------------------------------------------------------------------------------------------------

-- |
arrow :: Complex Double -> Complex Double -> Double -> Double -> Double -> Cairo.Render ()
arrow from to sl sw hw = do
    let (first:rest) = closePath $ Shapes.arrow from to sl sw hw
    vectorise Cairo.moveTo first
    forM_ rest $ vectorise Cairo.lineTo

-- |
-- Ugh, I hate underscores so much
-- TODO: Make polymorphic
circlearc :: Int -> Complex Double -> Double -> Double -> Double -> Double -> Cairo.Render ()
circlearc
    count    -- Number of small circles
    (ox:+oy) -- Centre of the 'arc' (pixels?)
    spread   -- Radius of the 'arc'
    radius   -- Radius of the small circles
    begin    -- Start angle of the arc
    extent = forM_ [1..count] $ \ n -> do
        let n' = fromIntegral n
        let θ  = begin + n'*extent/fromIntegral count
        Cairo.arc (ox - spread*cos θ) (oy - spread*sin θ) radius 0 τ
        Cairo.setSourceRGBA (0.5 * (1 + sin θ)) (0.1*n') (1/n') 0.95
        Cairo.fill

-- |
bezier :: Complex Double -> Complex Double -> Complex Double -> Cairo.Render ()
bezier (x1:+y1) (x2:+y2) (x3:+y3) = Cairo.curveTo x1 y1 x2 y2 x3 y3

-- |
-- TODO: Generic rounded polygon
roundrect :: Complex Double -> Complex Double -> Double -> Cairo.Render ()
roundrect centre@(cx:+cy) size@(dx:+dy) radius = forM_ (zip [real, imag, real, imag] [(-dx):+(-dy), (dx):+(-dy), (dx):+(dy), (-dx):+(dy)]) $ \(dir, delta@(dx':+dy')) -> do

  -- TODO: Finish refactoring

	--
	-- let dir = (signum (dx*dy))

	-- Line segment
	vectorise Cairo.moveTo (centre + delta + dir radius)
	vectorise Cairo.lineTo (centre - flipx size/2 - real radius)

	-- Curve

	-- -- First line segment
	-- vectorise Cairo.moveTo (centre - size/2       + real radius)
	-- vectorise Cairo.lineTo (centre - flipx size/2 - real radius)

	-- -- Curve
	-- let (cx':+cy') = (centre - flipx size/2 + ((-radius):+radius)) in Cairo.arc cx' cy' radius (3*π/2) (4*π/2)

	-- -- Second line segment
	-- vectorise Cairo.moveTo (centre - flipx size/2  + imag radius)
	-- vectorise Cairo.lineTo (centre + size/2        - imag radius)

	-- -- Curve
	-- let (cx':+cy') = (centre + size/2 - (radius:+radius)) in Cairo.arc cx' cy' radius 0 (π/2)

	-- -- Third line segment
	-- vectorise Cairo.moveTo (centre + size/2       - real radius)
	-- vectorise Cairo.lineTo (centre - flipy size/2 + real radius)

	-- -- Curve
	-- let (cx':+cy') = (centre - flipy size/2 + flipy (radius:+radius)) in Cairo.arc cx' cy' radius (π/2) π

	-- -- Fourth line segment
	-- vectorise Cairo.moveTo (centre - flipy size/2 - imag radius)
	-- vectorise Cairo.lineTo (centre - size/2       + imag radius)

	-- -- Curve
	-- let (cx':+cy') = (centre - size/2 + (radius:+radius)) in Cairo.arc cx' cy' radius π (3*π/2)

-- Images ----------------------------------------------------------------------------------------------------------------------------------

-- |
-- TODO: Factor out clip logic, document properly (cf. other anchored functions in this module)
-- TODO: Wrapper for images (?)
-- TODO: Factor out clip area (?)
anchoredImageWithClip :: (Complex Double -> Complex Double -> Cairo.Render ()) -> Complex Double -> Complex Double -> Cairo.Surface -> Cairo.Render ()
anchoredImageWithClip clip p anchor im = do
  size   <- imageSurfaceSize im
  centre <- return $ p - dotwise (*) anchor size
  clip centre size
  Cairo.clip
  vectorise Cairo.translate $ (centre - size*0.5)
  Cairo.setSourceSurface im 0 0
  -- vectorise Cairo.moveTo (centre - size*0.5)
  Cairo.paint
  vectorise Cairo.translate $ negate (centre - size*0.5)
  Cairo.resetClip

-- |
anchoredImage :: Complex Double -> Complex Double -> Cairo.Surface -> Cairo.Render ()
anchoredImage p anchor im = anchoredImageWithClip rectangle p anchor im

-- |
imageWithClip :: (Complex Double -> Complex Double -> Cairo.Render ()) -> Complex Double -> Cairo.Surface -> Cairo.Render ()
imageWithClip clip centre im = anchoredImageWithClip clip centre (0.0:+0.0) im

-- |
image :: Complex Double -> Cairo.Surface -> Cairo.Render ()
image p im = anchoredImageWithClip rectangle p (0.0:+0.0) im

-- Typography ------------------------------------------------------------------------------------------------------------------------------

-- | Renders the given string (with an arbitrary function 'draw'). The position of the centre of the bounding box is
--   given by the pin point 'p' and an 'anchor' (whose coordinates are normalised with respect to the text bounds).
-- TODO: Extract anchoring logic
-- TODO: Define anchor constants (eg. left:+top, centre:+bottom) (?)
anchoredText :: Complex Double -> Complex Double -> (String -> Cairo.Render a) -> String -> Cairo.Render a
anchoredText p anchor draw text = do
  extents <- textsize text
  vectorise Cairo.moveTo $ p - dotwise (*) (anchor + (0.5:+0.5)) extents
  draw text

-- |
centredText :: Complex Double -> (String -> Cairo.Render a) -> String -> Cairo.Render a
centredText centre draw text = anchoredText centre (0.0:+0.0) draw text


-- (C) Copyright Chris Banks 2011-2012

-- This file is part of The Continuous Pi-calculus Workbench (CPiWB).

--     CPiWB is free software: you can redistribute it and/or modify
--     it under the terms of the GNU General Public License as published by
--     the Free Software Foundation, either version 3 of the License, or
--     (at your option) any later version.

--     CPiWB is distributed in the hope that it will be useful,
--     but WITHOUT ANY WARRANTY; without even the implied warranty of
--     MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
--     GNU General Public License for more details.

--     You should have received a copy of the GNU General Public License
--     along with CPiWB.  If not, see <http://www.gnu.org/licenses/>.

{-# OPTIONS_GHC -w #-}

module CPi.ODE
    (-- * Important functions:
     dPdt,
     solveODE,
     evolveProcess,
     derivs,
     -- * Types:
     Solver,
     P',
     Expr(..),
     -- * Functions:
     timeSeries,
     timePoints,
     prettyODE,
     initials,
     speciesIn,
     speciesVars,
     diff,
     simp,
     simpP',
     (><)
    )where

--import qualified Data.List as L
import qualified Control.Exception as X
import qualified Data.Map as Map
import Data.Map (Map)
--import Data.Maybe

import qualified Numeric.GSL as GSL
import qualified Numeric.LinearAlgebra as LA

import BioCalcLib.Lib
import CPi.Lib
import CPi.Semantics

--------------------------------------------
-- Output the ODEs describing a CPi system
--------------------------------------------

-- Get the ODEs in a functional form.
xdot :: Env
     -> P' -- ^
     -> (Double -> [Double] -> [Double]) -- ^ ODEs in functional form
xdot env p = xdot'
    where
      -- the xdot function we need to return
      xdot' t xs = toODE p' vmap xs
      -- simplify the P'
      p' = Map.toList (simpP' env p)
      -- map species to indexed vars:
      defs2vars :: [(Species,Expr)] -> [Int] -> [(Species,Int)]
      defs2vars ((k,v):ks) (x:xs) = (k,x) : (defs2vars ks xs)
      defs2vars [] _ = []
      defs2vars _ [] = X.throw $ CpiException "CPi.ODE.xdot: run out of vars!"
      xvars :: [Int]
      xvars = [0..]
      vmap :: [(Species,Int)]
      vmap = defs2vars p' xvars
      -- get ODE expressions from P' expressions:
      toODE :: [(Species,Expr)] -> [(Species,Int)] -> [Double] -> [Double]
      toODE ps vmap xs = map (convert . snd) ps
          where
            convert (Var s') = xs!!(maybe (err s') id (lookup s' vmap))
            convert (Plus e e') = (convert e) + (convert e')
            convert (Times e e') = (convert e) * (convert e')
            convert (Num d) = d
            err s' = X.throw $ CpiException
                     ("Bug: bad lookup ("++(pretty s')++") in CPi.ODE.xdot.toODE")

{- FIXME: xdot and jac share a lot of inlined code. Refactor!
          vmap etc. can be replaced with 'speciesVars' (see below).
-}

-- Map species to labels:
speciesVars :: Env -> [a] -> P' -> Map Species a
speciesVars env vs p' = Map.fromList $ zip (Map.keys p') vs

-- Get the Jacobian in hmatrix form
jac :: Env -> P' -> (Double -> LA.Vector Double -> LA.Matrix Double)
jac env p = jac'
    where
      jac' t xs = toJac p' vmap (LA.toList xs)
      -- simplify the P'
      p' = Map.toList (simpP' env p)
      -- map species to indexed vars:
      defs2vars :: [(Species,Expr)] -> [Int] -> [(Species,Int)]
      defs2vars ((k,v):ks) (x:xs) = (k,x) : (defs2vars ks xs)
      defs2vars [] _ = []
      defs2vars _ [] = X.throw $ CpiException "CPi.ODE.xdot: run out of vars!"
      xvars :: [Int]
      xvars = [0..]
      vmap :: [(Species,Int)]
      vmap = defs2vars p' xvars
      -- get the jacobian matrix:
      toJac p vmap xs = (len><len) (map convert (map (simp env) (map diff' (cp (unzip p')))))
          where
            len = length p'
            cp (x,y) = [(a,b)|b<-y,a<-x]
            diff' (x,e) = diff x e
            convert (Var s') = xs!!(maybe (err s') id (lookup s' vmap))
            convert (Plus e e') = (convert e) + (convert e')
            convert (Times e e') = (convert e) * (convert e')
            convert (Num d) = d
            err s' = X.throw $ CpiException
                     ("Bug: bad lookup ("++(pretty s')++") in CPi.ODE.xdot.toODE")

-- Bringing in some infix funs from Numeric.LinearAlgebra:
-- Matrix constructor:
x><y = x LA.>< y
-- Vector accessor:
x@>y = x LA.! y

-- get the initial concentrations of the primes in process space:
initials :: Env -> Process -> P' -> [Double]
initials env proc p' = initials' (Map.toList (simpP' env p'))
    where
      initials' [] = []
      initials' ((s,_):ss) = (initconc proc s):(initials' ss)

-- Get the time points
timePoints :: (Int,(Double,Double)) -> LA.Vector Double
timePoints (r,(o,l)) = LA.linspace r (o,l)

-- | An ODE solver.
type Solver = Env -- ^ the environment
            -> Process -- ^ the syntactic process
            -> MTS -- ^ the MTS (from 'processMTS')
            -> P' -- ^ the symbolic process space (from 'dPdt')
            -> (Int,(Double,Double)) -- ^ time points to solve for (points(start,end))
            -> LA.Matrix Double -- ^ Matrix of numerical solution

-- FIXME: All solvers should be of this type ^
--        Need to refactor the following to conform.
--        (So we can pass different solvers around).

-- | Solve the ODEs using the internal solver.
--   Uses RKf45 explicit solver (general purpose, non-stiff).
solveODE :: Solver
solveODE env p mts dpdt (ts,(t0,tn))
    = GSL.odeSolve (xdot env dpdt) (initials env (wholeProc env p mts) dpdt) (timePoints (ts,(t0,tn)))

-- solve ODEs using Jacobian with hmatrix
solveODE' :: (Double -> [Double] -> [Double])
          -> (Double -> LA.Vector Double -> LA.Matrix Double)
          -> [Double]
          -> LA.Vector Double
          -> LA.Matrix Double
solveODE' odes jacob inits ts
   = GSL.odeSolveV (GSL.MSBDF jacob) hi epsAbs epsRel (l2v odes) (LA.fromList inits) ts
      where
        hi = (ts@>1 - ts@>0)/100
        epsAbs = 1e-06
        epsRel = 1e-06
        l2v f = \t -> LA.fromList . f t . LA.toList
solveODE'' odes inits ts
    = GSL.odeSolveV GSL.MSAdams hi epsAbs epsRel (l2v odes) (LA.fromList inits) ts
      where
        hi = (ts@>1 - ts@>0)/100
        epsAbs = 1e-06
        epsRel = 1e-06
        l2v f = \t -> LA.fromList . f t . LA.toList

-- plot the time series with GNUPlot
-- NOTE: deprecated by the CPi.Plot module.
--plotODE :: LA.Matrix Double -> LA.Vector Double -> IO ()
--plotODE x ts = Plot.mplot (ts : LA.toColumns x)

-- pretty print a Map of ODEs:
prettyODE env map = pp (Map.toList (simpP' env map))
    where
      pp [] = []
      pp ((x,y):z) = "d[" ++ pretty(nice x) ++ "]/dt ===> " ++ pretty y ++ "\n" ++ pp z
      nice x = maybe x id (revLookupDef env x)

---------------------------------------
-- Time series output of ODE solutions:
---------------------------------------

-- e.g. for sending to the model checker

-- A time series here is a list of time points with corresponding
-- map from Species to concentration.
type TimeSeries = [(Double, Map Species Double, Map Species Double)]

-- List of species in a process space (reduced to Defs if possible):
speciesIn :: Env -> P' -> [Species]
speciesIn env p = map
                  (\s->maybe s id (revLookupDef env s))
                  (map fst (Map.toList (simpP' env p)))

-- Get the raw time series:
timeSeries :: LA.Vector Double -- times
           -> LA.Matrix Double -- concs
           -> LA.Matrix Double -- derivs
           -> [Species]
           -> TimeSeries
timeSeries v cm dm ss
    = timeSeries' (LA.toList v)
      (map LA.toList (LA.toColumns (LA.tr cm)))
      (map LA.toList (LA.toColumns (LA.tr dm)))
      ss
          where
            timeSeries' [] _ _ _ = []
            timeSeries' _ [] _ _ = []
            timeSeries' _ _ [] _ = []
            timeSeries' _ _ _ [] = []
            timeSeries' (t:ts) (cs:css) (ds:dss) ss
                = (t, Map.fromList (zip ss cs), Map.fromList (zip ss ds))
                  : timeSeries' ts css dss ss

-----------------------------------
-- Process evolution
-----------------------------------

-- | Gives the syntactic process with initial conditions corresponding to the
-- | end conditions of the given process.

evolveProcess :: Env -- ^ the environment
              -> Process -- ^ the syntactic process
              -> MTS -- ^ the MTS (from 'processMTS')
              -> P' -- ^ the symbolic process space (from 'dPdt')
              -> (Int,(Double,Double)) -- ^ time points to solve for (points(start,end))
              -> Solver -- ^ the ODE solving function to use
              -> Process
evolveProcess env p mts p' ts solver
    = let matrix = solver env p mts p' ts
          vals = LA.toList $ last $ LA.toRows $ matrix
          ss = speciesIn env p'
          (Process _ net) = p
      in Process (zip ss vals) net

-----------------------------
-- Concentration derivatives
-----------------------------

derivs :: Env -- ^ the Environment
       -> P' -- ^ the symbolic process space (from 'dPdt')
       -> (Int,(Double,Double)) -- ^ time points to solve for (points(start,end))
       -> LA.Matrix Double -- ^ Matrix of numerical solution
       -> LA.Matrix Double -- ^ Matrix of concentration derivatives
derivs env p' ts solns
    = LA.fromLists $
      map
      (\(t,sol)->(xdot env p') t sol)
      (zip (LA.toList(timePoints ts)) (LA.toLists solns))

--------------------------------------------
-- Symbolic semantics
--------------------------------------------

-- Here we build the symbolic counterpart to the CPi immediate behaviour
-- substituting a variable for the real values.

-- We map to an expression, rather than a real value:
data Expr = Var Species
          | Plus Expr Expr
          | Times Expr Expr
          | Num Double
            deriving (Show,Eq,Ord)

instance Pretty Expr where
    pretty (Var x) = "[" ++ (pretty x) ++ "]"
    pretty (Plus x y) = pretty x ++ " + " ++ pretty y
    pretty (Times (Plus x y) (Plus x' y'))
        = "(" ++ pretty (Plus x y) ++ ") * (" ++ pretty (Plus x' y') ++ ")"
    pretty (Times (Plus x y) z)
        = "(" ++ pretty (Plus x y) ++ ") * " ++ pretty z
    pretty (Times x (Plus y z))
        = pretty x ++ " * (" ++ pretty (Plus y z) ++ ")"
    pretty (Times x y) = pretty x ++ " * " ++ pretty y
    pretty (Num k) = show k

-- | Symbolic representation of the process space.
type P' = Map Species Expr
type D' = Map (Species,Name,Concretion) Expr

-- pretty print our symbolic immediate behaviour:
prettyP' x = concat $ map (\(k,v)->((pretty k)++" |-> "++(pretty v)++"\n")) (Map.toList x)

-- Zero vectors:
p0' :: P'
p0' = Map.empty
d0' :: D'
d0' = Map.empty

-- vector constructors:
pVec' :: Env -> Species -> Expr -> P'
pVec' env s x
    = foldr (\s' -> \m -> Map.insertWith Plus s' x m) p0' (primes env s)
dVec' :: (Species,Name,Concretion) -> Expr -> D'
dVec' t x = Map.singleton t x

-- Vector addition in P' and D':
pplus' :: P' -> P' -> P'
pplus' x y = Map.unionWith (Plus) x y
dplus' :: D' -> D' -> D'
dplus' x y = Map.unionWith (Plus) x y

-- Scalar multiplication in P and D:
ptimes' :: P' -> Double -> P'
ptimes' p v = Map.map (Times (Num v)) p
dtimes' :: D' -> Double -> D'
dtimes' d v = Map.map (Times (Num v)) d

-- Vector subtraction in P and D:
pminus' :: P' -> P' -> P'
pminus' x y = x `pplus'` (y `ptimes'` (-1))
dminus' :: D' -> D' -> D'
dminus' x y = x `dplus'` (y `dtimes'` (-1))

-- The interaction potential
partial' :: Env -> MTS -> Process -> D'
partial' _ _ (Process [] _) = d0'
partial' env mts proc@(Process ps _) = foldr dplus' d0' (map partial'' ps)
    where
      partial'' (s,_) = foldr dplus' d0'
                        (map (\tr->dVec' (triple tr) (Var s)) (pots s))
      pots x = potentials (MTS (lookupTrans mts x))
      triple (TransSC s n c) = (s,n,c)
      triple _ = X.throw $ CpiException ("Bug: CPi.ODE.partial'.triple passed something other than a TransSC")

-- The interaction tensor
tensor' :: Env -> AffNet -> D' -> D' -> P'
tensor' env net ds1 ds2 = foldr pplus' p0' (map expr ds)
    where
      expr (x,y,p) = pVec' env p (expr' x y) `pminus'`
                     pVec' env (tri1(fst(x))) (expr' x y) `pminus'`
                     pVec' env (tri1(fst(y))) (expr' x y)
      ds = [(x,y,p)
            | x<-Map.toList ds1,
              y<-Map.toList ds2,
              p<-[maybe Nil id (pseudoapp (tri3(fst(x))) (tri3(fst(y))))],
              p/=Nil,
              aff net (tri2(fst(x)),tri2(fst(y))) /= Nothing]
      -- above: x,y are (Spec,Name,Conc),Concentration);
      -- a is Rate; p is Species result of pseudoapplication
      expr' ((s,n,c),e) ((s',n',c'),e')
          = maybe (Num 0.0)
            (\x-> (Times (Num x) (Times e e')))
            (aff net (n,n'))

-- | The symbolic immediate behaviour of a process (equivalent to the set of ODEs).
dPdt :: Env -- ^ the environment
     -> MTS -- ^ the MTS (generated by 'processMTS')
     -> Process -- ^ the syntactic process
     -> P' -- ^ the symbolic process space (=ODEs)
dPdt env mts proc = dPdt' env mts (wholeProc env proc mts)
    where
      dPdt' _ _ (Process [] _) = p0'
      dPdt' env mts p@(Process [(s,c)] net)
          = foldr pplus' p0' (map f taus)
            `pplus'`
            Map.map (\x->(Times (Num 0.5) x))
                   (tensor' env net (partial' env mts p) (partial' env mts p))
          where
            taus = [x|x<-lookupTrans mts s, tau x]
            tau (TransT _ _ _) = True
            tau _ = False
            f (TransT src r dst)
                = pVec' env dst (Times (Num k) (Var s)) `pplus'`
                  pVec' env s (Times (Num(-1*k)) (Var s))
                      where k = rate r
                            rate (TTau r) = r
            f _ = X.throw $ CpiException ("Bug: CPi.ODE.dPdt'.f passed something other than a TransT")
      dPdt' env mts (Process (p:ps) net)
          = (tensor' env net partH partT)
            `pplus'` (dPdt' env mts procH)
            `pplus'` (dPdt' env mts procT)
                where
                  partH = partial' env mts procH
                  partT = partial' env mts procT
                  procH = Process [p] net
                  procT = Process ps net

-- Symbolic differentiation of an Expr
diff :: Species -> Expr -> Expr
diff _ (Num _) = Num 0.0
diff x (Var s)
    | s==x
        = Num 1.0
    | otherwise
        = Num 0.0
diff x (Plus a b)
    = Plus (diff x a) (diff x b)
diff x (Times a b)
    = Plus (Times (diff x a) b) (Times a (diff x b))

-- Simple simplification of an expression
simp :: Env -> Expr -> Expr
simp env x
    | x==x' = x'
    | otherwise = simp env x'
    where
      x' = simp' x
      simp' (Num n) = Num n
      simp' (Var s) = Var (simpS env s)
      simp' (Plus (Num 0.0) b) = simp env b
      simp' (Plus a (Num 0.0)) = simp env a
      simp' (Plus (Num a) (Num b)) = Num (a+b)
      simp' (Plus a b) = Plus (simp env a) (simp env b)
      simp' (Times (Num 0.0) b) = Num 0.0
      simp' (Times a (Num 0.0)) = Num 0.0
      simp' (Times (Num 1.0) b) = simp env b
      simp' (Times a (Num 1.0)) = simp env a
      simp' (Times (Num (-1.0)) (Num b)) = Num (-b)
      simp' (Times (Num a) (Num (-1.0))) = Num (-a)
      simp' (Times (Num a) (Num b)) = Num (a*b)
      simp' (Times a b) = Times (simp env a) (simp env b)

-- Simplify dP/dt
simpP' :: Env -> P' -> P'
simpP' env p = Map.mapKeys (simpS env) (Map.map (simp env) p)

-- Simplify species
simpS :: Env -> Species -> Species
simpS env s = maybe s id (revLookupDef env s)


{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE AllowAmbiguousTypes #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE FunctionalDependencies #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE InstanceSigs #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE StandaloneDeriving #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE TemplateHaskell #-}
{-# LANGUAGE DeriveGeneric #-}
{-# LANGUAGE UndecidableInstances #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE PolyKinds #-}
{-# LANGUAGE PartialTypeSignatures #-}
{-# OPTIONS_GHC -Wno-partial-type-signatures #-}
{-# LANGUAGE StandaloneKindSignatures #-}
module Inference.Conjugate where

import           Control.Monad.Primitive        ( PrimMonad
                                                , PrimState
                                                )
import           Control.Monad.Reader           ( ReaderT
                                                , runReaderT
                                                )
import           Control.Monad.Reader.Class     ( MonadReader(ask) )
import           Control.Monad.State            ( State
                                                , StateT
                                                , evalState
                                                , evalStateT
                                                , execStateT
                                                , runStateT
                                                )
import           Control.Monad.State.Class      ( get
                                                , modify
                                                , put
                                                )
import           Control.Monad.Trans.Maybe      ( MaybeT(MaybeT)
                                                , runMaybeT
                                                )
import           Control.Monad.Writer
import           Data.Dynamic                   ( Dynamic
                                                , Typeable
                                                , fromDynamic
                                                , toDyn
                                                )
import           Data.Kind
import           Data.Maybe                     ( fromMaybe )
import           Data.MultiSet                 as MS
import qualified Data.Sequence                 as S
import           Data.Typeable                  ( Proxy(Proxy)
                                                , typeRep
                                                )
import qualified Data.Vector                   as V
import qualified Debug.Trace                   as DT
import           GHC.Float                      ( int2Double )
import           GHC.Generics
import           GHC.TypeNats
import           Lens.Micro
import           Lens.Micro.Extras
import           Lens.Micro.TH                  ( makeLenses )
import           Numeric.SpecFunctions          ( logBeta
                                                , logChoose
                                                , logFactorial
                                                , logGamma
                                                )
import           System.Random.MWC.Probability
                                         hiding ( Uniform )

----------------------------------------------------------------
-- type classes and families for distributions and conjugates --
----------------------------------------------------------------

-- | Describes a family of distributions with a fixed form.
-- For example, a 'Bernoulli' distribution is parameterized by a probability @p@
-- and produces binary samples
-- (@True@ with probability @p@, @False@ with probability @1-p@).
--
-- Its 'Distribution' instance is:
-- > instance Distribution Bernoulli where
-- >   type Params Bernoulli = Double
-- >   type Support Bernoulli = Bool
-- >   distSample _ = uncurry bernoulli
-- >   distLogP _ p True = log p
-- >   distLogP _ p False = log (1 - p)
class Distribution a where
  type Params a :: Type
  type Support a :: Type
  distSample :: (PrimMonad m) => a -> Params a -> Prob m (Support a)
  distLogP :: a -> Params a -> Support a -> Double

-- -- | Used as a type-lifted kind for conjugate distribution pairs.
-- data Conj a b = Conj a bb

-- | A type-level marker for treating a distribution as a prior.
newtype AsPrior p = AsPrior p

-- -- | A type-level marker for treating a distribution as a likelihood.
-- newtype AsLk l = AsLk l

-- | Marks two distributions as a conjugate pair of prior and likelihood.
-- The property of such a pair is that the posterior has the same form as the prior
-- (including the same 'Params' and 'Support'),
-- and that its parameters can be obtained analytically from the parameters of the prior
-- and a set of observations.
--
-- The class method 'updatePrior' returns the parameters of the posterior
-- given the prior parameters after a single observation.
class (Distribution p, Distribution l, Support p ~ Params l) => Conjugate p l where
  priorSingleton :: p -- ^ provides a singleton instance of the prior distribution
                      -- in order to make sampling from priors easier.
  updatePrior :: l -> Params p -> Support l -> Params p
  predLogP :: l -> Params p -> Support l -> Double

type family Hyper (a :: k) :: Type
type instance Hyper (AsPrior p) = Params p

type family Probs (a :: k) :: Type
type instance Probs (AsPrior p) = Support p

-- type family Value (a :: k) :: Type
-- type instance Value (Conj p l) = Support l

-- helper types for instantiating hyperparameters, parameters, and values
-- ----------------------------------------------------------------------

newtype HyperRep p = HyperRep { runHyper :: Hyper (AsPrior p) }
deriving instance Show (Hyper (AsPrior p)) => Show (HyperRep p)

type instance Hyper (a :: (Type -> Type) -> Type) = a HyperRep

newtype ProbsRep p = ProbsRep { runProbs :: Probs (AsPrior p)}
deriving instance Show (Probs (AsPrior p)) => Show (ProbsRep p)

type instance Probs (a :: (Type -> Type) -> Type) = a ProbsRep

-- newtype ValueRep p l = ValueRep { runValue :: Value (Conj p l) }
-- deriving instance Show (Value (Conj p l)) => Show (ValueRep p l)

-- type instance Value (a :: (Type -> Type -> Type) -> Type) = a ValueRep

-----------------------------------------------------
-- generic magic for conjugates and parameter sets --
-----------------------------------------------------

-- Jeffreys prior
-- ---------------

class Jeffreys a where
  jeffreysPrior :: Hyper a

class GJeffreys t where
  gjeffreysPrior :: forall p. t p

instance GJeffreys V1 where
  gjeffreysPrior = undefined -- ok

instance GJeffreys U1 where
  gjeffreysPrior = U1

-- base case: k is a conjugate distribution
instance (Jeffreys (AsPrior p)) => GJeffreys (K1 i (HyperRep p)) where
  gjeffreysPrior = K1 $ HyperRep $ jeffreysPrior @(AsPrior p)

-- recursive case: k is another record
instance (Jeffreys k, k HyperRep ~ Hyper k) => GJeffreys (K1 i (k HyperRep)) where
  gjeffreysPrior = K1 $ jeffreysPrior @k

instance (GJeffreys t) => GJeffreys (M1 i c (t :: Type -> Type)) where
  gjeffreysPrior = M1 gjeffreysPrior

instance (GJeffreys ta, GJeffreys tb) => GJeffreys (ta :*: tb) where
  gjeffreysPrior = gjeffreysPrior @ta :*: gjeffreysPrior @tb

instance (Generic (t HyperRep), GJeffreys (Rep (t HyperRep))) => Jeffreys (t :: (Type -> Type) -> Type) where
  jeffreysPrior = GHC.Generics.to (gjeffreysPrior @(Rep (t HyperRep)))

-- uniform prior
-- -------------

class Uniform a where
  uniformPrior :: Hyper a

class GUniform t where
  guniformPrior :: forall p. t p

instance GUniform V1 where
  guniformPrior = undefined -- ok

instance GUniform U1 where
  guniformPrior = U1

-- base case: k is a conjugate distribution
instance (Uniform (AsPrior p)) => GUniform (K1 i (HyperRep p)) where
  guniformPrior = K1 $ HyperRep $ uniformPrior @(AsPrior p)

-- recursive case: k is another record
instance (Uniform k, k HyperRep ~ Hyper k) => GUniform (K1 i (k HyperRep)) where
  guniformPrior = K1 $ uniformPrior @k

instance (GUniform t) => GUniform (M1 i c (t :: Type -> Type)) where
  guniformPrior = M1 guniformPrior

instance (GUniform ta, GUniform tb) => GUniform (ta :*: tb) where
  guniformPrior = guniformPrior @ta :*: guniformPrior @tb

instance (Generic (t HyperRep), GUniform (Rep (t HyperRep))) => Uniform (t :: (Type -> Type) -> Type) where
  uniformPrior = GHC.Generics.to (guniformPrior @(Rep (t HyperRep)))

-- sampling from prior
-- ------------------

class Prior a where
  sampleProbs :: (PrimMonad m) => Hyper a -> Prob m (Probs a)

class GPrior i o where
  gsampleProbs :: forall m p. PrimMonad m => i p -> Prob m (o p)

instance GPrior V1 V1 where
  gsampleProbs = undefined -- ok

instance GPrior U1 U1 where
  gsampleProbs _ = pure U1

-- base case: k is a conjugate distribution
instance (Prior (AsPrior p)) => GPrior (K1 i (HyperRep p)) (K1 i (ProbsRep p)) where
  gsampleProbs (K1 (HyperRep hyper)) =
    K1 . ProbsRep <$> sampleProbs @(AsPrior p) hyper

-- recursive case: k is another record
instance (Prior k, k HyperRep ~ Hyper k, k ProbsRep ~ Probs k) =>
         GPrior (K1 i (k HyperRep)) (K1 i (k ProbsRep)) where
  gsampleProbs (K1 hyper) = K1 <$> sampleProbs @k hyper

instance (GPrior ti to) => GPrior (M1 i c ti) (M1 i' c' to) where
  gsampleProbs (M1 x) = M1 <$> gsampleProbs x

instance (GPrior ia oa, GPrior ib ob) => GPrior (ia :*: ib) (oa :*: ob) where
  gsampleProbs (a :*: b) = (:*:) <$> gsampleProbs a <*> gsampleProbs b

instance (GPrior ia oa, GPrior ib ob) => GPrior (ia :+: ib) (oa :+: ob) where
  gsampleProbs (L1 a) = L1 <$> gsampleProbs a
  gsampleProbs (R1 b) = R1 <$> gsampleProbs b

instance ( Generic (a HyperRep)
         , Generic (a ProbsRep)
         , GPrior (Rep (a HyperRep)) (Rep (a ProbsRep))
         ) => Prior (a :: (Type -> Type) -> Type) where
  sampleProbs hyper = GHC.Generics.to <$> gsampleProbs (from hyper)

-----------------------
-- likelihood monads --
-----------------------

type Accessor r p = forall f . Lens' (r f) (f p)

class Monad m => RandomInterpreter m r | m -> r where
  type SampleCtx m a :: Constraint
  sampleValue :: (Conjugate p l, SampleCtx m l) => String -> l -> Accessor r p -> m (Support l)
  sampleConst :: (Distribution d, SampleCtx m d) => String -> d -> Params d -> m (Support d)
  permutationPlate :: (Ord a) => Int -> m a -> m [a]

newtype Trace (r :: (Type -> Type) -> Type)  = Trace { runTrace :: S.Seq Dynamic }
  deriving (Show)

observeValue
  :: (Conjugate p l, Typeable (Support l), Monad m)
  => String
  -> l
  -> Accessor r p
  -> Support l
  -> StateT (Trace r) m ()
observeValue _ _ _ val = modify $ \(Trace st) -> Trace $ st S.|> toDyn val

observeConst
  :: (Distribution d, Typeable (Support d), Monad m)
  => String
  -> d
  -> Params d
  -> Support d
  -> StateT (Trace r) m ()
observeConst _ _ _ val = modify $ \(Trace st) -> Trace $ st S.|> toDyn val

takeTrace :: Typeable a => Trace r -> Maybe (a, Trace r)
takeTrace (Trace t) = do
  (valDyn, rest) <- case S.viewl t of
    S.EmptyL         -> Nothing
    valDyn S.:< rest -> Just (valDyn, rest)
  val <- fromDynamic valDyn
  pure (val, Trace rest)

-- just sample
-- -----------

newtype SampleI m r a = SampleI (ReaderT (r ProbsRep) (Prob m) a)
  deriving (Functor, Applicative, Monad)

instance (PrimMonad m) => RandomInterpreter (SampleI m r) r where
  type SampleCtx (SampleI m r) a = ()
  sampleValue
    :: forall p l
     . (Conjugate p l)
    => String
    -> l
    -> Accessor r p
    -> SampleI m r (Support l)
  sampleValue _ lk getProbs = SampleI $ do
    probs <- ask
    lift $ distSample lk $ runProbs $ view getProbs probs
  sampleConst
    :: forall d
     . (Distribution d)
    => String
    -> d
    -> Params d
    -> SampleI m r (Support d)
  sampleConst _ dist params = SampleI $ lift $ distSample dist params
  permutationPlate = replicateM

sampleResult :: p ProbsRep -> SampleI m p a -> Gen (PrimState m) -> m a
sampleResult probs (SampleI a) = sample (runReaderT a probs)

-- sample and trace the execution process
-- --------------------------------------

newtype TraceI m r a = TraceI (ReaderT (r ProbsRep) (StateT (Trace r) (Prob m)) a)
  deriving (Functor, Applicative, Monad)

instance (PrimMonad m) => RandomInterpreter (TraceI m r) r where
  type SampleCtx (TraceI m r) l = Typeable (Support l)
  sampleValue
    :: forall p l
     . (Conjugate p l, Typeable (Support l))
    => String
    -> l
    -> Accessor r p
    -> TraceI m r (Support l)
  sampleValue _ lk getProbs = TraceI $ do
    probs <- ask
    val   <- lift $ lift $ distSample lk $ runProbs $ view getProbs probs
    modify $ \(Trace obs) -> Trace $ obs S.|> toDyn val
    pure val
  sampleConst
    :: forall d
     . (Distribution d, Typeable (Support d))
    => String
    -> d
    -> Params d
    -> TraceI m r (Support d)
  sampleConst _ dist params = TraceI $ do
    val <- lift $ lift $ distSample dist params
    modify $ \(Trace obs) -> Trace $ obs S.|> toDyn val
    pure val
  permutationPlate = replicateM

sampleTrace
  :: r ProbsRep -> TraceI m r a -> Gen (PrimState m) -> m (a, Trace r)
sampleTrace probs (TraceI a) = do
  let st = runReaderT a probs
      pr = runStateT st (Trace mempty)
  sample pr

-- evaluate the probability of a trace
-- -----------------------------------

newtype EvalTraceI r a = EvalTraceI (ReaderT (r ProbsRep) (StateT (Trace r, Double) Maybe) a)
  deriving (Functor, Applicative, Monad)

instance RandomInterpreter (EvalTraceI r) r where
  type SampleCtx (EvalTraceI r) l = Typeable (Support l)
  sampleValue
    :: forall p l
     . (Conjugate p l, Typeable (Support l))
    => String
    -> l
    -> Accessor r p
    -> EvalTraceI r (Support l)
  sampleValue _ lk getProbs = EvalTraceI $ do
    probs              <- ask
    (trace, totalLogP) <- get
    (val  , trace'   ) <- lift $ lift $ takeTrace trace
    let logP = distLogP lk (runProbs $ view getProbs probs) val
    put (trace', totalLogP + logP)
    pure val
  sampleConst
    :: forall d
     . (Distribution d, Typeable (Support d))
    => String
    -> d
    -> Params d
    -> EvalTraceI r (Support d)
  sampleConst _ dist params = EvalTraceI $ do
    (trace, totalLogP) <- get
    (val  , trace'   ) <- lift $ lift $ takeTrace trace
    let logP = distLogP dist params val
    put (trace', totalLogP + logP)
    pure val
  permutationPlate n submodel = EvalTraceI $ do
    probs                     <- ask
    (trace  , totalLogP     ) <- get
    (results, (trace', logP)) <- lift $ lift $ runTraceLogP
      probs
      trace
      (replicateM n submodel)
    let unique = MS.fromList results
        permutations =
          logFactorial n - sum (logFactorial . snd <$> MS.toOccurList unique)
    put (trace', totalLogP + logP + permutations)
    pure results

runTraceLogP
  :: r ProbsRep -> Trace r -> EvalTraceI r a -> Maybe (a, (Trace r, Double))
runTraceLogP probs trace (EvalTraceI model) = do
  runStateT (runReaderT model probs) (trace, 0)

evalTraceLogP :: r ProbsRep -> Trace r -> EvalTraceI r a -> Maybe (a, Double)
evalTraceLogP probs trace model = do
  (val, (_trace, logp)) <- runTraceLogP probs trace model
  pure (val, logp)

-- evaluate the predictive probability of a trace
-- -----------------------------------

newtype EvalPredTraceI r a = EvalPredTraceI (ReaderT (r HyperRep) (StateT (Trace r, Double) Maybe) a)
  deriving (Functor, Applicative, Monad)

instance RandomInterpreter (EvalPredTraceI r) r where
  type SampleCtx (EvalPredTraceI r) l = Typeable (Support l)
  sampleValue
    :: forall p l
     . (Conjugate p l, Typeable (Support l))
    => String
    -> l
    -> Accessor r p
    -> EvalPredTraceI r (Support l)
  sampleValue _ lk getHyper = EvalPredTraceI $ do
    hyper              <- ask
    (trace, totalLogP) <- get
    (val  , trace'   ) <- lift $ lift $ takeTrace trace
    let logP = predLogP @p lk (runHyper $ view getHyper hyper) val
    put (trace', totalLogP + logP)
    pure val
  sampleConst
    :: forall d
     . (Distribution d, Typeable (Support d))
    => String
    -> d
    -> Params d
    -> EvalPredTraceI r (Support d)
  sampleConst _ dist params = EvalPredTraceI $ do
    (trace, totalLogP) <- get
    (val  , trace'   ) <- lift $ lift $ takeTrace trace
    let logP = distLogP dist params val
    put (trace', totalLogP + logP)
    pure val
  permutationPlate n submodel = EvalPredTraceI $ do
    probs                     <- ask
    (trace  , totalLogP     ) <- get
    (results, (trace', logP)) <- lift $ lift $ runTracePredLogP
      probs
      trace
      (replicateM n submodel)
    let unique = MS.fromList results
        permutations =
          logFactorial n - sum (logFactorial . snd <$> MS.toOccurList unique)
    put (trace', totalLogP + logP + permutations)
    pure results

runTracePredLogP
  :: r HyperRep -> Trace r -> EvalPredTraceI r a -> Maybe (a, (Trace r, Double))
runTracePredLogP hyper trace (EvalPredTraceI model) = do
  runStateT (runReaderT model hyper) (trace, 0)

evalTracePredLogP
  :: r HyperRep -> Trace r -> EvalPredTraceI r a -> Maybe (a, Double)
evalTracePredLogP hyper trace model = do
  (val, (_trace, logp)) <- runTracePredLogP hyper trace model
  pure (val, logp)

-- update priors
-- -------------

newtype UpdatePriorsI r a = UpdatePriorsI (StateT (Trace r, r HyperRep) Maybe a)
  deriving (Functor, Applicative, Monad)

instance RandomInterpreter (UpdatePriorsI r) r where
  type SampleCtx (UpdatePriorsI r) l = Typeable (Support l)
  sampleValue
    :: forall p l
     . (Conjugate p l, Typeable (Support l))
    => String
    -> l
    -> Accessor r p
    -> UpdatePriorsI r (Support l)
  sampleValue _ lk accessor = UpdatePriorsI $ do
    (trace, priors) <- get
    (val  , trace') <- lift $ takeTrace trace
    let priors' :: r HyperRep
        priors' = over
          accessor
          (\(HyperRep pr) -> HyperRep $ updatePrior @p @l lk pr val)
          priors
    put (trace', priors')
    pure val
  sampleConst _ _ _ = UpdatePriorsI $ do
    (trace, priors) <- get
    (val  , trace') <- lift $ takeTrace trace
    put (trace', priors)
    pure val
  permutationPlate = replicateM

getPosterior
  :: r HyperRep -> Trace r -> UpdatePriorsI r a -> Maybe (r HyperRep)
getPosterior priors trace (UpdatePriorsI model) = do
  (_trace, posteriors) <- execStateT model (trace, priors)
  pure posteriors

-- show how a trace is generated by a model
-- ----------------------------------------

newtype ShowTraceI r a = ShowTraceI (MaybeT (WriterT String (State (Trace r))) a)
  deriving (Functor, Applicative, Monad)

showTraceItem
  :: forall l r
   . (Show (Support l), Typeable l, Typeable (Support l))
  => String
  -> ShowTraceI r (Support l)
showTraceItem name = ShowTraceI $ do
  trace         <- get
  (val, trace') <- MaybeT $ pure $ takeTrace trace
  put trace'
  let distName = show (typeRep (Proxy :: Proxy l))
  tell
    $  "Sampled value "
    <> show val
    <> " from a "
    <> distName
    <> " at "
    <> name
    <> ".\n"
  pure val

instance RandomInterpreter (ShowTraceI r) r where
  type SampleCtx (ShowTraceI r) l
    = (Typeable (Support l), Typeable l, Show (Support l))
  sampleValue
    :: forall p l
     . (Conjugate p l, Typeable (Support l), Typeable l, Show (Support l))
    => String
    -> l
    -> Accessor r p
    -> ShowTraceI r (Support l)
  sampleValue name _ acc = showTraceItem @l name
  sampleConst
    :: forall d
     . (Distribution d, SampleCtx (ShowTraceI r) d)
    => String
    -> d
    -> Params d
    -> ShowTraceI r (Support d)
  sampleConst name _ _ = showTraceItem @d name
  permutationPlate = replicateM

showTrace :: Trace r -> ShowTraceI r a -> (Maybe a, String)
showTrace trace (ShowTraceI model) =
  evalState (runWriterT (runMaybeT model)) trace

printTrace :: Trace r -> ShowTraceI r a -> IO ()
printTrace trace model = do
  let (res, txt) = showTrace trace model
  putStrLn txt
  case res of
    Nothing -> do
      putStrLn "Trace does not match the model (stops here)"
    Just _ -> pure ()

-- tracing interpreter (for debugging)
-- -----------------------------------

newtype TraceTraceI r a = TraceTraceI (State (Trace r) a)
  deriving (Functor, Applicative, Monad)

traceTraceItem
  :: forall l r
   . (Show (Support l), Typeable l, Typeable (Support l))
  => String
  -> TraceTraceI r (Support l)
traceTraceItem name = TraceTraceI $ do
  trace <- get
  let loc = show (typeRep (Proxy :: Proxy l)) <> " at " <> name
  case takeTrace trace of
    Just (val, trace') -> do
      put trace'
      DT.traceM $ "Sampled value " <> show val <> " from a " <> loc <> "."
      pure val
    Nothing -> error $ "Incompatible trace at " <> loc <> "."

instance RandomInterpreter (TraceTraceI r) r where
  type SampleCtx (TraceTraceI r) l
    = (Typeable (Support l), Typeable l, Show (Support l))
  sampleValue
    :: forall p l
     . (Conjugate p l, Typeable (Support l), Typeable l, Show (Support l))
    => String
    -> l
    -> Accessor r p
    -> TraceTraceI r (Support l)
  sampleValue name _ acc = traceTraceItem @l name
  sampleConst
    :: forall d
     . (Distribution d, SampleCtx (TraceTraceI r) d)
    => String
    -> d
    -> Params d
    -> TraceTraceI r (Support d)
  sampleConst name _ _ = traceTraceItem @d name
  permutationPlate = replicateM

traceTrace :: Trace r -> TraceTraceI r a -> a
traceTrace trace (TraceTraceI model) = evalState model trace

-------------------
-- distributions --
-------------------

-- Beta
-- ----

data Beta = Beta
  deriving (Eq, Ord, Show, Generic)

instance Distribution Beta where
  type Params Beta = (Double, Double)
  type Support Beta = Double
  distSample _ = uncurry beta
  distLogP _ (a, b) p =
    log (p ** (a - 1)) + log ((1 - p) ** (b - 1)) - logBeta a b

instance Jeffreys (AsPrior Beta) where
  jeffreysPrior = (0.5, 0.5)

instance Uniform (AsPrior Beta) where
  uniformPrior = (1, 1)

instance Prior (AsPrior Beta) where
  sampleProbs = distSample Beta

-- Bernoulli
-- ---------

data Bernoulli = Bernoulli
  deriving (Eq, Ord, Show, Generic)

instance Distribution Bernoulli where
  type Params Bernoulli = Double
  type Support Bernoulli = Bool
  distSample _ = bernoulli
  distLogP _ p True  = log p
  distLogP _ p False = log (1 - p)

-- Binomial
-- --------

newtype Binomial = Binomial Int
  deriving (Eq, Ord, Show, Generic)

instance Distribution Binomial where
  type Params Binomial = Double
  type Support Binomial = Int
  distSample (Binomial n) = binomial n
  distLogP (Binomial n) p k =
    logChoose n k + (log p * k') + (log (1 - p) * (n' - k'))
   where
    k' = int2Double k
    n' = int2Double n

-- Categorical
-- -----------

type Categorical :: Nat -> Type
data Categorical n = Categorical
  deriving (Eq, Ord, Show, Generic)

instance Distribution (Categorical n) where
  type Params (Categorical n) = V.Vector Double
  type Support (Categorical n) = Int
  distSample _ = categorical
  distLogP _ ps cat = log $ fromMaybe 0 $ ps V.!? cat

-- Dirichlet
-- ---------

type Dirichlet :: Nat -> Type
data Dirichlet n = Dirichlet
  deriving (Eq, Ord, Show, Generic)

instance Distribution (Dirichlet n) where
  type Params (Dirichlet n) = V.Vector Double
  type Support (Dirichlet n) = V.Vector Double
  distSample _ = dirichlet
  distLogP _ counts probs = logp + logz
   where
    logp = sum (V.zipWith (\a x -> log x * (a - 1)) counts probs)
    logz = logGamma (sum counts) - sum (logGamma <$> counts)

instance KnownNat n => Jeffreys (AsPrior (Dirichlet n)) where
  jeffreysPrior = V.replicate (fromIntegral $ natVal (Proxy :: Proxy n)) 0.5

instance KnownNat n => Uniform (AsPrior (Dirichlet n)) where
  uniformPrior = V.replicate (fromIntegral $ natVal (Proxy :: Proxy n)) 1

instance Prior (AsPrior (Dirichlet n)) where
  sampleProbs = distSample Dirichlet

-- Geometric (from 0)
-- ------------------

data Geometric0 = Geometric0
  deriving (Eq, Ord, Show, Generic)

instance Distribution Geometric0 where
  type Params Geometric0 = Double
  type Support Geometric0 = Int
  distSample _ = geometric0
   where
    geometric0 p = do
      coin <- bernoulli p
      if coin then pure 0 else (1 +) <$> geometric0 p
  distLogP _ p val | val >= 0  = (log (1 - p) * int2Double val) + log p
                   | otherwise = log 0

-- Geometric (from 1)
-- ------------------

data Geometric1 = Geometric1
  deriving (Eq, Ord, Show, Generic)

instance Distribution Geometric1 where
  type Params Geometric1 = Double
  type Support Geometric1 = Int
  distSample _ = geometric1
   where
    geometric1 p = do
      coin <- bernoulli p
      if coin then pure 1 else (1 +) <$> geometric1 p
  distLogP _ p val | val >= 1  = (log (1 - p) * int2Double (val - 1)) + log p
                   | otherwise = log 0

-----------------------------
-- conjugate distributions --
-----------------------------

-- beta bernoulli
-- --------------

instance Conjugate Beta Bernoulli where
  priorSingleton = Beta
  updatePrior _ (a, b) False = (a, b + 1)
  updatePrior _ (a, b) True  = (a + 1, b)
  predLogP _ (a, b) False = log a - log (a + b)
  predLogP _ (a, b) True  = log b - log (a + b)

-- B(1+a, b)   G(1+a) G(b) G(a+b)   G(1+a) G(a+b)   a G(a) G(a+b)        a
-- --------- = ------------------ = ------------- = ----------------- = ---
-- B(a, b)     G(a) G(b) G(a+b+1)   G(a) G(a+b+1)   G(a) (a+b) G(a+b)   a+b

-- beta binomial
-- -------------

instance Conjugate Beta Binomial where
  priorSingleton = Beta
  updatePrior (Binomial n) (a, b) x = (a + x', b + (n' - x'))
   where
    x' = int2Double x
    n' = int2Double n
  predLogP (Binomial n) (a, b) k =
    logChoose n k + logBeta (k' + a) (n' - k' + a) - logBeta a b
   where
    n' = int2Double n
    k' = int2Double k

-- beta geometric0
-- ---------------

instance Conjugate Beta Geometric0 where
  priorSingleton = Beta
  updatePrior _ (a, b) k = (a + 1, b + int2Double k)
  predLogP _ (a, b) k =
    (log a + logGamma (a + b) + logGamma (k' + b))
      - (logGamma b + logGamma (a + b + k' + 1))
    where k' = int2Double k

-- beta geometric1
-- ---------------

instance Conjugate Beta Geometric1 where
  priorSingleton = Beta
  updatePrior _ (a, b) k = (a + 1, b + int2Double (k - 1))
  predLogP _ (a, b) k =
    (log a + logGamma (a + b) + logGamma (k' + b))
      - (logGamma b + logGamma (a + b + k' + 1))
    where k' = int2Double k - 1

-- dirichlet categorical
-- ---------------------

instance Conjugate (Dirichlet n) (Categorical n) where
  priorSingleton = Dirichlet
  updatePrior _ counts obs
    | obs >= 0, obs < V.length counts
    = counts V.// [(obs, (counts V.! obs) + 1)]
    | otherwise
    = counts
  predLogP _ counts obs =
    log (fromMaybe 0 $ counts V.!? obs) - log (V.sum counts)


{-# LANGUAGE CPP #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE DeriveTraversable #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE UndecidableInstances #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE RankNTypes #-}
#if defined(__GLASGOW_HASKELL__) && __GLASGOW_HASKELL__ >= 702
{-# LANGUAGE Trustworthy #-}
{-# LANGUAGE DeriveGeneric #-}
#endif
{-# LANGUAGE DeriveDataTypeable #-}
#if defined(__GLASGOW_HASKELL__) && __GLASGOW_HASKELL__ < 708
{-# LANGUAGE StandaloneDeriving #-}
{-# LANGUAGE ScopedTypeVariables #-}
#endif
-----------------------------------------------------------------------------
-- |
-- License     :  BSD-style (see the file LICENSE)
-- Maintainer  :  Edward Kmett <ekmett@gmail.com>
-- Stability   :  provisional
-- Portability :  portable
--
-- Operations on affine spaces.
-----------------------------------------------------------------------------
module Linear.Affine where

import Control.Applicative
import Control.DeepSeq
import Control.Monad (liftM)
import Control.Lens
import Data.Binary as Binary
import Data.Bytes.Serial
import Data.Complex (Complex)
import Data.Data
import Data.Distributive
import Data.Foldable as Foldable
import Data.Functor.Bind
import Data.Functor.Classes
import Data.Functor.Rep as Rep
import Data.HashMap.Lazy (HashMap)
import Data.Hashable
import Data.IntMap (IntMap)
import Data.Ix
import Data.Map (Map)
import Data.Serialize as Cereal
import Data.Vector (Vector)
import Foreign.Storable
#if __GLASGOW_HASKELL__ >= 702
import GHC.Generics (Generic)
#endif
#if __GLASGOW_HASKELL__ >= 706
import GHC.Generics (Generic1)
#endif
import qualified Data.Vector.Generic.Mutable as M
import qualified Data.Vector.Generic as G
import qualified Data.Vector.Unboxed.Base as U
import Linear.Epsilon
import Linear.Metric
import Linear.Plucker
import Linear.Quaternion
import Linear.V
import Linear.V0
import Linear.V1
import Linear.V2
import Linear.V3
import Linear.V4
import Linear.Vector

#ifdef HLINT
{-# ANN module "HLint: ignore Unused LANGUAGE pragma" #-}
#endif

-- | An affine space is roughly a vector space in which we have
-- forgotten or at least pretend to have forgotten the origin.
--
-- > a .+^ (b .-. a)  =  b@
-- > (a .+^ u) .+^ v  =  a .+^ (u ^+^ v)@
-- > (a .-. b) ^+^ v  =  (a .+^ v) .-. q@
class Additive (Diff p) => Affine p where
  type Diff p :: * -> *

  infixl 6 .-.
  -- | Get the difference between two points as a vector offset.
  (.-.) :: Num a => p a -> p a -> Diff p a

  infixl 6 .+^
  -- | Add a vector offset to a point.
  (.+^) :: Num a => p a -> Diff p a -> p a

  infixl 6 .-^
  -- | Subtract a vector offset from a point.
  (.-^) :: Num a => p a -> Diff p a -> p a
  p .-^ v = p .+^ negated v
  {-# INLINE (.-^) #-}

-- | Compute the quadrance of the difference (the square of the distance)
qdA :: (Affine p, Foldable (Diff p), Num a) => p a -> p a -> a
qdA a b = Foldable.sum (fmap (join (*)) (a .-. b))
{-# INLINE qdA #-}

-- | Distance between two points in an affine space
distanceA :: (Floating a, Foldable (Diff p), Affine p) => p a -> p a -> a
distanceA a b = sqrt (qdA a b)
{-# INLINE distanceA #-}

#define ADDITIVEC(CTX,T) instance CTX => Affine T where type Diff T = T ; \
  (.-.) = (^-^) ; {-# INLINE (.-.) #-} ; (.+^) = (^+^) ; {-# INLINE (.+^) #-} ; \
  (.-^) = (^-^) ; {-# INLINE (.-^) #-}
#define ADDITIVE(T) ADDITIVEC((), T)

ADDITIVE([])
ADDITIVE(Complex)
ADDITIVE(ZipList)
ADDITIVE(Maybe)
ADDITIVE(IntMap)
ADDITIVE(Identity)
ADDITIVE(Vector)
ADDITIVE(V0)
ADDITIVE(V1)
ADDITIVE(V2)
ADDITIVE(V3)
ADDITIVE(V4)
ADDITIVE(Plucker)
ADDITIVE(Quaternion)
ADDITIVE(((->) b))
ADDITIVEC(Ord k, (Map k))
ADDITIVEC((Eq k, Hashable k), (HashMap k))
ADDITIVEC(Dim n, (V n))

-- | A handy wrapper to help distinguish points from vectors at the
-- type level
newtype Point f a = P (f a)
  deriving ( Eq, Ord, Show, Read, Monad, Functor, Applicative, Foldable
           , Eq1, Ord1, Show1, Read1
           , Traversable, Apply, Additive, Metric
           , Fractional , Num, Ix, Storable, Epsilon
           , Hashable
#if __GLASGOW_HASKELL__ >= 702
           , Generic
#endif
#if __GLASGOW_HASKELL__ >= 706
           , Generic1
#endif
#if __GLASGOW_HASKELL__ >= 708
           , Typeable, Data
#endif
           )

instance NFData (f a) => NFData (Point f a) where
  rnf (P x) = rnf x

instance Serial1 f => Serial1 (Point f) where
  serializeWith f (P p) = serializeWith f p
  deserializeWith m = P `liftM` deserializeWith m

instance Serial (f a) => Serial (Point f a) where
  serialize (P p) = serialize p
  deserialize = P `liftM` deserialize

instance Binary (f a) => Binary (Point f a) where
  put (P p) = Binary.put p
  get = P `liftM` Binary.get

instance Serialize (f a) => Serialize (Point f a) where
  put (P p) = Cereal.put p
  get = P `liftM` Cereal.get

#if __GLASGOW_HASKELL__ < 708
instance forall f. Typeable1 f => Typeable1 (Point f) where
  typeOf1 _ = mkTyConApp (mkTyCon3 "linear" "Linear.Affine" "Point") [] `mkAppTy`
              typeOf1 (undefined :: f a)

deriving instance (Data (f a), Typeable1 f, Typeable a) => Data (Point f a)
#endif

lensP :: Lens' (Point g a) (g a)
lensP afb (P a) = P <$> afb a
{-# INLINE lensP #-}

_Point :: Iso' (Point f a) (f a)
_Point = iso (\(P a) -> a) P
{-# INLINE _Point #-}

instance (t ~ Point g b) => Rewrapped (Point f a) t
instance Wrapped (Point f a) where
  type Unwrapped (Point f a) = f a
  _Wrapped' = _Point
  {-# INLINE _Wrapped' #-}

instance Bind f => Bind (Point f) where
  join (P m) = P $ join $ fmap (\(P m')->m') m

instance Distributive f => Distributive (Point f) where
  distribute = P . collect (\(P p) -> p)

instance Representable f => Representable (Point f) where
  type Rep (Point f) = Rep f
  tabulate f = P (tabulate f)
  {-# INLINE tabulate #-}
  index (P xs) = Rep.index xs
  {-# INLINE index #-}

type instance Index (Point f a) = Index (f a)
type instance IxValue (Point f a) = IxValue (f a)

instance Ixed (f a) => Ixed (Point f a) where
  ix l = lensP . ix l
  {-# INLINE ix #-}

instance Traversable f => Each (Point f a) (Point f b) a b where
  each = traverse
  {-# INLINE each #-}

instance R1 f => R1 (Point f) where
  _x = lensP . _x
  {-# INLINE _x #-}

instance R2 f => R2 (Point f) where
  _y = lensP . _y
  {-# INLINE _y #-}
  _xy = lensP . _xy
  {-# INLINE _xy #-}

instance R3 f => R3 (Point f) where
  _z = lensP . _z
  {-# INLINE _z #-}
  _xyz = lensP . _xyz
  {-# INLINE _xyz #-}

instance R4 f => R4 (Point f) where
  _w = lensP . _w
  {-# INLINE _w #-}
  _xyzw = lensP . _xyzw
  {-# INLINE _xyzw #-}

instance Additive f => Affine (Point f) where
  type Diff (Point f) = f
  P x .-. P y = x ^-^ y
  {-# INLINE (.-.) #-}
  P x .+^ v = P (x ^+^ v)
  {-# INLINE (.+^) #-}
  P x .-^ v = P (x ^-^ v)
  {-# INLINE (.-^) #-}

-- | Vector spaces have origins.
origin :: (Additive f, Num a) => Point f a
origin = P zero

-- | An isomorphism between points and vectors, given a reference
--   point.
relative :: (Additive f, Num a) => Point f a -> Iso' (Point f a) (f a)
relative p0 = iso (.-. p0) (p0 .+^)
{-# INLINE relative #-}

data instance U.Vector    (Point f a) =  V_P !(U.Vector    (f a))
data instance U.MVector s (Point f a) = MV_P !(U.MVector s (f a))
instance U.Unbox (f a) => U.Unbox (Point f a)

instance U.Unbox (f a) => M.MVector U.MVector (Point f a) where
  {-# INLINE basicLength #-}
  {-# INLINE basicUnsafeSlice #-}
  {-# INLINE basicOverlaps #-}
  {-# INLINE basicUnsafeNew #-}
  {-# INLINE basicUnsafeRead #-}
  {-# INLINE basicUnsafeWrite #-}
  basicLength (MV_P v) = M.basicLength v
  basicUnsafeSlice m n (MV_P v) = MV_P (M.basicUnsafeSlice m n v)
  basicOverlaps (MV_P v) (MV_P u) = M.basicOverlaps v u
  basicUnsafeNew n = MV_P `liftM` M.basicUnsafeNew n
  basicUnsafeRead (MV_P v) i = P `liftM` M.basicUnsafeRead v i
  basicUnsafeWrite (MV_P v) i (P x) = M.basicUnsafeWrite v i x
#if MIN_VERSION_vector(0,11,0)
  basicInitialize (MV_P v) = M.basicInitialize v
  {-# INLINE basicInitialize #-}
#endif

instance U.Unbox (f a) => G.Vector U.Vector (Point f a) where
  {-# INLINE basicUnsafeFreeze #-}
  {-# INLINE basicUnsafeThaw   #-}
  {-# INLINE basicLength       #-}
  {-# INLINE basicUnsafeSlice  #-}
  {-# INLINE basicUnsafeIndexM #-}
  basicUnsafeFreeze (MV_P v) = V_P `liftM` G.basicUnsafeFreeze v
  basicUnsafeThaw   ( V_P v) = MV_P `liftM` G.basicUnsafeThaw   v
  basicLength       ( V_P v) = G.basicLength v
  basicUnsafeSlice m n (V_P v) = V_P (G.basicUnsafeSlice m n v)
  basicUnsafeIndexM (V_P v) i = P `liftM` G.basicUnsafeIndexM v i