module GenomeTest (suite) where

import Test.Tasty
import Test.Tasty.HUnit
import Mind (Edge (..), NeuronIndex (..))
import qualified Mind as N (Network (..))
import Genome
import System.Random
import Data.Ix


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 = Input 0, sink = 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 = Input 1, sink = Output 1, weight = 4 }
      , Gene { source = Internal 0, sink = Output 1, weight = 4 }
      , Gene { source = Input 0, sink = 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 = Input 0, sink = 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 = Input 0, sink = Output 0, weight = 4 }
      , Gene { source = Input 0, sink = Output 1, weight = 4 }
      , Gene { source = Input 0, sink = 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 = Input 0, sink = 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 = Input 0, sink = Output 1, weight = 7.572357}
        , Gene { source = Input 0, sink = Output 1, weight = -1.4116564 }
        , Gene { source = Input 0, sink = 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 = Input 0, sink = Internal 0, weight = 1.0 }
          , Gene { source = Internal 0, sink = Internal 1, weight = 1.0 }
          , Gene { source = Internal 1, sink = 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 = Input 0, sink = Internal 0, weight = 1.0 }
          , Gene { source = Internal 0, sink = Internal 1, weight = 1.0 }
          , Gene { source = Internal 1, sink = Internal 2, weight = 1.0 }
          , Gene { source = Internal 2, sink = 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 = Input 0, sink = Internal 0, weight = 1.0 }
        , Gene { source = Internal 0, sink = Internal 1, weight = 1.0 }
        , Gene { source = Internal 0, sink = Internal 1, weight = 1.0 }
        , Gene { source = Internal 1, sink = Output 0, weight = 1.0 }
        ]

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

      (genes genome''') @?=
        [ Gene { source = Input 0, sink = Internal 0, weight = 1.0 }
        , Gene { source = Input 0, sink = Internal 0, weight = 1.0 }
        , Gene { source = Internal 0, sink = Internal 1, weight = 1.0 }
        , Gene { source = Internal 1, sink = Output 0, weight = 1.0 }
        ]
  , testCase "add new gene" $
    let 
      genome = Genome 
        { numInput = 1
        , numInternal = 2
        , numOutput = 1
        , genes = []
        }
      r = mkStdGen 5
      (genome', _) = mutateGenomeAddGene genome r
      new = last $ genes genome'
      -- checking sources
      validSource (Input x) = inRange (0, numInput genome') x
      validSource (Internal x) = inRange (0, numInternal genome') x
      validSource (Output _) = False
      -- checking sinks
      validSink (Input _) = False
      validSink (Internal x) = inRange(0, numInternal genome') x
      validSink (Output x) = inRange(0, numOutput genome') x
      w = weight new
    in do
      (numInput genome') @?= (numInput genome)
      (numInternal genome') @?= (numInternal genome)
      (numOutput genome') @?= (numOutput genome)
      (length $ genes genome') @?= 1 + (length $ genes genome)
      validSource (source new) @?= True
      validSink (sink new) @?= True
      (w >= -4) && (w <= 4) @?= True
  , testCase "remove a gene" $
    let
      r = mkStdGen 1 -- randomR (0, 2) produces (1, 2, 0, ...)
      genome = Genome 
        { numInput = 1
        , numInternal = 3
        , numOutput = 1
        , genes =
          [ Gene { source = Input 0, sink = Internal 0, weight = 1.0 }
          , Gene { source = Internal 0, sink = Internal 1, weight = 1.0 }
          , Gene { source = Internal 2, sink = Output 0, weight = 1.0 }
          ]
        }
      (genome', r') = mutateGenomeRemoveGene genome r
      (genome'', r'') = mutateGenomeRemoveGene genome r'
      (genome''', _) = mutateGenomeRemoveGene genome r''
    in do
      genes genome' @?= 
        [ Gene { source = Input 0, sink = Internal 0, weight = 1.0 }
        , Gene { source = Internal 2, sink = Output 0, weight = 1.0 }
        ]
      genes genome'' @?=
        [ Gene { source = Input 0, sink = Internal 0, weight = 1.0 }
        , Gene { source = Internal 0, sink = Internal 1, weight = 1.0 }
        ]
      genes genome''' @?=
        [ Gene { source = Internal 0, sink = Internal 1, weight = 1.0 }
        , Gene { source = Internal 2, sink = Output 0, weight = 1.0 }
        ]
  , testCase "parse genome into network" $
    let 
      genome = Genome 
        { numInput = 1
        , numInternal = 3
        , numOutput = 1
        , genes =
          [ Gene { source = Input 0, sink = Internal 0, weight = 1.0 }
          , Gene { source = Internal 0, sink = Internal 1, weight = 2.0 }
          , Gene { source = Internal 1, sink = Internal 2, weight = 3.0 }
          , Gene { source = Internal 2, sink = Output 0, weight = 4.0 }
          ]
        }
      network = parseGenome genome
      expected = N.Network
        { N.numInput = 1
        , N.internalNeurons =
          [ [ Edge (Input 0, 1.0) ]
          , [ Edge (Internal 0, 2.0) ]
          , [ Edge (Internal 1, 3.0) ]
          ]
        , N.outputNeurons = 
          [ [ Edge (Internal 2, 4.0) ]
          ]
        } 
    in network @?= expected
  ]