module MindTest (suite) where

import Test.Tasty
import Test.Tasty.HUnit
import Mind
import Data.Either

suite :: TestTree
suite = testGroup "mind tests" $
  [ neuronIndexTests
  , networkTests
  ]

neuronIndexTests :: TestTree
neuronIndexTests = testGroup "neuron index tests" $
  [ testCase "get input index" $ getNeuronIndex (Input 4) @?= 4
  , testCase "get internal index" $ getNeuronIndex (Internal 12) @?= 12
  , testCase "get output index" $ getNeuronIndex (Output 0) @?= 0
  ]

networkTests :: TestTree
networkTests = testGroup "network tests" $
  [ testCase "create empty network" $
    (createEmptyNetwork 3 2 1) @?= Network 3 [[], []] [[]]

  -- neuron connections
  , testCase "output network connection" $
    let network = Network 3 [[], []] [[]]
    in (connectNeurons network (Input 0) (Output 0) (negate 1.0))
    @?= (Just $ Network 3 [[], []] [[Edge (Input 0, (negate 1.0))]])
  , testCase "internal network connection" $
    let network = Network 3 [[], []] [[]]
    in (connectNeurons network (Internal 0) (Internal 1) (negate 1.0))
    @?= (Just $ Network 3 [[], [Edge (Internal 0, negate 1.0)]] [[]])
  , testCase "internal self-connection" $
    let network = Network 3 [[], []] [[]]
    in (connectNeurons network (Internal 0) (Internal 0) (negate 1.0))
    @?= (Just $ Network 3 [[Edge (Internal 0, negate 1.0)], []] [[]])
  , testCase "internal source out of range" $
    let network = Network 3 [[], []] [[]]
    in (connectNeurons network (Internal 5) (Internal 0) (negate 1.0))
    @?= Nothing
  , testCase "internal sink out of range" $
    let network = Network 3 [[], []] [[]]
    in (connectNeurons network (Internal 1) (Internal (negate 1)) (negate 1.0))
    @?= Nothing
  , testCase "input source out of range" $
    let network = Network 3 [[], []] [[]]
    in (connectNeurons network (Input (negate 1)) (Internal 0) (negate 1.0))
    @?= Nothing
  , testCase "input sink" $
    let network = Network 3 [[], []] [[]]
    in (connectNeurons network (Input 0) (Input 0) (negate 1.0))
    @?= Nothing
  , testCase "output source" $
    let network = Network 3 [[], []] [[]]
    in (connectNeurons network (Output 0) (Output 0) (negate 1.0))
    @?= Nothing
  , testCase "output sink out of range" $
    let network = Network 3 [[], []] [[]]
    in (connectNeurons network (Input 0) (Output 1) (negate 1.0))
    @?= Nothing

  -- network computations
  , testCase "single input, single output" $
    let 
      net = Network 1 [] [[Edge (Input 0, 2.0)]]
      Right (output, state) = compute net [negate 0.5] []
    in (output, state) @?=
    ( [tanh (2.0 * (negate 0.5))]
    , []
    )
  , testCase "multiple inputs, single output" $
    let 
      net = Network 4 [] [
        [ Edge (Input 0, negate 1.0)
        , Edge (Input 1, negate 2.0)
        , Edge (Input 2, 1.0)
        , Edge (Input 3, 2.0)
        ]]
      Right (output, state) = compute net [1, 2, 3, 5] []
    in (output, state) @?=
    ( 
      [tanh 
        ( (1 * (negate 1.0))
        + (2 * (negate 2.0))
        + (3 * 1.0)
        + (5 * 2.0)
        )
      ]
    , []
    )
  , testCase "multiple inputs, multiple outputs" $
    let 
      net = Network 4 [] 
        [ [ Edge (Input 0, negate 1.0)
          , Edge (Input 1, 1.0)
          ]
        , [ Edge (Input 2, negate 1.0)
          , Edge (Input 3, 1.0)
          ]
        ]
      Right (output, state) = compute net [1, 2, 3, 5] []
    in (output, state) @?=
    ( 
      [ tanh (2 - 1)
      , tanh (5 - 3)
      ]
    , []
    )
  , testCase "hidden neurons" $
    let 
      net = Network 4 
        -- hidden neurons
        [ [ Edge (Input 0, negate 1.0)
          , Edge (Input 1, 1.0)
          ]
        , [ Edge (Input 2, negate 1.0)
          , Edge (Input 3, 1.0)
          ]
        ] 
        -- output neurons
        [ [ Edge (Internal 0, negate 1.0)
          , Edge (Internal 1, 1.0)
          ]
        ]
      Right (output, state) = compute net [1, 2, 3, 5] [0, 0]
    in (output, state) @?=
    ( 
      [ tanh ( (tanh (5-3)) - (tanh (2-1)) ) ]
    , [ tanh (2-1)
      , tanh (5-3)
      ]
    )
  , testCase "computing with self-connection" $
    let
      net = Network 0
        -- hidden neurons
        [ [ Edge (Internal 0, negate 0.5) ] ]
        -- output neurons
        [ [ Edge (Internal 0, 2) ] ]
      Right result = compute net [] [1.0]
    in result @?=
    ( [ tanh $ 2 * (tanh (negate 0.5)) ]
    , [ tanh (negate 0.5) ]
    )
  , testCase "computation fails for bad input length" $
    let 
      net = Network 2 [[]] [[]]
      result = compute net (replicate 3 1.0) [1]
    in (isLeft result) @?= True
  , testCase "computation fails for bad state length" $
    let 
      net = Network 2 [[]] [[]]
      result = compute net (replicate 2 1.0) [1, 1]
    in (isLeft result) @?= True
  ]