{-# LANGUAGE CPP                   #-}
{-# LANGUAGE DataKinds             #-}
{-# LANGUAGE GADTs                 #-}
{-# LANGUAGE KindSignatures        #-}
{-# LANGUAGE TypeFamilies          #-}
{-# LANGUAGE TypeOperators         #-}
{-# LANGUAGE StandaloneDeriving    #-}
{-# LANGUAGE FlexibleContexts      #-}
{-# LANGUAGE ScopedTypeVariables   #-}
{-# LANGUAGE RankNTypes            #-}
{-# LANGUAGE UndecidableInstances  #-}
{-|
Module      : Grenade.Core.Shape
Description : Dependently typed shapes of data which are passed between layers of a network
Copyright   : (c) Huw Campbell, 2016-2017
License     : BSD2
Stability   : experimental

-}
module Grenade.Core.Shape (
    Shape (..)
  , S (..)
  , Sing (..)

  , randomOfShape
  , fromStorable
  ) where

import           Control.DeepSeq (NFData (..))
import           Control.Monad.Random ( MonadRandom, getRandom )

import           Data.Proxy
import           Data.Serialize
import           Data.Singletons
import           Data.Singletons.TypeLits
import           Data.Vector.Storable ( Vector )
import qualified Data.Vector.Storable as V

#if MIN_VERSION_base(4,11,0)
import           GHC.TypeLits hiding (natVal)
#else
import           GHC.TypeLits
#endif

import qualified Numeric.LinearAlgebra.Static as H
import           Numeric.LinearAlgebra.Static
import qualified Numeric.LinearAlgebra as NLA

-- | The current shapes we accept.
--   at the moment this is just one, two, and three dimensional
--   Vectors/Matricies.
--
--   These are only used with DataKinds, as Kind `Shape`, with Types 'D1, 'D2, 'D3.
data Shape
  = D1 Nat
  -- ^ One dimensional vector
  | D2 Nat Nat
  -- ^ Two dimensional matrix. Row, Column.
  | D3 Nat Nat Nat
  -- ^ Three dimensional matrix. Row, Column, Channels.

-- | Concrete data structures for a Shape.
--
--   All shapes are held in contiguous memory.
--   3D is held in a matrix (usually row oriented) which has height depth * rows.
data S (n :: Shape) where
  S1D :: ( KnownNat len )
      => R len
      -> S ('D1 len)

  S2D :: ( KnownNat rows, KnownNat columns )
      => L rows columns
      -> S ('D2 rows columns)

  S3D :: ( KnownNat rows
         , KnownNat columns
         , KnownNat depth
         , KnownNat (rows * depth))
      => L (rows * depth) columns
      -> S ('D3 rows columns depth)

deriving instance Show (S n)

-- Singleton instances.
--
-- These could probably be derived with template haskell, but this seems
-- clear and makes adding the KnownNat constraints simple.
-- We can also keep our code TH free, which is great.
data instance Sing (n :: Shape) where
  D1Sing :: Sing a -> Sing ('D1 a)
  D2Sing :: Sing a -> Sing b -> Sing ('D2 a b)
  D3Sing :: KnownNat (a * c) => Sing a -> Sing b -> Sing c -> Sing ('D3 a b c)

instance KnownNat a => SingI ('D1 a) where
  sing = D1Sing sing
instance (KnownNat a, KnownNat b) => SingI ('D2 a b) where
  sing = D2Sing sing sing
instance (KnownNat a, KnownNat b, KnownNat c, KnownNat (a * c)) => SingI ('D3 a b c) where
  sing = D3Sing sing sing sing

instance SingI x => Num (S x) where
  (+) = n2 (+)
  (-) = n2 (-)
  (*) = n2 (*)
  abs = n1 abs
  signum = n1 signum
  fromInteger x = nk (fromInteger x)

instance SingI x => Fractional (S x) where
  (/) = n2 (/)
  recip = n1 recip
  fromRational x = nk (fromRational x)

instance SingI x => Floating (S x) where
  pi = nk pi
  exp = n1 exp
  log = n1 log
  sqrt = n1 sqrt
  (**) = n2 (**)
  logBase = n2 logBase
  sin = n1 sin
  cos = n1 cos
  tan = n1 tan
  asin = n1 asin
  acos = n1 acos
  atan = n1 atan
  sinh = n1 sinh
  cosh = n1 cosh
  tanh = n1 tanh
  asinh = n1 asinh
  acosh = n1 acosh
  atanh = n1 atanh

--
-- I haven't made shapes strict, as sometimes they're not needed
-- (the last input gradient back for instance)
--
instance NFData (S x) where
  rnf (S1D x) = rnf x
  rnf (S2D x) = rnf x
  rnf (S3D x) = rnf x

-- | Generate random data of the desired shape
randomOfShape :: forall x m. ( MonadRandom m, SingI x ) => m (S x)
randomOfShape = do
  seed :: Int <- getRandom
  return $ case (sing :: Sing x) of
    D1Sing SNat ->
        S1D (randomVector  seed Uniform * 2 - 1)

    D2Sing SNat SNat ->
        S2D (uniformSample seed (-1) 1)

    D3Sing SNat SNat SNat ->
        S3D (uniformSample seed (-1) 1)

-- | Generate a shape from a Storable Vector.
--
--   Returns Nothing if the vector is of the wrong size.
fromStorable :: forall x. SingI x => Vector Double -> Maybe (S x)
fromStorable xs = case sing :: Sing x of
    D1Sing SNat ->
      S1D <$> H.create xs

    D2Sing SNat SNat ->
      S2D <$> mkL xs

    D3Sing SNat SNat SNat ->
      S3D <$> mkL xs
  where
    mkL :: forall rows columns. (KnownNat rows, KnownNat columns)
        => Vector Double -> Maybe (L rows columns)
    mkL v =
      let rows    = fromIntegral $ natVal (Proxy :: Proxy rows)
          columns = fromIntegral $ natVal (Proxy :: Proxy columns)
      in  if rows * columns == V.length v
             then H.create $ NLA.reshape columns v
             else Nothing

instance SingI x => Serialize (S x) where
  put i = (case i of
            (S1D x) -> putListOf put . NLA.toList . H.extract $ x
            (S2D x) -> putListOf put . NLA.toList . NLA.flatten . H.extract $ x
            (S3D x) -> putListOf put . NLA.toList . NLA.flatten . H.extract $ x
          ) :: PutM ()

  get = do
    Just i <- fromStorable . V.fromList <$> getListOf get
    return i

-- Helper function for creating the number instances
n1 :: ( forall a. Floating a => a -> a ) -> S x -> S x
n1 f (S1D x) = S1D (f x)
n1 f (S2D x) = S2D (f x)
n1 f (S3D x) = S3D (f x)

-- Helper function for creating the number instances
n2 :: ( forall a. Floating a => a -> a -> a ) -> S x -> S x -> S x
n2 f (S1D x) (S1D y) = S1D (f x y)
n2 f (S2D x) (S2D y) = S2D (f x y)
n2 f (S3D x) (S3D y) = S3D (f x y)

-- Helper function for creating the number instances
nk :: forall x. SingI x => Double -> S x
nk x = case (sing :: Sing x) of
  D1Sing SNat ->
    S1D (konst x)

  D2Sing SNat SNat ->
    S2D (konst x)

  D3Sing SNat SNat SNat ->
    S3D (konst x)


{-# LANGUAGE TypeApplications #-}
{-# OPTIONS_GHC -fno-warn-type-defaults #-}

module ScalarWaveOld where

import Control.Exception (assert)
import Control.Monad.Loops
import qualified Data.Vector.Unboxed as U
import Data.VectorSpace
import Dual
import EquationsOld
import Foreign
import qualified Numeric.LinearAlgebra as L
import SSSFN
import System.CPUTime

default (Int)

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

bndw :: Floating a => Coord a -> a
bndw (Coord u v) = swamp * cos (pi / 4 * swk * t) * cos (pi / 4 * swk * r)
  where
    t = v + u
    r = v - u
    swk = 1
    swamp = 0.1

guess1 :: Floating a => Coord a -> State a
-- guess1 c = State 1 0 (bndw c)
guess1 c = State 1 0 0

guess :: (L.Container L.Vector a, Floating a) => Grid (State a)
guess = gsample (\(u, v) -> guess1 (Coord u v))

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

eqnInt1 :: Fractional a => Eqn -> Coord a -> State a -> State a -> State a -> State a -> State a -> State a -> a
eqnInt1 eqn c s su sv suv suu svv =
  case eqn of
    EqnG01 -> eqnG01 c s su sv suv
    EqnG22 -> eqnG22 c s su sv suv
    EqnSW2 -> eqnSW2 c s su sv suv
    EqnG00 -> eqnG00 c s su suu
    EqnG11 -> eqnG11 c s sv svv

eqnAxis1 :: Fractional a => Eqn -> Coord a -> State a -> State a -> State a -> State a -> State a -> State a -> a
eqnAxis1 eqn c@(Coord u v) s@(State g q w) su@(State gu qu wu) sv@(State gv qv wv) suv@(State guv quv wuv) suu@(State guu quu wuu) svv@(State gvv qvv wvv) =
  case eqn of
    EqnG01 -> eqnAxisG01 c s su sv suv
    EqnG22 -> eqnAxisG22 c s su sv suv
    EqnSW2 -> eqnAxisSW2 c s su sv suv
    EqnG00 -> eqnAxisG00 c s su sv suu
    EqnG11 -> eqnAxisG11 c s su sv svv

eqnOrig1 :: Floating a => Eqn -> Coord a -> State a -> State a -> State a -> State a -> State a -> State a -> a
eqnOrig1 eqn c s@(State g q w) su sv suv suu svv =
  case eqn of
    EqnG01 -> g - 1
    EqnG22 -> q - 0
    EqnSW2 -> w - bndw c
    EqnG00 -> 0
    EqnG11 -> 0

-- i == 0, u == -1
eqnBndU1 :: (Eq a, Floating a) => Eqn -> Coord a -> State a -> State a -> State a -> State a -> State a -> State a -> a
eqnBndU1 eqn c@(Coord u v) s@(State g q w) _ sv _ _ svv =
  assert (u == -1) $
    case eqn of
      EqnG01 -> eqnG11 c s sv svv
      EqnG22 -> q - 0
      EqnSW2 -> w - bndw c
      EqnG00 -> 0
      EqnG11 -> 0

-- j == 0, v == -1
eqnBndV1 :: (Eq a, Floating a) => Eqn -> Coord a -> State a -> State a -> State a -> State a -> State a -> State a -> a
eqnBndV1 eqn c@(Coord u v) s@(State g q w) su _ _ suu _ =
  assert (v == -1) $
    case eqn of
      EqnG01 -> eqnG00 c s su suu
      EqnG22 -> q - 0
      EqnSW2 -> w - bndw c
      EqnG00 -> 0
      EqnG11 -> 0

eqnEndU1 :: Floating a => Eqn -> Coord a -> State a -> State a -> State a -> State a -> State a -> State a -> a
eqnEndU1 eqn c@(Coord u v) s@(State g q w) _ _ _ _ _ =
  case eqn of
    EqnG01 -> g - 1
    EqnG22 -> q - 0
    EqnSW2 -> w - bndw c
    EqnG00 -> 0
    EqnG11 -> 0

eqnEndV1 :: Floating a => Eqn -> Coord a -> State a -> State a -> State a -> State a -> State a -> State a -> a
eqnEndV1 eqn c@(Coord u v) (State g q w) _ _ _ _ _ =
  case eqn of
    EqnG01 -> g - 1
    EqnG22 -> q - 0
    EqnSW2 -> w - bndw c
    EqnG00 -> 0
    EqnG11 -> 0

equation1 :: (Eq a, Floating a) => Coord Int -> Eqn -> Coord a -> State a -> State a -> State a -> State a -> State a -> State a -> a
equation1 (Coord i j) =
  if
    | i == 0 && j == 0 -> eqnOrig1
    | i == 0 && j == np - 1 -> eqnEndU1
    | i == np - 1 && j == 0 -> eqnEndV1
    | i == 0 -> eqnBndU1
    | j == 0 -> eqnBndV1
    | i == j -> eqnAxis1
    | True -> eqnInt1

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

g2s :: Storable a => Grid a -> State (Grid a)
g2s (Grid xs) =
  let [g, q, w] = L.takesV (replicate 3 (np * np)) xs
   in State (Grid g) (Grid q) (Grid w)

s2g :: Storable a => State (Grid a) -> Grid a
s2g (State (Grid g) (Grid q) (Grid w)) = Grid (L.vjoin [g, q, w])

eqn1 ::
  forall a.
  (Eq a, L.Field a, U.Unbox a) =>
  Eqn ->
  Coord (Grid a) ->
  State (Grid a) ->
  Grid a
eqn1 eqn c s =
  Grid $
    L.fromList
      [ eqnRow i j
        | i <- [0 .. np - 1],
          j <- [0 .. np - 1]
      ]
  where
    su = (derivU #>) <$> s
    sv = (derivV #>) <$> s
    suv = (derivUV #>) <$> s
    suu = (derivUU #>) <$> s
    svv = (derivVV #>) <$> s
    eqnRow :: Int -> Int -> a
    eqnRow i j =
      let at :: Grid a -> a
          at x = x ! (i, j)
       in equation1
            (Coord i j)
            eqn
            (at <$> c)
            (at <$> s)
            (at <$> su)
            (at <$> sv)
            (at <$> suv)
            (at <$> suu)
            (at <$> svv)

eqns ::
  forall a.
  (Eq a, L.Field a, U.Unbox a) =>
  Coord (Grid a) ->
  State (Grid a) ->
  Grid a
eqns c s =
  Grid $
    L.fromList
      [ eqnRow (i, j, eqn)
        | eqn <- [EqnG01, EqnG22, EqnSW2],
          i <- [0 .. np - 1],
          j <- [0 .. np - 1]
      ]
  where
    su = (derivU #>) <$> s
    sv = (derivV #>) <$> s
    suv = (derivUV #>) <$> s
    suu = (derivUU #>) <$> s
    svv = (derivVV #>) <$> s
    eqnRow :: (Int, Int, Eqn) -> a
    eqnRow (i, j, eqn) =
      let at :: Grid a -> a
          at x = x ! (i, j)
       in equation1
            (Coord i j)
            eqn
            (at <$> c)
            (at <$> s)
            (at <$> su)
            (at <$> sv)
            (at <$> suv)
            (at <$> suu)
            (at <$> svv)

eqns5 ::
  forall a.
  (Eq a, L.Field a, U.Unbox a) =>
  Coord (Grid a) ->
  State (Grid a) ->
  Grid a
eqns5 c s =
  Grid $
    L.fromList
      [ eqnRow (i, j, eqn)
        | eqn <- [EqnG01, EqnG22, EqnSW2, EqnG00, EqnG11],
          i <- [0 .. np - 1],
          j <- [0 .. np - 1]
      ]
  where
    su = (derivU #>) <$> s
    sv = (derivV #>) <$> s
    suv = (derivUV #>) <$> s
    suu = (derivUU #>) <$> s
    svv = (derivVV #>) <$> s
    eqnRow :: (Int, Int, Eqn) -> a
    eqnRow (i, j, eqn) =
      let at :: Grid a -> a
          at x = x ! (i, j)
       in equation1
            (Coord i j)
            eqn
            (at <$> c)
            (at <$> s)
            (at <$> su)
            (at <$> sv)
            (at <$> suv)
            (at <$> suu)
            (at <$> svv)

op ::
  forall a.
  (L.Container L.Vector a, Eq a, L.Field a, U.Unbox a) =>
  Coord (Grid a) ->
  State (Grid a) ->
  Op a
op (Coord u v) s =
  Op $
    L.fromRows
      [ jacRow (i, j, eqn)
        | eqn <- [EqnG01, EqnG22, EqnSW2],
          i <- [0 .. np - 1],
          j <- [0 .. np - 1]
      ]
  where
    su = (derivU #>) <$> s
    sv = (derivV #>) <$> s
    suv = (derivUV #>) <$> s
    suu = (derivUU #>) <$> s
    svv = (derivVV #>) <$> s
    jacRow :: (Int, Int, Eqn) -> L.Vector a
    jacRow (i, j, eqn) =
      L.vjoin $
        [ let con :: Grid a -> Dual a
              der :: Grid a -> Dual a
              con x = Dual (x ! (i, j)) 0
              der x = Dual (x ! (i, j)) 1
              scon :: State (Grid a) -> State (Dual a)
              sder :: State (Grid a) -> State (Dual a)
              scon (State g q w) = State (con g) (con q) (con w)
              sder (State g q w) = case var of
                VarG -> State (der g) (con q) (con w)
                VarQ -> State (con g) (der q) (con w)
                VarW -> State (con g) (con q) (der w)
           in ( ( dual $
                    equation1
                      (Coord i j)
                      eqn
                      (Coord (con u) (con v))
                      (sder s)
                      (scon su)
                      (scon sv)
                      (scon suv)
                      (scon suu)
                      (scon svv)
                )
                  `L.scale` (getOp delta L.! (i * np + j))
              )
                `L.add` ( ( dual $
                              equation1
                                (Coord i j)
                                eqn
                                (Coord (con u) (con v))
                                (scon s)
                                (sder su)
                                (scon sv)
                                (scon suv)
                                (scon suu)
                                (scon svv)
                          )
                            `L.scale` (getOp derivU L.! (i * np + j))
                        )
                `L.add` ( ( dual $
                              equation1
                                (Coord i j)
                                eqn
                                (Coord (con u) (con v))
                                (scon s)
                                (scon su)
                                (sder sv)
                                (scon suv)
                                (scon suu)
                                (scon svv)
                          )
                            `L.scale` (getOp derivV L.! (i * np + j))
                        )
                `L.add` ( ( dual $
                              equation1
                                (Coord i j)
                                eqn
                                (Coord (con u) (con v))
                                (scon s)
                                (scon su)
                                (scon sv)
                                (sder suv)
                                (scon suu)
                                (scon svv)
                          )
                            `L.scale` (getOp derivUV L.! (i * np + j))
                        )
                `L.add` ( ( dual $
                              equation1
                                (Coord i j)
                                eqn
                                (Coord (con u) (con v))
                                (scon s)
                                (scon su)
                                (scon sv)
                                (scon suv)
                                (sder suu)
                                (scon svv)
                          )
                            `L.scale` (getOp derivUU L.! (i * np + j))
                        )
                `L.add` ( ( dual $
                              equation1
                                (Coord i j)
                                eqn
                                (Coord (con u) (con v))
                                (scon s)
                                (scon su)
                                (scon sv)
                                (scon suv)
                                (scon suu)
                                (sder svv)
                          )
                            `L.scale` (getOp derivVV L.! (i * np + j))
                        )
          | var <- [VarG, VarQ, VarW]
        ]

op5 ::
  forall a.
  (L.Container L.Vector a, Eq a, L.Field a, U.Unbox a) =>
  Coord (Grid a) ->
  State (Grid a) ->
  Op a
op5 (Coord u v) s =
  Op $
    L.fromRows
      [ jacRow (i, j, eqn)
        | eqn <- [EqnG01, EqnG22, EqnSW2, EqnG00, EqnG11],
          i <- [0 .. np - 1],
          j <- [0 .. np - 1]
      ]
  where
    su = (derivU #>) <$> s
    sv = (derivV #>) <$> s
    suv = (derivUV #>) <$> s
    suu = (derivUU #>) <$> s
    svv = (derivVV #>) <$> s
    jacRow :: (Int, Int, Eqn) -> L.Vector a
    jacRow (i, j, eqn) =
      L.vjoin $
        [ let con :: Grid a -> Dual a
              der :: Grid a -> Dual a
              con x = Dual (x ! (i, j)) 0
              der x = Dual (x ! (i, j)) 1
              scon :: State (Grid a) -> State (Dual a)
              sder :: State (Grid a) -> State (Dual a)
              scon (State g q w) = State (con g) (con q) (con w)
              sder (State g q w) = case var of
                VarG -> State (der g) (con q) (con w)
                VarQ -> State (con g) (der q) (con w)
                VarW -> State (con g) (con q) (der w)
           in ( ( dual $
                    equation1
                      (Coord i j)
                      eqn
                      (Coord (con u) (con v))
                      (sder s)
                      (scon su)
                      (scon sv)
                      (scon suv)
                      (scon suu)
                      (scon svv)
                )
                  `L.scale` (getOp delta L.! (i * np + j))
              )
                `L.add` ( ( dual $
                              equation1
                                (Coord i j)
                                eqn
                                (Coord (con u) (con v))
                                (scon s)
                                (sder su)
                                (scon sv)
                                (scon suv)
                                (scon suu)
                                (scon svv)
                          )
                            `L.scale` (getOp derivU L.! (i * np + j))
                        )
                `L.add` ( ( dual $
                              equation1
                                (Coord i j)
                                eqn
                                (Coord (con u) (con v))
                                (scon s)
                                (scon su)
                                (sder sv)
                                (scon suv)
                                (scon suu)
                                (scon svv)
                          )
                            `L.scale` (getOp derivV L.! (i * np + j))
                        )
                `L.add` ( ( dual $
                              equation1
                                (Coord i j)
                                eqn
                                (Coord (con u) (con v))
                                (scon s)
                                (scon su)
                                (scon sv)
                                (sder suv)
                                (scon suu)
                                (scon svv)
                          )
                            `L.scale` (getOp derivUV L.! (i * np + j))
                        )
                `L.add` ( ( dual $
                              equation1
                                (Coord i j)
                                eqn
                                (Coord (con u) (con v))
                                (scon s)
                                (scon su)
                                (scon sv)
                                (scon suv)
                                (sder suu)
                                (scon svv)
                          )
                            `L.scale` (getOp derivUU L.! (i * np + j))
                        )
                `L.add` ( ( dual $
                              equation1
                                (Coord i j)
                                eqn
                                (Coord (con u) (con v))
                                (scon s)
                                (scon su)
                                (scon sv)
                                (scon suv)
                                (scon suu)
                                (sder svv)
                          )
                            `L.scale` (getOp derivVV L.! (i * np + j))
                        )
          | var <- [VarG, VarQ, VarW]
        ]

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

getTime :: IO Double
getTime = do
  t <- getCPUTime
  return $ fromInteger t / 1.0e+12

main :: IO ()
main =
  do
    putStrLn "Spherically Symmetric Scalar Field in double-Null coordinates"
    let coords = Coord coordU (coordV @Double)
    let ini = State (gmap g guess) (gmap q guess) (gmap w guess)
    vars0 <- return (s2g ini)
    let done (n, vars) =
          let res = eqns coords (g2s vars)
           in -- let res = eqns5 coords (g2s vars)
              gmaxabs res < 1.0e-12 || n >= 5
    let iter (n, vars) = do
          let res = eqns coords (g2s vars)
          -- let res = eqns5 coords (g2s vars)
          t <- getTime
          putStrLn $
            "iter " ++ show n ++ " time " ++ show (round t) ++ " sec |res| "
              ++ show (gmaxabs res)
          let jac = op coords (g2s vars)
          let dvars = solve jac res
          -- let jac = op5 coords (g2s vars)
          -- let dvars = solveLS jac res
          return (n + 1, vars ^-^ dvars)
    (n1, vars1) <- iterateUntilM done iter (0, vars0)
    putStrLn $ "var.g " ++ show (gmap approx $ g $ g2s vars1)
    putStrLn $ "var.q " ++ show (gmap approx $ q $ g2s vars1)
    putStrLn $ "var.w " ++ show (gmap approx $ w $ g2s vars1)
    putStrLn $ "res.G01 " ++ show (gmap approx $ eqn1 EqnG01 coords (g2s vars1))
    putStrLn $ "res.G22 " ++ show (gmap approx $ eqn1 EqnG22 coords (g2s vars1))
    putStrLn $ "res.SW2 " ++ show (gmap approx $ eqn1 EqnSW2 coords (g2s vars1))
    putStrLn $ "res.G00 " ++ show (gmap approx $ eqn1 EqnG00 coords (g2s vars1))
    putStrLn $ "res.G11 " ++ show (gmap approx $ eqn1 EqnG11 coords (g2s vars1))
    let res1 = eqns coords (g2s vars1)
    t <- getTime
    putStrLn $
      "iter " ++ show n1 ++ " time " ++ show (round t) ++ " sec |res| "
        ++ show (gmaxabs res1)
    let res5 = eqns5 coords (g2s vars1)
    putStrLn $ "|res5| " ++ show (gmaxabs res5)
    -- let h = 2 / (fromIntegral np - 1) :: Double
    -- let alpha = 0.01 * h ^ 3
    -- let smooth = delta - alpha *^ (derivUU * derivUU + derivVV * derivVV)
    -- let vars1s = s2g $ (expfilter #>) <$> g2s vars1
    -- let res5s = eqns5 coords (g2s vars1s)
    -- putStrLn $ "|res5s| " ++ show (gmaxabs res5s)
    putStrLn "Done."
  where
    approx :: RealFrac a => a -> a
    approx x = fromInteger (round (1.0e+10 * x)) / 1.0e+10


-- |
-- Module      : Cartesian.Plane.Core
-- Description :
-- Copyright   : (c) Jonatan H Sundqvist, year
-- License     : MIT
-- Maintainer  : Jonatan H Sundqvist
-- Stability   : experimental|stable
-- Portability : POSIX (not sure)
--

-- Created date year

-- TODO | - Which constraints are appropriate (Num is probably too generic, should be Real, maybe RealFrac)
--        - Strictness, performance
--        - Rename (?)

-- SPEC | -
--        -

--------------------------------------------------------------------------------------------------------------------------------------------
-- API
--------------------------------------------------------------------------------------------------------------------------------------------
module Cartesian.Plane.Core where

--------------------------------------------------------------------------------------------------------------------------------------------
-- We'll need these
--------------------------------------------------------------------------------------------------------------------------------------------
-- import Data.List (sort, minimumBy)
-- import Data.Ord  (comparing)
-- import Data.Complex hiding (magnitude)

-- import           Control.Monad (when)
import           Control.Applicative

import           Control.Lens ((^.))

import Cartesian.Internal.Types
import Cartesian.Internal.Lenses (begin, end)
import Cartesian.Internal.Core

--------------------------------------------------------------------------------------------------------------------------------------------
-- Functions
--------------------------------------------------------------------------------------------------------------------------------------------

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

-- | Determines if a point lies within a polygon using the odd-even method.
--
-- TODO: Use epsilon (?)
-- TODO: How to treat points that lie on an edge
-- inside :: Num n => Polygon n -> Vector2D n -> Bool
-- inside polygon (Vector2D px py) = error "Cartesian.Plane.inside is still a work in progress"
--   where
--     edges = polygon ++ [head polygon] -- Close the loop
    -- between (Line (Vector ax ay) (Vector bx by)) = _

-- |
-- instance Convertible (Vector2D f, Vector3D f) where
  -- _

-- TODO: Use type families for this stuff (?)

-- |
-- to3D :: Num f => Vector2D f -> Vector3D f
-- to3D (Vector2D x' y') = Vector3D x' y' 0

-- -- |
-- from3D :: Num f => Vector3D f -> Vector2D f
-- from3D (Vector3D x' y' _) = Vector2D x' y'

-- -- | Perform some unary operation on a 2D vector as a 3D vector, converting the result back to 2D by discarding the z component.
-- -- TODO: Rename (?)
-- -- TODO: Loosen Num restriction (eg. to anything with a 'zero' value) (?)
-- in3D :: (Num f, Num f') => (Vector3D f -> Vector3D f') -> Vector2D f -> Vector2D f'
-- in3D f = from3D . f . to3D

-- | Same as in3D, but for binary operations.
-- _ :: _
-- _ f = from3D . (dotmap f) . to3D

-- Vector math -----------------------------------------------------------------------------------------------------------------------------

-- | Angle (in radians) between the positive X-axis and the vector
-- argument :: (Floating a, Eq a) => Vector a -> a
-- argument (Vector 0 0) = 0
-- argument (Vector x y) = atan $ y/x
--
--
-- arg :: (Floating a, Eq a) => Vector a -> a
-- arg = argument
--
--
-- -- | Vector -> (magnitude, argument)
-- polar :: (Floating a, Eq a) => Vector a -> (a, a)
-- polar v@(Vector x y) = (magnitude v, argument v)

-- Linear functions ------------------------------------------------------------------------------------------------------------------------

-- | Yields the intersection point of two finite lines. The lines are defined inclusively by
--   their endpoints. The result is wrapped in a Maybe value to account for non-intersecting
--   lines.
-- TODO: Refactor
-- TODO: Move equation solving to separate function (two linear functions)
-- TODO: Invariants, check corner cases
-- TODO: Deal with vertical lines
-- TODO: Factor out infinite-line logic
-- TODO: Decide how to deal with identical lines
-- TODO: Factor out domain logic (eg. write restrict or domain function)
-- TODO: Return Either instead of Maybe (eg. Left "parallel") (?)
-- TODO: Visual debugging functions
-- TODO: Math notes, MathJax or LaTex
-- TODO: Intersect for curves (functions) and single points (?)
-- TODO: Polymorphic, typeclass (lines, shapes, ranges, etc.) (?)
-- TODO: Intersect Rectangles
intersect :: RealFloat f => Line (Vector2D f) -> Line (Vector2D f) -> Maybe (Vector2D f)
intersect f' g' = do
  p <- mp
  indomain f' p
  indomain g' p
  where
    -- indomain :: RealFloat f => Line (Vector2D f) -> Vector2D f -> Maybe (Vector2D f)
    indomain h' = restrict (h'^.begin) (h'^.end) -- TODO: Rename

    -- mp :: Maybe (Vector2D f)
    mp = case (linear f', linear g') of
      (Just f, Nothing) -> let x' = g'^.begin.x in Just . Vector2D x' $ plotpoint f x'
      (Nothing, Just g) -> let x' = f'^.begin.x in Just . Vector2D x' $ plotpoint g x'
      (Just f,  Just g) -> linearIntersect f g
      _                 -> Nothing

-- | Gives the linear function overlapping the given segment, or Nothing if there is no such function
linear :: RealFloat f => Line (Vector2D f) -> Maybe (Linear f)
linear line = Linear <$> intercept line <*> slope line

-- | Applies a linear function to the given value
-- TODO: Rename (?)
plotpoint :: RealFloat f => Linear f -> f -> f
plotpoint f x' = slopeOf f*x' + interceptOf f

-- | Finds the intersection (if any) of two linear functions
-- TODO: Use Epsilon (?)
-- TODO: Rename (eg. 'solve') (?)
linearIntersect :: RealFloat f => Linear f -> Linear f -> Maybe (Vector2D f)
linearIntersect f g
  | slopeOf f == slopeOf g = Nothing
  | otherwise = let x' = (β-b)/(a-α) in Just $ Vector2D x' (a*x' + b)
  where
    (a, α) = (slopeOf     f, slopeOf     g)
    (b, β) = (interceptOf f, interceptOf g)

-- |
slope :: RealFloat f => Line (Vector2D f) -> Maybe f
slope (Line fr to)
  | dx == 0   = Nothing
  | otherwise = Just $ dy/dx
  where
    (Vector2D dx dy) = liftA2 (-) to fr

-- |
intercept :: RealFloat f => Line (Vector2D f) -> Maybe f
intercept line = do
  slope' <- slope line
  return $ y' - slope'*x'
  where
    (Vector2D x' y') = line^.begin

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

-- | Ensures that a given point lies within the domain and codomain
-- TODO: Make polymorphic
-- TODO: Let this function work on scalars, write another function for domain and codomain (?)
-- restrict domain codomain p = _
-- restrict :: (Applicative v, Traversable v, Num n, Ord n) => v n -> v n -> v n -> Maybe (v n)
-- restrict a b p@(Vector2D x' y')
--   | indomain && incodomain = Just p
--   | otherwise              = Nothing
--   where
--     (Vector2D lowx lowy)   = dotwise min a b
--     (Vector2D highx highy) = dotwise max a b
--     indomain   = between lowx highx x'
--     incodomain = between lowy highy y'


-- This module has functions that store pinwheels to disk and read them out to do convolution,
-- for the purpose of handling the case when filter size is greater than RAM size.
module STC.Binary
  ( makePlanBinary
  , computeInitialEigenVectorBinary
  , writeDFTPinwheel
  , convolutionBinary
  ) where

import           Array.UnboxedArray           as AU
import           Control.Monad                as M
import           Control.Monad.IO.Class
import           Control.Monad.Trans.Resource
import           Data.Array.Repa              as R
import           Data.Binary
import           Data.ByteString.Lazy         as BSL
import           Data.ByteString.Lazy         as BS
import           Data.Complex
import           Data.Conduit
import           Control.DeepSeq              (($!!))
import           Data.Conduit.List            as CL
import           Data.List                    as L
import           Data.Vector.Storable         as VS
import           Data.Vector.Unboxed          as VU
import           DFT.Plan
import           Filter.Pinwheel
import           FokkerPlanck.Interpolation
import           FokkerPlanck.Pinwheel        (cutoff)
import           GHC.Float
import           System.IO
import           System.Random
import           Types
import           Utils.Array
import           Utils.Parallel               (ParallelParams (..),
                                               parConduitIO)

makePlanBinary ::
     DFTPlan -> Bool -> FilePath -> Int -> Int -> Int -> Int -> IO DFTPlan
makePlanBinary oldPlan wisdomFlag wisdomFilePath numThetaFreqs numScaleFreqs rows cols = do
  let n = numThetaFreqs * numScaleFreqs * rows * cols
  xs <- M.replicateM n randomIO :: IO [Double]
  ys <- M.replicateM n randomIO :: IO [Double]
  let vecTemp = VS.fromList . L.zipWith (:+) xs $ ys
  when wisdomFlag (importFFTWWisdom wisdomFilePath)
  lock <- getFFTWLock
  plan <-
    fst <$>
    (dft1dGPlan
       lock
       oldPlan
       [cols, rows, numThetaFreqs, numScaleFreqs]
       [0, 1]
       vecTemp >>= \(plan, vec) ->
       idft1dGPlan
         lock
         plan
         [cols, rows, numThetaFreqs, numScaleFreqs]
         [0, 1]
         vec >>= \(plan, vec) ->
         dft1dGPlan
           lock
           plan
           [numThetaFreqs, numScaleFreqs, cols, rows]
           [0, 1]
           vec >>= \(plan, vec) ->
           idft1dGPlan
             lock
             plan
             [numThetaFreqs, numScaleFreqs, cols, rows]
             [0, 1]
             vec)
  exportFFTWWisdom wisdomFilePath
  return plan

computeInitialEigenVectorBinary ::
     Int
  -> Int
  -> [Double]
  -> [Double]
  -> [R2S1RPPoint]
  -> R2Z2Array
computeInitialEigenVectorBinary xLen yLen thetaFreqs scaleFreqs xs =
  let numThetaFreqs = L.length thetaFreqs
      numScaleFreqs = L.length scaleFreqs
      xShift = div xLen 2
      (xMin, xMax) =
        if odd xLen
          then (-xShift, xShift)
          else (-xShift, xShift - 1)
      yShift = div yLen 2
      (yMin, yMax) =
        if odd yLen
          then (-yShift, yShift)
          else (-yShift, yShift - 1)
      vec =
        toUnboxedVector .
        AU.accum (+) 0 ((xMin, yMin), (xMax, yMax)) .
        L.map
          (\(R2S1RPPoint (x, y, _, _)) ->
             ((x, y), 1 / (fromIntegral . L.length $ xs))) $
        xs
  in computeS .
     R.traverse
       (fromUnboxed (Z :. xLen :. yLen) vec)
       (const (Z :. numThetaFreqs :. numScaleFreqs :. xLen :. yLen)) $ \f idx@(Z :. tf :. sf :. i :. j) ->
       if tf == div numThetaFreqs 2 && sf == div numScaleFreqs 2
         then f (Z :. i :. j)
         else 0

{-# INLINE convertDFTPinwheel #-}
convertDFTPinwheel ::
     DFTPlan
  -> R.Array D DIM5 Double
  -> R.Array U DIM1 Double
  -> R.Array U DIM1 Double
  -> Double
  -> Double
  -> (Int, Int)
  -> (Int, Int)
  -> IO BS.ByteString
convertDFTPinwheel plan radialArr thetaFreqs scaleFreqs hollowRadius rMax (rows, cols) (t, s) =
  let (Z :. numThetaFreq :. numScaleFreq :. numTheta0Freq :. numScale0Freq :. _) =
        extent radialArr
      pinwheelArr =
        traverse2
          thetaFreqs
          scaleFreqs
          (\_ _ -> (Z :. numTheta0Freq :. numScale0Freq :. cols :. rows)) $ \ft0 fs0 (Z :. t0 :. s0 :. i :. j) ->
          pinwheelFunc
            (PinwheelHollow0 hollowRadius)
            (thetaFreqs R.! (Z :. t) - ft0 (Z :. t0))
            (scaleFreqs R.! (Z :. s) + fs0 (Z :. s0))
            rMax
            0
            (i - center cols)
            (j - center rows)
      interpolatedPinwheelArr =
        radialCubicInterpolation
          (R.slice radialArr (Z :. t :. s :. All :. All :. All))
          1
          pinwheelArr
  in do xx <-
          fmap
            (encode .
             L.map (\(x :+ y) -> double2Float x :+ double2Float y) . VS.toList) .
          dftExecute
            plan
            (DFTPlanID DFT1DG [cols, rows, numTheta0Freq, numScale0Freq] [0, 1]) .
          VS.convert . toUnboxed . computeS . rotate4D2 . makeFilter2D $
          interpolatedPinwheelArr
        return $!! xx

{-# INLINE writeDFTPinwheelSink #-}
writeDFTPinwheelSink ::
     FilePath -> ConduitT BS.ByteString Void (ResourceT IO) ()
writeDFTPinwheelSink filePath = do
  h <- liftIO $ openBinaryFile filePath WriteMode
  go h
  where
    go h = do
      z <- await
      case z of
        Nothing -> liftIO $ hClose h
        Just bs -> do
          liftIO . BSL.hPut h . encode . BS.length $ bs
          liftIO . BS.hPut h $ bs
          go h

writeDFTPinwheel ::
     ParallelParams
  -> DFTPlan
  -> R.Array D DIM5 Double
  -> Double
  -> Double
  -> [Double]
  -> [Double]
  -> Double
  -> (Int, Int)
  -> FilePath
  -> IO ()
writeDFTPinwheel parallelParams plan radialArr hollowRadius cutoffRadius thetaFreqs scaleFreqs rMax (rows, cols) filePath = do
  let (Z :. numThetaFreq :. numScaleFreq :. numTheta0Freq :. numScale0Freq :. _) =
        extent radialArr
      thetaFreqsArr = fromListUnboxed (Z :. L.length thetaFreqs) thetaFreqs
      scaleFreqsArr = fromListUnboxed (Z :. L.length scaleFreqs) scaleFreqs
  runConduitRes $
    CL.sourceList
      [(t, s) | t <- [0 .. numThetaFreq - 1], s <- [0 .. numScaleFreq - 1]] .|
    parConduitIO
      parallelParams
      (convertDFTPinwheel
         plan
         (cutoff (round cutoffRadius) radialArr)
         thetaFreqsArr
         scaleFreqsArr
         hollowRadius
         rMax
         (rows, cols)) .|
    writeDFTPinwheelSink filePath

{-# INLINE sourceBinary #-}
sourceBinary :: FilePath -> ConduitT () BS.ByteString (ResourceT IO) ()
sourceBinary filePath = do
  h <- liftIO $ openBinaryFile filePath ReadMode
  go h
  where
    go handle = do
      lenBS <- liftIO (BSL.hGet handle 8)
      if BSL.null lenBS
        then liftIO $ hClose handle
        else do
          let len = decode lenBS :: Int
          bs <- liftIO . BS.hGet handle $ len
          yield bs
          go handle

{-# INLINE convolution #-}
convolution ::
     Bool
  -> DFTPlan
  -> DIM4
  -> R.Array U DIM1 Double
  -> VS.Vector (Complex Double)
  -> BS.ByteString
  -> IO (VU.Vector (Complex Double))
convolution True plan dim@(Z :. cols :. rows :. numThetaFreq :. numScaleFreq) thetaFreqs input bs = do
  arr <-
    fmap (fromUnboxed dim . VS.convert) .
    dftExecute
      plan
      (DFTPlanID IDFT1DG [cols, rows, numThetaFreq, numScaleFreq] [0, 1]) .
    VS.zipWith (*) input $
    VS.fromList . L.map (\(x :+ y) -> float2Double x :+ float2Double y) . decode $
    bs
  return $!! toUnboxed . sumS . sumS . R.traverse2 arr thetaFreqs const $
    (\f ft idx@(Z :. _ :. _ :. i :. _) -> f idx * exp (0 :+ ft (Z :. i) * pi))
convolution False plan dim@(Z :. cols :. rows :. numThetaFreq :. numScaleFreq) thetaFreqs input bs = do
  x <-
    fmap (toUnboxed . sumS . sumS . fromUnboxed dim . VS.convert) .
    dftExecute
      plan
      (DFTPlanID IDFT1DG [cols, rows, numThetaFreq, numScaleFreq] [0, 1]) .
    VS.zipWith (*) input .
    VS.fromList . L.map (\(x :+ y) -> float2Double x :+ float2Double y) . decode $
    bs
  return $!! x

convolutionBinary ::
     Bool
  -> ParallelParams
  -> DFTPlan
  -> FilePath
  -> [Double]
  -> R.Array U DIM4 (Complex Double)
  -> IO (R.Array U DIM4 (Complex Double))
convolutionBinary sinkFlag parallelParams plan filterFilePath thetaFreqs inputArr = do
  let thetaFreqsArr = fromListUnboxed (Z :. (L.length thetaFreqs)) thetaFreqs
      (Z :. numThetaFreq :. numScaleFreq :. cols :. rows) = extent inputArr
  inputArrF <-
    dftExecute
      plan
      (DFTPlanID DFT1DG [cols, rows, numThetaFreq, numScaleFreq] [0, 1]) .
    VU.convert . toUnboxed . computeS . rotate4D2 $
    inputArr
  xs <-
    runConduitRes $
    sourceBinary filterFilePath .|
    parConduitIO
      parallelParams
      (convolution
         sinkFlag
         plan
         (Z :. cols :. rows :. numThetaFreq :. numScaleFreq)
         thetaFreqsArr
         inputArrF) .|
    CL.consume
  let arr = fromUnboxed (extent inputArr) . VU.concat $ xs
  return $
    if sinkFlag
      then computeS . R.traverse2 arr thetaFreqsArr const $ \f ft idx@(Z :. i :. _ :. _ :. _) ->
             f idx * exp (0 :+ ft (Z :. i) * pi)
      else arr


{-# language FlexibleContexts #-}
{-# language GADTs #-}
{-# language ScopedTypeVariables #-}
{-# language TypeApplications #-}

module Verify where

import Feldspar
import Feldspar.Software
import Feldspar.Software.Verify
import Feldspar.Array.Vector
import Feldspar.Array.Buffered

import Data.Bits (Bits)
import Data.Complex (Complex)

import Prelude hiding ((==), (/=), (>), (<=), length, div, reverse, sum)

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

inc :: Software ()
inc = do
  len :: SExp Length <- fget stdin
  ix  :: SExp Index  <- fget stdin
  arr :: Arr (SExp Word32) <- newArr len
--  assert (ix <= length arr) "ix out of bounds"
  val <- getArr arr ix
  setArr arr ix (val + 1)

copy :: Software ()
copy = do
  arr1 :: Arr (SExp Word32) <- newArr 10
  arr2 :: Arr (SExp Word32) <- unsafeFreezeArr arr1 >>= unsafeThawArr
  --
  setArr arr1 0 0
  val <- getArr arr1 0
  setArr arr2 0 val

inplace :: Software ()
inplace = do
  st  :: Store Software (SExp Word32) <- newInPlaceStore 10
  arr :: Manifest Software (SExp Word32) <- store st $ (1...10)
  brr <- store st $ reverse arr
  val <- shareM $ sum brr
  return ()
 where
  reverse :: Manifest Software (SExp Word32) -> Push Software (SExp Word32)
  reverse = pairwise (\ix -> (ix, 10-ix-1))

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

-- A super-simple verification example.
count :: Software ()
count = do
  printf "Enter a number: "
  n :: SExp Word32 <- fget stdin

  let total = iter n 0 (\i n -> hint (n == i) $ i + 1)
  total <- initRef total >>= unsafeFreezeRef
  assert (total == n) "Count is wrong"
  printf "The count is %d\n" total

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

rev :: Software ()
rev = do
  n <- fget stdin
  loc :: IArr (SExp Word32) <- newArr n >>= unsafeFreezeArr
  cpy :: Arr  (SExp Word32) <- newArr n
  assert (n > 0) "neg"
  for 0 1 (n-1) $ \i -> do
    setArr cpy i (loc ! (n-i-1))

rev_inplace :: Software ()
rev_inplace = do
    n <- fget stdin
    loc :: Arr (SExp Word32) <- newArr n
    vec <- unsafeFreezeArr loc >>= unsafeThawArr
    for 0 1 ((n `div` 2 :: SExp Word32)) $ \i -> do
      x <- getArr vec i
      y <- getArr vec (n-i-1)
      setArr loc i y
      setArr loc (n-i-1) x

rev_amazing :: Software ()
rev_amazing = do
    n <- fget stdin
    iarr :: IArr (SExp Word32) <- newArr n >>= unsafeFreezeArr
    arr <- unsafeThawArr iarr
    result <- manifestArr arr (pairwise @Software (\i -> (i, n-i-1)) (reverse iarr))
    return ()

------------------------------------------------------------
{-
test_scProd1 = do
    n <- fget stdin
    printf "result: %.3f\n" $
      (scProd (fmap i2n (0 ... n-1)) (fmap i2n (2 ... n+1)) :: Data Double)

test_scProd2 = do
    n  <- fget stdin
    v1 <- manifestFresh $ fmap i2n (0 ... n-1)
    v2 <- manifestFresh $ fmap i2n (2 ... n+1)
    printf "result: %.3f\n" (scProd v1 v2 :: Data Double)

map_inplace :: Run ()
map_inplace = do
    n   <- fget stdin
    loc <- newArr n
    vec <- manifest loc (0 ... n-1)
    manifestStore loc $ map (+1) vec
    vec <- unsafeFreezeArr loc
    printf "result: %d\n" $ sum vec

map2_inplace :: Run ()
map2_inplace = do
    n   <- fget stdin
    assert (n < maxBound) "oops"
    loc :: Arr (Data Word32) <- newArr (n+1)
    vec <- unsafeFreezeArr loc
    for (0, 1, Excl (n :: Data Word32)) $ \i -> do
      setArr loc i (arrIx vec i+1)

tail_inplace :: Run ()
tail_inplace = do
    n <- fget stdin
    loc :: Arr (Data Word32) <- newArr n
    vec <- unsafeFreezeArr loc
    let when cond x = iff cond x (return ())
    when (n > 0) $
      for (0, 1, Excl (n-1)) $ \i -> do
        setArr loc i (arrIx vec (i+1)+1)

filter_inplace :: Run ()
filter_inplace = do
    n <- fget stdin
    loc :: Arr (Data Word32) <- newArr n
    vec <- unsafeFreezeArr loc
    ref <- initRef 0
    let when cond x = iff cond x (return ())
    for (0, 1, Excl n) $ \i -> do
      let x = arrIx vec i
      when (x > 5) $ do
        j <- unsafeFreezeRef ref
        hint (j <= i)
        setArr loc j x
        setRef ref (j+1)

rev_inplace :: Software ()
rev_inplace = do
    n <- fget stdin
    loc :: Arr (Data Word32) <- newArr n
    vec <- unsafeFreezeArr loc >>= unsafeThawArr
    for (0, 1, Excl (n `div` 2 :: Data Word32)) $ \i -> do
      x <- getArr vec i
      y <- getArr vec (n-i-1)
      setArr loc i y
      setArr loc (n-i-1) x
-}
--------------------------------------------------------------------------------


-- Mandelbrot.hs
-- Caluclate Mandelbrot set in Haskell
--
-- vim: ft=haskell sw=2 ts=2 et
--

module Mandelbrot
  ( Pixel
  , mandelbrot ) where

import Data.Complex

-- Basic type defs
type Pixel = (Int, Int, Int)

-- Calculate pixel value
pixelValue :: (Integral a) => Double -> Double -> a -> a
pixelValue x y nMax =
  let z0 = x :+ y
      z = 0.0 :+ 0.0 in
    loop z z0 nMax
  where
    loop :: (Integral a) => Complex Double -> Complex Double -> a -> a
    loop z z0 nMax =
      nMax

valueToColor :: (Integral a) => a -> Pixel
valueToColor val =
  (0, 0, 0)

mandelbrot :: (Integral a, Floating b) => a -> a -> b -> b -> b -> [Pixel]
mandelbrot xMax yMax xCenter yCenter pixelSize =
  [(0, 0, 0)]


{-# LANGUAGE DuplicateRecordFields #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE ScopedTypeVariables #-}

-- |
-- Module      : PSO.Heisenberg
-- Description : Heisenberg Hamiltonian
-- Copyright   : (c) Tom Westerhout, 2017
-- License     : BSD3
-- Maintainer  : t.westerhout@student.ru.nl
-- Stability   : experimental
module PSO.Heisenberg
  ( -- * Introduction
    --
    -- | Let's start with a 1D chain of \(N\) spin-1/2 particles. Each of them
    -- is described by a state \(|\alpha_i\rangle\in\mathbb{C}^2\). Hilbert
    -- space of the whole system is then \(\bigotimes_{i=0}^{N-1}\mathbb{C}^2\).
    --
    -- We use the following Hamiltonian to describe this system:
    -- \[
    -- \hat{\mathcal{H}} = \sum_{i=0}^{N-2}
    --    \hat{\boldsymbol\sigma}_i \cdot \hat{\boldsymbol\sigma}_{i+1}
    --    = \sum_{i=0}^{N-2}
    --          \hat{\sigma}_i^{(x)}\hat{\sigma}_{i+1}^{(x)}
    --        + \hat{\sigma}_i^{(y)}\hat{\sigma}_{i+1}^{(y)}
    --        + \hat{\sigma}_i^{(z)}\hat{\sigma}_{i+1}^{(z)} \;.
    -- \]
    -- Only the nearest neighbours are connected. The boundary conditions used
    -- are called /open/ because the zeroth and N-1st spins don't interact.
    one2D
  , pauliX
  , pauliY
  , pauliZ
  , heisenberg1DOpen
    -- * Quantities of interest
    --
    -- | The quantity we're actually interested in is
    -- \[
    -- \frac{\langle\sigma| \hat{\mathcal{H}} |\psi\rangle}{%
    --   \langle\sigma|\psi\rangle} \;,
    -- \]
    -- where \(|\psi\rangle\) is described by the RBM \(\mathcal{W}\) and
    -- \(|\sigma\rangle\in\{\downarrow,\uparrow\}^N\) is a spin configuration.
    -- We call this quantity /local energy/ \(E_{loc}(\sigma;\mathcal{W})\).
    -- \[
    -- E_{loc}(\sigma;\mathcal{W}) = \sum_i \frac{\langle\sigma|
    --    \hat{\boldsymbol\sigma}_i\cdot\hat{\boldsymbol\sigma}_{i+1}
    --    |\psi\rangle}{\langle\sigma|\psi\rangle}
    -- \]
    -- Now
    -- \[
    -- \begin{aligned}
    -- &\hat{\boldsymbol\sigma}_1\cdot\hat{\boldsymbol\sigma}_2
    --    |\uparrow\uparrow\rangle = |\uparrow\uparrow\rangle \;,\\
    -- &\hat{\boldsymbol\sigma}_1\cdot\hat{\boldsymbol\sigma}_2
    --    |\downarrow\downarrow\rangle = |\downarrow\downarrow\rangle \;,\\
    -- &\hat{\boldsymbol\sigma}_1\cdot\hat{\boldsymbol\sigma}_2
    --    |\uparrow\downarrow\rangle = |\uparrow\downarrow\rangle
    --      - 2|\downarrow\uparrow\rangle \;,\\
    -- &\hat{\boldsymbol\sigma}_1\cdot\hat{\boldsymbol\sigma}_2
    --    |\downarrow\uparrow\rangle = |\downarrow\uparrow\rangle
    --      - 2|\uparrow\downarrow\rangle \;.
    -- \end{aligned}
    -- \]
  -- , locEnergyHH1DOpen
  ) where

import Data.Complex(Complex(..))
import Numeric.LinearAlgebra(C, Matrix, Herm, (><), trustSym, kronecker)
import qualified Numeric.LinearAlgebra as LA
import qualified Data.Vector.Storable as V

-- import PSO.Internal.Spin

one2D :: Matrix C
one2D = (2><2)
    [ 1.0,        0.0
    , 0.0,        1.0 ]

pauliX :: Matrix C
pauliX = (2><2)
    [ 0.0,        1.0
    , 1.0,        0.0 ]

pauliY :: Matrix C
pauliY = (2><2)
    [ 0.0,        0.0 :+ (-1.0)
    , 0.0 :+ 1.0, 0.0 ]

pauliZ :: Matrix C
pauliZ = (2><2)
    [ 1.0,        0.0
    , 0.0,        (-1.0) ]

rotate :: Int -> [a] -> [a]
rotate n xs = zipWith const (drop n (cycle xs)) xs

heisenberg1DOpen :: Int -> Herm C
heisenberg1DOpen n
  | n > 2 = trustSym $ heisenberg1DOpenImpl n pauliX
          + heisenberg1DOpenImpl n pauliY
          + heisenberg1DOpenImpl n pauliZ
  | n == 2 = trustSym $ pauliX `kronecker` pauliX
                      + pauliY `kronecker` pauliY
                      + pauliZ `kronecker` pauliZ
  | otherwise = undefined

heisenberg1DOpenImpl :: Int -> Matrix C -> Matrix C
heisenberg1DOpenImpl n pauli =
  let xs = (replicate (n-2) one2D) ++ (replicate 2 pauli)
      rotations = take (n-1) . iterate (rotate 1) $ xs
   in sum . map (foldr1 kronecker) $ rotations

-- locEnergyHH1DOpen ::
--      (Floating α, Eq α, Num (V.Vector α), LA.Numeric α)
--   => RBM α -> V.Vector α -> α
-- locEnergyHH1DOpen machine σ =
--   let θ = mkTheta machine σ
--       zipper i x y
--         | x == y    = 1
--         | otherwise = -1 + 2 * exp (logPoP2 machine θ σ (i, i + 1))
--    in V.sum $ V.izipWith zipper σ (V.tail σ)


{-# LANGUAGE CPP                   #-}
{-# LANGUAGE UndecidableInstances  #-}
{-# LANGUAGE PolyKinds             #-}
{-# LANGUAGE DataKinds             #-}
{-# LANGUAGE GADTs                 #-}
{-# LANGUAGE RankNTypes            #-}
{-# LANGUAGE TypeOperators         #-}
{-# LANGUAGE TypeFamilies          #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE FlexibleContexts      #-}
{-# LANGUAGE ViewPatterns          #-}
{-# LANGUAGE ScopedTypeVariables   #-}
{-# OPTIONS_GHC -Wno-incomplete-uni-patterns #-}

module Grenade.Recurrent.Layers.LSTM (
    LSTM (..)
  , LSTMWeights (..)
  , randomLSTM
  ) where

import           Control.Monad.Primitive           (PrimBase, PrimState)
import           System.Random.MWC hiding (create)
import GHC.TypeLits

-- import           Data.List ( foldl1' )
import           Data.Proxy
import           Data.Serialize

#if MIN_VERSION_base(4,9,0)
import           Data.Kind (Type)
#endif

import qualified Numeric.LinearAlgebra as LA
import           Numeric.LinearAlgebra.Static

import           Grenade.Core
import           Grenade.Utils.LinearAlgebra
import           Grenade.Recurrent.Core
import           Grenade.Layers.Internal.Update

-- | Long Short Term Memory Recurrent unit
--
--   This is a Peephole formulation, so the recurrent shape is
--   just the cell state, the previous output is not held or used
--   at all.
data LSTM :: Nat -> Nat -> Type where
  LSTM :: ( KnownNat input
          , KnownNat output
          ) => !(LSTMWeights input output) -- Weights
            -> !(LSTMWeights input output) -- Momentums
            -> LSTM input output

data LSTMWeights :: Nat -> Nat -> Type where
  LSTMWeights :: ( KnownNat input
                 , KnownNat output
                 ) => {
                   lstmWf :: !(L output input)  -- Weight Forget     (W_f)
                 , lstmUf :: !(L output output) -- Cell State Forget (U_f)
                 , lstmBf :: !(R output)        -- Bias Forget       (b_f)
                 , lstmWi :: !(L output input)  -- Weight Input      (W_i)
                 , lstmUi :: !(L output output) -- Cell State Input  (U_i)
                 , lstmBi :: !(R output)        -- Bias Input        (b_i)
                 , lstmWo :: !(L output input)  -- Weight Output     (W_o)
                 , lstmUo :: !(L output output) -- Cell State Output (U_o)
                 , lstmBo :: !(R output)        -- Bias Output       (b_o)
                 , lstmWc :: !(L output input)  -- Weight Cell       (W_c)
                 , lstmBc :: !(R output)        -- Bias Cell         (b_c)
                 } -> LSTMWeights input output

instance Show (LSTM i o) where
  show LSTM {} = "LSTM"

instance FoldableGradient (LSTMWeights input output) where
  mapGradient f (LSTMWeights wf uf bf wi ui bi wo uo bo wc bc) =
    LSTMWeights (dmmap f wf) (dmmap f uf) (dvmap f bf) (dmmap f wi) (dmmap f ui) (dvmap f bi) (dmmap f wo) (dmmap f uo) (dvmap f bo) (dmmap f wc) (dvmap f bc)
  squaredSums (LSTMWeights wf uf bf wi ui bi wo uo bo wc bc) =
    [ sumM . squareM $ wf
    , sumM . squareM $ uf
    , sumV . squareV $ bf
    , sumM . squareM $ wi
    , sumM . squareM $ ui
    , sumV . squareV $ bi
    , sumM . squareM $ wo
    , sumM . squareM $ uo
    , sumV . squareV $ bo
    , sumM . squareM $ wc
    , sumV . squareV $ bc
    ]

instance (KnownNat i, KnownNat o) => UpdateLayer (LSTM i o) where
  -- The gradients are the same shape as the weights and momentum
  -- This seems to be a general pattern, maybe it should be enforced.
  type Gradient (LSTM i o) = (LSTMWeights i o)

  -- Run the update function for each group matrix/vector of weights, momentums and gradients.
  -- Hmm, maybe the function should be used instead of passing in the learning parameters.
  runUpdate opt@OptSGD{} (LSTM w m) g =
    let MatrixResultSGD wf wf' = u lstmWf w m g
        MatrixResultSGD uf uf' = u lstmUf w m g
        VectorResultSGD bf bf' = v lstmBf w m g
        MatrixResultSGD wi wi' = u lstmWi w m g
        MatrixResultSGD ui ui' = u lstmUi w m g
        VectorResultSGD bi bi' = v lstmBi w m g
        MatrixResultSGD wo wo' = u lstmWo w m g
        MatrixResultSGD uo uo' = u lstmUo w m g
        VectorResultSGD bo bo' = v lstmBo w m g
        MatrixResultSGD wc wc' = u lstmWc w m g
        VectorResultSGD bc bc' = v lstmBc w m g
    in LSTM (LSTMWeights wf uf bf wi ui bi wo uo bo wc bc) (LSTMWeights wf' uf' bf' wi' ui' bi' wo' uo' bo' wc' bc')
      where
    -- Utility function for updating with the momentum, gradients, and weights.
    u :: forall x ix out. (KnownNat ix, KnownNat out) => (x -> L out ix) -> x -> x -> x -> MatrixResult out ix
    u e (e -> weights) (e -> momentum) (e -> gradient) =
      descendMatrix opt (MatrixValuesSGD weights gradient momentum)

    v :: forall x ix. (KnownNat ix) => (x -> R ix) -> x -> x -> x -> VectorResult ix
    v e (e -> weights) (e -> momentum) (e -> gradient) =
      descendVector opt (VectorValuesSGD weights gradient momentum)
  runUpdate _ d g = runUpdate defOptimizer d g

  -- There's a lot of updates here, so to try and minimise the number of data copies
  -- we'll create a mutable bucket for each.
  -- runUpdates rate lstm gs =
  --   let combinedGradient = foldl1' uu gs
  --   in  runUpdate rate lstm combinedGradient
  --     where
  --   uu :: (KnownNat i, KnownNat o) => LSTMWeights i o -> LSTMWeights i o -> LSTMWeights i o
  --   uu a b =
  --     let wf = u lstmWf a b
  --         uf = u lstmUf a b
  --         bf = v lstmBf a b
  --         wi = u lstmWi a b
  --         ui = u lstmUi a b
  --         bi = v lstmBi a b
  --         wo = u lstmWo a b
  --         uo = u lstmUo a b
  --         bo = v lstmBo a b
  --         wc = u lstmWc a b
  --         bc = v lstmBc a b
  --     in LSTMWeights wf uf bf wi ui bi wo uo bo wc bc
  --   u :: forall x ix out. (KnownNat ix, KnownNat out) => (x -> (L out ix)) -> x -> x -> L out ix
  --   u e (e -> a) (e -> b) = tr $ tr a + tr b

  --   v :: forall x ix. (x -> (R ix)) -> x -> x -> R ix
  --   v e (e -> a) (e -> b) = a + b

instance (KnownNat i, KnownNat o, KnownNat (i*o), KnownNat (o*o)) => RandomLayer (LSTM i o) where
  createRandomWith = randomLSTM

instance (KnownNat i, KnownNat o) => RecurrentUpdateLayer (LSTM i o) where
  -- The recurrent shape is the same size as the output.
  -- It's actually the cell state however, as this is a peephole variety LSTM.
  type RecurrentShape (LSTM i o) = S ('D1 o)

instance (KnownNat i, KnownNat o) => RecurrentLayer (LSTM i o) ('D1 i) ('D1 o) where

  -- The tape stores essentially every variable we calculate,
  -- so we don't have to run any forwards component again.
  type RecTape (LSTM i o) ('D1 i) ('D1 o) = (R o, R i, R o, R o, R o, R o, R o, R o, R o, R o, R o)
  -- Forward propagation for the LSTM layer.
  -- The size of the cell state is also the size of the output.
  runRecurrentForwards (LSTM lw _) (S1D cell) (S1D input) =
    let -- Forget state vector
        f_s = lstmBf lw + lstmWf lw #> input + lstmUf lw #> cell
        f_t = sigmoid f_s
        -- Input state vector
        i_s = lstmBi lw + lstmWi lw #> input + lstmUi lw #> cell
        i_t = sigmoid i_s
        -- Output state vector
        o_s = lstmBo lw + lstmWo lw #> input + lstmUo lw #> cell
        o_t = sigmoid o_s
        -- Cell input state vector
        c_s = lstmBc lw + lstmWc lw #> input
        c_x = tanh c_s
        -- Cell state
        c_t = f_t * cell + i_t * c_x
        -- Output (it's sometimes recommended to use tanh c_t)
        h_t = o_t * c_t
    in ((cell, input, f_s, f_t, i_s, i_t, o_s, o_t, c_s, c_x, c_t), S1D c_t, S1D h_t)

  -- Run a backpropogation step for an LSTM layer.
  -- We're doing all the derivatives by hand here, so one should
  -- be extra careful when changing this.
  --
  -- There's a test version using the AD library without hmatrix in the test
  -- suite. These should match always.
  runRecurrentBackwards (LSTM lw _) (cell, input, f_s, f_t, i_s, i_t, o_s, o_t, c_s, c_x, c_t) (S1D cellGrad) (S1D h_t') =
    let -- Reverse Mode AD Derivitives
        c_t' = h_t' * o_t + cellGrad

        f_t' = c_t' * cell
        f_s' = sigmoid' f_s * f_t'

        o_t' = h_t' * c_t
        o_s' = sigmoid' o_s * o_t'

        i_t' = c_t' * c_x
        i_s' = sigmoid' i_s * i_t'

        c_x' = c_t' * i_t
        c_s' = tanh' c_s * c_x'

        -- The derivatives to pass sideways (recurrent) and downwards
        cell'  = tr (lstmUf lw) #> f_s' + tr (lstmUo lw) #> o_s' + tr (lstmUi lw) #> i_s' + c_t' * f_t
        input' = tr (lstmWf lw) #> f_s' + tr (lstmWo lw) #> o_s' + tr (lstmWi lw) #> i_s' + tr (lstmWc lw) #> c_s'

        -- Calculate the gradient Matricies for the input
        lstmWf' = f_s' `outer` input
        lstmWi' = i_s' `outer` input
        lstmWo' = o_s' `outer` input
        lstmWc' = c_s' `outer` input

        -- Calculate the gradient Matricies for the cell
        lstmUf' = f_s' `outer` cell
        lstmUi' = i_s' `outer` cell
        lstmUo' = o_s' `outer` cell

        -- The biases just get the values, but we'll write it so it's obvious
        lstmBf' = f_s'
        lstmBi' = i_s'
        lstmBo' = o_s'
        lstmBc' = c_s'

        gradients = LSTMWeights lstmWf' lstmUf' lstmBf' lstmWi' lstmUi' lstmBi' lstmWo' lstmUo' lstmBo' lstmWc' lstmBc'
    in  (gradients, S1D cell', S1D input')

-- | Generate an LSTM layer with random Weights
--   one can also just call createRandom from UpdateLayer
--
--   Has forget gate biases set to 1 to encourage early learning.
--
--   https://github.com/karpathy/char-rnn/commit/0dfeaa454e687dd0278f036552ea1e48a0a408c9
--
randomLSTM :: forall m i o. (PrimBase m, KnownNat i, KnownNat o, KnownNat (i*o),KnownNat (o*o))
           => NetworkInitSettings -> Gen (PrimState m) -> m (LSTM i o)
randomLSTM (NetworkInitSettings m HMatrix _) gen = do
  let w = getRandomMatrix i o m gen
  let u = getRandomMatrix i o m gen
  let v = getRandomVector i o m gen

  let w0 = konst 0
      u0 = konst 0
      v0 = konst 0

  LSTM <$> (LSTMWeights <$> w <*> u <*> pure (konst 1) <*> w <*> u <*> v <*> w <*> u <*> v <*> w <*> v)
         <*> pure (LSTMWeights w0 u0 v0 w0 u0 v0 w0 u0 v0 w0 v0)

  where i = natVal (Proxy :: Proxy i)
        o = natVal (Proxy :: Proxy o)
randomLSTM (NetworkInitSettings _ cpu _) _ = error $ "CPU backend " ++ show cpu ++ " not supported by LSTM layer"

-- | Maths
--
-- TODO: Move to not here
--       Optimise backwards derivative
sigmoid :: Floating a => a -> a
sigmoid x = 1 / (1 + exp (-x))

sigmoid' :: Floating a => a -> a
sigmoid' x = logix * (1 - logix)
  where
    logix = sigmoid x

tanh' :: (Floating a) => a -> a
tanh' t = 1 - s ^ (2 :: Int)  where s = tanh t

instance (KnownNat i, KnownNat o) => Serialize (LSTM i o) where
  put (LSTM lw _) = do
      u (lstmWf lw)
      u (lstmUf lw)
      v (lstmBf lw)
      u (lstmWi lw)
      u (lstmUi lw)
      v (lstmBi lw)
      u (lstmWo lw)
      u (lstmUo lw)
      v (lstmBo lw)
      u (lstmWc lw)
      v (lstmBc lw)
    where
      u :: forall a b. (KnownNat a, KnownNat b) => Putter  (L b a)
      u = putListOf put . LA.toList . LA.flatten . extract
      v :: forall a. (KnownNat a) => Putter (R a)
      v = putListOf put . LA.toList . extract

  get = do
      w <- LSTMWeights <$> u <*> u <*> v <*> u <*> u <*> v <*> u <*> u <*> v <*> u <*> v
      return $ LSTM w (LSTMWeights w0 u0 v0 w0 u0 v0 w0 u0 v0 w0 v0)
    where
      u :: forall a b. (KnownNat a, KnownNat b) => Get  (L b a)
      u = let f = fromIntegral $ natVal (Proxy :: Proxy a)
          in  maybe (fail "Vector of incorrect size") return . create . LA.reshape f . LA.fromList =<< getListOf get
      v :: forall a. (KnownNat a) => Get (R a)
      v = maybe (fail "Vector of incorrect size") return . create . LA.fromList =<< getListOf get

      w0 = konst 0
      u0 = konst 0
      v0 = konst 0

-------------------- GNum instances --------------------

instance (KnownNat i, KnownNat o) => GNum (LSTM i o) where
  n |* (LSTM w m) = LSTM (n |* w) (n |* m)
  (LSTM w1 m1) |+ (LSTM w2 m2) = LSTM (w1 |+ w2) (m1 |+ m2)

instance (KnownNat i, KnownNat o) => GNum (LSTMWeights i o) where
  n |* (LSTMWeights wf uf bf wi ui bi wo uo bo wc bc) =
    LSTMWeights
      (fromRational n * wf)
      (fromRational n * uf)
      (fromRational n * bf)
      (fromRational n * wi)
      (fromRational n * ui)
      (fromRational n * bi)
      (fromRational n * wo)
      (fromRational n * uo)
      (fromRational n * bo)
      (fromRational n * wc)
      (fromRational n * bc)
  (LSTMWeights wf1 uf1 bf1 wi1 ui1 bi1 wo1 uo1 bo1 wc1 bc1) |+ (LSTMWeights wf2 uf2 bf2 wi2 ui2 bi2 wo2 uo2 bo2 wc2 bc2) =
    LSTMWeights (wf1 + wf2) (uf1 + uf2) (bf1 + bf2) (wi1 + wi2) (ui1 + ui2) (bi1 + bi2) (wo1 + wo2) (uo1 + uo2) (bo1 + bo2) (wc1 + wc2) (bc1 + bc2)


-- File created: 2009-08-06 20:47:49

module Haschoo.Evaluator.Standard.Vectors (procedures) where

import Control.Applicative ((<$>))
import Data.Array          (Array)
import Data.Array.IArray   ((!), bounds, elems)
import Data.Array.IO       (IOArray)
import Data.Array.MArray   ( readArray, writeArray
                           , newArray, newListArray
                           , getBounds, getElems)
import Data.Complex        (Complex((:+)))

import Haschoo.Types           (ScmValue(..), listToPair, pairToList)
import Haschoo.Utils           (ErrOr)
import Haschoo.Evaluator.Utils ( tooFewArgs, tooManyArgs, immutable
                               , notInt, notList)

procedures :: [(String, ScmValue)]
procedures = map (\(a,b) -> (a, ScmFunc a b))
   [ ("vector?", return . fmap ScmBool . scmIsVector)

   , ("make-vector", scmMakeVector)
   , ("vector",      scmVector)

   , ("vector-length", scmLength)

   , ("vector-ref",  scmRef)
   , ("vector-set!", scmSet)

   , ("vector->list", scmToList)
   , ("list->vector", scmFromList)

   , ("vector-fill!", scmFill)
   ]

scmIsVector :: [ScmValue] -> ErrOr Bool
scmIsVector [ScmMVector _] = Right True
scmIsVector [ScmVector  _] = Right True
scmIsVector [_]            = Right False
scmIsVector []             = tooFewArgs  "vector?"
scmIsVector _              = tooManyArgs "vector?"

scmMakeVector, scmVector :: [ScmValue] -> IO (ErrOr ScmValue)
scmMakeVector (ScmInt l : xs) =
   case xs of
        []  -> f (flip newArray (ScmComplex $ 0 :+ 42))
        [x] -> f (flip newArray x)
        _   -> return$ tooManyArgs "make-vector"
 where
   f mk = tryToLen "make-vector" l $ \i -> ScmMVector <$> mk (0, i-1)

scmMakeVector (_:_:_:_) = return$ tooManyArgs "make-vector"
scmMakeVector []        = return$ tooFewArgs  "make-vector"
scmMakeVector _         = return$ notInt      "make-vector"

scmVector xs = Right <$> toScmMVector xs

scmLength :: [ScmValue] -> IO (ErrOr ScmValue)
scmLength [x] = case x of
                     ScmVector  s -> Right . ScmInt . toInteger <$>  vLen s
                     ScmMVector s -> Right . ScmInt . toInteger <$> mvLen s
                     _            -> return$ notVector "vector-length"

scmLength [] = return$ tooFewArgs  "vector-length"
scmLength _  = return$ tooManyArgs "vector-length"

scmRef, scmSet :: [ScmValue] -> IO (ErrOr ScmValue)
scmRef [x, ScmInt i] = case x of
                            ScmVector  s -> f (vLen  s) (return . (s!))
                            ScmMVector s -> f (mvLen s) (readArray s)
                            _            -> return$ notVector "vector-ref"
 where
   f len = tryToIdx "vector-ref" len i

scmRef [_,_]   = return$ notInt      "vector-ref"
scmRef (_:_:_) = return$ tooManyArgs "vector-ref"
scmRef _       = return$ tooFewArgs  "vector-ref"

scmSet [v, ScmInt idx, x] =
   case v of
        ScmMVector v' -> do
           tryToIdx "vector-set!" (mvLen v') idx (\i -> writeArray v' i x)
           return (Right ScmVoid)

        ScmVector  _ -> return$ immutable "vector-set!"
        _            -> return$ notVector "vector-set!"

scmSet [_, _, ScmChar _] = return$ notInt      "vector-set!"
scmSet (_:_:_:_)         = return$ tooManyArgs "vector-set!"
scmSet _                 = return$ tooFewArgs  "vector-set!"

scmToList, scmFromList :: [ScmValue] -> IO (ErrOr ScmValue)
scmToList [x] = case x of
                     ScmVector  v -> f (elems v)
                     ScmMVector v -> f =<< getElems v
                     _            -> return$ notVector "vector->list"
 where
   f = fmap (Right . fst) . listToPair

scmToList [] = return$ tooFewArgs  "vector->list"
scmToList _  = return$ tooManyArgs "vector->list"

scmFromList [x] =
   case x of
        ScmList l       -> Right <$> toScmMVector l
        p@(ScmPair _ _) -> do
           l <- pairToList p
           case l of
                Left _             -> return$ notList "list->vector"
                Right (ScmList l') -> Right <$> toScmMVector l'

                _ -> error "list->vector :: the impossible happened"

        _ -> return$ notList "list->vector"

scmFromList [] = return$ tooFewArgs  "list->vector"
scmFromList _  = return$ tooManyArgs "list->vector"

scmFill :: [ScmValue] -> IO (ErrOr ScmValue)
scmFill [v, x] =
   case v of
        ScmMVector v' -> do
           (a,b) <- getBounds v'
           mapM_ (\i -> writeArray v' i x) [a..b]
           return (Right ScmVoid)

        ScmVector _ -> return$ immutable "vector-fill!"
        _           -> return$ notVector "vector-fill!"

scmFill (_:_:_) = return$ tooManyArgs "vector-fill!"
scmFill _       = return$ tooFewArgs  "vector-fill!"

------

notVector :: String -> ErrOr a
notVector = fail . ("Nonvector argument to " ++)

toScmMVector :: [ScmValue] -> IO ScmValue
toScmMVector v = ScmMVector <$> newListArray (0, length v - 1) v

tryToLen :: String -> Integer -> (Int -> IO a) -> IO (ErrOr a)
tryToLen = tryRange 0 (toInteger (maxBound :: Int))

tryToIdx :: String -> IO Int -> Integer -> (Int -> IO a) -> IO (ErrOr a)
tryToIdx s mlen i f = mlen >>= \len -> tryRange 0 (toInteger len - 1) s i f

tryRange :: Integer -> Integer -> String -> Integer -> (Int -> IO a)
         -> IO (ErrOr a)
tryRange a b s i f = if i >= a && i <= b
                        then Right <$> f (fromInteger i)
                        else return.Left $ "Out-of-range integer to " ++ s

vLen :: Array Int ScmValue -> IO Int
vLen = return . (+1) . snd . bounds

mvLen :: IOArray Int ScmValue -> IO Int
mvLen = fmap ((+1) . snd) . getBounds


-- H24.12.02 22:30 Close it.
module Main where
import System.Environment
import Text.ParserCombinators.Parsec hiding (spaces)
import Control.Monad
import Numeric
import Data.Ratio
import Data.Complex
import Data.Array
import Control.Monad.Error
import System.IO
import Data.IORef

main :: IO ()
main = do args <- getArgs
          case length args of
            0 -> runRepl
            1 -> runOnce $ args !! 0
            otherwise -> putStrLn "Program takes only 0 or 1 argument"

flushStr :: String -> IO ()
flushStr str = putStr str >> hFlush stdout

readPrompt :: String -> IO String
readPrompt prompt = flushStr prompt >> getLine

evalAndPrint :: Env -> String -> IO ()
evalAndPrint env expr = evalString env expr >>= putStrLn

evalString :: Env -> String -> IO String
evalString env expr = runIOThrows $ liftM show $
    (liftThrows $ readExpr expr) >>= eval env

until_ :: Monad m => (a -> Bool) -> m a -> (a -> m ()) -> m ()
until_ pred prompt action = do
  result <- prompt
  if pred result
     then return ()
     else action result >> until_ pred prompt action

runOnce :: String -> IO ()
runOnce expr = primitiveBindings >>= flip evalAndPrint expr

runRepl :: IO ()
runRepl = primitiveBindings >>=
    until_ (== "quit") (readPrompt "Lisp>>> ") . evalAndPrint

data LispVal = Atom String
             | List [LispVal]
             | DottedList [LispVal] LispVal
             | Number Integer
             | String String
             | Bool Bool
             | Character Char
             | Float Double
             | Ratio Rational
             | Complex (Complex Double)
             | Vector (Array Int LispVal)
             | PrimitiveFunc ([LispVal] -> ThrowsError LispVal)
             | Func {params :: [String], vararg :: (Maybe String),
                      body :: [LispVal], closure :: Env}

spaces :: Parser ()
spaces = skipMany1 space

symbol :: Parser Char
symbol = oneOf "!#$%&|*+-/:<=>?@^_~"
symbol1 :: Parser Char
symbol1 = oneOf "!$%&|*+-/:<=>?@^_~"

parseAtom :: Parser LispVal
parseAtom = do first <- letter <|> symbol1
               rest <- many (letter <|> digit <|> symbol)
               return . Atom $ first : rest

parseBool :: Parser LispVal
parseBool = liftM Bool ((char 't' >> return True)
                         <|> (char 'f' >> return False))

parseSigned :: (Real r) => ReadS r -> Parser [Char] -> Parser r
parseSigned reader lexer = liftM (fst . head . readSigned reader) (do
  sign <- char '-' <|> return '0'
  num <- lexer
  return (sign : num))

lexFloat :: Parser [Char]
lexFloat = do x <- many1 digit
              char '.'
              y <- many digit
              return $ x ++ "." ++ y

parseFloat :: Parser LispVal
parseFloat = (liftM Float $ parseSigned readFloat lexFloat)

parseDigits :: Parser LispVal
parseDigits = (liftM Number $ parseSigned readDec (many1 digit))

parseRatio :: Parser LispVal
parseRatio = do
  x <- parseSigned readDec (many1 digit)
  char '/'
  y <- parseSigned readDec (many1 digit)
  return $ Ratio (x % y)

toDouble :: LispVal -> Double
toDouble (Number n) = fromIntegral n
toDouble (Float f) = f

parseComplex :: Parser LispVal
parseComplex = do x <- (try parseFloat <|> parseDigits)
                  char '+'
                  y <- (try parseFloat <|> parseDigits)
                  char 'i'
                  return $ Complex (toDouble x :+ toDouble y)

parseNum :: Parser LispVal
parseNum = try parseComplex
           <|> try parseRatio
           <|> try parseFloat
           <|> parseDigits

parseNumber :: Parser LispVal
parseNumber = liftM Number
                (parseBin <|> parseOct <|> parseDig <|> parseHex)
  where
        parseBin = char 'b' >> parseSigned readBin (many1 $ oneOf "01")
        parseOct = char 'o' >> parseSigned readOct (many1 octDigit)
        parseDig = char 'd' >> parseSigned readDec (many1 digit)
        parseHex = char 'x' >> parseSigned readHex (many1 hexDigit)
        readBin  = readInt 2 (\x->x=='0'||x=='1') (\x->if x=='0' then 0 else 1)

parseSharpSyntax :: Parser LispVal
parseSharpSyntax = char '#' >> (parseBool
                     <|> parseNumber
                     <|> parseCharacter
                     <|> parseVector)

escapedChars :: Parser Char
escapedChars = char '\\' >> oneOf "\\\"nrt" >>= \x ->
               return $ case x of
                 'n' -> '\n'
                 'r' -> '\r'
                 't' -> '\t'
                 _ -> x

parseString :: Parser LispVal
parseString = do char '"'
                 x <- many (escapedChars <|> noneOf "\"\\")
                 char '"'
                 return $ String x

parseCharacter :: Parser LispVal
parseCharacter = liftM Character (char '\\' >>
                   option ' ' (try (string "space" >> return ' ')
                                <|> try (string "newline" >> return '\n')
                                <|> do { x <- anyChar;
                                         notFollowedBy alphaNum;
                                         return x}))

parseList :: Parser LispVal
parseList = do
  body <- sepEndBy parseExpr spaces
  if length body == 0 then return $ List []
  else parseTail >>= \tail-> return $
    case tail of
      Just (List l)         -> List (body ++ l)
      Just (DottedList h t) -> DottedList (body ++ h) t
      Just val              -> DottedList body val
      Nothing               -> List body

parseTail :: Parser (Maybe LispVal)
parseTail = liftM Just (char '.' >> skipMany space >> parseExpr)
            <|> return Nothing

makeSyntaxSugarParser :: Parser a -> String -> Parser LispVal
makeSyntaxSugarParser ch symbol =
  ch >> parseExpr >>= (\x -> return $ List [Atom symbol, x])

parseQuoted :: Parser LispVal
parseQuoted = makeSyntaxSugarParser (char '\'') "quote"

parseQusaiquote :: Parser LispVal
parseQusaiquote = makeSyntaxSugarParser (char '`') "qusaiquote"

parseUnquote :: Parser LispVal
parseUnquote = makeSyntaxSugarParser (char ',') "unquote"

parseUnquoteS :: Parser LispVal
parseUnquoteS = makeSyntaxSugarParser (string ",@") "unquote-splicing"

parens :: Parser a -> Parser a
parens = between (char '(' >> skipMany space) (char ')')

parseVector :: Parser LispVal
parseVector = do
  arrayValues <- parens (sepBy parseExpr spaces)
  return $ Vector (listArray (0,(length arrayValues - 1)) arrayValues)

parseExpr :: Parser LispVal
parseExpr = parseSharpSyntax
        <|> try parseNum
        <|> parseAtom
        <|> parseString
        <|> parseQuoted <|> parseQusaiquote
        <|> try parseUnquoteS <|> parseUnquote
        <|> parens parseList

readExpr :: String -> ThrowsError LispVal
readExpr input = case parse parseExpr "lisp" input of
                   Left err -> throwError $ Parser err
                   Right val -> return val

showVal :: LispVal -> String
showVal (String contents) = "\"" ++ contents ++ "\""
showVal (Atom name) = name
showVal (Number contents) = show contents
showVal (Float contents) = showFFloat Nothing contents ""
showVal (Ratio contents) = show (numerator contents) ++ "/"
                         ++ show (denominator contents)
showVal (Complex x) = showFFloat Nothing (realPart x) $ "+"
                    ++ showFFloat Nothing (imagPart x) "i"
showVal (Bool True) = "#t"
showVal (Bool False) = "#f"
showVal (Character '\n') = "#\\newline"
showVal (Character ' ') = "#\\space"
showVal (Character char) = "#\\" ++ [char]
showVal (List contents) = "(" ++ unwordsList contents ++ ")"
showVal (DottedList head tail) = "(" ++ unwordsList head
                               ++ " . " ++ showVal tail ++ ")"
showVal (Vector contents) = "#(" ++ unwordsList (elems contents) ++ ")"
showVal (PrimitiveFunc _) = "<primitive>"
showVal (Func {params = args, vararg = varargs, body = body, closure = env}) =
  "(lambda (" ++ unwords (map show args) ++
    (case varargs of
       Nothing -> ""
       Just arg -> " . " ++ arg) ++ ") ...)"

unwordsList :: [LispVal] -> String
unwordsList = unwords . map showVal

instance Show LispVal where show = showVal

eval :: Env -> LispVal -> IOThrowsError LispVal
eval env val@(String _) = return val
eval env val@(Number _) = return val
eval env val@(Bool _) = return val
eval env (Atom id) = getVar env id
eval env (List [Atom "quote", val]) = return val
eval env (List [Atom "if", pred, conseq, alt]) = do
  result <- eval env pred
  case result of
    Bool False -> eval env alt
    otherwise -> eval env conseq
eval env (List [Atom "set!", Atom var, form]) =
    eval env form >>= setVar env var
eval env (List [Atom "define", Atom var, form]) =
    eval env form >>= defineVar env var
eval env (List (Atom "define" : List (Atom var : params) : body)) =
    makeNormalFunc env params body >>= defineVar env var
eval env (List (Atom "define" : DottedList(Atom var : params) varargs : body)) =
    makeVarargs varargs env params body >>= defineVar env var
eval env (List (Atom "lambda" : List params : body)) =
    makeNormalFunc env params body
eval env (List (Atom "lambda" : DottedList params varargs : body)) =
    makeVarargs varargs env params body
eval env (List (Atom "lambda" : varargs@(Atom _) : body)) =
    makeVarargs varargs env [] body
eval env (List (function : args)) = do
    func <- eval env function
    argVars <- mapM (eval env) args
    apply func argVars
eval env badForm =
    throwError $ BadSpecialForm "Unrecognized special form" badForm

apply :: LispVal -> [LispVal] -> IOThrowsError LispVal
apply (PrimitiveFunc func) args = liftThrows $ func args
apply (Func params varargs body closure) args =
    if num params /= num args && varargs == Nothing
       then throwError $ NumArgs (num params) args
       else (liftIO $ bindVars closure $ zip params args)
                >>= bindVarArgs varargs >>= evalBody
    where remainingArgs = drop (length params) args
          num = toInteger . length
          evalBody env = liftM last $ mapM (eval env) body
          bindVarArgs arg env = case arg of
              Just argName -> liftIO $
                  bindVars env [(argName, List $ remainingArgs)]
              Nothing -> return env

makeFunc varargs env params body =
    return $ Func (map showVal params) varargs body env

makeNormalFunc = makeFunc Nothing

makeVarargs = makeFunc . Just . showVal

primitives :: [(String, [LispVal] -> ThrowsError LispVal)]
primitives = [ ("+", numericBinop (+)),
               ("-", numericBinop (-)),
               ("*", numericBinop (*)),
               ("/", numericBinop div),
               ("mod", numericBinop mod),
               ("quotient", numericBinop quot),
               ("remainder", numericBinop rem),
               ("boolean?", unaryOp booleanp),
               ("symbol?", unaryOp symbolp),
{-               ("char?", unaryOp charp),
               ("vector?", unaryOp vectorp),
--               ("procedure?", unaryOp procedurep),-}
               ("pair?", unaryOp pairp),
               ("number?", unaryOp numberp),{-
               ("complex?", unaryOp complexp),
               ("real?", unaryOp realp),
               ("rational?", unaryOp rationalp),
               ("integer?", unaryOp integerp),-}
               ("string?", unaryOp stringp),
--               ("port?", unaryOp portp),-}
               ("symbol->string", unaryOp symbol2string),
               ("string->symbol", unaryOp string2symbol),
               ("=", numBoolBinop (==)),
               ("/=", numBoolBinop (/=)),
               ("<", numBoolBinop (<)),
               (">", numBoolBinop (>)),
               ("<=", numBoolBinop (<=)),
               (">=", numBoolBinop (>=)),
               ("&&", boolBoolBinop (&&)),
               ("||", boolBoolBinop (||)),
               ("string=?", strBoolBinop (==)),
               ("string<?", strBoolBinop (<)),
               ("string>?", strBoolBinop (>)),
               ("string<=?", strBoolBinop (<=)),
               ("string>=?", strBoolBinop (>=)),
               ("car", car),
               ("cdr", cdr),
               ("cons", cons),
               ("eq?", eqv),
               ("eqv?", eqv),
               ("equal?", equal)
             ]

numericBinop :: (Integer -> Integer -> Integer) -> [LispVal]
                -> ThrowsError LispVal
numericBinop op singleVal@[_] = throwError $ NumArgs 2 singleVal
numericBinop op params = mapM unpackNum params >>= return . Number . foldl1 op

unpackNum :: LispVal -> ThrowsError Integer
unpackNum (Number n) = return n
unpackNum nan = throwError $ TypeMismatch "number" nan

unaryOp :: (LispVal -> LispVal) -> [LispVal] -> ThrowsError LispVal
unaryOp f [v] = return $ f v

symbolp, numberp, stringp, booleanp, pairp :: LispVal -> LispVal
symbolp (Atom _)   = Bool True
symbolp _          = Bool False
numberp (Number _) = Bool True
numberp _          = Bool False
stringp (String _) = Bool True
stringp _          = Bool False
booleanp (Bool _)  = Bool True
booleanp   _       = Bool False
pairp (List _)     = Bool True
pairp (DottedList _ _) = Bool True
pairp _            = Bool False

symbol2string, string2symbol :: LispVal -> LispVal
symbol2string (Atom s)   = String s
symbol2string _          = String ""
string2symbol (String s) = Atom s
string2symbol _          = Atom ""
--booleanp, symbolp, charp, vectorp, {-procedurep,-}
--  pairp, numberp, complexp, realp, rationalp, integerp, stringp{-, portp-} :: LispVal -> Bool
--booleanp

boolBinop :: (LispVal -> ThrowsError a) -> (a -> a -> Bool)
             -> [LispVal] -> ThrowsError LispVal
boolBinop unpacker op args = if length args /= 2
                             then throwError $ NumArgs 2 args
                             else do left <- unpacker $ args !! 0
                                     right <- unpacker $ args !! 1
                                     return $ Bool $ left `op` right

numBoolBinop = boolBinop unpackNum
strBoolBinop = boolBinop unpackStr
boolBoolBinop = boolBinop unpackBool

unpackStr :: LispVal -> ThrowsError String
unpackStr (String s) = return s
unpackStr (Number s) = return $ show s
unpackStr (Bool s) = return $ show s
unpackStr notStr = throwError $ TypeMismatch "string" notStr

unpackBool :: LispVal -> ThrowsError Bool
unpackBool (Bool b) = return b
unpackBool (Number 0) = return False
unpackBool (Number 1) = return True
unpackBool notBool = throwError $ TypeMismatch "boolean" notBool

car :: [LispVal] -> ThrowsError LispVal
car [List (x : xs)] = return x
car [DottedList (x : xs) _] = return x
car [badArg] = throwError $ TypeMismatch "pair" badArg
car badArgList = throwError $ NumArgs 1 badArgList

cdr :: [LispVal] -> ThrowsError LispVal
cdr [List (x : xs)] = return $ List xs
cdr [DottedList [xs] x] = return x
cdr [DottedList (_ : xs) x] = return $ DottedList xs x
cdr [badArg] = throwError $ TypeMismatch "pair" badArg
cdr badArgList = throwError $ NumArgs 1 badArgList

cons :: [LispVal] -> ThrowsError LispVal
cons [x, List []] = return $ List [x]
cons [x, List xs] = return $ List (x : xs)
cons [x, DottedList xs xlast] = return $ DottedList (x : xs) xlast
cons [x1, x2] = return $ DottedList [x1] x2
cons badArgList = throwError $ NumArgs 2 badArgList

eqv :: [LispVal] -> ThrowsError LispVal
eqv [(Bool x), (Bool y)] = return $ Bool $ x == y
eqv [(Number x), (Number y)] = return $ Bool $ x == y
eqv [(String x), (String y)] = return $ Bool $ x == y
eqv [(Atom x), (Atom y)] = return $ Bool $ x == y
eqv [(DottedList xs x), (DottedList ys y)] =
    eqv [List $ xs ++ [x], List $ ys ++ [y]]
eqv [(List x), (List y)] =
    return $ Bool $ (length x == length y) && (all eqvPair $ zip x y)
  where eqvPair (x1, x2) = case eqv [x1, x2] of
                             Left err -> False
                             Right (Bool val) -> val
eqv [_, _] = return $ Bool False
eqv badArgList = throwError $ NumArgs 2 badArgList

data Unpacker = forall a. Eq a => AnyUnpacker (LispVal -> ThrowsError a)
unpackEquals :: LispVal -> LispVal -> Unpacker -> ThrowsError Bool
unpackEquals x y (AnyUnpacker unpacker) =
  do
    unpacked1 <- unpacker x
    unpacked2 <- unpacker y
    return $ unpacked1 == unpacked2
  `catchError` (const $ return False)

equal :: [LispVal] -> ThrowsError LispVal
equal [x, y] = do
  primitiveEquals <- liftM or $ mapM (unpackEquals x y)
    [AnyUnpacker unpackNum, AnyUnpacker unpackStr, AnyUnpacker unpackBool]
  eqvEquals <- eqv [x, y]
  return $ Bool $ (primitiveEquals || let (Bool x) = eqvEquals in x)
equal badArgList = throwError $ NumArgs 2 badArgList

data LispError = NumArgs Integer [LispVal]
               | TypeMismatch String LispVal
               | Parser ParseError
               | BadSpecialForm String LispVal
               | NotFunction String String
               | UnboundVar String String
               | Default String

showError :: LispError -> String
showError (UnboundVar message varname) = message ++ ": " ++ varname
showError (BadSpecialForm message form) = message ++ ": " ++ show form
showError (NotFunction message func) = message ++ ": " ++ show func
showError (NumArgs expected found) = "Expected: " ++ show expected
                                   ++ ", found " ++ show found
showError (TypeMismatch expected found) = "Invalid type: expected " ++ expected
                                        ++ ", found " ++ show found
showError (Parser parseErr) = "Parse error at " ++ show parseErr

instance Show LispError where show = showError
instance Error LispError where
  noMsg = Default "An error has occurred"
  strMsg = Default
type ThrowsError = Either LispError

--trapError ::
trapError action = catchError action (return . show)

extractValue :: ThrowsError a -> a
extractValue (Right val) = val

type Env = IORef [(String, IORef LispVal)]

nullEnv :: IO Env
nullEnv = newIORef []

primitiveBindings :: IO Env
primitiveBindings =
    nullEnv >>= (flip bindVars $ map makePrimitiveFunc primitives)
  where makePrimitiveFunc (var, func) = (var, PrimitiveFunc func)

type IOThrowsError = ErrorT LispError IO

liftThrows :: ThrowsError a -> IOThrowsError a
liftThrows (Left err) = throwError err
liftThrows (Right val) = return val

runIOThrows :: IOThrowsError String -> IO String
runIOThrows action = runErrorT (trapError action) >>= return . extractValue

isBound :: Env -> String -> IO Bool
isBound envRef var =
    readIORef envRef >>= return . maybe False (const True) . lookup var

getVar :: Env -> String -> IOThrowsError LispVal
getVar envRef var = do
  env <- liftIO $ readIORef envRef
  maybe (throwError $ UnboundVar "Getting an unbound variable" var)
        (liftIO . readIORef)
        (lookup var env)

setVar :: Env -> String -> LispVal -> IOThrowsError LispVal
setVar envRef var value = do
  env <- liftIO $ readIORef envRef
  maybe (throwError $ UnboundVar "Setting an unbound variable" var)
        (liftIO . (flip writeIORef value))
        (lookup var env)
  return value

defineVar :: Env -> String -> LispVal -> IOThrowsError LispVal
defineVar envRef var value = do
  alreadyDefined <- liftIO $ isBound envRef var
  if alreadyDefined
     then setVar envRef var value >> return value
     else liftIO $ do
       valueRef <- newIORef value
       env <- readIORef envRef
       writeIORef envRef ((var, valueRef) : env)
       return value

bindVars :: Env -> [(String, LispVal)] -> IO Env
bindVars envRef bindings = readIORef envRef >>= extendEnv bindings >>= newIORef
  where extendEnv bindings env = liftM (++ env) (mapM addBinding bindings)
        addBinding (var, value) = do ref <- newIORef value
                                     return (var, ref)

-- ghc -package parsec -fglasgow-exts -o 8function.exe --make 8function.hs
-- 8function.exe
-- Lisp>>> (define (f x y) (+ x y))
-- Lisp>>> (f 1 2)
-- Lisp>>> (f 1 2 3)
-- Lisp>>> (f 1)
-- Lisp>>> (define (factorial x) (if (= x 1) 1 (* x (factorial (- x 1)))))
-- Lisp>>> (factorial 10)
-- Lisp>>> (define (counter inc) (lambda (x) (set! inc (+ x inc)) inc))
-- Lisp>>> (define my-count (counter 5))
-- Lisp>>> (my-count 3)
-- Lisp>>> (my-count 6)
-- Lisp>>> (my-count 5)
-- Lisp>>> quit


module School.Utils.Debug
( compareDoubleMatrixTrace
, db
, printConduit
, printConduitP
, printVar
, showIO
, sl
, trace
) where

import Conduit (ConduitM, mapC, peekC)
import Control.Monad.IO.Class (MonadIO, liftIO)
import Debug.Trace (trace, traceIO)
import Numeric.LinearAlgebra (Matrix, R, toLists)
import School.Types.FloatEq (compareDouble)

sl :: (Show a) => a -> String
sl a = show a ++ " "

showIO :: (MonadIO m) => String -> m ()
showIO = liftIO . traceIO

printConduitP :: (Monad m, Show a)
              => String
              -> ConduitM a a m ()
printConduitP str = do
  next <- peekC
  mapC (db id (str ++ ": " ++ show next))

printConduit :: (MonadIO m, Show a)
             => String
             -> ConduitM a a m ()
printConduit str = do
  next <- peekC
  liftIO $ putStrLn (str ++ ": " ++ show next)
  mapC id

printVar :: (Show a) => String -> a -> a
printVar str x = trace (str ++ ": " ++ show x) x

db :: a -> String -> a
db = flip trace

compareDoubleMatrixTrace :: Double
                         -> Matrix R
                         -> Matrix R
                         -> Bool
compareDoubleMatrixTrace prec m1 m2 =
  and $ zipWith (traceCompare prec) l1 l2
  where traceMsg c msg v = if c
           then trace msg v
           else v
        traceCompare p d1 d2 =
          let equal = compareDouble p d1 d2
          in traceMsg (not equal) ("UNEQ " ++ show d1 ++ " " ++ show d2) equal
        l1 = concat . toLists $ m1
        l2 = concat . toLists $ m2


{-# Language RecordWildCards #-}
-- |
-- Module    : GRN.EM
-- Copyright : (c) 2011 Jason Knight
-- License   : BSD3
--
-- Maintainer  : jason@jasonknight.us
-- Stability   : experimental
-- Portability : portable
--
-- Expectation Maximization Algorithm at work
--

module GRN.EM where

import GRN.StateTransition
import GRN.Render
import GRN.Sparse
import GRN.Types
import GRN.Utils

import qualified Data.Map as M
import Data.List
import Data.Char
import Data.Maybe
import Control.Monad
import Text.Printf
import System.Console.ParseArgs

import qualified Data.Vector as V
import qualified Data.Vector.Unboxed as U
import qualified Data.Vector.Generic as G

import Control.Parallel.Strategies
import Data.Vector.Strategies

import Data.Graph.Inductive
import Data.Graph.Analysis.Algorithms.Common (componentsOf)

import Statistics.Distribution
import Statistics.Distribution.Normal

{-import Numeric.LBFGSB -}

import Graphics.Gnuplot.Simple

smoothingVar = 0.2 -- Variance of smoothing gaussian

ssaMap :: [Gene] -> SSA -> SSAMap
ssaMap glist s = M.fromList $ zip glist (U.toList s)

-- | Returns a map Gene->Double of any gene that has a measurement recorded.
measureMap :: ParseData -> MeasureMap
measureMap p = M.map fromJust $ M.fromList $ filter (isJust.snd) $ map (\(x,gi)-> (x,measurement gi)) (M.toList p)

-- | Returns a list of Uncertain locations.
uncertainVals :: KmapSet -> [UncertainSpot]
uncertainVals ks = filter filtFun $ concatMap explode (M.elems ks)
  where explode (Kmap g upstream x) = map (\(pos,entry)->(g,pos,entry)) (M.toList x)

filtFun :: (Gene,[Int],Kentry) -> Bool
filtFun (g,pos,x) = case x of
  C _ -> True
  X -> True
  _ -> False

fillKentries :: [UncertainSpot] -> KmapSet -> Theta -> KmapSet
fillKentries spots ks mm = foldr fillKentry ks (zip spots (U.toList mm))

fillKentry :: (UncertainSpot,Double) -> KmapSet -> KmapSet
fillKentry ((g,pos,_),val) ks = M.adjust fill g ks
  where fill (Kmap g up x) = Kmap g up (M.insert pos (V 0 val) x)

-- truncated Norm:: Mu  -> Var    -> Query
truncatedNorm :: Double -> Double -> Double -> Double
truncatedNorm mu var pt = 1/var * std ((pt-mu)/var) / denom
  where std = density standard
        cstd = cumulative standard
        denom = cstd ((1-mu)/var) - cstd ((0-mu)/var)

-- Given a theta vector, ssa vector, and a set of measurements, multiply the
-- measurement lookups convolved with a measurement gaussian against each other
-- to get a joint lookup.
msLookup :: SSAMap -> MeasureMap -> Double
msLookup ssam mm = product $ M.elems $ M.mapWithKey draw mm
  where draw gene meas = truncatedNorm (ssam M.! gene) smoothingVar meas

runAndSplit :: Args String -> KmapSet -> Masks -> Attractors
runAndSplit args ks masks = map ((\(d,s)->(d, ssaMap (M.keys ks) (ssdToSSA s))).normalizeSSD) raw
  where base = simulateCSC args csc (U.replicate n (1/(fromIntegral n)))
        dok = kmapToDOK ks 0
        csc = dokToCSC dok
        n = 2^(M.size ks)
        raw = map (filtSSD base) masks
        filtSSD base mask = zero `U.update` inds
          where zero = U.replicate n 0.0
                inds = U.map (\i->(i,base U.! i)) mask

runAndSplit2 :: Args String -> KmapSet -> Masks -> Attractors
runAndSplit2 args ks masks = map ((\(d,g)->(d, ssaMap (M.keys ks) (genSSA g))).normalizeGraph) raw
  where ssd = simulateDOKUnif args.kmapToDOK ks $ 0
        raw = componentsOf $ convertProbsVG ssd stripG
        stripG = last . take 10.iterate stripTransNodes $ kmapToStateGraph ks

-- Given an unfilled KMS and a new theta set, fill the kmap, then generate an
-- Attractors list
generateNewAttractors :: Args String -> KmapSet -> Masks -> [UncertainSpot] -> Theta -> Attractors
generateNewAttractors args ks masks us theta = runAndSplit args newks masks
  where newks = fillKentries us ks theta

expectationRun :: EMData -> Theta -> Double
expectationRun EMData{..} new = emNormSum * (G.sum (V.generate (length emAttractors) eachAttractor))
  where newAgivenTheta = generateNewAttractors emArgs emUnfilledKMS emCompNodes emUncertainLocs new
        eachAttractor i = msLookup oldssam emMeasurements * w * log (newAgivenO * msLookup newssam emMeasurements)
          where (w,oldssam) = emAttractors !! i
                (newAgivenO, newssam) = newAgivenTheta !! i

normSum :: Attractors -> MeasureMap -> Double
normSum atracs mm = (-1) / (sum $ map (\(w,issam)-> w * msLookup issam mm) atracs)

posteriorProb :: EMData -> Theta -> U.Vector Double
posteriorProb EMData{..} new = G.convert $ G.map (*(1/normSum)) (V.generate (length emAttractors) eachAttractor)
  where normSum = sum $ map (\(w,issam)-> w * msLookup issam emMeasurements ) newAgivenTheta
        newAgivenTheta = generateNewAttractors emArgs emUnfilledKMS emCompNodes emUncertainLocs new
        eachAttractor i = msLookup newssam emMeasurements * newAgivenO
          where (newAgivenO, newssam) = newAgivenTheta !! i

{-fullEM :: EMData -> Int -> Theta -> [(Theta,EMData)]-}
{-fullEM emd 0 oldtheta = [(oldtheta,emd)]-}
{-fullEM emd maxiter oldtheta -}
  {-| thetaClose = (oldtheta,emd) : [(newtheta,newemd)]-}
  {-| otherwise = (oldtheta,emd) : fullEM newemd (maxiter-1) newtheta-}
  {-where newemd = emd { emAttractors=newAtracs, emNormSum=newNormSum } -}
        {-newAtracs = generateNewAttractors (emArgs emd) (emUnfilledKMS emd) -}
                        {-(emCompNodes emd) (emUncertainLocs emd) newtheta-}
        {-newNormSum = normSum newAtracs (emMeasurements emd)-}
        {-newtheta = optimize (expectationRun emd) oldtheta-}
        {-thetaClose = (<1e-5).U.sum.U.map (**2).U.zipWith (-) newtheta $ oldtheta-}

--data EMData = EMData {
--      emAttractors :: Attractors
--    , emUnfilledKMS :: KmapSet
--    , emComps :: [ColoredStateGraph]
--    , emCompNodes :: Masks
--    , emNormSum :: Double
--    , emArgs :: Args String
--    , emMeasurements :: MeasureMap
--    , emUncertainLocs :: [UncertainSpot]
--    }

{-emRun :: Args String -> ParseData -> IO () -}
{-emRun args p = do-}
  {-let-}
    {-gen = gotArg args "generate"-}
    {-ks = buildKmaps p -}
    {-ms = measureMap p-}
    {-uvals = uncertainVals ks-}

    {-ng = M.size ks-}
    {-n = length uvals-}
    {-printList :: [U.Vector Double] -> String-}
    {-printList x = concat $ intersperse "\n" (map (concatMap (printf "%7.3f").G.toList) $ x)-}

    {---startTheta = U.fromList [0.379,0.5,1.0]-}
    {-startTheta = U.replicate n 0.5-}

    {-newks = fillKentries uvals ks startTheta-}
    {-stripG = componentsOf . last . take 10.iterate stripTransNodes $ kmapToStateGraph newks-}
    {-stripNodes = map (U.fromList.sort.nodes) stripG-}
    {-startAtracs = runAndSplit args newks stripNodes-}
    {-norms = normSum startAtracs ms-}

    {-emdat = EMData startAtracs ks stripG stripNodes norms args ms uvals-}

    {-mapt = map (U.fromList.map fromIntegral.dec2bin n) [0..2^n-1]-}
    {---mapt = map (U.replicate n) [0,0.1,0.9,1.0]-}
    {---mapt = map (U.fromList) [replicate n 0.0, replicate n 0.001]-}
    {---prVal = map (expectationRun emdat) mapt-}
    {-prVal = map (posteriorProb emdat) mapt -}
    {-prVal2 = map (expectationRun emdat) mapt -}

    {-elOpt = optimize (expectationRun emdat) (U.replicate n 0.5) -}
    {-path = fullEM emdat 10 startTheta-}

    {-xs = clamped [0,0.05..1]-}

    {-clamped xs = map (\x-> if (x>1.0) then 1.0 else (if (x<0.0) then 0.0 else x)) xs-}

    {-printSeg :: String -> String -> IO ()-}
    {-printSeg title st = do-}
          {-let l = 2 + length title-}
              {-upperTitle = map toUpper title-}
              {-width = 80-}
              {-half = (width-l) `div` 2-}
              {-stars = replicate half '*'-}
          {-printf "\n%s %s %s\n\n%s\n\n%s\n\n" stars upperTitle stars st (replicate width '*')-}

  --print n

  --forM_ mapt $ \t -> do
  --  let path = fullEM emdat 10 t
  {-printSeg "theta path" (printList.map (\(t,e)->t) $ path)-}
  {-printSeg "posterior probs" (printList.map (\(t,e)->posteriorProb e t) $ path)-}

  {---printSeg "attractors" (concatMap (printf "%7s".map intToDigit.dec2bin ng.G.head) stripNodes)-}
  {-printSeg "attractors" (concatMap (printf "%7d".G.head) stripNodes)-}
  {-print $ M.keys ks-}

  {---printf "%d Evaluations\n" ((length xs)^2)-}
  {-plotFunc3d [XLabel "X", YLabel "Y", XRange (0,1), YRange (0,1)] [] xs xs -}
     {-(\x y -> (-1)*(expectationRun emdat (U.fromList [0.5,x,y])))-}

  {---putStrLn "Expectations:"-}
  {---mapM_ (printf "%7.3f") prVal2-}
  {---putStrLn ""-}

  {---print elOpt-}

  {---putStrLn "Posterior Probs:"-}
  {---mapM_ (sequence.printList.U.fromList.sort.U.toList) prVal-}
  {---putStrLn ""-}
  {---mapM_ (printf "%5.2f") (map U.sum prVal)-}
  {---putStrLn ""-}

  {---putStrLn ""-}
  {---putStrLn ""-}
  {---print $ generateNewAttractors args ks stripNodes uvals (last mapt)-}
  {----}
  {---putStrLn ""-}
  {---putStrLn "Expectations:"-}
  {---mapM_ (printf "%7.3f\n") prVal2-}
  {---putStrLn ""-}

  {---putStrLn ""-}
  {---putStrLn ""-}
  {---putStrLn "**** Optimal Posterior Probs ****"-}
  {---mapM_ (sequence.printList) [posteriorProb emdat elOpt]-}
  {---mapM_ (printf "%7.3d".(U.head)) stripNodes-}
  {---putStrLn ""-}
  {---putStrLn "*********************************"-}
  {----}
  {---putStrLn ""-}
  {---printf "Optimal Expectation: %5.2f\n" (expectationRun emdat elOpt)-}
  {---putStrLn ""-}
  {---putStrLn "**** Optimal Theta ****"-}
  {---mapM_ (sequence.printList) [elOpt]-}

  {-printSeg "uncertain locs" (show uvals)-}
  {-printSeg "measurements" (show ms)-}

  {-when gen $ do-}
    {-let newks2 = fillKentries uvals ks (fst.last $ path)-}
        {-finalGraph = simulate args $ kmapToStateGraph newks2-}
        {-stripG2 = componentsOf . last . take 10.iterate stripTransNodes $ finalGraph-}
        {-stripNodes2 = map (U.fromList.sort.nodes) stripG2-}

    {-printSeg "final attractors" (concatMap (printf "%7d".G.head) stripNodes2)-}
    {-drawStateGraph finalGraph args-}

    {---let newks2 = fillKentries uvals ks (fst.head $ path)-}
    {---    interGraph = kmapToStateGraph newks2-}
    {---    ssd = simulateDOK args $ kmapToDOK newks2-}

    {----- WATCH OUT TODO I just found a huge bug: StateTransition.simulate does-}
    {----- not respect edge weights when moving probabilities around-}