{-# LANGUAGE AllowAmbiguousTypes       #-}
{-# LANGUAGE DataKinds                 #-}
{-# LANGUAGE DeriveGeneric             #-}
{-# LANGUAGE PolyKinds                 #-}
{-# LANGUAGE FlexibleContexts          #-}
{-# LANGUAGE FlexibleInstances         #-}
{-# LANGUAGE GADTs                     #-}
{-# LANGUAGE NoMonomorphismRestriction #-}
{-# LANGUAGE OverloadedStrings         #-}
{-# LANGUAGE ScopedTypeVariables       #-}
{-# LANGUAGE TypeApplications          #-}
{-# LANGUAGE TypeFamilies              #-}
{-# LANGUAGE TypeOperators             #-}
{-# LANGUAGE StandaloneDeriving        #-}
{-# LANGUAGE UndecidableInstances        #-}

module BlueRipple.Model.MRP where

import qualified BlueRipple.Data.Keyed         as K
import qualified BlueRipple.Data.DataFrames    as BR
import           BlueRipple.Data.CountFolds
import qualified Control.Foldl                 as FL
import           Control.Monad                  ( join )
import qualified Control.Monad.State           as State
import qualified Data.Array                    as A
import           Data.Function                  ( on )
import qualified Data.List                     as L
import qualified Data.Set                      as S
import qualified Data.IntMap                   as IM
import qualified Data.Map                      as M
import           Data.Maybe                     ( isJust
                                                , fromJust
                                                )
import qualified Data.Text                     as T
import qualified Data.Profunctor               as P
import qualified Data.Serialize                as SE
import           Data.Serialize.Text ()
import qualified Data.Vector                   as V
import qualified Data.Vector.Storable          as VS

import qualified Frames                        as F
import qualified Frames.Melt                   as F
import qualified Frames.InCore                 as FI
import qualified Data.Vinyl                    as V
import qualified Data.Vinyl.TypeLevel          as V

import qualified Control.MapReduce             as MR
import qualified Frames.Transform              as FT
import qualified Frames.Folds                  as FF
import qualified Frames.MapReduce              as FMR
import qualified Frames.Enumerations           as FE
--import qualified Frames.Misc                  as FU
import qualified Frames.Serialize              as FS

import qualified Knit.Report                   as K hiding (elements)
import qualified Polysemy.Error                as P
                                                ( mapError )
import qualified Polysemy                      as P
                                                ( raise )
import           Text.Pandoc.Error             as PE
import qualified Text.Blaze.Colonnade          as BC
import qualified Text.Show

import qualified Numeric.LinearAlgebra         as LA

import qualified Data.IndexedSet               as IS
import qualified Numeric.GLM.ProblemTypes      as GLM
import qualified Numeric.GLM.ModelTypes        as GLM
import qualified Numeric.GLM.FunctionFamily    as GLM
import           Numeric.GLM.MixedModel        as GLM
import qualified Numeric.GLM.Bootstrap         as GLM
import qualified Numeric.GLM.Report            as GLM
import qualified Numeric.GLM.Predict           as GLM
import qualified Numeric.GLM.Confidence        as GLM
import qualified Numeric.SparseDenseConversions
                                               as SD

import qualified Relude.Extra as Relude

import qualified Statistics.Types              as ST
import           GHC.Generics                   ( Generic, Rep )

{- Moved to BlueRipple.Data.CountFolds
-- map reduce folds for counting
type Count = "Count" F.:-> Int
type Successes = "Successes" F.:-> Int
type MeanWeight = "MeanWeight" F.:-> Double
type VarWeight = "VarWeight" F.:-> Double
type WeightedSuccesses = "WeightedSuccesses" F.:-> Double
type UnweightedSuccesses =  "UnweightedSuccesses" F.:-> Int

binomialFold
  :: (F.Record r -> Bool) -> FL.Fold (F.Record r) (F.Record '[Count, Successes])
binomialFold testRow =
  let successesF = FL.premap (\r -> if testRow r then 1 else 0) FL.sum
  in  (\n s -> n F.&: s F.&: V.RNil) <$> FL.length <*> successesF

countFold
  :: forall k r d
   . ( Ord (F.Record k)
     , FI.RecVec (k V.++ '[Count, Successes])
     , k F.⊆ r
     , d F.⊆ r
     )
  => (F.Record d -> Bool)
  -> FL.Fold (F.Record r) [F.FrameRec (k V.++ '[Count, Successes])]
countFold testData = MR.mapReduceFold
  MR.noUnpack
  (FMR.assignKeysAndData @k)
  (FMR.foldAndAddKey $ binomialFold testData)

type CountCols = '[Count, UnweightedSuccesses, WeightedSuccesses, MeanWeight, VarWeight]

zeroCount :: F.Record CountCols
zeroCount = 0 F.&: 0 F.&: 0 F.&: 1 F.&: 0 F.&: V.RNil

weightedBinomialFold
  :: (F.Record r -> Bool)
  -> (F.Record r -> Double)
  -> FL.Fold (F.Record r) (F.Record CountCols)
weightedBinomialFold testRow weightRow =
  let wSuccessesF =
        FL.premap (\r -> if testRow r then weightRow r else 0) FL.sum
      successesF  = FL.premap (\r -> if testRow r then 1 else 0) FL.sum
      meanWeightF = FL.premap weightRow FL.mean
      varWeightF  = FL.premap weightRow FL.variance
      f s ws mw = if mw < 1e-6 then realToFrac s else ws / mw -- if meanweight is 0 but
  in  (\n s ws mw vw -> n F.&: s F.&: f s ws mw F.&: mw F.&: vw F.&: V.RNil)
      <$> FL.length
      <*> successesF
      <*> wSuccessesF
      <*> meanWeightF
      <*> varWeightF

weightedCountFold
  :: forall k r d
   . (Ord (F.Record k)
     , FI.RecVec (k V.++ CountCols)
     , k F.⊆ r
     , k F.⊆ (k V.++ r)
     , d F.⊆ (k V.++ r)
     )
  => (F.Record r -> Bool) -- ^ count this row?
  -> (F.Record d -> Bool) -- ^ success ?
  -> (F.Record d -> Double) -- ^ weight
  -> FL.Fold (F.Record r) (F.FrameRec (k V.++ CountCols))
weightedCountFold {- filterData testData weightData-} = weightedCountFoldGeneral (F.rcast @k)
{-  FMR.concatFold $ FMR.mapReduceFold
    (MR.filterUnpack filterData)
    (FMR.assignKeysAndData @k)
    (FMR.foldAndAddKey $ weightedBinomialFold testData weightData)
-}

weightedCountFoldGeneral
  :: forall k r d
   . (Ord (F.Record k)
     , FI.RecVec (k V.++ CountCols)
     , k F.⊆ (k V.++ r)
     , d F.⊆ (k V.++ r)
     )
  => (F.Record r -> F.Record k)
  -> (F.Record r -> Bool) -- ^ include this row?
  -> (F.Record d -> Bool) -- ^ success ?
  -> (F.Record d -> Double) -- ^ weight
  -> FL.Fold (F.Record r) (F.FrameRec (k V.++ CountCols))
weightedCountFoldGeneral getKey filterData testData weightData =
  FL.prefilter filterData $ FMR.concatFold $ FMR.mapReduceFold
    (MR.Unpack $ \r ->  [getKey r `V.rappend` r])
    (FMR.assignKeysAndData @k @d)
    (FMR.foldAndAddKey $ weightedBinomialFold testData weightData)

-}

getFraction r =
  let n = F.rgetField @Count r
  in  if n == 0
        then 0
        else realToFrac (F.rgetField @Successes r) / realToFrac n

getFractionWeighted r =
  let n = F.rgetField @Count r
  in  if n == 0 then 0 else F.rgetField @WeightedSuccesses r / realToFrac n

data RecordColsProxy k = RecordColsProxy deriving (Show, Enum, Bounded, A.Ix, Eq, Ord)
type instance GLM.GroupKey (RecordColsProxy k) = F.Record k

--type instance GLM.GroupKey (Proxy k) = F.Record k

recordToGroupKey
  :: forall k r . (k F.⊆ r) => F.Record r -> RecordColsProxy k -> F.Record k
recordToGroupKey r _ = F.rcast @k r

glmErrorToPandocError :: GLM.GLMError -> PE.PandocError
glmErrorToPandocError x = PE.PandocSomeError $ show x

type GroupCols gs cs = gs V.++ cs
type CountedCols gs cs = GroupCols gs cs V.++ CountCols

---
newtype  SimplePredictor ps = SimplePredictor { unSimplePredictor :: F.Record ps }

instance (Show (F.Record ps)) => Show (SimplePredictor ps) where
  show (SimplePredictor x) = "SimplePredictor " ++ show @String x

instance (Eq (F.Record ps)) => Eq (SimplePredictor ps) where
  (SimplePredictor x) == (SimplePredictor y) = x == y

instance (Ord (F.Record ps)) => Ord (SimplePredictor ps) where
  compare (SimplePredictor x) (SimplePredictor y) = compare x y

instance (Ord (F.Record ps), K.FiniteSet (F.Record ps)) => Enum (SimplePredictor ps) where
  toEnum n =
    let im = IM.fromList $ zip [0..] $ SimplePredictor <$> S.toAscList K.elements
    in fromMaybe (error "Bad n given to toEnum :: Int -> SimplePredictor") $ IM.lookup n im
  fromEnum a =
    let m = M.fromList $ zip (SimplePredictor <$> S.toAscList K.elements) [0..]
    in fromMaybe (error "Bad SimplePredictor given to fromEnum :: SimplePredictor -> Int.  This means something is very bad.") $ M.lookup a m

instance (K.FiniteSet (F.Record ps)) => Bounded (SimplePredictor ps) where
  minBound = fromMaybe (error "Empty Set of Records used as predictor") $ viaNonEmpty head $ SimplePredictor <$> S.toList K.elements
  maxBound = fromMaybe (error "Empty Set of Records used as predictor") $ viaNonEmpty last $ SimplePredictor <$> S.toList K.elements

allSimplePredictors :: K.FiniteSet (F.Record ps) => [SimplePredictor ps]
allSimplePredictors = SimplePredictor <$> S.toList K.elements

type SimpleEffect ps = GLM.WithIntercept (SimplePredictor ps)

simplePredictor :: forall ps rs. (ps F.⊆ rs
                                     , Eq (F.Record ps)
                                     )
                    => F.Record rs -> SimplePredictor ps -> Double
simplePredictor r p = if F.rcast @ps r == unSimplePredictor p then 1 else 0

predMap :: forall cs. (K.FiniteSet (F.Record cs), Ord (F.Record cs), cs F.⊆ cs)
  => F.Record cs -> M.Map (SimplePredictor cs) Double
predMap r =  M.fromList $ Relude.fmapToSnd (simplePredictor r) allSimplePredictors

catPredMaps :: forall cs.  (K.FiniteSet (F.Record cs), Ord (F.Record cs), cs F.⊆ cs)
  => M.Map (F.Record cs) (M.Map (SimplePredictor cs) Double)
catPredMaps = M.fromList $ fmap (\k -> (unSimplePredictor k,predMap (unSimplePredictor k))) allSimplePredictors

data  LocationHolder c f a =  LocationHolder { locName :: T.Text
                                             , locKey :: Maybe (F.Rec f LocationCols)
                                             , catData :: M.Map (F.Rec f c) a
                                             } deriving (Generic)

deriving instance (V.RMap c
                  , V.ReifyConstraint Show F.ElField c
                  , V.RecordToList c
                  , Show a) => Show (LocationHolder c F.ElField a)

instance (SE.Serialize a
         , Ord (F.Rec FS.SElField c)
         , SE.GSerializePut
           (Rep (F.Rec FS.SElField c))
         , SE.GSerializeGet (Rep (F.Rec FS.SElField c))
         , (Generic (F.Rec FS.SElField c))
         ) => SE.Serialize (LocationHolder c FS.SElField a)

lhToS :: (Ord (F.Rec FS.SElField c)
         , V.RMap c
         )
      => LocationHolder c F.ElField a -> LocationHolder c FS.SElField a
lhToS (LocationHolder n lkM cdm) = LocationHolder n (fmap FS.toS lkM) (M.mapKeys FS.toS cdm)

lhFromS :: (Ord (F.Rec F.ElField c)
           , V.RMap c
         ) => LocationHolder c FS.SElField a -> LocationHolder c F.ElField a
lhFromS (LocationHolder n lkM cdm) = LocationHolder n (fmap FS.fromS lkM) (M.mapKeys FS.fromS cdm)

type LocationCols = '[BR.StateAbbreviation]
locKeyPretty :: F.Record LocationCols -> T.Text
locKeyPretty r =
  let stateAbbr = F.rgetField @BR.StateAbbreviation r
  in stateAbbr

--type ASER = '[BR.SimpleAgeC, BR.SexC, BR.CollegeGradC, BR.SimpleRaceC]
predictionsByLocation ::
  forall ps r rs. ( V.RMap ps
                  , V.ReifyConstraint Show V.ElField ps
                  , V.RecordToList ps
                  , Ord (F.Record ps)
                  , Ord (SimplePredictor ps)
                  , FI.RecVec (ps V.++ CountCols)
                  , F.ElemOf (ps V.++ CountCols) Count
                  , F.ElemOf (ps V.++ CountCols) MeanWeight
                  , F.ElemOf (ps V.++ CountCols) UnweightedSuccesses
                  , F.ElemOf (ps V.++ CountCols) VarWeight
                  , F.ElemOf (ps V.++ CountCols) WeightedSuccesses
                  , K.FiniteSet (F.Record ps)
                  , Show (F.Record (LocationCols V.++ ps V.++ CountCols))
                  , Show (F.Record ps)
                  , Enum (SimplePredictor ps)
                  , Bounded (SimplePredictor ps)
                  , (ps V.++ CountCols) F.⊆ (LocationCols V.++ ps V.++ CountCols)
                  , ps F.⊆ (ps V.++ CountCols)
                  , ps F.⊆ (LocationCols V.++ ps V.++ CountCols)
                  , F.ElemOf rs BR.StateAbbreviation
                  , K.KnitEffects r
               )
  => GLM.MinimizeDevianceVerbosity
  -> F.FrameRec rs
  -> FL.Fold (F.Record rs) (F.FrameRec (LocationCols V.++ ps V.++ CountCols))
  -> [SimpleEffect ps]
  -> M.Map (F.Record ps) (M.Map (SimplePredictor ps) Double)
  -> K.Sem r [LocationHolder ps V.ElField Double]
predictionsByLocation verbosity ccesFrame countFold predictors catPredMap =
  P.mapError glmErrorToPandocError  $ K.wrapPrefix "predictionsByLocation" $ do
    K.logLE K.Diagnostic "Inferring"
    (mm, rc, ebg, bu, vb, bs) <- inferMR @LocationCols @ps @ps
                                 verbosity
                                 countFold
                                 predictors
                                 simplePredictor
                                 ccesFrame

    let states = FL.fold FL.set $ fmap (F.rgetField @BR.StateAbbreviation) ccesFrame
        allStateKeys = (F.&: V.RNil) <$> FL.fold FL.list states
        predictLoc l = LocationHolder (locKeyPretty l) (Just l) catPredMap
        toPredict = [LocationHolder "National" Nothing catPredMap] <> fmap predictLoc allStateKeys
        predict (LocationHolder n lkM cpms) = P.mapError glmErrorToPandocError $ do
          let predictFrom catKey predMap =
                let groupKeyM = fmap (`V.rappend` catKey) lkM --lkM >>= \lk -> return $ lk `V.rappend` catKey
                    emptyAsNationalGKM = case groupKeyM of
                                           Nothing -> Nothing
                                           Just k -> k <$ GLM.categoryNumberFromKey rc k (RecordColsProxy @(LocationCols V.++ ps))
                in (fmap snd . GLM.runLogOnGLMException Nothing)
                   $ GLM.predictFromBetaB mm (`M.lookup` predMap) (const emptyAsNationalGKM) rc ebg bu vb
          cpreds <- M.traverseWithKey predictFrom cpms
          return $ LocationHolder n lkM cpreds
    traverse predict toPredict

---

-- TODO: Add bootstraps back in, make optional
-- Maybe set up way to serialize all this at this level??

-- ls: are keys for locations (groups with no matching fixed-effect)
-- cs: are keys for fixedEffects, one for each constructor in b
-- b: Fixed Effect type
-- g: group type
inferMR
  :: forall ls cs ks b g f rs effs
   . ( Foldable f
     , Functor f
     , K.KnitEffects effs
     , F.RDeleteAll ls (ls V.++ ks V.++ CountCols) ~ (ks V.++ CountCols)
     , (ls V.++ ks V.++ CountCols) ~ (ls V.++ (ks V.++ CountCols))
     , Ord (F.Record ls)
     , ls F.⊆ (ls V.++ ks V.++ CountCols)
     , (ks V.++ CountCols) F.⊆ (ls V.++ ks V.++ CountCols)
     , ks F.⊆ (ks V.++ CountCols)
     , (ls V.++ cs) F.⊆ (ls V.++ (ks V.++ CountCols))
     , Show (F.Record (ls V.++ ks V.++ CountCols))
     , FI.RecVec (ls V.++ (ks V.++ CountCols))
     , F.ElemOf (ks V.++ CountCols) Count
     , F.ElemOf (ks V.++ CountCols) MeanWeight
     , F.ElemOf (ks V.++ CountCols) VarWeight
     , F.ElemOf (ks V.++ CountCols) WeightedSuccesses
     , F.ElemOf (ks V.++ CountCols) UnweightedSuccesses
     , F.ElemOf (ls V.++ ks V.++ CountCols) Count
     , F.ElemOf (ls V.++ ks V.++ CountCols) MeanWeight
     , F.ElemOf (ls V.++ ks V.++ CountCols) VarWeight
     , F.ElemOf (ls V.++ ks V.++ CountCols) WeightedSuccesses
     , F.ElemOf (ls V.++ ks V.++ CountCols) UnweightedSuccesses
     , Show (F.Record (ls V.++ cs))
     , K.FiniteSet (F.Record ks)
     , Ord b
     , Show b
     , Enum b
     , Bounded b
     , Ord (F.Record (GroupCols ls cs))
     , g ~ RecordColsProxy (GroupCols ls cs)
     )
  => GLM.MinimizeDevianceVerbosity
  -> FL.Fold (F.Record rs) (F.FrameRec (ls V.++ ks V.++ CountCols))
  -> [GLM.WithIntercept b] -- fixed effects to fit
  -> (F.Record (ls V.++ ks V.++ CountCols) -> b -> Double)  -- how to get a fixed effect value from a record
  -> f (F.Record rs)
  -> K.Sem
       effs
       ( GLM.MixedModel b g
       , GLM.RowClassifier g
       , GLM.EffectsByGroup g b
       , GLM.BetaVec
       , LA.Vector Double
       , [(GLM.BetaVec, LA.Vector Double)]
       )
inferMR verbosity cf fixedEffectList getFixedEffect rows =
  P.mapError glmErrorToPandocError
  $ (fmap snd . GLM.runLogOnGLMException Nothing)
  $ K.wrapPrefix "inferMR"
  $ do
    let
      addZeroCountsF = FMR.concatFold $ FMR.mapReduceFold
                       FMR.noUnpack
                       (FMR.splitOnKeys @ls)
                       ( FMR.makeRecsWithKey id
                         $ FMR.ReduceFold
                         $ const
                         $ K.addDefaultRec @ks zeroCount
                       )
      counted = FL.fold FL.list $ FL.fold addZeroCountsF $ FL.fold cf rows
--    K.logLE K.Diagnostic $ T.intercalate "\n" $ fmap (T.pack . show) counted
    let
      vCounts = VS.fromList $ fmap (F.rgetField @Count) counted
      designEffect mw vw = 1 + (vw / (mw * mw))
      vWeights = VS.fromList $ fmap
        (\r ->
            let mw = F.rgetField @MeanWeight r
                vw = F.rgetField @VarWeight r
            in  case (mw < 1e-12, vw < 1e-12) of
              (True, _) -> 0
              (False, True) -> 1
              _ -> 1/sqrt (designEffect mw vw)
        )
        counted -- VS.replicate (VS.length vCounts) 1.0

      fixedEffects = GLM.FixedEffects $ IS.fromList fixedEffectList
      groups       = IS.fromList [RecordColsProxy]
      (observations, fixedEffectsModelMatrix, rcM) = FL.fold
        (lmePrepFrame getFractionWeighted
         fixedEffects
         groups
         getFixedEffect
         (recordToGroupKey @(GroupCols ls cs))
        )
        counted
--        K.logLE K.Diagnostic $ "vCounts=" <> (T.pack $ show vCounts)
--        K.logLE K.Diagnostic $ "vWeights=" <> (T.pack $ show vWeights)
--        K.logLE K.Diagnostic $ "mX=" <> (T.pack $ show fixedEffectsModelMatrix)
--        K.logLE K.Diagnostic $ "vY=" <> (T.pack $ show observations)
    let regressionModelSpec = GLM.RegressionModelSpec
          fixedEffects
          fixedEffectsModelMatrix
          observations
    rowClassifier <- case rcM of
      Left  msg -> K.knitError msg
      Right x   -> return x
--        K.logLE K.Diagnostic $ "rc=" <> (T.pack $ show rowClassifier)
    let effectsByGroup =
          M.fromList [(RecordColsProxy, IS.fromList [GLM.Intercept])]
    fitSpecByGroup <- GLM.fitSpecByGroup @b @g fixedEffects
                      effectsByGroup
                      rowClassifier
    let lmmControls = GLM.LMMControls GLM.LMM_BOBYQA 1e-6 Nothing
        lmmSpec     = GLM.LinearMixedModelSpec
                      (GLM.MixedModelSpec regressionModelSpec fitSpecByGroup)
                      lmmControls
        cc = GLM.PIRLSConvergenceCriterion GLM.PCT_Deviance 1e-6 20
        glmmControls = GLM.GLMMControls GLM.UseCanonical 10 cc
        glmmSpec = GLM.GeneralizedLinearMixedModelSpec
                   lmmSpec
                   vWeights
                   (GLM.Binomial vCounts)
                   glmmControls
        mixedModel = GLM.GeneralizedLinearMixedModel glmmSpec
    randomEffectsModelMatrix <- GLM.makeZ fixedEffectsModelMatrix
                                fitSpecByGroup
                                rowClassifier
--        K.logLE K.Diagnostic $ "smZ=" <> (T.pack $ show randomEffectsModelMatrix)
    let randomEffectCalc = GLM.RandomEffectCalculated
          randomEffectsModelMatrix
          (GLM.makeLambda fitSpecByGroup)
        th0         = GLM.setCovarianceVector fitSpecByGroup 1 0
        -- mdVerbosity = MDVNone
    GLM.checkProblem mixedModel randomEffectCalc
    K.logLE K.Info "Fitting data..."
    ((th, pd, sigma2, beta, mzvu, mzvb, cs), vMuSol, cf) <- GLM.minimizeDeviance
                                                            verbosity
                                                            ML
                                                            mixedModel
                                                            randomEffectCalc
                                                            th0
    vb <- K.knitMaybe "b-vector (random effects coefficients) is zero in MRP.inferMR" $ GLM.useMaybeZeroVec Nothing Just mzvb
    GLM.report mixedModel
      randomEffectsModelMatrix
      beta
      (SD.toSparseVector vb)
    let fes = GLM.fixedEffectStatistics mixedModel sigma2 cs beta
    K.logLE K.Diagnostic $ "FixedEffectStatistics: " <> show fes
    epg <- GLM.effectParametersByGroup @g @b rowClassifier effectsByGroup vb
--        K.logLE K.Diagnostic
--          $  "EffectParametersByGroup: "
--          <> (T.pack $ show epg)
    gec <- GLM.effectCovariancesByGroup effectsByGroup mixedModel sigma2 th
    K.logLE K.Diagnostic
      $  "EffectCovariancesByGroup: "
      <> show gec
    rebl <- GLM.randomEffectsByLabel epg rowClassifier
    K.logLE K.Diagnostic
      $  "Random Effects:\n"
      <> GLM.printRandomEffectsByLabel rebl
    smCondVar <- GLM.conditionalCovariances mixedModel
                 cf
                 randomEffectCalc
                 th
                 beta
                 mzvu
    let bootstraps                           = []
    let GLM.FixedEffectStatistics _ mBetaCov = fes
    let f r = do
                let obs = getFractionWeighted r
                predictCVCI <- GLM.predictWithCI
                  mixedModel
                  (Just . getFixedEffect r)
                  (Just . recordToGroupKey @(GroupCols ls cs) r)
                  rowClassifier
                  effectsByGroup
                  beta
                  vb
                  (ST.mkCL 0.95)
                  (GLM.NaiveCondVarCI mBetaCov smCondVar)
                return (r, obs, predictCVCI)
    fitted <- traverse f (FL.fold FL.list counted)
    K.logLE K.Diagnostic
      $  "Fitted:\n"
      <> (T.intercalate "\n" $ fmap show fitted)
    fixedEffectTable <- GLM.printFixedEffects fes
    K.logLE K.Diagnostic $ "FixedEffects:\n" <> fixedEffectTable
    let GLM.FixedEffectStatistics fep _ = fes
    return
      (mixedModel, rowClassifier, effectsByGroup, beta, vb, bootstraps) -- fes, epg, rowClassifier, bootstraps)

lmePrepFrame
  :: forall p g rs
   . (GLM.PredictorC p, GLM.GroupC g)
  => (F.Record rs -> Double) -- ^ observations
  -> GLM.FixedEffects p
  -> IS.IndexedSet g
  -> (F.Record rs -> p -> Double) -- ^ predictors
  -> (F.Record rs -> g -> GLM.GroupKey g)  -- ^ classifiers
  -> FL.Fold
       (F.Record rs)
       ( LA.Vector Double
       , LA.Matrix Double
       , Either T.Text (GLM.RowClassifier g)
       ) -- ^ (X,y,(row-classifier, size of class))
lmePrepFrame observationF fe groupIndices getPredictorF classifierLabelF
  = let
      makeInfoVector
        :: M.Map g (M.Map (GLM.GroupKey g) Int)
        -> M.Map g (GLM.GroupKey g)
        -> Either T.Text (V.Vector (GLM.ItemInfo g))
      makeInfoVector indexMaps groupKeys =
        let
          g (grp, groupKey) =
            GLM.ItemInfo
              <$> (   maybe
                      (Left $ "Failed on " <> show (grp, groupKey))
                      Right
                  $   M.lookup grp indexMaps
                  >>= M.lookup groupKey
                  )
              <*> pure groupKey
        in  fmap V.fromList $ traverse g $ M.toList groupKeys
      makeRowClassifier
        :: Traversable f
        => M.Map g (M.Map (GLM.GroupKey g) Int)
        -> f (M.Map g (GLM.GroupKey g))
        -> Either T.Text (GLM.RowClassifier g)
      makeRowClassifier indexMaps labels = do
        let sizes = fmap M.size indexMaps
        indexed <- traverse (makeInfoVector indexMaps) labels
        return $ GLM.RowClassifier groupIndices
                                   sizes
                                   (V.fromList $ FL.fold FL.list indexed)
                                   indexMaps
      getPredictorF' _   GLM.Intercept     = 1
      getPredictorF' row (GLM.Predictor x) = getPredictorF row x
      predictorF row = LA.fromList $ case fe of
        GLM.FixedEffects indexedFixedEffects ->
          getPredictorF' row <$> IS.members indexedFixedEffects
        GLM.InterceptOnly -> [1]
      getClassifierLabels :: F.Record rs -> M.Map g (GLM.GroupKey g)
      getClassifierLabels r =
        M.fromList $ Relude.fmapToSnd (classifierLabelF r) $ IS.members groupIndices
      foldObs   = LA.fromList <$> FL.premap observationF FL.list
      foldPred  = LA.fromRows <$> FL.premap predictorF FL.list
      foldClass = FL.premap getClassifierLabels FL.list
      g (vY, mX, ls) =
        ( vY
        , mX
        , makeRowClassifier
          (snd $ State.execState (addAll ls) (M.empty, M.empty))
          ls
        )
    in
      g <$> ((,,) <$> foldObs <*> foldPred <*> foldClass)

addOne
  :: GLM.GroupC g
  => (g, GLM.GroupKey g)
  -> State.State (M.Map g Int, M.Map g (M.Map (GLM.GroupKey g) Int)) ()
addOne (grp, label) = do
  (nextIndexMap, groupIndexMaps) <- State.get
  let groupIndexMap = fromMaybe M.empty $ M.lookup grp groupIndexMaps
  case M.lookup label groupIndexMap of
    Nothing -> do
      let index         = fromMaybe 0 $ M.lookup grp nextIndexMap
          nextIndexMap' = M.insert grp (index + 1) nextIndexMap
          groupIndexMaps' =
            M.insert grp (M.insert label index groupIndexMap) groupIndexMaps
      State.put (nextIndexMap', groupIndexMaps')
      return ()
    _ -> return ()

addMany
  :: (GLM.GroupC g, Traversable h)
  => h (g, GLM.GroupKey g)
  -> State.State (M.Map g Int, M.Map g (M.Map (GLM.GroupKey g) Int)) ()
addMany = traverse_ addOne

addAll
  :: GLM.GroupC g
  => [M.Map g (GLM.GroupKey g)]
  -> State.State (M.Map g Int, M.Map g (M.Map (GLM.GroupKey g) Int)) ()
addAll = traverse_ (addMany . M.toList)

-- useful data folds
weightedSumF :: (Real w, Real v, Fractional y) => FL.Fold (w, v) y
weightedSumF =
  (/)
    <$> FL.premap (\(w, v) -> realToFrac w * realToFrac v) FL.sum
    <*> P.dimap fst realToFrac FL.sum

weightedSumRecF
  :: forall w v y rs
   . ( V.KnownField w
     , V.KnownField v
     , Fractional y
     , F.ElemOf rs w
     , F.ElemOf rs v
     , Real (V.Snd w)
     , Real (V.Snd v)
     )
  => FL.Fold (F.Record rs) y
weightedSumRecF =
  FL.premap (\r -> (F.rgetField @w r, F.rgetField @v r)) weightedSumF

sumProdIfRecF
  :: forall w v t y rs
   . ( V.KnownField w
     , V.KnownField v
     , V.KnownField t
     , Fractional y
     , F.ElemOf rs w
     , F.ElemOf rs v
     , F.ElemOf rs t
     , Real (V.Snd w)
     , Real (V.Snd v)
     )
  => (V.Snd t -> Bool)
  -> FL.Fold (F.Record rs) y
sumProdIfRecF test = FL.prefilter (test . F.rgetField @t) $ FL.premap
  (\r -> realToFrac (F.rgetField @w r) * realToFrac (F.rgetField @v r))
  FL.sum


{-# LANGUAGE GADTs,RebindableSyntax,CPP,FlexibleContexts,FlexibleInstances,ConstraintKinds #-}
{-# LANGUAGE StandaloneDeriving,DeriveDataTypeable #-}
{-# OPTIONS_GHC -dcore-lint #-}
{-
 - This test suite ensures that the rewrites that HerbiePlugin performs
 - give the correct results.
 -}

module Main
    where

import SubHask

import System.IO
-- import Data.Complex
-- import Linear.Quaternion
-- import Linear.V3
-- import Linear.Vector

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

test1a :: Double -> Double -> Double
test1a far near = -(2 * far * near) / (far - near)

{-# ANN test1b "NoHerbie" #-}
test1b :: Double -> Double -> Double
test1b far near = -(2 * far * near) / (far - near)

{-# ANN test1c "NoHerbie" #-}
test1c :: Double -> Double -> Double
test1c far near = if far < -1.7210442634149447e81
    then ((-2 * far) / (far - near)) * near
    else if far < 8.364504563556443e16
        then -2 * far * (near / (far - near))
        else ((-2 * far) / (far - near)) * near

{-
--------------------

test2a :: Double -> Double -> Double
test2a a b = a + ((b - a) / 2)

{-# ANN test2b "NoHerbie" #-}
test2b :: Double -> Double -> Double
test2b a b = a + ((b - a) / 2)

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

-- test3a :: Quaternion Double -> Quaternion Double -> Quaternion Double
-- test3a (Quaternion q0 (V3 q1 q2 q3)) (Quaternion r0 (V3 r1 r2 r3)) =
--     Quaternion (r0*q0+r1*q1+r2*q2+r3*q3)
--                (V3 (r0*q1-r1*q0-r2*q3+r3*q2)
--                    (r0*q2+r1*q3-r2*q0-r3*q1)
--                    (r0*q3-r1*q2+r2*q1-r3*q0))
--                ^/ (r0*r0 + r1*r1 + r2*r2 + r3*r3)
--
-- {-# ANN test3b "NoHerbie" #-}
-- test3b :: Quaternion Double -> Quaternion Double -> Quaternion Double
-- test3b (Quaternion q0 (V3 q1 q2 q3)) (Quaternion r0 (V3 r1 r2 r3)) =
--     Quaternion (r0*q0+r1*q1+r2*q2+r3*q3)
--                (V3 (r0*q1-r1*q0-r2*q3+r3*q2)
--                    (r0*q2+r1*q3-r2*q0-r3*q1)
--                    (r0*q3-r1*q2+r2*q1-r3*q0))
--                ^/ (r0*r0 + r1*r1 + r2*r2 + r3*r3)

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

data Yo a = Yo
    { yo_x2y :: a
    , yo_y2x :: a
    }

test4 :: Real a => a -> a -> Yo a
test4 x y = Yo
    { yo_x2y = x * x * y
    , yo_y2x = y * y * x
    }

test5 :: Float -> Float -> Float
test5 x y = (x * x) + (2 * x * y) + (y * y)

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

-- asinh_ :: Complex Double -> Complex Double
-- asinh_ x = log (x + sqrt (1.0+x*x))
--
-- acosh_ :: Complex Double -> Complex Double
-- acosh_ x = log (x + (x+1.0) * sqrt ((x-1.0)/(x+1.0)))

atanh_ :: Double -> Double
atanh_ x = 0.5 * log ((1.0+x) / (1.0-x))

--------------------------------------------------------------------------------
-}

#define mkTest(f1,f2,a,b) \
    putStrLn $ "mkTest: " ++ show (f1 (a) (b)); \
    putStrLn $ "mkTest: " ++ show (f2 (a) (b)); \
    putStrLn ""

#define mkTestB(f1,f2,a) \
    putStrLn $ "mkTest: " ++ show (f1 (a)); \
    putStrLn $ "mkTest: " ++ show (f2 (a)); \
    putStrLn ""

main = do
    mkTest(test1a,test1b,-2e90,6)
    mkTest(test1a,test1b,3,4)
    mkTest(test1a,test1b,2e90,6)

    mkTest(test1a,test1c,-2e90,6)
    mkTest(test1a,test1c,3,4)
    mkTest(test1a,test1c,2e90,6)

{-
    mkTest(test2a,test2b,1,2)

--     mkTest(test3a,test3b,(Quaternion 1 (V3 1 2 3)),(Quaternion 2 (V3 2 3 4)))

--     mkTestB(asinh,asinh_,5e-17::Complex Double)
--     mkTestB(acosh,acosh_,5e-17::Complex Double)
    mkTestB(atanh,atanh_,5e-17::Double)
-}

    putStrLn "done"


module Main where

import Control.DeepSeq (NFData)
import Data.Complex (Complex((:+)))
import Data.Function (on)

test1 :: Int
test1 = undefined

test2 :: a -> a -> Complex a
test2 = (:+)

test25 :: NFData a => a
test25 = undefined

test3 :: (b -> b -> c) -> (a -> b) -> a -> a -> c
test3 = on

test4 :: IO ()
test4 = putStrLn "Bar"

test5 :: [t] -> ()
test5 (_:_) = ()
test5 _ = error "test5"

-- hlint
test6 :: [Integer] -> [Integer]
test6 = map ((+ 1) . (* 2))


module Math.HiddenMarkovModel.Named (
   T(..),
   Discrete,
   Gaussian,
   fromModelAndNames,
   toCSV,
   fromCSV,
   ) where

import qualified Math.HiddenMarkovModel.Distribution as Distr
import qualified Math.HiddenMarkovModel.Private as HMM
import qualified Math.HiddenMarkovModel.CSV as HMMCSV
import Math.HiddenMarkovModel.Distribution (State(..))
import Math.HiddenMarkovModel.Utility (attachOnes)

import qualified Numeric.LinearAlgebra.Algorithms as Algo
import qualified Data.Packed.Vector as Vector

import qualified Text.CSV.Lazy.String as CSV
import Text.Printf (printf)

import qualified Control.Monad.Exception.Synchronous as ME
import qualified Control.Monad.Trans.State as MS
import Control.DeepSeq (NFData, rnf)
import Foreign.Storable (Storable)

import qualified Data.Map as Map
import qualified Data.List as List
import Data.Tuple.HT (swap)
import Data.Map (Map)

{- |
A Hidden Markov Model with names for each state.

Although 'nameFromStateMap' and 'stateFromNameMap' are exported
you must be careful to keep them consistent when you alter them.
-}
data T distr prob =
   Cons {
      model :: HMM.T distr prob,
      nameFromStateMap :: Map State String,
      stateFromNameMap :: Map String State
   }
   deriving (Show, Read)

type Discrete prob symbol = T (Distr.Discrete prob symbol) prob
type Gaussian a = T (Distr.Gaussian a) a

instance
   (NFData distr, NFData prob, Storable prob) =>
      NFData (T distr prob) where
   rnf hmm = rnf (model hmm, nameFromStateMap hmm, stateFromNameMap hmm)

fromModelAndNames :: HMM.T distr prob -> [String] -> T distr prob
fromModelAndNames md names =
   let m = Map.fromList $ zip [State 0 ..] names
   in  Cons {
          model = md,
          nameFromStateMap = m,
          stateFromNameMap = inverseMap m
       }

inverseMap :: Map State String -> Map String State
inverseMap =
   Map.fromListWith (error "duplicate label") .
   map swap . Map.toList

toCSV ::
   (Distr.CSV distr, Algo.Field prob, Show prob) =>
   T distr prob -> String
toCSV hmm =
   CSV.ppCSVTable $ snd $ CSV.toCSVTable $ HMMCSV.padTable "" $
      Map.elems (nameFromStateMap hmm) : HMM.toCells (model hmm)

fromCSV ::
   (Distr.CSV distr, Algo.Field prob, Read prob) =>
   String -> ME.Exceptional String (T distr prob)
fromCSV =
   MS.evalStateT parseCSV . map HMMCSV.fixShortRow . CSV.parseCSV

parseCSV ::
   (Distr.CSV distr, Algo.Field prob, Read prob) =>
   HMMCSV.CSVParser (T distr prob)
parseCSV = do
   names <- HMMCSV.parseStringList =<< HMMCSV.getRow
   let duplicateNames =
         Map.keys $ Map.filter (> (1::Int)) $
         Map.fromListWith (+) $ attachOnes names
    in HMMCSV.assert (null duplicateNames) $
          "duplicate names: " ++ List.intercalate ", " duplicateNames
   md <- HMM.parseCSV
   let n = length names
       m = Vector.dim (HMM.initial md)
    in HMMCSV.assert (n == m) $
          printf "got %d state names for %d state" n m
   return $ fromModelAndNames md names


{-# LANGUAGE ViewPatterns #-}

module Edge where

import Codec.Image.DevIL
import Data.Array.Unboxed
import Math.Probably.MCMC
import Math.Probably.Sampler
import Math.Probably.FoldingStats
import Control.Monad.State.Strict
import System.Cmd
import System.Environment
import Data.Array.IO
import System.IO

import Data.Array.Unboxed
import Numeric.LinearAlgebra hiding (find)
import qualified Math.Probably.PDF as PDF
import Data.List hiding (map)
import Data.Maybe
import Control.Applicative
import Data.Ord
import qualified Data.Vector as V
import qualified Graphics.Rendering.OpenGL.GLU.Tessellation as GLU
import Graphics.Rendering.OpenGL

import CVUtils
import qualified Graphics.UI.GLUT as GLUT

loadPoints :: IO [(R,R)]
loadPoints = fmap read $ readFile "edge.dat"

complexPolygon :: [(R,R)] -> GLU.ComplexPolygon GLfloat
complexPolygon points
  = let p2v (x,y) = Vertex3 (realToFrac x) (realToFrac y) 0
    in GLU.ComplexPolygon
         [GLU.ComplexContour $map (\v-> GLU.AnnotatedVertex (p2v v) 0) points]

noOpCombiner _newVertex _weightedProperties = 0.0 ::GLfloat

triangulateEdge :: [(R,R)] -> IO [[Vector R]]
triangulateEdge pts = do
   fmap getTriangles $ GLU.triangulate
            GLU.TessWindingPositive 0 (Normal3 0 0 0) noOpCombiner
            $ complexPolygon pts

type Triangles = [[Vector R]]

getTriangles (GLU.Triangulation tris) = map unTri tris where
   unTri
    (GLU.Triangle (GLU.AnnotatedVertex (Vertex3 x1 y1 _) _)
                  (GLU.AnnotatedVertex (Vertex3 x2 y2 _) _)
                  (GLU.AnnotatedVertex (Vertex3 x3 y3 _) _))
       = [fromList [realToFrac x1,realToFrac y1],
          fromList [realToFrac x2,realToFrac y2],
          fromList [realToFrac x3,realToFrac y3]]

--http://frame3dd.svn.sourceforge.net/viewvc/frame3dd/trunk/src/microstran/vec3.c?view=markup
cross a b = fromList [(a@>1)*(b@>2)-(a@>2)*(b@>1)
                     ,(a@>2)*(b@>0)-(a@>0)*(b@>2)
                     ,(a@>0)*(b@>1)-(a@>1)*(b@>0)]

crossZ x y = (x@>0)*(y@>1)-(x@>1)*(y@>0)

crossZp a b p = ((b@>0)-(a@>0))*((p@>1)-(a@>1))-((b@>1)-(a@>1))*((p@>0)-(a@>0))

--http://www.blackpawn.com/texts/pointinpoly/default.html

sameSide :: (Vector R,Vector R,Vector R,Vector R) -> Bool
sameSide(p1,p2, a,b)
 = let cp1 = crossZ (b-a) (p1-a)
       cp2 = crossZ (b-a) (p2-a)
    in cp1 * cp2 >= 0

sameSide' :: (Vector R,Vector R,Vector R,Vector R) -> Bool
sameSide'(p1,p2, a,b)
 = let cp1 = crossZp a b p1
       cp2 = crossZp a b p2
    in cp1 * cp2 >= 0

pointInTriangle :: Vector R -> [Vector R] -> Bool
pointInTriangle p [a,b,c]
    = sameSide'(p,a, b,c) && sameSide'(p,b, a,c)
        && sameSide'(p,c, a,b)

markBg :: Image -> IO Image
markBg im = do
    (progName,args) <-  GLUT.getArgsAndInitialize
    pts <- loadPoints
    tris <- triangulateEdge pts
 --   print tris
    print $ head tris
    --print $ cross (head tris!!0) (head tris!!1)

--    print $ pointInTriangle (fromList [100,100]) $ head tris
    mutIm <- thaw im
    ((loy,lox,_),(hiy,hix,_)) <- getBounds (mutIm::MImage)
    forM_ [(y,x,0) | x<- [lox..hix],
                     y<- [loy..hiy]] $ \ix@(y,x,_)-> do
      when (any (pointInTriangle (fromList [realToFrac x, realToFrac y])) tris) $ writeArray mutIm ix 255
    freeze (mutIm::MImage)
--    return im

loadTriangles :: IO [[Vector R]]
loadTriangles = do
    (progName,args) <-  GLUT.getArgsAndInitialize
    pts <- loadPoints
    triangulateEdge pts

writeVisMat :: String -> BitImage -> Image -> IO ()
writeVisMat nm vm im = do
   mutIm <- thaw im
   ((loy,lox,_),(hiy,hix,_)) <- getBounds (mutIm::MImage)
   forM_ [(y,x,0) | x<- [lox..hix],
                     y<- [loy..hiy]] $ \ix@(y,x,_)-> do
     when (not $ readBitImage vm x y) $ writeArray mutIm ix 255

   newIm <- freeze (mutIm::MImage)
   writeImage nm newIm
   return ()


module Main where

import Correlation (cacf)
import Data.Complex (Complex (..))
import Graphics.Rendering.Chart.Backend.Diagrams
import Graphics.Rendering.Chart.Easy
import Linear.Metric
import Linear.V3
import Linear.Vector
import MLS (mls)

j :: RealFloat a => Complex a
j = 0 :+ 1

main :: IO ()
main = do
  print "IDDQD"
  print code
  toFile def "example1_big.png" $ do
    layout_title .= "Amplitude Modulation"
    setColors [opaque blue, opaque red]
    plot (line "am" [signal' [0, (0.5) .. 400]])
    plot (points "am points" (signal' [0, 7 .. 400]))

signal' :: [Double] -> [(Double, Double)]
signal' xs = [(x, (sin (x * 3.14159 / 45) + 1) / 2 * (sin (x * 3.14159 / 5))) | x <- xs]

code = mls [True, False, True, False, False, True, False, False, True, False, False, True] [True, False, False, False, False, False, False, False, False, False, False]

base = length code

carrierFeq :: Double
carrierFeq = 1e10

codeInterval :: Double
codeInterval = 1e-6

-- Теорема Котельникова
df :: Double
df = (carrierFeq + fromIntegral base / codeInterval) * 3

signal :: Double -> Complex Double
signal t
  | t >= 0 && t < signalLength = a t * exp (- j * (2 * pi * carrierFeq * t :+ 0))
  | otherwise = 0
  where
    signalLength = codeInterval * fromIntegral base
    i = floor (t / codeInterval) `mod` base
    a t =
      if code !! i
        then 1
        else -1

echo :: (Double -> V3 Double) -> Double -> Complex Double
echo obj t = signal (t - d / c)
  where
    d = distance (carrier t) (obj t)
    c = 3e8

scene :: [V3 Double]
scene = [V3 0 10000 0]

zeroPoint :: V3 Double
zeroPoint = V3 0 0 0

vel :: V3 Double
vel = V3 1 0 0

carrier :: Double -> V3 Double
carrier t = zeroPoint + vel ^* t


{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE CPP #-}
{-# LANGUAGE ScopedTypeVariables #-}
#if __GLASGOW_HASKELL__ >= 800
  {-# OPTIONS_GHC -Wno-redundant-constraints #-}
#endif
-- |
-- Module      : Graphics.ColorSpace.Elevator
-- Copyright   : (c) Alexey Kuleshevich 2018-2019
-- License     : BSD3
-- Maintainer  : Alexey Kuleshevich <lehins@yandex.ru>
-- Stability   : experimental
-- Portability : non-portable
--
module Graphics.ColorSpace.Elevator (
  Elevator(..)
  , clamp01
  ) where

import qualified Data.Complex as C
import Data.Int
import Data.Typeable
import Data.Vector.Storable (Storable)
import Data.Vector.Unboxed (Unbox)
import Data.Word
import GHC.Float

-- | A class with a set of convenient functions that allow for changing precision of
-- channels within pixels, while scaling the values to keep them in an appropriate range.
--
-- >>> let rgb = PixelRGB 0.0 0.5 1.0 :: Pixel RGB Double
-- >>> eToWord8 <$> rgb
-- <RGB:(0|128|255)>
-- >>> eToWord16 <$> rgb
-- <RGB:(0|32768|65535)>
--
class (Eq e, Num e, Typeable e, Unbox e, Storable e) => Elevator e where

  -- | Values are scaled to @[0, 255]@ range.
  eToWord8 :: e -> Word8

  -- | Values are scaled to @[0, 65535]@ range.
  eToWord16 :: e -> Word16

  -- | Values are scaled to @[0, 4294967295]@ range.
  eToWord32 :: e -> Word32

  -- | Values are scaled to @[0, 18446744073709551615]@ range.
  eToWord64 :: e -> Word64

  -- | Values are scaled to @[0.0, 1.0]@ range.
  eToFloat :: e -> Float

  -- | Values are scaled to @[0.0, 1.0]@ range.
  eToDouble :: e -> Double

  -- | Values are scaled from @[0.0, 1.0]@ range.
  eFromDouble :: Double -> e

-- | Lower the precision
dropDown :: forall a b. (Integral a, Bounded a, Integral b, Bounded b) => a -> b
dropDown !e = fromIntegral $ fromIntegral e `div` ((maxBound :: a) `div`
                                                   fromIntegral (maxBound :: b))
{-# INLINE dropDown #-}

-- | Increase the precision
raiseUp :: forall a b. (Integral a, Bounded a, Integral b, Bounded b) => a -> b
raiseUp !e = fromIntegral e * ((maxBound :: b) `div` fromIntegral (maxBound :: a))
{-# INLINE raiseUp #-}

-- | Convert to fractional with value less than or equal to 1.
squashTo1 :: forall a b. (Fractional b, Integral a, Bounded a) => a -> b
squashTo1 !e = fromIntegral e / fromIntegral (maxBound :: a)
{-# INLINE squashTo1 #-}

-- | Convert to integral streaching it's value up to a maximum value.
stretch :: forall a b. (RealFrac a, Floating a, Integral b, Bounded b) => a -> b
stretch !e = round (fromIntegral (maxBound :: b) * clamp01 e)
{-# INLINE stretch #-}

-- | Clamp a value to @[0, 1]@ range.
clamp01 :: (Ord a, Floating a) => a -> a
clamp01 !x = min (max 0 x) 1
{-# INLINE clamp01 #-}

-- | Values between @[0, 255]]@
instance Elevator Word8 where
  eToWord8 = id
  {-# INLINE eToWord8 #-}
  eToWord16 = raiseUp
  {-# INLINE eToWord16 #-}
  eToWord32 = raiseUp
  {-# INLINE eToWord32 #-}
  eToWord64 = raiseUp
  {-# INLINE eToWord64 #-}
  eToFloat = squashTo1
  {-# INLINE eToFloat #-}
  eToDouble = squashTo1
  {-# INLINE eToDouble #-}
  eFromDouble = eToWord8
  {-# INLINE eFromDouble #-}

-- | Values between @[0, 65535]]@
instance Elevator Word16 where
  eToWord8 = dropDown
  {-# INLINE eToWord8 #-}
  eToWord16 = id
  {-# INLINE eToWord16 #-}
  eToWord32 = raiseUp
  {-# INLINE eToWord32 #-}
  eToWord64 = raiseUp
  {-# INLINE eToWord64 #-}
  eToFloat = squashTo1
  {-# INLINE eToFloat #-}
  eToDouble = squashTo1
  {-# INLINE eToDouble #-}
  eFromDouble = eToWord16
  {-# INLINE eFromDouble #-}

-- | Values between @[0, 4294967295]@
instance Elevator Word32 where
  eToWord8 = dropDown
  {-# INLINE eToWord8 #-}
  eToWord16 = dropDown
  {-# INLINE eToWord16 #-}
  eToWord32 = id
  {-# INLINE eToWord32 #-}
  eToWord64 = raiseUp
  {-# INLINE eToWord64 #-}
  eToFloat = squashTo1
  {-# INLINE eToFloat #-}
  eToDouble = squashTo1
  {-# INLINE eToDouble #-}
  eFromDouble = eToWord32
  {-# INLINE eFromDouble #-}

-- | Values between @[0, 18446744073709551615]@
instance Elevator Word64 where
  eToWord8 = dropDown
  {-# INLINE eToWord8 #-}
  eToWord16 = dropDown
  {-# INLINE eToWord16 #-}
  eToWord32 = dropDown
  {-# INLINE eToWord32 #-}
  eToWord64 = id
  {-# INLINE eToWord64 #-}
  eToFloat = squashTo1
  {-# INLINE eToFloat #-}
  eToDouble = squashTo1
  {-# INLINE eToDouble #-}
  eFromDouble = eToWord64
  {-# INLINE eFromDouble #-}

-- | Values between @[0, 18446744073709551615]@ on 64bit
instance Elevator Word where
  eToWord8 = dropDown
  {-# INLINE eToWord8 #-}
  eToWord16 = dropDown
  {-# INLINE eToWord16 #-}
  eToWord32 = dropDown
  {-# INLINE eToWord32 #-}
  eToWord64 = fromIntegral
  {-# INLINE eToWord64 #-}
  eToFloat = squashTo1
  {-# INLINE eToFloat #-}
  eToDouble = squashTo1
  {-# INLINE eToDouble #-}
  eFromDouble = stretch . clamp01
  {-# INLINE eFromDouble #-}

-- | Values between @[0, 127]@
instance Elevator Int8 where
  eToWord8 = fromIntegral . max 0
  {-# INLINE eToWord8 #-}
  eToWord16 = raiseUp . max 0
  {-# INLINE eToWord16 #-}
  eToWord32 = raiseUp . max 0
  {-# INLINE eToWord32 #-}
  eToWord64 = raiseUp . max 0
  {-# INLINE eToWord64 #-}
  eToFloat = squashTo1 . max 0
  {-# INLINE eToFloat #-}
  eToDouble = squashTo1 . max 0
  {-# INLINE eToDouble #-}
  eFromDouble = stretch . clamp01
  {-# INLINE eFromDouble #-}

-- | Values between @[0, 32767]@
instance Elevator Int16 where
  eToWord8 = dropDown . max 0
  {-# INLINE eToWord8 #-}
  eToWord16 = fromIntegral . max 0
  {-# INLINE eToWord16 #-}
  eToWord32 = raiseUp . max 0
  {-# INLINE eToWord32 #-}
  eToWord64 = raiseUp . max 0
  {-# INLINE eToWord64 #-}
  eToFloat = squashTo1 . max 0
  {-# INLINE eToFloat #-}
  eToDouble = squashTo1 . max 0
  {-# INLINE eToDouble #-}
  eFromDouble = stretch . clamp01
  {-# INLINE eFromDouble #-}

-- | Values between @[0, 2147483647]@
instance Elevator Int32 where
  eToWord8 = dropDown . max 0
  {-# INLINE eToWord8 #-}
  eToWord16 = dropDown . max 0
  {-# INLINE eToWord16 #-}
  eToWord32 = fromIntegral . max 0
  {-# INLINE eToWord32 #-}
  eToWord64 = raiseUp . max 0
  {-# INLINE eToWord64 #-}
  eToFloat = squashTo1 . max 0
  {-# INLINE eToFloat #-}
  eToDouble = squashTo1 . max 0
  {-# INLINE eToDouble #-}
  eFromDouble = stretch . clamp01
  {-# INLINE eFromDouble #-}

-- | Values between @[0, 9223372036854775807]@
instance Elevator Int64 where
  eToWord8 = dropDown . max 0
  {-# INLINE eToWord8 #-}
  eToWord16 = dropDown . max 0
  {-# INLINE eToWord16 #-}
  eToWord32 = dropDown . max 0
  {-# INLINE eToWord32 #-}
  eToWord64 = fromIntegral . max 0
  {-# INLINE eToWord64 #-}
  eToFloat = squashTo1 . max 0
  {-# INLINE eToFloat #-}
  eToDouble = squashTo1 . max 0
  {-# INLINE eToDouble #-}
  eFromDouble = stretch . clamp01
  {-# INLINE eFromDouble #-}

-- | Values between @[0, 9223372036854775807]@ on 64bit
instance Elevator Int where
  eToWord8 = dropDown . max 0
  {-# INLINE eToWord8 #-}
  eToWord16 = dropDown . max 0
  {-# INLINE eToWord16 #-}
  eToWord32 = dropDown . max 0
  {-# INLINE eToWord32 #-}
  eToWord64 = fromIntegral . max 0
  {-# INLINE eToWord64 #-}
  eToFloat = squashTo1 . max 0
  {-# INLINE eToFloat #-}
  eToDouble = squashTo1 . max 0
  {-# INLINE eToDouble #-}
  eFromDouble = stretch . clamp01
  {-# INLINE eFromDouble #-}

-- | Values between @[0.0, 1.0]@
instance Elevator Float where
  eToWord8 = stretch . clamp01
  {-# INLINE eToWord8 #-}
  eToWord16 = stretch . clamp01
  {-# INLINE eToWord16 #-}
  eToWord32 = stretch . clamp01
  {-# INLINE eToWord32 #-}
  eToWord64 = stretch . clamp01
  {-# INLINE eToWord64 #-}
  eToFloat = id
  {-# INLINE eToFloat #-}
  eToDouble = float2Double
  {-# INLINE eToDouble #-}
  eFromDouble = eToFloat
  {-# INLINE eFromDouble #-}

-- | Values between @[0.0, 1.0]@
instance Elevator Double where
  eToWord8 = stretch . clamp01
  {-# INLINE eToWord8 #-}
  eToWord16 = stretch . clamp01
  {-# INLINE eToWord16 #-}
  eToWord32 = stretch . clamp01
  {-# INLINE eToWord32 #-}
  eToWord64 = stretch . clamp01
  {-# INLINE eToWord64 #-}
  eToFloat = double2Float
  {-# INLINE eToFloat #-}
  eToDouble = id
  {-# INLINE eToDouble #-}
  eFromDouble = id
  {-# INLINE eFromDouble #-}

-- | Discards imaginary part and changes precision of real part.
instance (Num e, Elevator e, RealFloat e) => Elevator (C.Complex e) where
  eToWord8 = eToWord8 . C.realPart
  {-# INLINE eToWord8 #-}
  eToWord16 = eToWord16 . C.realPart
  {-# INLINE eToWord16 #-}
  eToWord32 = eToWord32 . C.realPart
  {-# INLINE eToWord32 #-}
  eToWord64 = eToWord64 . C.realPart
  {-# INLINE eToWord64 #-}
  eToFloat = eToFloat . C.realPart
  {-# INLINE eToFloat #-}
  eToDouble = eToDouble . C.realPart
  {-# INLINE eToDouble #-}
  eFromDouble = (C.:+ 0) . eFromDouble
  {-# INLINE eFromDouble #-}


{-# OPTIONS_GHC -fwarn-unused-imports #-}

module Main(main) where

import Control.Monad
import Control.Monad.Par
import Control.Monad.Par.Internal (runParAsync)
import Control.Monad.Par.Stream as S
import Control.Monad.Par.OpenList
import Control.DeepSeq
import Control.Exception

-- import Data.Array.Unboxed as U
import Data.Array.CArray as C
import Data.Complex
import Data.Int
import Data.Word
import Data.List (intersperse)
import Data.List.Split (chunk)

import System.CPUTime
import System.CPUTime.Rdtsc
import GHC.Conc as Conc
import GHC.IO (unsafePerformIO, unsafeDupablePerformIO, unsafeInterleaveIO)
import Debug.Trace

--------------------------------------------------------------------------------
-- Main script

main = do

  -- Generate 20 million elements:
  let s = countupWin 1024 (20 * 1000 * 1000) :: Par (WStream Int)
  measureRate $ runParAsync s

  putStrLn$ "Done with 5 million elements."


#!/usr/bin/env stack
-- stack runghc --package reanimate
{-# LANGUAGE OverloadedStrings #-}
module Main where

import           Codec.Picture.Types
import           Control.Exception
import           Control.Lens                    ()
import           Control.Monad
import           Data.Function
import           Data.List
import           Data.List.NonEmpty              (NonEmpty)
import qualified Data.List.NonEmpty              as NE
import           Data.Maybe
import           Data.Ratio
import qualified Data.Text                       as T
import           Data.Tuple
import qualified Data.Vector                     as V
import           Debug.Trace
import           Linear.Matrix                   hiding (trace)
import           Linear.Metric
import           Linear.V2
import           Linear.V3
import           Linear.Vector
import           Numeric.LinearAlgebra           hiding (polar, scale, (<>))
import qualified Numeric.LinearAlgebra           as Matrix
import           Numeric.LinearAlgebra.HMatrix   hiding (polar, scale, (<>))
import           Reanimate
import           Reanimate.Builtin.Documentation
import           Reanimate.Debug
import           Reanimate.Math.Balloon
import           Reanimate.Math.Common
import           Reanimate.Math.Compatible       (compatiblyTriangulateP)
import qualified Reanimate.Math.DCEL             as DCEL
import           Reanimate.Math.EarClip
import           Reanimate.Math.Polygon
import           Reanimate.Math.Render
import           Reanimate.Math.Smooth
import           Reanimate.Math.SSSP
import           Reanimate.Math.Visibility
import           Reanimate.Morph.Common
import           Reanimate.Morph.LeastDifference
import           Reanimate.Morph.LeastWork
import           Reanimate.Morph.Linear
import           Reanimate.Morph.Rigid
import           Reanimate.PolyShape             (svgToPolygons)
import           Text.Printf

-- p1 = centerPolygon $ shape2
--p1 = pScale 3.5 $ pAtCenter $ pAddPoints (0+2) (pSetOffset shape13 0)
-- p1 = setOffset (addPoints 2 shape13) 0
-- p2 = scalePolygon 0.5 $ centerPolygon shape20
-- p1 = centerPolygon shape2
--p2 = pScale 3.5 $ pAtCenter $ pAddPoints 0 (pSetOffset shape14 0)
-- p2 = setOffset shape14 0
p1 = pCopy $ pSetOffset (pCopy p1') 0
p2 = pCopy $ pSetOffset (pCopy p2') 0
(p1', p2') =
  -- normalizePolygons
  -- closestLinearCorrespondence
  (,)
  -- leastWork zeroStretchCosts defaultBendCosts
  -- leastWork defaultStretchCosts defaultBendCosts
  -- (pAtCenter $ unsafeSVGToPolygon 0.01 $ scale 6 $ latex "S")
  -- (pAtCenter $ unsafeSVGToPolygon 0.01 $ scale 6 $ latex "C")
  (pAtCenter $ unsafeSVGToPolygon 0.1 $ scale 6 $ latex "X")
  (pAtCenter $ unsafeSVGToPolygon 0.1 $ scale 6 $ latex "I")

p1_ = castPolygon p1
p2_ = castPolygon p2
polys = triangulate_ p1_ p2_
p1_circ = castPolygon (circumference (map fst polys') p1_)
p2_circ = castPolygon (circumference (map snd polys') p2_)

polys' = alignPolygons polys p1_ p2_
(p1s, p2s) = unzip $ compatTriagPairs [ (castPolygon a, castPolygon b) | (a,b) <- polys' ]

-- (p1s,p2s) = unzip (compatiblyTriangulateP p1 p2)

m2 = DCEL.buildMesh $ DCEL.polygonsMesh
        (map (fmap realToFrac) $ V.toList $ polygonPoints p2_circ)
        (map (map (fmap realToFrac) . V.toList . polygonPoints) p2s)

m1 = DCEL.buildMesh $ DCEL.polygonsMesh
        (map (fmap realToFrac) $ V.toList $ polygonPoints p1_circ)
        (map (map (fmap realToFrac) . V.toList . polygonPoints) p1s)

pipeline1 = last . take 20 . iterate
  ( uncurry DCEL.delaunayFlip .
    uncurry DCEL.splitInternalEdges .
   (\(a,b) -> (DCEL.meshSmoothPosition a, DCEL.meshSmoothPosition b)))
(m1final, m2final) = last $ take 30 $ iterate
  (pipeline1 .
    uncurry DCEL.splitLongestEdge .
    pipeline1
    ) (m1,m2)

-- (m1good, m2good) = last $ take 9 $ iterate
--   (pipeline1 .
--     -- uncurry DCEL.splitLongestEdge .
--     pipeline1
--     ) (m1,m2)

main :: IO ()
-- main = do
--   evaluate (last $ take 900 $ compatiblyTriangulateP p1 p2)
--   return ()
-- main = reanimate $ playTraces $ length $ take 29 p1s
-- main = reanimate $ playThenReverseA $ pauseAround 0.5 0.5 $ scene $ do

main = reanimate $ scene $ do
  newSpriteSVG_ $ mkBackgroundPixel rtfdBackgroundColor
  -- play
  --   $ playTraces
  --   $ traceA (mapA (withFillOpacity 0 . withStrokeColor "white") drawCircle)
  --   $ traceA (mapA (withFillOpacity 0 . withStrokeColor "white") drawBox)
  --   $ traceSVG ((withFillOpacity 0 . withStrokeColor "white") $ mkCircle 2)
  --   $ traceSVG ((withFillOpacity 0 . withStrokeColor "white") $ mkRect 2 2)
  --   $ 20
  -- play $ playTraces $ last $ take 35 $ compatiblyTriangulateP p1 p2
  -- fork $ newSpriteA $ drawCompatible p1 p2
  -- newSpriteSVG_ $ translate 5 0 $ lowerTransformations $ scale 2 $
  --   -- DCEL.renderMesh 0.02 m1final
  --   DCEL.renderMeshColored m1final
    -- renderMeshPair optMesh
  -- newSpriteSVG_ $ translate (-5) 0 $ lowerTransformations $ scale 2 $
  --   DCEL.renderMesh 0.02 m1good

  -- newSpriteSVG_ $
  --   translate (-3) 0 $ scale 1 $ mkGroup
  --   [ mkGroup []
  --   , withFillColor "grey" $ polygonShape p1
  --   , withFillColor "grey" $ polygonNumDots p1
  --   ]
  -- newSpriteSVG_ $
  --   translate 3 0 $ scale 1 $ mkGroup
  --   [ mkGroup []
  --   , withFillColor "grey" $ polygonShape p2
  --   , withFillColor "grey" $ polygonNumDots p2
  --   ]
  -- wait 1
  -- fork $ forM_ p1s $ \p1Piece -> do
  --   newSpriteSVG_ $
  --     translate (-3) 0 $ scale 1 $ mkGroup
  --     [ mkGroup []
  --     , withFillColor "white" $ polygonShape p1Piece
  --     -- , withFillColor "grey" $ polygonNumDots p1Piece
  --     ]
  --   wait (1/60)
  --   -- wait 1
  -- fork $ forM_ p2s $ \p1Piece -> do
  --   newSpriteSVG_ $
  --     translate 3 0 $ scale 1 $ mkGroup
  --     [ mkGroup []
  --     , withFillColor "white" $ polygonShape p1Piece
  --     -- , withFillColor "grey" $ polygonNumDots p1Piece
  --     ]
  --   wait (1/60)
  --   -- wait 1

  -- fork $ play $ mkAnimation 3 $ \t ->
  --   let points = interpolate prep t
  --   in  translate (-3) 0 $ scale 1 $ mkGroup
  --   [ mkGroup []
  --   , drawTrigs points (meshTriangles myMesh)
  --   -- , DCEL.renderMeshColored m1
  --   ]
  fork $ play $ mkAnimation 3 $ \t ->
    let points = interpolate optPrep t
    in  translate 3 0 $ lowerTransformations $ scale 1 $ mkGroup
      [ mkGroup []
      , drawTrigs points (meshTriangles optMesh)
      -- , DCEL.renderMeshColored m1final
      -- , withGroupOpacity 0.5 $ DCEL.renderMesh 0.02 m1final
      -- , withGroupOpacity 0.5 $ withFillOpacity 0 $ renderMeshPair optMesh
      -- , let V2 x y = V.head points in
      --   translate x y $ withFillColor "green" $
      --   mkCircle 0.03
      ]
        -- , scale 2 $ polygonNumDots p1
  -- play $ pauseAtEnd 1 $ mkAnimation 3 $ \t ->
  --   let points = interpolate bestPrep t
  --   in  translate 3 0 $ mkGroup [drawTrigs points (meshTriangles bestMesh)
  --       -- , scale 2 $ polygonNumDots p2
  --                                                                           ]
  -- play $ pauseAtEnd 1 $ mkAnimation 3 $ \t ->
  --   let points = interpolate prepRev t in
  --   mkGroup
  --   [ drawTrigs points (meshTriangles myMeshRev)
  --   -- , polygonNumDots shape2
  --   ]
  -- forM_ (myMesh : smoothMesh myMesh) $ \newMesh -> do
  --   let (minAngA, minAngB) = meshMinAngle newMesh
  --   txt <- newSpriteSVG $ withFillColor "white" $
  --     mkGroup
  --     [translate 5 2 $  center $
  --       latex $ T.pack $ printf "Min: %.1f" (minAngA/pi*180)
  --     ,translate 5 1 $  center $
  --       latex $ T.pack $ printf "Min: %.1f" (minAngB/pi*180) ]
  --   s <- newSpriteSVG $ renderAMesh newMesh
  --   wait (recip 60)
  --   destroySprite s
  --   destroySprite txt
  -- newSpriteSVG $ translate (-3) 0 $ scale 2 $ renderMesh p1 (map fst trigs)
  -- newSpriteSVG $ translate 3 0 $ scale 2 $ renderMesh p2 (map snd trigs)
  -- wait 1
 where
  prep      = prepare myMesh
  optPrep   = prepare optMesh
  -- bestPrep  = prepare bestMesh
  myMesh    = mkMesh p1 p2
  optMesh   = toRigidMesh  m1final m2final
  -- optMesh   = DCEL.toRigidMesh  m1good m2good
  -- bestMesh  = last $ smoothMesh myMesh
  prepRev   = prepare myMeshRev
  myMeshRev = reverseMesh myMesh

reverseMesh :: Mesh -> Mesh
reverseMesh mesh =
  mesh { meshPointsA = meshPointsB mesh, meshPointsB = meshPointsA mesh }

mkMesh :: Polygon -> Polygon -> Mesh
mkMesh a b = Mesh
  { meshPointsA   = V.map (fmap realToFrac) pointsA
  , meshPointsB   = V.map (fmap realToFrac) pointsB
  , meshOutline = V.map (fromJust . flip V.elemIndex pointsA) (polygonPoints a)
  , meshTriangles = relTrigs
  }
 where
  pointsA =
    V.fromList $ nub $ V.toList $ V.concat [ polygonPoints a | (a, b) <- trigs ]
  pointsB =
    V.fromList $ nub $ V.toList $ V.concat [ polygonPoints b | (a, b) <- trigs ]
  trigs = compatiblyTriangulateP a b
  mkRel arr p =
    ( fromJust $ V.elemIndex (pAccess p 0) arr
    , fromJust $ V.elemIndex (pAccess p 1) arr
    , fromJust $ V.elemIndex (pAccess p 2) arr
    )
  relTrigs = V.fromList [ mkRel pointsA a | (a, b) <- trigs ]
    -- data Mesh = Mesh (Vector P) (Vector (RelTrig, RelTrig))

testMesh2 :: Mesh
testMesh2 = mkMesh (pTranslate (V2 0 0) shape2) shape20

testPrep = prepare testMesh2

drawTrigs :: V.Vector (V2 Double) -> V.Vector RelTrig -> SVG
drawTrigs points trigs = mkGroup
  [ mkGroup
      [ withFillOpacity 1
        $ withStrokeWidth (defaultStrokeWidth * 0)
        $ withStrokeColor "grey"
        $ withFillColor "black"
        $ drawPolygon
        $ map (points V.!) [a, b, c]
      | (a, b, c) <- V.toList trigs
      ]
    -- , a == 0
  -- , mkGroup
  --   [ withFillColor "red" $ mkGroup
  --     [ drawPoint (points V.! a), drawPoint (points V.! b), drawPoint (points V.! c)]
  --   | (a,b,c) <- V.toList trigs
  --   -- , a == 0
  --   ]
  ]

drawTrigsLines :: V.Vector (V2 Double) -> V.Vector RelTrig -> SVG
drawTrigsLines points trigs = withFillOpacity 0 $ withStrokeColor "black" $ mkGroup
  [ mkLinePathClosed
    [ (aPx, aPy)
    , (bPx, bPy)
    , (cPx, cPy)]
  | (a,b,c) <- V.toList trigs
  -- , a==2 || b == 2 || c == 2
  , let V2 aPx aPy = points V.! a
        V2 bPx bPy = points V.! b
        V2 cPx cPy = points V.! c
  ]

drawPoint :: V2 Double -> SVG
drawPoint (V2 x y) = translate x y $ mkCircle 0.1

drawPolygon :: [V2 Double] -> SVG
drawPolygon lst = mkLinePathClosed [ (x, y) | V2 x y <- lst ]

drawCompatible :: Polygon -> Polygon -> Animation
drawCompatible a b = scene $ do
  let left  = -6
      right = 4
  newSpriteSVG $ translate left 0 $ mkGroup
    [ withFillColor "grey" $ polygonShape a
    , withFillColor "grey" $ polygonNumDots a
    ]
  newSpriteSVG $ translate right 0 $ mkGroup
    [ withFillColor "grey" $ polygonShape b
    , withFillColor "grey" $ polygonNumDots b
    ]
  return ()
  -- let compat = compatiblyTriangulateP a b
  -- forM_ compat $ \(l, r) -> do
  --   fork $ play $ staticFrame 1 $
  --     translate left 0 $ withStrokeColor "white" $ withStrokeWidth (defaultStrokeWidth*0.2) $
  --     withFillOpacity 0 $ polygonShape l
  --   fork $ play $ staticFrame 1 $
  --     translate right 0 $ withStrokeColor "white" $ withStrokeWidth (defaultStrokeWidth*0.2) $
  --     withFillOpacity 0 $ polygonShape r


{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE ViewPatterns     #-}
{-# LANGUAGE TypeFamilies     #-}
{-# LANGUAGE ScopedTypeVariables #-}

-- |
-- Module      : Main
-- Description : Test runner
-- Copyright   : (c) Tom Westerhout, 2017
-- License     : BSD3
-- Maintainer  : t.westerhout@student.ru.nl
-- Stability   : experimental

module Main where

import           Prelude hiding (map, zipWithM)

import           Debug.Trace
import qualified System.Random.MWC as MWC
-- import           Control.Lens hiding((<.>))
import           Control.Monad.Reader hiding (zipWithM)
import           Control.Monad.Primitive
import qualified Data.List            as L
import           Data.Complex
import           Data.Semigroup
import           Data.Vector.Storable (Vector, (!))
import qualified Data.Vector.Storable as V
import           System.Exit
import           System.IO hiding (hGetLine)
import           System.Environment(getArgs, getProgName)
import           Foreign.Storable
import           Data.Text (Text)
import           Data.Text.IO (hGetLine)
import qualified Data.Text.IO as T

import           Lens.Micro
import           Lens.Micro.Extras

import           PSO.Random
import           PSO.Swarm
-- import           PSO.Heisenberg
-- import           PSO.Energy
-- import           PSO.Neural
import           PSO.FromPython
import           NQS.Rbm

meanVariance :: (RealFloat a, Storable a) => V.Vector a -> (a, a)
meanVariance xs
  | n > 1     = (m, sumVar m xs / fromIntegral (n - 1))
  | otherwise = (m, 0)
    where n = V.length xs
          m = V.sum xs / fromIntegral (V.length xs)
          sumVar m' = V.sum . V.map ((^2) . subtract m')

type R = Float
type C = Complex Float
type X = Rbm C

fromPyFile :: FilePath -> IO X
fromPyFile name = withFile name ReadMode toRbm
  where toRight :: Either String (V.Vector (Complex Float)) -> V.Vector (Complex Float)
        toRight (Right x) = x
        toRight (Left x)  = error x
        toRbm h = do
          a <- trace ("a...") $ toRight <$> readVector <$> hGetLine h
          b <- trace ("b...") $ toRight <$> readVector <$> hGetLine h
          s <- hGetLine h
          T.putStrLn s
          let !w = trace ("w...") $ toRight $ readMatrix s
          return $ mkRbm (a, b, w)

randomRbm ::
  ( UniformDist m R
  )
  => Int
  -> Int
  -> (R, R)
  -> (R, R)
  -> (R, R)
  -> m X
randomRbm n m (lowV, highV) (lowH, highH) (lowW, highW) =
  do
    visible <- uniformVector n (lowV :+ lowV, highV :+ highV)
    hidden <- uniformVector m (lowH :+ lowH, highH :+ highH)
    weights <- uniformVector (n * m) (lowW :+ lowW, highW :+ highW)
    return $ trace ("mkRbm...") (mkRbm (visible, hidden, weights))

rbms :: (UniformDist m R) => m [X]
rbms = trace ("Creating rbms...") $ replicateM 20 (randomRbm 10 20 (-0.1, 0.1) (-0.1, 0.1) (-0.1, 0.1))

rbms' :: IO [X]
rbms' = Prelude.mapM fromPyFile $
  [ "test.txt"
  -- "/home/tom/src/tcm-swarm/10_spins/Weights_10_spins_2_density_-1.23981661958_energy"
  -- "/home/tom/src/tcm-swarm/10_spins/Weights_10_spins_2_density_-1.89541865304_energy"
  -- , "/home/tom/src/tcm-swarm/10_spins/Weights_10_spins_2_density_0.0887681979723_energy"
  -- , "/home/tom/src/tcm-swarm/10_spins/Weights_10_spins_2_density_1.40938434022_energy"
  -- , "/home/tom/src/tcm-swarm/10_spins/Weights_10_spins_2_density_-3.61205064542_energy"
  -- , "/home/tom/src/tcm-swarm/10_spins/Weights_10_spins_2_density_-4.1917579599_energy"
  -- , "/home/tom/src/tcm-swarm/10_spins/Weights_10_spins_2_density_-5.41817224673_energy"
  -- , "/home/tom/src/tcm-swarm/10_spins/Weights_10_spins_2_density_-5.74064288804_energy"
  -- , "/home/tom/src/tcm-swarm/10_spins/Weights_10_spins_2_density_-6.33997058402_energy"
  -- , "/home/tom/src/tcm-swarm/10_spins/Weights_10_spins_2_density_-6.48266135073_energy"
  -- , "/home/tom/src/tcm-swarm/10_spins/Weights_10_spins_2_density_-7.29704001394_energy"
  -- , "/home/tom/src/tcm-swarm/10_spins/Weights_10_spins_2_density_-7.4024108663_energy"
  -- , "/home/tom/src/tcm-swarm/10_spins/Weights_10_spins_2_density_-7.83407453451_energy"
  -- , "/home/tom/src/tcm-swarm/10_spins/Weights_10_spins_2_density_-8.17005224883_energy"
  -- , "/home/tom/src/tcm-swarm/10_spins/Weights_10_spins_2_density_-8.18318151693_energy"
  -- , "/home/tom/src/tcm-swarm/10_spins/Weights_10_spins_2_density_-8.75410430005_energy"
  -- , "/home/tom/src/tcm-swarm/10_spins/Weights_10_spins_2_density_-9.04188789775_energy"
  -- , "/home/tom/src/tcm-swarm/10_spins/Weights_10_spins_2_density_-9.34302369964_energy"
  -- , "/home/tom/src/tcm-swarm/10_spins/Weights_10_spins_2_density_-9.5917604763_energy"
  -- , "/home/tom/src/tcm-swarm/10_spins/Weights_10_spins_2_density_2.01304877555_energy"
  ]

{-
  [ "/home/tom/src/tcm-swarm/8_spins/input_singleflip_0.txt"
  , "/home/tom/src/tcm-swarm/8_spins/input_singleflip_1.txt"
  , "/home/tom/src/tcm-swarm/8_spins/input_singleflip_2.txt"
  , "/home/tom/src/tcm-swarm/8_spins/input_singleflip_3.txt"
  , "/home/tom/src/tcm-swarm/8_spins/input_singleflip_4.txt"
  , "/home/tom/src/tcm-swarm/8_spins/input_singleflip_5.txt"
  , "/home/tom/src/tcm-swarm/8_spins/input_singleflip_6.txt"
  , "/home/tom/src/tcm-swarm/8_spins/input_singleflip_7.txt"
  , "/home/tom/src/tcm-swarm/8_spins/input_singleflip_8.txt"
  , "/home/tom/src/tcm-swarm/8_spins/input_singleflip_9.txt"
  , "/home/tom/src/tcm-swarm/8_spins/input_singleflip_10.txt"
  , "/home/tom/src/tcm-swarm/8_spins/input_singleflip_11.txt"
  , "/home/tom/src/tcm-swarm/8_spins/input_singleflip_12.txt"
  , "/home/tom/src/tcm-swarm/8_spins/input_singleflip_13.txt"
  , "/home/tom/src/tcm-swarm/8_spins/input_singleflip_14.txt"
  , "/home/tom/src/tcm-swarm/8_spins/input_singleflip_15.txt"
  , "/home/tom/src/tcm-swarm/8_spins/input_singleflip_16.txt"
  , "/home/tom/src/tcm-swarm/8_spins/input_singleflip_17.txt"
  , "/home/tom/src/tcm-swarm/8_spins/input_singleflip_18.txt"
  , "/home/tom/src/tcm-swarm/8_spins/input_singleflip_19.txt"
  ]
-}

function :: QMState X -> ReaderT g IO (Vector C)
function x = lift $
  sampleMoments (x ^. pos) Heisenberg1DPeriodic 4 (1000, 7000 + 1000, 7) (Just 0) 4

instance {-# OVERLAPS #-} Ord (Vector C) where
  a <= b = let k x = realPart (x ! 3) / (realPart (x ! 1))^2
            in k a <= k b

process :: Swarm m (SwarmGuide X (Vector C)) (BeeGuide X (Vector C)) (QMState X) (Vector C)
        -> IO ()
process swarm = do
  putStrLn $ "At iteration #" <> show (swarm ^. guide . iteration)
  withFile "Energies.dat" AppendMode $ \h ->
    do
      let xs = view bees swarm
      hPutStrLn h (L.concat $ L.intersperse "\t" $ show . realPart . (! 0) . view val <$> xs)
  withFile "Variances.dat" AppendMode $ \h ->
    do
      let xs = view bees swarm
      hPutStrLn h (L.concat $ L.intersperse "\t" $ show . realPart . (! 1) . view val <$> xs)
  withFile "Kurtosis.dat" AppendMode $ \h ->
    do
      let xs = view bees swarm
          k x = realPart (x ! 3) / (realPart (x ! 1))^2
      hPutStrLn h (L.concat $ L.intersperse "\t" $ show . k . view val <$> xs)

instance Scalable Float (Rbm C) where
  scale λ x = map (* (λ :+ 0)) x

instance DeltaWell m R R => DeltaWell m R X where
  upDeltaWell κ p x = zipWithM (upDeltaWell κ) p x

update ::
  ( m ~ ReaderT g IO
  , Randomisable m Float
  )
  => PhaseUpdater m (SwarmGuide X r) (BeeGuide X r) (QMState X) r
update = PhaseUpdater $ deltaUpdater (1.9 * log 2 :: R)

runHeisenbergFromList ::
  ( m ~ ReaderT g IO
  -- , r ~ MeanVar Float
  -- , χ ~ Rbm (Complex Float)
  -- , DeltaWell m Float χ
  -- , Scalable Float χ
  , UniformDist m Int
  , Randomisable m Float
  , Randomisable m Bool
  )
  => g -> IO ()
runHeisenbergFromList gen = runReaderT go gen
  where go =
          do
            -- states <- fmap QMState <$> rbms
            states <- fmap QMState <$> lift rbms'
            optimiseNDFromList
                  states
                  update
                  function
                  (\s -> (s^.guide.iteration == 100))
                  (lift . process)
            return ()

{-
runHeisenberg ::
  ( m ~ ReaderT g IO
  -- , r ~ MeanVar Float
  -- , χ ~ Rbm (Complex Float)
  -- , DeltaWell m Float χ
  -- , Scalable Float χ
  , UniformDist m Int
  , Randomisable m Float
  , Randomisable m Bool
  )
  => Int -> Int -> g -> IO ()
runHeisenberg spins count gen = do
  let h  = heisenberg1DOpen spins
      initBounds = ((-2.0E-1) :+ (-2.0E-1), 2.0E-1 :+ 2.0E-1)
      n = spins
      newState = QMState <$> uniformRbm n (2 * spins) initBounds
      func x = uncurry MV <$> energyHH1DOpenMKLC (x ^. pos) 100 10000
      xs = optimiseND
        newState
        update
        func
        count
        (\s -> (s^.guide.iteration == 100))
  swarms <- runReaderT xs gen
  writeEnergies2TSV "Energies.dat"  (view (val . mean)) swarms
  writeEnergies2TSV "Variances.dat" (view (val . var)) swarms
  -- writeVariances2TSV "Variances.dat" swarms
  let swarm = last swarms
      (m, v) = meanVariance
                . V.fromList
                . map (view (val . mean))
                $ swarm ^. bees
      min = minimum $ view mean <$> view val <$> swarm ^. bees
      max = maximum $ view mean <$> view val <$> swarm ^. bees
  -- mapM_ (print . (!!0) . (view bees)) $ swarms
  putStrLn ""
  putStrLn $ "[+] After " ++ show (swarm ^. guide . iteration)
    ++ " iterations: "
  putStrLn $ "[+] E[<H>]         = " ++ show m
  putStrLn $ "[+] Min[<H>]       = " ++ show min
  putStrLn $ "[+] Max[<H>]       = " ++ show max
  putStrLn $ "[+] StdDev[<H>]    = " ++ show (sqrt v)
  putStrLn $ "[+] Best[<H>]      = " ++ show (swarm ^. guide . val . mean)
  putStrLn $ "[+] Best[Var[<H>]] = " ++ show (swarm ^. guide . val . var)
  putStrLn $ "[+] Actual states:   " ++ show (fst . LA.eigSH $ h)
-}

main :: IO ()
main = do
  args <- getArgs
  case args of
    [ reads -> [(spins :: Int, _)], reads -> [(count :: Int, _)], reads -> [(seed, _)]
      ] -> do g <- mkMWCGen (Just seed)
              runHeisenbergFromList g
    _ -> do name <- getProgName
            hPutStrLn stderr $ "usage: " ++ name ++ " <#spins> <#bees> <seed>"
            exitFailure


module Strategy.TestOrientation where

import Prelude hiding (Right, Left)
import Numeric.LinearAlgebra ((<>))
import qualified Control.Monad.Writer as Writer
import qualified Data.DList as DList

import Cube (Side(..), Color(..))
import qualified Cube as Cube
import qualified Rotation as Rotation
import qualified RotationPath as RotationPath
import qualified Strategy.Orientation as Orientation

testSides :: Bool
testSides =
  let expected = [
          (Top, White)
        , (Front, Green)
        , (Right, Red)
        , (Back, Blue)
        , (Left, Orange)
        , (Bottom, Yellow)
        ]
      actual = Orientation.sides Cube.solvedCube
  in expected == actual

testIsColorOnSide :: Bool
testIsColorOnSide =
  let t = Orientation.isColorOnSide White Top Cube.solvedCube
      f = Orientation.isColorOnSide Blue Bottom Cube.solvedCube
  in t == True && f == False

testForColorOnTop :: Bool
testForColorOnTop = and $ do
  color <- [Cube.White ..]
  let (prefix, suffix) = Orientation.forColorOnTop color Cube.solvedCube
      cube = Rotation.rotate prefix Cube.solvedCube
      isOnTop = Orientation.isColorOnSide color Cube.Top cube
      isIdentity = (suffix <> prefix) == Rotation.identity
  return $ isIdentity && isOnTop

testPutColorOnTop :: Bool
testPutColorOnTop =
  let expectedSides = [
          (Top, Yellow)
        , (Front, Blue)
        , (Right, Red)
        , (Back, Green)
        , (Left, Orange)
        , (Bottom, White)
        ]
      expectedRotations = [Rotation.topToFront <> Rotation.topToFront]
      (cube, rotationLog) = Writer.runWriter $
        Orientation.putColorOnTop Yellow Cube.solvedCube
      actualSides = Orientation.sides cube
  in actualSides == expectedSides && expectedRotations == DList.toList rotationLog

testWithColorOnTop :: Bool
testWithColorOnTop =
  let (cube, rotationLog) = Writer.runWriter $
        Orientation.withColorOnTop Green (RotationPath.rotate Rotation.topL) Cube.solvedCube
      [top, frn, rgt, bck, lft, bot] = [Cube.Top ..]
      [w, g, r, b, o, y] = [Cube.White ..]
      expectedCube = [
          (top,0,0,w),(top,0,1,w),(top,0,2,w),(top,1,0,w),(top,1,1,w),(top,1,2,w),(top,2,0,r),(top,2,1,r),(top,2,2,r)
        , (frn,0,0,g),(frn,0,1,g),(frn,0,2,g),(frn,1,0,g),(frn,1,1,g),(frn,1,2,g),(frn,2,0,g),(frn,2,1,g),(frn,2,2,g)
        , (rgt,0,0,y),(rgt,0,1,r),(rgt,0,2,r),(rgt,1,0,y),(rgt,1,1,r),(rgt,1,2,r),(rgt,2,0,y),(rgt,2,1,r),(rgt,2,2,r)
        , (bck,0,0,b),(bck,0,1,b),(bck,0,2,b),(bck,1,0,b),(bck,1,1,b),(bck,1,2,b),(bck,2,0,b),(bck,2,1,b),(bck,2,2,b)
        , (lft,0,0,o),(lft,0,1,o),(lft,0,2,w),(lft,1,0,o),(lft,1,1,o),(lft,1,2,w),(lft,2,0,o),(lft,2,1,o),(lft,2,2,w)
        , (bot,0,0,o),(bot,0,1,o),(bot,0,2,o),(bot,1,0,y),(bot,1,1,y),(bot,1,2,y),(bot,2,0,y),(bot,2,1,y),(bot,2,2,y)]
      expectedRotations = [Rotation.frontL]
  in cube == expectedCube && DList.toList rotationLog == expectedRotations

tests :: [(String, Bool)]
tests = [
    ("Strategy.Orientation.sides", testSides)
  , ("Strategy.Orientation.isColorOnSide", testIsColorOnSide)
  , ("Strategy.Orientation.forColorOnTop", testForColorOnTop)
  , ("Strategy.Orientation.putColorOnTop", testPutColorOnTop)
  , ("Strategy.Orientation.withColorOnTop", testWithColorOnTop)
  ]


{-# LANGUAGE BangPatterns #-}
module FokkerPlanck.MonteCarlo
  ( -- runMonteCarloFourierCoefficients
    runMonteCarloFourierCoefficientsGPU
  , solveMonteCarloR2S1
  ) where

import           Array.UnboxedArray              as UA
import           Control.DeepSeq
import           Control.Monad                   as M
import           Control.Concurrent.Async
import           Control.Monad.Parallel          as MP
import           Data.Array.Accelerate           (constant, use)
import           Data.Array.Accelerate.LLVM.PTX
import           Data.Array.Repa                 as R
import           Data.Binary
import           Data.ByteString.Lazy            as BL
import           Data.Complex
import           Data.DList                      as DL
import           Data.Ix
import           Data.List                       as L
import           Data.Vector.Unboxed             as VU
import           FokkerPlanck.BrownianMotion
import           FokkerPlanck.FourierSeries
import           FokkerPlanck.FourierSeriesGPU
import           FokkerPlanck.Histogram
import           Foreign.CUDA.Driver             as CUDA
import           Statistics.Distribution
import           Statistics.Distribution.Laplace
import           Statistics.Distribution.Normal
import           Statistics.Distribution.Poisson
import           Statistics.Distribution.Uniform
import           System.Directory
import           System.IO                       as IO
import           System.Random.MWC
import           Text.Printf
import           Utils.List
import           Utils.Parallel
import           Utils.Time
import GHC.Float

{-# INLINE runBatch #-}
runBatch :: Int -> Int -> Int -> (Int -> a -> [GenIO] -> IO a) -> a -> IO a
runBatch !numGens !numTrails !batchSize monterCarloHistFunc !init = do
  let (!numBatch, !numLeftover) = divMod numTrails batchSize
  printf
    "%d trails and batch size is %d, %d batch in total.\n"
    numTrails
    batchSize
    (numBatch +
     if numLeftover > 0
       then 1
       else 0)
  gensList <- M.replicateM numBatch (M.replicateM numGens createSystemRandom)
  x <-
    M.foldM
      (\y gens -> do
         printCurrentTime ""
         monterCarloHistFunc batchSize y gens)
      init
      gensList
  if numLeftover > 0
    then do
      gens <- M.replicateM numGens createSystemRandom
      monterCarloHistFunc numLeftover x gens
    else return x

{-# INLINE computeHistogramFromMonteCarloParallel #-}
computeHistogramFromMonteCarloParallel ::
     (NFData hist, Binary hist, Show hist)
  => FilePath
  -> (GenIO -> IO particle)
  -> ([particle] -> IO hist)
  -> (hist -> hist -> hist)
  -> Int
  -> hist
  -> [GenIO]
  -> IO hist
computeHistogramFromMonteCarloParallel !filePath pointsGenerator histFunc addHist !n !initHist !gens = do
  xs <- MP.mapM (M.replicateM (div n . L.length $ gens) . pointsGenerator) gens
  let tmpFilePath = (filePath L.++ "_tmp")
  hist <- L.foldl' addHist initHist <$> M.mapM histFunc xs
  unless
    (L.null filePath)
    (do encodeFile tmpFilePath hist
        copyFile tmpFilePath filePath)
  return hist

{-# INLINE computeHistogramFromMonteCarloParallelSingleGPU #-}
computeHistogramFromMonteCarloParallelSingleGPU ::
     (NFData hist, Binary hist, Show hist)
  => FilePath
  -> (GenIO -> IO particle)
  -> ([particle] -> hist)
  -> (hist -> hist -> hist)
  -> Int
  -> Int
  -> hist
  -> [GenIO]
  -> IO hist
computeHistogramFromMonteCarloParallelSingleGPU !filePath pointsGenerator histFunc addHist !deviceID !n !initHist !gens = do
  xs <- MP.mapM (M.replicateM (div n . L.length $ gens) . pointsGenerator) gens
  let tmpFilePath = (filePath L.++ "_tmp")
      !hist = addHist initHist . histFunc . L.concat $ xs
  unless
    (L.null filePath)
    (do encodeFile tmpFilePath hist
        copyFile tmpFilePath filePath)
  return hist

{-# INLINE computeHistogramFromMonteCarloParallelMultipleGPU #-}
computeHistogramFromMonteCarloParallelMultipleGPU ::
     (NFData hist, Binary hist, Show hist, NFData particle)
  => FilePath
  -> (GenIO -> IO particle)
  -> (PTX -> [particle] -> hist)
  -> (hist -> hist -> hist)
  -> [PTX]
  -> Int
  -> hist
  -> [GenIO]
  -> IO hist
computeHistogramFromMonteCarloParallelMultipleGPU !filePath pointsGenerator histFunc addHist !ptxs !n !initHist !gens = do
  xs <-
    L.concat <$>
    mapConcurrently
      (M.replicateM (div n . L.length $ gens) . pointsGenerator)
      gens
  let tmpFilePath = filePath L.++ "_tmp"
      !hist =
        L.foldl' addHist initHist .
        parZipWith rdeepseq histFunc ptxs . divideListN (L.length ptxs) $
        xs
  unless
    (L.null filePath)
    (do encodeFile tmpFilePath hist
        copyFile tmpFilePath filePath)
  return hist

-- {-# INLINE runMonteCarloFourierCoefficients #-}
-- runMonteCarloFourierCoefficients ::
--      Int
--   -> Int
--   -> Int
--   -> Double
--   -> Double
--   -> Double
--   -> Double
--   -> [Double]
--   -> [Double]
--   -> [Double]
--   -> [Double]
--   -> Double
--   -> FilePath
--   -> Histogram (Complex Double)
--   -> IO (Histogram (Complex Double))
-- runMonteCarloFourierCoefficients !numGens !numTrails !batchSize !thetaSigma !scaleSigma !maxScale !tao !phiFreqs !rhoFreqs !thetaFreqs !rFreqs !deltaLog !filePath !initHist = do
--   when
--     (maxScale <= 1)
--     (error $
--      printf "runMonteCarloFourierCoefficients error: maxScale(%f) <= 1" maxScale)
--   let !thetaDist = normalDistrE 0 thetaSigma
--       !scaleDist = normalDistrE 0 scaleSigma
--       pointsGenerator = generatePath thetaDist scaleDist maxScale tao 1
--       histFunc =
--         computeFourierCoefficients
--           phiFreqs
--           rhoFreqs
--           thetaFreqs
--           rFreqs
--           (log maxScale)
--           deltaLog
--       monterCarloHistFunc =
--         computeHistogramFromMonteCarloParallel
--           filePath
--           pointsGenerator
--           histFunc
--           addHistogramUnsafe
--   hist <- runBatch numGens numTrails batchSize monterCarloHistFunc initHist
--   return hist

{-# INLINE runMonteCarloFourierCoefficientsSingleGPU #-}
runMonteCarloFourierCoefficientsSingleGPU ::
     Int
  -> Int
  -> Int
  -> Int
  -> Double
  -> Double
  -> Double
  -> Double
  -> [Double]
  -> [Double]
  -> [Double]
  -> [Double]
  -> Double -> Double
  -> FilePath
  -> Histogram (Complex Double)
  -> IO (Histogram (Complex Double))
runMonteCarloFourierCoefficientsSingleGPU !deviceID !numGens !numTrails !batchSize !thetaSigma !scaleSigma !maxScale !tao !phiFreqs !rhoFreqs !thetaFreqs !rFreqs !deltaLog !initScale !filePath !initHist = do
  when
    (maxScale <= 1)
    (error $
     printf "runMonteCarloFourierCoefficients error: maxScale(%f) <= 1" maxScale)
  gen <- createSystemRandom
  initialise []
  dev <- device deviceID
  ctx <- CUDA.create dev []
  ptx <- createTargetFromContext ctx
  let !thetaDist = normalDistrE 0 thetaSigma
      !scaleDist = normalDistrE 0 scaleSigma
      -- !scaleDist = uniformDistrE (- (log maxScale)) (log maxScale)
      !poissonDist = poissonE 0
      !freqArr =
        computeFrequencyArray
          (L.map double2Float phiFreqs)
          (L.map double2Float rhoFreqs)
          (L.map double2Float thetaFreqs)
          (L.map double2Float rFreqs)
      pointsGenerator =
        generatePath
          thetaDist
          scaleDist
          poissonDist
          (maxScale ^ 2)
          maxScale
          tao
          initScale
          1
      histFuncSingleGPU =
        computeFourierCoefficientsGPU
          phiFreqs
          rhoFreqs
          thetaFreqs
          rFreqs
          (use freqArr)
          (constant . double2Float $ maxScale)
          (constant . double2Float . log $ maxScale)
          ptx
      monterCarloHistFuncSingleGPU =
        computeHistogramFromMonteCarloParallelSingleGPU
          filePath
          pointsGenerator
          histFuncSingleGPU
          addHistogramUnsafe
          deviceID
  hist <-
    runBatch numGens numTrails batchSize monterCarloHistFuncSingleGPU initHist
  -- deepseq hist (destroy ctx)
  return hist

-- {-# INLINE runMonteCarloFourierCoefficientsMultipleGPU #-}
-- runMonteCarloFourierCoefficientsMultipleGPU ::
--      [Int]
--   -> Int
--   -> Int
--   -> Int
--   -> Double
--   -> Double
--   -> Double
--   -> Double
--   -> [Double]
--   -> [Double]
--   -> [Double]
--   -> [Double]
--   -> Double
--   -> Double
--   -> FilePath
--   -> Histogram (Complex Double)
--   -> IO (Histogram (Complex Double))
-- runMonteCarloFourierCoefficientsMultipleGPU !deviceIDs !numGens !numTrails !batchSize !thetaSigma !scaleSigma !maxScale !tao !phiFreqs !rhoFreqs !thetaFreqs !rFreqs !deltaLog !initScale !filePath !initHist = do
--   when
--     (maxScale <= 1)
--     (error $
--      printf "runMonteCarloFourierCoefficients error: maxScale(%f) <= 1" maxScale)
--   initialise []
--   devs <- M.mapM device deviceIDs
--   ctxs <- M.mapM (\dev -> CUDA.create dev []) devs
--   ptxs <- M.mapM createTargetFromContext ctxs
--   let !thetaDist = normalDistrE 0 thetaSigma
--       !scaleDist = normalDistrE 0 scaleSigma
--       !freqArr = computeFrequencyArray phiFreqs rhoFreqs thetaFreqs rFreqs
--       pointsGenerator = generatePath thetaDist scaleDist maxScale tao initScale
--       histFuncMultipleGPU =
--         computeFourierCoefficientsGPU
--           phiFreqs
--           rhoFreqs
--           thetaFreqs
--           rFreqs
--           (log maxScale)
--           deltaLog
--           (use freqArr)
--           (constant maxScale)
--           (constant (log maxScale))
--           (constant (deltaLog :+ 0))
--       monterCarloHistFuncMultipleGPU =
--         computeHistogramFromMonteCarloParallelMultipleGPU
--           filePath
--           pointsGenerator
--           histFuncMultipleGPU
--           addHistogramUnsafe
--           ptxs
--   hist <-
--     runBatch numGens numTrails batchSize monterCarloHistFuncMultipleGPU initHist
--   deepseq hist (M.mapM destroy ctxs)
--   return hist

{-# INLINE runMonteCarloFourierCoefficientsGPU #-}
runMonteCarloFourierCoefficientsGPU ::
     [Int]
  -> Int
  -> Int
  -> Int
  -> Double
  -> Double
  -> Double
  -> Double
  -> Double
  -> Double
  -> [Double]
  -> [Double]
  -> [Double]
  -> [Double]
  -> Double
  -> Double
  -> FilePath
  -> IO (Histogram (Complex Double))
runMonteCarloFourierCoefficientsGPU !deviceIDs !numGens !numTrails !batchSize !thetaSigma !scaleLambda !poissonLambda !sigma !tao !deltaT !phiFreqs !rhoFreqs !thetaFreqs !rFreqs !periodEnv !stdR2 !filePath = do
  when
    (periodEnv <= 1)
    (error $
     printf
       "runMonteCarloFourierCoefficients error: periodEnv(%f) <= 1"
       periodEnv)
  initialise []
  devs <- M.mapM device deviceIDs
  ctxs <- M.mapM (\dev -> CUDA.create dev []) devs
  ptxs <- M.mapM createTargetFromContext ctxs
  let !initHist =
        emptyHistogram
          [ L.length phiFreqs
          , L.length rhoFreqs
          , L.length thetaFreqs
          , L.length rFreqs
          ]
          0
      !thetaDist = normalDistrE 0 (thetaSigma * sqrt deltaT)
      !scaleDist = normalDistrE 0 (scaleLambda * sqrt deltaT)
      !poissonDist = poissonE (poissonLambda / deltaT)
      !freqArr =
        computeFrequencyArray
          (L.map double2Float phiFreqs)
          (L.map double2Float rhoFreqs)
          (L.map double2Float thetaFreqs)
          (L.map double2Float rFreqs)
      !maxRho =  sqrt periodEnv / 1 / sqrt 2
      !maxR = sqrt periodEnv / 1 / sqrt 2
      pointsGenerator =
        generatePath
          thetaDist
          scaleDist
          poissonDist
          maxRho
          maxR
          (tao / deltaT)
          deltaT
          stdR2
      -- pointsGenerator =
      --   generatePath'
      --     thetaSigma
      --     scaleDist
      --     poissonLambda
      --     maxRho
      --     maxR
      --     tao
      --     deltaT
      --     stdR2
      -- pointsGenerator =
      --   generatePath' thetaSigma scaleDist poissonLambda periodEnv (tao / deltaT) deltaT
      histFuncSingleGPU =
        computeFourierCoefficientsGPU
          phiFreqs
          rhoFreqs
          thetaFreqs
          rFreqs
          (use freqArr)
          (constant . double2Float $ sigma)
          (constant . double2Float $ log periodEnv)
      monterCarloHistFuncGPU =
        computeHistogramFromMonteCarloParallelMultipleGPU
          filePath
          pointsGenerator
          histFuncSingleGPU
          addHistogramUnsafe
          ptxs
  runBatch numGens numTrails batchSize monterCarloHistFuncGPU initHist

{-# INLINE countR2S1 #-}
countR2S1 ::
     GenIO
  -> Double
  -> (Int, Int)
  -> (Int, Int)
  -> Int
  -> [DList Particle]
  -> IO (Histogram Double)
countR2S1 randomGen thetaSigma xRange@(!xMin, !xMax) yRange@(!yMin, !yMax) !numOrientations !xs = do
  let !deltaTheta = 2 * pi / (fromIntegral numOrientations)
  ys <-
    fmap
      (L.filter
         (\((_, x, y), _) ->
            (inRange xRange x) && (inRange yRange y) && (x /= 0 || y /= 0)) .
       L.map
         (\(Particle phi rho theta' _ _) ->
            let !x = rho * cos phi
                !y = rho * sin phi
                !theta =
                  if theta' < 0
                    then theta' + 2 * pi
                    else theta'
             in ((floor $ theta / deltaTheta, round x, round y), 1)) .
       L.concat) .
    M.mapM
      (\particle@(Particle phi rho theta r _) -> do
         let delta = 1
             n = Prelude.floor $ r / delta
         zs <-
           M.mapM
             (\i -> do
                let thetaDist =
                      normalDistr
                        0
                        (thetaSigma * sqrt (delta * fromIntegral i / r))
                    (Particle a b c d _) =
                      FokkerPlanck.BrownianMotion.moveParticle
                        1
                        (Particle
                           phi
                           rho
                           theta
                           (Prelude.fromIntegral i * delta)
                           1)
                deltaThetaDiffusion <- genContVar thetaDist randomGen
                return (Particle a b (c `thetaPlus` deltaThetaDiffusion) d 1))
             [1 .. n - 1]
         return .
           L.concatMap
             (\(Particle phi' rho' theta' r' v) ->
                [ (Particle phi' rho' theta' r' v)
                , (Particle (-phi') rho' (-theta') r' v)
                ]) $
           (FokkerPlanck.BrownianMotion.moveParticle
              1
              (Particle phi rho theta r 1)) :
           zs) .
    DL.toList . DL.concat $
    xs
  return .
    Histogram
      [(yMax - yMin + 1), (xMax - xMin + 1), numOrientations]
      (L.length ys) .
    toUnboxedVector .
    UA.accum (+) 0 ((0, xMin, yMin), (numOrientations - 1, xMax, yMax)) $
    ys

{-# INLINE solveMonteCarloR2S1 #-}
solveMonteCarloR2S1 ::
     Int
  -> Int
  -> Int
  -> Int
  -> Int
  -> Int
  -> Double
  -> Double
  -> Double
  -> Double
  -> FilePath
  -> IO (R.Array U DIM3 Double)
solveMonteCarloR2S1 numGens numTrails batchSize xLen yLen numOrientations thetaSigma tao r initSpeed histFilePath = do
  gen <- createSystemRandom
  let !xShift = div xLen 2
      xRange' =
        if odd xLen
          then (-xShift, xShift)
          else (-xShift, xShift - 1)
      !yShift = div yLen 2
      yRange' =
        if odd yLen
          then (-yShift, yShift)
          else (-yShift, yShift - 1)
      thetaDist = normalDistrE 0 thetaSigma
      scaleDist = normalDistrE 0 0
      poissonDist = poissonE 0
      pointsGenerator =
        generatePath
          thetaDist
          scaleDist
          poissonDist
          (r^2)
          r -- (sqrt . fromIntegral $ xLen ^ 2 + yLen ^ 2)
          tao
          initSpeed
          1
      histFunc = countR2S1 gen thetaSigma xRange' yRange' numOrientations
      monterCarloHistFunc =
        computeHistogramFromMonteCarloParallel
          histFilePath
          pointsGenerator
          histFunc
          addHistogramUnsafe
  hist <-
    runBatch
      numGens
      numTrails
      batchSize
      monterCarloHistFunc
      (emptyHistogram [yLen, xLen, numOrientations] 0)
  return . getNormalizedHistogramArr $ hist


module Language.Definition where

import qualified Data.Complex as C

-- Data Type
-- @todo: the number tower needs to be implemented
--
-- @fixme: It doesn't seem like the DottedList is implemented
-- correctly. I copied it from the tutorial
data LispVal = Atom String
             | List [LispVal]
             | DottedList [LispVal] LispVal
             | Char Char
             | String String
             | Bool Bool
             | Number Integer
             | Float Double
             | Ratio Rational
             | Complex (C.Complex Double)
             deriving Show


module Sound.Plot where

import qualified Data.Stream as S

import           Numeric.LinearAlgebra

import           Sound.Types

import           Graphics.Rendering.Plot (Series)
import qualified Graphics.Rendering.Plot as P
import           Graphics.Rendering.Plot.Gtk (PlotHandle)
import qualified Graphics.Rendering.Plot.Gtk as P

plot :: Int -> [Audio] -> IO PlotHandle
plot n audio = P.display $ do
  P.plot (P.Line, map (fromList . S.take n) audio :: [Series])
  P.withTextDefaults $ P.setFontFamily "Ubuntu"
  P.xlabel "samples"
  P.ylabel "amplitude"
  P.yrange P.Linear (-1) 1


import Paraiso
import Data.Complex

main = do
  args <- getArgs
  let arch = if "--cuda" `elem` args then
               CUDA 128 128
             else
               X86
  putStrLn $ compile arch code
    where
      code = do
         parallel 16384 $ do
           r <- allocate
           x <- allocate
           r =$ Rand (0.0::Double) 4.0
           x =$ Rand 0 1
           cuda $ do
             sequential 65536 $ do
               r =$ r * x * (1-x)
           output [r,x]


{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE TupleSections #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE TypeSynonymInstances #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE ConstraintKinds #-}
{-# LANGUAGE CPP #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE DefaultSignatures #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE AllowAmbiguousTypes #-}
{-# LANGUAGE UndecidableInstances #-}

{-# OPTIONS_GHC -Wall #-}

#include "ConCat/Ops.inc"

-- | Commutative monoid intended to be used with a multiplicative monoid

module ConCat.Additive where

import Prelude hiding (zipWith)
import Data.Monoid (Monoid(..), Sum(..), Product(..))
import Data.Semigroup (Semigroup(..))
import Data.Complex hiding (magnitude)
import Data.Ratio
import Foreign.C.Types (CSChar, CInt, CShort, CLong, CLLong, CIntMax, CFloat, CDouble)
import GHC.Generics (U1(..),Par1(..),(:*:)(..),(:.:)(..))
import GHC.TypeLits (KnownNat)

import Data.Key(Zip(..))
import Data.Pointed
import Data.Functor.Rep (Representable(..))
import Data.Vector.Sized (Vector)
import Data.Finite.Internal

import ConCat.Misc
import ConCat.Rep (HasRep(abst),inAbst,inAbst2)
import qualified ConCat.Rep
import ConCat.Orphans ()

-- | Commutative monoid intended to be used with a multiplicative monoid
class Additive a where
  zero  :: a
  infixl 6 ^+^
  (^+^) :: a -> a -> a
  default zero :: (Pointed h, Additive b, a ~ h b) => a
  zero = pointNI zero
  default (^+^) :: (Zip h, Additive b, a ~ h b) => Binop a
  (^+^) = zipWithNI (^+^)
  {-# INLINE zero #-}
  {-# INLINE (^+^) #-}

-- These definitions and the corresponding Catify rewrites in AltCat prevent the point and zipWith methods from getting inlined too soon.
-- See 2018-04-09 notes.
pointNI :: Pointed h => a -> h a
pointNI = point
{-# INLINE [0] pointNI #-}

zipWithNI :: Zip h => (a -> b -> c) -> h a -> h b -> h c
zipWithNI = zipWith
{-# INLINE [0] zipWithNI #-}

instance Additive () where
  zero = ()
  () ^+^ () = ()

#define ScalarType(t) \
  instance Additive (t) where { zero = 0 ; (^+^) = (+) }

ScalarType(Int)
ScalarType(Integer)
ScalarType(Float)
ScalarType(Double)
ScalarType(CSChar)
ScalarType(CInt)
ScalarType(CShort)
ScalarType(CLong)
ScalarType(CLLong)
ScalarType(CIntMax)
ScalarType(CDouble)
ScalarType(CFloat)

instance Integral a => Additive (Ratio a) where
  zero  = 0
  (^+^) = (+)

instance (RealFloat v, Additive v) => Additive (Complex v) where
  zero  = zero :+ zero
  (^+^) = (+)

-- The 'RealFloat' constraint is unfortunate here. It's due to a
-- questionable decision to place 'RealFloat' into the definition of the
-- 'Complex' /type/, rather than in functions and instances as needed.

instance (Additive u,Additive v) => Additive (u,v) where
  zero              = (zero,zero)
  (u,v) ^+^ (u',v') = (u^+^u',v^+^v')

instance (Additive u,Additive v,Additive w)
    => Additive (u,v,w) where
  zero                   = (zero,zero,zero)
  (u,v,w) ^+^ (u',v',w') = (u^+^u',v^+^v',w^+^w')

instance (Additive u,Additive v,Additive w,Additive x)
    => Additive (u,v,w,x) where
  zero                        = (zero,zero,zero,zero)
  (u,v,w,x) ^+^ (u',v',w',x') = (u^+^u',v^+^v',w^+^w',x^+^x')

type AddF f = (Pointed f, Zip f)

instance KnownNat n => Additive (Finite n) where
  zero = 0
  (^+^) = (+)

#if 1

#define AdditiveFunctor(f) instance (AddF (f), Additive v) => Additive ((f) v)

AdditiveFunctor((->) a)
AdditiveFunctor(Sum)
AdditiveFunctor(Product)
AdditiveFunctor(U1)
AdditiveFunctor(Par1)
AdditiveFunctor(f :*: g)
AdditiveFunctor(g :.: f)

#else

instance Additive v => Additive (a -> v)
instance Additive v => Additive (Sum     v)
instance Additive v => Additive (Product v)

instance Additive v => Additive (U1 v)
instance Additive v => Additive (Par1 v)
instance (Additive v, AddF f, AddF g) => Additive ((f :*: g) v)
instance (Additive v, AddF f, AddF g) => Additive ((g :.: f) v)

#endif

-- instance (Eq i, Additive v) => Additive (Arr i v) where
--   zero = point zero
--   as ^+^ bs = fmap (uncurry (^+^)) (zipC (as,bs))

-- TODO: Define and use zipWithC (^+^) as bs.

instance (Additive v, KnownNat n) => Additive (Vector n v)

-- Maybe is handled like the Maybe-of-Sum monoid
instance Additive a => Additive (Maybe a) where
  zero                = Nothing
  Nothing ^+^ b'      = b'
  a' ^+^ Nothing      = a'
  Just a' ^+^ Just b' = Just (a' ^+^ b')

-- -- Memo tries
-- instance (HasTrie u, Additive v) => Additive (u :->: v) where
--   zero  = pure   zero
--   (^+^) = liftA2 (^+^)

-- Experiment
instance Additive Bool where
  -- zero = undefined
  -- _ ^+^ _ = undefined
  zero = False
  (^+^) = (||)
  {-# INLINE zero #-}
  {-# INLINE (^+^) #-}

{--------------------------------------------------------------------
    Monoid wrapper
--------------------------------------------------------------------}

-- | Monoid under group addition.  Alternative to the @Sum@ in
-- "Data.Monoid", which uses 'Num' instead of 'Additive'.
newtype Add a = Add { getAdd :: a }
  deriving (Eq, Ord, Read, Show, Bounded)

instance HasRep (Add a) where
  type Rep (Add a) = a
  abst = Add
  repr = getAdd

instance Functor Add where fmap = inAbst

instance Applicative Add where
  pure  = abst
  (<*>) = inAbst2 ($)

instance Additive a => Semigroup (Add a) where
  (<>) = inAbst2 (^+^)

instance Additive a => Monoid (Add a) where
  mempty  = abst zero
  mappend = (<>)

instance Additive a => Additive (Add a) where
  zero  = mempty
  (^+^) = mappend

-- sumA' :: (Foldable h, Additive a) => h a -> a
-- sumA' = getAdd . foldMap Add

-- Enables translation of sumA to jamPF in AltCat.
type SummableF h = (Representable h, Eq (Rep h), Zip h, Pointed h, Foldable h)

class    SummableF h => Summable h
instance SummableF h => Summable h

-- The constraint ‘Representable h’
--   is no smaller than the instance head
-- (Use UndecidableInstances to permit this)

sumA :: (
  -- Summable h, Additive a
  Foldable h, Additive a
  ) => h a -> a
sumA = getAdd . foldMap Add
{-# OPINLINE sumA #-}