{-# LANGUAGE DataKinds, RankNTypes, TypeFamilies #-}
{-# OPTIONS_GHC -Wno-incomplete-uni-patterns #-}
module TestSingleGradient (testTrees) where

import Prelude

import qualified Data.Array.ShapedS as OS
import qualified Data.Strict.Vector as Data.Vector
import qualified Data.Vector.Generic as V
import           Numeric.LinearAlgebra (Vector)
import           Test.Tasty
import           Test.Tasty.HUnit hiding (assert)

import HordeAd hiding (sumElementsVectorOfDual)
import HordeAd.Core.DualClass (IsPrimal, dAdd, dScale)

import TestCommon

testTrees :: [TestTree]
testTrees = [ testDReverse0
            , testDReverse1
            , testPrintDf
            , testDForward
            , testDFastForward
            , quickCheckForwardAndBackward
            , readmeTests
            , readmeTestsV
            ]

dReverse0
  :: HasDelta r
  => (DualNumberVariables 'DModeGradient r
      -> DualMonadGradient r (DualNumber 'DModeGradient r))
  -> [r]
  -> ([r], r)
dReverse0 f deltaInput =
  let ((results, _, _, _), value) =
        dReverse 1 f (V.fromList deltaInput, V.empty, V.empty, V.empty)
  in (V.toList results, value)

fX :: DualMonad 'DModeGradient Float m
   => DualNumberVariables 'DModeGradient Float
   -> m (DualNumber 'DModeGradient Float)
fX variables = do
  let x = var0 variables 0
  return x

fX1Y :: DualMonad 'DModeGradient Float m
     => DualNumberVariables 'DModeGradient Float
     -> m (DualNumber 'DModeGradient Float)
fX1Y variables = do
  let x = var0 variables 0
      y = var0 variables 1
  x1 <- x +\ 1
  x1 *\ y

fXXY :: DualMonad 'DModeGradient Float m
     => DualNumberVariables 'DModeGradient Float
     -> m (DualNumber 'DModeGradient Float)
fXXY variables = do
  let x = var0 variables 0
      y = var0 variables 1
  xy <- x *\ y
  x *\ xy

fXYplusZ :: DualMonad 'DModeGradient Float m
         => DualNumberVariables 'DModeGradient Float
         -> m (DualNumber 'DModeGradient Float)
fXYplusZ variables = do
  let x = var0 variables 0
      y = var0 variables 1
      z = var0 variables 2
  xy <- x *\ y
  xy +\ z

fXtoY :: DualMonad 'DModeGradient Float m
      => DualNumberVariables 'DModeGradient Float
      -> m (DualNumber 'DModeGradient Float)
fXtoY variables = do
  let x = var0 variables 0
      y = var0 variables 1
  x **\ y

freluX :: DualMonad 'DModeGradient Float m
       => DualNumberVariables 'DModeGradient Float
       -> m (DualNumber 'DModeGradient Float)
freluX variables = do
  let x = var0 variables 0
  reluAct x

testDReverse0 :: TestTree
testDReverse0 = testGroup "Simple dReverse application tests" $
  map (\(txt, f, v, expected) ->
        testCase txt $ dReverse0 f v @?= expected)
    [ ("fX", fX, [99], ([1.0],99.0))
    , ("fX1Y", fX1Y, [3, 2], ([2.0,4.0],8.0))
    , ("fXXY", fXXY, [3, 2], ([12.0,9.0],18.0))
    , ("fXYplusZ", fXYplusZ, [1, 2, 3], ([2.0,1.0,1.0],5.0))
    , ( "fXtoY", fXtoY, [0.00000000000001, 2]
      , ([2.0e-14,-3.2236188e-27],9.9999994e-29) )
    , ("fXtoY2", fXtoY, [1, 2], ([2.0,0.0],1.0))
    , ("freluX", freluX, [-1], ([0.0],0.0))
    , ("freluX2", freluX, [0], ([0.0],0.0))
    , ("freluX3", freluX, [0.0001], ([1.0],1.0e-4))
    , ("freluX4", freluX, [99], ([1.0],99.0))
    , ("fquad", fquad, [2, 3], ([4.0,6.0],18.0))
    , ("scalarSum", vec_omit_scalarSum_aux, [1, 1, 3], ([1.0,1.0,1.0],5.0))
    ]

vec_omit_scalarSum_aux
  :: DualMonad d r m
  => DualNumberVariables d r -> m (DualNumber d r)
vec_omit_scalarSum_aux vec = returnLet $ foldlDual' (+) 0 vec

sumElementsV
  :: DualMonad d r m
  => DualNumberVariables d r -> m (DualNumber d r)
sumElementsV variables = do
  let x = var1 variables 0
  returnLet $ sumElements0 x

altSumElementsV
  :: DualMonad d r m
  => DualNumberVariables d r -> m (DualNumber d r)
altSumElementsV variables = do
  let x = var1 variables 0
  returnLet $ altSumElements0 x

-- hlint would complain about spurious @id@, so we need to define our own.
id2 :: a -> a
id2 x = x

sinKonst
  :: DualMonad d r m
  => DualNumberVariables d r -> m (DualNumber d r)
sinKonst variables = do
  let x = var1 variables 0
  return $ sumElements0 $
    sin x + (id2 $ id2 $ id2 $ konst1 1 2)

sinKonstOut
  :: ( DualMonad d r m
     , Floating (Out (DualNumber d (Vector r))) )
  => DualNumberVariables d r -> m (DualNumber d r)
sinKonstOut variables = do
  let x = var1 variables 0
  return $ sumElements0 $
    unOut $ sin (Out x) + Out (id2 $ id2 $ id2 $ konst1 1 2)

sinDelayed :: (Floating a, IsPrimal d a) => DualNumber d a -> DualNumber d a
sinDelayed (D u u') = delayD (sin u) (dScale (cos u) u')

plusDelayed :: (Floating a, IsPrimal d a)
            => DualNumber d a -> DualNumber d a -> DualNumber d a
plusDelayed (D u u') (D v v') = delayD (u + v) (dAdd u' v')

sinKonstDelay
  :: DualMonad d r m
  => DualNumberVariables d r -> m (DualNumber d r)
sinKonstDelay variables = do
  let x = var1 variables 0
  return $ sumElements0 $
    sinDelayed x `plusDelayed` (id2 $ id2 $ id2 $ konst1 1 2)

powKonst
  :: DualMonad d r m
  => DualNumberVariables d r -> m (DualNumber d r)
powKonst variables = do
  let x = var1 variables 0
  return $ sumElements0 $
    x ** (sin x + (id2 $ id2 $ id2 $ konst1 (sumElements0 x) 2))

powKonstOut
  :: ( DualMonad d r m
     , Floating (Out (DualNumber d (Vector r))) )
  => DualNumberVariables d r -> m (DualNumber d r)
powKonstOut variables = do
  let x = var1 variables 0
  return $ sumElements0 $
    x ** unOut (sin (Out x)
                + Out (id2 $ id2 $ id2 $ konst1 (sumElements0 x) 2))

powKonstDelay
  :: DualMonad d r m
  => DualNumberVariables d r -> m (DualNumber d r)
powKonstDelay variables = do
  let x = var1 variables 0
  return $ sumElements0 $
    x ** (sinDelayed x
          `plusDelayed` (id2 $ id2 $ id2 $ konst1 (sumElements0 x) 2))

sinKonstS
  :: forall d r m. DualMonad d r m
  => DualNumberVariables d r -> m (DualNumber d r)
sinKonstS variables = do
  let x = varS variables 0
  return $ sumElements0 $ fromS1
    ((sin x + (id2 $ id2 $ id2 $ konstS 1))
       :: DualNumber d (OS.Array '[2] r))

sinKonstOutS
  :: forall r d m. ( DualMonad d r m
                   , Floating (Out (DualNumber d (OS.Array '[2] r))) )
  => DualNumberVariables d r -> m (DualNumber d r)
sinKonstOutS variables = do
  let x = varS variables 0
  return $ sumElements0 $ fromS1
    (unOut (sin (Out x) + Out (id2 $ id2 $ id2 $ konstS 1))
       :: DualNumber d (OS.Array '[2] r))

sinKonstDelayS
  :: forall d r m. DualMonad d r m
  => DualNumberVariables d r -> m (DualNumber d r)
sinKonstDelayS variables = do
  let x = varS variables 0
  return $ sumElements0 $ fromS1
    ((sinDelayed x `plusDelayed` (id2 $ id2 $ id2 $ konstS 1))
       :: DualNumber d (OS.Array '[2] r))

dReverse1
  :: (r ~ Float, d ~ 'DModeGradient)
  => (DualNumberVariables d r -> DualMonadGradient r (DualNumber d r))
  -> [[r]]
  -> ([[r]], r)
dReverse1 f deltaInput =
  let ((_, results, _, _), value) =
        dReverse 1 f
          (V.empty, V.fromList (map V.fromList deltaInput), V.empty, V.empty)
  in (map V.toList $ V.toList results, value)

testDReverse1 :: TestTree
testDReverse1 = testGroup "Simple dReverse application to vectors tests" $
  map (\(txt, f, v, expected) ->
        testCase txt $ dReverse1 f v @?= expected)
    [ ("sumElementsV", sumElementsV, [[1, 1, 3]], ([[1.0,1.0,1.0]],5.0))
    , ("altSumElementsV", altSumElementsV, [[1, 1, 3]], ([[1.0,1.0,1.0]],5.0))
    , ( "sinKonst", sinKonst, [[1, 3]]
      , ([[0.5403023,-0.9899925]],2.982591) )
    , ( "sinKonstOut", sinKonstOut, [[1, 3]]
      , ([[0.5403023,-0.9899925]],2.982591) )
    , ( "sinKonstDelay", sinKonstDelay, [[1, 3]]
      , ([[0.5403023,-0.9899925]],2.982591) )
    , ( "powKonst", powKonst, [[1, 3]]
      , ([[108.7523,131.60072]],95.58371) )
    , ( "powKonstOut", powKonstOut, [[1, 3]]
      , ([[108.7523,131.60072]],95.58371) )
    , ( "powKonstDelay", powKonstDelay, [[1, 3]]
      , ([[108.7523,131.60072]],95.58371) )
    ]

testPrintDf :: TestTree
testPrintDf = testGroup "Pretty printing test" $
  map (\(txt, f, v, expected) ->
        testCase txt $ prettyPrintDf False f
          (V.empty, V.fromList (map V.fromList v), V.empty, V.empty)
        @?= expected)
    [ ( "sumElementsV", sumElementsV, [[1 :: Float, 1, 3]]
      , unlines
        [ "let0 DeltaId_0 = SumElements0 (Var1 (DeltaId 0)) 3"
        , "in Var0 (DeltaId 0)" ] )
    , ( "altSumElementsV", altSumElementsV, [[1, 1, 3]]
      , unlines
        [ "let0 DeltaId_0 = Add0"
        , "  (Index0 (Var1 (DeltaId 0)) 2 3)"
        , "  (Add0"
        , "     (Index0 (Var1 (DeltaId 0)) 1 3)"
        , "     (Add0 (Index0 (Var1 (DeltaId 0)) 0 3) Zero0))"
        , "in Var0 (DeltaId 0)" ] )
    , ( "sinKonst", sinKonst, [[1, 3]]
      , unlines
        [ "in SumElements0"
        , "  (Add1"
        , "     (Scale1 [ 0.5403023 , -0.9899925 ] (Var1 (DeltaId 0)))"
        , "     (Konst1 Zero0 2))"
        , "  2" ] )
    , ( "sinKonstOut", sinKonstOut, [[1, 3]]
      , unlines
        [ "in SumElements0"
        , "  (Outline1"
        , "     PlusOut"
        , "     [ [ 0.84147096 , 0.14112 ] , [ 1.0 , 1.0 ] ]"
        , "     [ Outline1 SinOut [ [ 1.0 , 3.0 ] ] [ Var1 (DeltaId 0) ]"
        , "     , Konst1 Zero0 2"
        , "     ])"
        , "  2" ] )
    , ( "powKonst", powKonst, [[1, 3]]
      , unlines
        [ "in SumElements0"
        , "  (Add1"
        , "     (Scale1 [ 4.8414707 , 130.56084 ] (Var1 (DeltaId 0)))"
        , "     (Scale1"
        , "        [ 0.0 , 103.91083 ]"
        , "        (Add1"
        , "           (Scale1 [ 0.5403023 , -0.9899925 ] (Var1 (DeltaId 0)))"
        , "           (Konst1 (SumElements0 (Var1 (DeltaId 0)) 2) 2))))"
        , "  2" ] )
    , ( "powKonstOut", powKonstOut, [[1, 3]]
      , unlines
        [ "in SumElements0"
        , "  (Add1"
        , "     (Scale1 [ 4.8414707 , 130.56084 ] (Var1 (DeltaId 0)))"
        , "     (Scale1"
        , "        [ 0.0 , 103.91083 ]"
        , "        (Outline1"
        , "           PlusOut"
        , "           [ [ 0.84147096 , 0.14112 ] , [ 4.0 , 4.0 ] ]"
        , "           [ Outline1 SinOut [ [ 1.0 , 3.0 ] ] [ Var1 (DeltaId 0) ]"
        , "           , Konst1 (SumElements0 (Var1 (DeltaId 0)) 2) 2"
        , "           ])))"
        , "  2" ] )
    ]

testDForward :: TestTree
testDForward =
 testGroup "Simple dForward application tests" $
  map (\(txt, f, v, expected) ->
        let vp = listsToParameters v
        in testCase txt $ dForward f vp vp @?= expected)
    [ ("fquad", fquad, ([2 :: Double, 3], []), (26.0, 18.0))
    , ( "atanReadmeM", atanReadmeM, ([1.1, 2.2, 3.3], [])
      , (7.662345305800865, 4.9375516951604155) )
    , ( "vatanReadmeM", vatanReadmeM, ([], [1.1, 2.2, 3.3])
      , (7.662345305800865, 4.9375516951604155) )
    ]

testDFastForward :: TestTree
testDFastForward =
 testGroup "Simple dFastForward application tests" $
  map (\(txt, f, v, expected) ->
        let vp = listsToParameters v
        in testCase txt $ dFastForward f vp vp @?= expected)
    [ ("fquad", fquad, ([2 :: Double, 3], []), (26.0, 18.0))
    , ( "atanReadmeM", atanReadmeM, ([1.1, 2.2, 3.3], [])
      , (7.662345305800865, 4.9375516951604155) )
    , ( "vatanReadmeM", vatanReadmeM, ([], [1.1, 2.2, 3.3])
      , (7.662345305800865, 4.9375516951604155) )
    ]

-- The formula for comparing derivative and gradient is due to @awf
-- at https://github.com/Mikolaj/horde-ad/issues/15#issuecomment-1063251319
quickCheckForwardAndBackward :: TestTree
quickCheckForwardAndBackward =
  testGroup "Simple QuickCheck of gradient vs derivative vs perturbation"
    [ quickCheckTest0 "fquad" fquad (\(x, y, _z) -> ([x, y], [], [], []))
    , quickCheckTest0 "atanReadmeM" atanReadmeM
             (\(x, y, z) -> ([x, y, z], [], [], []))
    , quickCheckTest0 "vatanReadmeM" vatanReadmeM
             (\(x, y, z) -> ([], [x, y, z], [], []))
    , quickCheckTest0 "sinKonst" sinKonst  -- powKonst NaNs immediately
             (\(x, _, z) -> ([], [x, z], [], []))
    , quickCheckTest0 "sinKonstOut" sinKonstOut
             (\(x, _, z) -> ([], [x, z], [], []))
    , quickCheckTest0 "sinKonstDelay" sinKonstDelay
             (\(x, _, z) -> ([], [x, z], [], []))
    , quickCheckTest0 "sinKonstS" sinKonstS
             (\(x, _, z) -> ([], [], [], [x, z]))
    , quickCheckTest0 "sinKonstOutS" sinKonstOutS
             (\(x, _, z) -> ([], [], [], [x, z]))
    , quickCheckTest0 "sinKonstDelayS" sinKonstDelayS
             (\(x, _, z) -> ([], [], [], [x, z]))
   ]

-- A function that goes from `R^3` to `R^2`, with a representation
-- of the input and the output tuple that is convenient for interfacing
-- with the library.
atanReadmeOriginal :: RealFloat a => a -> a -> a -> Data.Vector.Vector a
atanReadmeOriginal x y z =
  let w = x * sin y
  in V.fromList [atan2 z w, z * x]

-- Here we instantiate the function to dual numbers
-- and add a glue code that selects the function inputs from
-- a uniform representation of objective function parameters
-- represented as delta-variables (`DualNumberVariables`).
atanReadmeVariables
  :: IsScalar d r
  => DualNumberVariables d r -> Data.Vector.Vector (DualNumber d r)
atanReadmeVariables variables =
  let x : y : z : _ = vars variables
  in atanReadmeOriginal x y z

-- According to the paper, to handle functions with non-scalar results,
-- we dot-product them with dt which, for simplicity, we here set
-- to a record containing only ones. We could also apply the dot-product
-- automatically in the library code (though perhaps we should
-- emit a warning too, in case the user just forgot to apply
-- a loss function and that's the only reason the result is not a scalar?).
-- For now, let's perform the dot product in user code.

-- Here is the function for dot product with ones, which is just the sum
-- of elements of a vector.
sumElementsOfDualNumbers
  :: IsScalar d r
  => Data.Vector.Vector (DualNumber d r) -> DualNumber d r
sumElementsOfDualNumbers = V.foldl' (+) 0

-- Here we apply the function.
atanReadmeScalar
  :: IsScalar d r
  => DualNumberVariables d r -> DualNumber d r
atanReadmeScalar = sumElementsOfDualNumbers . atanReadmeVariables

-- Here we introduce a single delta-let binding (`returnLet`) to ensure
-- that if this code is used in a larger context and repeated,
-- no explosion of delta-expressions can happen.
-- If the code above had any repeated non-variable expressions
-- (e.g., if @w@ appeared twice) the user would need to make it monadic
-- and apply @returnLet@ already there.
atanReadmeM
  :: DualMonad d r m
  => DualNumberVariables d r -> m (DualNumber d r)
atanReadmeM = returnLet . atanReadmeScalar

-- The underscores and empty vectors are placeholders for the vector,
-- matrix and arbitrary tensor components of the parameters tuple,
-- which we here don't use (above we construct a vector output,
-- but it's a vector of scalar parameters, not a single parameter
-- of rank 1).
atanReadmeDReverse :: HasDelta r
                   => Domain0 r -> (Domain0 r, r)
atanReadmeDReverse ds =
  let ((result, _, _, _), value) =
        dReverse 1 atanReadmeM (ds, V.empty, V.empty, V.empty)
  in (result, value)

readmeTests :: TestTree
readmeTests = testGroup "Simple tests for README"
  [ testCase " Float (1.1, 2.2, 3.3)"
    $ atanReadmeDReverse (V.fromList [1.1 :: Float, 2.2, 3.3])
      @?= (V.fromList [3.0715904, 0.18288425, 1.1761366], 4.937552)
  , testCase " Double (1.1, 2.2, 3.3)"
    $ atanReadmeDReverse (V.fromList [1.1 :: Double, 2.2, 3.3])
      @?= ( V.fromList [ 3.071590389300859
                       , 0.18288422990948425
                       , 1.1761365368997136 ]
          , 4.9375516951604155 )
  ]

-- And here's a version of the example that uses vector parameters
-- (quite wasteful in this case) and transforms intermediate results
-- via a primitive differentiable type of vectors instead of inside
-- vectors of primitive differentiable scalars.

vatanReadmeM
  :: DualMonad d r m
  => DualNumberVariables d r -> m (DualNumber d r)
vatanReadmeM variables = do
  let xyzVector = var1 variables 0
      [x, y, z] = map (index0 xyzVector) [0, 1, 2]
      v = seq1 $ atanReadmeOriginal x y z
  returnLet $ sumElements0 v

vatanReadmeDReverse :: HasDelta r
                    => Domain1 r -> (Domain1 r, r)
vatanReadmeDReverse dsV =
  let ((_, result, _, _), value) =
        dReverse 1 vatanReadmeM (V.empty, dsV, V.empty, V.empty)
  in (result, value)

readmeTestsV :: TestTree
readmeTestsV = testGroup "Simple tests of vector-based code for README"
  [ testCase "V Float (1.1, 2.2, 3.3)"
    $ vatanReadmeDReverse (V.singleton $ V.fromList [1.1 :: Float, 2.2, 3.3])
      @?= ( V.singleton $ V.fromList [3.0715904, 0.18288425, 1.1761366]
          , 4.937552 )
  , testCase "V Double (1.1, 2.2, 3.3)"
    $ vatanReadmeDReverse (V.singleton $ V.fromList [1.1 :: Double, 2.2, 3.3])
      @?= ( V.singleton $ V.fromList [ 3.071590389300859
                                     , 0.18288422990948425
                                     , 1.1761365368997136 ]
          , 4.9375516951604155 )
  ]


module Main where

import           Data.Complex
import qualified ExampleSeeds                  as E
import           Graphics.Gloss
import           Seed                           ( drawSeed )

main :: IO ()
main = animate FullScreen (greyN 0.05) (drawSeed E.heartSeed)


-- |
-- Module      : BattleHack.Piano
-- Description :
-- Copyright   : (c) Jonatan H Sundqvist, 2015
-- License     : MIT
-- Maintainer  : Jonatan H Sundqvist
-- Stability   : experimental
-- Portability : POSIX

-- Created September 12 2015

-- TODO | - Split up into several modules
--        -

-- SPEC | -
--        -

--------------------------------------------------------------------------------------------------------------------------------------------
-- GHC pragmas
--------------------------------------------------------------------------------------------------------------------------------------------

--------------------------------------------------------------------------------------------------------------------------------------------
-- API
--------------------------------------------------------------------------------------------------------------------------------------------
module BattleHack.Piano where

--------------------------------------------------------------------------------------------------------------------------------------------
-- We'll need these
--------------------------------------------------------------------------------------------------------------------------------------------
import Control.Lens hiding (inside)
import Data.List    (find)
import Data.Complex

import BattleHack.Types
import BattleHack.Lenses
import BattleHack.Utilities.Vector
import BattleHack.Utilities.General

--------------------------------------------------------------------------------------------------------------------------------------------
-- Data
--------------------------------------------------------------------------------------------------------------------------------------------
-- |
naturals :: Integral n => [n]
naturals = [0, 2, 4, 5, 7, 9, 11]

-- |
accidentals :: Integral n => [n]
accidentals = [1, 3, 6, 8, 10]

-- | The indeces of every natural key in order, starting at C0 (index 0)
allnaturals :: [Int]
allnaturals = zipWith (\i key -> div i 7 * 12 + key) [0..] $ cycle naturals

-- | The indeces of every accidental key in order, starting at C#0 (index 1)
allaccidentals :: [Int]
allaccidentals = zipWith (\i key -> div i 5 * 12 + key) [0..] $ cycle accidentals

-- |
chordlayout :: [KeyLayout]
chordlayout = [KeyRight, KeyAccidental, KeyBoth, KeyAccidental, KeyLeft, KeyRight, KeyAccidental, KeyBoth, KeyAccidental, KeyBoth, KeyAccidental, KeyLeft]

-- | Horizontal offset for each key in the octave, relative to the very first key
keysteps :: RealFloat r => [r]
keysteps = scanl1 (+) [0, 0, 1, 0, 1, 1, 0, 1, 0, 1, 0, 1]

--------------------------------------------------------------------------------------------------------------------------------------------
-- Functions
--------------------------------------------------------------------------------------------------------------------------------------------
-- |
keyorigin :: PianoSettings -> Int -> Vector
keyorigin piano i = piano-->origin + (sx*shiftx :+ 0)
  where
    (sx:+_)  = piano-->keysize
    shiftx   = 7*octaveFromKeyIndex i + (keysteps !! (i `mod` 12))

-- |
keylayout :: Int -> KeyLayout
keylayout i = chordlayout !! (i `mod` 12)

-- | Is the point inside the bounds of the given key?
--   Note that the function assumes that the key is pinned to the origin.
-- TODO: Bugs ahead (swat them!)
-- TODO: Rename (?)
inside :: PianoSettings -> KeyLayout -> Vector -> Bool
inside piano KeyLeft       p = insideBounds piano p && not (insideLeft piano p)
inside piano KeyRight      p = insideBounds piano p && not (insideRight piano p)
inside piano KeyBoth       p = insideBounds piano p && not (insideLeft piano p || insideRight piano p)
inside piano KeyAccidental p = insideMiddle piano p --  --insideLeft   piano p || insideRight piano p

-- | Is the point within the rectangular bounding box of the key?
insideBounds :: PianoSettings -> Vector -> Bool
insideBounds piano (x:+y) = let (dx:+dy) = piano-->keysize in between 0 dx x && between 0 dy y

-- | Is the point within the left indent of the key?
insideLeft :: PianoSettings -> Vector -> Bool
insideLeft  piano (x:+y) = let (dx:+dy) = dotwise (*) (piano-->indent :+ piano-->mid) (piano-->keysize) in between 0 dx x && between 0 dy y

-- | Is the point within the right indent of the key?
-- TODO: Buggy!
insideRight :: PianoSettings -> Vector -> Bool
insideRight piano p = let shiftx = piano-->keysize.real * (piano-->indent - 1) in insideLeft piano $ p + (shiftx:+0)
-- insideRight piano p = let shiftx = piano-->keysize.real * (1 - piano-->indent) in insideLeft piano $ p - (piano-->keysize.real:+0) + ((piano-->keysize.real*piano-->indent):+0)
-- insideRight piano (x:+y) = between ()

-- |
insideMiddle :: PianoSettings -> Vector -> Bool
insideMiddle piano (x:+y) = let (ox:+oy, dx:+dy) = keybounds piano KeyAccidental in between ox (ox+dx) x && between oy (oy+dy) y

-- |
-- TODO: Move to Piano
-- TODO: Rename (?)
layout :: RealFloat r => Complex r -> r -> r -> KeyLayout -> [Complex r]
layout (sx:+sy) indent mid which = case which of
  KeyRight      -> [nw,  nei,  rmi,         rm, se, sw]
  KeyBoth       -> [nwi, nei,  rmi,         rm, se, sw, lm, lmi]
  KeyLeft       -> [nwi, ne,   se,          sw, lm, lmi]
  KeyAccidental -> [nei, rmi,  (sx:+0)+lmi, (sx:+0)+nwi]
  where
    indent' = sx*indent
    mid'    = sy*mid

    nw = 0:+0   -- North west
    ne = sx:+0  -- North east
    se = sx:+sy -- South east
    sw = 0:+sy  -- South west

    nwi = indent':+0       -- North west indented
    nei = (sx-indent'):+0  -- North east indented

    lmi = indent':+mid'      -- Left  middle indent
    rmi = (sx-indent'):+mid' -- Right middle indent

    lm = 0:+mid'  -- Left middle
    rm = sx:+mid' -- Right middle

-- | Rectangular bounds of a key (currently as a topleft, size tuple)
-- TODO: Use Bounding Box type (?)
keybounds :: PianoSettings -> KeyLayout -> (Vector, Vector)
keybounds piano layout = case layout of
  KeyAccidental -> (piano-->keysize.real' - (indent':+0), (2*indent'):+mid')
  _             -> (0:+0, piano-->keysize)
  where
    (indent':+mid') = dotwise (*) (piano-->keysize) (piano-->indent:+piano-->mid)

-- |
-- TODO: Simplify
-- TODO: Don't hard-code the range
-- TODO: This is a pretty dumb algorithm for finding hovered-over keys
findKeyAt :: Vector -> PianoSettings -> Maybe Int
findKeyAt p piano' = find (insideFromKeyIndex piano' p) (zipWith const [0..] (piano'-->keys))

-- | Is the point inside the key at the given index?
-- TODO: Rename (?)
insideFromKeyIndex :: PianoSettings -> Vector -> Int -> Bool
insideFromKeyIndex piano' p i = inside piano' (keylayout i) (p-keyorigin piano' i)

-- |
pitchFromKeyIndex :: RealFloat r => Int -> r
pitchFromKeyIndex i = let i' = i+4+49 in 440.0*2.0**((fromIntegral i' - 49)/12.0) -- TODO: Make sure this is correct, elaborate on the meaning of the different index conversions
-- pitchFromKeyIndex i = 440*2**((fromIntegral $ i+3 + 12*4-49)/12)

-- | Which octave does the key belong to (the 0th octave starts with C0, the 1st with C1; no overlap between octaves)
octaveFromKeyIndex :: RealFloat r => Int -> r
octaveFromKeyIndex = fromIntegral . (`div` 12)

-- |
-- TODO: Refactor
-- TODO: Use Unicode (?)
notenameFromKeyIndex :: Int -> String
notenameFromKeyIndex i = ["C", "C#", "D", "D#", "E", "F", "F#", "G", "G#", "A", "A#", "B"] !! mod i 12


{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE DataKinds         #-}
{-# LANGUAGE TypeFamilies      #-}
{-# LANGUAGE RecordWildCards   #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeSynonymInstances #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE GADTs             #-}
{-# LANGUAGE FlexibleContexts  #-}

module Main where

import           Lens.Micro.Extras (view)
import qualified Data.Text.IO as Text (readFile)
import           Data.Text (lines)
import           Data.Maybe (mapMaybe, fromMaybe)
import           Data.List (zip4) --, foldl')

import           Control.Monad (forM_)
import           Control.Monad.Random (evalRand)
import           System.Random (mkStdGen, RandomGen)
import qualified System.Random as SR (split)

import           Data.Semigroup ((<>))
import           Options.Applicative (Parser, option, auto, optional, long, short, value,
                                      strOption, helper, idm, execParser, info, (<**>), (<|>),
                                      flag')
import qualified Data.ByteString as B
import           Data.Serialize (Serialize, runPut, put, runGet, get) --, Get
import           System.Directory (createDirectoryIfMissing)

import           Numeric.LinearAlgebra.Static ((#), (&), R, unwrap, konst) --, zipWithVector)
import           Numeric.LinearAlgebra (toList)

--import           GHC.TypeLits
import           Data.Singletons (Sing, SomeSing(SomeSing), toSing)
import           Data.Singletons.TypeLits (KnownNat, Nat, Sing(SNat))
import           Data.Singletons.Prelude.List (Head, Last)

import           Grenade (Network,  FullyConnected, Tanh, Logit, Shape(D1), --Relu,
                         LearningParameters(LearningParameters), randomNetwork, S(S1D),
                         Gradients)
--import           Grenade (Network((:~>), NNil),  FullyConnected(FullyConnected), FullyConnected'(FullyConnected'), Tanh, Logit, Shape(D1), runNet, Relu,
--                         LearningParameters(LearningParameters), randomNetwork, S(S1D))

import           GrenadeExtras (binaryNetError, trainOnBatchesEpochs, normalize, hotvector)
import           GrenadeExtras.Zip (Zip)
--import           GrenadeExtras.Orphan()
import           GrenadeExtras.GradNorm (GradNorm)

import           Song (Song, line2song, genre)
import           Song.Grenade (song2TupleRn) --, SongSD)
import           Shuffle (shuffle) -- TODO: use the "standard" shuffleM http://hackage.haskell.org/package/random-shuffle

--import           Debug.Trace (trace)

type LogisticRegression = Network '[ FullyConnected 75 1, Logit ]
                                  '[ 'D1 75, 'D1 1, 'D1 1 ]

type OneHiddenLayer n = Network '[ FullyConnected 75 n, Tanh, FullyConnected n 1, Logit ]
                                '[ 'D1 75, 'D1 n, 'D1 n, 'D1 1, 'D1 1 ]

type DiscreteSepContinuousNet = Network '[ Zip ('D1 21) ('D1 2) (FullyConnected 21 2) ('D1 54) ('D1 2) (FullyConnected 54 2),
                                           Tanh, FullyConnected 4 1,
                                           Logit ]
                                        '[ 'D1 75, 'D1 4, 'D1 4, 'D1 1, 'D1 1 ]

type HugeNetwork = Network '[ FullyConnected 75 40, Tanh, FullyConnected 40 10, Tanh, FullyConnected 10 3, Tanh, FullyConnected 3 1, Logit ]
                           '[ 'D1 75, 'D1 40, 'D1 40, 'D1 10, 'D1 10, 'D1 3, 'D1 3, 'D1 1, 'D1 1 ]

--type FFNet = Network '[ FullyConnected 75 100, Tanh, FullyConnected 100 40, Tanh, FullyConnected 40 20, Relu, FullyConnected 20 1, Logit ]
--                     '[ 'D1 75, 'D1 100, 'D1 100, 'D1 40, 'D1 40, 'D1 20, 'D1 20, 'D1 1, 'D1 1 ]

--netLoad :: KnownNat n => FilePath -> IO (OneHiddenLayer n)
netLoad :: Serialize b => FilePath -> IO b
netLoad modelPath = do
  modelData <- B.readFile modelPath
  either fail return $ runGet get modelData

createKFoldPartitions :: Int -> [a] -> (Int, [([a],[a])])
createKFoldPartitions size ks = (parts, take parts . separateTrainAndValidLists . splitInSizes $ ks)
  where
    len = length ks
    parts = floor $ fromIntegral len / (fromIntegral size :: Double)

    splitInSizes :: [a] -> [[a]]
    splitInSizes [] = []
    splitInSizes xs =
      let (start, finish) = splitAt size xs
       in start : splitInSizes finish

    separateTrainAndValidLists :: [[a]] -> [([a], [a])]
    separateTrainAndValidLists = fmap (\(t,v)->(t [], v)) . separate' ([]++)

    separate' :: ([a]->[a]) -> [[a]] -> [([a]->[a], [a])]
    separate' acc [x]    = [(acc, x)]
    separate' acc (x:xs) = (acc . (concat xs++), x) : separate' (acc . (x++)) xs
    separate' _   []     = error "this should never happen, the size of the partition is too big to even hold an element"

saveScores :: FilePath -> [NetScore] -> IO ()
saveScores logsPath scores = writeFile logsPath (headScores<>"\n"<>scoresStr) -- TODO: catch IO Errors!
  where
    scoresStr :: String
    scoresStr = mconcat . fmap netScore2String $ zip [(0::Int)..] scores

    netScore2String (i, NetScore{..}) = show i <> "\t"
                                     <> show trainClassError <> "\t"
                                     <> show trainingError   <> "\t"
                                     <> show testClassError  <> "\t"
                                     <> show gradientNorm    <> "\n"

    headScores = "epoch\ttrain_classification_error\ttrain_error\ttest_classification_error\tgradient_norm"

-- Reading data and running code

filenameDataset :: String
filenameDataset = "./msd_genre_dataset.txt"

getSongs :: String -> IO [Song]
getSongs filename = do
  --putStrLn "Reading data file ..." -- <- this doesn't work by the lazy nature of haskell
  file <- Text.readFile filename
  --let listLines = drop 10 $ Data.Text.lines file
  let listLines = Data.Text.lines file
      all_songs = mapMaybe line2song listLines
  return $ filter inRightGenre all_songs

 where inRightGenre song = let g = view genre song
                            in g == "dance and electronica" || g == "jazz and blues"

labelSong :: String -> Double
labelSong "dance and electronica" = 1
labelSong _                       = 0

data StoppingCondition = StoppingCondition { -- stop training if either:
  maxEpochs           :: Int,           -- maximum total epochs has been reached
  gradientSmallerthan :: (Double, Int), -- gradient is smaller than `d` for at least `i` consecutive epochs
  stuckInARange       :: (Double, Int)  -- or, the error hasn't changed more than `d` in the last `i` consecutive epochs
  }

data TypeOfModelToTrain = LogitRegModel | OneHiddenLayerModel Integer | ArbitraryNNModel Int

data ModelsParameters =
  ModelsParameters
    Float -- Size of test set
    Int   -- Size of batch
    Bool  -- Normalize data
    StoppingCondition
    LearningParameters -- Parameters for Gradient Descent
    (Maybe FilePath) -- Load path
    (Maybe FilePath) -- Save path
    (Maybe FilePath) -- Logs path
    (Maybe Double)   -- K-fold size of validation set
    TypeOfModelToTrain

modelsParameters :: Parser ModelsParameters
modelsParameters =
  ModelsParameters <$> option auto (long "test-set"   <> short 't' <> value 0.15)
                   <*> option auto (long "batch-size" <> short 'b' <> value 40)
                   <*> option auto (long "normalize"  <> value True)
                   <*> ( StoppingCondition
                         <$> option auto (long "epochs" <> short 'e' <> value 300)
                         <*> option auto (long "gradient-smaller" <> short 'g' <> value (0.085, 3))
                         <*> option auto (long "max-error-change" <> value (0.004, 15))
                       )
                   <*> (LearningParameters
                       <$> option auto (long "train_rate" <> short 'r' <> value 0.01)
                       <*> option auto (long "momentum" <> value 0.9)
                       <*> option auto (long "l2" <> value 0) -- 0.0005
                       )
                   <*> optional (strOption (long "load"))
                   <*> optional (strOption (long "save"))
                   <*> optional (strOption (long "logs"))
                   <*> optional (option auto (long "k-fold-val-size"))
                   <*> (   flag' LogitRegModel (long "logit-reg")
                       <|> (OneHiddenLayerModel <$> option auto (long "one-hidden-layer" <> short 'h'))
                       <|> (ArbitraryNNModel    <$> option auto (long "arbitrary-nn"))
                       )

addLog :: KnownNat n => R n -> R n
addLog a = signum a * log (abs a+1)

discreteToOneHotVector :: R 3 -> R 21
discreteToOneHotVector featDiscre = (fromMaybe (konst 0) $ hotvector (round a)::R 8)
                                  # (fromMaybe (konst 0) $ hotvector (round b)::R 12)
                                  & c
  where (a,b,c) = case toList $ unwrap featDiscre of
                    [a',b',c'] -> (a',b',c')
                    _          -> error $  "This is weird, there should be only three discrete "
                                        <> "features (R 3), aka. this never happens but the compiler "
                                        <> "asks me to give an exahustive list for pattern matching, doh'"

takeWhileCond :: ([a] -> [b]) -> (b -> Bool) -> [a] -> [a]
takeWhileCond toCond stopCond xs = fmap snd . takeWhile (stopCond . fst) $ zip (toCond xs) xs

data NetScore = NetScore {
  trainClassError :: Double, -- Classification training error
  trainingError   :: Double, -- (optimization) training error
  testClassError  :: Double, -- Classification testing error
  gradientNorm    :: Double  -- accumulated norm of the gradients on an epoch
} deriving (Show)

netScore :: (Last shapes ~ 'D1 1, Foldable t)
         => t (S (Head shapes), S ('D1 1))
         -> t (S (Head shapes), S ('D1 1))
         -> (Double, Network layers shapes)
         -> NetScore
netScore trainSet testSet (gradNorm, nn) =
  let (trainCE, trainE) = binaryNetError nn trainSet
      (testCE, _)       = binaryNetError nn testSet
   in
      NetScore { trainClassError = trainCE,
                 trainingError   = trainE,
                 testClassError  = testCE,
                 gradientNorm    = gradNorm }

takeWhileCondFunc :: Show a => StoppingCondition -> [(a, NetScore)] -> [(a, NetScore)]
takeWhileCondFunc (StoppingCondition epochs (grad, giters) (stuck, siters)) =
 -- trace ("hi" <> show (head (takeWhileCond maxErrorChangeInSIters (<stuck) xs))) .
  take (epochs+1) . takeWhileCond consecutiveSmallerThanGrad (<giters)
                  . takeWhileCond maxErrorChangeInSIters (>stuck)
  where
    consecutiveSmallerThanGrad :: [(a, NetScore)] -> [Int]
    consecutiveSmallerThanGrad = (0:) . countConsecutive (0::Int) . fmap (trainClassError . snd)
      where
        countConsecutive _    []     = []
        countConsecutive cond (x:xs) =
          let condVal = if x<grad then cond+1 else 0
           in condVal : countConsecutive condVal xs

    maxErrorChangeInSIters :: [(a, NetScore)] -> [Double]
    maxErrorChangeInSIters = (replicate siters (stuck+1) <>) . fmap findMaxChange . groupOnSIters . fmap (trainClassError . snd)
      where groupOnSIters :: [a] -> [[a]]
            groupOnSIters []         = []
            groupOnSIters xs@(_:xs') = take siters xs : groupOnSIters xs'
            findMaxChange :: [Double] -> Double
            findMaxChange xs = maximum xs - minimum xs

loadModel (Just loadFile) _       = netLoad loadFile
loadModel Nothing         seedNet = return $ evalRand randomNetwork seedNet

main :: IO ()
main = do
  mparams@(ModelsParameters _ _ _ _ _ load _ _ _ modelType) <- execParser (info (modelsParameters <**> helper) idm)

  let (seedNet, randSeed) = SR.split (mkStdGen 487239842)

  case modelType of
    LogitRegModel ->
      (loadModel load seedNet :: IO LogisticRegression) >>= mainWithRandmNet mparams randSeed

    OneHiddenLayerModel sizeHidden ->
      let (singSizeHidden :: SomeSing Nat) = toSing sizeHidden
      in case singSizeHidden of
           SomeSing (SNat :: Sing n) ->
             (loadModel load seedNet :: IO (OneHiddenLayer n)) >>=
             mainWithRandmNet mparams randSeed

    ArbitraryNNModel arbModelNum ->
      case arbModelNum of
        1 -> (loadModel load seedNet :: IO DiscreteSepContinuousNet) >>= mainWithRandmNet mparams randSeed
        2 -> (loadModel load seedNet :: IO HugeNetwork) >>= mainWithRandmNet mparams randSeed
        _ -> putStrLn $ "Sorry but there is no arbitrary neural net number " <> show arbModelNum

mainWithRandmNet :: (Head shapes ~ 'D1 75,
                     Last shapes ~ 'D1 1,
                     Show (Network layers shapes),
                     Serialize (Network layers shapes),
                     GradNorm (Gradients layers),
                     Num (Gradients layers),
                     RandomGen g)
                  => ModelsParameters -> g -> Network layers shapes -> IO ()
mainWithRandmNet mparams randSeed net0 = do
  let (seedShuffle, seedTraining) = SR.split randSeed
      ModelsParameters testPerc batchSize norm stopCond rate _ save logs kfoldValPer modelType = mparams
      modelName = case modelType of
                    LogitRegModel                  -> fromMaybe "logitReg"  logs
                    OneHiddenLayerModel sizeHidden -> fromMaybe "hidden"    logs <> "-" <> show sizeHidden
                    ArbitraryNNModel arbModelNum   ->
                      fromMaybe "arbitrary" logs <> "-" <> case arbModelNum of
                                                             1 -> "separated_discrete_continuous"
                                                             2 -> "huge_network"
                                                             _ -> error $ "There is no arbitrary Neural Net model #" <> show arbModelNum

  putStrLn $ "Model name: " <> modelName

  let normFun :: KnownNat m => [R m] -> [R m]
      normFun = if norm then normalize else id

  -- Loading songs
  songsRaw <- fmap (song2TupleRn labelSong) <$> getSongs filenameDataset
  let songsDiscrete :: [R 3]
      songsDiscrete = fmap (fst . fst) songsRaw
      songsDisOneHotVector :: [R 21]
      songsDisOneHotVector = fmap discreteToOneHotVector songsDiscrete
      --songsFloat    :: [R 27]
      songsFloat    = fmap (snd . fst) songsRaw
      --songsLog      :: [R 27]
      songsLog      = fmap addLog songsFloat
      --songsLabel    :: [R 1]
      songsLabel    = fmap snd songsRaw
      --songs :: [(S ('D1 57), S ('D1 1))]
      songs = (\(l1,l2,l3,o)->(S1D (l1#l2#l3), S1D o)) <$> zip4 songsDisOneHotVector (normFun songsFloat) (normFun songsLog) songsLabel

  --print . toList . unwrap $ foldl1' (zipWithVector max) songsDiscrete
  --print . toList . unwrap $ foldl1' (zipWithVector min) songsDiscrete

  let sizeSongs = length songsRaw
      -- Shuffling songs
      shuffledSongs       = evalRand (shuffle songs) seedShuffle
      (testSet, trainSet) = splitAt (round $ fromIntegral sizeSongs * testPerc) shuffledSongs

  {-
   -case net0 of
   -  x@(FullyConnected (FullyConnected' i o) _) :~> _ -> do
   -    let subnet = x :~> NNil :: Network '[FullyConnected 75 n] '[ 'D1 75, 'D1 n ]
   -    print $ runNet subnet (S1D 0) -- with this we get the biases
   -    --print . sqrt . normSquared $ net0
   -    print i -- biases for each output neuron
   -    print o -- weights between neurons
   -}

  -- Pretraining scores
  --print $ head shuffledSongs
  putStrLn $ "Total Size: " <> show sizeSongs
  putStrLn $ "Test Size:  " <> show (length testSet)
  putStrLn $ "Batch Size: " <> show batchSize

  case kfoldValPer of
    Nothing -> trainNet net0 trainSet testSet rate batchSize seedTraining stopCond logs save
    Just valPer -> do
      let valSize             = floor $ fromIntegral sizeSongs * valPer
          (parts, kfoldSets)  = createKFoldPartitions valSize trainSet
      putStrLn $ "Validation Size: " <> show valSize
      putStrLn $ "Total parts for kfold: " <> show parts

      forM_ (zip [(1::Int)..] kfoldSets) $ \(kfoldNum, (trainSet', valSet)) -> do
        let dirSaveNet = "kfold-val" <> show valPer <> "/" <> modelName <> "/" <> "kfoldnet" <> "-" <> show kfoldNum
            logsKfoldN = dirSaveNet <> ".txt"
            saveKfoldN = dirSaveNet <> "/" <> fromMaybe "nnet-hidden" save
        putStrLn ""
        createDirectoryIfMissing True dirSaveNet -- TODO: catch exceptions
        putStrLn $ "Results to be saved in " <> logsKfoldN
        putStrLn $ "Nets to be saved in " <> saveKfoldN
        trainNet net0 trainSet' valSet rate batchSize seedTraining stopCond (Just logsKfoldN) (Just saveKfoldN)

trainNet :: (Last shapes ~ 'D1 1,
             RandomGen g,
             Num (Gradients layers),
             GradNorm (Gradients layers),
             Show (Network layers shapes),
             Serialize (Network layers shapes))
         => Network layers shapes
         -> [(S (Head shapes), S (Last shapes))]
         -> [(S (Head shapes), S ('D1 1))]
         -> LearningParameters
         -> Int
         -> g
         -> StoppingCondition
         -> Maybe FilePath
         -> Maybe FilePath
         -> IO ()
trainNet net0 trainSet testSet rate batchSize seedTraining stopCond logs save = do
  -- Training Net
  let netsInf = (read "Infinity", net0) : evalRand (trainOnBatchesEpochs net0 rate trainSet batchSize) seedTraining
      netsScoresInf = fmap (netScore trainSet testSet) netsInf

      (nets, netsScores) = unzip . takeWhileCondFunc stopCond $ zip (fmap snd netsInf) netsScoresInf
      --net = last nets

  putStrLn $ "Epoch"
           <> "\tTraining classification error"
           <> "\tTraining error"
           <> "\tTesting error"
           <> "\tGradNorm"

  -- Showing results of training net
  forM_ (zip3 [(0::Integer)..] netsScores nets) $ \(epoch, NetScore trainCE trainE testE gradNorm, currentNet) -> do
    putStrLn $ show epoch
             <> "\t" <> show trainCE
             <> "\t" <> show trainE
             <> "\t" <> show testE
             <> "\t" <> show gradNorm
    -- saving current trained net
    case save of
      Just saveFile -> let saveFileBin = saveFile <> "-e_" <> show epoch
                                                  <> "-trainCE_" <> show trainCE
                                                  <> "-testE_"   <> show testE
                                                  <> ".bin"
                        in B.writeFile saveFileBin $ runPut (put currentNet)
      Nothing -> return ()

  case logs of
    Just logsPath -> saveScores logsPath netsScores
    Nothing       -> return ()

  --case save of
  --  Just saveFile -> B.writeFile saveFile $ runPut (put net)
  --  Nothing       -> return ()


module Lib (
    readExpr
  , eval
  , valueToPrint
) where

import Control.Monad
import Control.Monad.Error
import Control.Monad.Error.Class
import Data.Array
import Data.Complex
import Data.Ratio
import Numeric

import Text.ParserCombinators.Parsec hiding (spaces)

data LispVal = Atom String
             | Vector (Array Int LispVal)
             | List [LispVal]
             | DottedList [LispVal] LispVal
             | Number Integer
             | Real Double
             | Rational Rational
             | Complex (Complex Double)
             | String String
             | Character Char
             | Bool Bool

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

type ThrowsError = Either LispError

showVals :: [LispVal] -> String
showVals = unwords . map show

instance Show LispVal where
    show (Atom name) = name
    show (Vector contents) = "(" ++ (showVals $ elems contents) ++ ")"
    show (List contents) = "(" ++ showVals contents ++ ")"
    show (DottedList init last) = "(" ++ showVals init ++ " . " ++ show last ++ ")"
    show (Number n) = show n
    show (Real x) = show x
    show (Rational r) = show r
    show (Complex c) = show c
    show (String content) = content
    show (Character c) = charToString c
    show (Bool True) = "#t"
    show (Bool False) = "#f"

instance Show LispError where
    show (UnboundVar message varname) = message ++ ": " ++ varname
    show (BadSpecialForm message form) = message ++ ": " ++ show form
    show (NotFunction message func) = message ++ ": " ++ show func
    show (NumArgs expected found) = "Expected " ++ show expected ++ " arguments. Found values " ++ showVals found
    show (TypeMismatch expected found) = "Invalid type: expected " ++ expected ++ ", found " ++ show found
    show (Parser parseErr) = "Parse error at " ++ show parseErr

instance Error LispError where
     noMsg = Default "An error has occurred"
     strMsg = Default

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

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

spaces :: Parser ()
spaces = skipMany1 space

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

parseVector :: Parser LispVal
parseVector = do
    string "#("
    x <- parseVector'
    char ')'
    return x

parseDatum :: Parser LispVal
parseDatum = do
    char '.'
    spaces
    parseExpr

parseListContent :: Parser LispVal
parseListContent = do
    list <- sepEndBy parseExpr spaces
    datum <- optionMaybe parseDatum
    return $ case datum of
        Nothing -> List list
        Just datum -> DottedList list datum

parseList :: Parser LispVal
parseList = do
    char '('
    list <- parseListContent
    char ')'
    return list

parseQuoted :: Parser LispVal
parseQuoted = do
    char '\''
    x <- parseExpr
    return $ List [Atom "quote", x]

parseQuasiQuoted :: Parser LispVal
parseQuasiQuoted = do
    char '`'
    x <- parseExpr
    return $ List [Atom "quasiquote", x]

parseUnquoted :: Parser LispVal
parseUnquoted = do
    char ','
    x <- parseExpr
    return $ List [Atom "unquote", x]

parseDecimal :: Parser LispVal
parseDecimal = do
    x <- many1 digit
    (return . Number . read) x

parseDec :: Parser LispVal
parseDec = do
    try $ string "#d"
    x <- many1 digit
    (return . Number . read) x

hex2dig x = fst . head $ readHex x

parseHex :: Parser LispVal
parseHex = do
    try $ string "#h"
    x <- many1 hexDigit
    (return . Number . hex2dig) x

oct2dig x = fst . head $ readOct x

parseOct :: Parser LispVal
parseOct = do
    try $ string "#o"
    x <- many1 octDigit
    (return . Number . oct2dig) x

bin2dig = bin2dig' 0
bin2dig' value "" = value
bin2dig' value (x:xs) = let newValue = 2 * value + (if x == '0' then 0 else 1) in
                        bin2dig' newValue xs

parseBin :: Parser LispVal
parseBin = do
    try $ string "#b"
    let binDigit = oneOf "10"
    x <- many1 binDigit
    (return . Number . bin2dig) x

parseNumber :: Parser LispVal
parseNumber = parseDecimal
          <|> parseDec
          <|> parseHex
          <|> parseOct
          <|> parseBin

parseReal :: Parser LispVal
parseReal = do
    x <- many1 digit
    char '.'
    y <- many1 digit
    (return . Real . fst . head . readFloat) (y ++ "." ++ x)

parseRational :: Parser LispVal
parseRational = do
    x <- many1 digit
    char '/'
    y <- many1 digit
    (return . Rational) ((read x) % (read y))

toDouble :: LispVal -> Double
toDouble (Real r) = realToFrac r
toDouble (Number n) = fromIntegral n

parseComplex :: Parser LispVal
parseComplex = do
    x <- try parseReal <|> parseDecimal
    char '+'
    y <- try parseReal <|> parseDecimal
    char 'i'
    (return . Complex) (toDouble x :+ toDouble y)

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

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

charToString :: Char -> String
charToString x = [x]

parseCharacter :: Parser LispVal
parseCharacter = do
    string "#\\"
    value <- try (string "newline" <|> string "space")
        <|> do
            x <- anyChar
            notFollowedBy alphaNum
            return $ charToString x
    return $ Character $ case value of
        "space" -> ' '
        "newline" -> '\n'
        otherwise -> value !! 0

parseBool :: Parser LispVal
parseBool = do
    char '#'
    val <- oneOf "ft"
    return $ case val of
        'f' -> Bool False
        't' -> Bool True

parseExpr :: Parser LispVal
parseExpr = try parseAtom
        <|> try parseVector
        <|> try parseList
        <|> try parseQuoted
        <|> try parseQuasiQuoted
        <|> try parseUnquoted
        <|> try parseReal
        <|> try parseRational
        <|> try parseComplex
        <|> try parseNumber
        <|> try parseString
        <|> try parseCharacter
        <|> try parseBool

primitives :: [(String, [LispVal] -> ThrowsError LispVal)]
primitives = [("+", numberOp (+)),
              ("-", numberOp (-)),
              ("*", numberOp (*)),
              ("/", numberOp div),
              ("mod", numberOp mod),
              ("quotient", numberOp quot),
              ("remainder", numberOp rem),
              ("=", numberOpBool (==)),
              ("<", numberOpBool (<)),
              (">", numberOpBool (>)),
              ("/=", numberOpBool (/=)),
              ("<=", numberOpBool (<=)),
              (">=", numberOpBool (>=)),
              ("&&", boolOpBool (&&)),
              ("||", boolOpBool (||)),
              ("string=?", stringOpBool (==)),
              ("string<?", stringOpBool (<)),
              ("string>?", stringOpBool (>)),
              ("string<=?", stringOpBool (<=)),
              ("string>=?", stringOpBool (>=)),
              ("symbol?", isAtom),
              ("vector?", isVector),
              ("list?", isList),
              ("number?", isNumber),
              ("real?", isReal),
              ("rational?", isRational),
              ("complex?", isComplex),
              ("string?", isString),
              ("character?", isCharacter),
              ("bool?", isBool),
              ("string->symbol", string2symbol),
              ("symbol->string", symbol2string)]

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

unpackString :: LispVal -> ThrowsError String
unpackString (String n) = return n
unpackString notString = throwError $ TypeMismatch "string" notString

unpackBool :: LispVal -> ThrowsError Bool
unpackBool (Bool n) = return n
unpackBool notBool = throwError $ TypeMismatch "bool" notBool

getNumberValues :: [LispVal] -> ThrowsError [Integer]
getNumberValues lispvals = mapM unpackNumber lispvals

numberOp :: (Integer -> Integer -> Integer) -> [LispVal] -> ThrowsError LispVal
numberOp op [] = throwError $ NumArgs 2 []
numberOp op arg@[x] = throwError $ NumArgs 2 arg
numberOp op args = getNumberValues args >>= return . Number . foldl1 op

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

numberOpBool = opBool unpackNumber
stringOpBool = opBool unpackString
boolOpBool = opBool unpackBool

isAtom :: [LispVal] -> ThrowsError LispVal
isAtom [Atom _] = return $ Bool True
isAtom [List (Atom "quoted":xs)] = return $ Bool True
isAtom [_] = return $ Bool False
isAtom x = throwError $ NumArgs 1 x

isVector :: [LispVal] -> ThrowsError LispVal
isVector [Vector _] = return $ Bool True
isVector [_] =return $  Bool False
isVector x = throwError $ NumArgs 1 x

isList :: [LispVal] -> ThrowsError LispVal
isList [List _] = return $ Bool True
isList [DottedList _ _] = return $ Bool True
isList [_] = return $ Bool False
isList x = throwError $ NumArgs 1 x

isNumber :: [LispVal] -> ThrowsError LispVal
isNumber [Number _] = return $ Bool True
isNumber [_] = return $ Bool False
isNumber x = throwError $ NumArgs 1 x

isReal :: [LispVal] -> ThrowsError LispVal
isReal [Number _] = return $ Bool True
isReal [_] = return $ Bool False
isReal x = throwError $ NumArgs 1 x

isRational :: [LispVal] -> ThrowsError LispVal
isRational [Number _] = return $ Bool True
isRational [_] = return $ Bool False
isRational x = throwError $ NumArgs 1 x

isComplex :: [LispVal] -> ThrowsError LispVal
isComplex [Number _] = return $ Bool True
isComplex [_] = return $ Bool False
isComplex x = throwError $ NumArgs 1 x

isString :: [LispVal] -> ThrowsError LispVal
isString [String _] = return $ Bool True
isString [_] = return $ Bool False
isString x = throwError $ NumArgs 1 x

isCharacter :: [LispVal] -> ThrowsError LispVal
isCharacter [Character _] = return $ Bool True
isCharacter [_] = return $ Bool False
isCharacter x = throwError $ NumArgs 1 x

isBool :: [LispVal] -> ThrowsError LispVal
isBool [Bool _] = return $ Bool True
isBool [_] = return $ Bool False
isBool x = throwError $ NumArgs 1 x

string2symbol :: [LispVal] -> ThrowsError LispVal
string2symbol [String s] = return $ Atom s
string2symbol [x] = throwError $ TypeMismatch "string" x
string2symbol x = throwError $ NumArgs 1 x

symbol2string :: [LispVal] -> ThrowsError LispVal
symbol2string [Atom s] = return $ String s
symbol2string [x] = throwError $ TypeMismatch "symbol" x
symbol2string x = throwError $ NumArgs 1 x

valueToPrint = extractValue . trapError

trapError action = catchError action (return . show)

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

apply :: String -> [LispVal] -> ThrowsError LispVal
apply func args = maybe (throwError $ NotFunction "Unrecognized primitive function args" func)
                        ($ args)
                        (lookup func primitives)

eval :: LispVal -> ThrowsError LispVal
eval atom@(Atom _) = return atom
eval vector@(Vector _) = return vector
eval list@(List [Atom "quote", value]) = return value
eval list@(List (Atom func : args)) = mapM eval args >>= apply func
eval list@(List _) = return list
eval dotted@(DottedList _ _) = return dotted
eval number@(Number _) = return number
eval real@(Real _) = return real
eval rational@(Rational _) = return rational
eval complex@(Complex _) = return complex
eval string@(String _) = return string
eval char@(Character _) = return char
eval bool@(Bool _) = return bool

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


{-# OPTIONS_GHC  -fno-warn-unused-binds -fno-warn-unused-matches -fno-warn-name-shadowing -fno-warn-missing-signatures #-}
{-# LANGUAGE FlexibleInstances, MultiParamTypeClasses, UndecidableInstances, FlexibleContexts, TypeSynonymInstances #-}

---------------------------------------------------------------------------------------------------
---------------------------------------------------------------------------------------------------
-- |
-- | Module : Fequency Moments
-- | Creator: Xiao Ling
-- | Created: 12/17/2015
-- |
---------------------------------------------------------------------------------------------------
---------------------------------------------------------------------------------------------------

module FeqMoments where

import Prelude hiding (replicate)
import Control.Monad.Random.Class
import Control.Monad.Random
import Control.Monad.State

import Data.Conduit
import Data.Foldable (toList)
import qualified Data.Sequence as S
import qualified Data.Conduit.List as Cl
import Data.Sequence (Seq,(|>),update,empty)

import Core
import Statistics

{-----------------------------------------------------------------------------
  Types
------------------------------------------------------------------------------}

{-----------------------------------------------------------------------------
  Approximate Median
------------------------------------------------------------------------------}


{-# LANGUAGE CPP #-}
{-# OPTIONS_GHC -fno-warn-orphans #-}

module Instances where

import Control.Monad(replicateM)

#if !MIN_VERSION_base(4,8,0)
import Control.Applicative((<$>), (<*>), pure)
#endif

import Test.Tasty.QuickCheck
import Text.Namelist.Types
import Data.Complex(Complex((:+)))
import Data.CaseInsensitive(mk)
import Data.Char(toUpper)
import Numeric(showFFloat)

arbitrarySafeDouble :: Gen Double
arbitrarySafeDouble = suchThat (arbitrary :: Gen Double) $ \d ->
    let (n, _:f) = break (== '.') $ showFFloat Nothing d ""
    in d == read (n ++ f) / 10 ** fromIntegral (length f)

arbitraryIndex :: Gen Index
arbitraryIndex = do
    i <- choose (1, maxBound)
    return $ Index i

arbitraryRange :: Gen Index
arbitraryRange = do
    a <- oneof [Just <$> choose (1, maxBound - 1),          pure Nothing]
    b <- oneof [Just <$> choose (maybe 1 succ a, maxBound), pure Nothing]
    s <- oneof [Just . unSmall <$> arbitrary,               pure Nothing]
    return $ Range a b s
  where
    unSmall (Small a) = a

instance Arbitrary Index where
    arbitrary = oneof
        [ arbitraryIndex
        , arbitraryRange
        ]

arbitraryName :: Gen String
arbitraryName = sized $ \i -> do
    a  <- elements alpha
    as <- replicateM (min i 30) (elements an_)
    return $ a : as
  where
    an_   = '_': ['0' .. '9'] ++ alpha
    alpha = map toUpper lower ++ lower
    lower = ['a'..'z']

arbitraryKeyName :: Gen Key
arbitraryKeyName = oneof
    [ Key . mk <$> arbitraryName
    , Indexed  <$> (mk <$> arbitraryName) <*> listOf1 arbitrary
    ]

arbitrarySub :: Gen Key
arbitrarySub = sized $ \depth ->
    if depth < 1
    then arbitraryKeyName
    else Sub <$> (mk <$> arbitraryName) <*> arbitrary

instance Arbitrary Key where
    arbitrary = oneof
        [ arbitraryKeyName
        , arbitrarySub
        ]

newtype Scalar = Scalar { unScalar :: Value }

instance Arbitrary Scalar where
    arbitrary = Scalar <$> oneof
        [ Integer <$> arbitrary
        , Real    <$> arbitrarySafeDouble
        , Complex <$> ((:+) <$> arbitrarySafeDouble <*> arbitrarySafeDouble)
        , Logical <$> arbitrary
        , String  <$> arbitrary
        ]

instance Arbitrary Value where
    arbitrary = oneof
        [ unScalar <$> arbitrary
        , pure Null
        , sized $ \i -> do
            as <- map unScalar <$> replicateM (max 1 i) arbitrary
            a  <- unScalar <$> arbitrary
            return $ Array (reverse $ a:as)
        , sized $ \i -> (:*) <$> pure i <*> (unScalar <$> arbitrary)
        ]

instance Arbitrary Pair where
    arbitrary = (:=) <$> arbitrary <*> arbitrary

scaleDouble :: (Double -> Double) -> Gen a -> Gen a
scaleDouble f = scale (round . f . fromIntegral)

instance Arbitrary Group where
    arbitrary = Group <$> (mk <$> arbitraryName) <*> arbitrary

newtype Namelist = Namelist { getNamelist :: [Group] }

instance Show Namelist where
    show (Namelist gs) = show gs

instance Arbitrary Namelist where
    arbitrary = Namelist <$> scaleDouble sqrt arbitrary


{-# LANGUAGE DataKinds             #-}
{-# LANGUAGE DeriveAnyClass        #-}
{-# LANGUAGE DeriveGeneric         #-}
{-# LANGUAGE FlexibleContexts      #-}
{-# LANGUAGE FlexibleInstances     #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE RankNTypes            #-}
{-# LANGUAGE ScopedTypeVariables   #-}
{-# LANGUAGE Strict                #-}
{-# LANGUAGE TypeFamilies          #-}
{-# LANGUAGE TypeOperators         #-}
{-|
Module      : Grenade.Layers.LeakyTanh
Description : Hyperbolic tangent nonlinear layer
Copyright   : (c) Huw Campbell, 2016-2017
License     : BSD2
Stability   : experimental
-}
module Grenade.Layers.LeakyTanh
  ( LeakyTanh(..)
  , SpecLeakyTanh (..)
  , specLeakyTanh1D
  , specLeakyTanh2D
  , specLeakyTanh3D
  , specLeakyTanh
  , leakyTanhLayer
  ) where

import           Control.DeepSeq                (NFData (..))
import           Data.Constraint                (Dict (..))
import           Data.Proxy
import           Data.Reflection                (reifyNat)
import           Data.Serialize
import           Data.Singletons
import           GHC.Generics                   (Generic)
import           GHC.TypeLits
import qualified Numeric.LinearAlgebra.Static   as LAS
import           Unsafe.Coerce                  (unsafeCoerce)

import           Grenade.Core
import           Grenade.Dynamic
import           Grenade.Dynamic.Internal.Build
import           Grenade.Types
import           Grenade.Utils.Conversion
import           Grenade.Utils.Vector

-- | A LeakyTanh layer. A layer which can act between any shape of the same dimension, performing a tanh function. The maximum value is given in Promille, i.e. 995 is 0.995, and is positive and negative. I.e. LeakyTanh(x) = min 0.995 (max (-0.995) x).
data LeakyTanh (v :: Nat) = LeakyTanh
  deriving (Generic,NFData,Show)

instance UpdateLayer (LeakyTanh maxVal) where
  type Gradient (LeakyTanh maxVal) = ()
  runUpdate _ l@LeakyTanh{} _ = l

instance RandomLayer (LeakyTanh maxVal) where
  createRandomWith _ _ = return LeakyTanh

instance Serialize (LeakyTanh maxVal) where
  put _ = return ()
  get = return LeakyTanh

instance (a ~ b, SingI a, KnownNat maxVal) => Layer (LeakyTanh maxVal) a b where
  type Tape (LeakyTanh maxVal) a b = S a
  runForwards _ (S1DV v) = (S1DV v, S1DV $ mapVector (tanhMax maxVal) v)
    where
      maxVal = (/ 1000) $ fromIntegral $ max 0 $ min 1000 $ natVal (Proxy :: Proxy maxVal)
  runForwards _ (S2DV v) = (S2DV v, S2DV $ mapVector (tanhMax maxVal) v)
    where
      maxVal = (/ 1000) $ fromIntegral $ max 0 $ min 1000 $ natVal (Proxy :: Proxy maxVal)
  runForwards _ (S1D a) = (S1D a, S1D $ LAS.dvmap (tanhMax maxVal) a)
    where
      maxVal = (/ 1000) $ fromIntegral $ max 0 $ min 1000 $ natVal (Proxy :: Proxy maxVal)
  runForwards _ (S2D a) = (S2D a, S2D $ LAS.dmmap (tanhMax maxVal) a)
    where
      maxVal = (/ 1000) $ fromIntegral $ max 0 $ min 1000 $ natVal (Proxy :: Proxy maxVal)
  runForwards _ (S3D a) = (S3D a, S3D $ LAS.dmmap (tanhMax maxVal) a)
    where
      maxVal = (/ 1000) $ fromIntegral $ max 0 $ min 1000 $ natVal (Proxy :: Proxy maxVal)
  runBackwards _ (S1DV v) (S1DV gs) = ((), S1DV $ zipWithVector (\t g -> tanhMax' maxVal t * g) v gs)
    where
      maxVal = (/ 1000) $ fromIntegral $ max 0 $ min 1000 $ natVal (Proxy :: Proxy maxVal)
  runBackwards _ (S2DV v) (S2DV gs) = ((), S2DV $ zipWithVector (\t g -> tanhMax' maxVal t * g) v gs)
    where
      maxVal = (/ 1000) $ fromIntegral $ max 0 $ min 1000 $ natVal (Proxy :: Proxy maxVal)
  runBackwards _ (S1D a) (S1D g) = ((), S1D $ LAS.dvmap (tanhMax' maxVal) a * g)
    where
      maxVal = (/ 1000) $ fromIntegral $ max 0 $ min 1000 $ natVal (Proxy :: Proxy maxVal)
  runBackwards _ (S2D a) (S2D g) = ((), S2D $ LAS.dmmap (tanhMax' maxVal) (tanh' a) * g)
    where
      maxVal = (/ 1000) $ fromIntegral $ max 0 $ min 1000 $ natVal (Proxy :: Proxy maxVal)
  runBackwards _ (S3D a) (S3D g) = ((), S3D $ LAS.dmmap (tanhMax' maxVal) (tanh' a) * g)
    where
      maxVal = (/ 1000) $ fromIntegral $ max 0 $ min 1000 $ natVal (Proxy :: Proxy maxVal)
  runBackwards l x y = runBackwards l x (toLayerShape x y)

tanhMax :: (Ord a, Fractional a) => a -> a -> a
tanhMax m v
  | v > m = leaky * (v - m) + m
  | v < -m = leaky * (v - m) - m
  | otherwise = v
  where
    leaky = 0.05

tanhMax' :: (Ord a, Fractional a, Floating a) => a -> a -> a
tanhMax' m v
  | v > m = leaky
  | v < -m = -leaky
  | otherwise = max 0.005 (tanh' v)
  where
    leaky = 0.05

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

-------------------- DynamicNetwork instance --------------------

instance (KnownNat maxVal) => FromDynamicLayer (LeakyTanh maxVal) where
  fromDynamicLayer inp _ LeakyTanh = SpecNetLayer $ SpecLeakyTanh maxVal (tripleFromSomeShape inp)
    where maxVal = (/1000) $ fromIntegral $ max 0 $ min 1000 $ natVal (Proxy :: Proxy maxVal)

instance ToDynamicLayer SpecLeakyTanh where
  toDynamicLayer _ _ (SpecLeakyTanh maxVal (rows, cols, depth)) =
     reifyNat rows $ \(_ :: (KnownNat rows) => Proxy rows) ->
     reifyNat cols $ \(_ :: (KnownNat cols) => Proxy cols) ->
     reifyNat depth $ \(_ :: (KnownNat depth) => Proxy depth) ->
     reifyNat (round $ 1000 * maxVal) $ \(_ :: (KnownNat maxVal) => Proxy maxVal) ->
     case (rows, cols, depth) of
         (_, 1, 1) -> case (unsafeCoerce (Dict :: Dict()) :: Dict ()) of
           Dict -> return $ SpecLayer (LeakyTanh :: LeakyTanh maxVal) (sing :: Sing ('D1 rows)) (sing :: Sing ('D1 rows))
         (_, _, 1) -> case (unsafeCoerce (Dict :: Dict()) :: Dict ()) of
           Dict -> return $ SpecLayer (LeakyTanh :: LeakyTanh maxVal) (sing :: Sing ('D2 rows cols)) (sing :: Sing ('D2 rows cols))
         _         -> case (unsafeCoerce (Dict :: Dict()) :: Dict (KnownNat (rows GHC.TypeLits.* depth))) of
           Dict -> return $ SpecLayer (LeakyTanh :: LeakyTanh maxVal) (sing :: Sing ('D3 rows cols depth)) (sing :: Sing ('D3 rows cols depth))

-- | Create a specification for a LeakyTanh layer.
specLeakyTanh1D :: RealNum -> Integer -> SpecNet
specLeakyTanh1D maxVal i = specLeakyTanh3D maxVal (i, 1, 1)

-- | Create a specification for a LeakyTanh layer.
specLeakyTanh2D :: RealNum -> (Integer, Integer) -> SpecNet
specLeakyTanh2D maxVal (i, j) = specLeakyTanh3D maxVal (i, j, 1)

-- | Create a specification for a LeakyTanh layer.
specLeakyTanh3D :: RealNum -> (Integer, Integer, Integer) -> SpecNet
specLeakyTanh3D maxVal = SpecNetLayer . SpecLeakyTanh maxVal

-- | Create a specification for a LeakyTanh layer.
specLeakyTanh :: RealNum -> (Integer, Integer, Integer) -> SpecNet
specLeakyTanh maxVal = SpecNetLayer . SpecLeakyTanh maxVal

-- | Add a LeakyTanh layer to your build.
leakyTanhLayer :: RealNum -> BuildM ()
leakyTanhLayer maxVal
  | maxVal <= 0 || maxVal > 1 = error "The maxVal parameter of tanhMaxLayer has to be in (0,1]"
  | otherwise = buildGetLastLayerOut >>= buildAddSpec . SpecNetLayer . SpecLeakyTanh maxVal

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

instance GNum (LeakyTanh maxVal) where
  _ |* _ = LeakyTanh
  _ |+ _ = LeakyTanh
  sumG _ = LeakyTanh


--
-- Copyright (c) 2009-2011, ERICSSON AB
-- All rights reserved.
--
-- Redistribution and use in source and binary forms, with or without
-- modification, are permitted provided that the following conditions are met:
--
--     * Redistributions of source code must retain the above copyright notice,
--       this list of conditions and the following disclaimer.
--     * Redistributions in binary form must reproduce the above copyright
--       notice, this list of conditions and the following disclaimer in the
--       documentation and/or other materials provided with the distribution.
--     * Neither the name of the ERICSSON AB nor the names of its contributors
--       may be used to endorse or promote products derived from this software
--       without specific prior written permission.
--
-- THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
-- AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
-- IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
-- DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
-- FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
-- DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
-- SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
-- CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
-- OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
-- OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
--

module Feldspar.Core.Frontend.Complex
where

import Data.Complex (Complex)

import Feldspar.Core.Constructs
import Feldspar.Core.Constructs.Complex
import Feldspar.Core.Frontend.Num

complex :: (Numeric a, RealFloat a) => Data a -> Data a -> Data (Complex a)
complex = sugarSymF MkComplex

realPart :: (Numeric a, RealFloat a) => Data (Complex a) -> Data a
realPart = sugarSymF RealPart

imagPart :: (Numeric a, RealFloat a) => Data (Complex a) -> Data a
imagPart = sugarSymF ImagPart

conjugate :: (Numeric a, RealFloat a) => Data (Complex a) -> Data (Complex a)
conjugate = sugarSymF Conjugate

mkPolar :: (Numeric a, RealFloat a) => Data a -> Data a -> Data (Complex a)
mkPolar = sugarSymF MkPolar

cis :: (Numeric a, RealFloat a) => Data a -> Data (Complex a)
cis = sugarSymF Cis

magnitude :: (Numeric a, RealFloat a) => Data (Complex a) -> Data a
magnitude = sugarSymF Magnitude

phase :: (Numeric a, RealFloat a) => Data (Complex a) -> Data a
phase = sugarSymF Phase

polar :: (Numeric a, RealFloat a) => Data (Complex a) -> (Data a, Data a)
polar c = (magnitude c, phase c)

infixl 6 +.

(+.) :: (Numeric a, RealFloat a) => Data a -> Data a -> Data (Complex a)
(+.) = complex

iunit :: (Numeric a, RealFloat a) => Data (Complex a)
iunit = 0 +. 1


-- Rummikub.hs
-- Alex Striff

module Rummikub where

import Prelude hiding
  (Enum, succ, pred, toEnum, fromEnum, enumFrom, enumFromThen, enumFromTo, enumFromThenTo)
import Prelude.SafeEnum
import Control.Applicative
import Control.Arrow
import Control.Monad
import Data.Function
import Data.Functor
import Data.Maybe
import Data.Function.Pointless
import Data.Traversable
import Data.Graph
import Data.List
import Data.List.Split
import Text.Read
import Numeric.LinearAlgebra hiding ((<#), (<>))

newtype TileNum = TileNum Int deriving (Eq, Ord)

instance Show TileNum where
  show (TileNum n) = show n

instance Bounded TileNum where
  minBound = TileNum 1
  maxBound = TileNum 12

instance DownwardEnum TileNum where
  pred t@(TileNum n)
    | minBound < t && t <= maxBound = TileNum <$> pred n
    | otherwise = Nothing
  precedes (TileNum n) (TileNum m) = minBound <= n && n < m && m <= maxBound

instance UpwardEnum TileNum where
  succ t@(TileNum n)
    | minBound <= t && t < maxBound = TileNum <$> succ n
    | otherwise = Nothing
  succeeds (TileNum n) (TileNum m) = maxBound >= n && n > m && m >= minBound

instance Enum TileNum where
  toEnum n
    | minBound <= tn && tn <= maxBound = Just tn
    | otherwise = Nothing
    where tn = TileNum n
  fromEnum (TileNum n) = Just n

tilenum :: Int -> Maybe TileNum
tilenum = toEnum

data Color = Black | Red | Orange | Blue | AnyColor

instance Show Color where
  show Black    = "B"
  show Red      = "r"
  show Orange   = "o"
  show Blue     = "b"
  show AnyColor = "*"

instance Eq Color where
  Black    /= Black    = False
  Red      /= Red      = False
  Orange   /= Orange   = False
  Blue     /= Blue     = False
  _ /= _ = True

instance Ord Color where
  compare Black    Red      = LT
  compare Red      Orange   = LT
  compare Orange   Blue     = LT
  compare Blue     Black    = LT
  compare Black    Black    = EQ
  compare Red      Red      = EQ
  compare Orange   Orange   = EQ
  compare Blue     Blue     = EQ
  compare AnyColor _        = EQ
  compare _        AnyColor = EQ
  compare _        _        = GT

color :: Char -> Maybe Color
color s
  | s == 'B'  = Just Black
  | s == 'r'  = Just Red
  | s == 'o'  = Just Orange
  | s == 'b'  = Just Blue
  | s == '*'  = Just AnyColor
  | otherwise = Nothing

data Tile = NumTile TileNum Color | Joker deriving (Ord)

instance Show Tile where
  show (NumTile n c) = show c ++ show n
  show Joker = "j"

instance Eq Tile where
  (NumTile n c) == (NumTile m d) = n == m && c == d
  _ == _ = True -- Comparision with Joker

(<#)   :: Tile -> Tile -> Bool
(NumTile n c) <# (NumTile m d) = succ n == pure m && c == d
Joker         <# (NumTile m _) = minBound < m
(NumTile n _) <# Joker         = n < maxBound
Joker         <# Joker         = True

(<#<)  :: Tile -> Tile -> Bool
(NumTile n c) <#< (NumTile m d) = (succ <=< succ) n == pure m && c == d
Joker         <#< (NumTile m _) = pure minBound < pred m
(NumTile n _) <#< Joker         = succ n < pure maxBound
Joker         <#< Joker         = True

(<##<)  :: Tile -> Tile -> Bool
(NumTile n c) <##< (NumTile m d) = (succ <=< succ <=< succ) n == pure m && c == d
Joker         <##< (NumTile m _) = pure minBound < (pred <=< pred) m
(NumTile n _) <##< Joker         = (succ <=< succ) n < pure maxBound
Joker         <##< Joker         = True

(<@)   :: Tile -> Tile -> Bool
(NumTile n c) <@ (NumTile m d) = c /= d && n == m
_ <@ _ = True -- Comparison with Joker

(<@<)  :: Tile -> Tile -> Bool
(<@<) = (<@)

tileNum :: Tile -> Maybe TileNum
tileNum (NumTile n _) = pure n
tileNum Joker = Nothing

tileColor (NumTile _ c) = c
tileColor _ = AnyColor

tile :: String -> Maybe Tile
tile [] = Nothing
tile s@(c:n)
  | s == "j"  = Just Joker
  | otherwise = liftA2 NumTile (tilenum =<< readMaybe n) (color c)

tiles :: String -> Maybe [Tile]
tiles = traverse tile . splitOn " "

allRel :: (b -> [(a, a)]) -> (a -> a -> Bool) -> b -> Bool
allRel p r = all (uncurry r) . p

dpairs :: [a] -> [(a, a)]
dpairs = sequence <=< ap zip (tail <$> tails)

neighbors :: [a] -> [(a, a)]
neighbors = ap zip tail

acquaintances :: [a] -> [(a, a)]
acquaintances [_] = []
acquaintances xs  = ap zip (tail . tail) xs

randos :: [a] -> [(a, a)]
randos xs
  | length xs < 4 = []
  | otherwise     = ap zip (tail . tail . tail) xs

validNumRun :: [Tile] -> Bool
validNumRun = (. (&)) . flip all $
  [ (3 <=) . length
  , allRel neighbors (<#)
  , allRel acquaintances (<#<)
  , allRel randos (<##<)]

validColorRun :: [Tile] -> Bool
validColorRun = (. (&)) . flip all $
  [ liftA2 (||) (3 ==) (4 ==) . length
  , allRel dpairs ((/=) `on` tileColor)
  , allRel neighbors (==) . mapMaybe tileNum]

validRun :: [Tile] -> Bool
validRun = liftA2 (||) validNumRun validColorRun

type Run    = [Tile]
type Board  = [Run]
type Frag   = Run
type PBoard = ([Frag], Board)

showBoard :: Board -> String
showBoard = join . intersperse "\n" . fmap (join . intersperse " " . fmap show)

rightFriends :: Tile -> Tile -> Bool
rightFriends = (liftA2 . liftA2) (||) (<#) (<@)

leftFriends :: Tile -> Tile -> Bool
leftFriends = flip rightFriends

splitBin :: (a -> b -> Bool) -> a -> [b] -> [([b], [b])]
splitBin = (.) $ liftA2 fmap (flip splitAt) . findIndices

leftFriendSplits :: Tile -> Frag -> [(Frag, Frag)]
leftFriendSplits  = splitBin leftFriends
rightFriendSplits = splitBin rightFriends

friendSplits :: Frag -> Run -> ([(Frag, Frag)], [(Frag, Frag)])
friendSplits = liftA2 (&&&) (leftFriendSplits . head) (rightFriendSplits . last)

leftSplitRuns :: Frag -> (Frag, Frag) -> [Frag]
leftSplitRuns fs (xs, (f:ys))
  | null ys   = [left]
  | otherwise = [left, ys]
  where left  = xs ++ [f] ++ fs

rightSplitRuns :: Frag -> (Frag, Frag) -> [Frag]
rightSplitRuns fs (xs, xs')
  | null xs   = [right]
  | otherwise = [xs, right]
  where right = fs ++ xs'

fragPlays' :: Run -> Frag -> [[Frag]]
fragPlays' ts fs = (leftSplitRuns fs <$> ls) ++ (rightSplitRuns fs <$> rs)
  where (ls, rs) = friendSplits fs ts

fragPlays :: Run -> Frag -> [PBoard]
fragPlays = fmap sortRuns .: fragPlays'

frags :: PBoard -> [Frag]
frags = fst

runs :: PBoard -> Board
runs = snd

pboard :: [Frag] -> Board -> PBoard
pboard = (,)

parsePBoard :: [String] -> [String] -> Maybe PBoard
parsePBoard = liftA2 pboard `on` traverse tiles

uniq :: Eq a => [a] -> [a]
uniq = nub -- lame O(n^2)

altTopBot :: Int -> [Int] -- altTopBot 4 == [1, 5, 2, 4, 3]
altTopBot n
  | n > 0 = take n . join . ap (zipWith $ \a b -> [a, b]) reverse $ [1..n]
  | otherwise = []

sumSet :: Int -> [[Int]]
sumSet n = sumSet' n [[]] where
  sumSet' n ss
    | n > 0 = [1..n] >>= \a -> sumSet' (n - a) ((a:) <$> ss)
    | otherwise = ss

fragments :: [a] -> [[[a]]]
fragments xs = reverse $ flip splitPlaces xs <$> sumSet (length xs)

sortRuns :: [Frag] -> PBoard
sortRuns = foldl (\(a, b) r -> if validRun r then (a, r:b) else (r:a, b)) mempty

addSplitRuns :: PBoard -> PBoard -> PBoard
addSplitRuns = mappend

playFrags :: PBoard -> PBoard
playFrags p = addSplitRuns (pboard mempty . runs $ p) . sortRuns . frags $ p

select :: Eq b => (a -> [b]) -> a -> [(b, [b])]
select = (ap (zipWith $ liftA2 (.) (,) delete) repeat .)

dropUnfinished :: [PBoard] -> [PBoard]
dropUnfinished = filter $ null . frags

liftFrag :: Functor f => ([Frag] -> f [Frag]) -> PBoard -> f PBoard
liftFrag f = uncurry (<&>) . (f *** flip pboard)

fragment :: PBoard -> [(Frag, [Frag])]
fragment p = do
  (frag, fs)   <- select frags p -- select a fragment
  ffs          <- fragments frag -- smash it
  (frag', fs') <- select id ffs  -- select one of the pieces
  return $ (frag', fs' ++ fs)

stepPlay :: ([PBoard] -> [PBoard]) -> PBoard -> [PBoard]
stepPlay filt p
  | null $ frags p = return p
  | otherwise = do
  (frag, fs)   <- fragment p     -- select a fragment (by smashing)
  (run,  rs)   <- select runs p  -- select a run
  -- fp           <- filt . uniq $ fragPlays run frag -- play the fragment in the run
  fp           <- filt $ fragPlays run frag -- play the fragment in the run
  return $ addSplitRuns fp $ pboard fs rs

stepPlay' :: Int -> ([PBoard] -> [PBoard]) -> PBoard -> [PBoard]
stepPlay' n = foldr1 (<=<) . take n . repeat . stepPlay

fuse :: Eq a => [a] -> [a] -> [[a]]
fuse x@(a:as) y@(b:bs)
  | last as == b = [x ++ bs]
  | a == last bs = [y ++ as]
  | otherwise    = [x, y]

-- TODO: solve problem of turning
-- [[b1], [b2], [b3]] into [[b1 b2 b3]] without duplicates or using uniq/nub

-- fuseFrags' :: [Frag] -> [[Frag]]
-- fuseFrags' fl = do
--   (f, f')   <- dpairs fl
--   return $ fuse f f' ++ (fl \\ [f, f'])

groupFrags' :: [Frag] -> [[Frag]]
groupFrags' fl = do
  (f, f')   <- dpairs fl
  fg        <- fragPlays' f f'
  return $ fg ++ (fl \\ [f, f'])

groupFrags :: PBoard -> [PBoard]
groupFrags = liftFrag groupFrags'

-- Graph stuff

type NGraph node key = (Graph, Vertex -> (node, key, [key]), key -> Maybe Vertex)

graph :: Ord key => NGraph node key -> Graph
graph (g, _, _) = g

nodeFromVertex :: Ord key => NGraph node key -> (Vertex -> (node, key, [key]))
nodeFromVertex (_, nfv, _) = nfv

vertexFromKey :: Ord key => NGraph node key -> (key -> Maybe Vertex)
vertexFromKey (_, _, vfk) = vfk

split2 :: Int -> Int -> [a] -> ([a], [a])
split2 i j xs = (ys, zs)
  where (ys, ys') = splitAt i xs
        zs = drop j ys'

constrain1 :: Int -> [Bool] -> (Int, Int) -> [[Bool]]
constrain1 m xs (i, h)
  | m' > l = [[]]
  | otherwise = fmap (\i -> sandwich (split2 i m' xs) ts) [a..b]
  where m' = m - h + 1
        l = length xs
        ts = replicate m' True
        a = max 0 $ i - (m' - 1)
        b = min i $ l - m'
        sandwich (x, y) z = x ++ z ++ y

-- True represents a constrained vertex that will be propagated
-- m is the minimum number of tiles required to form a run
-- TODO: Rewrite to always be minimal
constrain :: Int -> [Bool] -> [[Bool]]
constrain m xs
  | length xs < m  = [[]]
  | not (or xs) = [xs]
  | otherwise = nub
    . join -- . fmap (minimumBy howConstrained)
    . groupBy ((==) `on` parts)
    . foldr1 (liftA2 $ zipWith (||))
    . fmap (constrain1 m xs)
    $ boundaryTrues xs
    where trues = length . filter id
          howConstrained = compare `on` trues
          sigbool b = if b then 1 else -1
          parts bs = ((*) . sigbool . head $ bs)
            . length . filter (uncurry (/=)) . ap zip tail $ bs

boundaryTrues :: [Bool] -> [(Int, Int)]
boundaryTrues bs = reverse . (\(_, _, ts) -> ts)
  . foldl go (head bs, 0, []) $ groups
  where go (b, i, ts) h = (not b, i+h, if b then (i+h-1,h):(i,h):ts else ts)
        groups = length <$> group bs

-- Needs tests to check all cases
-- TODO: add more counting stuff to deal with cases like splitting
-- *oo...oo* (or maybe even *oo..oo*)
growthComponents :: Int -> [Bool] -> [[Bool]]
growthComponents m bs = flip splitPlaces bs . reverse $ maxBound:gs -- maxBound avoids calculating length
  where gc (c, n, i, gs, j) True  = (True, 1, i + 1, maybe gs (:gs) j, Nothing)
        gc (c, n, i, gs, j) False = if n == 2 * (m - 1)
                                    then (False, 1, m, gs, Just $ i - (m - 1))
                                    else (c, n', i + 1, gs, j)
                                    where n' = if c then n + 1 else 0
        (_, _, i, gs, _) = foldl gc (False, 1, 1, [], Nothing) bs

showConstraint :: [Bool] -> String
showConstraint = fmap $ \b -> if b then 'o' else '.'

parseConstraint :: String -> Maybe [Bool]
parseConstraint = traverse parseConstraint'
  where parseConstraint' c
          | c == '.'  = Just False
          | c == 'o'  = Just True
          | otherwise = Nothing

parseConstraint' = fromJust . parseConstraint

mpow :: Matrix Z -> Int -> Matrix Z
mpow = mconcat .: flip replicate

adjacencyMatrix :: (a -> a -> Bool) -> [a] -> Matrix Z
adjacencyMatrix r = ap (join (><) . length) (fmap edge . join (liftA2 r))
  where edge b = (if b then 1 else 0) :: Z

paths :: Int -> Matrix Z -> Int -> Z
paths i = liftA2 ((+) `on` sumElements . (!! i)) toColumns toRows .: mpow

minpaths :: Int -> Int -> Matrix Z -> Z
minpaths m i a = sum $ fmap (paths i a) [m..cols a]

tileGraph :: (Tile -> Tile -> Bool) -> [Tile]
          -> NGraph Tile Int
tileGraph r ts = graphFromEdges
  $ zipWith (\t i -> (t, i, fmap snd . filter fst
    $ zipWith (\t' j -> (r t t', j)) ts [0..])) ts [0..]

-- Here m is the number of tiles
partitionRuns m = join (***) (concat . fmap (foldMap (:[])))
  . partition ((m <=) . length) . components . graph

constraints :: Ord key => Int -> [NGraph node key] -> [[[(Vertex, Bool)]]]
constraints m gs = fmap ((\(a, b) -> (fmap . fmap) (id &&& (`elem` b)) a)
  . (fst *** outcasts . fmap snd)) . select id
  . fmap ((id *** concat) . partRuns . comps) $ gs
  where comps = fmap (foldMap (:[])) . components . graph
        partRuns = partition ((m <=) . length)
        outcasts = foldl1 intersect


{-# LANGUAGE TypeOperators, FlexibleContexts, DeriveDataTypeable, FlexibleInstances, GADTs, ScopedTypeVariables, UndecidableInstances, OverlappingInstances, MultiParamTypeClasses, NoMonomorphismRestriction, RankNTypes, BangPatterns #-}

module Target.Prelude where

{-
(S.Seed,  powInt, sample, primUnit, printC, solveODE,solveODEs, runTests, assertIt, expect, optimise, choose, listToVec, pack, Consolable(..), Askable(..), fillV, Always(..),Never(..),Mostly(..),SamplerDensity(..), runIO, observeSig, (!$!), packL, Plot(..),Radian(..),ObservedSignal(..), csvFile, withNats,  packV, vecToList, fillM, mdims)
-}

import qualified Math.Probably.Sampler as S
--import qualified Math.Probably.MALA as MALA
import qualified Math.Probably.NelderMead as NM
import qualified Math.Probably.PDF as PDF
--import Math.Probably.BFGS
import System.Random.Mersenne.Pure64
import Math.Probably.FoldingStats
--import Math.Probably.GlobalRandoms
import qualified Numeric.LinearAlgebra as L
import Data.List
import Data.Char (toLower)
import Data.Maybe (fromJust)
import Text.Printf
import Data.Record
import Data.Record.Combinators ((!!!))
import Data.Kind
import Data.TypeFun
import qualified Data.Vector.Storable as V
import qualified Data.Vector.Generic as VG

import qualified Numeric.SpecFunctions as SM
import Statistics.Test.KolmogorovSmirnov
import qualified Data.Vector.Unboxed as UV
import Control.Monad.Trans (lift)
import Data.Array

import Debug.Trace
import Control.Monad

import Data.List
import System.Exit
import System.IO.Unsafe
import Unsafe.Coerce
import Foreign.StablePtr
import Foreign.Storable
import Foreign.Storable.Tuple
import qualified Data.Text as T
import qualified Data.Text.IO as TIO

import Data.STRef
import Control.Monad.ST
import qualified Data.Vector.Storable.Mutable as VSM

runIO :: IO a -> IO a
runIO =   id -- silence

  --return $ S.Samples [k $ L.fromList init0]

choose_arr = array ((0,0), (500,500)) [((n,k), SM.choose n k) | n <- [0..500], k<- [0..500]]
choose n k | n< 0 || k < 0 = realToFrac $ SM.choose n k
           | n <501 && k<501 = realToFrac $ choose_arr!(n,k)
           | otherwise = realToFrac $ SM.choose n k

primUnit = \seed -> let (x, nseed) = randomDouble seed
                    in (realToFrac x, nseed)

--sample :: (Seed -> (a, Seed)) -> IO a
sample s = S.sampleIO s

solveODE :: (Double -> Double -> Double) -> Double -> Double -> Double -> (Double -> Double)
solveODE f tmax dt y0 = y where
  ts = V.enumFromStepN 0 dt (round $ tmax/dt)
  ys = V.scanl intf y0 ts
  intf ylast t = ylast + dt * f ylast t
  y = pack dt 0 ys
--  y t = let ix = round $ t/dt
--        in  ys L.@> ix

nats :: [Int]
nats = [0..]

withNats :: [a] -> [(Int,a)]
withNats = zip nats

solveODEs :: ([Double] -> Double -> [Double]) -> Double -> Double -> [Double] -> [Double -> Double]
solveODEs f tmax dt y0 = y where
  is = [0..(round $ tmax/dt)]
  yss = transpose $ scanl intf y0 is
--   yss = L.toRows $ L.trans $ L.fromRows $ map (L.fromList) $ scanl intf y0 is
  intf ylast t = zipWith (+) ylast $ map (dt *) $ f ylast $ (*dt) . realToFrac $ t
  y  = flip map yss $ \ys -> let yv = L.fromList ys in pack dt 0 yv

decide :: (Real b,Floating b, V.Storable b, RealFrac b) => (b -> (((V.Vector b) -> (((S.Prob a) -> (((((V.Vector b) -> ((a -> b)))) -> (V.Vector b))))))))
decide = \(tol) -> \((ini::V.Vector b)) -> \((prob::S.Prob a)) -> \((util::((V.Vector b) -> ((a -> b))))) -> (optimise tol ini) (\((vaction::V.Vector b)) -> expect (fmap (util vaction) prob))

traceIt s x = trace (s++": "++show x) x

{-pmap :: (a -> b) -> ((c -> ((a,c))) -> (c -> ((b,c))))
pmap = \((f::a -> b)) -> \((sam::c -> ((a,c)))) -> bindP sam (\((x::a)) -> returnP (f x))

bindP :: (s -> ((a,s))) -> ((a -> (s -> ((b,s)))) -> (s -> ((b,s))))
bindP = \((f::s -> ((a,s)))) -> \((g::a -> (s -> ((b,s))))) -> \((s0::s)) -> let {((x::a),(s1::s)) = f s0} in g x s1

returnP :: a -> (b -> ((a,b)))
returnP = \((x::a)) -> \((s::b)) -> (x,s)
-}

{-for ::  Int -> (Int -> ((Int -> (s -> ((a,s)))) -> (s -> ((([a]),s)))))
for n m s | n < m = s n `bindP` \x ->
                    for (n+1) m s `bindP` \xs ->
                    returnP (x:xs)
          | otherwise = ((s n) `bindP` (\v -> returnP (v:[]))) -}

--uniform a b = (\x->(realToFrac x)*(b-a)+a) `pmap` ran0

expect :: Fractional c => S.Prob c -> c
expect (S.Samples xs) = runStat meanF xs
expect s@(S.Sampler _) = unsafePerformIO $ do
  xs <- S.sampleNIO 100 s
  return $ expect $ S.Samples xs

optimise :: (Real a , Floating a, RealFrac a, V.Storable a) => a -> V.Vector a -> (V.Vector a -> a) -> V.Vector a
optimise nmtol ini fitfunPos =

      let fitfun =  negate . realToFrac . fitfunPos . V.map realToFrac
          traceIt x = trace (show x) x
               --liftIO $ print (inivec, nmtol, fitfun inivec)
          inisimplex  = NM.genInitial fitfun [] (const 0.1) $ V.map realToFrac ini
          --check here if differences is too small
          --liftIO $ print inisimplex
          nmRes = NM.goNm fitfun [] (realToFrac nmtol) 100 5000 $ inisimplex
                 --liftIO $ putStrLn $ "bye from optimse" ++
                  --                    show (NM.centroid nmRes)
      in V.map realToFrac $ NM.centroid nmRes

--decide :: Double -> ((Vector Double) -> ((S.Prob a) -> (((Vector Double) -> (a -> Double)) -> (Vector Double))))
--decide = \((tol::Double)) -> \((ini::Vector Double)) -> \((prob::S.Prob a)) -> \((util::(Vector Double) -> (a -> Double))) -> (optimise tol ini) (\((vaction::Vector Double)) -> expect (fmap (util vaction) prob))

xs !$! ix = if length xs > ix
               then xs !! ix
               else error $ "!$! out of boiunds: " ++show (ix, length xs)

powInt :: Double -> Int -> Double
powInt x p = x ^ p

{-binomial n p =
   (for 1 n $ const $ pmap (<p) ran0) `bindP` \bools ->
   returnP $ length $ [t | t@True <- bools]

unormal =  ran0 `bindP` \u1 ->
           ran0 `bindP` \u2 ->
           returnP $ sqrt(-2*log(u1))*cos(2*pi*u2) -}

csvFile :: FromCSV a => T.Text -> IO [a]
csvFile fnm = do
  lns <- fmap (tail . T.lines) $ TIO.readFile (T.unpack fnm)
  return $ map (fromCSV . T.splitOn (T.pack ",")) lns

class FromCSV a where
  fromCSV :: [T.Text] -> a

instance FromCSV (X (Id KindStar)) where
  fromCSV _ = X

instance (Name lab, FromCSVCell a, FromCSV (b (Id KindStar)))
     => FromCSV ((b :& lab ::: a) (Id KindStar)) where
 fromCSV (t:ts) = (fromCSV ts) :& (name::lab) := (fromCSVcell t)

class FromCSVCell a where
  fromCSVcell :: T.Text -> a

instance FromCSVCell Double where
  fromCSVcell  = read . T.unpack

instance FromCSVCell T.Text where
  fromCSVcell  = id

printC :: Consolable a => T.Text -> a -> IO ()
printC s x = showC x >>= \s1-> putStrLn (T.unpack s++" => "++s1)

class Consolable a where
  showC :: a -> IO String
  showSamples :: [a] -> IO String

instance Consolable Double where
 showC x = return $ printf "%.3g" x
 showSamples xs =   let (mean,sd) = runStat meanSDF xs
                    in  return $ printf "%.3g +/- %.3g" mean sd

instance Consolable Bool where
 showC x = return $ show x
 showSamples bs =
  let ntrue = length $ filter id bs
      ntotal = length bs
  in return $ show ntrue ++"/"++show ntotal

instance Consolable Int where
 showC x = return $ show x
 showSamples xs =   let (mean::Double,sd) = runStat meanSDF $ map realToFrac xs
                    in  return $ printf "%.3g +/- %.3g" mean sd

instance Consolable [Char] where
 showC x = return x
 showSamples (x:_) = return x

instance Consolable (Double -> Double) where
 showC _ = return "<signal>"
 showSamples _ = return "<signals>"

instance Consolable T.Text where
 showC x = return $ T.unpack x
 showSamples (x:_) = return $ T.unpack x

instance Consolable a => Consolable [a] where
 showC x = do ss <- mapM showC x
              return $ "["++intercalate "," ss++"]"
 showSamples xs = do sxs <- mapM showC xs
                     return $ "oneOf "++show sxs

instance (Consolable a, V.Storable a) => Consolable (L.Vector a) where
 showC x = do ss <- mapM showC $ L.toList x
              return $ "<"++intercalate "," ss++">"
 showSamples xs = do sxs <- mapM showC xs
                     return $ "oneOf "++show sxs

instance (Consolable a, Consolable b) => Consolable (a,b) where
 showC (x,y) = do sx <- showC x
                  sy <- showC y
                  return $ "("++sx++","++sy++")"
 showSamples xs = do sxs <- mapM showC xs
                     return $ "oneOf" ++ show sxs

instance Consolable (X (Id KindStar)) where
 showC x = return ""
 showSamples _ = return ""

instance (Consolable a, Show lab, Consolable (b (Id KindStar)))
     => Consolable ((b :& lab ::: a) (Id KindStar)) where
 showC (rest :& lab := val) = do
      vs <- showC val
      rest <- showC rest
      return $ if null rest
                  then uncap (show lab) ++ "=>"++vs
                  else uncap (show lab) ++ "=>"++vs ++ ";"++rest
 showSamples recs = do
    let vs= map (\(rest :& lab := val) -> val) recs
    ss<- showSamples vs
    more <- showSamples (map getRest recs)
    return ("\n"++"   "++uncap (getLab (head recs)) ++ " => " ++ ss++more)

uncap [] = []
uncap (c:cs) = toLower c : cs

getRest (rest :& lab := val) = rest
getLab (rest :& lab := val) = show lab

sampleFrom n sam@(S.Sampler _) =  do
  sequence $ replicate n $ sam
sampleFrom n sam@(S.Samples xs) =  return xs

instance Consolable a => Consolable (S.Prob a) where
 showC sam@(S.Sampler _) = do
  xs <- S.sampleIO $ sequence $ replicate 100 $ sam
  showSamples xs
 showC sam@(S.Samples xs) = do
  showSamples xs
---unormal = S.gaussD 0 1
data Plot  =
   Plot ([(T.Text,T.Text)]) (S.Prob ([Radian]))
   |PlotRow ([(T.Text,T.Text)]) ([Plot])
   |PlotColumn ([(T.Text,T.Text)]) ([Plot])
   |PlotGrid ([(T.Text,T.Text)]) Int Int ([(T.Text,Plot)])
   |PlotStack ([(T.Text,T.Text)]) ([(T.Text,Plot)])
data Radian  =
   Lines ([(Double,Double)])
   |Points ([(Double,Double)])
   |Bars ([(Double,Double)])
   |Histogram ([Double])
   |Timeseries ((Double -> Double))
   |Options ([(T.Text,T.Text)]) ([Radian])
data ObservedSignal a =
    ObservedSignal Double Double (V.Vector Double)
   |ObservedXYSignal (V.Vector ((Double,Double)))
   deriving Show

observeSig :: (Double -> Double) -> ObservedSignal Double
observeSig  sig = unsafePerformIO $ do
    let ref = sig (0/0)
    deRefStablePtr $ unsafeCoerce ref

pack :: Double -> Double -> L.Vector Double -> (Double -> Double)
pack dt t0 ys = unsafePerformIO $ do
     putStrLn $ "packing.."
     let obs = ObservedSignal dt t0 ys
     ptr <- newStablePtr obs
     return $ \t -> if isNaN t
                       then unsafeCoerce ptr
                       else let ix = round $ (t-t0)/dt
                            in  ys L.@> ix

packL :: Double -> Double -> [Double] -> (Double->Double)
packL dt t0 ys  = pack dt t0 (V.fromList ys)

--packV dt t0 ys t = let ix = round $ (t-t0)/dt
--                  in  ys V.! ix

data Always  =
   Always (S.Prob Bool)
data Never  =
   Never (S.Prob Bool)
data Mostly  =
   Mostly (S.Prob Bool)
data SamplerDensity  =
     SamplerDensity (S.Prob Double) (Double -> Double)
   | SamplerMass (S.Prob Int) (Int -> Double)
   | SamplesSame (S.Prob Double) (S.Prob Double)
class Assertable a where
  assertIt :: a -> S.Prob Bool

instance Assertable Bool where
  assertIt = return

instance Assertable Always where
  assertIt (Always s) = do
     bs <- sequence $ replicate 100 $ s
     return $ all id bs

instance Assertable Never where
  assertIt (Never s) = do
     bs <- sequence $ replicate 100 $ s
     return $ all (id . not) bs

instance Assertable Mostly where
  assertIt (Mostly s) = do
     bs <- sequence $ replicate 100 $ s
     let yeas = length $ filter id bs
     let nays = length $ filter not bs
     return $ yeas > nays

ncdf = 2000

instance Assertable SamplerDensity where
  assertIt (SamplerDensity sam pdf) = do
      samples@(s0:_) <- sampleFrom 500 sam
      let (lo',hi') = runStat (both (minFrom s0) (maxFrom s0)) samples
          rng' = hi' - lo'
          lo = lo' - 3*rng'
          hi = hi' + 3*rng'
          rng = hi - lo
          dx = rng/ncdf
          xs = map ((+lo) . (*dx)) [0..(ncdf-1)]
          ys = map ((*dx) . exp . pdf) xs
          --prs = zip ys (tail ys)
          cumv =  listArray (0,(ncdf-1)) $ scanl (+) 0 ys
      let cdf x = cumv ! (round $ (x-lo)/dx)
      return $ kolmogorovSmirnovTestCdf cdf 0.05 (UV.fromList samples) == NotSignificant

  assertIt (SamplesSame sam1 sam2) = do
      samples1 <- sampleFrom 500 sam1
      samples2 <- sampleFrom 500 sam2
      return $ kolmogorovSmirnovTest2 0.05 (UV.fromList samples1)
                                           (UV.fromList samples2)
                                    == NotSignificant

runTests :: [(String,S.Prob Bool)] -> IO ()
runTests tsts = do
  seed <- S.getSeedIO
  rt 0 0 seed tsts
 where
  rt succ 0 _ [] = do
     putStrLn $ "All ("++show succ++") tests pass"
  rt succ fail _ [] = do
     putStrLn $ show fail++"/"++show (succ+fail)++" tests fail"
     exitWith $ ExitFailure 1
  rt s f seed ((nm, sam):rest) = do
     let (b,nseed) = sam1 seed sam
     if b
        then rt (s+1) f nseed rest
        else do putStrLn $ "FAIL: "++ nm
                rt s (f+1) nseed rest

sam1 seed (S.Samples xs) = (head xs, seed)
sam1 seed (S.Sampler f) = f seed

--myRound :: Real a => a-> Int

listToVec xs = L.fromList xs

vecToList xs = L.toList xs

fillV = L.buildVector

fillM (n,m) f = L.buildMatrix n m f

svd m = let (m1, v, m2) = L.svd m
        in [m1, L.diag v, m2]

mdims m = (L.rows m, L.cols m)

class Askable a where
  ask :: Double -> Double -> [String] -> a -> IO [String]

instance Consolable a => Askable a where
   ask _ _ lns x = do
      ans  <- fmap lines $ showC x
      return $ map (('>':) . (' ':)) $ init lns ++
                                       [last lns ++ " => " ++ head ans] ++
                                       tail' ans

tail' [] = []
tail' xs = tail xs

-- run-time reflection on records

{-class Has a where
  has :: String -> a -> Bool

instance Has (X (Id KindStar)) where
  has _ _ = False

instance (Show lab, Name lab, Has (b (Id KindStar)))
     => Has ((b :& lab ::: a) (Id KindStar)) where
  has nm (rest :& lab := val) = nm == show lab || has nm rest -}

instance L.Element Int

class (Num a, Ord a, L.Element a) => BayNum a where
  unround :: a -> Double

instance BayNum Double where
  unround = id

instance BayNum Int where
  unround = realToFrac

bayError = error . T.unpack

addSTRef ref x = modifySTRef ref (+x)

addSTRef ref !x = modifySTRef ref (+x)

consSTRef ref !x = modifySTRef ref (x:)

incrSTRef ref size = do
  x <- readSTRef ref
  let y = x+size
  y `seq` writeSTRef ref y
  return x

addMV vref ix x = do
  y <- VSM.read vref ix
  VSM.write vref ix (x+y)

showReal :: Double -> T.Text
showReal = T.pack . show

(*%) :: Double -> L.Matrix Double -> L.Matrix Double
(*%) = L.scale

transL :: [[a]] -> [[a]]
transL = transpose

unsafeFreeze = V.unsafeFreeze

vcons = V.cons
slice = V.slice -- FIXME: are we sure this uses inicies correctly?
vmap = V.map


{-|

 Module      :  HF
 Copyright   :  Copyright (c) David Schlegel
 License     :  BSD
 Maintainer  :  David Schlegel
 Stability   :  experimental
 Portability :  Haskell

 Module "HF".
-}

module HF (
{-| Module "HF.Data" provides all important datastructures and corresponding functions as well as dictionary lists.	-}
module HF.Data,
{-|In "ReadWrite" important functions are provided to

	* Read input (Geometry information, Basis set information) and convert to internal datastructures, provided in "Data"
	* Write output (Geometry, structure and basis set information)
-}
module HF.ReadWrite,
{-|Gaussian Integral evaluation plays an important role in quantum chemistry. In "HF.Gauss", functions will be provided to compute the most important integrals involving gaussian-type orbitals (GTOs). -}
module HF.Gauss)where
{-
--Total Packages included
import System.Process
import System.IO
import System.IO.Unsafe
import Data.List
import Data.List.Split
import Numeric.Container hiding (linspace)
import Data.Maybe
import Numeric.LinearAlgebra
import Numeric.LinearAlgebra.Data hiding (linspace)
import Text.Printf
-}

--Some Dictionaries
--dict1 = [("CARBON", "C"), ("NITROGEN", "N"), ("OXYGEN", "O"), ("HYDROGEN", "H")]
--dict2 = [("CARBON", 6.0), ("NITROGEN", 7.0), ("OXYGEN", 8.0), ("HYDROGEN", 1.0)]
--dict3 = [("S", 0), ("P", 1), ("D", 2), ("F", 3)]

import HF.Data

import HF.ReadWrite

import HF.Gauss

--import Matrices
import Numeric.Container hiding (linspace)