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

-- 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 = [ (Nitrate, 2)
             , (Phosphor, 3)
             , (Photosynthesis, 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 = True
                          , 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 = 1
    loop' :: Int -> Int -> [Plant] -> Environment -> IO ()
    loop' loopAmount curLoop plants e = unless (loopAmount == 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 <- fmap (+0.01) <$> fitness plants -- fitness should be at least 0.01 for mating to work
          let fps = zip plants fs -- gives us plants & their fitness in a tuple
              sumFitness = sum fs
          when (curLoop `mod` printEvery == 0) $ do
            liftIO $ printPopulation stringe fps
            liftIO $ hFlush stdout
          -- 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
      hFlush stdout
      when (curLoop `mod` printEvery == 0) $ do
        threadDelay $ 100*1000 -- sleep x*1000ns (=x ~ ms)
      loop' loopAmount (curLoop+1) newPlants e

main :: IO ()
main = do
  hSetBuffering stdin NoBuffering
  --hSetBuffering stdout NoBuffering
  randomCompounds <- makeHead (Substrate Photosynthesis) <$> generateTreeFromList 30 (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
  printEnvironment env
  writeFile "poison.twopi" $ generateDotFromPoisonTree "poison" 0.5 poisonedTree
  putStr "\ESC[?1049h"
  loop 200 emptyPlants env
  putStrLn "Simulation ended. Press key to exit."
  _ <- getChar
  putStr "\ESC[?1049l"

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.1,0.2)
              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)] -> 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 _,_) -> 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 = "_"
      putStr (plot curE)
    putStrLn ""

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