chemodiversity/app/Main.hs

{-# LANGUAGE BangPatterns #-}
module Main where

import Text.Printf
import Control.Monad.Reader
import Numeric.LinearAlgebra
import Data.List
import System.Random
import Control.Concurrent
import Control.Parallel.Strategies
import qualified Debug.Trace as Debug
import System.IO

import ArbitraryEnzymeTree
import Environment
import Evaluation

-- Example definitions
-- -------------------

-- Enzymes

-- pps :: Enzyme -- uses Phosphor from Substrate to produce PP
-- pps = Enzyme "PPS" [(Substrate Phosphor,1)] ((Substrate Phosphor,-1),(Produced PP,1)) Nothing
--
-- fpps :: Enzyme -- PP -> FPP
-- fpps = makeSimpleEnzyme (Produced PP) (Produced FPP)

-- Environment

exampleEnvironment :: Int -> [Enzyme] -> [(Predator,Probability)] -> [(Compound,Amount)] -> Environment
exampleEnvironment addedC es pred tox =
  Environment
    { soil = [ (PPM, 10)
             ]
    , predators = pred -- [ (greenfly, 0.1) ]
    , metabolismIteration = 100
    , maxCompound = maxCompoundWithoutGeneric + addedC
    , toxicCompounds = tox --[(Produced FPP,0.1)] ++ tox
    , possibleEnzymes = es -- [pps,fpps] ++ es
    , settings = Settings { automimicry = False
                          , predatorsRandom = False
                          , numPlants = 150
                          }
    }

-- Plants

-- examplePlants :: [Plant]
-- examplePlants = (\g -> Plant g defaultAbsorption) <$> genomes
--  where
--    enzymes = [pps, fpps]
--    quantity = [1,2] :: [Quantity]
--    activation = [0.7, 0.9, 1]
--
--    genomes = do
--      e <- permutations enzymes
--      e' <- subsequences e
--      q <- quantity
--      a <- activation
--      return $ (,,) <$> e' <*> [q] <*> [a]
--
--    defaultAbsorption = fmap ( limit Phosphor 2
--                             . limit Nitrate 1
--                             . limit Sulfur 0
--                             ) <$> asks soil
--    -- custom absorbtion with helper-function:
--    limit :: Nutrient -> Amount -> (Nutrient, Amount) -> (Nutrient, Amount)
--    limit n a (n', a')
--      | n == n'   = (n, min a a') -- if we should limit, then we do ;)
--      | otherwise = (n', a')

-- Running the simulation
-- ----------------------

loop :: Int -> [Plant] -> Environment -> IO ()
loop loopAmount ps env = loop' loopAmount 0 ps env

  where
    -- cache enzyme colorful-strings
    stringe :: [(Enzyme, String)]
    stringe = (\e -> case Data.List.find (\(t,_) -> (t==) . fst . snd . synthesis $ e) toxins of
               Just (_,toxicity) -> (e,"\ESC[38;5;" ++ show (16 + 36*5 + 6*floor (5*(1-toxicity)) + 0) ++ "m" -- yellow -> red rainbow for tocixity 0 -> 1
                                       ++ padded 50 (show (enzymeName e)) ++ "\ESC[0m")
               Nothing           -> (e, padded 50 (show (enzymeName e)))
             ) <$> possibleEnzymes env
    toxins :: [(Compound, Amount)]
    toxins = toxicCompounds env
    padded i str = take i $ str ++ repeat ' '
    printEvery = 10
    addedConstFitness = 0.1
    loop' :: Int -> Int -> [Plant] -> Environment -> IO ()
    loop' loopAmount curLoop plants e = unless (loopAmount+1 == curLoop) $ do
      when (curLoop `mod` printEvery == 0) $ do
        putStr "\ESC[2J\ESC[H"
        printEnvironment e
        putStrLn ""
        putStrLn $ "Generation " ++ show curLoop ++ " of " ++ show loopAmount ++ ":"
      newPlants <- flip runReaderT e $ do
          (!fs,cs) <- unzip . fmap (\(f,c) -> (f,c)) <$> fitness plants
          let fps = zip plants fs -- gives us plants & their fitness in a tuple
              sumFitness = sum fs
          when (curLoop `mod` printEvery == 0) $ liftIO $ do
            printPopulation stringe (zip3 plants fs cs)
            putStrLn $ "Population statistics: VarC = " ++ (padded 50 . show . varianceOfProducedCompounds $ cs) 
                                         ++ " DistC = " ++ (padded 50 . show . meanOfDistinctCompounds $ cs)
            hFlush stdout
            threadDelay $ 100*1000 -- sleep x*1000ns (=x ~ ms)
          -- generate x new plants.
          np <- asks (numPlants . settings)
          sequence . flip fmap [1..np] $ \_ -> do
                parent' <- liftIO $ randomRIO (0,sumFitness)
                let
                    -- if we only have one parent in our list, take it.
                    findParent :: Double -> [(Plant,Double)] -> Plant
                    findParent _ [(last,_)] = last
                    -- otherwise count down x to find the parent in the list
                    findParent x ((p,f):ps)
                      | x < f     = p
                      | otherwise = findParent (x-f) ps
                    parent = findParent parent' fps
                haploMate parent
      loop' loopAmount (curLoop+1) newPlants e

