module Main where

{-
  | = Example application: trigonometry cheating

  Find the trigonometric expression of cos(x) through sin(x)
  using our automatic programming method.
-}

import qualified Data.Vector as V
import           Data.List ( foldl' )
import           Control.Monad ( foldM )
import           Math.Probable.Random  -- From `probable` package
                 ( vectorOf
                 , double
                 )

import           AI.MEP

config = defaultConfig {
  -- Functions available to genetically produced programs
  c'ops = V.fromList [
       ('*', (*)),
       ('+', (+)),
       ('/', (/)),
       ('-', (-)),
       ('s', \x _ -> sin x)
     ]
  -- Chromosome length
  , c'length = 50
  }

-- | Absolute value distance between two scalar values
dist :: Double -> Double -> Double
dist x y = if isNaN x || isNaN y
  -- Large distance
  then 10000
  else abs $ x - y

-- Could be optimized
sum' :: Num a => [V.Vector a] -> V.Vector a
sum' xss = foldl' (V.zipWith (+)) base xss
  where
    len = V.length $ head xss
    base = V.replicate len 0

main :: IO ()
main = do
  -- A vector of 50 random numbers between 0 and 1 (including 1)
  xs <- runRandIO (vectorOf 50 double)

  -- Scale the values to the interval of (-pi, pi]
  let xs' = V.map ((2*pi *). subtract 0.5) xs
      -- Target function f to approximate
      function x = (cos x)^2
      -- Pairs (x, f(x))
      dataset = V.map (\x -> (x, function x)) xs'

  -- Randomly create a population of chromosomes
  pop <- runRandIO $ initialize config

  let -- The loss function which depends on the dataset
      loss evalf = (V.singleton i', loss')
        where
          (xs, ys) = unzip $ V.toList dataset
          -- Distances resulting from multiple expression evaluation
          dss = zipWith (\x y -> V.map (dist y). evalf. V.singleton $ x) xs ys
          -- Cumulative distances for each index
          dcumul = sum' dss
          -- Select index minimizing cumulative distances
          i' = V.minIndex dcumul
          -- The loss value with respect to the index of the best expression
          loss' = dcumul V.! i'

  -- Evaluate the initial population
  let popEvaluated = evaluateGeneration loss pop

  putStrLn $ "Average loss in the initial population " ++ show (avgLoss popEvaluated)

  -- Declare how to produce the new generation
  let nextGeneration = evolve config loss (mutation3 config) crossover binaryTournament

  -- Specify the I/O loop, which logs every 5 generation
      runIO pop i = do
        newPop <- runRandIO $ foldM (\xg _ -> nextGeneration xg) pop [1..generations]
        putStrLn $ "Population " ++ show (i * generations) ++ ": average loss " ++ show (avgLoss newPop)
        return newPop
          where generations = 5

  -- Final generation
  final <- foldM runIO popEvaluated [1..20]
  let best = last final
  print best
  putStrLn "Interpreted expression:"
  putStrLn $ generateCode best