{-# LANGUAGE DataKinds #-}
{-# LANGUAGE ExistentialQuantification #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE PolyKinds #-}

module CSimpleSpec where

import Commons
import Control.Applicative (liftA2)
import Control.Concurrent
import Control.Monad (replicateM_, unless, when)
import Data.Array
import Data.Complex (Complex (..))
import qualified Data.IntMap.Strict as IM
import Data.List (intercalate, sort, tails)
import Data.List.Split (splitOn)
import qualified Data.Map.Strict as Map
import Data.Maybe (fromJust, mapMaybe)
import qualified Data.String.Interpolate as I
import qualified Data.Text as T
import qualified Data.Text.IO as TIO
import Debug.Trace (traceShowId)
import GHC.IO.Exception (ExitCode (..))
import HashedExpression.Codegen
import HashedExpression.Codegen.CSimple
import HashedExpression.Embed.FFTW
import HashedExpression.Internal.Expression hiding (InnerProductSpaceOp (..))
import HashedExpression.Internal.Inner
import HashedExpression.Internal.Node
import HashedExpression.Internal.Normalize (normalize)
import HashedExpression.Internal.Utils
import HashedExpression.Interp
import HashedExpression.Prettify (showExp, showExpDebug)
import HashedExpression.Value
import System.FilePath
import System.Process (readProcess, readProcessWithExitCode)
import Test.Hspec
import Test.QuickCheck
import Var
import Prelude hiding ((!!))

-- | Compute the inner offset:
-- e.g localOffset [3, 4, 5] [2, 0, 1] = 2 * (4 * 5) + 0 * (5) + 1
localOffset :: [Int] -> [Int] -> Int
localOffset shape indices
  | length shape == length indices =
    sum . zipWith (*) indices . map product . tail . tails $ shape
  | otherwise = error $ "shape and indices are not compatible" ++ show shape ++ show indices

-- | Generate a fully working C program that compute the expression and
-- print out the result, mostly used for testing
singleExpressionCProgram ::
  (DimensionType d, NumType et) => ValMaps -> Expression d et -> Code
singleExpressionCProgram valMaps expr =
  [ "#include <math.h>", --
    "#include <stdio.h>",
    "#include <stdlib.h>",
    if containsFTNode $ exMap expr
      then T.pack $ fftUtils
      else "",
    "int main()" --
  ]
    ++ scoped (initMemory ++ assignVals ++ codes ++ printValue ++ releaseMemory)
  where
    (mp, n) = unwrap expr
    bound = product (retrieveShape n mp)
    et = retrieveElementType n mp
    codeGen = initCodegen (CodegenInit mp []) CSimpleConfig
    [i, j, k, nooffset] = ["i", "j", "k", "0"]
    initMemory = [[I.i|double *ptr = malloc(sizeof(double) * #{cMemSize codeGen});|]]
    -- assign value to variables
    assignVals = assigningValues codeGen valMaps
    -- codes to compute
    codes = evaluating codeGen [n]
    -- print the value of expression
    printValue
      | et == R = for i bound [[I.i|printf("%f ", #{(!!) codeGen n i});|]]
      | et == C =
        for i bound [[I.i|printf("%f ", #{(!!) codeGen n i});|]]
          ++ [[I.i|printf("\\n");|]]
          ++ for i bound [[I.i|printf("%f ", #{imAt codeGen n i});|]]
    releaseMemory = ["free(ptr);"]
    -------------------------------------------------------------------------------
    assigningValues :: CSimpleCodegen -> ValMaps -> Code
    assigningValues CSimpleCodegen {..} valMaps = concatMap codesForVar vars
      where
        [i, j, k, nooffset] = ["i", "j", "k", "0"]
        vars :: [(Int, String)]
        vars =
          let toVar nId
                | Var varName <- retrieveNode nId mp = Just (nId, varName)
                | otherwise = Nothing
           in mapMaybe toVar . IM.keys $ mp
        codesForVar :: (Int, String) -> Code
        codesForVar (n, varName) =
          case Map.lookup varName valMaps of
            Just (VScalar val) -> [[I.i|#{n !! nooffset} = #{val};|]]
            Just (V1D array1d) ->
              let assignIndex id = [[I.i|#{n !! (showT id)} = #{array1d ! id};|]]
               in concatMap assignIndex $ indices array1d
            Just (V2D array2d) ->
              let shape = retrieveShape n mp
                  assignIndex (id1, id2) =
                    [[I.i|#{n !! showT (localOffset shape [id1, id2])} = #{array2d ! (id1, id2)};|]]
               in concatMap assignIndex $ indices array2d
            Just (V3D array3d) ->
              let shape = retrieveShape n mp
                  assignIndex (id1, id2, id3) =
                    [[I.i|#{n !! showT (localOffset shape [id1, id2, id3])} = #{array3d ! (id1, id2, id2)};|]]
               in concatMap assignIndex $ indices array3d
            _ -> []

hasFFTW :: IO Bool
hasFFTW = do
  fileName <- generate $ vectorOf 10 $ elements ['A' .. 'Z']
  let cFilePath = "C" </> fileName <.> "c"
      executablePath = "C" </> fileName
      codes = ["#include <fftw3.h>", "int main() {}"]
      program = T.intercalate "\n" . map T.pack $ codes
  TIO.writeFile cFilePath program
  (exitCode, output, _) <- readProcessWithExitCode "gcc" [cFilePath, "-o", executablePath, "-lm", "-lfftw3"] ""
  readProcessWithExitCode "rm" [cFilePath] ""
  readProcessWithExitCode "rm" [executablePath] ""
  return $ exitCode == ExitSuccess

-- |
evaluateCodeC :: (DimensionType d, NumType et) => Bool -> Expression d et -> ValMaps -> IO (ExitCode, String)
evaluateCodeC withFT exp valMaps = do
  readProcessWithExitCode "mkdir" ["C"] ""
  fileName <- generate $ vectorOf 10 $ elements ['A' .. 'Z']
  let cFilePath = "C" </> fileName <.> "c"
      executablePath = "C" </> fileName
      program = singleExpressionCProgram valMaps exp
      libs
        | withFT = ["-lm", "-lfftw3"]
        | otherwise = ["-lm"]
  TIO.writeFile cFilePath (T.intercalate "\n" program)
  readProcess "gcc" ([cFilePath, "-o", executablePath] ++ libs) ""
  (exitCode, output, _) <- readProcessWithExitCode executablePath [] ""
  readProcess "rm" [cFilePath] ""
  readProcess "rm" [executablePath] ""
  return (exitCode, output)

-- | Parse output of the C program
readR :: String -> [Double]
readR = map read . filter (not . null) . splitOn " "

-- |
readC :: String -> ([Double], [Double])
readC str = (readR rePart, readR imPart)
  where
    [rePart, imPart] = splitOn "\n" str

-- |
prop_CEqualInterpScalarR :: SuiteScalarR -> Expectation
prop_CEqualInterpScalarR (Suite exp valMaps) =
  if containsFTNode $ exMap exp
    then do
      hasFTLib <- hasFFTW
      when hasFTLib $ proceed True
    else proceed False
  where
    proceed withFT = do
      let normalizedExp = normalize exp
      (exitCode, outputSimple) <- evaluateCodeC withFT normalizedExp valMaps
      let resultNormalize = read . head . splitOn " " $ outputSimple
      let resultInterpNormalize = eval valMaps normalizedExp
      resultNormalize `shouldApprox` resultInterpNormalize

-- |
prop_CEqualInterpScalarC :: SuiteScalarC -> Expectation
prop_CEqualInterpScalarC (Suite exp valMaps) =
  if containsFTNode $ exMap exp
    then do
      hasFTLib <- hasFFTW
      when hasFTLib $ proceed True
    else proceed False
  where
    proceed withFT = do
      let normalizedExp = normalize exp
      (exitCode, outputCodeC) <- evaluateCodeC withFT normalizedExp valMaps
      let ([im], [re]) = readC outputCodeC
      let resultNormalize = im :+ re
      let resultInterpNormalize = eval valMaps normalizedExp
      resultNormalize `shouldApprox` resultInterpNormalize

-- |
prop_CEqualInterpOneR :: SuiteOneR -> Expectation
prop_CEqualInterpOneR (Suite exp valMaps) =
  if containsFTNode $ exMap exp
    then do
      hasFTLib <- hasFFTW
      when hasFTLib $ proceed True
    else proceed False
  where
    proceed withFT = do
      let normalizedExp = normalize exp
      (exitCode, outputSimple) <- evaluateCodeC withFT normalizedExp valMaps
      let resultNormalize = listArray (0, defaultDim1D - 1) $ readR outputSimple
      let resultInterpNormalize = eval valMaps normalizedExp
      resultNormalize `shouldApprox` resultInterpNormalize

-- |
prop_CEqualInterpOneC :: SuiteOneC -> Expectation
prop_CEqualInterpOneC (Suite exp valMaps) =
  if containsFTNode $ exMap exp
    then do
      hasFTLib <- hasFFTW
      when hasFTLib $ proceed True
    else proceed False
  where
    proceed withFT = do
      let normalizedExp = normalize exp
      (exitCode, outputCodeC) <- evaluateCodeC withFT normalizedExp valMaps
      let (re, im) = readC outputCodeC
          resultNormalize = listArray (0, defaultDim1D - 1) $ zipWith (:+) re im
          resultInterpNormalize = eval valMaps normalizedExp
      resultNormalize `shouldApprox` resultInterpNormalize

-- |
prop_CEqualInterpTwoR :: SuiteTwoR -> Expectation
prop_CEqualInterpTwoR (Suite exp valMaps) =
  if containsFTNode $ exMap exp
    then do
      hasFTLib <- hasFFTW
      when hasFTLib $ proceed True
    else proceed False
  where
    proceed withFT = do
      let normalizedExp = normalize exp
      (exitCode, outputSimple) <- evaluateCodeC withFT normalizedExp valMaps
      let resultNormalize = listArray ((0, 0), (default1stDim2D - 1, default2ndDim2D - 1)) $ readR outputSimple
      let resultInterpNormalize = eval valMaps normalizedExp
      resultNormalize `shouldApprox` resultInterpNormalize

-- |
prop_CEqualInterpTwoC :: SuiteTwoC -> Expectation
prop_CEqualInterpTwoC (Suite exp valMaps) =
  if containsFTNode $ exMap exp
    then do
      hasFTLib <- hasFFTW
      when hasFTLib $ proceed True
    else proceed False
  where
    proceed withFT = do
      let normalizedExp = normalize exp
      (exitCode, outputCodeC) <- evaluateCodeC withFT (normalize exp) valMaps
      let (re, im) = readC outputCodeC
      let resultNormalize = listArray ((0, 0), (default1stDim2D - 1, default2ndDim2D - 1)) $ zipWith (:+) re im
      let resultInterpNormalize = eval valMaps normalizedExp
      resultNormalize `shouldApprox` resultInterpNormalize

-- | Spec
spec :: Spec
spec =
  describe "C simple spec" $ do
    specify
      "Evaluate hash interp should equal to C code evaluation (Expression Scalar R)"
      $ property prop_CEqualInterpScalarR
    specify
      "Evaluate hash interp should equal to C code evaluation (Expression Scalar C)"
      $ property prop_CEqualInterpScalarC
    specify
      "Evaluate hash interp should equal to C code evaluation (Expression One R)"
      $ property prop_CEqualInterpOneR
    specify
      "Evaluate hash interp should equal to C code evaluation (Expression One C)"
      $ property prop_CEqualInterpOneC
    specify
      "Evaluate hash interp should equal to C code evaluation (Expression Two R)"
      $ property prop_CEqualInterpTwoR
    specify
      "Evaluate hash interp should equal to C code evaluation (Expression Two C)"
      $ property prop_CEqualInterpTwoC


{-# LANGUAGE GADTs #-}
{-# LANGUAGE FlexibleContexts #-}

module HVX.Primitives
-- TODO(mh): Currently primitives that should generate implicit constraints are
-- unsupported. The should be supported shortily. (2014-06-04)
  ( apply
  , hadd
  , (+~)
  , hmul
  , (*~)
  , habs
  , neg
  , hlog
  , hexp
  , logsumexp
  , hmax
  , hmin
  , norm
  , berhu
  , huber
  , quadform
  , powBaseP0
  , powBaseP01
  , powBaseP1
  , powBaseP1InfEven
  , powBaseP1InfNotInt
  ) where

import Numeric.LinearAlgebra hiding (i)
import Numeric.LinearAlgebra.LAPACK

import HVX.Internal.DCP
import HVX.Internal.Matrix
import HVX.Internal.Primitives
import HVX.Internal.Util

apply :: (Vex vf, Mon mf, Vex ve, Mon me, Vex (ApplyVex vf mf ve me), Mon (ApplyMon mf me))
  => Fun vf mf -> Expr ve me -> Expr (ApplyVex vf mf ve me) (ApplyMon mf me)
apply = EFun

hadd :: (Vex v1, Mon m1, Vex v2, Mon m2, Vex (AddVex v1 v2), Mon (AddMon m1 m2))
  => Expr v1 m1 -> Expr v2 m2 -> Expr (AddVex v1 v2) (AddMon m1 m2)
hadd = EAdd

infixl 6 +~
(+~) :: (Vex v1, Mon m1, Vex v2, Mon m2, Vex (AddVex v1 v2), Mon (AddMon m1 m2))
  => Expr v1 m1 -> Expr v2 m2 -> Expr (AddVex v1 v2) (AddMon m1 m2)
(+~) = hadd

-- Constructors that enforce DCP constraints.
hmul :: (Vex v1, Mon m1, Vex (ApplyVex Affine Nonmon v1 m1), Mon (ApplyMon Nonmon m1))
  => Expr Affine Const -> Expr v1 m1 -> Expr (ApplyVex Affine Nonmon v1 m1) (ApplyMon Nonmon m1)
hmul (EConst a) e = apply (Mul a) e
hmul _ _ = error "the left argument of a multiply must be a constant"

infixl 7 *~
(*~) :: (Vex v1, Mon m1, Vex (ApplyVex Affine Nonmon v1 m1), Mon (ApplyMon Nonmon m1))
  => Expr Affine Const -> Expr v1 m1 -> Expr (ApplyVex Affine Nonmon v1 m1) (ApplyMon Nonmon m1)
(*~) (EConst a) e = apply (Mul a) e
(*~) _ _ = error "the left argument of a multiply must be a constant"

habs :: (Vex v1, Mon m1, Vex (ApplyVex Convex Nonmon v1 m1), Mon (ApplyMon Nonmon m1))
  => Expr v1 m1 -> Expr (ApplyVex Convex Nonmon v1 m1) (ApplyMon Nonmon m1)
habs = apply Abs

neg :: (Vex v1, Mon m1, Vex (ApplyVex Affine Noninc v1 m1), Mon (ApplyMon Noninc m1))
  => Expr v1 m1 -> Expr (ApplyVex Affine Noninc v1 m1) (ApplyMon Noninc m1)
neg = apply Neg

hlog :: (Vex v1, Mon m1, Vex (ApplyVex Concave Nondec v1 m1), Mon (ApplyMon Nondec m1))
  => Expr v1 m1 -> Expr (ApplyVex Concave Nondec v1 m1) (ApplyMon Nondec m1)
hlog = apply Log

hexp :: (Vex v1, Mon m1, Vex (ApplyVex Convex Nondec v1 m1), Mon (ApplyMon Nondec m1))
  => Expr v1 m1 -> Expr (ApplyVex Convex Nondec v1 m1) (ApplyMon Nondec m1)
hexp = apply Exp

logsumexp :: (Vex v1, Mon m1, Vex (ApplyVex Convex Nondec v1 m1), Mon (ApplyMon Nondec m1))
  => Expr v1 m1 -> Expr (ApplyVex Convex Nondec v1 m1) (ApplyMon Nondec m1)
logsumexp = apply LogSumExp

hmax :: (Vex v1, Mon m1, Vex (ApplyVex Convex Nondec v1 m1), Mon (ApplyMon Nondec m1))
  => Expr v1 m1 -> Expr (ApplyVex Convex Nondec v1 m1) (ApplyMon Nondec m1)
hmax = apply Max

hmin :: (Vex v1, Mon m1, Vex (ApplyVex Concave Nondec v1 m1), Mon (ApplyMon Nondec m1))
  => Expr v1 m1 -> Expr (ApplyVex Concave Nondec v1 m1) (ApplyMon Nondec m1)
hmin = apply Min

norm :: (Vex v1, Mon m1, Vex (ApplyVex Convex Nonmon v1 m1), Mon (ApplyMon Nonmon m1))
  => Double -> Expr v1 m1 -> Expr (ApplyVex Convex Nonmon v1 m1) (ApplyMon Nonmon m1)
norm p
  | p == infinity = error "Internal: Infinity norm should become max . abs."
  | 1 <= p = apply (Norm p)
  | otherwise = error "Internal: Norm only supports p >= 1."

berhu :: (Vex v1, Mon m1, Vex (ApplyVex Convex Nonmon v1 m1), Mon (ApplyMon Nonmon m1))
  => Double -> Expr v1 m1 -> Expr (ApplyVex Convex Nonmon v1 m1) (ApplyMon Nonmon m1)
berhu m
  | 0 < m = apply (Berhu m)
  | otherwise = error "Internal: Berhu only supports m >= 0."

huber :: (Vex v1, Mon m1, Vex (ApplyVex Convex Nonmon v1 m1), Mon (ApplyMon Nonmon m1))
  => Double -> Expr v1 m1 -> Expr (ApplyVex Convex Nonmon v1 m1) (ApplyMon Nonmon m1)
huber m
  | 0 < m = apply (Huber m)
  | otherwise = error "Internal: Huber only supports m >= 0."

quadform :: (Mon m1, Vex (ApplyVex Convex Nonmon Affine m1), Mon (ApplyMon Nonmon m1))
  => Expr Affine Const -> Expr Affine m1 -> Expr (ApplyVex Convex Nonmon Affine m1) (ApplyMon Nonmon m1)
quadform (EConst a) e
  | rows a == cols a
    && fpequalsMat a (trans a)
    && 0 <= maxElement (eigOnlyS a) = apply (Quadform a) e
  | otherwise = error "Matrices in quadratic forms must be positive semidefinite."
quadform _ _ = error "The matrix sandwitched by the quadratic form must be constant."

powBaseP0 :: (Vex v1, Mon m1, Vex (ApplyVex Affine Const v1 m1), Mon (ApplyMon Const m1))
  => Double -> Expr v1 m1 -> Expr (ApplyVex Affine Const v1 m1) (ApplyMon Const m1)
powBaseP0 p
  | p `fpequals` 0 = apply PowBaseP0
  | otherwise = error "Internal: PowBaseP0 only supports p == 0."

powBaseP01 :: (Vex v1, Mon m1, Vex (ApplyVex Concave Nondec v1 m1), Mon (ApplyMon Nondec m1))
  => Double -> Expr v1 m1 -> Expr (ApplyVex Concave Nondec v1 m1) (ApplyMon Nondec m1)
powBaseP01 p
  | 0 < p && p < 1 = apply (PowBaseP01 p)
  | otherwise = error "Internal: PowBaseP01 only supports 0 < p < 1."

powBaseP1 :: (Vex v1, Mon m1, Vex (ApplyVex Affine Nondec v1 m1), Mon (ApplyMon Nondec m1))
  => Double -> Expr v1 m1 -> Expr (ApplyVex Affine Nondec v1 m1) (ApplyMon Nondec m1)
powBaseP1 p
  | p `fpequals` 1 = apply PowBaseP1
  | otherwise = error "Internal: PowBaseP1 only supports p == 1."

powBaseP1InfEven :: (Vex v1, Mon m1, Vex (ApplyVex Convex Nonmon v1 m1), Mon (ApplyMon Nonmon m1))
  => Double -> Expr v1 m1 -> Expr (ApplyVex Convex Nonmon v1 m1) (ApplyMon Nonmon m1)
powBaseP1InfEven p
  | 1 < p && isInteger p && even intP = apply (PowBaseP1InfEven intP)
  | otherwise = error "Internal: PowBaseP1InfEven only supports even p > 1."
  where intP = round p :: Integer

powBaseP1InfNotInt :: (Vex v1, Mon m1, Vex (ApplyVex Convex Nondec v1 m1), Mon (ApplyMon Nondec m1))
  => Double -> Expr v1 m1 -> Expr (ApplyVex Convex Nondec v1 m1) (ApplyMon Nondec m1)
powBaseP1InfNotInt p
  | 1 < p && not (isInteger p) = apply (PowBaseP1InfNotInt p)
  | otherwise = error "Internal: PowBaseP1InfNotInt only supports non integral p > 1."


{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE NoMonomorphismRestriction #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}

module TicTacToe.AiInterface where

import Data.Maybe
import Control.Monad.State hiding (state)
import Control.Monad.Reader
import Control.Exception.Base (assert)
import System.Random (StdGen, split, randomR, Random)
import System.Random.Shuffle (shuffle')

import Numeric.LinearAlgebra

import TicTacToe.Core (move, result)
import TicTacToe.Domain (Result(..), Board, Move, CellPos)

class Player a where
  predictedAction :: BoardMatrix -> a -> Action

data GameState = GameState {
  reward' :: Int,
  done' :: Bool,
  board :: BoardMatrix
}

type BoardMatrix = Matrix Z
type Action = (Int, Int)

--step :: (Player a) => BoardMatrix -> Action -> a -> GameState
--step m action opponent =
--  case aiResult m action of
--    Error   -> err
--    Win     -> win
--    Draw    -> draw
--    _       -> opponentMoveToState (fromJust (aiMove' m action)) opponent
--
--opponentMoveToState :: (Player a) => BoardMatrix -> a -> GameState
--opponentMoveToState m opponent =
--  let pAction = predictedAction m opponent
--  in
--    case aiResult m pAction of
--      Win         -> loss
--      Draw        -> draw
--      Unfinished  -> unfinished m pAction
--
--aiMove' :: BoardMatrix -> Action -> Maybe BoardMatrix
--aiMove' matrix action =
--  let board           = matrixToBoard matrix
--      (pl, cellPos)   = actionToPlayerCellPos action
--  in
--    case move board pl cellPos of
--      Just b  -> Just $ boardToMatrix b
--      _       -> Nothing
--
--aiResult :: BoardMatrix -> Action -> Result
--aiResult matrix action =
--  result board pl cellPos
--  where
--    board           = matrixToBoard matrix
--    (pl, cellPos)   = actionToPlayerCellPos action
--
--actionToPlayerCellPos :: Action -> (Move, CellPos)
--actionToPlayerCellPos = assert False undefined
--
--matrixToBoard :: BoardMatrix -> Board
--matrixToBoard = assert False undefined
--
--boardToMatrix :: Board -> BoardMatrix
--boardToMatrix = assert False undefined
--
--matrixToState :: BoardMatrix -> GameState
--matrixToState = assert False undefined
--
--emptyGameState :: GameState
--emptyGameState = assert False undefined
--
--err :: GameState
--err = assert False undefined
--
--loss :: GameState
--loss = assert False undefined
--
--win :: GameState
--win = assert False undefined
--
--draw :: GameState
--draw = assert False undefined
--
--unfinished :: BoardMatrix -> Action -> GameState
--unfinished m a = assert False undefined

-- NN stuff

type Input  = Vector R
type W1     = Matrix R
type W2     = Matrix R
type W3     = Matrix R
type Output = Vector R
type B      = Matrix R

type Hidden1  = Matrix R
type Hidden2  = Matrix R

type Reward = R

type ActionIndex = Int

data Model = Model {
  w1        :: W1,
  w2        :: W2,
  w3        :: W3,
  b1        :: B,
  b2        :: B,
  b3        :: B
}

newtype Target = Target {
  unTarget :: Model
}

data NeuralNetwork = NN {
  epsilon       :: Double,
  learningRate  :: Double,
  model         :: Model,
  targetModel   :: Target,
  input         :: Input,
  buffer        :: [BufferData],
  adamIter      :: Int,
  adamW1        :: AdamState,
  adamW2        :: AdamState,
  adamW3        :: AdamState,
  adamB1        :: AdamState,
  adamB2        :: AdamState,
  adamB3        :: AdamState,
  randomGen     :: StdGen
}

data AdamState = AdamState {
  mass      :: Matrix R,
  velocity  :: Matrix R
}

data BufferData = BufferD {
  state     :: Input,
  action    :: ActionIndex,
  reward    :: Reward,
  nextState :: Input,
  done      :: Bool
}

data HyperParameters = Hyp {
  player            :: Move,
  epsilonDecay      :: Double,
  epsilonMin        :: Double,
  gamma             :: Double,
  learningRateDecay :: Double,
  episodes          :: Int,
  batchSize         :: Int,
  bufferSize        :: Int,
  stateSize         :: Int,
  adamP             :: AdamParameters,
  winReward         :: Reward,
  drawReward        :: Reward,
  lossReward        :: Reward,
  unfinishedReward  :: Reward
}

data AdamParameters = AdamParameters {
  eps       :: Double,
  beta1     :: Double,
  beta2     :: Double
}

type Network' a = ReaderT HyperParameters (State NeuralNetwork) a

newtype Network a = Network {
  unNetwork :: Network' a
} deriving (Monad, Applicative, Functor, MonadReader HyperParameters,
  MonadState NeuralNetwork)

learn1 :: Network ()
learn1 = do
  state <- gets input
  actionIndex <- act state
  (nextState, reward, done) <- aiMove actionIndex
  return ()

aiMove :: ActionIndex -> Network (Input, Reward, Bool)
aiMove action = do
  currentState <- gets input
  let stateAfterAction = update currentState action
  if (win stateAfterAction) then do
      reward <- asks winReward
      return (stateAfterAction, reward, True)
    else do
      opponentAction <- act $ reverseState $ stateAfterAction
      let stateAfterOpponent = update stateAfterAction opponentAction
      tmp stateAfterOpponent
        where
          tmp :: Input -> Network (Input, Reward, Bool)
          tmp state
            | loss state  = do
              reward <- asks lossReward
              return (state, reward, True)
            | draw state  = do
              reward <- asks drawReward
              return (state, reward, True)
            | otherwise   = do
              reward <- asks unfinishedReward
              return (state, reward, False)

update :: Input -> ActionIndex -> Input
update = assert False undefined

reverseState :: Input -> Input
reverseState = assert False undefined

win :: Input -> Bool
win = assert False undefined

loss :: Input -> Bool
loss = assert False undefined

draw :: Input -> Bool
draw = assert False undefined

replay :: Network ()
replay = do
  epsMin  <- asks epsilonMin
  epsD    <- asks epsilonDecay
  bSize   <- asks batchSize
  buffer  <- gets buffer
  sample  <- randomSample bSize buffer
  forM_ sample (\bufferD -> do
    target <- calcTarget bufferD
    model <- gets model
    model' <- fit model (state bufferD) target
    s <- get
    put s{ model=model' }
    s@NN{epsilon = eps} <- get
    when (eps > epsMin) $ put s{epsilon = eps * epsD})

calcTarget :: BufferData -> Network Output
calcTarget (BufferD state action reward nextState done) = do
  model <- gets model
  targetModel <- gets targetModel
  gamma <- asks gamma
  predCurRew <- predict model state
  if done then
      return $ fromList $ replace (toList predCurRew) action reward
    else do
      predFutRew <- predict model nextState
      targetPredFutRew <- predict (unTarget targetModel) nextState
      return $ fromList $ replace (toList predCurRew) action $ reward + gamma * targetPredFutRew ! maxIndex predFutRew
    where
      replace xs index value =
        let (before, _:after) = splitAt (index - 1) xs
        in before ++ (value:after)

fit :: Model -> Input -> Output -> Network Model
fit m@(Model w1 w2 w3 b1 b2 b3) state target = do
  let state' = asColumn state -- [9x1]
      target' = asColumn target -- [9x1]
      hidden1 = relu $ b1 + (w1 <> state') -- [24x9] x [9x1] = [24x1]
      hidden2 = relu $ add b2 $ w2 <> hidden1 -- [24x24] x [24x1] = [24x1]
      prediction = add b3 $ w3 <> hidden2 -- [9x24] x [24x1] = [9x1]
      loss = meanSquaredError target' prediction -- [9x1]
      g1 = w1 <> state' + b1 -- [24x9] x [9x1] = [24x1]
      g2 = w2 <> hidden1 + b2 -- [24x24] x [24x1] = [24x1]
      dg3 = - (target' - prediction) -- [9x1]
      dg2 = drelu g1 * (tr' dg3 <> w3) -- [1x9] x [9x24] = [1x24]
      dg1 = drelu g2 * dg2 <> w2 -- [1x24] x [24x24] = [1x24]
      dw3 = dg3 <> tr' hidden2 -- [9x1] x [1x24] = [9x24]
      db3 = dg3 -- [9x1]
      dw2 = tr' dg2 <> tr' hidden1 -- [24x1] x [1x24] = [24x24]
      db2 = tr' dg2 -- [24x1]
      dw1 = tr' dg1 <> tr' state' -- [24x1] x [1x9] = [24x9]
      db1 = tr' dg1 -- [24x1]
      dm  = Model dw1 dw2 dw3 b1 b2 b3
      drelu g = fromLists $ map (map cutOff) $ toLists g
      cutOff el = if el <= 0 then 0 else 1
  adam m dm

adam :: Model -> Model -> Network Model
adam (Model w1 w2 w3 b1 b2 b3) (Model dw1 dw2 dw3 db1 db2 db3) = do
  adamP   <- asks adamP
  t       <- gets adamIter
  learningRate  <- gets learningRate
  w1A <- gets adamW1
  w2A <- gets adamW2
  w3A <- gets adamW3
  b1A <- gets adamB1
  b2A <- gets adamB2
  b3A <- gets adamB3
  let (w1', w1A') = adam' w1 dw1 w1A adamP t learningRate
      (w2', w2A') = adam' w2 dw2 w2A adamP t learningRate
      (w3', w3A') = adam' w3 dw3 w3A adamP t learningRate
      (b1', b1A') = adam' b1 db1 b1A adamP t learningRate
      (b2', b2A') = adam' b2 db2 b2A adamP t learningRate
      (b3', b3A') = adam' b3 db3 b3A adamP t learningRate
      t'          = t + 1
      model'      = Model w1 w2 w3 b1 b2 b3
  state   <- get
  put state{
    adamW1  = w1A',
    adamW2  = w2A',
    adamW3  = w3A',
    adamB1  = b1A',
    adamB2  = b2A',
    adamB3  = b3A',
    adamIter  = t'
    }
  return model'

adam' :: Matrix R -> Matrix R -> AdamState -> AdamParameters -> Int -> Double
  -> (Matrix R, AdamState)
adam' x dx (AdamState m v) (AdamParameters eps beta1 beta2) t learningRate =
  let m'  = scale beta1 m + scale (1 - beta1) dx
      mt  = m' / scalar (1 - beta1^t)
      v'  = scale beta2 v + scale (1 - beta2) (dx^2)
      vt  = v' / scalar (1 - beta2^t)
      x'  = x + scale (- learningRate) mt / (sqrtm vt + scalar eps)
  in  (x', AdamState m' v')

meanSquaredError :: Floating a => a -> a -> a
meanSquaredError target prediction =
  0.5 * ((target - prediction)**2)

act :: Input -> Network ActionIndex
act state = do
  eps <- gets epsilon
  (rNumber :: Double) <- randomNumber (0,1)
  if rNumber <= eps then
      randomElement $ possibleActions state
    else do
      model <- gets model
      maxIndex <$> predict model state

possibleActions :: Input -> [ActionIndex]
possibleActions =
  find (\el -> el > (- 0.5) && el < 0.5)

predict :: Model -> Input -> Network Output
predict (Model w1 w2 w3 b1 b2 b3) input = do
  let input' = asColumn input -- [9x1]
      hidden1 = relu $ add b1 $ w1 <> input' -- [24x9] x [9x1] = [24x1]
      hidden2 = relu $ add b2 $ w2 <> hidden1 -- [24x24] x [24x1] = [24x1]
  return (flatten $ add b3 $ w3 <> hidden2) -- [9x24] x [24x1] = [9x1]

remember :: BufferData -> Network ()
remember bData = do
  s@NN{ buffer=buffer }   <- get
  bSize                   <- asks bufferSize
  if length buffer < bSize
    then put s{ buffer = bData : buffer }
    else put s{ buffer = bData : init buffer }

relu :: Matrix R -> Matrix R
relu =
  fromLists . map (map (max 0)) . toLists

initHyp :: HyperParameters
initHyp = assert False undefined

initNN :: Reader HyperParameters NeuralNetwork
initNN = assert False undefined

randomSample :: Int -> [a] -> Network [a]
randomSample n xs =
  take n <$> randomShuffle xs

randomShuffle :: [a] -> Network [a]
randomShuffle xs = do
  s@NN{ randomGen=gen }  <- get
  let shuffled  = shuffle' xs (length xs) gen
      (gen', _) = split gen
  put s{ randomGen=gen' }
  return shuffled

randomElement :: [a] -> Network a
randomElement xs = do
  randomIndex :: Int <- randomNumber (0, length xs)
  return $ xs !! randomIndex

randomNumber :: (Random a) => (a, a) -> Network a
randomNumber (beg, end) = do
  s@NN{ randomGen=gen }  <- get
  let (rNumber, gen') = randomR (beg, end) gen
  put s{randomGen=gen'}
  return rNumber


{-# LANGUAGE ExtendedDefaultRules #-}
{-# LANGUAGE OverloadedStrings    #-}
{-# LANGUAGE FlexibleContexts     #-}
{-# LANGUAGE ScopedTypeVariables  #-}
{-# LANGUAGE BangPatterns         #-}

module TestChapter3MDP (
       testChapter3
      )where

import           Control.Monad
import           Control.Monad.IO.Class
import           Control.Monad.Trans.State

import           Data.Configurator
import           Data.Configurator.Types
import           Data.Text (Text)
import           Numeric.LinearAlgebra (Vector, Matrix)
import qualified Numeric.LinearAlgebra as LA
import           System.Console.AsciiProgress(Options(..), Stats(..),
                                              displayConsoleRegions, complete,
                                              getProgressStats, def, newProgressBar, tickN)
-- project
import           Utils
import           Chapter3MDP

testChapter3 :: FilePath -> IO ()
testChapter3 configPath = do
  print "Chapter 3 Experiment Starting "
  (config, _) <- autoReload autoConfig [Required configPath]
  (bGridWorld :: Bool) <- require config "enable.bGridWorld"
  when bGridWorld (doGridWorldTest config)

doGridWorldTest :: Config -> IO ()
doGridWorldTest config = do
  (size::Int) <- require config "gridWorld.worldSize"
  (policies::[String]) <- require config "gridWorld.policy"
  (discountGamma::Double) <- require config "gridWorld.discount"
  (learningAccuracy::Double) <- require config "gridWorld.learningAccuracy"
  (specialPositions::[[Int]]) <- require config "gridWorld.specialPositions"
  (specialTransitPos::[[Int]]) <- require config "gridWorld.specialTransitPos"
  (specialRewards::[Double]) <- require config "gridWorld.specialRewards"
  let specials = zipWith3 (\ src dst r -> ((src!!0, src!!1), (dst!!0, dst!!1), r))
                                          specialPositions specialTransitPos specialRewards
  -- do experiments
  mapM_ (goOnePolicyTest size discountGamma learningAccuracy specials) policies
  where
  goOnePolicyTest size discountGamma learningAccuracy specials aPolicy = displayConsoleRegions $ do
    print ("Will do experiment with policy " ++ aPolicy)
    putStrLn ("Special positions: " ++ show specials)
    let thePolicy = read aPolicy
        world = createWorld size thePolicy discountGamma specials
    pg <- newProgressBar def { pgWidth = 80
                             , pgTotal = 100
                             , pgOnCompletion = Just "Done :percent after :elapsed seconds"
                             }
    loop pg world 1
    print "finish chapter3 experiments"
    where
    loop pg world count = do
      let (diff, w') = runState step world
      when (count `mod` 10 == 0)
           (do
            print ("Iteration " ++ show count)
            putStr $ showStateValues (_maxSize w') (_stateValues w')
           )
      case diff < learningAccuracy of
         False -> do
           stat <- getProgressStats pg
           let ticked = stCompleted stat
               percent = floor (100.0 / (diff / learningAccuracy))
           when (ticked < percent) (tickN pg (fromInteger $ percent - ticked))
           loop pg w' (count + 1)
         True -> complete pg >> (putStr $ showStateValues (_maxSize w') (_stateValues w'))


module Main where

import           Control.Monad.ST               ( runST )
import qualified Data.Vector                   as V
import           Statistics.BBVI.Examples

main :: IO ()
main =
  let xs = runST $ genMixedMem
  in  (putStrLn . unlines . fmap show . take 100 . V.toList) xs
        >> putStrLn (show . snd $ mixedMemFit xs)


{-# options_ghc -Wwarn #-}
{-# language TypeOperators #-}
{-# language FlexibleInstances #-}
module Main where

import Solitaire.Yukon

import System.Directory
import System.Environment
import System.Microtimer

import qualified Data.Vector as V

import Statistics.Sample
import Data.List (sort)
import Options.Generic
import LinearRegression

data CliArgs w = CliArgs
  { numTrials  :: w ::: Int <!> "10"  <?> "Number of trials to run"
  , limitGames :: w ::: Int <!> "100" <?> "Limit number of games to test"
  , step       :: w ::: Float <!> "0.1" <?> "Step size"
  }
  deriving stock (Generic)

deriving instance Show (CliArgs Unwrapped)
instance ParseRecord (CliArgs Wrapped)

main :: IO ()
main = do
  args :: CliArgs Unwrapped <- unwrapRecord "Optimize heuristics"

  let
    limitGames = args ^. #limitGames
    numTrials = args ^. #numTrials

  fileNames <- take limitGames <$> listDirectory "games/Solvable"

  games <- for fileNames $ \name -> do
    fileContents <- readFile $ "games/Solvable/" <> name
    pure $ read fileContents

  let
    optimizeConfig = OptimizeConfig { numTrials, games }

    initialWeights = mempty
      & at "numFaceUp" ?~ 2
      & at "numFaceDown" ?~ 5
      & at "totalRunScore" ?~ (-0.5)

    test = testWeights optimizeConfig

  runTime <- test initialWeights

  prettyPrint runTime
  pure ()

data PartialDArgs = PartialDArgs
  { step :: Float
  , field :: Text
  , runTime :: RunTime
  , weights :: Weights
  , function :: Weights -> IO RunTime
  }
  deriving stock (Generic)

partialDerivative :: PartialDArgs -> IO RunTime
partialDerivative args = do
  let
    w_1 :: Weights
    w_1 = args ^. #weights & ix (args ^. #field) -~ (args ^. #step)

    w_2 :: Weights
    w_2 = args ^. #weights & ix (args ^. #field) +~ (args ^. #step)

    r0 = args ^. #runTime

  r1 <- args ^. #function $ w_1
  r2 <- args ^. #function $ w_2

  pure $ undefined

data Uncertain a = a :+- a

data Line = Line
  { m :: Uncertain Float
  , b :: Uncertain Float
  }
  deriving stock (Generic)

linearRegression :: [(Weight, RunTime)] -> Line
linearRegression = undefined

data OptimizeConfig = OptimizeConfig
  { numTrials :: !Int
  , games :: ![Game]
  }
  deriving stock (Generic)

data RunTime = RunTime
  { mean :: !Double
  , standardDeviation :: !Double
  }
  deriving stock (Generic)

instance Pretty RunTime where
  prettyExpr run =
    prettyExpr $ (Empty :: Map Text Text)
      & at "mean" ?~ (run ^. #mean . to formatSeconds . packed)
      & at "std dev" ?~ (run ^. #standardDeviation . to formatSeconds . packed)

testWeights :: OptimizeConfig -> Weights -> IO RunTime
testWeights config weights = do
  let
    games = config ^. #games
    numTrials = config ^. #numTrials

  times <- for games $ \game -> do
    let
      run = runGame @Yukon game
        & (runReaderT ?? weights)
        & time_

    sequenceA $ replicate numTrials run

  let
    sortedData = join times & sort & V.fromList
    (mean, variance) = meanVariance sortedData
    standardDeviation = sqrt variance

  pure $ RunTime { mean, standardDeviation }

updateWeightsAtRandom
  :: forall k m. (Ord k, MonadRandom m)
  => Weights
  -> m Weights
updateWeightsAtRandom map = do
  let
    n = length map
    keys = (map ^.. ifolded . asIndex)

  delta <- randomDelta n

  pure $ flip execState map $
    ifor_ keys $ \i key -> do
      ix key += (delta ^?! ix i)


module Numeric.FFT.Plan ( plan, planFromFactors ) where

import qualified Control.Monad          as CM
import           Control.Monad.IO.Class
import           Criterion.Measurement              (measure)
import           Criterion.Measurement.Types hiding (measure)
import           Data.Complex
import           Data.Function          (on)
import qualified Data.IntMap.Strict     as IM
import           Data.List              ((\\))
import qualified Data.List              as L
import qualified Data.Set               as S
import qualified Data.Vector            as V
import           Data.Vector.Unboxed
import           Prelude                hiding (any, concatMap, enumFromTo,
                                         filter, length, map, maximum, null,
                                         reverse, scanl, sum, zip, zipWith,
                                         (++))
import qualified Prelude                as P
import           System.Directory
import           System.Environment
import           System.FilePath

import           Numeric.FFT.Execute
import           Numeric.FFT.Special
import           Numeric.FFT.Types
import           Numeric.FFT.Utils

-- | Number of plans to test empirically.
nTestPlans :: Int
nTestPlans = 50

-- | Plan calculation for a given problem size.
plan :: Int -> IO Plan
plan 1 = return $ Plan V.empty Nothing (SpecialBase 1)
plan n = do
  wis <- readWisdom n
  let fixRader p = case plBase p of
        bpl@(RaderBase _ _ _ _ csz _) -> do
          cplan <- liftIO $ plan csz
          return $ p { plBase = bpl { raderConvPlan = cplan } }
        _ -> return p
  pret <- case wis of
    Just (p, _) -> planFromFactors n p
    Nothing -> do
      let ps = testPlans n nTestPlans
          v = generate n (\i -> sin (2 * pi * fromIntegral i / 511) :+ 0)
      tps <- CM.forM ps $ \p -> do
        ptest <- liftIO $ planFromFactors n p >>= fixRader
        let niters = fromIntegral $ 50000 `div` n
        (meas, _) <- measure (nf (execute ptest Forward) v) niters
        let t = measTime meas / fromIntegral (measIters meas)
            _iters = measIters meas
        return (t, p)
      let (rest, resp) = L.minimumBy (compare `on` fst) tps
      liftIO $ writeWisdom n resp rest
      liftIO $ planFromFactors n resp
  fixRader pret

-- | Get execution time for plan for a given problem size.
planTime :: Int -> IO Double
planTime n = do
  wis <- readWisdom n
  case wis of
    Just (_, t) -> return t
    Nothing -> do
      -- Force generation of wisdom if it's not already there.
      _ <- plan n
      Just (_, t) <- readWisdom n
      return t

-- | Plan calculation for a given problem factorisation.
planFromFactors :: Int -> (Int, Vector Int) -> IO Plan
planFromFactors n (lastf, fs) = do
  -- Base transform.
  (base, mextraperm) <- makeBase lastf

  -- Include permutation of base transform if needed.
  let perm = case (digperm, mextraperm) of
        (Just dp, Just ep) -> Just $ dupperm n ep %.% dp
        (Nothing, Just ep) -> Just $ dupperm n ep
        (Just dp, Nothing) -> Just dp
        (Nothing, Nothing) -> Nothing

  return $ Plan dlinfo perm base
  where
    -- Input data "digit reversal" permutation.
    digperm = digrev n fs

    -- Size information for Danielson-Lanczos steps.
    wfacs = map (n `div`) $ scanl (*) 1 fs
    vwfacs = convert wfacs
    vfs = convert fs
    dmatps = V.zipWith (dmat 1) vwfacs vfs
    dmatms = V.zipWith (dmat (-1)) vwfacs vfs
    dlinfo = V.reverse $ V.zip4 vwfacs vfs dmatps dmatms

    -- Calculate diagonal matrix entries used in Danielson-Lanczos steps.
    dmat :: Int -> Int -> Int -> VVVCD
    dmat sign wfac split =
      let ns = wfac `div` split
          w = omega $ sign * wfac
      in V.generate split $
         \r -> V.generate split $
           \c -> map ((w^(ns*r*c) *) . ((w^^) . (c *))) $ enumFromN 0 ns

-- | Read from wisdom for a given problem size.
readWisdom :: Int -> IO (Maybe ((Int, Vector Int), Double))
readWisdom n = do
  home <- getEnv "HOME"
  let wisf = home </> ".fft-plan" </> show n
  ex <- doesFileExist wisf
  if ex then
    (do wist <- readFile wisf
        let ((wisb, wisfs), wistim) = read wist :: ((Int, [Int]), Double)
        return $ Just ((wisb, fromList wisfs), wistim))
    else return Nothing

-- | Write wisdom for a given problem size.
writeWisdom :: Int -> (Int, Vector Int) -> Double -> IO ()
writeWisdom n (b, fs) tim = do
  home <- getEnv "HOME"
  let wisd = home </> ".fft-plan"
      wisf = wisd </> show n
  createDirectoryIfMissing True wisd
  writeFile wisf $ show ((b, toList fs), tim) P.++ "\n"

-- | Make base transform for a given sub-problem size.
makeBase :: Int -> IO (BaseTransform, Maybe VI)
makeBase sz
  | sz `IM.member` specialBases = return (SpecialBase sz, Nothing)
  | isPrime sz                  = makeRaderBase sz
  | otherwise                   = return (makeDFTBase sz, Nothing)

-- | Generate digit reversal permutation using elementary "modulo"
-- permutations: last digit is not permuted to match with using a
-- simple DFT or instance of Rader's algorithm at the "bottom" of the
-- overall algorithm.
digrev :: Int -> VI -> Maybe VI
digrev n fs
  | null fs   = Nothing
  | otherwise = Just $ V.foldl1' (%.%) $ V.map (dupperm n) subperms
  where
    vfs = convert fs

    -- Sizes of the individual permutations that we need, one per
    -- factor.
    sizes = V.scanl div n vfs

    -- Partial sub-permutations, one per factor.
    subperms = V.reverse $ V.zipWith perm sizes vfs

    -- Generate a single "modulo" permutation.
    perm sz fac = concatMap doone $ enumFromN 0 fac
      where doone i = generate (sz `div` fac) (\j -> j * fac + i)

-- | Pre-computation plan for basic DFT transform.
makeDFTBase :: Int -> BaseTransform
makeDFTBase sz = DFTBase sz wsfwd wsinv
  where w = omega sz
        wsfwd = generate sz (w ^)
        wsinv = map (1 /) wsfwd

-- | Pre-compute plan for prime-length Rader FFT transform.
makeRaderBase :: Int -> IO (BaseTransform, Maybe VI)
makeRaderBase sz = do
  -- Forward FFT with embedded plan calculation.
  let fft p xs = execute p Forward xs

  -- Times for unpadded and padded convolution transforms.
  unpadtime <- planTime sz1
  padtime <- planTime pow2sz

  -- Should we use a padded or an unpadded transform for the
  -- convolution?
  let pad = padtime < unpadtime
      csz = if pad then pow2sz else sz1

  -- Plan for convolution transforms.
  cplan <- plan csz

  -- FFT transforms of b sequences for use in convolution, one for
  -- forward transform and one for inverse transform.  Either just the
  -- basic root of unity powers for the convolution (if we're not
  -- padding), or the powers cyclically repeated to make a vector of
  -- the next power-of-two length.
  let convb =
        fft cplan $ if pad
                    then generate csz (\idx -> bs ! (idx `mod` sz1))
                    else bs
      convbinv =
        fft cplan $ if pad
                    then generate csz (\idx -> bsinv ! (idx `mod` sz1))
                    else bsinv

  return (RaderBase sz outperm convb convbinv csz cplan, Just inperm)
  where
    -- Convolution length.
    sz1 = sz - 1

    -- Convolution length padded to next greater power of two.
    pow2sz = if sz1 == 2^ log2 sz1
             then sz1
             else 2 ^ (1 + log2 (2 * sz1 - 3))

    -- Group generator and inverse group generator.
    g = primitiveRoot sz
    ig = invModN sz g

    -- Input value permutation according to group generator indexing.
    inperm = 0 `cons` iterateN sz1 (\n -> (g * n) `mod` sz) 1

    -- Index vector based on inverse group generator ordering.
    outperm = iterateN sz1 (\n -> (ig * n) `mod` sz) 1

    -- Root of unity powers based on inverse group generator indexing,
    -- for forward and inverse transform.
    w = omega sz
    bs = backpermute (map (w ^^) $ enumFromTo 0 sz1) outperm
    bsinv = backpermute (map ((w ^^) . negate) $ enumFromTo 0 sz1) outperm

-- | Base transform type with heuristic ordering.
data BaseType = Special Int | Rader Int deriving (Eq, Show)

-- | Newtype wrapper for custom sorting.
newtype SPlan = SPlan (BaseType, Vector Int) deriving (Eq, Show)

-- | Base transform size.
bSize :: BaseType -> Int
bSize (Special b) = b
bSize (Rader b)   = b

-- | Heuristic ordering for base transform types: special bases come
-- first, then prime bases using Rader's algorithm, ordered according
-- to size compensating for padding needed in the Rader's algorithm
-- convolution.
instance Ord BaseType where
  compare (Special _)  (Rader _)    = LT
  compare (Rader _)    (Special _)  = GT
  compare (Special s1) (Special s2) = compare s2 s1
  compare (Rader r1)   (Rader r2)   = case (isPow2 $ r1 - 1, isPow2 $ r2 - 1) of
    (True, True)   -> compare r1 r2
    (True, False)  -> compare r1 (2 * r2)
    (False, True)  -> compare (2 * r1) r2
    (False, False) -> compare r1 r2

-- | Heuristic ordering for full plans, based first on base type, then
-- on the maximum size of Danielson-Lanczos step.
instance Ord SPlan where
  compare (SPlan (b1, fs1)) (SPlan (b2, fs2)) = case compare b1 b2 of
    LT -> LT
    EQ -> compare (maximum fs2) (maximum fs1)
    GT -> GT

-- | Generate test plans for a given input size, sorted in heuristic
-- order.
testPlans :: Int -> Int -> [(Int, Vector Int)]
testPlans n nplans = L.take nplans $ L.map clean $ L.sort okplans
  where vfs = allFactors n
        bs = usableBases n vfs
        doone b = basePlans n vfs b
        clean (SPlan (b, fs)) = (bSize b, fs)
        allplans = P.concatMap doone bs
        okplans = case L.filter (not . ridiculous) allplans of
          []  -> L.filter (not . reallyRidiculous) allplans
          oks -> oks
        ridiculous (SPlan (_, fs)) = any (> 128) fs
        reallyRidiculous (SPlan (_, fs)) =
          any (> 128) $ filter (not . isPrime) fs

-- | List plans from a single base.
basePlans :: Int -> Vector Int -> BaseType -> [SPlan]
basePlans _n vfs bt = if null lfs
                     then [SPlan (bt, empty)]
                     else P.map (\v -> SPlan (bt, v)) $ leftOvers lfs
  where lfs = fromList $ toList vfs \\ toList (allFactors b)
        b = bSize bt

-- | Produce all distinct permutations and compositions constructable
-- from a given list of factors.
leftOvers :: Vector Int -> [Vector Int]
leftOvers fs =
  if null fs
  then []
  else S.toList $ L.foldl' go S.empty (multisetPerms fs)
  where n = length fs
        go fset perm = foldl' doone fset (enumFromN 0 (2^(n - 1)))
          where doone s i = S.insert (makeComp perm i) s

-- | Usable base transform sizes.
usableBases :: Int -> Vector Int -> [BaseType]
usableBases n fs = P.map Special bs P.++ P.map Rader ps
  where bs = toList $ filter ((== 0) . (n `mod`)) specialBaseSizes
        ps = toList $ filter isPrime $ filter (> maxPrimeSpecialBaseSize) fs


{-# LANGUAGE GeneralizedNewtypeDeriving #-}

module ML where

import AI.HNN.FF.Network
import Data.Either (rights)
import Numeric.LinearAlgebra
import System.IO
import System.IO.Unsafe
import Text.ParserCombinators.Parsec

newtype KmH = KmH { kmh :: Double }
   deriving (Num, Floating, Fractional, Real, RealFloat, RealFrac, Eq, Ord)

newtype Rad = Rad { rad :: Double }
   deriving (Num, Floating, Fractional, Real, RealFloat, RealFrac, Eq, Ord)

newtype Force = Force { force :: Double }
   deriving (Num, Floating, Fractional, Real, RealFloat, RealFrac, Eq, Ord)

newtype Sec = Sec { sec :: Double }
   deriving (Num, Floating, Fractional, Real, RealFloat, RealFrac, Eq, Ord)

instance Read KmH where
   readsPrec i = map (\(x,s') -> (KmH x,s')) . readsPrec i

instance Show KmH where
   showsPrec i (KmH d) = showsPrec i d
   show (KmH d) = show d

instance Read Rad where
   readsPrec i = map (\(x,s') -> (Rad x,s')) . readsPrec i

instance Show Rad where
   showsPrec i (Rad d) = showsPrec i d
   show (Rad d) = show d

instance Read Force where
   readsPrec i = map (\(x,s') -> (Force x,s')) . readsPrec i

instance Show Force where
   showsPrec i (Force d) = showsPrec i d
   show (Force d) = show d

instance Read Sec where
   readsPrec i = map (\(x,s') -> (Sec x,s')) . readsPrec i

instance Show Sec where
   showsPrec i (Sec d) = showsPrec i d
   show (Sec d) = show d

data Example = Example { speedX0  :: KmH, speedX1 :: KmH
                       , speedY0  :: KmH, speedY1 :: KmH
                       , angle0   :: Rad, angle1  :: Rad
                       , accel0   :: Force
                       , steer0   :: Force
                       , timeStep :: Sec
                       } deriving Show

deg2rad :: Double -> Rad
deg2rad deg = Rad (deg * 0.017453292519943295)

rad2deg :: Rad -> Double
rad2deg (Rad rad) = rad / 0.017453292519943295

kmh2ms :: KmH -> Double
kmh2ms (KmH kmh) = kmh / 3.6

ms2kmh :: Double -> KmH
ms2kmh ms = KmH (ms * 3.6)

examplesIO :: String -> IO [Example]
examplesIO fn = readFile fn >>= return . map make . rights . map parseLine . lines
   where make :: [Double] -> Example
         make xs = Example { speedX0  = KmH $ xs !! 0, speedX1 = KmH $ xs !! 1
                           , speedY0  = KmH $ xs !! 2, speedY1 = KmH $ xs !! 3
                           , angle0   = Rad $ xs !! 4, angle1  = Rad $ xs !! 5
                           , accel0   = Force $ (xs !! 6) - (xs !! 7)
                           , steer0   = Force $ xs !! 8
                           , timeStep = Sec $ xs !! 9
                           }

         parseLine :: String -> Either ParseError [Double]
         parseLine input = parse lineFmt source input
            where source = "(string: " ++ input ++ ")"

         lineFmt :: GenParser Char st [Double]
         lineFmt = many number

         number :: GenParser Char st Double
         number = do _ <- many space
                     x <- many1 (noneOf "() \n")
                     _ <- many space
                     return (read x)

examples :: [Example]
examples = unsafePerformIO $ examplesIO "bla2"

example2sample :: Example -> Sample Double
example2sample ex = (fromList [{-kmh2ms (speedX0 ex),-} force (accel0 ex)], fromList [kmh2ms (speedX1 ex - speedX0 ex)])

samples :: Samples Double
samples = map example2sample examples

test :: IO ()
test = do   n <- createNetwork 1 [2] 1
            --mapM_ (putStrLn . show . output n tanh . fst) samples
            putStrLn "------------------"
            let n' = trainNTimes 100000 0.8 tanh tanh' n samples
            --mapM_ (putStrLn . show . present n) samples
            putStrLn ("Average error: " ++ show (avgErr n samples))
   where present n (inp, outp) = (toList inp, head $ toList obs, head $ toList outp, head $ diff)
            where obs  = output n tanh inp
                  diff = map (uncurry (-)) (zip (toList obs) (toList outp))
         avgErr n samples = sum (map err samples) / fromIntegral (length samples)
            where err (inp, outp) = abs ((head $ toList $ output n tanh inp) - (head $ toList $ outp)) / abs (head $ toList $ output n tanh inp)


{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE EmptyDataDecls #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE DeriveDataTypeable #-}
{-# LANGUAGE ViewPatterns #-}
-- |
-- Module     : Data.Matrix.Symmetric.Mutable
-- Copyright  : Copyright (c) 2012 Aleksey Khudyakov <alexey.skladnoy@gmail.com>
-- License    : BSD3
-- Maintainer : Aleksey Khudyakov <alexey.skladnoy@gmail.com>
-- Stability  : experimental
--
-- Symmetric and hermitian matrices
module Data.Matrix.Symmetric.Mutable (
    -- * Data types
    MSymmetric
  , MHermitian
    -- ** Implementation
  , MSymmetricRaw(..)
  , IsSymmetric
  , IsHermitian
    -- * Function
  , new
  , symmetricAsDense
  , hermitianAsDense
  , symmIndex
    -- * Complex number
  , NumberType
  , IsReal
  , IsComplex
  , castSymmetric
  , Conjugate(..)
  ) where

import Control.Monad
import Control.Monad.Primitive

import Data.Complex  (Complex,conjugate)
import Data.Typeable (Typeable)
import           Data.Vector.Storable.Internal
import qualified Data.Vector.Generic.Mutable as M

-- import Foreign.Ptr
import Foreign.Marshal.Array ( advancePtr )
import Foreign.ForeignPtr
import Foreign.Storable

import Data.Internal
import Data.Vector.Storable.Strided.Mutable
import Data.Matrix.Generic.Mutable
import Data.Matrix.Dense.Mutable   (MMatrix(..))

import Unsafe.Coerce

----------------------------------------------------------------
-- Data types
----------------------------------------------------------------

-- | Symmetric/hermitian matrix. Whether it's symmetric of hermitian
--   is determined by type tag. See 'IsSymmetric' and 'IsHermitian'.
--
--   Storage takes n² elements and data is stored in column major
--   order. Fields are
--
-- * Order of matrix
--
-- * Leading dimension size
--
-- * Pointer to data
data MSymmetricRaw tag s a = MSymmetricRaw
                             {-# UNPACK #-} !Int -- Order of matrix
                             {-# UNPACK #-} !Int -- Leading dimension size
                             {-# UNPACK #-} !(ForeignPtr a)
                      deriving (Typeable)

-- | Type tag for symmetric matrices.
data IsSymmetric

-- | Type tag for hermitian matrices.
data IsHermitian

-- | Mutable symmetric matrix
type MSymmetric = MSymmetricRaw IsSymmetric

-- | Mutable hermitian matrix
type MHermitian = MSymmetricRaw IsHermitian

instance Storable a => IsMMatrix (MSymmetricRaw IsSymmetric) a where
  basicRows (MSymmetricRaw n _ _) = n
  {-# INLINE basicRows #-}
  basicCols (MSymmetricRaw _ n _) = n
  {-# INLINE basicCols #-}
  basicIsIndexMutable _ _     = True
  {-# INLINE basicIsIndexMutable #-}
  basicUnsafeRead  (MSymmetricRaw _ lda fp) (symmIndex -> (!i,!j))
    = unsafePrimToPrim
    $ withForeignPtr fp (`peekElemOff` (i + j*lda))
  {-# INLINE basicUnsafeRead #-}
  basicUnsafeWrite (MSymmetricRaw _ lda fp) (symmIndex -> (!i,!j)) x
    = unsafePrimToPrim
    $ withForeignPtr fp $ \p -> pokeElemOff p (i + j*lda) x
  {-# INLINE basicUnsafeWrite #-}
  basicCloneShape = new . rows
  {-# INLINE basicCloneShape #-}
  basicClone = cloneSym
  {-# INLINE basicClone #-}

instance (Conjugate a, Storable a) => IsMMatrix (MSymmetricRaw IsHermitian) a where
  basicRows (MSymmetricRaw n _ _) = n
  {-# INLINE basicRows #-}
  basicCols (MSymmetricRaw _ n _) = n
  {-# INLINE basicCols #-}
  basicIsIndexMutable _ _     = True
  {-# INLINE basicIsIndexMutable #-}
  basicUnsafeRead  (MSymmetricRaw _ lda fp) (!i,!j)
    = unsafePrimToPrim
    $ case () of
       _| i > j     -> conjugateNum `liftM`  withForeignPtr fp (`peekElemOff` (j + i*lda))
        | otherwise ->                       withForeignPtr fp (`peekElemOff` (i + j*lda))
  {-# INLINE basicUnsafeRead #-}
  basicUnsafeWrite (MSymmetricRaw _ lda fp) (!i,!j) x
    = unsafePrimToPrim
    $ case () of
       _| i > j     -> withForeignPtr fp $ \p -> pokeElemOff p (j + i*lda) (conjugateNum x)
        | otherwise -> withForeignPtr fp $ \p -> pokeElemOff p (i + j*lda) x
  {-# INLINE basicUnsafeWrite #-}
  basicCloneShape = new . rows
  {-# INLINE basicCloneShape #-}
  basicClone = cloneSym
  {-# INLINE basicClone #-}

-- | Choose index so upper part of matrix is accessed
symmIndex :: (Int,Int) -> (Int,Int)
{-# INLINE symmIndex #-}
symmIndex (i,j)
  | i > j     = (j,i)
  | otherwise = (i,j)

-- | Allocate new matrix. It works for both symmetric and hermitian
--   matrices.
new :: (PrimMonad m, Storable a)
    => Int                      -- ^ Matrix order
    -> m (MSymmetricRaw tag (PrimState m) a)
{-# INLINE new #-}
new n = do
  fp <- unsafePrimToPrim $ mallocVector $ n * n
  return $ MSymmetricRaw n n fp

-- | Convert to dense matrix. Dense matrix will use same buffer as
--   symmetric
symmetricAsDense
  :: (PrimMonad m, Storable a)
  => MSymmetricRaw IsSymmetric (PrimState m) a
  -> m (MMatrix (PrimState m) a)
symmetricAsDense (MSymmetricRaw n lda fp) = do
  let m = MMatrix n n lda fp
  forM_ [0 .. n-2] $ \i -> do
    let row = MVector (n-i-1) lda $ updPtr (`advancePtr` (  i + (i+1) * lda)) fp
        col = MVector (n-i-1) 1   $ updPtr (`advancePtr` (1+i +    i  * lda)) fp
    M.move col row
  return m

-- | Convert to dense matrix. Dense matrix will use same buffer as
--   symmetric
hermitianAsDense
  :: (PrimMonad m, Storable a, Conjugate a)
  => MSymmetricRaw IsSymmetric (PrimState m) a
  -> m (MMatrix (PrimState m) a)
hermitianAsDense (MSymmetricRaw n lda fp) = do
  let m = MMatrix n n lda fp
  forM_ [0 .. n-2] $ \i -> do
    let len = n - i - 1
        row = MVector (n-i-1) lda $ updPtr (`advancePtr` (  i + (i+1) * lda)) fp
        col = MVector (n-i-1) 1   $ updPtr (`advancePtr` (1+i +    i  * lda)) fp
    forM_ [0 .. len - 1] $ \j -> do
      M.unsafeWrite col j . conjugateNum =<< M.unsafeRead row j
  return m

----------------------------------------------------------------
-- Real/Complex distinction
----------------------------------------------------------------

-- | Type tag for real numbers
data IsReal

-- | Type tag for complex numbers
data IsComplex

type family NumberType a :: *

type instance NumberType Float       = IsReal
type instance NumberType Double      = IsReal
type instance NumberType (Complex a) = IsComplex

-- | Cast between symmetric and hermitian matrices is data parameter
--   is real.
castSymmetric :: (NumberType a ~ IsReal)
              => MSymmetricRaw tag s a -> MSymmetricRaw tag' s a
{-# INLINE castSymmetric #-}
castSymmetric = unsafeCoerce

-- | Conjugate which works for both real (noop) and complex values.
class Conjugate a where
  conjugateNum :: a -> a

instance Conjugate Float where
  conjugateNum = id
  {-# INLINE conjugateNum #-}

instance Conjugate Double where
  conjugateNum = id
  {-# INLINE conjugateNum #-}

instance RealFloat a => Conjugate (Complex a) where
  conjugateNum = conjugate
  {-# INLINE conjugateNum #-}

----------------------------------------------------------------
-- Helpers
----------------------------------------------------------------

-- Get n'th column of matrix as mutable vector. Internal since part of
-- vector contain junk
unsafeGetCol :: Storable a => MSymmetricRaw tag s a -> Int -> MVector s a
{-# INLINE unsafeGetCol #-}
unsafeGetCol (MSymmetricRaw n lda fp) i
  = MVector n 1 $ updPtr (`advancePtr` (i*lda)) fp

cloneSym :: (Storable a, PrimMonad m)
         => MSymmetricRaw tag (PrimState m) a
         -> m (MSymmetricRaw tag (PrimState m) a)
{-# INLINE cloneSym #-}
cloneSym m@(MSymmetricRaw n _ _) = do
  q <- new n
  forM_ [0 .. n - 1] $ \i ->
    M.unsafeCopy (unsafeGetCol q i) (unsafeGetCol m i)
  return q


module Parser
  (readExpr) where

import Control.Monad
import Text.ParserCombinators.Parsec hiding (spaces)
import Data.Ratio
import Numeric
import Data.Complex

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

readExpr :: String -> String
readExpr input = case parse parseExpr "lisp" input of
                   Left err -> "No match. " ++ show err
                   Right val -> "Found Val " ++ show val

spaces :: Parser()
spaces = skipMany1 space

data LispVal = Atom String
             | List [LispVal]
             | DottedList [LispVal] LispVal
             | Number Integer
             | String String
             | Bool Bool
             | Character Char
             | Float Double
             | Ratio Rational
             | Complex (Complex Double)
             deriving Show

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

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

parseCharacter :: Parser LispVal
parseCharacter = do
  string "#\\"
  x <- many $ anyChar
  return $ case x of
             "space" -> Character ' '
             "newline" -> Character '\n'
             _ -> Character (x !! 0)

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

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

parseList :: Parser LispVal
parseList = liftM List $ sepBy parseExpr spaces

parseDottedList :: Parser LispVal
parseDottedList = do
    head <- endBy parseExpr spaces
    tail <- char '.' >> spaces >> parseExpr
    return $ DottedList head tail

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

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

parseDecimal :: Parser LispVal
parseDecimal = many1 digit >>= (return . Number . read)

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

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

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

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

oct2dig x = fst $ readOct x !! 0
hex2dig x = fst $ readHex x !! 0
bin2dig  = bin2dig' 0
bin2dig' digint "" = digint
bin2dig' digint (x:xs) = let old = 2 * digint + (if x == '0' then 0 else 1) in
                             bin2dig' old xs

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

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

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

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

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

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

parseExpr :: Parser LispVal
parseExpr = parseAtom
         <|> parseString
         <|> try parseComplex
         <|> try parseRatio
         <|> try parseFloat
         <|> try parseNumber
         <|> try parseBool
         <|> try parseCharacter
         <|> parseQuoted
         <|> parseBackquote
         <|> parseCommaList
         <|> do char '('
                x <- try parseList <|> parseDottedList
                char ')'
                return x


module TensorUtils where

import Prelude hiding ((.), id)
import System.Random
import Control.Monad

import Numeric.LinearAlgebra.Array
import Numeric.LinearAlgebra.Array.Util

import CategoricDefinitions
import Autodiff.GAD
import Autodiff.D
import Autodiff.Dual
import Autodiff.Additive

import OnesLike

{-
All tensor stuff is pretty much ad-hoc, a complete rewrite is eventually needed
-}

class ArrShow a where
    arrShow :: a -> String

instance {-# OVERLAPPABLE #-} Show a => ArrShow a where
    arrShow = show

instance  ArrShow Tensor where
    arrShow = formatFixed 2

instance (ArrShow a, ArrShow b) => ArrShow (a, b) where
    arrShow (a, b) = let ls = replicate 10 '-' ++ "\n"
                     in ls ++ arrShow a ++ ", \n" ++ arrShow b ++ "\n" ++ ls

sh :: ArrShow a => a -> IO ()
sh a = putStr $ arrShow a ++ "\n"

infixl 8 −|
--(−|) :: Name → [Array Double ] → Array Double
(−|) = index

axes :: Tensor -> String -> (String, [Int])
axes t axesNames = let toSum = filter (\idx -> head (iName idx) `elem` axesNames) (dims t)
                   in (map (head . iName) toSum, map iDim toSum)

class TensorContractable k where
    sumAxes :: String -> Tensor `k` Tensor

instance TensorContractable (->) where
    sumAxes axesNames = \t -> let (names, lengths) = axes t axesNames
                                  u = lengths # repeat 1 ! names
                              in t*u

instance (Scalable k Tensor, TensorContractable k)
    => TensorContractable (GADType k) where
    sumAxes axesNames = GAD $ \t -> let (names, lengths) = axes t axesNames
                                        u = lengths # repeat 1 ! names
                                    in (t*u, scale u)

instance (TensorContractable k) => TensorContractable (DualType k) where
    sumAxes axesNames = Dual (sumAxes axesNames)

instance TensorContractable DType where
    sumAxes axesNames = D (sumAxes axesNames)

instance TensorContractable (->+) where
    sumAxes axesNames = AddFun (sumAxes axesNames)

--meanAxes :: _ => String -> Tensor -> Tensor
--meanAxes axesNames t = let (names, lengths) = axes t axesNames
--                           n = product lengths
--                           v = scalar . fromIntegral $ n
--                       in divide (sumAxes axesNames t, v)

-- rand array with given shape and axis names
randomTensor :: [Int] -> String -> IO Tensor
randomTensor xs cs = do
    rs <- replicateM (product xs) randomIO
    return $ xs # rs ! cs


{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE NamedFieldPuns #-}
{-# LANGUAGE ScopedTypeVariables #-}

-- |
module BrainWall.Random where

import BrainWall.Edge
import qualified BrainWall.Main.GenerateHole as GenHole
import qualified BrainWall.Main.GenerateFigure as GenFig
import BrainWall.Json
import BrainWall.Polygon
import BrainWall.Polygon.ContainsEdge
import BrainWall.Problem
import BrainWall.V2
import Control.Monad
import qualified Data.Aeson as Aeson
import qualified Data.Aeson.Types as Aeson
import Data.List
import Data.Ord
import Data.Ratio
import qualified Data.Set as Set
import qualified Data.Text as T
import qualified Data.Vector as V
import qualified Data.Vector.Storable as VS
import Debug.Trace
import Numeric.GSL.Minimization
import Numeric.LinearAlgebra
import System.Directory
import System.Random.MWC as MWC
import System.Random.Stateful

chooseFrom :: StatefulGen g m => V.Vector x -> g -> m x
chooseFrom v gen = do
  idx <- uniformRM (0, V.length v - 1) gen
  return $ v V.! idx

genFigureFromPolygon :: StatefulGen g m => Polygon Integer -> g -> m Figure
genFigureFromPolygon polygon gen = do
  edges <- replicateM 50 genEdge
  return $ figureFromEdges $ filter isGood edges --
  where
    vertices = unPolygon polygon
    genEdge = do
      v1 <- chooseFrom vertices gen
      v2 <- chooseFrom vertices gen
      return $ Edge v1 v2
    isGood edge@(Edge v1 v2) = v1 /= v2 && containsEdge edge polygon

data EpsilonType
  = ForwardEpsilon -- The mutated figure will be the target figure
  | ReverseEpsilon -- The mutated figure will be the initial figure
  deriving (Enum, Bounded, Eq, Show)

data InitialPositionType
  = StartFigurePosition
  | RandomPosition
  deriving (Enum, Bounded, Eq, Show)

data ObjectiveType
  = Crumple
  | Stretch
  deriving (Enum, Bounded, Eq, Show)

data FigureSource
  = HandDrawn
  | PolyTree
  | FromHoleGenerator
  deriving (Enum, Bounded, Eq, Show)

randomGenerateFigure :: StatefulGen g IO => FigureSource -> g -> IO Figure
randomGenerateFigure PolyTree gen =
  figureFromEdges . GenFig.bfEdges <$> GenFig.generateFigure GenFig.defaultOptions gen
randomGenerateFigure FromHoleGenerator gen = do
  poly <- GenHole.generateHole GenHole.defaultOptions gen
  let problemHole = case GenHole.polyToHole poly of
        Just p  -> p
        Nothing -> error "error: no polygon"
  GenHole.genFigureInHole (GenHole.optionsFigure GenHole.defaultOptions) problemHole gen
randomGenerateFigure HandDrawn gen = do
  examples <- listDirectory "examples"
  let problems = V.fromList $ map ("examples/" ++) $ filter isProblem examples
  problemPath <- chooseFrom problems gen
  Just value <- Aeson.decodeFileStrict problemPath
  let figure = case Aeson.parse (\(Aeson.Object obj) -> obj Aeson..: T.pack "figure" >>= decodeFigure ) value of
        Aeson.Success p -> p
        Aeson.Error str -> error str
  pure figure
  where
    isProblem "" = False
    isProblem ".problem" = True
    isProblem (_:xs) = isProblem xs

randomMutateFigure ::
  forall m g.
  StatefulGen g m =>
  InitialPositionType ->
  EpsilonType ->
  ObjectiveType ->
  Figure ->
  (Integer, Integer) ->
  g ->
  m (Figure, Epsilon)
randomMutateFigure initialPositionType epsilonType objectiveType startFigure (canvasWidth, canvasHeight) gen = do
  initialVec <- case initialPositionType of
    StartFigurePosition -> pure $ V.map (fmap fromIntegral) $ figureVertices startFigure
    RandomPosition -> V.replicateM numVertices genPosition
  objective <- case objectiveType of
    Crumple -> do
      targetPositions <- V.replicateM numVertices genPosition
      pure $ negateObjective $ targetObjective targetPositions
    Stretch -> do
      vertexPairs <- genVertexPairs (V.length $ figureVertices startFigure) numPairs gen
      pure $ negateObjective $ distanceSumObjective vertexPairs
  return $ randomMutateFigureWithObjective initialVec epsilonType startFigure objective
  where
    numPairs = 50
    numVertices = V.length $ figureVertices startFigure
    genPosition :: m (V2 Double)
    genPosition = do
      x <- uniformRM (0, fromIntegral canvasWidth) gen
      y <- uniformRM (0, fromIntegral canvasHeight) gen
      return $ V2 x y

genVertexPairs :: forall m g. StatefulGen g m => Int -> Int -> g -> m (V.Vector (Int, Int))
genVertexPairs numVertices numPairs gen = do
  pairs <- replicateM numPairs genPair
  return $ V.fromList $ Set.toAscList $ Set.fromList $ filter (\(x, y) -> x /= y) pairs
  where
    genPair :: m (Int, Int)
    genPair = do
      x <- uniformRM (0, numVertices - 1) gen
      y <- uniformRM (0, numVertices - 1) gen
      return $ if x < y then (x, y) else (y, x)

data PositionObjectiveWithGradientHessian = PositionObjectiveWithGradientHessian
  { posObjective :: V.Vector (V2 Double) -> Double,
    posGradient :: V.Vector (V2 Double) -> VS.Vector Double,
    posHessian :: V.Vector (V2 Double) -> Matrix Double
  }

targetObjective :: V.Vector (V2 Double) -> PositionObjectiveWithGradientHessian
targetObjective targetPositions =
  PositionObjectiveWithGradientHessian
    { posObjective,
      posGradient,
      posHessian
    }
  where
    posObjective positions = V.sum $ V.zipWith dist targetPositions positions
    numVertices = V.length targetPositions
    posGradient positions =
      VS.convert $
        V.concat
          [ V.zipWith (\(V2 tx _) (V2 x _) -> 2 * (x - tx)) targetPositions positions,
            V.zipWith (\(V2 _ ty) (V2 _ y) -> 2 * (y - ty)) targetPositions positions
          ]
    posHessian _ = build (2 * numVertices, 2 * numVertices) $ \i j -> if i == j then 2 else 0
    dist v1 v2 = sq (v2X v1 - v2X v2) + sq (v2Y v1 - v2Y v2)
    sq x = x * x

negateObjective :: PositionObjectiveWithGradientHessian -> PositionObjectiveWithGradientHessian
negateObjective = scaleObjective (-1)

scaleObjective :: Double -> PositionObjectiveWithGradientHessian -> PositionObjectiveWithGradientHessian
scaleObjective mult oldObj =
  PositionObjectiveWithGradientHessian
    { posObjective = (*) mult . posObjective oldObj,
      posGradient = VS.map ((*) mult) . posGradient oldObj,
      posHessian = scale mult . posHessian oldObj
    }

distanceSumObjective :: V.Vector (Int, Int) -> PositionObjectiveWithGradientHessian
distanceSumObjective vertexPairs =
  PositionObjectiveWithGradientHessian
    { posObjective,
      posGradient,
      posHessian
    }
  where
    posObjective positions = V.sum $ V.map go vertexPairs
      where
        go (i, j) = dist posI posJ
          where
            posI = positions V.! i
            posJ = positions V.! j
    posGradient positions = VS.fromList $ map (go v2X) [0 .. numVertices - 1] ++ map (go v2Y) [0 .. numVertices - 1]
      where
        numVertices = V.length positions
        go get k = V.sum $ V.map go2 vertexPairs
          where
            go2 (i, j)
              | posI == posJ = 0
              | k == i = (get posI - get posJ) / dist posI posJ
              | k == j = (get posJ - get posI) / dist posI posJ
              | otherwise = 0
              where
                posI = positions V.! i
                posJ = positions V.! j
    posHessian positions = fromBlocks [[makeHessian goXX, makeHessian goXY], [makeHessian goXY, makeHessian goYY]]
      where
        numVertices = V.length positions
        get = (positions V.!)
        makeHessian aux =
          fromLists [[aux rowIdx colIdx | colIdx <- [0 .. numVertices - 1]] | rowIdx <- [0 .. numVertices - 1]]

        goXX rowIdx colIdx
          | rowIdx == colIdx =
            V.sum $
              V.map (\(i, j) -> let d = dist (get i) (get j) in if d == 0 then 0 else 1 / d - sq (v2X $ get i .-. get j) / cube d) $
                V.filter (\(i, j) -> i == rowIdx || j == rowIdx) vertexPairs
          | otherwise =
            V.sum $
              V.map (\(i, j) -> let d = dist (get i) (get j) in if d == 0 then 0 else - sq (v2Y $ get i .-. get j) / cube d) $
                V.filter (\(i, j) -> i == min colIdx rowIdx && j == max colIdx rowIdx) vertexPairs
        goYY rowIdx colIdx
          | rowIdx == colIdx =
            V.sum $
              V.map (\(i, j) -> let d = dist (get i) (get j) in if d == 0 then 0 else 1 / d - sq (v2Y $ get i .-. get j) / cube d) $
                V.filter (\(i, j) -> i == rowIdx || j == rowIdx) vertexPairs
          | otherwise =
            V.sum $
              V.map (\(i, j) -> let d = dist (get i) (get j) in if d == 0 then 0 else - sq (v2X $ get i .-. get j) / cube d) $
                V.filter (\(i, j) -> i == min colIdx rowIdx && j == max colIdx rowIdx) vertexPairs
        goXY rowIdx colIdx
          | rowIdx == colIdx =
            V.sum $
              V.map (\(i, j) -> let d = dist (get i) (get j) in if d == 0 then 0 else - (v2X $ get i .-. get j) * (v2Y $ get i .-. get j) / cube d) $
                V.filter (\(i, j) -> i == rowIdx || j == rowIdx) vertexPairs
          | otherwise =
            V.sum $
              V.map (\(i, j) -> let d = dist (get i) (get j) in if d == 0 then 0 else (v2X $ get i .-. get j) * (v2Y $ get i .-. get j) / cube d) $
                V.filter (\(i, j) -> i == min colIdx rowIdx && j == max colIdx rowIdx) vertexPairs

    dist v1 v2 = sqrt $ sq (v2X v1 - v2X v2) + sq (v2Y v1 - v2Y v2)
    sq x = x * x
    cube x = x * x * x

randomMutateFigureWithObjective :: V.Vector (V2 Double) -> EpsilonType -> Figure -> PositionObjectiveWithGradientHessian -> (Figure, Epsilon)
randomMutateFigureWithObjective initialPos epsilonType startFigure PositionObjectiveWithGradientHessian {posObjective, posGradient, posHessian} =
  traceShow ("rows", rows path) $
    traceShow ("objective", objective solution, epsilon) $
      traceShow ("distances", V.zip4 distances finalDistances roundedDistances epsilons) $
        traceShow (newPositions, newMultipliers) $
          traceShow ("gradients", lagrangian (newPositions, newMultipliers), lagrangianGradient (newPositions, newMultipliers), lagrangianHessian (newPositions, newMultipliers) #> lagrangianGradient (newPositions, newMultipliers)) $
            (Figure roundedPositions (figureEdges startFigure), epsilon)
  where
    initialVec = VS.convert $ V.map v2X initialPos V.++ V.map v2Y initialPos V.++ V.replicate numEdges 0
    lagrangian :: (V.Vector (V2 Double), V.Vector Double) -> Double
    lagrangian (positions, multipliers) =
      posObjective positions
        + V.sum (V.zipWith3 (\mult edge realDist -> mult * (distIdx positions edge - fromIntegral realDist) / sqrt (fromIntegral realDist)) multipliers (figureEdges startFigure) distances)
    lagrangianGradient :: (V.Vector (V2 Double), V.Vector Double) -> VS.Vector Double
    lagrangianGradient (positions, multipliers) =
      ( posGradient positions
          + VS.concat
            [ VS.map (edgeGradient v2X) (VS.enumFromN 0 numVertices),
              VS.map (edgeGradient v2Y) (VS.enumFromN 0 numVertices)
            ]
      )
        VS.++ V.convert (V.zipWith (\edge realDist -> (distIdx positions edge - fromIntegral realDist) / sqrt (fromIntegral realDist)) (figureEdges startFigure) distances)
      where
        edgeGradient get vIdx =
          V.sum $
            V.zipWith3
              ( \(v1Idx, v2Idx) mult realDist ->
                  case () of
                    _
                      | vIdx == v1Idx -> mult * 2 * (get (positions V.! vIdx) - get (positions V.! v2Idx)) / sqrt (fromIntegral realDist)
                      | vIdx == v2Idx -> mult * 2 * (get (positions V.! vIdx) - get (positions V.! v1Idx)) / sqrt (fromIntegral realDist)
                      | otherwise -> 0
              )
              (figureEdges startFigure)
              multipliers
              distances
    lagrangianHessian (positions, multipliers) =
      fromBlocks
        [ [topLeftHessian, tr' bottomLeftHessian],
          [bottomLeftHessian, multHessian]
        ]
      where
        multHessian :: Matrix Double
        multHessian = build (numEdges, numEdges) $ \_ _ -> 0
        xMultHessian get = build (numEdges, numVertices) $ \eIdx vIdx ->
          let realDist = distances V.! round eIdx
              (v1Idx, v2Idx) = figureEdges startFigure V.! round eIdx
           in case () of
                _
                  | round vIdx == v1Idx -> 2 / sqrt (fromIntegral realDist) * (get (positions V.! v1Idx) - get (positions V.! v2Idx))
                  | round vIdx == v2Idx -> 2 / sqrt (fromIntegral realDist) * (get (positions V.! v2Idx) - get (positions V.! v1Idx))
                  | otherwise -> 0
        topLeftHessian =
          fromBlocks
            [ [xxHessian, 0],
              [0, xxHessian]
            ]
            + posHessian positions
        bottomLeftHessian = fromBlocks [[xMultHessian v2X, xMultHessian v2Y]]
        xxHessian :: Matrix Double
        xxHessian = build (numVertices, numVertices) $ \i j ->
          V.sum $
            V.zipWith3
              ( \(v1Idx, v2Idx) mult realDist ->
                  case () of
                    _
                      | round i == v1Idx && round j == v1Idx -> mult * 2 / sqrt (fromIntegral realDist)
                      | round i == v2Idx && round j == v2Idx -> mult * 2 / sqrt (fromIntegral realDist)
                      | round i == v1Idx && round j == v2Idx -> - mult * 2 / sqrt (fromIntegral realDist)
                      | round i == v2Idx && round j == v1Idx -> - mult * 2 / sqrt (fromIntegral realDist)
                      | otherwise -> 0
              )
              (figureEdges startFigure)
              multipliers
              distances

    objective = VS.sum . VS.map sq . lagrangianGradient . splitParams . VS.convert
    -- (solution, path) =
    --   minimizeV
    --     NMSimplex2
    --     1e-9
    --     1000000
    --     (VS.replicate numParams 1)
    --     objective
    --     initialVec
    (solution, path) =
      minimizeVD
        VectorBFGS2
        1e-9
        1000
        0.01
        0.1
        objective
        ((\(positions, multipliers) -> scale 2 $ lagrangianHessian (positions, multipliers) #> lagrangianGradient (positions, multipliers)) . splitParams . VS.convert)
        initialVec
    -- initialVec = posInitial (V.map (fmap fromIntegral) $ figureVertices startFigure) VS.++ VS.replicate numEdges 0
    (newPositions, newMultipliers) = splitParams $ VS.convert solution
    numRotations = 200 :: Int
    numShifts = 20 :: Int
    (roundedPositions, roundedDistances, epsilons, epsilon) =
      minimumBy (comparing $ \(_, _, _, x4) -> x4) $
        [tryShiftRotation (V2 shiftX shiftY) angle | angle <- angles, shiftX <- shifts, shiftY <- shifts]
      where
        angles = [fromIntegral x / fromIntegral numRotations * pi / 2 | x <- [0 .. (numRotations - 1)]]
        shifts = [fromIntegral x / fromIntegral numShifts | x <- [0 .. (numShifts - 1)]]
    tryShiftRotation shift angle = (pos, roundedDists, epss, eps)
      where
        pos = V.map (fmap round) shifted
        shifted = V.map (.+. shift) rotated
        rotated = V.map (\(V2 x y) -> (V2 (cos angle * x - sin angle * y) (sin angle * x + cos angle * y))) newPositions
        roundedDists = V.map (distIdx pos) (figureEdges startFigure)
        eps :: Epsilon
        eps = ceiling $ V.maximum epss * fromIntegral epsilonDenominator
        epss = case epsilonType of
          ForwardEpsilon -> V.zipWith getEpsilon distances roundedDists
          ReverseEpsilon -> V.zipWith (flip getEpsilon) distances roundedDists
        reallyBig = fromIntegral (1000000000000 :: Int)
        getEpsilon refDist finalDistance = if refDist == 0 then reallyBig else
          (case compare finalDistance refDist of
            EQ -> 0
            GT -> finalDistance % refDist - 1
            LT -> negate $ finalDistance % refDist - 1)
    sq x = x * x
    numVertices = V.length $ figureVertices startFigure
    numEdges = V.length $ figureEdges startFigure
    splitParams params = (positions, multipliers)
      where
        multipliers = V.drop (2 * numVertices) params
        xs = V.take numVertices params
        ys = V.take numVertices $ V.drop numVertices params
        positions = V.zipWith V2 xs ys
    distances = V.map (distIdx (figureVertices startFigure)) (figureEdges startFigure)
    finalDistances = V.map (distIdx newPositions) (figureEdges startFigure)
    dist v1 v2 = sq (v2X v1 - v2X v2) + sq (v2Y v1 - v2Y v2)
    distIdx positions (v1Idx, v2Idx) = dist v1 v2
      where
        v1 = positions V.! v1Idx
        v2 = positions V.! v2Idx


-- TODO: better naming conventions

{-# LANGUAGE TypeFamilies, TypeOperators #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE DataKinds, PolyKinds, ConstraintKinds #-}
{-# LANGUAGE FlexibleInstances, FlexibleContexts, UndecidableInstances #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE PartialTypeSignatures #-}

module GrammarNNDAG where

import Control.Monad (foldM)
import Data.Function (fix)

import Generics.SOP hiding (SingI)
import Generics.SOP.NP
import Generics.SOP.NS
import Generics.SOP.Constraint
import Generics.SOP.Dict
import Data.Vinyl hiding (Dict)
import Data.Singletons.Prelude hiding (All, And)
import Data.Singletons.Prelude.List (Length)

import Utils
import List
import Grammar
import TensorHMatrix
import DAGIO
import TensorDAG
import VinylUtils
import Random

import Data.Default
import DefaultM

import Numeric.LinearAlgebra (Numeric)
import qualified Data.Vector.Storable as Vector

type family Size (s :: k -> *) t :: k
type IntegralS (s :: k -> *) = IntegralK ('KProxy :: KProxy k)

class SingI (Size s t) => KnownSize s t
instance SingI (Size s t) => KnownSize s t

class All2 (KnownSize s) (Code t) => KnownSizes s t
instance All2 (KnownSize s) (Code t) => KnownSizes s t

class SingI (FromInteger (Length (Code t)) :: k) => KnownCode (s :: k -> *) t
instance SingI (FromInteger (Length (Code t)) :: k) => KnownCode (s :: k -> *) t

class HasNodes t where
  getNodes :: t -> [Some Node]

newtype Repr a dim = Repr { runRepr :: Node (Tensor '[dim] a) }

instance HasNodes (Repr a dim) where
  getNodes (Repr node) = [Some node]

newtype ReprT a s t = ReprT { runReprT :: Node (Tensor '[Size s t] a) }

instance HasNodes (ReprT a s t) where
  getNodes (ReprT node) = [Some node]

instance (Default a, Usable a, IntegralS s, KnownSize s t) => DefaultM IO (ReprT a s t) where
  defM = ReprT <$> defM

--newtype Linear a outDim inDim = Linear (Node (Tensor '[outDim, inDim] a))
newtype LinearT a s tOut tIn = LinearT (Node (Tensor '[Size s tOut, Size s tIn] a))

instance HasNodes (LinearT a s tOut tIn) where
  getNodes (LinearT node) = [Some node]

--instance (Default a, Usable a, IntegralN outDim, SingI outDim, SingI inDim) => DefaultM IO (Linear a outDim inDim) where
--  defM = Linear <$> defM

instance (Default a, Usable a, IntegralS s, KnownSize s tIn, KnownSize s tOut) => DefaultM IO (LinearT a s tOut tIn) where
  defM = LinearT <$> defM

--linear :: Numeric a => Linear a outDim inDim -> Repr a inDim -> IO (Node (Tensor a '[outDim]))
--linear (Linear m) (Repr v) = makeMV m v

linearT :: (SingI (Size s tOut), IntegralS s, Usable a) => LinearT a s tOut tIn -> ReprT a s tIn -> IO (ReprT a s tOut)
linearT (LinearT m) (ReprT v) = ReprT <$> makeMV m v

data EncodeParams' a s parent children = EncodeParams' (ReprT a s parent) (NP (LinearT a s parent) children)

instance HasNodes (EncodeParams' a s parent children) where
  getNodes (EncodeParams' bias weights) = getNodes bias ++ concat (collapse_NP $ liftA_NP' (K . getNodes) weights)

instance (Default a, Usable a, IntegralS s, KnownSize s parent, SListI children, All (KnownSize s) children) => DefaultM IO (EncodeParams' a s parent children) where
  defM = EncodeParams' <$> defM <*> sequence'_NP (cpure_NP (Proxy::Proxy (KnownSize s)) (Comp defM))

encodeParent' :: (Usable a, Floating a, Floating (Vector.Vector a), IntegralS s, KnownSize s parent) =>
  EncodeParams' a s parent children -> NP (ReprT a s) children -> IO (ReprT a s parent)
encodeParent' (EncodeParams' (ReprT bias) weights) inputs = do
  let linearT' m v = Comp $ K <$> linearT m v
  children <- map runReprT . collapse_NP <$> (sequence'_NP $ liftA2_NP' linearT' weights inputs)
  parent <- foldM (makeBinary (+)) bias children
  ReprT <$> makeUnary tanh parent

newtype EncodeParams a s t = EncodeParams { runEncodeParams :: NP (EncodeParams' a s t) (Code t) }

instance HasNodes (EncodeParams a s t) where
  getNodes (EncodeParams params) = concat . collapse_NP $ liftA_NP' (K . getNodes) params

encodeParent :: (Floating a, Floating (Vector.Vector a), Usable a, IntegralS s, KnownSize s t) =>
  EncodeParams a s t -> SOP (ReprT a s) (Code t) -> IO (ReprT a s t)
encodeParent (EncodeParams params) (SOP sop) = collapse_NS $ liftA2_NS' encode params sop
  where encode p cs = K $ encodeParent' p cs

data Encoding a s t where
  Primitive :: ReprT a s t -> Encoding a s t
  Generic :: ReprT a s t -> SOP (Encoding a s) (Code t) -> Encoding a s t

getRepr :: Encoding a s t -> ReprT a s t
getRepr (Primitive repr) = repr
getRepr (Generic repr _) = repr

instance (Default a, Usable a, IntegralS s, KnownSize s t, KnownSizes s t) => DefaultM IO (EncodeParams a s t) where
  defM = EncodeParams <$> sequence'_NP (cpure_NP (Proxy::Proxy (All (KnownSize s))) (Comp defM))

newtype Encoder a s t = Encoder { runEncoder :: t -> IO (Encoding a s t) }

encodeRec :: forall ts a s t. (Floating a, Floating (Vector.Vector a), Usable a, KnownSize s t, IntegralS s) =>
  Rec (Encoder a s) ts -> Dict (Contained ts) t -> EncodeParams a s t -> Encoder a s t

encodeRec encoders Dict params = Encoder f where
  encoders' = cpure_POP (Proxy::Proxy (Find ts)) (Fn $ Comp . runEncoder (index find encoders) . unI)
  f t = do
    children <- sequence_SOP' (ap_SOP encoders' (from t))
    let childReprs = liftA_SOP' getRepr children
    parent <- encodeParent params childReprs
    return $ Generic parent children

makeEncoders :: forall p g a s. (Floating a, Floating (Vector.Vector a), Usable a, IntegralS s, All (KnownSize s) g) =>
  Complete p g -> NP (EncodeParams a s) g -> Rec (Encoder a s) p -> Rec (Encoder a s) (p :++ g)
makeEncoders complete params prim = fix f where
  f encoders = rAppend prim $ np2rec (cliftA2_NP (Proxy::Proxy (KnownSize s)) (encodeRec encoders) (unAll_NP complete) params)

newtype AnyEncoder a s ts = AnyEncoder { runAnyEncoder :: forall t. Find ts t => t -> IO (Encoding a s t) }

makeEncoder :: (Floating a, Floating (Vector.Vector a), Usable a, IntegralS s, All (KnownSize s) g) =>
  Complete p g -> NP (EncodeParams a s) g -> Rec (Encoder a s) p -> AnyEncoder a s (p :++ g)
makeEncoder complete params prim = AnyEncoder $ runEncoder (index find encoders)
  where encoders = makeEncoders complete params prim

newtype LinearIn a s t t' = LinearIn (Node (Tensor '[Size s t', Size s t] a))

instance HasNodes (LinearIn a s t t') where
  getNodes (LinearIn node) = [Some node]

instance (Default a, Usable a, IntegralS s, KnownSize s t, KnownSize s t') => DefaultM IO (LinearIn a s t t') where
  defM = LinearIn <$> defM

linearIn :: (IntegralS s, KnownSize s t', Usable a) => LinearIn a s t t' -> ReprT a s t -> IO (ReprT a s t')
linearIn (LinearIn m) (ReprT v) = ReprT <$> makeMV m v

data AffineIn a s t t' = AffineIn (ReprT a s t') (LinearIn a s t t')

instance HasNodes (AffineIn a s t t') where
  getNodes (AffineIn bias linear) = getNodes bias ++ getNodes linear

instance (Default a, Usable a, IntegralS s, KnownSize s t, KnownSize s t') => DefaultM IO (AffineIn a s t t') where
  defM = AffineIn <$> defM <*> defM

affineIn :: (Floating a, Floating (Vector.Vector a), Usable a, IntegralS s, KnownSize s t') => AffineIn a s t t' -> ReprT a s t -> IO (ReprT a s t')
affineIn (AffineIn (ReprT bias) weights) input = do
  ReprT l <- linearIn weights input
  b <- makeBinary (+) bias l
  t <- makeUnary tanh b
  return $ ReprT t

data Affine' a outDim inDim =
  Affine' (Node (Tensor '[outDim] a)) (Node (Tensor '[outDim, inDim] a))

instance HasNodes (Affine' a outDim inDim) where
  getNodes (Affine' bias linear) = [Some bias, Some linear]

instance (Default a, Usable a, IntegralN outDim, SingI outDim, SingI inDim) => DefaultM IO (Affine' a outDim inDim) where
  defM = Affine' <$> defM <*> defM

sigmoid x = 1 / (1 + exp (-x))

affine' :: (Floating a, Floating (Vector.Vector a), Usable a, IntegralN outDim, SingI outDim) => Affine' a outDim inDim -> Repr a inDim -> IO (Repr a outDim)
affine' (Affine' bias weight) (Repr input) = do
  l <- makeMV weight input
  b <- makeBinary (+) bias l
  Repr <$> makeUnary sigmoid b

data DecodeParams a s t =
  DecodeParams (Affine' a (Length (Code t)) (Size s t)) (POP (AffineIn a s t) (Code t))

instance HasNodes (DecodeParams a s t) where
  getNodes (DecodeParams aff params) = getNodes aff ++ concat (collapse_POP' $ liftA_POP' (K . getNodes) params)

instance (Default a, Usable a, KnownCode s t, KnownSize s t, KnownSizes s t) => DefaultM IO (DecodeParams a s t) where
  defM = DecodeParams <$> defM <*> sequence'_POP (cpure_POP (Proxy::Proxy (KnownSize s)) (Comp defM))

decodeParent :: forall a s t. (Real a, Floating a, Floating (Vector.Vector a), Usable a, All2 (KnownSize s) (Code t), KnownCode s t) =>
  DecodeParams a s t -> ReprT a s t -> IO (SOP (ReprT a s) (Code t))
decodeParent (DecodeParams aff params) parent = do
  Repr node <- affine' aff (Repr $ runReprT parent)
  Vector v <- evalNode node
  let weights = map toRational $ Vector.toList v

  let children = np2ns . unPOP $ cliftA_POP (Proxy::Proxy (KnownSize s)) (Comp . flip affineIn parent) params

  child <- sample $ zip children weights
  sequence'_SOP $ SOP child

newtype Decoder a s t = Decoder { runDecoder :: ReprT a s t -> IO t }

decodeRec :: forall ts a s t. (Real a, Floating a, Floating (Vector.Vector a), Usable a, KnownSizes s t, KnownCode s t) =>
  Rec (Decoder a s) ts -> Dict (Contained ts) t -> DecodeParams a s t -> Decoder a s t

decodeRec decoders Dict params = Decoder f where
  decoders' = cpure_POP (Proxy::Proxy (Find ts)) (Fn $ Comp . (I <$>) . runDecoder (index find decoders))
  f repr = do
    childReprs <- decodeParent params repr
    children <- sequence_SOP' (ap_SOP decoders' childReprs)
    return $ to children

makeDecoders :: forall p g a s. (Real a, Floating a, Floating (Vector.Vector a), Usable a, All (And (KnownSizes s) (KnownCode s)) g) =>
  Complete p g -> NP (DecodeParams a s) g -> Rec (Decoder a s) p -> Rec (Decoder a s) (p :++ g)
makeDecoders complete params prim = fix f where
  f decoders = rAppend prim $ np2rec (cliftA2_NP (Proxy::Proxy (And (KnownSizes s) (KnownCode s))) (decodeRec decoders) (unAll_NP complete) params)

newtype AnyDecoder a s ts = AnyDecoder { runAnyDecoder :: forall t. Find ts t => ReprT a s t -> IO t }

makeDecoder :: (Real a, Floating a, Floating (Vector.Vector a), Usable a, All (And (KnownSizes s) (KnownCode s)) g) =>
  Complete p g -> NP (DecodeParams a s) g -> Rec (Decoder a s) p -> AnyDecoder a s (p :++ g)
makeDecoder complete params prim = AnyDecoder $ runDecoder (index find decoders)
  where decoders = makeDecoders complete params prim

newtype AutoDecoder a s t = AutoDecoder { runAutoDecoder :: Encoding a s t -> IO (Node a) }

autoDecodeRec :: forall ts a s t. (Floating a, Floating (Vector.Vector a), Usable a, KnownCode s t, All2 (KnownSize s) (Code t)) =>
  Rec (AutoDecoder a s) ts -> Dict (Contained ts) t -> DecodeParams a s t -> AutoDecoder a s t

autoDecodeRec autoDecoders Dict (DecodeParams aff params) = AutoDecoder f where
  autoDecoders' = cpure_POP (Proxy::Proxy (Find ts)) (Fn $ Comp . (K <$>) . runAutoDecoder (index find autoDecoders))
  f :: Encoding a s t -> IO (Node a)
  f (Generic parent children) = do
    Repr node <- affine' aff (Repr $ runReprT parent)
    let i = ns2int (unSOP children)

    prob <- makeSelect i node
    log_score <- makeUnary (negate . log) prob

    let childReprs = liftA_SOP getRepr children
    let decodings = cliftA_POP (Proxy::Proxy (KnownSize s)) (Comp . flip affineIn parent) params

    let dist :: forall t'. KnownSize s t' => (IO :.: ReprT a s) t' -> ReprT a s t' -> (IO :.: K (Node a)) t'
        dist dIO (ReprT c) = Comp $ do
          ReprT d <- unComp dIO
          diff <- makeBinary (-) d c
          norm <- makeDot diff diff
          return $ K norm

    norms <- collapse_SOP <$> (sequence_SOP' $ cliftA2_SOP (Proxy::Proxy (KnownSize s)) dist decodings childReprs)

    childNodes <- collapse_SOP <$> sequence_SOP' (ap_SOP autoDecoders' children)

    foldM (makeBinary (+)) log_score (norms ++ childNodes)

primAutoDecoder :: Num a => AutoDecoder a s t
primAutoDecoder = AutoDecoder f where
  f (Primitive _) = makeSource 0

makeAutoDecoders :: forall proxy p g a s. (Real a, Floating a, Floating (Vector.Vector a), Usable a, All (And (KnownSizes s) (KnownCode s)) g) =>
  Complete p g -> NP (DecodeParams a s) g -> Rec (AutoDecoder a s) p -> Rec (AutoDecoder a s) (p :++ g)
makeAutoDecoders complete params prim = fix f where
  --prim = np2rec $ pure_NP primitiveAutoDecoder :: Rec _ p
  f autoDecoders = rAppend prim $ np2rec (cliftA2_NP (Proxy::Proxy (And (KnownSizes s) (KnownCode s))) (autoDecodeRec autoDecoders) (unAll_NP complete) params)

newtype AnyAutoDecoder a s ts = AnyAutoDecoder { runAnyAutoDecoder :: forall t. Find ts t => Encoding a s t -> IO (Node a) }

makeAutoDecoder :: (Real a, Floating a, Floating (Vector.Vector a), Usable a, All (And (KnownSizes s) (KnownCode s)) g) =>
  Complete p g -> NP (DecodeParams a s) g -> Rec (AutoDecoder a s) p -> AnyAutoDecoder a s (p :++ g)
makeAutoDecoder complete params prim = AnyAutoDecoder $ runAutoDecoder (index find autoDecoders)
  where autoDecoders = makeAutoDecoders complete params prim

newtype AutoEncoder a s t = AutoEncoder { runAutoEncoder :: t -> IO (Node a) }

makeAutoEncoders :: forall p g a s. (Real a, Floating a, Floating (Vector.Vector a), Usable a, All (KnownSize s) g, All (And (KnownSizes s) (KnownCode s)) g) =>
  Complete p g -> NP (EncodeParams a s) g -> Rec (Encoder a s) p -> NP (DecodeParams a s) g -> Rec (AutoDecoder a s) p -> Rec (AutoEncoder a s) (p :++ g)

makeAutoEncoders complete encodeParams primEncoders decodeParams primAutoDecoders = autoEncoders where
  encoders = makeEncoders complete encodeParams primEncoders
  autoDecoders = makeAutoDecoders complete decodeParams primAutoDecoders

  makeAutoEncoder (Encoder e) (AutoDecoder d) = AutoEncoder f where
    f t = e t >>= d

  autoEncoders = rZipWith makeAutoEncoder encoders autoDecoders

newtype AnyAutoEncoder a s ts = AnyAutoEncoder { runAnyAutoEncoder :: forall t. Find ts t => t -> IO (Node a) }

makeAutoEncoder :: (Real a, Floating a, Floating (Vector.Vector a), Usable a, All (KnownSize s) g, All (And (KnownSizes s) (KnownCode s)) g) =>
  Complete p g -> NP (EncodeParams a s) g -> Rec (Encoder a s) p -> NP (DecodeParams a s) g -> Rec (AutoDecoder a s) p -> AnyAutoEncoder a s (p :++ g)
makeAutoEncoder complete encodeParams primEncoders decodeParams primAutoDecoders = AnyAutoEncoder $ runAutoEncoder (index find autoEncoders)
  where autoEncoders = makeAutoEncoders complete encodeParams primEncoders decodeParams primAutoDecoders