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

  , mutateGenomeAddInternal
  , mutateGenomeRemoveInternal
  , mutateGenomeAddGene 
  , mutateGenomeRemoveGene 
  ) where


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


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 (idx, r') = randomR (0, length xs - 1) r
  in (xs !! idx, 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
    idx :: Int
    (idx, r') = randomR (0, nInput+nInternal-1) r
    source' = 
      if idx < nInput
        then Input idx
        else Internal $ idx - 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
    idx :: Int
    (idx, r') = randomR (0, nInternal+nOutput-1) r
    sink' = if idx < nInternal
      then Internal idx
      else Output $ idx - 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'

-- add new internal neuron
mutateGenomeAddInternal :: RandomGen a => Genome -> a -> (Genome, a)
mutateGenomeAddInternal genome r = (genome { numInternal = 1 + numInternal genome }, r)


-- remove an internal neuron, decrementing sources and incrementing sinks
mutateGenomeRemoveInternal :: RandomGen a => Genome -> a -> (Genome, a)
mutateGenomeRemoveInternal g r =
  let
    (idx, r') = randomR (0, numInternal g - 1) r
    remove = Internal idx 
    genes' = map
      (\(Gene so si w) ->
        let
          so' = if (isInternal so) && (so >= remove)
            then decrementNeuron g so else so
          si' = if (si < remove) || (not $ isInternal si) then si
            else if si == remove
              then
                let si'' = incrementNeuron g si
                in if isInternal si'' then si else si''
              else decrementNeuron g si 
        in Gene {source=so', sink=si', weight=w}
      )
      (genes g)
  in
    (g { numInternal = (numInternal g - 1), genes = genes' }, r')

validn :: Genome -> NeuronIndex -> NeuronIndex
validn g (Input x)
  | (inRange (0, numInput g - 1) x) = Input x
  | otherwise = error "out of range Input!"
validn g (Internal x)
  | (inRange (0, numInternal g - 1) x) = Internal x
  | otherwise = error "out of range Internal!"
validn g (Output x)
  | (inRange (0, numOutput g - 1) x) = Output x
  | otherwise = error "out of range Output!"

incrementNeuron :: Genome -> NeuronIndex -> NeuronIndex
incrementNeuron g (Input x)
  | (x >= numInput g - 1) = validn g (Internal 0)
  | otherwise = validn g (Input $ x+1)
incrementNeuron g (Internal x)
  | (x >= numInternal g - 1) = validn g (Output 0)
  | otherwise = validn g (Internal $ x+1)
incrementNeuron g (Output x)
  | (x >= numOutput g - 1) = error "cannot increment past the end of outputs!"
  | otherwise = validn g (Output $ x+1)

decrementNeuron :: Genome -> NeuronIndex -> NeuronIndex
decrementNeuron g (Input x)
  | (x <= 0) = error "cannot decrement past the first Input!"
  | otherwise = validn g (Input $ x-1)
decrementNeuron g (Internal x)
  | (x <= 0) = validn g (Input $ numInput g - 1)
  | otherwise = validn g (Internal $ x-1)
decrementNeuron g (Output x)
  | (x <= 0) = validn g (Internal $ numInternal g - 1)
  | otherwise = validn g (Output $ x-1)

isInternal :: NeuronIndex -> Bool
isInternal (Internal _) = True
isInternal _ = False


-- add a new, random gene
mutateGenomeAddGene :: RandomGen a => Genome -> a -> (Genome, a)
mutateGenomeAddGene g r =
  let
    (so, r') = randomSource g r
    (si, r'') = randomSink g r'
    (w, r''') = randomR (-4, 4) r''
    gene = Gene so si w
  in (g {genes = gene:(genes g)}, r''')

randomSource :: RandomGen a => Genome -> a -> (NeuronIndex, a)
randomSource g r =
  let
    (idx, r') = randomR (0, numInput g + numInternal g - 1) r
    result = if idx < numInput g then Input idx else Internal $ idx - numInput g
  in (result, r')

randomSink :: RandomGen a => Genome -> a -> (NeuronIndex, a)
randomSink g r =
  let
    (idx, r') = randomR (0, numInternal g + numOutput g - 1) r
    result = if idx < numInternal g then Internal idx else Output $ idx - numInput g
  in (result, r')


-- remove a random gene
mutateGenomeRemoveGene :: RandomGen a => Genome -> a -> (Genome, a)
mutateGenomeRemoveGene g r =
  let
    (idx, r') = randomR (0, (length $ genes g) - 1) r
    (front, _:back) = splitAt idx (genes g)
  in (g {genes = front ++ back}, r')