text:
{-# 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

max_stars_repo_name: aezarebski/epi-sim

max_issues_repo_name: aezarebski/epi-sim

max_forks_repo_name: aezarebski/epi-sim


text:
-- |

module Math.Algebra.Roots (zeros) where

import Data.Complex
import Polynomial.Roots

-- roots of a polynomial --------------------------------------------------------
-- no dealing with the complex problem of calculating roots, we will use
-- the Polynomial.Roots module of the `dsp` package, complex roots will look a
-- bit 'odd' to many as the output format is x :+ y for x + iy
zeros :: RealFloat a => [Complex a]
zeros = roots 1e-16 1000 [1,3,-10]
-- note that this uses Laguerre's method
-- >>> zeros
-- [0.5 :+ 0.0,(-0.2) :+ 0.0]

ext: hs

max_stars_repo_name: karetsu/swh

max_issues_repo_name: karetsu/swh

max_forks_repo_name: karetsu/swh


text:
module Main where

import qualified Numeric.LinearAlgebra as LA
import qualified MachineLearning as ML
import qualified MachineLearning.Optimization as Opt
import qualified MachineLearning.NeuralNetwork as NN
import qualified MachineLearning.NeuralNetwork.TopologyMaker as TM
import qualified MachineLearning.TerminalProgress as TP

main = do
  putStrLn "\n== Neural Networks (Digits Recognition) ==\n"

  -- Step 1. Data loading.
  -- Step 1.1 Training Data loading.
  (x, y) <- pure ML.splitToXY <*> LA.loadMatrix "digits_classification/optdigits.tra"
  -- Step 1.1 Testing Data loading.
  (xTest, yTest) <- pure ML.splitToXY <*> LA.loadMatrix "digits_classification/optdigits.tes"

  -- Step 2. Initialize Neural Network.
  let nnt = TM.makeTopology TM.ARelu TM.LSoftmax (LA.cols x) 10 [100, 100]
      model = NN.NeuralNetwork nnt
  -- Step 3. Initialize theta with randon values.
      initTheta = NN.initializeTheta 5191711 nnt

      lambda = NN.L2 $ 5 / (fromIntegral $ LA.rows x)

  -- Step 4. Train the Neural Network.
  (thetaNN, optPath) <- TP.learnWithProgressBar (Opt.minimize (Opt.BFGS2 0.03 0.7) model 1e-7 5 lambda x y) initTheta 20

  -- Step 5. Making predictions and checking accuracy on training and test sets.
  let accuracyTrain = NN.calcAccuracy y (NN.hypothesis model x thetaNN)
      accuracyTest = NN.calcAccuracy yTest (NN.hypothesis model xTest thetaNN)

  -- Step 6. Printing results.

  putStrLn $ "\nNumber of iterations to learn the Neural Network: " ++ show (LA.rows optPath)

  putStrLn $ "\nAccuracy on train set (%): " ++ show (accuracyTrain*100)
  putStrLn $ "Accuracy on test set (%): " ++ show (accuracyTest*100)

max_stars_repo_name: aligusnet/mltool

max_issues_repo_name: aligusnet/mltool

max_forks_repo_name: aligusnet/mltool


text:
module Polestar.Core.PrettyPrint
  (prettyPrintTypeP
  ,prettyPrintType
  ,prettyPrintTermP
  ,prettyPrintTerm
  ) where
import Polestar.Core.Type
import Polestar.Core.Parse (NameBinding(..),unaryFns,binaryFns)
import Data.List (find,intersperse)
import Data.Complex

varNames :: [NameBinding] -> [String]
varNames = foldr (\x ys -> case x of
                             NVarBind n -> n:ys
                             _ -> ys) []

tyVarNames :: [NameBinding] -> [String]
tyVarNames = foldr (\x ys -> case x of
                               NTyVarBind n -> n:ys
                               _ -> ys) []

rename :: [String] -> String -> String
rename ctx name | name `notElem` ctx = name
                | otherwise = loop 1
  where loop :: Int -> String
        loop i = case name ++ show i of
          m | m `notElem` ctx -> m
          _ -> loop (i + 1)

showPrimType :: PrimType -> ShowS
showPrimType p = case p of
  PTyNat -> showString "Nat"
  PTyInt -> showString "Int"
  PTyReal -> showString "Real"
  PTyBool -> showString "Bool"
  PTyUnit -> showString "Unit"

-- <3>: <primitive type> | <type variable> | <0>
-- <2>: <3> -> <2>
-- <1>: <2> & .. & <2>
-- <0>: forall a <: <1>. <0> | <1>
prettyPrintTypeP :: Int -> [NameBinding] -> Type -> ShowS
prettyPrintTypeP p ctx t = case t of
  TyPrim p -> showPrimType p
  TyArr s t -> showParen (p > 2) $ prettyPrintTypeP 3 ctx s . showString " -> " . prettyPrintTypeP 2 (NAnonymousBind : ctx) t
  TyTuple components -> showParen True $ foldr (.) id $ intersperse (showString ", ") $ map (prettyPrintTypeP 0 ctx) components

prettyPrintType :: Type -> String
prettyPrintType t = prettyPrintTypeP 0 [] t ""

unaryFnName :: UnaryFn -> String
unaryFnName f | Just (name,_) <- find ((== f) . snd) unaryFns = name
              | otherwise = "<name not found: " ++ show f ++ ">"

binaryFnName :: BinaryFn -> String
binaryFnName f | Just (name,_) <- find ((== f) . snd) binaryFns = name
               | otherwise = "<name not found: " ++ show f ++ ">"

-- <2>: <primitive> | <variable> | (<0>)
-- <1>: <1> <2> | <1> [ty] | <1> as <type>
-- <0>: \x:t. <0> | ?a. <0> | if <0> then <0> else <0>
prettyPrintTermP :: Int -> [NameBinding] -> Term -> ShowS
prettyPrintTermP p ctx t = case t of
  TmPrim (PVNat x) -> shows x
  TmPrim (PVInt x) -> shows x
  TmPrim (PVReal x) -> shows x
  TmPrim (PVBool x) -> if x then showString "true" else showString "false"
  TmPrim PVUnit -> showString "unit"
  TmPrim (PVUnary f) -> showString (unaryFnName f)
  TmPrim (PVBinary f) -> showString (binaryFnName f)
  TmAbs (Id name) ty body -> showParen (p > 0) $ showChar '\\' . showString name' . showChar ':' . prettyPrintTypeP 1 ctx ty . showString ". " . prettyPrintTermP 0 (NVarBind name' : ctx) body
    where name' = rename (varNames ctx) name
  TmRef i | i < length ctx -> case ctx !! i of
              NVarBind n -> showString n
              n -> showString "<invalid variable reference #" . shows i . showString ": " . shows n . showChar '>'
          | otherwise -> showString "<invalid variable reference #" . shows i . showString ", index out of range>"
  TmApp u v -> showParen (p > 1) $ prettyPrintTermP 1 ctx u . showChar ' ' . prettyPrintTermP 2 ctx v
  TmLet (Id name) def body -> showParen (p > 0) $ showString "let " . showString name' . showString " = " . prettyPrintTermP 0 ctx def . showString " in " . prettyPrintTermP 0 (NVarBind name' : ctx) body
    where name' = rename (varNames ctx) name
  TmTypedLet (Id name) ty def body -> showParen (p > 0) $ showString "let " . showString name' . showString " = " . prettyPrintTermP 0 ctx def . {- showString " : " . prettyPrintTypeP 1 ctx ty . -} showString " in " . prettyPrintTermP 0 (NVarBind name' : ctx) body
    where name' = rename (varNames ctx) name
  TmIf cond then_ else_ -> showParen (p > 0) $ showString "if " . prettyPrintTermP 0 ctx cond . showString " then " . prettyPrintTermP 0 ctx then_ . showString " else " . prettyPrintTermP 0 ctx else_
  TmTuple components -> showParen True $ foldr (.) id $ intersperse (showString ", ") $ map (prettyPrintTermP 0 ctx) components
  TmProj tuple i -> showParen (p > 1) $ prettyPrintTermP 1 ctx tuple . showString " @" . shows i
  TmIterate n init succ -> showParen (p > 1) $ showString "iterate " . prettyPrintTermP 2 ctx n . showChar ' ' . prettyPrintTermP 2 ctx init . showChar ' ' . prettyPrintTermP 2 ctx succ

prettyPrintTerm :: Term -> String
prettyPrintTerm t = prettyPrintTermP 0 [] t ""

