{-# LANGUAGE OverloadedStrings #-}

{-# LANGUAGE RecordWildCards #-}
import           Control.Exception                  (evaluate)
import           Control.Monad
import qualified Data.Aeson                         as Json
import qualified Data.ByteString                    as B
import qualified Data.ByteString.Builder            as BBuilder
import           Data.Either                        (isRight, isLeft)
import           Data.Maybe                         (fromJust, isJust,
                                                     isNothing)
import qualified Data.Vector                        as V
import           Epidemic
import qualified Epidemic.Model.BDSCOD              as BDSCOD
import qualified Epidemic.Model.InhomogeneousBDSCOD as InhomBDSCOD
import           Epidemic.Data.Events
import           Epidemic.Data.Newick
import           Epidemic.Data.Observations ( Observation(..)
                                            , ReconstructedTree(..)
                                            , reconstructedTree
                                            , PointProcessEvents(..)
                                            , pointProcessEvents
                                            , reconstructedTreeEvents
                                            , observedEvents
                                            , aggregated)
import           Epidemic.Data.Parameter
import           Epidemic.Data.Population
import           Epidemic.Data.Simulation          (SimulationState (..),
                                                     TerminationHandler (..),
                                                     genIOFromFixed,
                                                     genIOFromSystem,
                                                     genIOFromWord32)
import           Epidemic.Data.Time
import           Epidemic.Utility
import           Statistics.Sample
import qualified Statistics.Distribution as Dist
import Statistics.Distribution.Binomial (binomial)
import qualified System.Random.MWC                  as MWC
import           Test.Hspec

-- | check if the value x seems like a plausible draw from a Binom(n,p) distribution.
binomialTest :: Int -> Int -> Probability -> Probability -> Bool
binomialTest x n p a =
  let binom = binomial n p
      cdfVal = Dist.cumulative binom (fromIntegral x)
  in a < cdfVal && cdfVal < (1 - a)

indSampWithRemoval :: AbsoluteTime -> Person -> Bool -> EpidemicEvent
indSampWithRemoval a b c = IndividualSample a b c True

-- | Helper function for converting from Either to Maybe monad.
either2Maybe x = case x of
  Right v -> Just v
  Left _  -> Nothing

-- | y is within n% of x from x.
withinNPercent n x y = x - d < y && y < x + d
  where
    d = n * x / 100

p1 = Person (Identifier 1)

p2 = Person (Identifier 2)

p3 = Person (Identifier 3)

p4 = Person (Identifier 4)

p5 = Person (Identifier 5)

p6 = Person (Identifier 6)

p7 = Person (Identifier 7)

-- | The first set of test data does not have any catastrophe events.
demoFullEvents01 =
  [ Infection (AbsoluteTime 1) p1 p2
  , Infection (AbsoluteTime 2) p1 p3
  , indSampWithRemoval (AbsoluteTime 3) p1 True
  , Infection (AbsoluteTime 4) p2 p4
  , Infection (AbsoluteTime 5) p2 p5
  , indSampWithRemoval (AbsoluteTime 6) p4 True
  , Infection (AbsoluteTime 7) p3 p6
  , indSampWithRemoval (AbsoluteTime 8) p2 False
  , Removal (AbsoluteTime 9) p3
  , Infection (AbsoluteTime 10) p5 p7
  , indSampWithRemoval (AbsoluteTime 11) p6 False
  , indSampWithRemoval (AbsoluteTime 12) p5 True
  , Removal (AbsoluteTime 13) p7
  ]

demoSampleEvents01 =
  [ Infection (AbsoluteTime 1) p1 p2
  , indSampWithRemoval (AbsoluteTime 3) p1 True
  , Infection (AbsoluteTime 4) p2 p4
  , indSampWithRemoval (AbsoluteTime 6) p4 True
  , indSampWithRemoval (AbsoluteTime 8) p2 False
  , indSampWithRemoval (AbsoluteTime 11) p6 False
  , indSampWithRemoval (AbsoluteTime 12) p5 True
  ]

-- | The second set of test data is the same as the first but includes a
-- catastrophe event.
demoFullEvents02 =
  [ Infection (AbsoluteTime 1) p1 p2
  , Infection (AbsoluteTime 2) p1 p3
  , indSampWithRemoval (AbsoluteTime 3) p1 True
  , Infection (AbsoluteTime 4) p2 p4
  , Infection (AbsoluteTime 5) p2 p5
  , indSampWithRemoval (AbsoluteTime 6) p4 True
  , Infection (AbsoluteTime 7) p3 p6
  , indSampWithRemoval (AbsoluteTime 8) p2 False
  , Removal (AbsoluteTime 9) p3
  , Infection (AbsoluteTime 10) p5 p7
  , PopulationSample (AbsoluteTime 11) (asPeople [p5]) True
  , indSampWithRemoval (AbsoluteTime 12) p6 False
  , Removal (AbsoluteTime 13) p7
  ]

demoSampleEvents02 =
  [ Infection (AbsoluteTime 1) p1 p2
  , indSampWithRemoval (AbsoluteTime 3) p1 True
  , Infection (AbsoluteTime 4) p2 p4
  , indSampWithRemoval (AbsoluteTime 6) p4 True
  , indSampWithRemoval (AbsoluteTime 8) p2 False
  , PopulationSample (AbsoluteTime 11) (asPeople [p5]) True
  , indSampWithRemoval (AbsoluteTime 12) p6 False
  ]

-- | Another test set to test that catastrophes are handled correctly.
demoFullEvents03 =
  [ Infection (AbsoluteTime 1) p1 p4
  , Infection (AbsoluteTime 2) p1 p2
  , indSampWithRemoval (AbsoluteTime 3) p1 True
  , Infection (AbsoluteTime 4) p2 p3
  , Infection (AbsoluteTime 5) p4 p5
  , PopulationSample (AbsoluteTime 6) (asPeople [p2, p3, p4]) True
  ]

demoSampleEvents03 =
  [ Infection (AbsoluteTime 1) p1 p4
  , Infection (AbsoluteTime 2) p1 p2
  , indSampWithRemoval (AbsoluteTime 3) p1 True
  , Infection (AbsoluteTime 4) p2 p3
  , PopulationSample (AbsoluteTime 6) (asPeople [p2, p3, p4]) True
  ]

-- | Another test to make sure that disasters are handled.
demoFullEvents04 =
  [ Infection (AbsoluteTime 1) p1 p4
  , Infection (AbsoluteTime 2) p1 p2
  , indSampWithRemoval (AbsoluteTime 3) p1 True
  , Infection (AbsoluteTime 4) p2 p3
  , Infection (AbsoluteTime 5) p4 p5
  , PopulationSample (AbsoluteTime 6) (asPeople [p2, p3, p4]) True
  , Infection (AbsoluteTime 7) p5 p6
  , Infection (AbsoluteTime 8) p5 p7
  , PopulationSample (AbsoluteTime 9) (asPeople [p5, p6]) False
  ]

demoSampleEvents04 =
  [ Infection (AbsoluteTime 1) p1 p4
  , Infection (AbsoluteTime 2) p1 p2
  , indSampWithRemoval (AbsoluteTime 3) p1 True
  , Infection (AbsoluteTime 4) p2 p3
  , PopulationSample (AbsoluteTime 6) (asPeople [p2, p3, p4]) True
  , PopulationSample (AbsoluteTime 9) (asPeople [p5, p6]) False
  ]

eventHandlingTests = do
  describe "Post-simulation processing" $ do
    it "Extracting observed events" $ do
      let demoEvents =
            [ PopulationSample (AbsoluteTime 0.5) (asPeople []) True -- Because the first event is a null event it can be ignored!
            , Infection (AbsoluteTime 1.0) p1 p2
            , PopulationSample (AbsoluteTime 1.5) (asPeople []) True
            , PopulationSample (AbsoluteTime 2.0) (asPeople [p1, p2]) True
            ]
      (length demoEvents == 4) `shouldBe` True
      ((length <$> observedEvents (tail demoEvents)) == (Right 2)) `shouldBe`
        True
      ((length <$> observedEvents (demoEvents)) == (Right 2)) `shouldBe` True
      (observedEvents (demoEvents) == observedEvents (tail demoEvents)) `shouldBe`
        True
      (maybeEpidemicTree (demoEvents) == maybeEpidemicTree (tail demoEvents)) `shouldBe`
        True
  describe "Catastrophe definitions" $ do
    it "Check we can find a catastrophe" $ do
      (noScheduledEvent (AbsoluteTime 0) (AbsoluteTime 1) (Timed [])) `shouldBe`
        True
      (noScheduledEvent
         (AbsoluteTime 0)
         (AbsoluteTime 1)
         (Timed [(AbsoluteTime 2, 0.5)])) `shouldBe`
        True
      (noScheduledEvent
         (AbsoluteTime 0)
         (AbsoluteTime 1)
         (Timed [(AbsoluteTime 0.5, 0.5)])) `shouldBe`
        False
      (noScheduledEvent
         (AbsoluteTime 0)
         (AbsoluteTime 1)
         (Timed [(AbsoluteTime 2, 0.6), (AbsoluteTime 0.5, 0.5)])) `shouldBe`
        False
    it "Check we can find a particular catastrophe" $ do
      (firstScheduled (AbsoluteTime 1) (Timed [])) `shouldBe` Nothing
      (firstScheduled (AbsoluteTime 1) (Timed [(AbsoluteTime 2, 0.5)])) `shouldBe`
        Just (AbsoluteTime 2, 0.5)
      (firstScheduled (AbsoluteTime 1) (Timed [(AbsoluteTime 0.5, 0.5)])) `shouldBe`
        Nothing
      (firstScheduled
         (AbsoluteTime 1)
         (Timed [(AbsoluteTime 2, 0.6), (AbsoluteTime 0.5, 0.5)])) `shouldBe`
        Just (AbsoluteTime 2, 0.6)
      isNothing
        (asTimed
           [ (AbsoluteTime 2, 0.6 :: Rate)
           , (AbsoluteTime 0.5, 0.5)
           , (AbsoluteTime 1.5, 0.4)
           ]) `shouldBe`
        True
      (firstScheduled
         (AbsoluteTime 1)
         (Timed
            [ (AbsoluteTime 2, 0.6)
            , (AbsoluteTime 0.5, 0.5)
            , (AbsoluteTime 1.5, 0.4)
            ])) `shouldBe`
        Just (AbsoluteTime 2, 0.6)
    it "Works on a very specific case it seems to not like" $ do
      (noScheduledEvent
         (AbsoluteTime 2.28)
         (AbsoluteTime (2.28 + 0.42))
         (Timed [(AbsoluteTime 2.3, 0.9)])) `shouldBe`
        False
  describe "Disaster definitions" $ do
    it "Disasters are handled correctly" $ do
      ((Right $ [Observation e | e <- demoSampleEvents04]) ==
       (observedEvents demoFullEvents04)) `shouldBe`
        True
    it "Disasters can be simulated" $ do
      (Right demoSim) <-
        simulationWithFixedSeed
          (fromJust
             (BDSCOD.configuration
                (AbsoluteTime 0)
                (TimeDelta 4)
                False
                Nothing
                ( 1.3
                , 0.1
                , 0.1
                , [(AbsoluteTime 3, 0.5)]
                , 0.2
                , [(AbsoluteTime 3.5, 0.5)]
                , 1.0)))
          (allEvents BDSCOD.randomEvent)
      length demoSim > 1 `shouldBe` True
  describe "Extracting observed events" $ do
    it "unseq obs still extracted if no seq obs" $ do
      let noSequencedEvents =
            [ Infection (AbsoluteTime 4.1) p1 p2
            , Infection (AbsoluteTime 4.3) p2 p3
            , Infection (AbsoluteTime 4.5) p2 p4
            , indSampWithRemoval (AbsoluteTime 5.3) p1 False
            , StoppingTime (AbsoluteTime 6.0)
            ]
          expectedObs =
            [ Observation
                (indSampWithRemoval (AbsoluteTime 5.3) p1 False)
            ]
      isRight (observedEvents noSequencedEvents) `shouldBe` True
      ((Right expectedObs) == (observedEvents noSequencedEvents)) `shouldBe`
        True

helperFuncTests = do
  describe "Helpers in Utility" $ do
    it "the asTimed function works" $ do
      (isJust $ asTimed []) `shouldBe` True
      (isJust $ asTimed [(AbsoluteTime 0, 1)]) `shouldBe` True
      (isJust $ asTimed [(AbsoluteTime 0, 1), (AbsoluteTime 1, 3)]) `shouldBe`
        True
      (isJust $ asTimed [(AbsoluteTime 0, 3), (AbsoluteTime 1, 1)]) `shouldBe`
        True
      (isJust $ asTimed [(AbsoluteTime 1, 3), (AbsoluteTime 0, 1)]) `shouldBe`
        False
    let demoTimed =
          fromJust $
          asTimed
            [ (AbsoluteTime 0, 1.2)
            , (AbsoluteTime 1, 3.1)
            , (AbsoluteTime 2, 2.7)
            ]
     in do it "the cadlagValue function works" $ do
             (isJust $ cadlagValue demoTimed (AbsoluteTime (-1.0))) `shouldBe`
               False
             ((== 1.2) . fromJust $ cadlagValue demoTimed (AbsoluteTime 0.0)) `shouldBe`
               True
             ((== 1.2) . fromJust $ cadlagValue demoTimed (AbsoluteTime 0.5)) `shouldBe`
               True
             ((== 3.1) . fromJust $ cadlagValue demoTimed (AbsoluteTime 1.5)) `shouldBe`
               True
           it "the diracDeltaValue function works" $ do
             ((== 1.2) . fromJust $ diracDeltaValue demoTimed (AbsoluteTime 0)) `shouldBe`
               True
             (isJust $ diracDeltaValue demoTimed (AbsoluteTime 1)) `shouldBe`
               True
             (isJust $ diracDeltaValue demoTimed (AbsoluteTime 0.9)) `shouldBe`
               False
             (isJust $ diracDeltaValue demoTimed (AbsoluteTime 1.1)) `shouldBe`
               False
           it "the hasTime function works" $ do
             (hasTime demoTimed (AbsoluteTime 0)) `shouldBe` True
             (hasTime demoTimed (AbsoluteTime 0.5)) `shouldBe` False
             (hasTime demoTimed (AbsoluteTime 1)) `shouldBe` True
             (hasTime demoTimed (AbsoluteTime 1.5)) `shouldBe` False
           it "the nextTime function works" $ do
             (AbsoluteTime 0 ==
              (fromJust $ nextTime demoTimed (AbsoluteTime (-1)))) `shouldBe`
               True
             (AbsoluteTime 1 == (fromJust $ nextTime demoTimed (AbsoluteTime 0))) `shouldBe`
               True
             (AbsoluteTime 1 ==
              (fromJust $ nextTime demoTimed (AbsoluteTime 0.5))) `shouldBe`
               True
           it "the nextTime function handles the last time correctly" $ do
             isJust (nextTime demoTimed (AbsoluteTime 1.9)) `shouldBe` True
             isJust (nextTime demoTimed (AbsoluteTime 2.0)) `shouldBe` True
             isJust (nextTime demoTimed (AbsoluteTime 2.1)) `shouldBe` True
             isJust (nextTime demoTimed (AbsoluteTime 10.0)) `shouldBe` True
           it "the maybeNextTimed function works as expected" $ do
             let (Just tA) = asTimed [(AbsoluteTime 1, (1.1 :: Double)), (AbsoluteTime 3, 2.3)]
                 (Just tB) = asTimed [(AbsoluteTime 2, (1 :: Int))]
             maybeNextTimed tA tB (AbsoluteTime 0.5) == Just (AbsoluteTime 1.0,Left 1.1) `shouldBe` True
             maybeNextTimed tA tB (AbsoluteTime 1.5) == Just (AbsoluteTime 2.0,Right 1) `shouldBe` True
             maybeNextTimed tA tB (AbsoluteTime 2.5) == Just (AbsoluteTime 3.0,Left 2.3) `shouldBe` True
             isNothing (maybeNextTimed tA tB (AbsoluteTime 3.5)) `shouldBe` True

    it "shifted times work" $
      let sf =
            fromJust $
            asTimed [(AbsoluteTime (-1.0), 2.0), (AbsoluteTime 1, 3.0)]
          val1 = cadlagValue sf (AbsoluteTime 0)
          val2 = cadlagValue sf (AbsoluteTime (-2.0))
          val3 = cadlagValue sf (AbsoluteTime 1.5)
       in do isJust val1 `shouldBe` True
             val1 == Just 2.0 `shouldBe` True
             (not $ isJust val2) `shouldBe` True
             isJust val3 `shouldBe` True
             val3 == Just 3.0 `shouldBe` True
    it "the asTimed function returns nothing as expected" $ do
      (isJust $ asTimed [(AbsoluteTime 0.0, -1)]) `shouldBe` True
      (isJust $ asTimed [(AbsoluteTime 0.0, 1), (AbsoluteTime 1.0, -1)]) `shouldBe` True
      (isNothing $ asTimed [(AbsoluteTime 2.0, 1), (AbsoluteTime 1.0, -1)]) `shouldBe` True
      (isNothing $ asTimed [(AbsoluteTime 1.0, 1), (AbsoluteTime 1.0, -1)]) `shouldBe` True

bdscodTests =
  describe "Test BDSCOD simulations look reasonable" $
  do it "number of removed individual samples looks right" $
       let
         startTime = AbsoluteTime 0.0
         duration = TimeDelta 4.0
         atLeastCherry = True
         remProb = 0.5
         sampRate = 0.45
         occrRate = 0.45
         paramTuple = (2.0,0.1,sampRate,[],occrRate,[],remProb)
         (Just simConfig) =
           BDSCOD.configuration startTime duration atLeastCherry Nothing paramTuple

         -- It is helpful to be able to count the number of removed and
         -- non-removed individual samples.
         numIndSamp = length . filter isIndividualSample
         numIndSampSeq = length . filter indSampSeq . filter isIndividualSample
         numRemIndSamp = length . filter indSampRemoved . filter isIndividualSample

         -- carry out a hypothesis test that the observation looks reasonable in
         -- the sense that the correct proportion of individuals have been
         -- removed after sampling and that the correct proportion of sampled
         -- individuals have their viral genome sequenced.
         bTest a = do eEvents <- simulationWithSystem simConfig (allEvents BDSCOD.randomEvent)
                      isRight eEvents `shouldBe` True
                      let (Right events) = eEvents
                          x = numRemIndSamp events
                          y = numIndSampSeq events
                          n = numIndSamp events
                      -- print ((x, n), (y, n))
                      binomialTest x n remProb a `shouldBe` True
                      binomialTest y n (sampRate / (sampRate + occrRate)) a `shouldBe` True

       in do
         -- this test will occassionally fail because it relies on random
         -- numbers.
         replicateM_ 10 (bTest 0.01) -- 10 tests at 0.01-level.

simTypeTests =
  describe "Test Data.Simulation PRNG helpers" $
  do it "check genIOFromFixed always gives same result" $
       do g1 <- genIOFromFixed
          u11 <- MWC.uniform g1 :: IO Double
          u12 <- MWC.uniform g1 :: IO Double

          g2 <- genIOFromFixed
          u21 <- MWC.uniform g2 :: IO Double
          u22 <- MWC.uniform g2 :: IO Double

          u11 == u21 `shouldBe` True
          u11 /= u22 `shouldBe` True
          u12 == u22 `shouldBe` True

     it "check genIOFromSystem always gives different results" $
       do g1 <- genIOFromSystem
          u11 <- MWC.uniform g1 :: IO Double
          u12 <- MWC.uniform g1 :: IO Double

          g2 <- genIOFromSystem
          u21 <- MWC.uniform g2 :: IO Double
          u22 <- MWC.uniform g2 :: IO Double

          u11 /= u12 `shouldBe` True
          u11 /= u21 `shouldBe` True
          u11 /= u22 `shouldBe` True
          u12 /= u21 `shouldBe` True
          u12 /= u22 `shouldBe` True
          u21 /= u22 `shouldBe` True

     it "check genIOFromWord32 works as expected" $
       do g1 <- genIOFromWord32 1
          u11 <- MWC.uniform g1 :: IO Double
          u12 <- MWC.uniform g1 :: IO Double

          g2 <- genIOFromWord32 1
          u21 <- MWC.uniform g2 :: IO Double
          u22 <- MWC.uniform g2 :: IO Double

          g3 <- genIOFromWord32 2
          u31 <- MWC.uniform g3 :: IO Double
          u32 <- MWC.uniform g3 :: IO Double

          u11 == u21 `shouldBe` True
          u11 /= u22 `shouldBe` True
          u12 == u22 `shouldBe` True

          u11 /= u12 `shouldBe` True
          u11 /= u31 `shouldBe` True
          u11 /= u32 `shouldBe` True
          u12 /= u31 `shouldBe` True
          u12 /= u32 `shouldBe` True
          u31 /= u32 `shouldBe` True

inhomExpTests =
  describe "Test the inhomogeneous exponential variate generator" $
  let rate1 = 2.0
      sF1 = fromJust $ asTimed [(AbsoluteTime 0, rate1)]
      mean1 = 1 / rate1
      var1 = 1 / (rate1 ** 2.0)
      sF2 =
        fromJust $ asTimed [(AbsoluteTime 0, 1e-10), (AbsoluteTime 1, rate1)]
      mean2 = 1 / rate1 + 1
      var2 = var1
   in do it "check we can get a positive variate out" $ do
           gen <- genIOFromSystem
           u1 <- MWC.uniform gen :: IO Double
           (u1 > 0) `shouldBe` True
           (Just x1) <- inhomExponential sF1 (AbsoluteTime 0) gen
           (x1 > AbsoluteTime 0) `shouldBe` True
           (x1 < AbsoluteTime 100) `shouldBe` True
           True `shouldBe` True
         it "check the mean and variance look sensible" $ do
           gen <- genIOFromSystem
           xBoxed <-
             V.replicateM 20000 (inhomExponential sF1 (AbsoluteTime 0) gen)
           let x = fmap (\(Just (AbsoluteTime t)) -> t) xBoxed
           withinNPercent 5 (mean x) mean1 `shouldBe` True
           withinNPercent 5 (variance x) var1 `shouldBe` True
         it "check the mean and variance look sensible with delay" $ do
           gen <- genIOFromSystem
           xBoxed <-
             V.replicateM 20000 (inhomExponential sF2 (AbsoluteTime 0) gen)
           let x = fmap (\(Just (AbsoluteTime t)) -> t) xBoxed
           withinNPercent 5 (mean x) mean2 `shouldBe` True
           withinNPercent 5 (variance x) var2 `shouldBe` True

illFormedTreeTest :: SpecWith ()
illFormedTreeTest =
  describe "Prevent the simulator returning a broken tree" $ do
    let simDuration = TimeDelta 0.2
        simLambda = 3.2
        simMu = 0.3
        simPsi = 0.3
        simRho = 0.15
        simRhoTime = AbsoluteTime 2.6
        simOmega = 0.3
        simNu = 0.15
        simNuTime = AbsoluteTime 3.0
        simRemProb = 1.0
        simParams =
          ( simLambda
          , simMu
          , simPsi
          , [(simRhoTime, simRho)]
          , simOmega
          , [(simNuTime, simNu)]
          , simRemProb)
        simConfig = BDSCOD.configuration (AbsoluteTime 0) simDuration True Nothing simParams
     in do it "stress testing the observed events function" $ do
             isJust simConfig `shouldBe` True
             null (observedEvents []) `shouldBe` True
             (Right simEvents) <-
               simulationWithFixedSeed (fromJust simConfig) (allEvents BDSCOD.randomEvent)
             any isReconTreeLeaf simEvents `shouldBe` True
             let eitherObsEs = observedEvents simEvents
             isRight eitherObsEs `shouldBe` True
             either (const False) (\oes -> length oes > 1) eitherObsEs `shouldBe` True
           it "check leaves are recognised correctly" $ do
             let (absT, person) = (AbsoluteTime 1.0, Person (Identifier 1))
                 infEvent = Infection absT person person
                 remEvent = Removal absT person
                 sampIndv = indSampWithRemoval absT person
                 popSampEmpty = PopulationSample absT $ asPeople []
                 popSampPerson = PopulationSample absT $ asPeople [person]
             isReconTreeLeaf infEvent `shouldBe` False
             isReconTreeLeaf remEvent `shouldBe` False
             isReconTreeLeaf (sampIndv True) `shouldBe` True
             isReconTreeLeaf (sampIndv False) `shouldBe` False
             isReconTreeLeaf (popSampEmpty True) `shouldBe` False
             isReconTreeLeaf (popSampEmpty False) `shouldBe` False
             isReconTreeLeaf (popSampPerson True) `shouldBe` True
             isReconTreeLeaf (popSampPerson False) `shouldBe` False

resultAA = demoSampleEvents01

resultAB =
  [ Infection (AbsoluteTime 1) p1 p2
  , indSampWithRemoval (AbsoluteTime 3) p1 True
  , Infection (AbsoluteTime 4) p2 p4
  , indSampWithRemoval (AbsoluteTime 6) p4 True
  , PopulationSample (AbsoluteTime 12) (asPeople [p2, p6]) False
  , indSampWithRemoval (AbsoluteTime 12) p5 True
  , PopulationSample (AbsoluteTime 17.0) (asPeople []) False
  ]

resultBA =
  [ Infection (AbsoluteTime 1) p1 p2
  , indSampWithRemoval (AbsoluteTime 3) p1 True
  , Infection (AbsoluteTime 4) p2 p4
  , indSampWithRemoval (AbsoluteTime 6) p4 True
  , indSampWithRemoval (AbsoluteTime 8) p2 False
  , indSampWithRemoval (AbsoluteTime 11) p6 False
  , PopulationSample (AbsoluteTime 13) (asPeople [p5]) True
  ]

resultBB =
  [ Infection (AbsoluteTime 1.0) p1 p2
  , indSampWithRemoval (AbsoluteTime 3.0) (Person (Identifier 1)) True
  , Infection (AbsoluteTime 4.0) p2 p4
  , indSampWithRemoval (AbsoluteTime 6.0) (Person (Identifier 4)) True
  , PopulationSample (AbsoluteTime 12) (asPeople [p2, p6]) False
  , PopulationSample (AbsoluteTime 13) (asPeople [p5]) True
  , PopulationSample (AbsoluteTime 17) (asPeople []) False
  ]

aggregationTests =
  describe "Aggregation functionality tests" $ do
    it "check it does nothing unless it needs to" $
      let demoObs1 = [Observation ee | ee <- demoFullEvents01]
          demoObs2 = [Observation ee | ee <- demoSampleEvents01]
      in do
        (aggregated [] [] demoObs1) == demoObs1 `shouldBe` True
        (aggregated [] [] demoObs2) == demoObs2 `shouldBe` True
    it "check relevant intervals are processed correctly" $
      let demoObs = [Observation ee | ee <- demoSampleEvents01]
          demoSeqInts = asConsecutiveIntervals1 [AbsoluteTime 10, AbsoluteTime 13]
          demoUnseqInts = asConsecutiveIntervals1 [AbsoluteTime 7, AbsoluteTime 12, AbsoluteTime 17]
      in do
        aggregated [] [] demoObs == (map Observation resultAA) `shouldBe` True
        aggregated [] demoUnseqInts demoObs == (map Observation resultAB) `shouldBe` True
        aggregated demoSeqInts [] demoObs == (map Observation resultBA) `shouldBe` True
        aggregated demoSeqInts demoUnseqInts demoObs == (map Observation resultBB) `shouldBe` True

inhomogeneousBDSTest =
  describe "InhomogeneousBDS module tests" $ do
    it "Check the observedEvents filters out removals" $
      let demoAllEvents =
            [ Infection (AbsoluteTime 0.1) p1 p2
            , indSampWithRemoval (AbsoluteTime 0.2) p1 True
            , Removal (AbsoluteTime 0.3) p3
            , indSampWithRemoval (AbsoluteTime 0.4) p2 True
            ]
          demoObsEvents =
            [ Infection (AbsoluteTime 0.1) p1 p2
            , indSampWithRemoval (AbsoluteTime 0.2) p1 True
            , indSampWithRemoval (AbsoluteTime 0.4) p2 True
            ]
          compObsEvents = observedEvents demoAllEvents
       in do (compObsEvents == (Right [Observation e | e <- demoObsEvents])) `shouldBe` True

helperTypeTests = do
  describe "Helpers for working with the types" $ do
    it "the asTimed function works" $ do
      (isJust $ asTimed []) `shouldBe` True
      (isJust $ asTimed [(AbsoluteTime 0, 1)]) `shouldBe` True
      (isJust $ asTimed [(AbsoluteTime 0, 1), (AbsoluteTime 1, 3)]) `shouldBe`
        True
      (isJust $ asTimed [(AbsoluteTime 0, 3), (AbsoluteTime 1, 1)]) `shouldBe`
        True
      (isJust $ asTimed [(AbsoluteTime 1, 3), (AbsoluteTime 0, 1)]) `shouldBe`
        False
    let demoTimed =
          fromJust $
          asTimed
            [ (AbsoluteTime 0, 1.2)
            , (AbsoluteTime 1, 3.1)
            , (AbsoluteTime 2, 2.7)
            ]
     in do it "the cadlagValue function works" $ do
             (isJust $ cadlagValue demoTimed (AbsoluteTime (-1.0))) `shouldBe`
               False
             ((== 1.2) . fromJust $ cadlagValue demoTimed (AbsoluteTime 0.0)) `shouldBe`
               True
             ((== 1.2) . fromJust $ cadlagValue demoTimed (AbsoluteTime 0.5)) `shouldBe`
               True
             ((== 3.1) . fromJust $ cadlagValue demoTimed (AbsoluteTime 1.5)) `shouldBe`
               True
           it "the diracDeltaValue function works" $ do
             ((== 1.2) . fromJust $ diracDeltaValue demoTimed (AbsoluteTime 0)) `shouldBe`
               True
             (isJust $ diracDeltaValue demoTimed (AbsoluteTime 1)) `shouldBe`
               True
             (isJust $ diracDeltaValue demoTimed (AbsoluteTime 0.9)) `shouldBe`
               False
             (isJust $ diracDeltaValue demoTimed (AbsoluteTime 1.1)) `shouldBe`
               False
           it "the hasTime function works" $ do
             (hasTime demoTimed (AbsoluteTime 0)) `shouldBe` True
             (hasTime demoTimed (AbsoluteTime 0.5)) `shouldBe` False
             (hasTime demoTimed (AbsoluteTime 1)) `shouldBe` True
             (hasTime demoTimed (AbsoluteTime 1.5)) `shouldBe` False
           it "the nextTime function works" $ do
             (AbsoluteTime 0 ==
              (fromJust $ nextTime demoTimed (AbsoluteTime (-1)))) `shouldBe`
               True
             (AbsoluteTime 1 == (fromJust $ nextTime demoTimed (AbsoluteTime 0))) `shouldBe`
               True
             (AbsoluteTime 1 ==
              (fromJust $ nextTime demoTimed (AbsoluteTime 0.5))) `shouldBe`
               True
    it "shifted times work" $
      let sf =
            fromJust $
            asTimed [(AbsoluteTime (-1.0), 2.0), (AbsoluteTime 1, 3.0)]
          val1 = cadlagValue sf (AbsoluteTime 0)
          val2 = cadlagValue sf (AbsoluteTime (-2.0))
          val3 = cadlagValue sf (AbsoluteTime 1.5)
       in do isJust val1 `shouldBe` True
             val1 == Just 2.0 `shouldBe` True
             (not $ isJust val2) `shouldBe` True
             isJust val3 `shouldBe` True
             val3 == Just 3.0 `shouldBe` True

jsonTests = do
  describe "Converting to and from JSON" $ do
    it "Conversion of Timed Rate" $ do
      let demoObj =
            Timed [(AbsoluteTime 0.0, 1.0), (AbsoluteTime 1.0, 1.0)] :: Timed Rate
          (Timed demoVals) = demoObj
          demoJson = "[[0,1],[1,1]]"
          encodedObj = Json.encode demoObj
          decodedJson = Json.decode demoJson :: Maybe (Timed Rate)
       in do True `shouldBe` True
             let (Timed foo) = demoObj
              in demoVals == foo `shouldBe` True
             encodedObj == demoJson `shouldBe` True
             isJust decodedJson `shouldBe` True
             let (Timed bar) = fromJust decodedJson
              in demoVals == bar `shouldBe` True

equalBuilders :: BBuilder.Builder -> BBuilder.Builder -> Bool
equalBuilders a b = BBuilder.toLazyByteString a == BBuilder.toLazyByteString b

rightEqualBuilders a b = case (a, b) of
  (Right av, Right bv) -> equalBuilders av bv
  _ -> False

newickTests =
  let p1 = Person (Identifier 1)
      p2 = Person (Identifier 2)
      p3 = Person (Identifier 3)
      absT = AbsoluteTime
      ps = asPeople [p1, p2]
      maybeEpiTree =
        maybeEpidemicTree
          [ Infection (AbsoluteTime 1) p1 p2
          , Infection (AbsoluteTime 2) p2 p3
          , PopulationSample (AbsoluteTime 3) (asPeople [p1, p3]) True
          , Removal (AbsoluteTime 4) p2
          ]
      maybeEpiTree' =
        maybeEpidemicTree
          [ Infection (AbsoluteTime 1) p1 p2
          , Infection (AbsoluteTime 2) p2 p3
          , PopulationSample (AbsoluteTime 3) (asPeople [p1, p3]) True
          , indSampWithRemoval (AbsoluteTime 4) p2 True
          ]
      maybeEpiTree'' =
        maybeEpidemicTree
          [ Infection (AbsoluteTime 1) p1 p2
          , Infection (AbsoluteTime 2) p2 p3
          , PopulationSample (AbsoluteTime 3) (asPeople [p1, p3]) False
          , indSampWithRemoval (AbsoluteTime 4) p2 True
          ]
   in describe "Writing to Newick" $ do
        it "equalBuilders works as expected" $ do
          equalBuilders (BBuilder.charUtf8 ':') (BBuilder.charUtf8 ':') `shouldBe`
            True
          equalBuilders (BBuilder.charUtf8 'a') (BBuilder.charUtf8 ':') `shouldBe`
            False
        it "derivedFrom works as expected" $ do
          let p1 = Person (Identifier 1)
              p2 = Person (Identifier 2)
              p3 = Person (Identifier 3)
              e = [Infection (AbsoluteTime 0.3) p1 p2]
          derivedFrom p1 e == derivedFrom p2 e `shouldBe` True
          derivedFrom p1 e /= derivedFrom p3 e `shouldBe` True
          derivedFrom p1 e /= [] `shouldBe` True
          null (derivedFrom p3 e) `shouldBe` True
          derivedFrom p1 e == e `shouldBe` True
          let foo =
                derivedFrom
                  (Person (Identifier 1))
                  [ Infection
                      (AbsoluteTime 0.3)
                      (Person (Identifier 1))
                      (Person (Identifier 2))
                  , indSampWithRemoval
                      (AbsoluteTime 0.7)
                      (Person (Identifier 1))
                      True
                  ]
          let bar =
                derivedFrom
                  (Person (Identifier 2))
                  [ Infection
                      (AbsoluteTime 0.3)
                      (Person (Identifier 1))
                      (Person (Identifier 2))
                  , indSampWithRemoval
                      (AbsoluteTime 0.7)
                      (Person (Identifier 1))
                      True
                  ]
          foo == bar `shouldBe` True
        it "maybeEpidemicTree works as expected: 1" $ do
          let e1 = Removal (AbsoluteTime 1) (Person (Identifier 1))
          maybeEpidemicTree [e1] == Right (Leaf e1) `shouldBe` True
          let t1 =
                maybeEpidemicTree
                  [ Infection
                      (AbsoluteTime 0.3)
                      (Person (Identifier 1))
                      (Person (Identifier 2))
                  , indSampWithRemoval
                      (AbsoluteTime 0.6)
                      (Person (Identifier 2))
                      True
                  , indSampWithRemoval
                      (AbsoluteTime 0.7)
                      (Person (Identifier 1))
                      True
                  ]
          let t2 =
                Right
                  (Branch
                     (Infection
                        (AbsoluteTime 0.3)
                        (Person (Identifier 1))
                        (Person (Identifier 2)))
                     (Leaf
                        (indSampWithRemoval
                           (AbsoluteTime 0.7)
                           (Person (Identifier 1))
                           True))
                     (Leaf
                        (indSampWithRemoval
                           (AbsoluteTime 0.6)
                           (Person (Identifier 2))
                           True)))
          t1 == t2 `shouldBe` True
          maybeEpidemicTree
            [ Infection
                (AbsoluteTime 0.3)
                (Person (Identifier 1))
                (Person (Identifier 2))
            ] ==
            Right
              (Branch
                 (Infection
                    (AbsoluteTime 0.3)
                    (Person (Identifier 1))
                    (Person (Identifier 2)))
                 (Shoot (Person (Identifier 1)))
                 (Shoot (Person (Identifier 2)))) `shouldBe`
            True
          maybeEpidemicTree
            [ Infection
                (AbsoluteTime 0.3)
                (Person (Identifier 1))
                (Person (Identifier 2))
            , indSampWithRemoval (AbsoluteTime 0.7) (Person (Identifier 1)) True
            ] ==
            Right
              (Branch
                 (Infection
                    (AbsoluteTime 0.3)
                    (Person (Identifier 1))
                    (Person (Identifier 2)))
                 (Leaf
                    (indSampWithRemoval
                       (AbsoluteTime 0.7)
                       (Person (Identifier 1))
                       True))
                 (Shoot (Person (Identifier 2)))) `shouldBe`
            True
          let trickyEvents =
                [ Infection
                    (AbsoluteTime 0.3)
                    (Person (Identifier 1))
                    (Person (Identifier 2))
                , Infection
                    (AbsoluteTime 0.4)
                    (Person (Identifier 2))
                    (Person (Identifier 3))
                , indSampWithRemoval
                    (AbsoluteTime 0.6)
                    (Person (Identifier 3))
                    True
                , indSampWithRemoval
                    (AbsoluteTime 0.7)
                    (Person (Identifier 1))
                    True
                ]
          isRight (maybeEpidemicTree trickyEvents) `shouldBe` True
        it "maybeEpidemicTree works as expected: 2" $ do
          let p1 = Person (Identifier 1)
              p2 = Person (Identifier 2)
              demoEvents =
                [ PopulationSample (AbsoluteTime 0.5) (asPeople []) True -- Because the first event is a null event it can be ignored!
                , Infection (AbsoluteTime 1.0) p1 p2
                , PopulationSample (AbsoluteTime 1.5) (asPeople []) True
                , PopulationSample (AbsoluteTime 2.0) (asPeople [p1, p2]) True
                ]
          (length demoEvents == 4) `shouldBe` True
          (maybeEpidemicTree demoEvents == maybeEpidemicTree (tail demoEvents)) `shouldBe`
            True

        it "asNewickString works for ReconstructedTree: 1" $ do
          isRight
            (asNewickString
               (AbsoluteTime 0, Person (Identifier 1))
               (RLeaf
                  (Observation (indSampWithRemoval
                     (AbsoluteTime 1)
                     (Person (Identifier 1))
                     True)))) `shouldBe`
            True
          let trickyEvents =
                [ Infection
                    (AbsoluteTime 0.3)
                    (Person (Identifier 1))
                    (Person (Identifier 2))
                , Infection
                    (AbsoluteTime 0.4)
                    (Person (Identifier 2))
                    (Person (Identifier 3))
                , indSampWithRemoval
                    (AbsoluteTime 0.6)
                    (Person (Identifier 3))
                    True
                , indSampWithRemoval
                    (AbsoluteTime 0.7)
                    (Person (Identifier 1))
                    True
                ]
              et = maybeEpidemicTree trickyEvents :: Either String EpidemicTree
              rt = reconstructedTree =<< et :: Either String ReconstructedTree
              maybeNewickPair = asNewickString (AbsoluteTime 0, Person (Identifier 1)) =<< rt
              newickTarget = BBuilder.stringUtf8 "(p1:0.39999999999999997,p3:0.3):0.3;"
          isRight maybeNewickPair `shouldBe` True
          equalBuilders newickTarget (fromJust $ either2Maybe maybeNewickPair) `shouldBe`
            True
          let catasNewick =
                (asNewickString
                   (AbsoluteTime 0, Person (Identifier 1))
                   (RLeaf
                      (Observation (PopulationSample
                         (AbsoluteTime 1)
                         (asPeople
                            [Person (Identifier 1), Person (Identifier 2)])
                         True))))
          let catasTarget = BBuilder.stringUtf8 "p1&p2:1.0;"
          equalBuilders catasTarget (fromJust . either2Maybe $ catasNewick) `shouldBe` True

        it "asNewickString works for ReconstructedTree: 2" $ do

          let toNS = asNewickString (absT 0.0, p1)
              leafBaseCase1 = toNS $ RLeaf (Observation (IndividualSample (absT 1.0) p1 True True))
              leafBaseCase1Sol = Right $ BBuilder.stringUtf8 "p1:1.0;"
          rightEqualBuilders leafBaseCase1 leafBaseCase1Sol `shouldBe` True

          let leafBaseCase2 = toNS $ RLeaf (Observation (IndividualSample (absT 1.0) p1 False True))
              leafBaseCase3 = toNS $ RLeaf (Observation (IndividualSample (absT 1.0) p1 True False))
          isLeft leafBaseCase2 `shouldBe` True
          isLeft leafBaseCase3 `shouldBe` True

          let subTree = RLeaf (Observation (IndividualSample (absT 2.0) p2 True True))
              leafBaseCase4 = toNS $ RBurr (Observation (IndividualSample (absT 1.0) p1 True True)) (Just subTree)
              leafBaseCase4Sol = Right $ BBuilder.stringUtf8 "(p2:1.0)p1:1.0;"
              leafBaseCase5 = toNS $ RBurr (Observation (IndividualSample (absT 1.0) p1 True True)) Nothing
              leafBaseCase5Sol = Right $ BBuilder.stringUtf8 "p1:1.0;"
              leafBaseCase6 = toNS $ RBurr (Observation (Removal undefined undefined)) (Just subTree)
          rightEqualBuilders leafBaseCase4 leafBaseCase4Sol `shouldBe` True
          rightEqualBuilders leafBaseCase4 leafBaseCase5 `shouldBe` False
          rightEqualBuilders leafBaseCase5 leafBaseCase5Sol `shouldBe` True
          isLeft leafBaseCase6 `shouldBe` True

terminationTests1 =
  describe "Termination handling tests: InhomogeneousBDSCOD" $ do
    let duration = TimeDelta 2.0
        startTime = AbsoluteTime 0.0
        birthRateSpec = [(AbsoluteTime 0.0, 1.5), (AbsoluteTime 0.5, 0.5)]
        deathRateSpec = [(AbsoluteTime 0.0, 0.4)]
        sampRateSpec = [(AbsoluteTime 0.0, 0.1)]
        occRateSpec = [(AbsoluteTime 0.0, 0.1)]
        seqSched = [(AbsoluteTime 0.9, 0.1)]
        unseqSched = [(AbsoluteTime 0.5, 0.4), (AbsoluteTime 0.75, 0.5)]
        sampOccRemProbSpec = [(AbsoluteTime 0.0, 1.0)]
        ratesAndProbs = (birthRateSpec,deathRateSpec,sampRateSpec,seqSched,occRateSpec,unseqSched, sampOccRemProbSpec)
        conf maybeTH = fromJust $ InhomBDSCOD.configuration startTime duration True maybeTH ratesAndProbs
        -- We need one simulation configuration for each of the termination
        -- handlers that we want to test.
        simConfigNothing  = conf Nothing
        simConfigNever = conf (Just (const False, const ()))
        simConfigAlways = conf (Just (const True, const ()))
        numDeadThreshold = 3
        simConfigSometimes = conf (Just ((>numDeadThreshold) . InhomBDSCOD.ipNumRemovedByDeath,
                                         \es -> length [() | Removal _ _ <- es]))
        allEventsFunc = allEvents InhomBDSCOD.randomEvent
    it "test simulation works without hander" $
      do
        (Right esNothing) <- simulationWithFixedSeed simConfigNothing allEventsFunc
        -- There should always be at least one event in the simulation.
        (length esNothing > 0) `shouldBe` True
    it "test simulation works with handler that does not trigger" $
      do
        (Right esNothing) <- simulationWithFixedSeed simConfigNothing allEventsFunc
        (Right esNever) <- simulationWithFixedSeed simConfigNever allEventsFunc
        -- If the handler never triggers this should look the same as not having
        -- the event handler.
        (all id $ zipWith (==) esNothing esNever) `shouldBe` True
    it "test simulation works with handler that always triggers" $
      do
        -- If the handler always triggers this should always return the summary.
        replicateM_ 30
          (do esAlways <- simulationWithSystem simConfigAlways allEventsFunc
              esAlways == Left (Just ()) `shouldBe` True
           )
    it "test simulation works with handler that sometimes triggers" $
      do
        -- If the handler only sometimes triggers then we need to test both
        -- branches.
        replicateM_ 30
          (do esSometimes <- simulationWithSystem simConfigSometimes allEventsFunc
              case esSometimes of
                -- If the termination handler did not trigger the number of
                -- removals should not exceed the threshold allowed by the
                -- termination handler
                (Right es) -> do length [() | Removal _ _ <- es] <= numDeadThreshold `shouldBe` True
                -- If the termination handler did trigger then we should know
                -- exactly how many removals there was.
                (Left (Just n)) -> do n == numDeadThreshold + 1 `shouldBe` True
                (Left Nothing) -> True `shouldBe` False -- this branch should not be reached.
                )

main :: IO ()
main =
  hspec $ do
    eventHandlingTests
    helperFuncTests
    inhomExpTests
    illFormedTreeTest
    inhomogeneousBDSTest
    helperTypeTests
    jsonTests
    newickTests
    aggregationTests
    simTypeTests
    terminationTests1
    bdscodTests


import Control.Arrow
import System.Directory
import Data.Char
import qualified Data.Text as T
import qualified Data.Text.IO as T
import Data.List
import qualified Data.Set as S
import qualified Data.Map.Strict as M
import qualified Data.IntMap.Strict as IM
import NLP.Stemmer
import Data.Ord
import Data.Function
import Data.SVM
import qualified Numeric.LinearAlgebra.Data as LA
import qualified Numeric.LinearAlgebra as LA

-- Ratings are embedded in the file names, e.g., 100_7.txt has a rating of 7/10. Extract it.
getRating :: (Fractional a, Read a) => String -> a
getRating = (\x -> (x-5.5)/4.5) . read . takeWhile (/= '.') . drop 1 . dropWhile (/= '_')

-- Get the rating and text from all files in a directory.
readSamples :: (Fractional a, Read a) => String -> IO [(a, T.Text)]
readSamples dir = mapM ((sequence .) $ getRating &&& T.readFile . (dir++)) . filter ((/= '.') . head) =<< getDirectoryContents dir

-- All of the training data, which is located across two directories.
trainingSamples :: (Fractional a, Read a) => IO [(a, T.Text)]
trainingSamples = do
  pos <- readSamples "aclImdb/train/pos/"
  neg <- readSamples "aclImdb/train/neg/"
  return $ takeEvery 5 (pos ++ neg)
validationSamples :: (Fractional a, Read a) => IO [(a, T.Text)]
validationSamples = do
  pos <- readSamples "aclImdb/train/pos/"
  neg <- readSamples "aclImdb/train/neg/"
  return $ takeEvery 10 (undefined:(pos ++ neg))

-- Words to ignore
stoplist :: S.Set T.Text
stoplist = S.fromList $ map T.pack ["i", "me", "my", "she", "her", "he", "him", "his", "they", "them", "their", "it", "who", "this", "that", "which",
  "", "br", "the", "a", "an", "is", "was", "be", "have", "are", "had", "were", "do", "did",
  "and", "or", "to", "in", "for", "with", "as", "one", "at", "by", "about", "from", "other", "into", "when", "would", "than",
  "also", "how", "becaus", "could", "after", "thing", "dont", "year", "two", "after", "peopl", "get", "where", "did", "off"]

minimalStoplist = S.fromList $ map T.pack ["the", "a", "an", "to", "is", "that", ""]
negateList = S.fromList $ map T.pack ["not", "isnt"]

-- Word splitter which ignores puctuation and doesn't care about capitalization or apostrophes.
words' :: T.Text -> [T.Text]
words' = map (stem' English) . filter (not . T.null) . T.split (not . isAlpha) . T.toLower . T.filter (/= '\'')

data PreFilterMode = Flip | Drop | Pass

-- Apply +/- modifier to word list
preFilter :: Num a => PreFilterMode -> [T.Text] -> [(a, T.Text)]
preFilter _ [] = []
preFilter Flip (w:ws)
  | w `S.member` negateList = alterHead (first negate) (preFilter Flip ws)
  | w `S.member` minimalStoplist = preFilter Flip ws
  | T.pack "un" `T.isPrefixOf` w = (-1, T.drop 2 w) : preFilter Flip ws
  | T.pack "nt" `T.isSuffixOf` w = (-1, T.dropEnd 2 w) : preFilter Flip ws
  | otherwise = (1, w) : preFilter Flip ws
  where alterHead f [] = []
        alterHead f (x:xs) = (f x) : xs
preFilter Drop (w:ws)
  | w `S.member` negateList = preFilter Drop ws
  | w `S.member` minimalStoplist = preFilter Drop ws
  | T.pack "un" `T.isPrefixOf` w = (1, T.drop 2 w) : preFilter Drop ws
  | T.pack "nt" `T.isSuffixOf` w = (1, T.dropEnd 2 w) : preFilter Drop ws
  | otherwise = (1, w) : preFilter Drop ws
preFilter Pass ws = zip (repeat 1) ws

-- Unify the number of occurences
bag :: Num a => [(a, T.Text)] -> [(a, T.Text)]
bag = map (sum . map fst &&& snd . head) . groupBy (curry $ uncurry (==) . both snd) . sortBy (comparing snd)

-- The stemming method from NLP.Stemmer, wrapped up to use Text insead of String
stem' :: Stemmer -> T.Text -> T.Text
stem' lang = T.pack . stem lang . T.unpack

-- The set of words in a dataset that meet the common threshold
common :: PreFilterMode -> Int -> [(a, T.Text)] -> S.Set T.Text
common mode thresh dataset = S.fromList . map snd . filter (\w -> fst w * thresh >= datalen) $ (df mode (`S.notMember` stoplist) dataset)
  where datalen = length dataset

-- Number of documents with each word not excluded by filter
df :: PreFilterMode -> (T.Text -> Bool) -> [(a, T.Text)] -> [(Int, T.Text)]
df mode wf = map (length &&& head) . group . sort . concatMap (map head . group . sort . filter wf . map snd . preFilter mode . words' . snd)

df' :: PreFilterMode -> (T.Text -> Bool) -> [(a, T.Text)] -> [(Float, T.Text)]
df' mode wf = map (log . balance . sum . map fst &&& snd . head) .
              groupBy ((==) `on` snd) . sortBy (comparing snd) .
              concatMap (bag . filter (wf . snd) . preFilter mode . words' . snd)

-- Calculates the Delta Inverse Document Frequency of the dataset
didf :: (Num a, Ord a) => PreFilterMode -> S.Set T.Text -> [(a, T.Text)] -> M.Map T.Text Float
didf mode legalWords dataset = uncurry (M.unionWith (+)) . both (M.fromList . map swap) .
    (                     df' mode (`S.member` legalWords) . filter ((>0) . fst) &&&
     map (first negate) . df' mode (`S.member` legalWords) . filter ((<0) . fst)) $ dataset

-- Converts a review into a vector for linear algebra
t2v :: PreFilterMode -> M.Map T.Text Int -> M.Map T.Text Double -> T.Text -> LA.Vector Double
t2v mode lut weights t = LA.fromList $ map (flip (IM.findWithDefault 0) (t2v' mode lut weights t)) [0 .. length lut - 1]

-- Converts a review into a vector for SVM, which is a different format
t2v' :: PreFilterMode -> M.Map T.Text Int -> M.Map T.Text Double -> T.Text -> IM.IntMap Double
t2v' mode lut weights t = IM.fromList . map (\(x, w) -> (lut M.! w, realToFrac $ weights M.! w * x)) .
                          filter ((`M.member` lut) . snd) . bag . preFilter mode . words' $ t

-- Converts a dataset into a term frequency matrix for linear algebra
data2m :: PreFilterMode -> M.Map T.Text Int -> M.Map T.Text Double -> [(a, T.Text)] -> LA.Matrix Double
data2m mode lut weights = LA.fromRows . map (t2v mode lut weights . snd)

-- Converts a matrix into a SVM problem
m2p :: [(Double, a)] -> LA.Matrix Double -> [(Double, IM.IntMap Double)]
m2p dataset = zip (map (clamp . fst) dataset) . map (\row -> IM.fromList $ map (id &&& (row LA.!)) [0 .. LA.size row - 1]) . LA.toRows

-- Applies a SVM model to a problem
checkProblem :: Model -> [(Double, IM.IntMap Double)] -> IO [(Double, Double)]
checkProblem model = sequence . map (sequence . second (predict model))

showReport :: (Ord a, Ord b, Num a, Num b) => [(a, b)] -> String
showReport results = show a ++ "/" ++ show n ++ " (" ++ show (100 * fromIntegral a / fromIntegral n) ++ "%)"
  where a = (length $ filter (\(a, b) -> (a>0) == (b>0)) results)
        n = (length results)

printReport :: (Ord a, Ord b, Num a, Num b) => (c, [(a, b)], [(a, b)]) -> IO ()
printReport (_, trainingResults, testingResults) = do
  putStrLn $ "\ttraining accuracy: " ++ showReport trainingResults ++ "\tvalidation accuracy: " ++ showReport testingResults

runExperiment :: (PreFilterMode, Int, Int, KernelType, Algorithm, [(Double, T.Text)], [(Double, T.Text)]) ->
                 IO (S.Set T.Text, [(Double, Double)], [(Double, Double)])
runExperiment (filterMode, commonThreshold, nSingularValues, svmKernel, svmAlgorithm, trainingDataset, testingDataset) = do
  let legalWords = common filterMode commonThreshold trainingDataset
  let weights = realToFrac <$> didf filterMode legalWords trainingDataset
  let wordIndexes = M.fromList $ zip (S.toList legalWords) [0..]
  let (u, s, v) = LA.thinSVD (data2m filterMode wordIndexes weights trainingDataset)
  let u' = u LA.?? (LA.All, LA.Take nSingularValues)
  let s' = LA.subVector 0 nSingularValues s
  let v' = v LA.?? (LA.All, LA.Take nSingularValues)
  let testingDataset' = m2p testingDataset $ data2m filterMode wordIndexes weights testingDataset LA.<> v'
  let trainingDataset' = m2p trainingDataset (u' LA.<> LA.diag s')
  svmModel <- withPrintFn (const $ return ()) $ train svmAlgorithm svmKernel trainingDataset'
  resultsTraining <- checkProblem svmModel trainingDataset'
  resultsValidation <- checkProblem svmModel testingDataset'
  return (legalWords, resultsTraining, resultsValidation)

main :: IO ()
main = do
  let thresh = 200
  let nsv = 300
  let kern = Linear --RBF 0.01
  let alg = CSvc 1

  dataset <- trainingSamples
  validationDataset <- validationSamples

  putStrLn "Our Method (Term Frequency Negation):"
  printReport =<< runExperiment (Flip, thresh, nsv, kern, alg, dataset, validationDataset)
  putStrLn "Stoplist & Stem Modifiers Without Negation:"
  printReport =<< runExperiment (Drop, thresh, nsv, kern, alg, dataset, validationDataset)
  putStrLn "Modifiers Treated Normally:"
  printReport =<< runExperiment (Pass, thresh, nsv, kern, alg, dataset, validationDataset)

-- Assorted utility functions

crop :: (Num a, Ord a) => a -> a
crop = max (-1) . min 1

clamp :: (Num a, Ord a) => a -> a
clamp x = if x > 0 then 1 else -1

balance :: (Ord a, Fractional a) => a -> a
balance x = if x > 0 then 1 + x else 1 / (1 - x)

takeEvery :: Int -> [a] -> [a]
takeEvery n xs = case drop (n-1) xs of (y:ys) -> y : takeEvery n ys; [] -> []

dropEvery :: Int -> [a] -> [a]
dropEvery n xs = case splitAt (n-1) xs of (x, (y:ys)) -> x ++ dropEvery n ys; (x, []) -> x

r2 :: Floating a => [(a, a)] -> (a, (a, a))
r2 points = let n = fromIntegral (length points)
                ex = sum (map fst points) / n
                ey = sum (map snd points) / n
                sx = sqrt $ sum (map ((^2) . ((-)ex) . fst) points) / n
                sy = sqrt $ sum (map ((^2) . ((-)ey) . snd) points) / n
                cov = sum (map (\(x, y) -> (x - ex) * (y - ey)) points) / n
            in (cov / (sx * sy), (sx / sy, ex * sy / sx - ey))

both :: (a -> b) -> (a, a) -> (b, b)
both f (x, y) = (f x, f y)

swap :: (a, b) -> (b, a)
swap (x, y) = (y, x)


module Svg where

import Data.Maybe
import Control.Applicative
import Data.Complex

data ViewBox = ViewBox {
  viewBoxX :: String,
  viewBoxY :: String,
  viewBoxWidth :: String,
  viewBoxHeight :: String
  } deriving Eq

viewBoxToString :: ViewBox -> String
viewBoxToString vb = "viewBox=\"" ++ viewBoxX vb ++ " " ++ viewBoxY vb ++ " " ++
                     viewBoxWidth vb ++ " " ++ viewBoxHeight vb ++ "\""

data WidthHeight = WidthHeight {
  width :: String,
  height :: String
  } deriving Eq

switchOrientation :: WidthHeight -> WidthHeight
switchOrientation (WidthHeight w h) = WidthHeight h w

widthHeightToString :: WidthHeight -> String
widthHeightToString (WidthHeight w h) = "width=\"" ++ w ++ "\" height=\"" ++ h ++ "\""

a4PortraitWidthHeight :: WidthHeight
a4PortraitWidthHeight = WidthHeight "210mm" "297mm"

data Stroke = Stroke {
  strokeColorMaybe :: Maybe String,
  strokeWidthMaybe :: Maybe Double
  } deriving Eq

strokeToString :: Stroke -> String
strokeToString (Stroke cm wm) = maybeToAttr cm colorToString ++
                                maybeToAttr wm widthToString where
                                  colorToString c = "stroke=\"" ++ c ++ "\""
                                  widthToString w = "stroke-width=\"" ++ show w ++ "\""

strokeWidthMap :: (Double -> Double) -> Stroke -> Stroke
strokeWidthMap f (Stroke cm wm) = Stroke cm (f <$> wm)

maybeToAttr :: Maybe a -> (a -> String) -> String
maybeToAttr maybe toStr =
  fromMaybe "" (((' ' :) . toStr) <$> maybe)

maybePrefix :: Maybe String -> String
maybePrefix m = fromMaybe "" ((++ ":") <$> m)

prologue :: Maybe String -> String -> Maybe WidthHeight -> Maybe ViewBox -> String
prologue svgPrefixMaybe xlinkPrefix widthHeightMaybe viewBoxMaybe =
  "<" ++  maybePrefix svgPrefixMaybe  ++
  "svg xmlns" ++ (fromMaybe "" ((':' :) <$> svgPrefixMaybe)) ++
  "=\"http://www.w3.org/2000/svg\" xmlns:" ++
  xlinkPrefix ++ "=\"http://www.w3.org/1999/xlink\"" ++
  maybeToAttr widthHeightMaybe widthHeightToString ++
  maybeToAttr viewBoxMaybe viewBoxToString ++
  "><" ++ maybePrefix svgPrefixMaybe ++ "g transform=\"scale(1 -1)\">"

epilogue :: Maybe String -> String
epilogue svgPrefixMaybe = let pref = maybePrefix svgPrefixMaybe in
  "</" ++ pref ++ "g></" ++ pref ++ "svg>"

drawGraph :: Maybe String -> Double -> Double -> Stroke -> [(Double, Double)] -> String
drawGraph svgPrefixMaybe xScale yScale stroke pnts =
  let scaledPnts = map (\(x, y) -> (x * xScale, y * yScale)) pnts
      pairToString (x, y) = show x ++ " " ++ show y
  in
   "<" ++ maybePrefix svgPrefixMaybe ++ "path" ++
   strokeToString stroke ++
   " fill=\"none\" d=\"M " ++ pairToString (head scaledPnts) ++
   " L" ++ (foldr (\s1 s2 -> (' ' : s1) ++ s2) "" (map pairToString (tail scaledPnts))) ++
   "\" />"

drawComplex :: Show a => Complex a -> Maybe String -> Stroke -> String
drawComplex (a :+ b) svgPrefixMaybe stroke =
  "<" ++ maybePrefix svgPrefixMaybe ++ "line" ++ strokeToString stroke ++
  " x1=\"0\" y1=\"0\" x2=\"" ++ show a ++ "\" y2=\"" ++ show b ++ "\" />"


-- |
-- Module      : SpriteClip.Render
-- Description :
-- Copyright   : (c) Jonatan H Sundqvist, 2015
-- License     : MIT
-- Maintainer  : Jonatan H Sundqvist
-- Stability   : experimental|stable
-- Portability : POSIX (not sure)
--

-- Created September 7 2015

-- TODO | - UI and graphics settings
--        - Factor out logic (and settings, input mode, etc.) (should not be embedded in rendering functions)

-- SPEC | -
--        -

--------------------------------------------------------------------------------------------------------------------------------------------
-- GHC Pragmas
--------------------------------------------------------------------------------------------------------------------------------------------
{-# LANGUAGE OverloadedStrings #-}

--------------------------------------------------------------------------------------------------------------------------------------------
-- API
--------------------------------------------------------------------------------------------------------------------------------------------
module SpriteClip.Render (module SpriteClip.Render,
                          module Render,
                          module Palette,
                          module Southpaw.Picasso.RenderUtils,
                          module Cairo) where

--------------------------------------------------------------------------------------------------------------------------------------------
-- We'll need these
--------------------------------------------------------------------------------------------------------------------------------------------
import Data.Complex
import Data.Maybe
import qualified Data.Set as S

import Text.Printf
import Control.Monad (forM_, when, liftM)
import Control.Lens

import qualified Graphics.Rendering.Cairo as Cairo
import qualified Graphics.UI.Gtk          as Gtk

import           Southpaw.Picasso.Palette (alphapart, redpart, greenpart, bluepart)
import qualified Southpaw.Picasso.Palette     as Palette
import qualified Southpaw.Picasso.Render      as Render
import           Southpaw.Picasso.RenderUtils hiding (grid, closePath)

import qualified Southpaw.Cartesian.Plane.BoundingBox as BBox
import           Southpaw.Cartesian.Plane.Utilities

import           SpriteClip.Types
import           SpriteClip.Lenses
import qualified SpriteClip.Logic  as Logic

--------------------------------------------------------------------------------------------------------------------------------------------
-- Functions
--------------------------------------------------------------------------------------------------------------------------------------------
-- | Renders a frame
render :: AppState -> Cairo.Render ()
render appstate = do
  --
  canvasSize <- imageSurfaceSize sheet'

  -- Render.imageWithClip (const $ const $ Render.circle (appstate ^. mouse) 120) (50:+50) (appstate ^. sheet)
  Render.image (100:+100) (appstate ^. sheet)

  --
  -- Render.image (canvasSize*0.5) (appstate ^. sheet) -- Render sprite sheet
  -- Cairo.liftIO $ print (length $ take 5 sections)
  forM_ sections' (section mouse') -- Render sprite sections with markers
  alignments sections' mouse'      -- Render alignments
  axes mouse' canvasSize           -- Render cursor axes
  debugHUD appstate                -- Render debug overlay
  where
    sections' = (appstate ^. cutouts) ++ maybe [] (\click -> [BBox.fromCorners click mouse']) (appstate ^. pin)                -- TODO: Refactor
    mouse'    = Logic.applyOptionals [(gridsnap, dotmap $ Logic.nearestMultiple 20.0), (square, smallest)] (appstate ^. mouse) -- TODO: Rename (eg. cursor or pointer)
    sheet'    = appstate ^. sheet
    gridsnap  = S.member "shift_r"    $ appstate ^. keys --
    square    = S.member "control_r"  $ appstate ^. keys --
    smallest (re:+im) = min re im :+ min re im           --

-- |
-- TODO: Options
section ::  Complex Double -> BoundingBox Double -> Cairo.Render ()
section mouse bbox = do
  -- Outline
  Render.rectangle (_centre bbox) (_size bbox)
  choose Palette.magenta
  Cairo.stroke

  -- Markers
  sectionMarkers bbox mouse

-- | Renders the markers for a sheet section (respresented as a 'BoundingBox')
-- TODO: Options
-- TODO: Move out hover logic (?)
-- TODO: Refactor
sectionMarkers :: BoundingBox Double -> Complex Double -> Cairo.Render ()
sectionMarkers bbox mouse = forM_ markerBunches $ \ (radius, colour, offsets) ->
    forM_ offsets $ \pt -> do
      choose colour
      Render.circle (centre pt) (r radius pt)
      Cairo.fill
  where
    scale     = max 0.3 $ min 1.6 (let szx:+szy = _size bbox in min (abs szx) (abs szy)/120.0) -- Scale depending on section size
    focused r = Logic.within (scale*r) mouse . centre                                          -- Does the cursor lie on the marker?
    r r' pt   = scale*r'*(if focused r' pt then 1.4 else 1.0)                                  --
    (rectx:+recty)  = _size bbox
    centre (dx:+dy) = _centre bbox + ((dx*rectx/2):+(dy*recty/2))
    markerBunches   = [(12, Palette.orange & alphapart .~ 0.88, Logic.corners),
                       (10, Palette.peru   & alphapart .~ 0.88, Logic.midline),
                       ( 7, Palette.plum   & alphapart .~ 0.88, Logic.centre)]

-- | Renders the alignments for a sprite section (respresented as a 'BoundingBox')
-- TODO: Don't assume the point represents the mouse
alignments :: [BoundingBox Double] -> Complex Double -> Cairo.Render ()
alignments bboxes mouse@(mx:+my) = do

  -- TODO: Implement actual snapping
  -- TODO: Coalesce duplicate snaps

  forM_ (Logic.uniqueAlignments 0.0 6.0 bboxes mouse) (uncurry referenceline)

  -- TODO: Save and push state to prevent interference, rather than resetting manually
  Cairo.newPath        --
  Cairo.setDash [] 0.0 -- Disable dashes
  where
    crosshairs (fr:+om) = Cairo.setLineWidth 1.0 >> Cairo.setDash [] 0.0 >> choose Palette.black >> Render.crosshairs (fr:+om) (22:+22) >> Cairo.stroke
    cairosetup d        = choose (markTheSpot d) >> Cairo.setLineWidth 1 >> Cairo.setDash [8, 8] 0
    markTheSpot d       = let i = if d /= 0.0 then 0.4 else 0.0 in (i, i, i, 1.0) --

    referenceline (Vertical d)   (fr:+om) = cairosetup d >> Render.line (fr:+om) (fr:+my) >> Cairo.stroke >> crosshairs (fr:+om)
    referenceline (Horizontal d) (fr:+om) = cairosetup d >> Render.line (fr:+om) (mx:+om) >> Cairo.stroke >> crosshairs (fr:+om)

-- | Renders the X and Y axes meeting at the given origin point.
-- TODO: Options
-- TODO: Hide mouse, use axes instead
-- TODO: Use crosshairs with 'hollow middle' at the origin point (?)
axes :: Complex Double -> Complex Double -> Cairo.Render ()
axes (ox:+oy) (width:+height) = do
  Cairo.setLineWidth 1
  Cairo.setDash [5, 5] 14

  choose Palette.green
  Cairo.moveTo     0 oy
  Cairo.lineTo width oy
  Cairo.stroke

  choose Palette.red
  Cairo.moveTo ox      0
  Cairo.lineTo ox height
  Cairo.stroke

-- |
debugHUD :: AppState -> Cairo.Render ()
debugHUD appstate = do
  -- Debug info (HUD)
  Cairo.moveTo 20 20
  choose Palette.darkblue
  Cairo.setFontSize 16
  -- Cairo.fontOptionsSetAntilias Cairo.AntialiasDefault
  Cairo.selectFontFace ("Helvetica" :: String) Cairo.FontSlantNormal Cairo.FontWeightNormal
  Cairo.showText $ (printf "Mouse=(%.02f, %.02f)" mx my :: String)
  maybe pass (\(px:+py) -> Cairo.showText (printf " | Size=(%.02f, %.02f)" (abs $ px-mx) (abs $ py-my) :: String)) (appstate ^. pin)
  where
    (mx:+my) = appstate ^. mouse
    msize    = liftM (dotmap abs . subtract (mx:+my)) (appstate ^. pin)
    pass     = return ()


{-# LANGUAGE OverloadedStrings, ScopedTypeVariables #-}

module Main where

import Control.Monad (forM, forM_)
import Data.List (intercalate)
import Statistics.Sample (mean, stdDev)
import Turtle
import qualified Data.Text as T
import qualified Data.Vector as V

data Platform = PlatformLinux | PlatformMacOS deriving (Eq)

main :: IO ()
main = do
    uname <- strict $ inshell "uname" empty
    platform <-
        if "Linux" `T.isInfixOf` uname then return PlatformLinux
        else if "Darwin" `T.isInfixOf` uname then return PlatformMacOS
        else failWithMsg "Currently only Linux and macOS are supported."

    -- Required for pcregrep on Linux.
    when (platform == PlatformLinux) $ export "LC_ALL" "en_US.UTF-8"

    let logFile :: (IsString a) => a
        logFile = "./tmp/bench.sh.log"

    export "logfile" (fromString logFile)
    testfile logFile >>= flip when (rm logFile)

    procs "echo" ["Compiling C program... "] empty
    shells "mkdir -p ./tmp" empty
    shells "gcc -O2 -Wall bench-lmdb.c -llmdb -o bench-lmdb >> $logfile 2>&1" empty

    procs "echo" ["Compiling Haskell programs... "] empty
    shells "stack build >> $logfile 2>&1" empty

    let c_executable = "./bench-lmdb"
    hs_plain_executable <- T.strip <$> strict (inshell "stack exec -- which bench-lmdb-plain 2> /dev/null" empty)
    hs_streamly_executable <- T.strip <$> strict (inshell "stack exec -- which bench-lmdb-streamly 2> /dev/null" empty)

    echoHrule
    answerIsYes "Perform read benchmarks?" >>= flip when (do
        echo "Creating databases using C program (unless they already exist)..."
        let tmpDir :: (IsString a) => a
            tmpDir = "./tmp/read"
        mktree tmpDir

        let mb = 1000000.0 :: Double
        let pairCounts = [31250000 :: Int, 62500000, 125000000]
        let kvFactors = [1 :: Int, 2, 4]
        forM_ pairCounts $ \pairCount -> forM_ kvFactors $ \kvFactor -> do
            let dbPath = format (""%fp%"/test_"%d%"_"%d%"") tmpDir pairCount kvFactor
            (_, existingDbPairCount) <-
                shellStrict (format ("mdb_stat -e "%s%" 2> /dev/null | pcregrep -o1 'Entries: (.+)$'") dbPath) empty
            if T.strip existingDbPairCount == format (""%d%"") pairCount then
                procs "echo" [format ("    Database "%s%" already exists.") dbPath] empty
            else do
                procs "echo" ["-n", format ("    Creating database "%s%"... ") dbPath] empty
                testdir (fromText dbPath) >>= flip when (shells (format ("rm -r "%s%" 2> /dev/null") dbPath) empty)
                let chunkSize :: Int = round $ mb / (8 * fromIntegral kvFactor + 8 * fromIntegral kvFactor)
                shells (format (""%s%" write "%s%" "%d%" "%d%" "%d%" "%d%"")
                    c_executable dbPath kvFactor kvFactor pairCount chunkSize) empty
                echo "success."

        echo "Measuring read-cursor..."

        let csvFile :: (IsString a) => a
            csvFile = "./tmp/read/read-cursor.csv"
        shells (format ("rm -f "%fp%"") csvFile) empty

        output csvFile . return . unsafeTextToLine . T.pack $
            "pairCount,kvFactor,"
            ++ "c_mean,c_std,"
            ++ "hsPlain_mean,hsPlain_std,"
            ++ "hsStreamly_mean,hsStreamly_std,"
            ++ "hsStreamlySafe_mean,hsStreamlySafe_std,"
            ++ "hsPlain_diff_c_mean,hsPlain_diff_c_std,"
            ++ "hsStreamly_diff_hsPlain_mean,hsStreamly_diff_hsPlain_std,"
            ++ "hsStreamlySafe_diff_hsStreamly_mean,hsStreamlySafe_diff_hsStreamly_std"

        forM_ pairCounts $ \pairCount -> forM_ kvFactors $ \kvFactor -> do
            let dbPath = format (""%fp%"/test_"%d%"_"%d%"") tmpDir pairCount kvFactor
            let expectedInOutput = format ("Pair count:       "%d%"") pairCount
            let args = format ("read-cursor "%s%"") dbPath
            let warmCount = 2
            let timeCount = 5

            -- Mean and standard deviation of nanoseconds per pair.
            let toStats :: [NominalDiffTime] -> (Double, Double)
                toStats secs =
                    let  nsPerPair = V.fromList $ map (realToFrac . (/fromIntegral pairCount) . (*1e9)) secs
                    in (mean nsPerPair, stdDev nsPerPair)

            procs "echo" [format ("    Timing database "%s%" with C...") dbPath] empty
            (cMean, cStd) <-
                toStats <$>
                timeCommand Nothing (format (""%s%" "%s%"") c_executable args)
                    [expectedInOutput] warmCount timeCount

            procs "echo" [format ("    Timing database "%s%" with Haskell (plain)...") dbPath] empty
            (hsPlainMean, hsPlainStd) <-
                toStats <$>
                timeCommand Nothing (format (""%s%" "%s%"") hs_plain_executable args)
                    [expectedInOutput] warmCount timeCount

            procs "echo" [format ("    Timing database "%s%" with Haskell (streamly)...") dbPath] empty
            (hsStreamlyMean, hsStreamlyStd) <-
                toStats <$>
                timeCommand Nothing (format (""%s%" "%s%"") hs_streamly_executable args)
                    [expectedInOutput] warmCount timeCount

            procs "echo" [format ("    Timing database "%s%" with Haskell (streamly, safe)...") dbPath] empty
            (hsStreamlySafeMean, hsStreamlySafeStd) <-
                toStats <$>
                timeCommand Nothing (format (""%s%" read-cursor-safe "%s%"") hs_streamly_executable dbPath)
                    [expectedInOutput] warmCount timeCount

            let diffStd :: Double -> Double -> Double
                diffStd std1 std2 =
                    let two = 2 :: Int
                    in sqrt $ std1^two + std2^two

            let line = intercalate "," $ [show pairCount, show kvFactor] ++ map show [cMean, cStd,
                    hsPlainMean, hsPlainStd, hsStreamlyMean, hsStreamlyStd, hsStreamlySafeMean, hsStreamlySafeStd,
                    hsPlainMean - cMean, diffStd hsPlainStd cStd,
                    hsStreamlyMean - hsPlainMean, diffStd hsStreamlyStd hsPlainStd,
                    hsStreamlySafeMean - hsStreamlyMean, diffStd hsStreamlySafeStd hsStreamlyStd]
            append csvFile . return . unsafeTextToLine . T.pack $ line
            procs "echo" [format ("    Appended results to "%s%"") csvFile] empty
        )

    echoHrule
    answerIsYes "Perform write benchmarks?" >>= flip when (do
        let tmpDir :: (IsString a) => a
            tmpDir = "./tmp/write"
        shells (format ("rm -rf "%fp%"") tmpDir) empty
        mktree tmpDir
        let csvFile :: (IsString a) => a
            csvFile = "./tmp/write/write.csv"
        output csvFile . return . unsafeTextToLine . T.pack $
            "pairCount,kvFactor,"
            ++ "c_mean,c_std,"
            ++ "hsPlain_mean,hsPlain_std,"
            ++ "hsStreamly_mean,hsStreamly_std,"
            ++ "hsPlain_div_c_mean,hsPlain_div_c_std,"
            ++ "hsStreamly_div_c_mean,hsStreamly_div_c_std,"
            ++ "hsStreamly_div_hsPlain_mean,hsStreamly_div_hsPlain_std"

        let mb = 1000000.0 :: Double
        let pairCounts = [1000000 :: Int, 10000000]
        let kvFactors = [1 :: Int, 2]
        forM_ pairCounts $ \pairCount -> forM_ kvFactors $ \kvFactor -> do
            let dbPath = format (""%fp%"/test_"%d%"_"%d%"") tmpDir pairCount kvFactor
            let chunkSize :: Int = round $ mb / (8 * fromIntegral kvFactor + 8 * fromIntegral kvFactor)
            let args  = format ("write "%s%" "%d%" "%d%" "%d%" "%d%"") dbPath kvFactor kvFactor pairCount chunkSize
            let warmCount = 1
            let timeCount = 5
            procs "echo" [format ("    Writing database "%s%" with C...") dbPath] empty
            cTimes <- timeCommand (Just dbPath) (format (""%s%" "%s%"") c_executable args)
                        [] warmCount timeCount
            procs "echo" [format ("    Writing database "%s%" with Haskell (plain)...") dbPath] empty
            hsPlainTimes <- timeCommand (Just dbPath) (format (""%s%" "%s%"") hs_plain_executable args)
                        [] warmCount timeCount
            procs "echo" [format ("    Writing database "%s%" with Haskell (streamly)...") dbPath] empty
            hsStreamlyTimes <- timeCommand (Just dbPath) (format (""%s%" "%s%"") hs_streamly_executable args)
                        [] warmCount timeCount

            let hsPlainDivCTimes = [a / b | a <- hsPlainTimes, b <- cTimes]
            let hsStreamlyDivCTimes = [a / b | a <- hsStreamlyTimes, b <- cTimes]
            let hsStreamlyDivHsPlainTimes = [a / b | a <- hsStreamlyTimes, b <- hsPlainTimes]

            let mean' :: [NominalDiffTime] -> Double
                mean' xs = mean . V.fromList $ map realToFrac xs

            let std' :: [NominalDiffTime] -> Double
                std' xs = stdDev . V.fromList $ map realToFrac xs

            let line = intercalate "," $ [show pairCount, show kvFactor] ++ map show [
                    mean' cTimes, std' cTimes,
                    mean' hsPlainTimes, std' hsPlainTimes,
                    mean' hsStreamlyTimes, std' hsStreamlyTimes,
                    mean' hsPlainDivCTimes, std' hsPlainDivCTimes,
                    mean' hsStreamlyDivCTimes, std' hsStreamlyDivCTimes,
                    mean' hsStreamlyDivHsPlainTimes, std' hsStreamlyDivHsPlainTimes ]
            append csvFile . return . unsafeTextToLine . T.pack $ line
            procs "echo" [format ("    Appended results to "%s%"") csvFile] empty
        )

failWithMsg :: Line -> IO a
failWithMsg msg = echo msg >> exit (ExitFailure 1)

echoHrule :: IO ()
echoHrule = echo "----------------------------------------"

answerIsYes :: String -> IO Bool
answerIsYes question = do
    procs "echo" ["-n", T.pack $ question ++ " (y/N) "] empty
    line <- readline
    case line of
        Just answer | answer == "y" -> return True
        _ -> return False

timeCommand' :: Text -> [Text] -> IO NominalDiffTime
timeCommand' cmdAndArgs expectedInOut = do
    (out, t) <- time . strict $ inshell cmdAndArgs empty
    if all (`T.isInfixOf` out) expectedInOut then
        return t
    else
        error $ "all of " ++ show expectedInOut ++ " were not found in output "
            ++ show out ++ " of command " ++ show cmdAndArgs

timeCommand :: Maybe T.Text -> Text -> [Text] -> Int -> Int -> IO [NominalDiffTime]
timeCommand forceRemovalPath cmdAndArgs expectedInOut warmCount timeCount = do
    let removePath = maybe (return ()) (\fp' -> shells (format ("rm -rf "%s%"") fp') empty) forceRemovalPath
    let echoCount c = procs "echo" ["-n", format (""%d%" ") c] empty
    when (warmCount > 0) $ do
        procs "echo" ["-n", "        Warming up... "] empty
        forM_ [1..warmCount] $ \c ->
            echoCount c >> removePath >> timeCommand' cmdAndArgs expectedInOut
        echo "done."
    procs "echo" ["-n", "        Timing... "] empty
    times <- forM [1..timeCount] $ \c ->
        echoCount c >> removePath >> timeCommand' cmdAndArgs expectedInOut
    removePath
    echo "done."
    return times


{-# LANGUAGE DataKinds #-}
{-# LANGUAGE IncoherentInstances #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE TypeSynonymInstances #-}
{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# OPTIONS_GHC -fno-warn-orphans #-}
module Data.C
  ( CVal (..)
  , CTy(..)
  , IsFun
  , eraseTy
  , getCTyR
  , getTy
  , eitherTy
  , pairTy
  , cTySpec
  , (<::)
  , tagFld
  , valFld
  , inlFld
  , inrFld
  , fstFld
  , sndFld
  , cCase
  , cTagl
  , cTagr
  , declVar
  , Int
  , Bool
  , Float
  , Double
  , String
  , Either
  ) where

import Control.Monad.CGen
import Type.Reflection ( Typeable, TypeRep, typeRep )
import Data.Complex

data CTy a where
  CUnit   :: CTy ()
  CInt    :: CTy Int
  CBool   :: CTy Bool
  CFlt    :: CTy Float
  CDbl    :: CTy Double
  CCplx   :: CTy (Complex Double)
  CStr    :: CTy String
  CPair   :: (CVal a, CVal b) => CTy a -> CTy b -> CTy (a,b)
  CEither :: (CVal a, CVal b) => CTy a -> CTy b -> CTy (Either a b)
  CVec    :: CVal a => CTy a -> CTy [a]

eraseTy :: CTy a -> ECTy
eraseTy CUnit         = ECUnit
eraseTy CInt          = ECInt
eraseTy CBool         = ECBool
eraseTy CFlt          = ECFlt
eraseTy CDbl          = ECDbl
eraseTy CCplx         = ECCplx
eraseTy CStr          = ECStr
eraseTy (CPair   l r) = ECPair (eraseTy l) (eraseTy r)
eraseTy (CEither l r) = ECEither (eraseTy l) (eraseTy r)
eraseTy (CVec    l  ) = ECVec (eraseTy l)

getTy :: forall a t. CVal a => t a -> ECTy
getTy _ = eraseTy (getCTy :: CTy a)

getCTyR :: CTy a -> TypeRep a
getCTyR (CPair _ _) = typeRep
getCTyR (CEither _ _) = typeRep
getCTyR (CVec _) = typeRep
getCTyR CUnit = typeRep
getCTyR CInt = typeRep
getCTyR CBool = typeRep
getCTyR CFlt = typeRep
getCTyR CDbl = typeRep
getCTyR CCplx = typeRep
getCTyR CStr = typeRep

eitherTy :: (CVal a, CVal b) => t a -> t b -> TypeRep (Either a b)
eitherTy _ _ = typeRep

pairTy :: (CVal a, CVal b) => t a -> t b -> TypeRep (a, b)
pairTy _ _ = typeRep

cTySpec :: forall a t st. CVal a => t a -> CGen st ([CDeclSpec], [CDerivedDeclr])
cTySpec _ = typeSpec $ eraseTy $ (getCTy :: CTy a)

type family IsFun f where
  IsFun (a -> b) = 'True
  IsFun a = 'False

class (Eq a, Ord a, Show a, Typeable a, IsFun a ~ 'False) => CVal a where
  getCTy  :: CTy a
  cVal    :: a -> CGen st CExpr

instance CVal () where
  getCTy = CUnit
  cVal () = cTySpec CUnit *> pure (cVar cUnit)

instance CVal Int where
  getCTy = CInt
  cVal i = pure $ cInt i

instance CVal Bool where
  getCTy = CBool
  cVal True = pure $ cInt 1
  cVal False = pure $ cInt 0

instance CVal Float where
  getCTy = CFlt
  cVal i = pure $ cFlt i

instance CVal Double where
  getCTy = CDbl
  cVal i = pure $ cDbl i

instance Ord (Complex Double) where
  compare i j =
    case compare (realPart i) (realPart j) of
      EQ -> compare (imagPart i) (imagPart j)
      o -> o

instance CVal (Complex Double) where
  getCTy = CCplx
  cVal i = pure $
    CBinary CAddOp (cDbl (realPart i))
    (CBinary CMulOp (cDbl (imagPart i)) (cVar $ internalIdent "I") undefNode)
    undefNode

instance (CVal a, CVal b) => CVal (a, b) where
  getCTy = CPair getCTy getCTy
  cVal (x, y) = do
    vx <- cVal x
    vy <- cVal y
    (t, q) <- cTySpec $ getCTy @(a,b)
    cVar <$> newVar t q (Just $ ini [mkInit vx, mkInit vy])
      where
        mkInit e = ([], CInitExpr e undefNode)

instance (CVal a, CVal b) => CVal (Either a b) where
  getCTy = CEither getCTy getCTy
  cVal (Left x) = do
    vx <- cVal x
    (t, q) <- cTySpec $ getCTy @(Either a b)
    cVar <$> newVar t q (Just $ ini [tagL, ([], ini [([mkInL], iniE vx)])])
      where
        tagL = ([], iniE $ cVar cTagl)
        mkInL = CMemberDesig inlFld undefNode
  cVal (Right x) = do
    vx <- cVal x
    (t, q) <- cTySpec $ getCTy @(Either a b)
    cVar <$> newVar t q (Just $ ini [tagR, ([], ini [([mkInR], iniE vx)])])
      where
        tagR = ([], iniE $ cVar cTagr)
        mkInR = CMemberDesig inrFld undefNode

ini :: CInitializerList NodeInfo -> CInitializer NodeInfo
ini = (`CInitList` undefNode)
iniE :: CExpression NodeInfo -> CInitializer NodeInfo
iniE = (`CInitExpr` undefNode)

instance (CVal a) => CVal [a] where
  getCTy = CVec getCTy
  cVal _xs = error "FIXME: unimplemented"

instance {-# OVERLAPPING #-} CVal String where
  getCTy = CStr
  cVal s = pure $ cStr s

(<::) :: forall a t st. CVal a => Ident -> t a -> CGen st [CBlockItem]
v <:: _ = do
  (t, q) <- cTySpec (getCTy :: CTy a)
  pure $ [CBlockDecl $ varDecl v t q Nothing]

cCase :: CExpr -> [CBlockItem] -> [CBlockItem] -> CStat
cCase e l r = CSwitch e cases undefNode
  where
    cases = cComp $ zipWith3 mkCase [True, False] [l, r] [cTagl, cTagr]
    mkCase bb ss tag = CBlockStmt $ CCase (cVar tag) (caseCode bb ss) undefNode
    caseCode b ss = cComp $ ss ++ mbreak
      where
        mbreak = if b then [CBlockStmt $ CBreak undefNode] else []

declVar :: forall a t st. CVal a => t a -> CGen st (CExpr, [CBlockItem])
declVar e
  | getTy e == ECUnit = pure (cVar cUnit, [])
  | otherwise = do
      v <- freshVar
      dv <- v <:: e
      pure (cVar v, dv)


{-# LANGUAGE Trustworthy #-}
{-# LANGUAGE DeriveDataTypeable, DeriveGeneric, OverloadedStrings,
    RecordWildCards, ScopedTypeVariables #-}

-- |
-- Module      : Criterion.Report
-- Copyright   : (c) 2009-2014 Bryan O'Sullivan
--
-- License     : BSD-style
-- Maintainer  : bos@serpentine.com
-- Stability   : experimental
-- Portability : GHC
--
-- Reporting functions.

module Criterion.Report
    (
      formatReport
    , report
    , tidyTails
    -- * Rendering helper functions
    , TemplateException(..)
    , loadTemplate
    , includeFile
    , getTemplateDir
    , vector
    , vector2
    ) where

import Control.Exception (Exception, IOException, throwIO)
import Control.Monad (mplus)
import Control.Monad.IO.Class (MonadIO(liftIO))
import Control.Monad.Reader (ask)
import Criterion.Monad (Criterion)
import Criterion.Types
import Data.Aeson.Encode (encodeToTextBuilder)
import Data.Aeson.Types (toJSON)
import Data.Data (Data, Typeable)
import Data.Foldable (forM_)
import GHC.Generics (Generic)
import Paths_criterion (getDataFileName)
import Statistics.Function (minMax)
import System.Directory (doesFileExist)
import System.FilePath ((</>), (<.>), isPathSeparator)
import Text.Hastache (MuType(..))
import Text.Hastache.Context (mkGenericContext, mkStrContext, mkStrContextM)
import qualified Control.Exception as E
import qualified Data.Text as T
import qualified Data.Text.IO as T
import qualified Data.Text.Lazy as TL
import qualified Data.Text.Lazy.Builder as TL
import qualified Data.Text.Lazy.IO as TL
import qualified Data.Vector.Generic as G
import qualified Data.Vector.Unboxed as U
import qualified Language.Javascript.Flot as Flot
import qualified Language.Javascript.JQuery as JQuery
import qualified Text.Hastache as H

-- | Trim long flat tails from a KDE plot.
tidyTails :: KDE -> KDE
tidyTails KDE{..} = KDE { kdeType   = kdeType
                        , kdeValues = G.slice front winSize kdeValues
                        , kdePDF    = G.slice front winSize kdePDF
                        }
  where tiny     = uncurry subtract (minMax kdePDF) * 0.005
        omitTiny = G.length . G.takeWhile ((<= tiny) . abs)
        front    = omitTiny kdePDF
        back     = omitTiny . G.reverse $ kdePDF
        winSize  = G.length kdePDF - front - back

-- | Return the path to the template and other files used for
-- generating reports.
getTemplateDir :: IO FilePath
getTemplateDir = getDataFileName "templates"

-- | Write out a series of 'Report' values to a single file, if
-- configured to do so.
report :: [Report] -> Criterion ()
report reports = do
  Config{..} <- ask
  forM_ reportFile $ \name -> liftIO $ do
    td <- getTemplateDir
    tpl <- loadTemplate [td,"."] template
    TL.writeFile name =<< formatReport reports tpl

-- | Format a series of 'Report' values using the given Hastache
-- template.
formatReport :: [Report]
             -> T.Text    -- ^ Hastache template.
             -> IO TL.Text
formatReport reports template = do
  templates <- getTemplateDir
  let context "report"    = return $ MuList $ map inner reports
      context "json"      = return $ MuVariable (encode reports)
      context "include"   = return $ MuLambdaM $ includeFile [templates]
      context "js-jquery" = fmap MuVariable $ TL.readFile =<< JQuery.file
      context "js-flot"   = fmap MuVariable $ TL.readFile =<< Flot.file Flot.Flot
      context _           = return $ MuNothing
      encode v = TL.toLazyText . encodeToTextBuilder . toJSON $ v
      inner r@Report{..} = mkStrContextM $ \nym ->
                         case nym of
                           "name"     -> return . MuVariable . H.htmlEscape .
                                         TL.pack $ reportName
                           "json"     -> return $ MuVariable (encode r)
                           "number"   -> return $ MuVariable reportNumber
                           "iters"    -> return $ vector "x" iters
                           "times"    -> return $ vector "x" times
                           "cycles"   -> return $ vector "x" cycles
                           "kdetimes" -> return $ vector "x" kdeValues
                           "kdepdf"   -> return $ vector "x" kdePDF
                           "kde"      -> return $ vector2 "time" "pdf" kdeValues kdePDF
                           ('a':'n':_)-> mkGenericContext reportAnalysis $
                                         H.encodeStr nym
                           _          -> mkGenericContext reportOutliers $
                                         H.encodeStr nym
          where [KDE{..}]   = reportKDEs
                iters       = measure measIters reportMeasured
                times       = measure measTime reportMeasured
                cycles      = measure measCycles reportMeasured
      config = H.defaultConfig {
                 H.muEscapeFunc = H.emptyEscape
               , H.muTemplateFileDir = Just templates
               , H.muTemplateFileExt = Just ".tpl"
               }
  H.hastacheStr config template context

-- | Render the elements of a vector.
--
-- For example, given this piece of Haskell:
--
-- @'mkStrContext' $ \\name ->
-- case name of
--   \"foo\" -> 'vector' \"x\" foo@
--
-- It will substitute each value in the vector for @x@ in the
-- following Hastache template:
--
-- > {{#foo}}
-- >  {{x}}
-- > {{/foo}}
vector :: (Monad m, G.Vector v a, H.MuVar a) =>
          String                -- ^ Name to use when substituting.
       -> v a
       -> MuType m
{-# SPECIALIZE vector :: String -> U.Vector Double -> MuType IO #-}
vector name v = MuList . map val . G.toList $ v
    where val i = mkStrContext $ \nym ->
                  if nym == name
                  then MuVariable i
                  else MuNothing

-- | Render the elements of two vectors.
vector2 :: (Monad m, G.Vector v a, G.Vector v b, H.MuVar a, H.MuVar b) =>
           String               -- ^ Name for elements from the first vector.
        -> String               -- ^ Name for elements from the second vector.
        -> v a                  -- ^ First vector.
        -> v b                  -- ^ Second vector.
        -> MuType m
{-# SPECIALIZE vector2 :: String -> String -> U.Vector Double -> U.Vector Double
                       -> MuType IO #-}
vector2 name1 name2 v1 v2 = MuList $ zipWith val (G.toList v1) (G.toList v2)
    where val i j = mkStrContext $ \nym ->
                    case undefined of
                      _| nym == name1 -> MuVariable i
                       | nym == name2 -> MuVariable j
                       | otherwise    -> MuNothing

-- | Attempt to include the contents of a file based on a search path.
-- Returns 'B.empty' if the search fails or the file could not be read.
--
-- Intended for use with Hastache's 'MuLambdaM', for example:
--
-- @context \"include\" = 'MuLambdaM' $ 'includeFile' ['templateDir']@
--
-- Hastache template expansion is /not/ performed within the included
-- file.  No attempt is made to ensure that the included file path is
-- safe, i.e. that it does not refer to an unexpected file such as
-- \"@/etc/passwd@\".
includeFile :: (MonadIO m) =>
               [FilePath]       -- ^ Directories to search.
            -> T.Text          -- ^ Name of the file to search for.
            -> m T.Text
{-# SPECIALIZE includeFile :: [FilePath] -> T.Text -> IO T.Text #-}
includeFile searchPath name = liftIO $ foldr go (return T.empty) searchPath
    where go dir next = do
            let path = dir </> H.decodeStr name
            T.readFile path `E.catch` \(_::IOException) -> next

-- | A problem arose with a template.
data TemplateException =
    TemplateNotFound FilePath   -- ^ The template could not be found.
    deriving (Eq, Read, Show, Typeable, Data, Generic)

instance Exception TemplateException

-- | Load a Hastache template file.
--
-- If the name is an absolute or relative path, the search path is
-- /not/ used, and the name is treated as a literal path.
--
-- This function throws a 'TemplateException' if the template could
-- not be found, or an 'IOException' if no template could be loaded.
loadTemplate :: [FilePath]      -- ^ Search path.
             -> FilePath        -- ^ Name of template file.
             -> IO T.Text
loadTemplate paths name
    | any isPathSeparator name = T.readFile name
    | otherwise                = go Nothing paths
  where go me (p:ps) = do
          let cur = p </> name <.> "tpl"
          x <- doesFileExist cur
          if x
            then T.readFile cur `E.catch` \e -> go (me `mplus` Just e) ps
            else go me ps
        go (Just e) _ = throwIO (e::IOException)
        go _        _ = throwIO . TemplateNotFound $ name


-- |
-- Module      : Southpaw.Interactive.Types
-- Description :
-- Copyright   : (c) Jonatan H Sundqvist, 2015
-- License     : MIT
-- Maintainer  : Jonatan H Sundqvist
-- Stability   : experimental|stable
-- Portability : POSIX (not sure)
--

-- Created October 15 2015

-- TODO | -
--        -

-- SPEC | -
--        -

--------------------------------------------------------------------------------------------------------------------------------------------
-- GHC Pragmas
--------------------------------------------------------------------------------------------------------------------------------------------

--------------------------------------------------------------------------------------------------------------------------------------------
-- API
--------------------------------------------------------------------------------------------------------------------------------------------
module Southpaw.Interactive.Types where

--------------------------------------------------------------------------------------------------------------------------------------------
-- We'll need these
--------------------------------------------------------------------------------------------------------------------------------------------
import Data.Complex
import Data.IORef
import Graphics.UI.Gtk
import qualified Graphics.Rendering.Cairo as Cairo

--------------------------------------------------------------------------------------------------------------------------------------------
-- Types
--------------------------------------------------------------------------------------------------------------------------------------------

--------------------------------------------------------------------------------------------------------------------------------------------

-- |
-- TODO: Input state (?)
-- TODO: Layout type (?)
data App layout state = App { _window :: Window, _layout :: layout, _appsize :: Complex Int, _state :: IORef state } --
-- data AppState = AppState { _game :: Game, _selected :: Maybe Int, _path :: [Complex Double] }                     --

-- |
-- TODO: Fill out
data WidgetSettings self state = WidgetSettings { _settingsize :: Complex Int, _settingeventmasks :: [EventMask], _settinglisteners :: EventMap self state }

-- |
data FullCanvasLayout = FullCanvasLayout { _canvas :: DrawingArea }

-- data WindowState = Layout { _window :: Window, _canvas :: DrawingArea, _size :: Complex Int }

-- | High-level app setup (flags, config, cleanup, events, etc.)
-- TODO: Use more than one type variable (?)
-- TODO: Lenses
-- TODO: Key maps, event maps
-- TODO: Rename type (?)
-- TODO: Assume stateref (?)
-- TODO: Wrapper for all key press events (?)
data EventMap self s = EventMap { onmousedown :: Maybe (IORef s -> EventM EButton Bool),
	                                onmouseup   :: Maybe (IORef s -> EventM EButton Bool),
														      -- onresize,
														      -- ondrawIO
														      -- onquit :: IO s
														      onmousemotion :: Maybe (IORef s -> EventM EMotion Bool),
														      onkeydown     :: Maybe (IORef s -> EventM EKey Bool),
														      onkeyup       :: Maybe (IORef s -> EventM EKey Bool),
														      onanimate     :: Maybe (Int,  self -> IORef s -> IO Bool), -- TODO: Figure out how to deal with FPS
														      ondraw        :: Maybe (s -> Cairo.Render ()),
														      ondelete      :: Maybe (IORef s -> EventM EAny Bool) }

-- Classes ---------------------------------------------------------------------------------------------------------------------------------


module World
where

import Galaxy
import ZipperGalaxy
import ZipperGalaxyUtils
import Civilization
import Statistics
import Utils
import System.Random
import Control.Monad.State
import Data.Maybe (mapMaybe)
import Data.List (foldl')
import qualified Data.Edison.Assoc.StandardMap as E
import Terrain
import Good

data World = World { galaxy  :: Galaxy Terrain
                   , time    :: GalaxyTime
                   , empires :: E.FM CivKey Empire
                   , player  :: Player
                   }

data GalaxyTime = GalaxyTime { year   :: Int
                             , day    :: Int
                             , hour   :: Int
                             , minute :: Int
                             }

data Relation a = Relation { relationship :: Int
                           , contracts :: [Contract a] }

data Contract a = Embargo a

data Player = Player {playername :: String }

testRandomWorld :: Galaxy Terrain -> Int -> Int -> Maybe World
testRandomWorld g v numciv =
  let r = mkStdGen v
  in evalState (createWorld g numciv) r

nullGalaxyTime = GalaxyTime 0 0 0 0

nullPlayer = Player ""

habitablePlanetsZ :: Galaxy Terrain -> [GalaxyZipper Terrain]
habitablePlanetsZ g = filter sustainsLifeZ (ZipperGalaxyUtils.allBodies g)

civnames :: [String]
civnames = ["humans",
            "aliens",
            "ant insect animals",
            "slimy aliens",
            "cosmic hive snakes"
           ]

colonyStartPopulation :: ResourceUnit
colonyStartPopulation = 100

createWorld :: Galaxy Terrain -> Int -> Rnd (Maybe World)
createWorld g numciv = do
  let ps = habitablePlanetsZ g
  if null ps
    then return Nothing
    else do
      ps' <- shuffle ps
      let cs = take numciv ps'
      if length cs < numciv
        then return Nothing
        else do
          let cols = map ((:[]) . stdColony "colony" colonyStartPopulation) (map zipperToNames cs)
          let empires = zipWith Empire civnames (map namedsToMap cols)
          let collocs = pairToLists civnames (repeat "colony")
          return $ Just $ World
               (foldl' (\g' (cl, pn) -> setColonyOnPlanet cl pn g')
                       g
                       (zip collocs (map planetname (mapMaybe satelliteInZipper cs))))
               nullGalaxyTime
               (namedsToMap empires)
               nullPlayer