main :: IO ()
main = do
  hSetBuffering stdin NoBuffering
  --hSetBuffering stdout NoBuffering
  randomCompounds <- makeHead (Substrate PPM) <$> generateTreeFromList 20 (toEnum <$> [(maxCompoundWithoutGeneric+1)..] :: [Compound]) -- generate roughly x compounds
  ds <- randoms <$> newStdGen
  probs <- randomRs (0.2,0.7) <$> newStdGen
  let poisonedTree     = poisonTree ds randomCompounds
      poisonCompounds  = foldMap (\(a,b) -> [(b,a) | a > 0.2]) poisonedTree
  predators <- generatePredators 0.5 poisonedTree
  let env              = exampleEnvironment (getTreeSize randomCompounds) (generateEnzymeFromTree randomCompounds) (zip predators probs) poisonCompounds
      emptyPlants      = replicate (numPlants . settings $ env) emptyPlant
  enzs <- randomRs (0,length (possibleEnzymes env) - 1) <$> newStdGen
  let startPlants      = randomGenome 10 enzs (possibleEnzymes env) emptyPlants
  printEnvironment env
  writeFile "poison.twopi" $ generateDotFromPoisonTree "poison" 0.5 poisonedTree
  putStr "\ESC[?1049h"
  loop 2000 startPlants env
  putStrLn "Simulation ended. Press key to exit."
  _ <- getChar
  putStr "\ESC[?1049l"

randomGenome :: Int -> [Int] -> [Enzyme] -> [Plant] -> [Plant]
randomGenome num inds enzs []     = []
randomGenome num inds enzs (p:ps) = p { genome = genes} : randomGenome num r enzs ps
  where
    i' = take num inds
    r  = drop num inds
    enzymes = (enzs!!) <$> i'
    genes = (\e -> (e,1,1)) <$> enzymes


generatePredators :: Double -> EnzymeTree s (Double,Compound) -> IO [Predator]
generatePredators threshold t = do
    ps <- mapM generatePredators' $ getSubTrees t
    return $ filter ((/= []) . irresistance) $ concat ps -- filter out predators that are resistant to everything because this does not make sense in our model.
  where
      generatePredators' :: EnzymeTree s (Double, Compound) -> IO [Predator]
      generatePredators' t = do -- not fully resistant to t, but fully resistant to everything in ts
          let comps = foldMap (\(a,b) -> [(a,b) | a > threshold]) t
          amount <- randomRIO (0,length comps + 1) :: IO Int
          forM [1..amount] $ \_ -> do
              impact <- randomRIO (0.2,0.7)
              rands <- randoms <$> newStdGen
              let unresists = foldMap (\((a,b),r) -> [b | r*2 < a]) $ zip comps rands
              return $ Predator unresists impact 1

printEnvironment :: Environment -> IO ()
printEnvironment (Environment soil pred metaIter maxComp toxic possEnz settings) =
  do
    putStrLn "Environment:"
    putStrLn $ "Soil:      " ++ show soil
    putStrLn $ "Predators: " ++ show pred
    putStrLn $ "PSM Iters: " ++ show metaIter
    putStrLn $ "Compounds: " ++ show ((toEnum <$> [0..maxComp]) :: [Compound])
    putStrLn $ "Toxic:     " ++ show toxic
    putStrLn $ "Settings:  " ++ show settings

printPopulation :: [(Enzyme,String)] -> [(Plant,Double,Vector Amount)] -> IO ()
printPopulation es ps = do
  let padded i str = take i $ str ++ repeat ' '
  putStr $ padded 50 "Population:"
  forM_ ps $ \(_,f,_) -> putStr (printColor f '█')
  putStrLn colorOff
  forM_ es $ \(e,s) -> do
    putStr s
    forM_ ps $ \(Plant g _,_,cs) -> do
      let curE = sum $ map (\(_,q,a) -> fromIntegral q*a)
                     . filter (\(e',_,_) -> e == e')
                     $ g
          plot x
           | x > 2     = 'O'
           | x > 1     = '+'
           | x > 0.7   = 'ö'
           | x > 0.5   = 'o'
           | x > 0     = '.'
           | otherwise = '_'
          amount = min 2 $ cs ! fromEnum (fst . snd . synthesis $ e)
      putStr $ printColor (amount/2) (plot curE)
    putStrLn colorOff

printColor :: Double -> Char -> String
printColor x c
  | x*x < 0.5 = "\ESC[38;5;" ++ show (16 + 36*5 + 6*floor (5*2*x') + 0) ++ "m" ++ [c] ++ ""
  | otherwise = "\ESC[38;5;" ++ show (16 + 36*floor (5*2*(1-x')) + 6*5 + 0) ++ "m" ++ [c] ++ ""
  -- 32 bit
  -- | x*x < 0.5 = "\ESC[38;2;255;" ++ (show . floor $ 255*2*x') ++ ";0m" ++ [c] ++ ""
  -- | otherwise = "\ESC[38;2;" ++ (show . floor $ 255*2*(1-x')) ++ ";255;0m" ++ [c] ++ ""
   where x' = x*x

colorOff :: String
colorOff = "\ESC[0m"

generateEnzymeFromTree :: EnzymeTree s Compound -> [Enzyme]
generateEnzymeFromTree t = (makeSimpleEnzyme c . getElement <$> sts)
                           ++ concatMap generateEnzymeFromTree sts
  where
    c   = getElement t
    sts = getSubTrees t