max_stars_repo_name: minoki/polestar

max_issues_repo_name: minoki/polestar

max_forks_repo_name: minoki/polestar


text:
{-# LANGUAGE NoMonomorphismRestriction, ScopedTypeVariables, FlexibleContexts, TypeFamilies #-}

module Math.Probably.PCA where

import Numeric.LinearAlgebra
import Numeric.LinearAlgebra.Data
import Numeric.LinearAlgebra.HMatrix ((#>))
import Math.Probably.FoldingStats
import qualified Data.Vector.Unboxed as U
import qualified Data.Vector.Storable as VS
import Foreign.Storable.Tuple ()
import Control.Spoon
import qualified Data.Map.Strict as M
import           Data.Maybe (fromJust)
import           Data.List
import           Data.Ord
import           Data.Number.Erf
import Debug.Trace

empiricalCovariance :: [Vector Double] -> Matrix Double
empiricalCovariance xs
   = let k = length xs
         --barxk = L.scale (recip $ realToFrac $ k+1) $ sum xs
         --m1 = sum $ map (\xv-> L.outer xv xv) xs
         --m2 = L.scale (realToFrac $ k+1) $ L.outer barxk barxk
         xmn = empiricalMean xs
         f xi = outerSame $ xi - xmn
--     in L.scale (recip $ realToFrac k) $ m1 - m2
     in scale (recip $ realToFrac k - 1 ) $ sum $ map f xs

empiricalCovariance' :: Vector Double -> [Vector Double] -> Matrix Double
empiricalCovariance' xmn xs
   = let k = length xs
         --barxk = L.scale (recip $ realToFrac $ k+1) $ sum xs
         --m1 = sum $ map (\xv-> L.outer xv xv) xs
         --m2 = L.scale (realToFrac $ k+1) $ L.outer barxk barxk
         f xi = outerSame $ xi - xmn
--     in L.scale (recip $ realToFrac k) $ m1 - m2
     in scale (recip $ realToFrac k - 1 ) $ sumWith f xs

outerSame v = outer v v

--sumWithVs :: (Vector Double -> Vector Double) -> [Vector Double] -> Vector Double
--sumWithVs f = foldl1' (\acc v -> acc + f v)

sumWith :: Num a => (Vector Double -> a) -> [Vector Double] -> a
sumWith f (v:vs) = foldl' (\acc v -> acc + f v) (f v) vs

empiricalMean :: [Vector Double] -> Vector Double
empiricalMean vecs = scale (recip $ realToFrac $ length vecs) $ sumWith id vecs

-- copied from https://github.com/albertoruiz/hmatrix/blob/master/examples/pca2.hs

type Vec = Vector Double
type Mat = Matrix Double

-- Vector with the mean value of the columns of a matrix
mean :: Mat -> Vec
mean a = constant (recip . fromIntegral . rows $ a) (rows a) <> a

-- covariance matrix of a list of observations stored as rows
cov :: Mat -> Mat
cov x = (trans xc <> xc) / fromIntegral (rows x - 1)
    where xc = x - asRow (mean x)

covN :: Mat -> Mat
covN x = (trans x <> x) / fromIntegral (rows x)

cov0 :: Mat -> Mat
cov0 x = (trans x <> x) / fromIntegral (rows x - 1)
    --where xc = x - asRow (mean x)

type Stat = (Vec, [Double], Mat)
-- 1st and 2nd order statistics of a dataset (mean, eigenvalues and eigenvectors of cov)
stat :: Mat -> Stat
stat x = (m, toList s, trans v) where
    m = mean x
    (s,v) = eigSH' (cov x)

stat0 :: Mat -> Stat
stat0 x = (m, toList s, trans v) where
    m = VS.map (const 0) s
    c = cov0 x
    (s,v) = eigSH' $ trace ("cov00="++(show (c!0!0))) $ c

statSVD :: Mat -> Stat
statSVD x = (m, [], evecs) where
    m = mean x
    c = cov x
    (_, evals, evecCols) = svd c
    evecs = fromRows $ map evecSigns $ toColumns evecCols
    evecSigns ev = let maxelidx = maxIndex $ cmap abs ev
                       sign = signum (ev ! maxelidx)
                   in cmap (sign *) ev

statThinSVD :: Mat -> Stat
statThinSVD x = (m, [], evecs) where
    m = mean x
    c = cov x
    (_, evals, evecCols) = thinSVD c
    evecs = fromRows $ map evecSigns $ toColumns evecCols
    evecSigns ev = let maxelidx = maxIndex $ cmap abs ev
                       sign = signum (ev ! maxelidx)
                   in cmap (sign *) ev

pcaNSVD :: Int -> Stat -> (Vec -> Vec , Vec -> Vec)
pcaNSVD  n (m,_,evecs) = (encode,decode)
  where
    encode x = VS.take n $ evecs #> (x - m)
    decode x = error "not sure right now"

statVs :: [Vector Double] -> Stat
statVs vs = (m, toList s, trans v) where
    m = empiricalMean vs
    (s,v) = eigSH' (empiricalCovariance' m vs)

mbStat :: Mat -> Maybe Stat
mbStat = spoon . stat

pca :: Double -> Stat -> (Vec -> Vec , Vec -> Vec)
pca prec (m,s,v) = (encode,decode)
  where
    encode x = vp <> (x - m)
    decode x = x <> vp + m
    vp = takeRows n v
    n = 1 + (length $ fst $ span (< (prec'*sum s)) $ cumSum s)
    cumSum = tail . scanl (+) 0.0
    prec' = if prec <=0.0 || prec >= 1.0
                then error "the precision in pca must be 0<prec<1"
                else prec

pcaN :: Int -> Stat -> (Vec -> Vec , Vec -> Vec)
pcaN n (m,s,v) = (encode,decode)
  where
    encode x = vp <> (x - m)
    decode x = x <> vp + m
    vp = takeRows n v

type Lens a b = ( (a->b), (a-> b -> a) )

get :: Lens a b -> a -> b
get (f,_) x = f x

set :: Lens a b -> a -> b -> a
set (_,f) x y = f x y

idlens :: Lens a a
idlens = (id,  (\_ x -> x))

--, pcaFeatures
normaliseBy :: (Eq b, Ord b) => Lens a Vec -> Lens a b -> Int -> [a] -> [a]
normaliseBy veclens idlens ncomponents = equaliseVars veclens
                                         . rankTransform veclens idlens
                                         . pcaFeatures veclens ncomponents
                                         . equaliseVars veclens

normaliseNoRank :: Lens a Vec -> Int -> [a] -> [a]
normaliseNoRank veclens  ncomponents = equaliseVars veclens
                                         . pcaFeatures veclens ncomponents
                                         . equaliseVars veclens

equaliseVars ::Lens a Vec -> [a] -> [a]
equaliseVars veclens xs = equal_vars where
  decorr_vars :: VS.Vector (Double,Double)
  decorr_vars = uncurry (VS.zipWith (,)) $ runStat meanSDF $ map (get veclens) xs

  equal_vars =  map (\x -> set veclens x (g $ get veclens x)) xs

  g :: Vec -> Vec
  g fl = VS.zipWith (\x (mn,sd) -> (x-mn)/sd) fl decorr_vars

vecMeanSds :: [Vector Double] -> Vector (Double, Double)
vecMeanSds xs = uncurry (VS.zipWith (,)) $ runStat meanSDF  xs

vecMeans :: [Vector Double] -> Vector Double
vecMeans xs = scale (recip $ realToFrac $ length xs) $ sumWith id xs

rankTransform :: (Eq b, Ord b) => Lens a Vec -> Lens a b -> [a] -> [a]
rankTransform veclens idlens nmfeatures = rankTransformed where
  nfeatures = VS.length $ get veclens $ head $ nmfeatures
  nposters = realToFrac $ length nmfeatures
  sortations = flip map [0..(nfeatures-1)] $ \ix -> toRankMap $ map (get idlens) $ sortBy (comparing ((VS.! ix) . get veclens)) nmfeatures -- decorrelated
  rankTransformed = flip map nmfeatures $ \x ->
                      let nm = get idlens x
                          f sortation = invnormcdf $ realToFrac ((getRank sortation nm)+1) / (nposters+1)
                          newv =  VS.fromList $ map f sortations
                      in set veclens x newv

pcaFeatures :: Lens a Vec -> Int -> [a] -> [a]
pcaFeatures veclens ncomp nmfeatures = pca'd where
  input_data = map (get veclens ) nmfeatures
  pcaStats@(_,eigVals,_) = stat $ fromRows input_data
  (enc,_) = trace ("eigenVals ="++show eigVals ) $ pcaN ncomp pcaStats
--  pca_red = PCA.pcaReduceN input_data 5
  setFeatures x =
      let oldv = get veclens x
          newNormV = enc oldv
      in set veclens x newNormV
  pca'd = map setFeatures nmfeatures

type RankMap a = M.Map a Int

toRankMap :: (Ord a, Eq a) => [a] -> RankMap a
toRankMap = M.fromList . (`zip` [0..])

getRank :: (Ord a, Eq a) => RankMap a -> a -> Int
getRank m x = fromJust $ M.lookup x m

newtype NumUV a = NumUV { unNumUV :: U.Vector a } deriving (Eq, Ord)

numUmap f = NumUV . U.map f . unNumUV
numUlength = U.length . unNumUV

instance (U.Unbox a, Num a) => Num (NumUV a) where
    (+) = binvop (+)
    (-) = binvop (-)
    (*) = binvop (*)
    abs = numUmap abs
    signum = numUmap signum
    fromInteger x = NumUV $ U.fromList [fromInteger x]

instance (U.Unbox a, Fractional a) => Fractional (NumUV a) where
    (/) = binvop (/)
    fromRational x = NumUV $ U.fromList [fromRational x]
instance (U.Unbox a, Floating a) => Floating (NumUV a) where
    pi = NumUV $ U.fromList [pi]
    sqrt = numUmap sqrt
    log = numUmap log
    exp = numUmap exp
    sin = numUmap sin
    cos = numUmap cos
    tan = numUmap tan
    asin = numUmap asin
    acos = numUmap acos
    atan = numUmap atan
    sinh = numUmap sinh
    cosh = numUmap cosh
    tanh = numUmap tanh
    asinh = numUmap asinh
    acosh = numUmap acosh
    atanh = numUmap atanh

binvop f xs ys
  | numUlength ys == 1 = let y = (unNumUV ys) U.! 0 in numUmap (\x-> f x y) xs
  | numUlength xs == 1 = let x = (unNumUV xs) U.! 0 in numUmap (\y-> f x y) ys
  | otherwise = NumUV $ U.zipWith f (unNumUV xs) (unNumUV ys)

max_stars_repo_name: glutamate/probably-baysig

max_issues_repo_name: silky/probably-baysig

max_forks_repo_name: silky/probably-baysig


text:
{-# LANGUAGE ViewPatterns #-}
{-# OPTIONS_GHC -Wall #-}

import qualified Numeric.Sundials.ARKode.ODE as ARK
import qualified Numeric.Sundials.CVode.ODE  as CV
import           Numeric.LinearAlgebra
import           Numeric.Sundials.ODEOpts (ODEOpts(..))

import           Plots as P
import qualified Diagrams.Prelude as D
import           Diagrams.Backend.Rasterific

import           Control.Lens

import           Data.Functor.Compose

import           Test.Hspec

lorenz :: Double -> [Double] -> [Double]
lorenz _t u = [ sigma * (y - x)
              , x * (rho - z) - y
              , x * y - beta * z
              ]
  where
    rho = 28.0
    sigma = 10.0
    beta = 8.0 / 3.0
    x = u !! 0
    y = u !! 1
    z = u !! 2

_lorenzJac :: Double -> Vector Double -> Matrix Double
_lorenzJac _t u = (3><3) [ (-sigma), rho - z, y
                        , sigma   , -1.0   , x
                        , 0.0     , (-x)   , (-beta)
                        ]
  where
    rho = 28.0
    sigma = 10.0
    beta = 8.0 / 3.0
    x = u ! 0
    y = u ! 1
    z = u ! 2

brusselator :: Double -> [Double] -> [Double]
brusselator _t x = [ a - (w + 1) * u + v * u * u
                   , w * u - v * u * u
                   , (b - w) / eps - w * u
                   ]
  where
    a = 1.0
    b = 3.5
    eps = 5.0e-6
    u = x !! 0
    v = x !! 1
    w = x !! 2

brussJac :: Double -> Vector Double -> Matrix Double
brussJac _t x = tr $
  (3><3) [ (-(w + 1.0)) + 2.0 * u * v, w - 2.0 * u * v, (-w)
         , u * u                     , (-(u * u))     , 0.0
         , (-u)                      , u              , (-1.0) / eps - u
         ]
  where
    y = toList x
    u = y !! 0
    v = y !! 1
    w = y !! 2
    eps = 5.0e-6

brusselatorWithJacobian :: Vector Double -> Bool -> CV.SolverResult
brusselatorWithJacobian ts usejac = CV.odeSolveRootVWith' opts
                      (\t v -> vector $ brusselator t (toList v))
                      (if usejac then Just brussJac else Nothing)
                      (vector [1.2, 3.1, 3.0])
                      [] 0
                      ts
  where
    opts = ODEOpts { maxNumSteps = 10000
                   , minStep     = 1.0e-12
                   , maxFail     = 10
                   , odeMethod = CV.BDF
                   , stepControl = CV.XX' 1.0e-6 1.0e-10 1 1
                   , initStep = Nothing
                   }

brussRoot :: CV.SolverResult
brussRoot = CV.odeSolveRootVWith' opts
                      (\t v -> vector $ brusselator t (toList v))
                      Nothing
                      (vector [1.2, 3.1, 3.0])
                      events 100
                      (vector [0.0, 0.1 .. 10.0])
  where
    events =
      [ CV.EventSpec { CV.eventCondition = brussRootFn
                     , CV.eventUpdate =
                         \_ev x -> let y = toList x in vector [(y!!0) + 0.5 , (y!!1), (y!!2)]
                     , CV.eventDirection = CV.AnyDirection
                     }
      ]
    opts = ODEOpts { maxNumSteps = 10000
                   , minStep     = 1.0e-12
                   , maxFail     = 10
                   , odeMethod = CV.BDF
                   , stepControl = CV.XX' 1.0e-6 1.0e-10 1 1
                   , initStep = Nothing
                   }

brussRootFn :: Double -> Vector Double -> Double
brussRootFn _ v = case xs of
                    [y1, _y2, y3] -> y1 - y3
                    _            -> error "brusselator root function RHS not defined"
  where
    xs = toList v

exponential :: CV.SolverResult
exponential = CV.odeSolveRootVWith' opts
                      (\t y -> vector [y ! 0])
                      Nothing
                      (vector [1])
                      events 100
                      (vector [ fromIntegral k / 100 | k <- [0..(22::Int)]])
  where
    events =
      [ CV.EventSpec { CV.eventCondition = \t y -> y ! 0 - 1.1
                     , CV.eventUpdate = \ev y -> vector [ 2 ]
                     , CV.eventDirection = CV.Upwards
                     }
      ]
    opts = ODEOpts { maxNumSteps = 10000
                   , minStep     = 1.0e-12
                   , maxFail     = 10
                   , odeMethod = CV.BDF
                   , stepControl = CV.XX' 1.0e-6 1.0e-10 1 1
                   , initStep = Nothing
                   }

-- A sine wave that only changes direction once it reaches ±0.9.
-- Illustrates event-specific reset function
boundedSine :: CV.SolverResult
boundedSine = CV.odeSolveRootVWith'
  opts
  (\_t y -> vector [y ! 1, - y ! 0]) -- ODE RHS
  Nothing
  (vector [0, 1]) -- initial conditions
  events
  100 -- maximum number of events
  (vector [ 2 * pi * k / 360 | k <- [0..360]]) -- solution times
  where
    opts = ODEOpts { maxNumSteps = 10000
                   , minStep     = 1.0e-12
                   , maxFail     = 10
                   , odeMethod = CV.ADAMS
                   , stepControl = CV.XX' 1.0e-6 1.0e-10 1 1
                   , initStep = Nothing
                   }
    events =
      [ CV.EventSpec { CV.eventCondition = \_t y -> y ! 0 - 0.9
                     , CV.eventUpdate = \_ y -> vector [ y ! 0, - abs (y ! 1) ]
                     , CV.eventDirection = CV.Upwards
                     }
      , CV.EventSpec { CV.eventCondition = \_t y -> y ! 0 + 0.9
                     , CV.eventUpdate = \_ y -> vector [ y ! 0, abs (y ! 1) ]
                     , CV.eventDirection = CV.Downwards
                     }
      ]

stiffish :: Double -> [Double] -> [Double]
stiffish t v = [ lamda * u + 1.0 / (1.0 + t * t) - lamda * atan t ]
  where
    lamda = -100.0
    u = v !! 0

stiffishV :: Double -> Vector Double -> Vector Double
stiffishV t v = fromList [ lamda * u + 1.0 / (1.0 + t * t) - lamda * atan t ]
  where
    lamda = -100.0
    u = v ! 0

_stiffJac :: Double -> Vector Double -> Matrix Double
_stiffJac _t _v = (1><1) [ lamda ]
  where
    lamda = -100.0

predatorPrey :: Double -> [Double] -> [Double]
predatorPrey _t v = [ x * a - b * x * y
                    , d * x * y - c * y - e * y * z
                    , (-f) * z + g * y * z
                    ]
  where
    x = v!!0
    y = v!!1
    z = v!!2
    a = 1.0
    b = 1.0
    c = 1.0
    d = 1.0
    e = 1.0
    f = 1.0
    g = 1.0

roberts :: Double -> Vector Double -> Vector Double
roberts t v = vector $ robertsAux t (toList v)
  where
    robertsAux _ [y1, y2, y3] =
      [ -0.04 * y1 + 1.0e4 * y2 * y3
      , 0.04 * y1 - 1.0e4 * y2 * y3 - 3.0e7 * (y2)^(2 :: Int)
      , 3.0e7 * (y2)^(2 :: Int)
      ]
    robertsAux _ _ = error "roberts RHS not defined"

robertsJac :: Double -> Vector Double -> Matrix Double
robertsJac _t (toList -> [y1, y2, y3]) = (3 >< 3)
  [ -0.04, 1.0e4 * y3, 1.0e4 * y2
  , 0.04, -1.0e4*y3 - 3.0e7*2*y2, -1.0e4*y2
  , 0, 3.0e7*2*y2, 0
  ]

ts :: [Double]
ts = take 12 $ map (* 10.0) (0.04 : ts)

solve :: CV.SolverResult
solve = CV.odeSolveRootVWith' opts
                      roberts Nothing (vector [1.0, 0.0, 0.0])
                      events 100
                      (vector (0.0 : ts))
  where
    opts = ODEOpts { maxNumSteps = 10000
                   , minStep     = 1.0e-12
                   , maxFail     = 10
                   , odeMethod = CV.BDF
                   , stepControl = CV.ScXX' 1.0 1.0e-4 1.0 1.0 (vector [1.0e-8, 1.0e-14, 1.0e-6])
                   , initStep = Nothing
                   }
    events =
      [ CV.EventSpec { CV.eventCondition = \_t y -> y ! 0 - 0.0001
                     , CV.eventUpdate = const id
                     , CV.eventDirection = CV.AnyDirection
                     }
      , CV.EventSpec { CV.eventCondition = \_t y -> y ! 2 - 0.01
                     , CV.eventUpdate = const id
                     , CV.eventDirection = CV.AnyDirection
                     }
      ]

solve2 :: CV.SolverResult
solve2 = CV.odeSolveRootVWith' opts
                      roberts Nothing (vector [1.0, 0.0, 0.0])
                      events 100
                      (vector (0.0 : ts))
  where
    opts = ODEOpts { maxNumSteps = 10000
                   , minStep     = 1.0e-12
                   , maxFail     = 10
                   , odeMethod = CV.BDF
                   , stepControl = CV.ScXX' 1.0 1.0e-4 1.0 1.0 (vector [1.0e-8, 1.0e-14, 1.0e-6])
                   , initStep = Nothing
                   }
    events =
      [ CV.EventSpec { CV.eventCondition = \_t y -> y ! 0 - 0.0001
                     , CV.eventUpdate = upd
                     , CV.eventDirection = CV.AnyDirection
                     }
      , CV.EventSpec { CV.eventCondition = \_t y -> y ! 2 - 0.01
                     , CV.eventUpdate = upd
                     , CV.eventDirection = CV.AnyDirection
                     }
      ]
    upd _ _ = vector [1.0, 0.0, 0.0]

solve1 :: CV.SolverResult
solve1 = CV.odeSolveRootVWith' opts
                      roberts Nothing (vector [1.0, 0.0, 0.0])
                      events 100
                      (vector (0.0 : ts))
  where
    opts = ODEOpts { maxNumSteps = 10000
                   , minStep     = 1.0e-12
                   , maxFail     = 10
                   , odeMethod = CV.BDF
                   , stepControl = CV.ScXX' 1.0 1.0e-4 1.0 1.0 (vector [1.0e-8, 1.0e-14, 1.0e-6])
                   , initStep = Nothing
                   }
    events =
      [ CV.EventSpec { CV.eventCondition = \t _y -> t - 1.0
                     , CV.eventUpdate = \t y -> vector [2.0, y!1, y!2]
                     , CV.eventDirection = CV.AnyDirection
                     }
      ]

robertsonWithJacobian :: Vector Double -> Bool -> CV.SolverResult
robertsonWithJacobian ts usejac = CV.odeSolveRootVWith' opts
                      roberts (if usejac then Just robertsJac else Nothing) (vector [1.0, 0.0, 0.0])
                      [] 0
                      ts
  where
    opts = ODEOpts { maxNumSteps = 10000
                   , minStep     = 1.0e-12
                   , maxFail     = 10
                   , odeMethod = CV.BDF
                   , stepControl = CV.ScXX' 1.0 1.0e-4 1.0 1.0 (vector [1.0e-8, 1.0e-14, 1.0e-6])
                   , initStep = Nothing
                   }

lSaxis :: [[Double]] -> P.Axis B D.V2 Double
lSaxis xs = P.r2Axis &~ do
  let zs = xs!!0
      us = xs!!1
      vs = xs!!2
      ws = xs!!3
  P.linePlot' $ zip zs us
  P.linePlot' $ zip zs vs
  P.linePlot' $ zip zs ws

lSaxis2 :: [[Double]] -> P.Axis B D.V2 Double
lSaxis2 xs = P.r2Axis &~ do
  let zs = xs!!0
      us = xs!!1
      vs = xs!!2
  P.linePlot' $ zip zs us
  P.linePlot' $ zip zs vs

kSaxis :: [(Double, Double)] -> P.Axis B D.V2 Double
kSaxis xs = P.r2Axis &~ do
  P.linePlot' xs

main :: IO ()
main = do
  let res1 = ARK.odeSolve brusselator [1.2, 3.1, 3.0] (fromList [0.0, 0.1 .. 10.0])
  renderRasterific "diagrams/brusselator.png"
                   (D.dims2D 500.0 500.0)
                   (renderAxis $ lSaxis $ [0.0, 0.1 .. 10.0]:(toLists $ tr res1))

  let res1a = ARK.odeSolve brusselator [1.2, 3.1, 3.0] (fromList [0.0, 0.1 .. 10.0])
  renderRasterific "diagrams/brusselatorA.png"
                   (D.dims2D 500.0 500.0)
                   (renderAxis $ lSaxis $ [0.0, 0.1 .. 10.0]:(toLists $ tr res1a))

  let res2 = ARK.odeSolve stiffish [0.0] (fromList [0.0, 0.1 .. 10.0])
  renderRasterific "diagrams/stiffish.png"
                   (D.dims2D 500.0 500.0)
                   (renderAxis $ kSaxis $ zip [0.0, 0.1 .. 10.0] (concat $ toLists res2))

  let res2a = ARK.odeSolveV (ARK.SDIRK_5_3_4') Nothing 1e-6 1e-10 stiffishV (fromList [0.0]) (fromList [0.0, 0.1 .. 10.0])

  let res2b = ARK.odeSolveV (ARK.TRBDF2_3_3_2') Nothing 1e-6 1e-10 stiffishV (fromList [0.0]) (fromList [0.0, 0.1 .. 10.0])

  let maxDiffA = maximum $ map abs $
                 zipWith (-) ((toLists $ tr res2a)!!0) ((toLists $ tr res2b)!!0)

  let res2c = CV.odeSolveV (CV.BDF) Nothing 1e-6 1e-10 stiffishV (fromList [0.0]) (fromList [0.0, 0.1 .. 10.0])

  let maxDiffB = maximum $ map abs $
                 zipWith (-) ((toLists $ tr res2a)!!0) ((toLists $ tr res2c)!!0)

  let maxDiffC = maximum $ map abs $
                 zipWith (-) ((toLists $ tr res2b)!!0) ((toLists $ tr res2c)!!0)

  let res3 = ARK.odeSolve lorenz [-5.0, -5.0, 1.0] (fromList [0.0, 0.01 .. 20.0])

  renderRasterific "diagrams/lorenz.png"
                   (D.dims2D 500.0 500.0)
                   (renderAxis $ kSaxis $ zip ((toLists $ tr res3)!!0) ((toLists $ tr res3)!!1))

  renderRasterific "diagrams/lorenz1.png"
                   (D.dims2D 500.0 500.0)
                   (renderAxis $ kSaxis $ zip ((toLists $ tr res3)!!0) ((toLists $ tr res3)!!2))

  renderRasterific "diagrams/lorenz2.png"
                   (D.dims2D 500.0 500.0)
                   (renderAxis $ kSaxis $ zip ((toLists $ tr res3)!!1) ((toLists $ tr res3)!!2))

  let res4 = CV.odeSolve predatorPrey [0.5, 1.0, 2.0] (fromList [0.0, 0.01 .. 10.0])

  renderRasterific "diagrams/predatorPrey.png"
                   (D.dims2D 500.0 500.0)
                   (renderAxis $ kSaxis $ zip ((toLists $ tr res4)!!0) ((toLists $ tr res4)!!1))

  renderRasterific "diagrams/predatorPrey1.png"
                   (D.dims2D 500.0 500.0)
                   (renderAxis $ kSaxis $ zip ((toLists $ tr res4)!!0) ((toLists $ tr res4)!!2))

  renderRasterific "diagrams/predatorPrey2.png"
                   (D.dims2D 500.0 500.0)
                   (renderAxis $ kSaxis $ zip ((toLists $ tr res4)!!1) ((toLists $ tr res4)!!2))

  let res4a = ARK.odeSolve predatorPrey [0.5, 1.0, 2.0] (fromList [0.0, 0.01 .. 10.0])

  let maxDiffPpA = maximum $ map abs $
                   zipWith (-) ((toLists $ tr res4)!!0) ((toLists $ tr res4a)!!0)

  let cond5 =
        case solve of
          CV.SolverSuccess events _ _ -> do
            length events `shouldBe` 2
            (abs (CV.eventTime (events!!0) - 0.2640208751331032) / 0.2640208751331032 < 1.0e-8) `shouldBe` True
            (abs (CV.eventTime (events!!1) - 2.0786731062254436e7) / 2.0786731062254436e7 < 1.0e-8) `shouldBe` True
          CV.SolverError _ _ ->
            error "Root finding error!"

  let cond6 =
        case solve1 of
          CV.SolverSuccess events _ _ -> do
            length events `shouldBe` 1
            (abs (CV.eventTime (events!!0) - 1.0) / 1.0 < 1.0e-10) `shouldBe` True
          CV.SolverError _ _ ->
            error "Root finding error!"

  let cond7 =
        case solve2 of
          CV.SolverSuccess {} ->
            error "Solver returned Success"
          CV.SolverError _ _ ->
            True

  case brussRoot of
    CV.SolverSuccess events m _diagn -> do
      renderRasterific
        "diagrams/brussRoot.png"
        (D.dims2D 500.0 500.0)
        (renderAxis $ lSaxis $ toLists $ tr m)
    CV.SolverError m n ->
      expectationFailure $ show n

  let boundedSineSpec = do
        case boundedSine of
          CV.SolverSuccess events m _ -> do
            renderRasterific
              "diagrams/boundedSine.png"
              (D.dims2D 500.0 500.0)
              (renderAxis $ lSaxis2 $ toLists $ tr m)
            length events `shouldBe` 3
            map CV.rootDirection events `shouldBe` [CV.Upwards, CV.Downwards, CV.Upwards]
            map CV.eventIndex events `shouldBe` [0, 1, 0]
            all ((< 1e-8) . abs) (zipWith (-)
              (map CV.eventTime events)
              [1.1197660081724263,3.3592952656818404,5.5988203973243])
                `shouldBe` True
          CV.SolverError m n ->
            expectationFailure "Solver error"
  let exponentialSpec = do
        case exponential of
          CV.SolverSuccess events _m _diagn -> do
            length events `shouldBe` 1
            (abs (CV.eventTime (events!!0) - log 1.1) < 1e-4) `shouldBe` True
            CV.rootDirection (events!!0) `shouldBe` CV.Upwards
            CV.eventIndex (events!!0) `shouldBe` 0
          CV.SolverError m n ->
            expectationFailure $ show n

      robertsonJac = do
        let ts = vector [0, 1 .. 10]
            CV.SolverSuccess _ m1 _ = robertsonWithJacobian ts True
            CV.SolverSuccess _ m2 _ = robertsonWithJacobian ts False
        (norm_2 (m1-m2) < 1e-4) `shouldBe` True

      brusselatorJac = do
        let ts = [0.0, 0.1 .. 10.0]
            CV.SolverSuccess _ m1 _ = brusselatorWithJacobian (vector ts) True
            CV.SolverSuccess _ m2 _ = brusselatorWithJacobian (vector ts) False
        (norm_2 (m1-m2) < 1e-3) `shouldBe` True

  hspec $ describe "Compare results" $ do
    it "Robertson should fail" $ cond7
    it "Robertson time only" $ cond6
    it "Robertson from SUNDIALS manual" $ cond5
    it "Robertson with explicit Jacobian up to t=10" robertsonJac
    it "Brusselator with explicit Jacobian" brusselatorJac
    it "for SDIRK_5_3_4' and TRBDF2_3_3_2'" $ maxDiffA < 1.0e-6
    it "for SDIRK_5_3_4' and BDF" $ maxDiffB < 1.0e-6
    it "for TRBDF2_3_3_2' and BDF" $ maxDiffC < 1.0e-6
    it "for CV and ARK for the Predator Prey model" $ maxDiffPpA < 1.0e-3
    it "Bounded sine events" $ boundedSineSpec
    it "Exponential events" $ exponentialSpec

max_stars_repo_name: regnat/hmatrix-sundials

max_issues_repo_name: regnat/hmatrix-sundials

max_forks_repo_name: regnat/hmatrix-sundials


text:
{-# LANGUAGE TupleSections #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE CPP #-}
{-# LANGUAGE RecordWildCards #-}

-- |
-- Module      :  Pact.Repl.Lib
-- Copyright   :  (C) 2016 Stuart Popejoy
-- License     :  BSD-style (see the file LICENSE)
-- Maintainer  :  Stuart Popejoy <stuart@kadena.io>
--
-- Built-ins for repl functionality. Not part of standard library
-- for blockchain execution, but instead for testing and dev.
--

module Pact.Repl.Lib where

import Control.Arrow ((&&&))
import Control.Concurrent.MVar
import Control.Lens
import Control.Monad.Reader
import Control.Monad.Catch

import Data.Aeson (eitherDecode,toJSON)
import qualified Data.ByteString.Lazy as BSL
import Data.Default
import qualified Data.Map as M
import Data.Semigroup (Endo(..))
import qualified Data.Set as S
import Data.Text (Text, pack, unpack)
import qualified Data.Text as Text
import Data.Text.Encoding
import qualified Data.Vector as V

#if defined(ghcjs_HOST_OS)
import qualified Pact.Analyze.Remote.Client as RemoteClient
import Data.Maybe
#else

import Criterion
import Criterion.Types

import Control.Monad.State.Strict (get)

import Statistics.Types (Estimate(..))

import qualified Pact.Analyze.Check as Check
import qualified Pact.Types.Crypto as Crypto
import Pact.Types.Util (fromText')
#endif

import Pact.Typechecker
import qualified Pact.Types.Typecheck as TC
-- intentionally hidden unused functions to prevent lib functions from consuming gas
import Pact.Native.Internal hiding (defRNative,defGasRNative,defNative)
import qualified Pact.Native.Internal as Native
import Pact.Types.Runtime
import Pact.Eval
import Pact.Persist.Pure
import Pact.PersistPactDb
import Pact.Types.Logger
import Pact.Types.Pretty
import Pact.Repl.Types
import Pact.Native.Capabilities (evalCap)
import Pact.Gas.Table
import Pact.Types.PactValue

initLibState :: Loggers -> Maybe String -> IO LibState
initLibState loggers verifyUri = do
  m <- newMVar (DbEnv def persister
                (newLogger loggers "Repl")
                def 0 def)
  createSchema m
  return (LibState m Noop def def verifyUri M.empty M.empty)

-- | Native function with no gas consumption.
type ZNativeFun e = FunApp -> [Term Ref] -> Eval e (Term Name)

zeroGas :: Eval e a -> Eval e (Gas,a)
zeroGas = fmap (0,)

defZNative :: NativeDefName -> ZNativeFun e -> FunTypes (Term Name) -> [Example] -> Text -> NativeDef
defZNative name fun = Native.defNative name $ \fi as -> zeroGas $ fun fi as

defZRNative :: NativeDefName -> RNativeFun e -> FunTypes (Term Name) -> [Example] -> Text -> NativeDef
defZRNative name fun = Native.defNative name (reduced fun)
    where reduced f fi as = mapM reduce as >>= zeroGas . f fi

replDefs :: NativeModule
replDefs = ("Repl",
     [
      defZRNative "load" load (funType tTyString [("file",tTyString)] <>
                              funType tTyString [("file",tTyString),("reset",tTyBool)])
      [LitExample "(load \"accounts.repl\")"]
      "Load and evaluate FILE, resetting repl state beforehand if optional RESET is true."
     ,defZRNative "format-address" formatAddr (funType tTyString [("scheme", tTyString), ("public-key", tTyString)])
      []
      "Transform PUBLIC-KEY into an address (i.e. a Pact Runtime Public Key) depending on its SCHEME."
     ,defZRNative "env-keys" setsigs (funType tTyString [("keys",TyList tTyString)])
      ["(env-keys [\"my-key\" \"admin-key\"])"]
      "Set transaction signature KEYS."
     ,defZRNative "env-data" setmsg (funType tTyString [("json",json)])
      ["(env-data { \"keyset\": { \"keys\": [\"my-key\" \"admin-key\"], \"pred\": \"keys-any\" } })"]
      "Set transaction JSON data, either as encoded string, or as pact types coerced to JSON."

     ,defZRNative "continue-pact" continuePact
      (funType tTyString [("step",tTyInteger)] <>
       funType tTyString [("step",tTyInteger),("rollback",tTyBool)] <>
       funType tTyString [("step",tTyInteger),("rollback",tTyBool),("pact-id",tTyString)] <>
       funType tTyString [("step",tTyInteger),("rollback",tTyBool),("pact-id",tTyString),("yielded",tTyObject (mkSchemaVar "y"))])
      [LitExample "(continue-pact 1)", LitExample "(continue-pact 1 true)",
       LitExample "(continue-pact 1 false \"[pact-id-hash]\"))",
       LitExample "(continue-pact 2 1 false \"[pact-id-hash]\" { \"rate\": 0.9 })"]
      ("Continue previously-initiated pact identified STEP, " <>
       "optionally specifying ROLLBACK (default is false), " <>
       "PACT-ID of the pact to be continued (defaults to the pact initiated in the current transaction, if one is present), and " <>
       "YIELDED value to be read with 'resume' (if not specified, uses yield in most recent pact exec, if any).")

     ,defZRNative "pact-state" pactState
      (funType (tTyObject TyAny) [] <> funType (tTyObject TyAny) [("clear",tTyBool)])
      [LitExample "(pact-state)", LitExample "(pact-state true)"]
      ("Inspect state from most recent pact execution. Returns object with fields " <>
      "'pactId': pact ID; 'yield': yield result or 'false' if none; 'step': executed step; " <>
      "'executed': indicates if step was skipped because entity did not match. " <>
      "With CLEAR argument, erases pact from repl state.")

     ,defZRNative "env-entity" setentity
      (funType tTyString [] <> funType tTyString [("entity",tTyString)])
      [LitExample "(env-entity \"my-org\")", LitExample "(env-entity)"]
      ("Set environment confidential ENTITY id, or unset with no argument.")
     ,defZRNative "begin-tx" (tx Begin) (funType tTyString [] <>
                                        funType tTyString [("name",tTyString)])
      [LitExample "(begin-tx \"load module\")"] "Begin transaction with optional NAME."
     ,defZRNative "commit-tx" (tx Commit) (funType tTyString []) [LitExample "(commit-tx)"] "Commit transaction."
     ,defZRNative "rollback-tx" (tx Rollback) (funType tTyString []) [LitExample "(rollback-tx)"] "Rollback transaction."
     ,defZRNative "expect" expect (funType tTyString [("doc",tTyString),("expected",a),("actual",a)])
      ["(expect \"Sanity prevails.\" 4 (+ 2 2))"]
      "Evaluate ACTUAL and verify that it equals EXPECTED."
     ,defZNative "expect-failure" expectFail (funType tTyString [("doc",tTyString),("exp",a)])
      ["(expect-failure \"Enforce fails on false\" (enforce false \"Expected error\"))"]
      "Evaluate EXP and succeed only if it throws an error."
     ,defZNative "bench" bench' (funType tTyString [("exprs",TyAny)])
      [LitExample "(bench (+ 1 2))"] "Benchmark execution of EXPRS."
     ,defZRNative "typecheck" tc (funType tTyString [("module",tTyString)] <>
                                 funType tTyString [("module",tTyString),("debug",tTyBool)])
       []
       "Typecheck MODULE, optionally enabling DEBUG output."
     ,defZRNative "env-gaslimit" setGasLimit (funType tTyString [("limit",tTyInteger)])
       []
       "Set environment gas limit to LIMIT."
     ,defZRNative "env-gas" envGas (funType tTyInteger [] <> funType tTyString [("gas",tTyInteger)])
       []
       "Query gas state, or set it to GAS."
     ,defZRNative "env-gasprice" setGasPrice (funType tTyString [("price",tTyDecimal)])
       []
       "Set environment gas price to PRICE."
     ,defZRNative "env-gasrate" setGasRate (funType tTyString [("rate",tTyInteger)])
       []
       "Update gas model to charge constant RATE."
     ,defZRNative "env-gasmodel" setGasModel (funType tTyString [("model",tTyString)])
       []
       "Update gas model to the model named MODEL."
     ,defZRNative "verify" verify (funType tTyString [("module",tTyString)])
       []
       "Verify MODULE, checking that all properties hold."

     ,defZRNative "sig-keyset" sigKeyset (funType tTyKeySet [])
       []
       "Convenience function to build a keyset from keys present in message signatures, using 'keys-all' as the predicate."
     ,defZRNative "print" print' (funType tTyString [("value",a)])
       []
       "Output VALUE to terminal as unquoted, unescaped text."
     ,defZRNative "env-hash" envHash (funType tTyString [("hash",tTyString)])
       ["(env-hash (hash \"hello\"))"]
       "Set current transaction hash. HASH must be an unpadded base64-url encoded BLAKE2b 256-bit hash."
     ,defZNative "test-capability" testCapability
      (funType tTyString [("capability", TyFun $ funType' tTyBool [])])
      [LitExample "(test-capability (MY-CAP))"] $
     "Specify and request grant of CAPABILITY. Once granted, CAPABILITY and any composed capabilities are in scope " <>
     "for the rest of the transaction. Allows direct invocation of capabilities, which is not available in the " <>
     "blockchain environment."
     ,defZRNative "mock-spv" mockSPV
       (funType tTyString [("type",tTyString),("payload",tTyObject TyAny),("output",tTyObject TyAny)])
      [LitExample "(mock-spv \"TXOUT\" { 'proof: \"a54f54de54c54d89e7f\" } { 'amount: 10.0, 'account: \"Dave\", 'chainId: \"1\" })"]
       "Mock a successful call to 'spv-verify' with TYPE and PAYLOAD to return OUTPUT."
     , envChainDataDef
     ])
     where
       json = mkTyVar "a" [tTyInteger,tTyString,tTyTime,tTyDecimal,tTyBool,
                         TyList (mkTyVar "l" []),TySchema TyObject (mkSchemaVar "o") def,tTyKeySet]
       a = mkTyVar "a" []

invokeEnv :: (MVar (DbEnv PureDb) -> IO b) -> MVar LibState -> IO b
invokeEnv f e = withMVar e $ \ls -> f $! _rlsPure ls
{-# INLINE invokeEnv #-}

repldb :: PactDb LibState
repldb = PactDb {

    _readRow = \d k -> invokeEnv $ _readRow pactdb d k
  , _writeRow = \wt d k v -> invokeEnv $ _writeRow pactdb wt d k v
  , _keys = \t -> invokeEnv $ _keys pactdb t
  , _txids = \t tid -> invokeEnv $ _txids pactdb t tid
  , _createUserTable = \t m -> invokeEnv $ _createUserTable pactdb t m
  , _getUserTableInfo = \t -> invokeEnv $ _getUserTableInfo pactdb t
  , _beginTx = \tid -> invokeEnv $ _beginTx pactdb tid
  , _commitTx = invokeEnv $ _commitTx pactdb
  , _rollbackTx = invokeEnv $ _rollbackTx pactdb
  , _getTxLog = \d t -> invokeEnv $ _getTxLog pactdb d t

}

load :: RNativeFun LibState
load _ [TLitString fn] = setop (Load (unpack fn) False) >> return (tStr $ "Loading " <> fn <> "...")
load _ [TLitString fn, TLiteral (LBool r) _] = setop (Load (unpack fn) r) >> return (tStr $ "Loading " <> fn <> "...")
load i as = argsError i as

modifyLibState :: (LibState -> (LibState,a)) -> Eval LibState a
modifyLibState f = view eePactDbVar >>= \m -> liftIO $ modifyMVar m (return . f)

setLibState :: (LibState -> LibState) -> Eval LibState ()
setLibState f = modifyLibState (f &&& const ())

viewLibState :: (LibState -> a) -> Eval LibState a
viewLibState f = modifyLibState (id &&& f)

setop :: LibOp -> Eval LibState ()
setop v = setLibState $ set rlsOp v

setenv :: Setter' (EvalEnv LibState) a -> a -> Eval LibState ()
setenv l v = setop $ UpdateEnv $ Endo (set l v)

mockSPV :: RNativeFun LibState
mockSPV i as = case as of
  [TLitString spvType, TObject payload _, TObject out _] -> do
    setLibState $ over rlsMockSPV (M.insert (SPVMockKey (spvType,payload)) out)
    return $ tStr $ "Added mock SPV for " <> spvType
  _ -> argsError i as

formatAddr :: RNativeFun LibState
#if !defined(ghcjs_HOST_OS)
formatAddr i [TLitString scheme, TLitString cryptoPubKey] = do
  let eitherEvalErr :: Either String a -> String -> (a -> b) -> Eval LibState b
      eitherEvalErr res effectStr transformFunc =
        case res of
          Left e  -> evalError' i $ prettyString effectStr <> ": " <> prettyString e
          Right v -> return (transformFunc v)
  sppk  <- eitherEvalErr (fromText' scheme)
           "Invalid PPKScheme"
           Crypto.toScheme
  pubBS <- eitherEvalErr (parseB16TextOnly cryptoPubKey)
           "Invalid Public Key format"
           Crypto.PubBS
  addr  <- eitherEvalErr (Crypto.formatPublicKeyBS sppk pubBS)
           "Unable to convert Public Key to Address"
           toB16Text
  return (tStr addr)
formatAddr i as = argsError i as
#else
formatAddr i _ = evalError' i "Address formatting not supported in GHCJS"
#endif

setsigs :: RNativeFun LibState
setsigs i [TList ts _ _] = do
  ks <- forM ts $ \t -> case t of
          (TLitString s) -> return s
          _ -> argsError i (V.toList ts)
  setenv eeMsgSigs (S.fromList (map (PublicKey . encodeUtf8) (V.toList ks)))
  return $ tStr "Setting transaction keys"
setsigs i as = argsError i as

setmsg :: RNativeFun LibState
setmsg i as = case as of
  [TLitString j] ->
    case eitherDecode (BSL.fromStrict $ encodeUtf8 j) of
      Left f -> evalError' i ("Invalid JSON: " <> prettyString f)
      Right v -> go v
  [TObject (Object om _ _ _) _] -> go (toJSON (fmap toPactValueLenient om))
  [a] -> go (toJSON a)
  _ -> argsError i as
  where go v = setenv eeMsgBody v >> return (tStr "Setting transaction data")

continuePact :: RNativeFun LibState
continuePact i as = case as of
  [TLitInteger step] ->
    go step False Nothing Nothing
  [TLitInteger step,TLitBool rollback] ->
    go step rollback Nothing Nothing
  [TLitInteger step,TLitBool rollback,TLitString pid] ->
    go step rollback (Just pid) Nothing
  [TLitInteger step,TLitBool rollback,TLitString pid,TObject (Object o _ _ _) _] ->
    go step rollback (Just pid) (Just o)
  _ -> argsError i as
  where
    go :: Integer -> Bool -> Maybe Text -> Maybe (ObjectMap (Term Name)) -> Eval LibState (Term Name)
    go step rollback pid userResume = do
      pe <- use evalPactExec
      (pactId, y) <- unwrapExec pid userResume pe

      let pactStep = PactStep (fromIntegral step) rollback pactId y
      viewLibState (view rlsPacts) >>= \pacts ->
        case M.lookup pactId pacts of
          Nothing -> evalError' i $ "Invalid pact id: " <> pretty pactId
          Just PactExec{..} -> do
            evalPactExec .= Nothing
            local (set eePactStep $ Just pactStep) $ resumePact (_faInfo i) Nothing

    unwrapExec mpid mobj Nothing = do
      pid <- case mpid of
        Nothing -> evalError' i
          "continue-pact: No pact id supplied and no pact exec in context"
        Just pid -> return $ PactId pid
      y <- maybe (return Nothing) (toYield Nothing) mobj
      return (pid, y)
    unwrapExec mpid mobj (Just ex) = do
      let pid = maybe (_pePactId ex) PactId mpid
      y <- case mobj of
        Nothing -> return $ _peYield ex
        Just o -> case _peYield ex of
          Just (Yield _ p) -> toYield p o
          Nothing -> toYield Nothing o
      return (pid, y)

    toYield p = fmap (Just . flip Yield p) . enforcePactValue'

setentity :: RNativeFun LibState
setentity i as = case as of
  [TLitString s] -> do
    setenv eeEntity $ Just (EntityName s)
    return (tStr $ "Set entity to " <> s)
  [] -> do
    setenv eeEntity Nothing
    return (tStr "Unset entity")
  _ -> argsError i as

pactState :: RNativeFun LibState
pactState i as = case as of
  [] -> go False
  [TLitBool clear] -> go clear
  _ -> argsError i as
  where
    go clear = do
      e <- use evalPactExec
      when clear $ evalPactExec .= Nothing
      case e of
        Nothing -> evalError' i "pact-state: no pact exec in context"
        Just PactExec{..} -> return $ toTObject TyAny def $
          [("yield",maybe (toTerm False) (toTObjectMap TyAny def . fmap fromPactValue . _yData) _peYield)
          ,("executed",toTerm _peExecuted)
          ,("step",toTerm _peStep)
          ,("pactId",toTerm _pePactId)]

txmsg :: Maybe Text -> Maybe TxId -> Text -> Term Name
txmsg n tid s = tStr $ s <> " Tx " <> pack (show tid) <> maybe "" (": " <>) n

tx :: Tx -> RNativeFun LibState
tx Begin i as = do
  tname <- case as of
             [TLitString n] -> return $ Just n
             [] -> return Nothing
             _ -> argsError i as
  setop $ Tx (_faInfo i) Begin tname
  setLibState $ set rlsTxName tname
  return (tStr "")
tx t i [] = do
  tname <- modifyLibState (set rlsTxName Nothing &&& view rlsTxName)
  setop (Tx (_faInfo i) t tname)
  return (tStr "")
tx _ i as = argsError i as

recordTest :: Text -> Maybe (FunApp,Text) -> Eval LibState ()
recordTest name failure = setLibState $ over rlsTests (++ [TestResult name failure])

testSuccess :: Text -> Text -> Eval LibState (Term Name)
testSuccess name msg = recordTest name Nothing >> return (tStr msg)

testFailure :: FunApp -> Text -> Text -> Eval LibState (Term Name)
testFailure i name msg = recordTest name (Just (i,msg)) >> return (tStr msg)

expect :: RNativeFun LibState
expect i [TLitString a,b,c] =
  if b `termEq` c
  then testSuccess a $ "Expect: success: " <> a
  else testFailure i a $ renderCompactText' $
       "FAILURE: " <> pretty a <> ": expected " <> pretty b <> ":" <> pretty (typeof' b) <>
       ", received " <> pretty c <> ":" <> pretty (typeof' c)
expect i as = argsError i as

expectFail :: ZNativeFun LibState
expectFail i as@[a,b] = do
  a' <- reduce a
  case a' of
    TLitString msg -> do
      r <- catch (Right <$> reduce b) (\(_ :: SomeException) -> return $ Left ())
      case r of
        Right v -> testFailure i msg $ "FAILURE: " <> msg <> ": expected failure, got result = " <> pack (show v)
        Left _ -> testSuccess msg $ "Expect failure: success: " <> msg
    _ -> argsError' i as
expectFail i as = argsError' i as

bench' :: ZNativeFun LibState
#if !defined(ghcjs_HOST_OS)
bench' i as = do
  e <- ask
  s <- get
  (r :: Either SomeException Report) <-
      try $ liftIO $ benchmark' $ whnfIO $ runEval s e $ do
                !ts <- mapM reduce as
                return $! toTerm (length ts)
  case r of
    Left ex -> evalError' i (viaShow ex)
    Right rpt -> do
           let mean = estPoint (anMean (reportAnalysis rpt))
               sd = estPoint (anStdDev (reportAnalysis rpt))
               (reg,_) = splitAt 1 $ anRegress (reportAnalysis rpt)
               val = case reg of
                       [] -> mean
                       (r':_) -> case M.lookup "iters" (regCoeffs r') of
                                  Nothing -> mean
                                  Just t -> estPoint t
               tps = 1/val
               tperr = (1/(val - (sd/2))) - (1/(val + (sd/2)))
           liftIO $ putStrLn $ show (round tps :: Integer) ++ "/s, +-" ++ show (round tperr :: Integer) ++ "/s"
           return (tStr "Done")
#else
bench' i _ = evalError' i "Benchmarking not supported in GHCJS"
#endif

tc :: RNativeFun LibState
tc i as = case as of
  [TLitString s] -> go s False
  [TLitString s,TLiteral (LBool d) _] -> go s d
  _ -> argsError i as
  where
    go modname dbg = do
      md <- getModule i (ModuleName modname Nothing)
      r :: Either TC.CheckerException ([TC.TopLevel TC.Node],[TC.Failure]) <-
        try $ liftIO $ typecheckModule dbg md
      case r of
        Left (TC.CheckerException ei e) -> evalError ei ("Typechecker Internal Error: " <> prettyString e)
        Right (_,fails) -> case fails of
          [] -> return $ tStr $ "Typecheck " <> modname <> ": success"
          _ -> do
            setop $ TcErrors $ map (\(TC.Failure ti s) -> renderInfo (TC._tiInfo ti) ++ ":Warning: " ++ s) fails
            return $ tStr $ "Typecheck " <> modname <> ": Unable to resolve all types"

verify :: RNativeFun LibState
verify i as = case as of
  [TLitString modName] -> do
    md <- getModule i (ModuleName modName Nothing)
    -- reading all modules from db here, but should be fine in repl
    modules <- getAllModules i
#if defined(ghcjs_HOST_OS)
    uri <- fromMaybe "localhost" <$> viewLibState (view rlsVerifyUri)
    renderedLines <- liftIO $
                     RemoteClient.verifyModule modules md uri
#else
    modResult <- liftIO $ Check.verifyModule modules md
    let renderedLines = Check.renderVerifiedModule modResult
#endif
    setop $ TcErrors $ Text.unpack <$> renderedLines
    return $ tStr $ "Verification of " <> modName <> " failed"

  _ -> argsError i as

sigKeyset :: RNativeFun LibState
sigKeyset _ _ = view eeMsgSigs >>= \ss -> return $ toTerm $ KeySet (S.toList ss) (Name (asString KeysAll) def)

print' :: RNativeFun LibState
print' _ [v] = setop (Print v) >> return (tStr "")
print' i as = argsError i as

envHash :: RNativeFun LibState
envHash i [TLitString s] = case fromText' s of
  Left err -> evalError' i $ "Bad hash value: " <> pretty s <> ": " <> prettyString err
  Right h -> do
    setenv eeHash h
    return $ tStr $ "Set tx hash to " <> s
envHash i as = argsError i as

setGasLimit :: RNativeFun LibState
setGasLimit _ [TLitInteger l] = do
  setenv (eeGasEnv . geGasLimit) (fromIntegral l)
  return $ tStr $ "Set gas limit to " <> tShow l
setGasLimit i as = argsError i as

envGas :: RNativeFun LibState
envGas _ [] = use evalGas >>= \g -> return (tLit $ LInteger $ fromIntegral g)
envGas _ [TLitInteger g] = do
  evalGas .= fromIntegral g
  return $ tStr $ "Set gas to " <> tShow g
envGas i as = argsError i as

setGasPrice :: RNativeFun LibState
setGasPrice _ [TLiteral (LDecimal d) _] = do
  setenv (eeGasEnv . geGasPrice) (wrap (wrap d))
  return $ tStr $ "Set gas price to " <> tShow d
setGasPrice i as = argsError i as

setGasRate :: RNativeFun LibState
setGasRate _ [TLitInteger r] = do
  let model = constGasModel $ fromIntegral r
  setenv (eeGasEnv . geGasModel) model
  return $ tStr $ "Set gas model to " <> gasModelDesc model
setGasRate i as = argsError i as

setGasModel :: RNativeFun LibState
setGasModel _ [] = do
  model <- asks (_geGasModel . _eeGasEnv)
  return $ tStr $ "Current gas model is '" <> gasModelName model <> "': " <> gasModelDesc model
setGasModel _ as = do
  let mMod = case as of
        [TLitString "table"] -> Just $ tableGasModel defaultGasConfig
        [TLitString "fixed", TLitInteger r] -> Just $ constGasModel (fromIntegral r)
        _ -> Nothing
  case mMod of
    Nothing -> return $ tStr "Unrecognized model, perhaps try (env-gasmodel \"table\") or (env-gasmodel \"fixed\" 1)"
    Just model -> do
      setenv (eeGasEnv . geGasModel) model
      return $ tStr $ "Set gas model to " <> gasModelDesc model

-- | This is the only place we can do an external call to a capability,
-- using 'evalCap False'.
testCapability :: ZNativeFun ReplState
testCapability _ [ c@TApp{} ] = do
  cap <- evalCap False $ _tApp c
  return . tStr $ case cap of
    Nothing -> "Capability granted"
    Just cap' -> "Capability granted: " <> tShow cap'
testCapability i as = argsError' i as

-- | Modify existing env chain data with new data, replacing just those
-- environment items only.
envChainDataDef :: NativeDef
envChainDataDef = defZRNative "env-chain-data" envChainData
    (funType tTyString [("new-data", tTyObject TyAny)])
    ["(env-chain-data { \"chain-id\": \"TestNet00/2\", \"block-height\": 20 })"]
    "Update existing entries 'chain-data' with NEW-DATA, replacing those items only."
  where
    envChainData :: RNativeFun LibState
    envChainData i as = case as of
      [TObject (Object (ObjectMap ks) _ _ _) _] -> do
        pd <- view eePublicData

        ud <- foldM (go (_faInfo i)) pd (M.toList ks)
        setenv eePublicData ud
        return $ tStr "Updated public metadata"
      _ -> argsError i as

    go i pd ((FieldKey k), (TLiteral (LInteger l) _)) = case Text.unpack k of
      "gas-limit"    -> pure $ set (pdPublicMeta . pmGasLimit) (wrap (wrap l)) pd
      "block-height" -> pure $ set pdBlockHeight (fromIntegral l) pd
      "block-time"   -> pure $ set pdBlockTime (fromIntegral l) pd
      t              -> evalError i $ "envChainData: bad public metadata key: " <> prettyString t

    go i pd ((FieldKey k), (TLiteral (LDecimal l) _)) = case Text.unpack k of
      "gas-price" -> pure $ set (pdPublicMeta . pmGasPrice) (wrap (wrap l)) pd
      t           -> evalError i $ "envChainData: bad public metadata key: " <> prettyString t

    go i pd ((FieldKey k), (TLiteral (LString l) _)) = case Text.unpack k of
      "chain-id" -> pure $ set (pdPublicMeta . pmChainId) (ChainId l) pd
      "sender"   -> pure $ set (pdPublicMeta . pmSender) l pd
      t          -> evalError i $ "envChainData: bad public metadata key: " <> prettyString t
    go i _ as = evalError i $ "envChainData: bad public metadata values: " <> pretty as

max_stars_repo_name: wsdjeg/pact

max_issues_repo_name: wsdjeg/pact

max_forks_repo_name: wsdjeg/pact