module Genome
  ( Gene (..)
  , Genome (..)
  , mutateGeneSource
  , mutateGeneSink
  , mutateGeneWeight
  , mutateGene
  ) where


import Mind (NeuronIndex (..))
import System.Random


data Gene = Gene { source :: NeuronIndex, sink :: NeuronIndex, weight :: Float } deriving (Eq, Show)
data Genome = Genome { numInput :: Int, numInternal :: Int, numOutput :: Int, genes :: [Gene] } deriving (Eq, Show)


-- choose a random list element
randomChoice :: RandomGen a => [b] -> a -> (b, a)
randomChoice xs r =
  let (index, r') = randomR (0, length xs - 1) r
  in (xs !! index, r')


-- pick a new random source for the gene
mutateGeneSource :: RandomGen a => Genome -> Gene -> a -> (Gene, a)
mutateGeneSource genome g r =
  let 
    nInput = numInput genome
    nInternal = numInternal genome
    index :: Int
    (index, r') = randomR (0, nInput+nInternal-1) r
    source' = 
      if index < nInput
        then Input index
        else Internal $ index - nInput
  in ( g { source = source' }, r')


-- pick a new random sink for the gene
mutateGeneSink :: RandomGen a => Genome -> Gene -> a -> (Gene, a)
mutateGeneSink genome g r =
  let
    nInternal = numInternal genome
    nOutput = numOutput genome
    index :: Int
    (index, r') = randomR (0, nInternal+nOutput-1) r
    sink' = if index < nInternal
      then Internal index
      else Output $ index - nInternal
  in ( g { sink = sink' }, r' )


-- pick a new random weight, in the range (-w, w)
-- where w is twice the absolute value of the current weight
mutateGeneWeight :: RandomGen a => Genome -> Gene -> a -> (Gene, a)
mutateGeneWeight _ g r =
  let
    w = 2 * (abs $ weight g)
    (weight', r') = randomR (negate w, w) r
  in ( g { weight = weight' }, r' )


-- randomly mutate gene (modify source, sink, or weight
mutateGene :: RandomGen a => Genome -> Gene -> a -> (Gene, a)
mutateGene genome g r =
  let (f, r') = randomChoice [mutateGeneSource, mutateGeneSink, mutateGeneWeight] r
  in f genome g r'