{-# 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)

-- Predator

greenfly :: Predator         -- 20% of plants die to greenfly, but the fly is
greenfly = Predator [] 0.2   -- killed by any toxic Component

-- 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
    }

-- 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 = loop' loopAmount 0

  where
    loop' :: Int -> Int -> [Plant] -> Environment -> IO ()
    loop' loopAmount curLoop plants e = unless (loopAmount == curLoop) $ do
      putStr "\ESC[2J\ESC[H"
      printEnvironment e
      putStrLn ""
      putStrLn $ "Generation " ++ show curLoop ++ " of " ++ show loopAmount ++ ":"
      newPlants <- flip runReaderT e $ do
          ! fs <- sequence (fitness <$> plants)
          let fps = zip plants fs -- gives us plants & their fitness in a tuple
              sumFitness = sum fs
          pe <- asks possibleEnzymes
          tc <- asks toxicCompounds
          liftIO $ printPopulation tc pe fps
          liftIO $ hFlush stdout
          -- generate 100 new plants.
          sequence . flip fmap [1..100] $ \_ -> 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
      threadDelay $ 100*1000 -- sleep 100ms
      loop' loopAmount (curLoop+1) newPlants e

main :: IO ()
main = do
  hSetBuffering stdin NoBuffering
  hSetBuffering stdout NoBuffering
  randomCompounds <- makeHead (Substrate Photosynthesis) <$> generateTreeFromList 40 (toEnum <$> [(maxCompoundWithoutGeneric+1)..] :: [Compound]) -- generate roughly x compounds
  ds <- randoms <$> newStdGen
  probs <- randomRs (0.2,0.7) <$> newStdGen
  let emptyPlants      = replicate 100 emptyPlant
      poisonedTree     = poisonTree ds randomCompounds
      poisonCompounds  = foldMap (\(a,b) -> if a > 0.5 then [(b,a)] else []) $ poisonedTree
  predators <- generatePredators 0.5 poisonedTree
  let env              = exampleEnvironment (getTreeSize randomCompounds) (generateEnzymeFromTree randomCompounds) (zip predators probs) poisonCompounds
  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"

  -- fitness <- runReaderT (sequence $ (\a -> do p <- absorbNutrients a >>= produceCompounds a; (,,) a p <$> deterPredators p) <$> emptyPlants) exampleEnvironment
  -- mapM_ (printf "%15.15s, " . show . toEnum @Compound) [0..maxCompoundWithoutGeneric]
  -- putStrLn "Fitness"
  -- forM_ fitness $ \(p, c, f) -> do
  --   mapM_ (printf "%15.2f, ") (toList c)
  --   printf "%15.2f" f
  --   putStr "\n"

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) -> if a > threshold then [(a,b)] else []) 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) -> if r*2 < a then [b] else []) $ zip comps rands
              return $ Predator unresists impact

printEnvironment :: Environment -> IO ()
printEnvironment (Environment soil pred metaIter maxComp toxic possEnz) =
  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

printPopulation :: [(Compound, Amount)] -> [Enzyme] -> [(Plant,Double)] -> IO ()
printPopulation toxins 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 -> do
    putStr $ case Data.List.find (\(t,_) -> (t==) . fst . snd . synthesis $ e) toxins of
               Just (_,toxicity) -> "\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           -> padded 50 (show (enzymeName e))
    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