module GenomeTest (suite) where

import Test.Tasty
import Test.Tasty.HUnit
import qualified Mind as M
import Genome
import System.Random

suite :: TestTree
suite = testGroup "genome tests" $
  [ mutationTests
  ]


mutationTests :: TestTree
mutationTests = testGroup "mutations" $
  [ testCase "mutating the source of a gene" $
    let
      rand = mkStdGen 1 -- randomR generates sequence (1, 2, 0, ...)
      genome = Genome { numInput = 2, numInternal = 1, numOutput = 5, genes = [] }
      sourceGene = Gene { source = M.Input 0, sink = M.Output 1, weight = 4 }
      mutatedGenes = fst $ foldl
        (\(list, r) g ->
          let (g', r') = mutateGeneSource genome g r
          in (list ++ [g'], r')
        )
        ([], rand)
        (replicate 3 sourceGene)
    in
      mutatedGenes @?=
      [ Gene { source = M.Input 1, sink = M.Output 1, weight = 4 }
      , Gene { source = M.Internal 0, sink = M.Output 1, weight = 4 }
      , Gene { source = M.Input 0, sink = M.Output 1, weight = 4 }
      ]
  , testCase "mutating the sink of a gene" $
    let
      rand = mkStdGen 1 -- randomR generates sequence (1, 2, 0, ...)
      genome = Genome { numInput = 2, numInternal = 1, numOutput = 2, genes = [] }
      sourceGene = Gene { source = M.Input 0, sink = M.Output 1, weight = 4 }
      mutatedGenes = fst $ foldl
        (\(list, r) g ->
          let (g', r') = mutateGeneSink genome g r
          in (list ++ [g'], r')
        )
        ([], rand)
        (replicate 3 sourceGene)
    in
      mutatedGenes @?=
      [ Gene { source = M.Input 0, sink = M.Output 0, weight = 4 }
      , Gene { source = M.Input 0, sink = M.Output 1, weight = 4 }
      , Gene { source = M.Input 0, sink = M.Internal 0, weight = 4 }
      ]
  , testCase "mutating the weight of a gene" $
    let
      rand = mkStdGen 0 -- randomR generates sequence (7.572357,-1.4116564,-7.2413177, ...)
      genome = Genome { numInput = 2, numInternal = 1, numOutput = 2, genes = [] }
      sourceGene = Gene { source = M.Input 0, sink = M.Output 1, weight = 4 }
      mutatedGenes = fst $ foldl
        (\(list, r) g ->
          let (g', r') = mutateGeneWeight genome g r
          in (list ++ [g'], r')
        )
        ([], rand)
        (replicate 3 sourceGene)
      expected =
        [ Gene { source = M.Input 0, sink = M.Output 1, weight = 7.572357}
        , Gene { source = M.Input 0, sink = M.Output 1, weight = -1.4116564 }
        , Gene { source = M.Input 0, sink = M.Output 1, weight = -7.2413177  }
        ]
      approxEqual a b = abs (a-b) < 0.0001
    in do
      (map source mutatedGenes) @?= (map source expected)
      (map sink mutatedGenes) @?= (map sink expected)
      True @?= foldl (&&) True (zipWith approxEqual
        (map weight mutatedGenes)
        (map weight expected)
        )
  , testCase "insert new internal neuron" $
    let 
      genome = Genome 
        { numInput = 1
        , numInternal = 2
        , numOutput = 1
        , genes =
          [ Gene { source = M.Input 0, sink = M.Internal 0, weight = 1.0 }
          , Gene { source = M.Internal 0, sink = M.Internal 1, weight = 1.0 }
          , Gene { source = M.Internal 1, sink = M.Output 0, weight = 1.0 }
          ]
        }
      r = mkStdGen 5
      (genome', r') = mutateGenomeAddInternal genome r
    in do
      r' @?= r
      (numInput genome') @?= (numInput genome)
      (numInternal genome') @?= (1 + numInternal genome)
      (numOutput genome') @?= (numOutput genome)
      (genes genome') @?= (genes genome)
  , testCase "remove internal neuron" $
    let 
      r = mkStdGen 1 -- randomR (0, 2) produces (1, 2, 0, ...)
      genome = Genome 
        { numInput = 1
        , numInternal = 3
        , numOutput = 1
        , genes =
          [ Gene { source = M.Input 0, sink = M.Internal 0, weight = 1.0 }
          , Gene { source = M.Internal 0, sink = M.Internal 1, weight = 1.0 }
          , Gene { source = M.Internal 1, sink = M.Internal 2, weight = 1.0 }
          , Gene { source = M.Internal 2, sink = M.Output 0, weight = 1.0 }
          ]
        }
      (genome', r') = mutateGenomeRemoveInternal genome r
      (genome'', r'') = mutateGenomeRemoveInternal genome r'
      (genome''', _) = mutateGenomeRemoveInternal genome r''
    in do
      (numInput genome') @?= (numInput genome)
      (numInput genome'') @?= (numInput genome)
      (numInput genome''') @?= (numInput genome)

      (numInternal genome') @?= (numInternal genome - 1)
      (numInternal genome'') @?= (numInternal genome - 1)
      (numInternal genome''') @?= (numInternal genome - 1)

      (numOutput genome') @?= (numOutput genome)
      (numOutput genome'') @?= (numOutput genome)
      (numOutput genome''') @?= (numOutput genome)

      (genes genome') @?=
        [ Gene { source = M.Input 0, sink = M.Internal 0, weight = 1.0 }
        , Gene { source = M.Internal 0, sink = M.Internal 1, weight = 1.0 }
        , Gene { source = M.Internal 0, sink = M.Internal 1, weight = 1.0 }
        , Gene { source = M.Internal 1, sink = M.Output 0, weight = 1.0 }
        ]

      (genes genome'') @?=
        [ Gene { source = M.Input 0, sink = M.Internal 0, weight = 1.0 }
        , Gene { source = M.Internal 0, sink = M.Internal 1, weight = 1.0 }
        , Gene { source = M.Internal 1, sink = M.Output 0, weight = 1.0 }
        , Gene { source = M.Internal 1, sink = M.Output 0, weight = 1.0 }
        ]

      (genes genome''') @?=
        [ Gene { source = M.Input 0, sink = M.Internal 0, weight = 1.0 }
        , Gene { source = M.Input 0, sink = M.Internal 0, weight = 1.0 }
        , Gene { source = M.Internal 0, sink = M.Internal 1, weight = 1.0 }
        , Gene { source = M.Internal 1, sink = M.Output 0, weight = 1.0 }
        ]
  ]