{-# LANGUAGE GADTs #-}
module Evaluation
  ( sumProducedCompounds
  , numDistinctCompounds
  , percentagePoisonous
  , sumCompounds
  , from
  , meanAndVar
  , sumWith
  , countWith
  , byCompound
  , byProducts
  ) where

import Control.Foldl as F
import Numeric.LinearAlgebra as LA
import Numeric.LinearAlgebra.Devel as LA
import Data.List as L

import Environment

eps :: Amount
eps = 0.01

-- | sum of produced compounds ignoring everything defined as 'Nutrient'.
sumProducedCompounds :: Functor f => f (LA.Vector Amount) -> f Amount
-- sumProducedCompounds :: [LA.Vector Amount] -> [Amount]
-- cut off numNutrients from the start of the vector and then sum.
sumProducedCompounds = fmap sumElements . cutoffNutrients

-- | sum of all compounds
sumCompounds :: Functor f => f (LA.Vector Amount) -> f Amount
--sumCompounds :: [LA.Vector Amount] -> [Amount]
sumCompounds = fmap sumElements

-- | count compound as active if it is over eps.
numDistinctCompounds :: Functor f => f (LA.Vector Amount) -> f Amount
--numDistinctCompounds :: [LA.Vector Amount] -> [Amount]
numDistinctCompounds comps = sumElements . LA.cmap (\x -> if abs x < eps then 0 else 1) <$> comps

countWith :: Functor f => (LA.Vector Bool) -> f (LA.Vector Amount) -> f Amount
countWith toSelect = fmap $ sumElements . LA.zipVectorWith (\selected _ -> if selected then 1 else 0) toSelect
--                                      apply selection to set data to 1 or 0
--                        sum up 1 or 0s
--                 for all data

sumWith :: Functor f => (LA.Vector Bool) -> f (LA.Vector Amount) -> f Amount
sumWith toSelect = fmap $ sumElements . LA.zipVectorWith (\selected d -> if selected then d else 0) toSelect
--                                      apply selection to set data to 1 or 0
--                        sum up data or 0s
--                 for all data

percentagePoisonous :: Functor f => [Int] -> f (LA.Vector Amount) -> f Amount
percentagePoisonous poisons = fmap percentage
  where
      percentage v = (\(tox,total) -> tox / total) $ foldl1' pfold [(if i `L.elem` poisons then v!i else 0, v!i) | i <- [1..LA.size v]]
      pfold (a,b) (a',b') = (a+a', b+b')

-- | helper function for Foldl-Package.
--
--   Usage: @F.mean `from` sumCompounds $ v@ where v is a Set/List/Vector/... of Vector of Compounds.
infixr 9 `from`
from :: Foldable f => F.Fold a b -> (c -> f a) -> c -> b
from f w b = F.fold f (w b)

-- | helper to get mean and variance in a single pass.
meanAndVar :: F.Fold Amount (Amount,Amount)
meanAndVar = (,) <$> F.mean <*> F.variance

byProducts :: [LA.Vector Amount] -> [[Amount]]
byProducts = transpose . fmap toList . cutoffNutrients

byCompound :: [LA.Vector Amount] -> [[Amount]]
byCompound = transpose . fmap toList

-- | sum of produced compounds ignoring everything defined as 'Nutrient'.
cutoffNutrients :: Functor f => f (LA.Vector Amount) -> f (LA.Vector Amount)
-- sumProducedCompounds :: [LA.Vector Amount] -> [Amount]
-- cut off numNutrients from the start of the vector and then sum.
cutoffNutrients = fmap (\v -> LA.subVector numNutrients (size v - numNutrients) v)
    where
        numNutrients = fromEnum (maxBound :: Nutrient) + 1 --enum starts at 0, subVector-indexing at 1
