text:
{-# LANGUAGE DeriveGeneric #-}

{-|
Module      : HABQTlib.Data

This module contains data types and helper functions for working with quantum
state vectors and density matrices.
-}
module HABQTlib.Data
  ( Dim
  , Rank
  , NumberOfParticles
  , QBitNum
  , Weight
  , MHMCiter
  , OptIter
  , OutputVerb(..)
  , DensityMatrix(..)
  , truncateRank
  , PureStateVector(..)
  , pureStateLikelihood
  , svToDM
  , WeighedDensityMatrix(..)
  , mkWDM
  , mkWDM1
  , WeighedPureStateVector(..)
  , (<+>)
  , fidelity
  , fidelityDM
  , validSV
  , validDM
  , validPartNum
  , validRank
  , validMHMCiter
  , validOptIter
  , validQBitN
  ) where

import Control.Newtype.Generics (Newtype, unpack)
import Data.Bool.HT (select)
import Data.Complex
import Data.Validation
import GHC.Generics (Generic)
import qualified Numeric.LinearAlgebra as LA
import Numeric.LinearAlgebra (Matrix)

-- | Dimension of Hilbert space.
type Dim = Int

-- | Rank of mixed state.
type Rank = Int

-- | Number of particles per rank.
type NumberOfParticles = Int

-- | Number of quantum bits.
type QBitNum = Int

-- | Number of MHMC iterations to perform when resampling.
type MHMCiter = Int

-- | Number of optimisation steps to perform when searching for optimal
-- measurment.
type OptIter = Int

-- | Weight associated with a particle.
type Weight = Double

-- | Output verbosity settings.
data OutputVerb
  = NoOutput -- ^ No stdout output
  | FidOutput -- ^ Only output fidelities and weights of hierarchical mean estimates
  | FullOutput -- ^ Full output, including resampling diagnostic information
  deriving (Eq, Show, Ord)

-- | Density matrix are stored as hmatrix matrices of complex doubles.
newtype DensityMatrix = DensityMatrix
  { getDensityMatrix :: Matrix (Complex Double)
  } deriving (Eq, Show, Generic)

-- | Pure state vectors are stored as hmatrix matrices of complex doubles.
-- Such matrices only have one column.
newtype PureStateVector = PureStateVector
  { getStateVector :: Matrix (Complex Double)
  } deriving (Eq, Show, Generic)

instance Newtype DensityMatrix

instance Newtype PureStateVector

-- | Check whether a pure state vector is properly normed.
validSV :: PureStateVector -> Validation [String] PureStateVector
validSV =
  validate
    ["State vector must have unit norm."]
    (\x -> abs (1 - LA.norm_2 (unpack x)) < 1e-12)

-- | Verify that density matrix is Hermitian and has trace 1.
validDM :: DensityMatrix -> Validation [String] DensityMatrix
validDM =
  let traceU :: LA.Matrix (Complex Double) -> Bool
      traceU dm =
        (abs (1 - (magnitude . LA.sumElements . LA.takeDiag) dm) < 1e-12)
      hermU dm = (LA.norm_2 (dm - LA.tr dm) < 1e-6)
      both = ((&&) <$> traceU <*> hermU) . unpack
      dmM = ["Density matrix must be Hermitian and have trace of 1."]
   in validate dmM both

qnM :: [String]
qnM = ["Number of quantum bits must be a positive integer."]

-- | Verify that number of quantum bits is positive.
validQBitN :: QBitNum -> Validation [String] QBitNum
validQBitN = validate qnM (> 0)

miM :: [String]
miM = ["Number of MHMC iterations must be a positive integer."]

-- | Verify that number of MHMC iterations is a positive integer.
validMHMCiter :: MHMCiter -> Validation [String] MHMCiter
validMHMCiter = validate miM (> 0)

pnM :: [String]
pnM = ["Number of particles per rank must be a positive integer."]

-- | Verify that particle number is a positive integer.
validPartNum :: NumberOfParticles -> Validation [String] NumberOfParticles
validPartNum = validate pnM (> 0)

rM :: [String]
rM = ["Rank must be a positive integer."]

-- | Verify that rank is a positive integer. Setting rank to be higher than
-- the dimension of space creates poinless performance overhead, but isn't
-- prevented by validation.
validRank :: Rank -> Validation [String] Rank
validRank = validate rM (> 0)

oiM :: [String]
oiM = ["Number of POVM optimisation iterations must be a positive integer."]

-- | Verify that number of POVM optimisation iterations is positive.
validOptIter :: OptIter -> Validation [String] OptIter
validOptIter = validate oiM (> 0)

-- | Weighed density matrix where weight is stored separately as first
-- coordinate of a tuple.
newtype WeighedDensityMatrix = WeighedDensityMatrix
  { getWDM :: (Weight, DensityMatrix)
  } deriving (Eq, Show, Generic)

-- | A shorter alias for curried 'WeighedDensityMatrix' constructor.
mkWDM :: Weight -> DensityMatrix -> WeighedDensityMatrix
mkWDM w dm = WeighedDensityMatrix (w, dm)

-- | Alias for @'mkWDM' 1@.
mkWDM1 :: DensityMatrix -> WeighedDensityMatrix
mkWDM1 = mkWDM 1

-- | Weighed state vector where weight is stored separately as first coordinate
-- of a tuple.
newtype WeighedPureStateVector = WeighedPureStateVector
  { getWSV :: (Weight, PureStateVector)
  } deriving (Eq, Show, Generic)

instance Newtype WeighedDensityMatrix

instance Newtype WeighedPureStateVector

-- | Fidelity (probability of measurement) between pure states.
fidelity :: PureStateVector -> PureStateVector -> Double
fidelity (PureStateVector sv1) (PureStateVector sv2) =
  let ips = magnitude (LA.atIndex (LA.tr sv1 LA.<> sv2) (0, 0)) ^ (2 :: Int)
   in select ips [(ips < 0, 0), (ips > 1, 1)]

-- | Fidelity (probability of measurement) between mixed states.
fidelityDM :: DensityMatrix -> DensityMatrix -> Double
fidelityDM (DensityMatrix dm1) (DensityMatrix dm2) =
  let (u1, s1) = LA.leftSV dm1
      (u2, s2) = LA.leftSV dm2
      ss1 = LA.cmap (\x -> sqrt x :+ 0) s1
      ss2 = LA.cmap (\x -> sqrt x :+ 0) s2
   in LA.norm_nuclear
        (u1 LA.<> LA.diag ss1 LA.<> LA.tr u1 LA.<> u2 LA.<> LA.diag ss2 LA.<>
         LA.tr u2) ^
      (2 :: Int)

-- | Calculate the density matrix of a given pure state.
svToDM :: PureStateVector -> DensityMatrix
svToDM (PureStateVector sv) = DensityMatrix $ sv LA.<> LA.tr sv

infix 8 <+>

-- | Given two weighed density matrixes, compute their mixture. Associative
-- operation.
(<+>) :: WeighedDensityMatrix -> WeighedDensityMatrix -> WeighedDensityMatrix
wdm0 <+> wdm1 =
  WeighedDensityMatrix
    (w0 + w1, DensityMatrix $ LA.scale c0 dm0 + LA.scale c1 dm1)
  where
    up = fmap unpack . unpack
    (w0, dm0) = up wdm0
    (w1, dm1) = up wdm1
    cs = LA.fromList [w0 :+ 0, w1 :+ 0]
    csn = LA.scale (1 / LA.sumElements cs) cs
    c0 = csn LA.! 0
    c1 = csn LA.! 1

-- | Probability of obtaining a measurement result when projecting a system in
-- mixed state determined by a density matrix onto a pure state.
pureStateLikelihood :: PureStateVector -> DensityMatrix -> Double
pureStateLikelihood (PureStateVector sv) (DensityMatrix dm) =
  let p =
        LA.magnitude . LA.sumElements . LA.takeDiag $ LA.tr sv LA.<> dm LA.<> sv
   in select p [(p < 0, 0), (p > 1, 1)]

-- | Set smallest eigenvalues of a weighed density matrix to zero until
-- specified rank is reached.
truncateRank :: Rank -> WeighedDensityMatrix -> WeighedDensityMatrix
truncateRank targetRank (WeighedDensityMatrix (w, DensityMatrix dm)) =
  let (u, s, _) = LA.svd dm
      st = LA.real $ LA.subVector 0 targetRank s
      ut = u LA.¿ [0 .. (targetRank - 1)]
      stn = LA.scale (1 / LA.sumElements st) st
   in WeighedDensityMatrix
        (w, DensityMatrix $ ut LA.<> LA.diag stn LA.<> LA.tr ut)

max_stars_repo_name: Belinsky-L-V/HABQT

max_issues_repo_name: Belinsky-L-V/HABQT

max_forks_repo_name: Belinsky-L-V/HABQT


text:
module Main where

import Softbody
import qualified Data.Array.Accelerate as A
import Data.Array.Accelerate.LLVM.Native as CPU
import Data.Complex
import Graphics.Gloss
import Graphics.Gloss.Data.ViewPort
import Graphics.Gloss.Data.Color

hs_new_v = CPU.runN new_v

update_v :: C -> V -> Double -> Double -> V
update_v c v t d = hs_new_v c v (A.fromList A.Z $ [t :+ 0]) (A.fromList A.Z $ [d :+ 0])

---

display_v :: (V, C) -> Picture
display_v (v, c) = pictures [nodePictures, edgePictures]
  where
    complexToFloatTuple (real :+ imag) = (realToFrac real, realToFrac imag)
    nodeList = A.toList v

    rectangleAt x y = translate x y $ rectangleSolid 10 10
    nodePictures = pictures (map (\(pi, delta, alpha, rho) -> uncurry rectangleAt $ complexToFloatTuple pi) nodeList)

    nodePairsWithDistance = zip [(pi1, pi2) | (pi1, _, _, _) <- nodeList, (pi2, _, _, _) <- nodeList] (A.toList c)
    edgeColor n = mixColors 10 n black red
    edgePicture ((pi1, pi2), d) = color (edgeColor . realToFrac . abs $ (magnitude $ pi2 - pi1) - d) (line $ complexToFloatTuple <$> [pi1, pi2])
    edgePictures = pictures (edgePicture <$> nodePairsWithDistance)

tick_v :: ViewPort -> Float -> (V, C) -> (V, C)
tick_v _ t (v, c) = (update_v c v (realToFrac t * 8) 1, c)

main :: IO ()
main = simulate
  (InWindow "https://aearnus.github.io/" (800, 600) (0, 0))
  (greyN 0.5)
  120
  (my_v, my_c)
  display_v
  tick_v

ext: hs

max_stars_repo_name: Aearnus/softbody-accelerate

max_issues_repo_name: Aearnus/softbody-accelerate

max_forks_repo_name: Aearnus/softbody-accelerate


text:
-----------------------------------------------------------------------------
-- |
-- Module      : Berp.Base.Truth
-- Copyright   : (c) 2010 Bernie Pope
-- License     : BSD-style
-- Maintainer  : florbitous@gmail.com
-- Stability   : experimental
-- Portability : ghc
--
-- Implementation of the truth predicate on Python objects.
--
-----------------------------------------------------------------------------

module Berp.Base.Truth (truth) where

import {-# SOURCE #-} Berp.Base.HashTable as HashTable (sizeIO)
import {-# SOURCE #-} Berp.Base.HashSet as HashSet (sizeIO)
import Berp.Base.SemanticTypes (Object (..))
import Data.Complex (Complex (..))
import System.IO.Unsafe (unsafePerformIO)
import Data.IORef (readIORef)

-- XXX incomplete
truth :: Object -> Bool
truth (Bool { object_bool = b }) = b
truth (Integer { object_integer = i }) = i /= 0
truth (Float { object_float = f }) = f /= 0
truth (Complex { object_complex = c }) = c /= (0 :+ 0)
truth None = False
truth obj@(Tuple {}) = object_length obj /= 0
truth obj@(String {}) = object_string obj /= []
truth obj@(List {}) = unsafePerformIO $ do
   numElems <- readIORef $ object_list_num_elements obj
   return (numElems /= 0)
truth obj@(Dictionary {}) = unsafePerformIO $ do
   numElems <- HashTable.sizeIO $ object_hashTable obj
   return (numElems /= 0)
truth obj@(Set {}) = unsafePerformIO $ do
   numElems <- HashSet.sizeIO $ object_hashSet obj
   return (numElems /= 0)
truth _other = True

{-
   Python Language Reference, section 5.10 "Boolean operations":

   In the context of Boolean operations, and also when expressions
   are used by control flow statements, the following values are
   interpreted as false: False, None, numeric zero of all types,
   and empty strings and containers (including strings, tuples,
   lists, dictionaries, sets and frozensets). All other values
   are interpreted as true. User-defined objects can customize their
   truth value by providing a __bool__() method.
-}

max_stars_repo_name: ppelleti/berp

max_issues_repo_name: ppelleti/berp

max_forks_repo_name: bjpop/berp


text:
#!/usr/bin/env stack
-- stack runghc --package reanimate
{-# LANGUAGE OverloadedStrings #-}
module Main (main) where

import           Data.Complex
import           Graphics.SvgTree
import           Linear.V2
import           Reanimate
import           Reanimate.Ease
import           Codec.Picture

-- layer 3
main :: IO ()
main = reanimate $ setDuration 30 $ sceneAnimation $ do
    _ <- newSpriteSVG $ mkBackgroundPixel (PixelRGBA8 252 252 252 0xFF)
    play $ fourierA (fromToS 0 5)      -- Rotate 15 times
      # setDuration 50
      # signalA (reverseS . powerS 2 . reverseS) -- Start fast, end slow
      # pauseAtEnd 2
    play $ fourierA (constantS 0)       -- Don't rotate at all
      # setDuration 10
      # reverseA
      # signalA (powerS 2)                       -- Start slow, end fast
      # pauseAtEnd 2

-- layer 2
fourierA :: (Double -> Double) -> Animation
fourierA genPhi = animate $ \t ->
    let circles = setFourierLength (t*piFourierLen) piFourier
        coeffs = fourierCoefficients $ rotateFourier (genPhi t) circles
    in mkGroup
    [ drawCircles coeffs
    , withStrokeColor "green" $
      withStrokeLineJoin JoinRound $
      withFillOpacity 0 $
      withStrokeWidth (defaultStrokeWidth*2) $
      mkLinePath $ mkFourierOutline circles
    , let x :+ y = sum coeffs in
      translate x y $ withFillColor "red" $ mkCircle (defaultStrokeWidth*3)
    ]

drawCircles :: [Complex Double] -> SVG
drawCircles [] = mkGroup []
drawCircles ( x :+ y : xs) =
  translate x y $ drawCircles' xs

drawCircles' :: [Complex Double] -> SVG
drawCircles' circles = mkGroup
    [ worker circles
    , withStrokeColor "black" $
      withStrokeLineJoin JoinRound $
      withFillOpacity 0 $
      mkLinePath [ (x, y) | x :+ y <- scanl (+) 0 circles ]]
  where
    worker [] = None
    worker (x :+ y : rest) =
      let radius = sqrt(x*x+y*y) in
      mkGroup
      [ withStrokeColor "dimgrey" $
        withFillOpacity 0 $
        mkCircle radius
      , translate x y $ worker rest ]

-- layer 1
newtype Fourier = Fourier {fourierCoefficients :: [Complex Double]}

piFourier :: Fourier
piFourier = mkFourier $ lineToPoints 500 $
  toLineCommands $ extractPath $ scale 15 $
  center $ latexAlign "\\pi"

piFourierLen :: Double
piFourierLen = sum $ map magnitude $ drop 1 $ take 500 $ fourierCoefficients piFourier

pointAtFourier :: Fourier -> Complex Double
pointAtFourier = sum . fourierCoefficients

mkFourier :: [RPoint] -> Fourier
mkFourier points = Fourier $ findCoefficient 0 :
    concat [ [findCoefficient n, findCoefficient (-n)] | n <- [1..] ]
  where
    findCoefficient :: Int -> Complex Double
    findCoefficient n =
        sum [ toComplex point * exp (negate (fromIntegral n) * 2 *pi * i*t) * deltaT
            | (idx, point) <- zip [0::Int ..] points, let t = fromIntegral idx/nPoints ]
    i = 0 :+ 1
    toComplex (V2 x y) = x :+ y
    deltaT = recip nPoints
    nPoints = fromIntegral (length points)

setFourierLength :: Double -> Fourier -> Fourier
setFourierLength _ (Fourier []) = Fourier []
setFourierLength len0 (Fourier (first:lst)) = Fourier $ first : worker len0 lst
  where
    worker _len [] = []
    worker len (c:cs) =
      if magnitude c < len
        then c : worker (len - magnitude c) cs
        else [c * realToFrac (len / magnitude c)]

rotateFourier :: Double -> Fourier -> Fourier
rotateFourier phi (Fourier coeffs) =
    Fourier $ worker coeffs (0::Integer)
  where
    worker [] _ = []
    worker (x:rest) 0 = x : worker rest 1
    worker [left] n = worker [left,0] n
    worker (left:right:rest) n =
      let n' = fromIntegral n in
      left * exp (negate n' * 2 * pi * i * phi') :
      right * exp (n' * 2 * pi * i * phi') :
      worker rest (n+1)
    i = 0 :+ 1
    phi' = realToFrac phi

mkFourierOutline :: Fourier -> [(Double, Double)]
mkFourierOutline fourier =
    [ (x, y)
    | idx <- [0 .. granularity]
    , let x :+ y = pointAtFourier $ rotateFourier (idx/granularity) fourier
    ]
  where
    granularity = 500

max_stars_repo_name: TristanCacqueray/reanimate

max_issues_repo_name: TristanCacqueray/reanimate

max_forks_repo_name: TristanCacqueray/reanimate


text:
{-# LANGUAGE RecordWildCards #-}
{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE TypeSynonymInstances #-}
{-# LANGUAGE FlexibleInstances #-}
module Statistics.BBVI.Propagator
  ( DistCell(..)
  , DistCells
  , DistCellss
  , SampleVector
  , SampleDouble
  , Gradient
  , Memory
  , Time
  , mergeGeneric
  , mergeGenerics
  , mergeGenericss
  , dist
  , time
  , defaultDistCell
  )
where

import           Statistics.BBVI.Class
import qualified Data.Vector                   as V
import           Data.Propagator

-- | a sample with a vector type e.g. sample from a dirichlet
type SampleVector = V.Vector Double

-- | a sample with a real type e.g. number sampled from a gaussian
type SampleDouble = Double

type Gradient = V.Vector Double

type Memory = V.Vector Double

-- | number of times a distribution cell has been updated
type Time = Int

-- | distribution cell
data DistCell a
  =
    U !Memory !Gradient -- ^ update to a distribution cell
  | Node !Time !Memory !a -- ^ distribution cell with a distribution
                          -- of type a
  deriving (Show, Eq, Ord, Read)

-- !(Gradient -> DistCell a -> (Memory, Gradient))

-- | helper function for initializing distribution cells
defaultDistCell
  :: DistUtil a
  => a -- ^ initial point for distribution
  -> DistCell a -- ^ distribution cell
defaultDistCell d = Node 1 (V.replicate (nParams d) 0) d

-- | time accessor for a distribution cell
time :: DistCell a -> Time
time (Node t _ _) = t
time (U{}       ) = error "called time on update propnode!"
-- memory (Node _ m _) = m

-- | distribution accessor for a distribution cell
dist :: DistCell a -> a
dist (Node _ _ d) = d
dist (U{}       ) = error "called dist on update propnode!"

-- | single cell representing a vector of distribution cells
type DistCells a = V.Vector (DistCell a)

-- | single cell representing an array (vector of vector) of distribution cells
type DistCellss a = V.Vector (V.Vector (DistCell a))

-- | generic merge for DistCells: useful for customizing
-- "quiesence" thresholds with cellWith
mergeGeneric
  :: DistUtil a
  => Time -- ^ maximum number of updates to a cell
  -> Double -- ^ threshold for information to be considered "new"; if
            -- the l2-norm of the change in the distribution's
            -- parameters is less than this threshold, the cell
            -- remains unchanged
  -> DistCell a -- ^ current cell
  -> DistCell a -- ^ proposed update to cell
  -> Change (DistCell a)
mergeGeneric maxStep delta !x1 !x2 = m x1 x2
 where
  m no@(Node t memory d) (U memUp gradUp)
    | norm gradUp < delta = Change False no
    | t >= maxStep        = Change False no
    | otherwise           = Change True updateNode
   where
    updateNode = Node (t + 1)
                      (V.zipWith (+) memory memUp)
                      (fromParamVector newQ)
      where newQ = V.zipWith (+) (toParamVector d) gradUp
  m (U _ _) _ = Contradiction mempty "Trying to update a gradient"
  -- CAREFUL: below is dangerous if I start doing the ideas i thought
  -- about: changing maxstep and elta node for local optmizations
  m no1@(Node t1 _m1 d1) no2@(Node t2 _m2 d2)
    | t1 >= maxStep
    = Change False no1
    | (t2 > t1)
      && (norm (V.zipWith (-) (toParamVector d1) (toParamVector d2)) >= delta)
    = Change True no2
    | otherwise
    = Change False no1

-- | generalization of 'mergeGeneric' to cells of vectors of distributions
mergeGenerics
  :: DistUtil a
  => Time -- ^ maximum number of updates to the cells
  -> Double -- ^ threshold for information to be considered "new"; if
            -- the l2-norm of change in the distribution's parameters
            -- is less than this threshold for all distributions, the
            -- cell remains unchanged
  -> DistCells a -- ^ current cells
  -> DistCells a -- ^ proposed update to cells
  -> Change (DistCells a)
mergeGenerics m d x1 x2 = V.sequence . V.zipWith (mergeGeneric m d) x1 $ x2

-- | generalization of 'mergeGeneric' to cells of arrays of distributions
mergeGenericss
  :: DistUtil a
  => Time -- ^ maximum number of updates to the cells
  -> Double -- ^ threshold for information to be considered "new"; if
            -- the l2-norm of change in the distribution's parameters
            -- is less than this threshold for all distributions, the
            -- cell remains unchanged
  -> DistCellss a -- ^ current cell
  -> DistCellss a -- ^ proposed update to cell
  -> Change (DistCellss a)
mergeGenericss m d v1 v2 = V.sequence . V.zipWith (mergeGenerics m d) v1 $ v2

instance DistUtil a => Propagated (DistCell a) where
  merge = mergeGeneric 1000000 1e-16

instance DistUtil a => Propagated (DistCells a) where
  merge ns updates = V.sequence $ V.zipWith merge ns updates

instance DistUtil a => Propagated (DistCellss a) where
  merge ns updates = V.sequence $ V.zipWith merge ns updates

norm :: V.Vector Double -> Double
norm = sqrt . V.sum . V.map (^ (2 :: Int))

max_stars_repo_name: massma/propagator-bbvi

max_issues_repo_name: massma/propagator-bbvi

max_forks_repo_name: massma/propagator-bbvi


text:
module CNC.DeclarativeTests where

import CNC.Declarative
import CNC.Geometry(eps)

import Data.Complex
import Control.Monad

fig1 = do [p1, p2] <- declarePoints 2
          xSize p1 p2 1

points1 = [0 :+ 0, 1 :+ 0]

fig2 = do [p1, p2, p3] <- declarePoints 3
          len p1 p2 1
          xAngle p1 p2 (pi/2)
          xAngle p2 p3 (-pi/2)
          len p2 p3 1

points2 = [0 :+ 0, 0 :+ 1.0, 0 :+ 0]

fig3 = do
  [p1,p2,p3,p4] <- declarePoints 4
  xSize p1 p2 1
  xAngle p2 p3 ( pi / 3)
  len p2 p3 2
  len p4 p3 2
  xAngle p3 p4 (-pi / 3)
--  xSize p3 p4 2

points3 = [0, 1, 1 + mkPolar 2 (pi/3), 3] --1 + mkPolar 2 (- 2*pi/3) + 2]

tests = [ (fig1,points1),
          (fig2, points2),
          (fig3, points3)
        ]

declarative_tests = do mapM_ run' tests
  where run' (fig, points) = do
          let solved = figure fig
          print solved
          when (any ((> 0.1) . magnitude) $ zipWith (-) (unPath solved) points) $ fail $ "test failed, should be " ++ show points

max_stars_repo_name: akamaus/gcodec

max_issues_repo_name: akamaus/gcodec

max_forks_repo_name: akamaus/gcodec


text:
{-# LANGUAGE FlexibleContexts, UndecidableInstances #-}

module School.Train.TrainState
( HandlerStore(..)
, TrainState(..)
, def
) where

import Data.Default.Class (Default(..))
import Numeric.LinearAlgebra (Container, Vector)
import School.Types.PingPong (PingPong, pingPongSingleton)
import School.Unit.UnitParams (UnitParams(..))

data HandlerStore a = CostList [a]
                    | NoStore deriving (Eq, Show)

data TrainState a = TrainState
  { handlerStore :: HandlerStore a
  , iterationCount :: Int
  , learningRate :: a
  , paramDerivs :: [UnitParams a]
  , paramList :: PingPong (UnitParams a)
  } deriving (Show)

instance (Container Vector a, Eq a, Num a) => Eq (TrainState a) where
  s1 == s2 = handlerStore s1 == handlerStore s2
          && iterationCount s1 == iterationCount s2
          && learningRate s1 == learningRate s2
          && paramDerivs s1 == paramDerivs s2
          && paramList s1 == paramList s2

instance (Num a) => Default (TrainState a) where
  def = TrainState { handlerStore = NoStore
                   , iterationCount = 0
                   , learningRate = 0
                   , paramDerivs = []
                   , paramList = pingPongSingleton EmptyParams
                   }

max_stars_repo_name: jfulseca/School

max_issues_repo_name: jfulseca/School

max_forks_repo_name: jfulseca/School


text:
module Elliptic
  where
import Carlson
import Data.Complex
import Internal

ellipticF' :: Double -> Cplx -> Cplx -> Cplx
ellipticF' err phi m
  | phi == 0 =
    toCplx 0
  | m == 1 && abs(realPart phi) == pi/2 =
    toCplx (0/0)
  | m == 1 && abs(realPart phi) < pi/2 =
    atanh(sin phi)
  | abs(realPart phi) <= pi/2 =
    if m == 0
      then
        phi
      else
        let sine = sin phi in
        let sine2 = sine*sine in
        let (cosine2, oneminusmsine2) = (1 - sine2, 1 - m*sine2) in
        sine * carlsonRF' err cosine2 oneminusmsine2 1
  | otherwise =
    let (phi', k) = getPhiK phi in
    2 * fromIntegral k * ellipticF' err (pi/2) m + ellipticF' err phi' m

ellipticF :: Cplx -> Cplx -> Cplx
ellipticF = ellipticF' 1e-15

ellipticE' :: Double -> Cplx -> Cplx -> Cplx
ellipticE' err phi m
  | phi == 0 =
    toCplx 0
  | abs(realPart phi) <= pi/2 =
    case m of
      0 -> phi
      1 -> sin phi
      _ ->
        let sine = sin phi in
        let sine2 = sine*sine in
        let (cosine2, oneminusmsine2) = (1 - sine2, 1 - m*sine2) in
        sine * (carlsonRF' err cosine2 oneminusmsine2 1 -
          m * sine2 / 3 * carlsonRD' err cosine2 oneminusmsine2 1)
  | otherwise =
    let (phi', k) = getPhiK phi in
    2 * fromIntegral k * ellipticE' err (pi/2) m + ellipticE' err phi' m

ellipticE :: Cplx -> Cplx -> Cplx
ellipticE = ellipticE' 1e-15

ellipticPI' :: Double -> Cplx -> Cplx -> Cplx -> Cplx
ellipticPI' err phi n m
  | phi == 0 =
    toCplx 0
  | phi == pi/2 && n == 1 =
    0/0
  | phi == pi/2 && m == 0 =
    pi/2/sqrt(1-n)
  | phi == pi/2 && m == n =
    ellipticE' err (pi/2) m / (1-m)
  | phi == pi/2 && n == 0 =
    ellipticF' err (pi/2) m
  | abs(realPart phi) <= pi/2 =
    let sine = sin phi in
    let sine2 = sine*sine in
    let (cosine2, oneminusmsine2) = (1 - sine2, 1 - m*sine2) in
    sine * (carlsonRF' err cosine2 oneminusmsine2 1 +
      n * sine2 / 3 * carlsonRJ' err cosine2 oneminusmsine2 1 (1-n*sine2))
  | otherwise =
    let (phi', k) = getPhiK phi in
    2 * fromIntegral k * ellipticPI' err (pi/2) n m + ellipticPI' err phi' n m

ellipticPI :: Cplx -> Cplx -> Cplx -> Cplx
ellipticPI = ellipticPI' 1e-15

jacobiZeta' :: Double -> Cplx -> Cplx -> Cplx
jacobiZeta' err phi m =
  if m == 1
    then
      if abs(realPart phi) <= pi/2
        then sin phi
        else let (phi',_) = getPhiK phi in sin phi'
    else
      ellipticE' err phi m -
        ellipticE' err (pi/2) m / ellipticF' err (pi/2) m *
        ellipticF' err phi m

jacobiZeta :: Cplx -> Cplx -> Cplx
jacobiZeta = jacobiZeta' 1e-15

max_stars_repo_name: stla/elliptic

max_issues_repo_name: stla/elliptic

max_forks_repo_name: stla/elliptic


text:
{-# LANGUAGE AllowAmbiguousTypes       #-}
{-# LANGUAGE DataKinds                 #-}
{-# LANGUAGE FlexibleContexts          #-}
{-# LANGUAGE GADTs                     #-}
{-# LANGUAGE OverloadedStrings         #-}
{-# LANGUAGE PolyKinds                 #-}
{-# LANGUAGE ScopedTypeVariables       #-}
{-# LANGUAGE TypeApplications          #-}
{-# LANGUAGE TypeOperators             #-}
{-# LANGUAGE UndecidableInstances      #-}

module BlueRipple.Model.Preference where

import qualified Control.Foldl                 as FL
import qualified Control.Lens                  as L
import qualified Control.Monad.Except          as X
import           Control.Monad.IO.Class         ( MonadIO(liftIO) )
import qualified Colonnade                     as C
import qualified Text.Blaze.Colonnade          as BC
import qualified Text.Blaze.Html               as BH
import qualified Text.Blaze.Html5.Attributes   as BHA
import qualified Data.List                     as L
import qualified Data.Map                      as M
import qualified Data.Array                    as A
import qualified Data.Vector                   as VB
import qualified Data.Vector.Storable          as VS

import qualified Text.Pandoc.Error             as PA

import qualified Data.Profunctor               as PF
import qualified Data.Text                     as T
import qualified Data.Time.Calendar            as Time
import qualified Data.Time.Clock               as Time
import qualified Data.Time.Format              as Time
import qualified Data.Vinyl                    as V
import qualified Text.Printf                   as PF
import qualified Frames                        as F
import qualified Frames.Melt                   as F
import qualified Frames.CSV                    as F
import qualified Frames.InCore                 as F
                                         hiding ( inCoreAoS )

import qualified Pipes                         as P
import qualified Pipes.Prelude                 as P
import qualified Statistics.Types              as S
import qualified Statistics.Distribution       as S
import qualified Statistics.Distribution.StudentT      as S

import qualified Numeric.LinearAlgebra         as LA

import qualified Graphics.Vega.VegaLite        as GV
import           Graphics.Vega.VegaLite.Configuration as FV
import qualified Graphics.Vega.VegaLite.Compat as FV

import qualified Frames.Visualization.VegaLite.Data
                                               as FV
import qualified Frames.Visualization.VegaLite.StackedArea
                                               as FV
import qualified Frames.Visualization.VegaLite.LineVsTime
                                               as FV
import qualified Frames.Visualization.VegaLite.ParameterPlots
                                               as FV

import qualified Frames.Visualization.VegaLite.Correlation
                                               as FV

import qualified Frames.Transform              as FT
import qualified Frames.Folds                  as FF
import qualified Frames.MapReduce              as MR
import qualified Frames.Enumerations           as FE

import qualified Knit.Report                   as K
import qualified Knit.Report.Input.MarkDown.PandocMarkDown    as K
import           Polysemy.Error                 (Error)
import qualified Text.Pandoc.Options           as PA

import           Data.String.Here               ( here, i )
import qualified Relude.Extra as Relude
import           BlueRipple.Configuration
import           BlueRipple.Utilities.KnitUtils
import           BlueRipple.Utilities.TableUtils
import           BlueRipple.Data.DataFrames
import           BlueRipple.Data.PrefModel
import           BlueRipple.Data.PrefModel.SimpleAgeSexRace
import           BlueRipple.Data.PrefModel.SimpleAgeSexEducation
import qualified BlueRipple.Model.PreferenceBayes as PB
import qualified BlueRipple.Model.TurnoutAdjustment as TA

modeledResults :: ( MonadIO (K.Sem r)
                  , K.KnitEffects r
                  , Show tr
                  , Show b
                  , Enum b
                  , Bounded b
                  , A.Ix b
                  , FL.Vector (F.VectorFor b) b)
               => DemographicStructure dr tr HouseElections b
               -> (F.Record LocationKey -> Bool)
               -> F.Frame dr
               -> F.Frame tr
               -> F.Frame HouseElections
               -> M.Map Int Int
               -> K.Sem r (M.Map Int (PreferenceResults b FV.ParameterEstimate))
modeledResults ds locFilter dFrame tFrame eFrame years = flip traverse years $ \y -> do
  K.logLE K.Info $ "inferring " <> show (dsCategories ds) <> " for " <> show y
  preferenceModel ds locFilter y dFrame eFrame tFrame

-- PreferenceResults to list of group names and predicted D votes
-- But we want them as a fraction of D/D+R
data VoteShare = ShareOfAll | ShareOfD

modeledDVotes :: forall b. (A.Ix b, Bounded b, Enum b, Show b)
  => VoteShare -> PreferenceResults b Double -> [(T.Text, Double)]
modeledDVotes vs pr =
  let
    summed = FL.fold
             (votesAndPopByDistrictF @b)
             (F.rcast <$> votesAndPopByDistrict pr)
    popArray =
      F.rgetField @(PopArray b) summed
    turnoutArray =
      F.rgetField @(TurnoutArray b) summed
    predVoters = zipWith (*) (A.elems turnoutArray) $ fmap realToFrac (A.elems popArray)
    allDVotes  = F.rgetField @DVotes summed
    allRVotes  = F.rgetField @RVotes summed
    dVotes b =
      realToFrac (popArray A.! b)
      * (turnoutArray A.! b)
      * (modeled pr A.! b)
    allPredictedD = FL.fold FL.sum $ fmap dVotes Relude.universe --[minBound..maxBound]
    scale = case vs of
      ShareOfAll -> (realToFrac allDVotes/realToFrac (allDVotes + allRVotes))/allPredictedD
      ShareOfD -> 1/allPredictedD
  in
    fmap (\b -> (show b, scale * dVotes b))
    [(minBound :: b) .. maxBound]

data DeltaTableRow =
  DeltaTableRow
  { dtrGroup :: T.Text
  , dtrPop :: Int
  , dtrFromPop :: Int
  , dtrFromTurnout :: Int
  , dtrFromOpinion :: Int
  , dtrTotal :: Int
  , dtrPct :: Double
  } deriving (Show)

deltaTable
  :: forall dr tr e b r
   . (A.Ix b
     , Bounded b
     , Enum b
     , Show b
     , MonadIO (K.Sem r)
     , K.KnitEffects r
     )
  => DemographicStructure dr tr e b
  -> (F.Record LocationKey -> Bool)
  -> F.Frame e
  -> Int -- ^ year A
  -> Int -- ^ year B
  -> PreferenceResults b FV.ParameterEstimate
  -> PreferenceResults b FV.ParameterEstimate
  -> K.Sem r ([DeltaTableRow], (Int, Int), (Int, Int))
deltaTable ds locFilter electionResultsFrame yA yB trA trB = do
  let
    groupNames = show <$> dsCategories ds
    getPopAndTurnout
      :: Int -> PreferenceResults b FV.ParameterEstimate -> K.Sem r (A.Array b Int, A.Array b Double)
    getPopAndTurnout y tr = do
      resultsFrame <- knitX $ dsPreprocessElectionData ds y electionResultsFrame
      let
        totalDRVotes =
          let filteredResultsF = F.filterFrame (locFilter . F.rcast) resultsFrame
          in FL.fold (FL.premap (\r -> F.rgetField @DVotes r + F.rgetField @RVotes r) FL.sum) filteredResultsF
        totalRec = FL.fold
          votesAndPopByDistrictF
          ( fmap
              (F.rcast
                @'[PopArray b, TurnoutArray b, DVotes, RVotes]
              )
          $ F.filterFrame (locFilter . F.rcast)
          $ F.toFrame
          $ votesAndPopByDistrict tr
          )
        totalCounts = F.rgetField @(PopArray b) totalRec
        unAdjTurnout = nationalTurnout tr
      tDelta <- liftIO $ TA.findDeltaA totalDRVotes totalCounts unAdjTurnout
      let adjTurnout = TA.adjTurnoutP tDelta unAdjTurnout
      return (totalCounts, adjTurnout)

  (popA, turnoutA) <- getPopAndTurnout yA trA
  (popB, turnoutB) <- getPopAndTurnout yB trB
{-  K.logLE K.Info $ (T.pack $ show yA) <> "->" <> (T.pack $ show yB)
  K.logLE K.Info $ T.pack $ show turnoutA
  K.logLE K.Info $ T.pack $ show turnoutB -}
  let
    pop        = FL.fold FL.sum popA
    probsArray = fmap FV.value . modeled
    probA      = probsArray trA
    probB      = probsArray trB
    modeledVotes popArray turnoutArray probArray =
      let dVotes b =
              round
                $ realToFrac (popArray A.! b)
                * (turnoutArray A.! b)
                * (probArray A.! b)
          rVotes b =
              round
                $ realToFrac (popArray A.! b)
                * (turnoutArray A.! b)
                * (1.0 - probArray A.! b)
      in  FL.fold
            ((,) <$> FL.premap dVotes FL.sum <*> FL.premap rVotes FL.sum) Relude.universe
    makeDTR b =
      let pop0     = realToFrac $ popA A.! b
          dPop     = realToFrac $ (popB A.! b) - (popA A.! b)
          turnout0 = realToFrac $ turnoutA A.! b
          dTurnout = realToFrac $ (turnoutB A.! b) - (turnoutA A.! b)
          prob0    = realToFrac (probA A.! b)
          dProb    = realToFrac $ (probB A.! b) - (probA A.! b)
          dtrCombo = dPop * dTurnout * (2 * dProb) / 4 -- the rest is accounted for in other terms, we spread this among them
          dtrN =
              round
                $ dPop
                * (turnout0 + dTurnout / 2)
                * (2 * (prob0 + dProb / 2) - 1)
                + (dtrCombo / 3)
          dtrT =
              round
                $ (pop0 + dPop / 2)
                * dTurnout
                * (2 * (prob0 + dProb / 2) - 1)
                + (dtrCombo / 3)
          dtrO =
              round
                $ (pop0 + dPop / 2)
                * (turnout0 + dTurnout / 2)
                * (2 * dProb)
                + (dtrCombo / 3)
          dtrTotal = dtrN + dtrT + dtrO
      in  DeltaTableRow (show b)
                        (popB A.! b)
                        dtrN
                        dtrT
                        dtrO
                        dtrTotal
                        (realToFrac dtrTotal / realToFrac pop)
    groupRows = fmap makeDTR [minBound ..]
    addRow (DeltaTableRow g p fp ft fo t _) (DeltaTableRow _ p' fp' ft' fo' t' _)
      = DeltaTableRow g
                      (p + p')
                      (fp + fp')
                      (ft + ft')
                      (fo + fo')
                      (t + t')
                      (realToFrac (t + t') / realToFrac (p + p'))
    totalRow = FL.fold
      (FL.Fold addRow (DeltaTableRow "Total" 0 0 0 0 0 0) id)
      groupRows
    dVotesA = modeledVotes popA turnoutA probA
    dVotesB = modeledVotes popB turnoutB probB
  return (groupRows ++ [totalRow], dVotesA, dVotesB)

deltaTableColonnade :: C.Colonnade C.Headed DeltaTableRow T.Text
deltaTableColonnade =
  C.headed "Group" dtrGroup
    <> C.headed "Population (k)" (show . (`div` 1000) . dtrPop)
    <> C.headed "+/- From Population (k)"
                (show . (`div` 1000) . dtrFromPop)
    <> C.headed "+/- From Turnout (k)"
                (show . (`div` 1000) . dtrFromTurnout)
    <> C.headed "+/- From Opinion (k)"
                (show . (`div` 1000) . dtrFromOpinion)
    <> C.headed "+/- Total (k)" (show . (`div` 1000) . dtrTotal)
    <> C.headed "+/- %Vote" (toText @String . PF.printf "%2.2f" . (* 100) . dtrPct)

deltaTableColonnadeBlaze :: CellStyle DeltaTableRow T.Text -> C.Colonnade C.Headed DeltaTableRow BC.Cell
deltaTableColonnadeBlaze cas =
  C.headed "Group" (toCell cas "Group" "" (textToStyledHtml . dtrGroup))
  <> C.headed "Population (k)" (toCell cas "Population" "Population" (textToStyledHtml . show . (`div` 1000) . dtrPop))
  <>  C.headed "+/- From Population (k)"
  (toCell cas "FromPop" "+/- From Population" (numberToStyledHtml "%d" . (`div` 1000) . dtrFromPop))
  <> C.headed "+/- From Turnout (k)"
  (toCell cas "FromTurnout" "+/- From Turnout" (numberToStyledHtml "%d" . (`div` 1000) . dtrFromTurnout))
  <> C.headed "+/- From Opinion (k)"
  (toCell cas "FromOpinion" "+/- From Opinion" (numberToStyledHtml "%d" .  (`div` 1000) . dtrFromOpinion))
--     (\tr -> (numberCell "%.0d" (opinionHighlight (dtrGroup tr)) . (`div` 1000) $ dtrFromOpinion tr))
  <> C.headed "+/- Total (k)" (toCell cas "Total" "Total" (numberToStyledHtml "%d" . (`div` 1000) . dtrTotal))
  <> C.headed "+/- %Vote" (toCell cas "PctVote" "% Vote" (numberToStyledHtml "%2.2f" . (* 100) . dtrPct))

type X = "X" F.:-> Double
type ScaledDVotes = "ScaledDVotes" F.:-> Int
type ScaledRVotes = "ScaledRVotes" F.:-> Int
type PopArray b = "PopArray" F.:-> A.Array b Int
type TurnoutArray b = "TurnoutArray" F.:-> A.Array b Double

votesAndPopByDistrictF
  :: forall b
   . (A.Ix b, Bounded b, Enum b)
  => FL.Fold
       (F.Record '[PopArray b, TurnoutArray b, DVotes, RVotes])
       (F.Record '[PopArray b, TurnoutArray b, DVotes, RVotes])
votesAndPopByDistrictF =
  let voters r = A.listArray (minBound, maxBound)
                 $ zipWith (*) (A.elems $ F.rgetField @(TurnoutArray b) r) (fmap realToFrac $ A.elems $ F.rgetField @(PopArray b) r)
      g r = A.listArray (minBound, maxBound)
            $ zipWith (/) (A.elems $ F.rgetField @(TurnoutArray b) r) (fmap realToFrac $ A.elems $ F.rgetField @(PopArray b) r)
      recomputeTurnout r = F.rputField @(TurnoutArray b) (g r) r
  in PF.dimap (F.rcast @'[PopArray b, TurnoutArray b, DVotes, RVotes]) recomputeTurnout
    $    FF.sequenceRecFold
    $    FF.FoldRecord (PF.dimap (F.rgetField @(PopArray b)) V.Field FE.sumTotalNumArray)
    V.:& FF.FoldRecord (PF.dimap voters V.Field FE.sumTotalNumArray)
    V.:& FF.FoldRecord (PF.dimap (F.rgetField @DVotes) V.Field FL.sum)
    V.:& FF.FoldRecord (PF.dimap (F.rgetField @RVotes) V.Field FL.sum)
    V.:& V.RNil

data Aggregation c b where
  Aggregation :: (Enum b, Bounded b, Eq b, Ord b, A.Ix b,
                  Enum c, Bounded c, Eq c, Ord c, A.Ix c) => (c -> [b]) -> Aggregation c b

aggregateFold :: forall c b a x. Aggregation c b -> FL.Fold a x -> A.Array b a -> A.Array c x
aggregateFold (Aggregation children) fld arr =
  let cs = Relude.universe
      expanded :: [[a]]
      expanded = fmap (fmap (arr A.!) . children) cs
      folded = fmap (FL.fold fld) expanded
  in A.listArray (minBound,maxBound) folded

aggregateFold2 :: Aggregation c b -> FL.Fold (a,w) x -> A.Array b a -> A.Array b w -> A.Array c x
aggregateFold2 (Aggregation children) fld arrA arrW =
  let cs = Relude.universe
      as = fmap (fmap (arrA A.!) . children) cs
      ws = fmap (fmap (arrW A.!) . children) cs
      folded = FL.fold fld . uncurry zip <$> zip as ws
  in A.listArray (minBound, maxBound) folded

weightedFold :: (Num a, Fractional a) => FL.Fold (a,a) a
weightedFold =
  let sumWeightsF = FL.premap snd FL.sum
      weightedSumF = FL.premap (uncurry (*)) FL.sum
  in fmap (uncurry (/)) $ (,) <$> weightedSumF <*> sumWeightsF

popWeightedAggregate :: (Num d, Fractional d) => Aggregation c b -> A.Array b Int -> A.Array b d -> A.Array c d
popWeightedAggregate agg popArray datArray = aggregateFold2 agg weightedFold datArray (fmap realToFrac popArray)

aggregateRecord
  :: forall b c. Aggregation c b
  -> F.Record [StateAbbreviation
              , CongressionalDistrict
              , PopArray b -- population by group
              , TurnoutArray b -- adjusted turnout by group
              , DVotes
              , RVotes
              ]
  -> F.Record [StateAbbreviation
              , CongressionalDistrict
              , PopArray c -- population by group
              , TurnoutArray c -- adjusted turnout by group
              , DVotes
              , RVotes
              ]
aggregateRecord agg r =
  let
    f :: F.Record [PopArray b, TurnoutArray b] -> F.Record [PopArray c, TurnoutArray c]
    f x =
      let popB = F.rgetField @(PopArray b) x
          turnoutB = F.rgetField @(TurnoutArray b) x
      in aggregateFold agg FL.sum popB F.&: popWeightedAggregate agg popB turnoutB F.&: V.RNil
  in F.rcast $ FT.transform f r

cByB :: Aggregation c b -> LA.Matrix Double
cByB (Aggregation children) =
  let allBs = Relude.universe
      toZeroOne y = if y then 1 else 0
      getCol c = LA.fromList $ fmap (toZeroOne . (`elem` children c)) allBs
  in LA.fromColumns $ fmap getCol Relude.universe

aggregatePreferenceResults :: Fractional a => Aggregation c b -> PreferenceResults b a -> PreferenceResults c a
aggregatePreferenceResults agg pr =
  let prVBPAD' = fmap (aggregateRecord agg) (votesAndPopByDistrict pr)
      prTurnout' = popWeightedAggregate agg (nationalVoters pr) (nationalTurnout pr)
      prVoters' = aggregateFold agg FL.sum (nationalVoters pr)
      prModeled' = popWeightedAggregate agg (nationalVoters pr) (modeled pr)
      mCB = cByB agg
      prCovar' = LA.tr mCB LA.<> covariances pr LA.<> mCB
  in PreferenceResults prVBPAD' prTurnout' prVoters' prModeled' prCovar'

data PreferenceResults b a = PreferenceResults
  {
    votesAndPopByDistrict :: [F.Record [ StateAbbreviation
                                       , CongressionalDistrict
                                       , PopArray b -- population by group
                                       , TurnoutArray b -- adjusted turnout by group
                                       , DVotes
                                       , RVotes
                                       ]]
    , nationalTurnout :: A.Array b Double
    , nationalVoters :: A.Array b Int
    , modeled :: A.Array b a
    , covariances :: LA.Matrix Double
  }

instance Functor (PreferenceResults b) where
  fmap f (PreferenceResults v nt nv m c) = PreferenceResults v nt nv (fmap f m) c

preferenceModel
  :: forall dr tr b r
   . ( Show tr
     , Show b
     , Enum b
     , Bounded b
     , A.Ix b
     , FL.Vector (F.VectorFor b) b
     , K.KnitEffects r
     , MonadIO (K.Sem r)
     )
  => DemographicStructure dr tr HouseElections b
  -> (F.Record LocationKey -> Bool)
  -> Int
  -> F.Frame dr
  -> F.Frame HouseElections
  -> F.Frame tr
  -> K.Sem
       r
       (PreferenceResults b FV.ParameterEstimate)
preferenceModel ds locFilter year identityDFrame houseElexFrame turnoutFrame =
  do
    -- reorganize data from loaded Frames
    resultsFlattenedFrameFull <- knitX
      $ dsPreprocessElectionData ds year houseElexFrame
    let resultsFlattenedFrame = F.filterFrame (locFilter . F.rcast) resultsFlattenedFrameFull
    filteredTurnoutFrame <- knitX
      $ dsPreprocessTurnoutData ds year turnoutFrame
    let year' = year --if (year == 2018) then 2017 else year -- we're using 2017 for now, until census updated ACS data
    longByDCategoryFrame <- knitX
      $ dsPreprocessDemographicData ds year' identityDFrame

    -- turn long-format data into Arrays by demographic category, beginning with national turnout
    turnoutByGroupArray <-
      K.knitMaybe "Missing or extra group in turnout data?" $ FL.foldM
        (FE.makeArrayMF (F.rgetField @(DemographicCategory b))
                        (F.rgetField @VotedPctOfAll)
                        (flip const)
        )
        filteredTurnoutFrame

    -- now the populations in each district
    let votersArrayMF = MR.mapReduceFoldM
          (MR.generalizeUnpack MR.noUnpack)
          (MR.generalizeAssign $ MR.splitOnKeys @LocationKey)
          (MR.foldAndLabelM
            (fmap (FT.recordSingleton @(PopArray b))
                  (FE.recordsToArrayMF @(DemographicCategory b) @PopCount)
            )
            V.rappend
          )
    -- F.Frame (LocationKey V.++ (PopArray b))
    populationsFrame <-
      K.knitMaybe "Error converting long demographic data to arrays!"
      $   F.toFrame
      <$> FL.foldM votersArrayMF longByDCategoryFrame

    -- and the total populations in each group
    let addArray :: (A.Ix k, Num a) => A.Array k a -> A.Array k a -> A.Array k a
        addArray a1 a2 = A.accum (+) a1 (A.assocs a2)
        zeroArray :: (A.Ix k, Bounded k, Enum k, Num a) => A.Array k a
        zeroArray = A.listArray (minBound, maxBound) $ L.repeat 0
        popByGroupArray = FL.fold (FL.premap (F.rgetField @(PopArray b)) (FL.Fold addArray zeroArray id)) populationsFrame

    let
      resultsWithPopulationsFrame =
        catMaybes $ fmap F.recMaybe $ F.leftJoin @LocationKey resultsFlattenedFrame
                                                              populationsFrame

    K.logLE K.Info "Computing Ghitza-Gelman turnout adjustment for each district so turnouts produce correct number D+R votes."
    resultsWithPopulationsAndGGAdjFrame <- fmap F.toFrame $ flip traverse resultsWithPopulationsFrame $ \r -> do
      let tVotesF x = F.rgetField @DVotes x + F.rgetField @RVotes x -- Should this be D + R or total?
      ggDelta <- ggTurnoutAdj r tVotesF turnoutByGroupArray
      K.logLE K.Diagnostic $
        "Ghitza-Gelman turnout adj="
        <> show ggDelta
        <> "; Adj Turnout=" <> show (TA.adjTurnoutP ggDelta turnoutByGroupArray)
      return $ FT.mutate (const $ FT.recordSingleton @(TurnoutArray b) $ TA.adjTurnoutP ggDelta turnoutByGroupArray) r

    let onlyOpposed r =
          (F.rgetField @DVotes r > 0) && (F.rgetField @RVotes r > 0)
        opposedFrame = F.filterFrame onlyOpposed resultsWithPopulationsAndGGAdjFrame
        numCompetitiveRaces = FL.fold FL.length opposedFrame

    K.logLE K.Info
      $ "After removing races where someone is running unopposed we have "
      <> show numCompetitiveRaces
      <> " contested races."

    totalVoteDiagnostics @b resultsWithPopulationsAndGGAdjFrame opposedFrame

    let
      scaleInt s n = round $ s * realToFrac n
      mcmcData =
        fmap
        (\r ->
           ( F.rgetField @DVotes r
           , VB.fromList $ fmap round (adjVotersL (F.rgetField @(TurnoutArray b) r) (F.rgetField @(PopArray b) r))
           )
        )
        $ FL.fold FL.list opposedFrame
      numParams = length $ dsCategories ds
    (cgRes, _, _) <- liftIO $ PB.cgOptimizeAD mcmcData (VB.fromList $ (const 0.5) <$> dsCategories ds)
    let cgParamsA = A.listArray (minBound :: b, maxBound) $ VB.toList cgRes
        cgVarsA = A.listArray (minBound :: b, maxBound) $ VS.toList $ PB.variances mcmcData cgRes
        npe cl b =
          let
            x = cgParamsA A.! b
            sigma = sqrt $ cgVarsA A.! b
            dof = realToFrac $ numCompetitiveRaces - L.length (A.elems cgParamsA)
            interval = S.quantile (S.studentTUnstandardized dof 0 sigma) (1.0 - (S.significanceLevel cl/2))
          in FV.ParameterEstimate x (x - interval/2.0, x + interval/2.0)
--          in FV.NamedParameterEstimate (T.pack $ show b) pEstimate
        parameterEstimatesA = A.listArray (minBound :: b, maxBound) $ fmap (npe S.cl95) Relude.universe

    K.logLE K.Info $ "MLE results: " <> show (A.elems parameterEstimatesA)
-- For now this bit is diagnostic.  But we should chart the correlations
-- and, perhaps, the eigenvectors of the covariance??
    let cgCovar = PB.covar mcmcData cgRes -- TODO: make a chart out of this
        (cgEv, cgEvs) = PB.mleCovEigens mcmcData cgRes
    K.logLE K.Diagnostic $ "sigma = " <> show (fmap sqrt $ cgVarsA)
    K.logLE K.Diagnostic $ "Covariances=" <> toText (PB.disps 3 cgCovar)
    K.logLE K.Diagnostic $ "Correlations=" <> toText (PB.disps 3 $ PB.correlFromCov cgCovar)
    K.logLE K.Diagnostic $ "Eigenvalues=" <> show cgEv
    K.logLE K.Diagnostic $ "Eigenvectors=" <> toText (PB.disps 3 cgEvs)

    return $ PreferenceResults
      (F.rcast <$> FL.fold FL.list opposedFrame)
      turnoutByGroupArray
      popByGroupArray
      parameterEstimatesA
      cgCovar

ggTurnoutAdj :: forall b rs r. (A.Ix b
                               , F.ElemOf rs (PopArray b)
                               , MonadIO (K.Sem r)
                               ) => F.Record rs -> (F.Record rs -> Int) -> A.Array b Double -> K.Sem r Double
ggTurnoutAdj r totalVotesF unadjTurnoutP = do
  let population = F.rgetField @(PopArray b) r
      totalVotes = totalVotesF r
  liftIO $ TA.findDeltaA totalVotes population unadjTurnoutP

adjVotersL :: A.Array b Double -> A.Array b Int -> [Double]
adjVotersL turnoutPA popA = zipWith (*) (A.elems turnoutPA) (realToFrac <$> A.elems popA)

totalVoteDiagnostics :: forall b rs f r
                        . (A.Ix b
                          , Foldable f
                          , F.ElemOf rs (PopArray b)
                          , F.ElemOf rs (TurnoutArray b)
                          , F.ElemOf rs Totalvotes
                          , F.ElemOf rs DVotes
                          , F.ElemOf rs RVotes
                          , K.KnitEffects r
                        )
  => f (F.Record rs) -- ^ frame with all rows
  -> f (F.Record rs) -- ^ frame with only rows from competitive races
  -> K.Sem r ()
totalVoteDiagnostics allFrame opposedFrame = K.wrapPrefix "VoteSummary" $ do
  let allVoters r = FL.fold FL.sum
                    $ zipWith (*) (A.elems $ F.rgetField @(TurnoutArray b) r) (fmap realToFrac $ A.elems $ F.rgetField @(PopArray b) r)
      allVotersF = FL.premap allVoters FL.sum
      allVotesF  = FL.premap (F.rgetField @Totalvotes) FL.sum
      allDVotesF = FL.premap (F.rgetField @DVotes) FL.sum
      allRVotesF = FL.premap (F.rgetField @RVotes) FL.sum
  --      allDRVotesF = FL.premap (\r -> F.rgetField @DVotes r + F.rgetField @RVotes r) FL.sum
      (totalVoters, totalVotes, totalDVotes, totalRVotes) = FL.fold
        ((,,,) <$> allVotersF <*> allVotesF <*> allDVotesF <*> allRVotesF)
        allFrame
      (totalVotersCD, totalVotesCD, totalDVotesCD, totalRVotesCD) = FL.fold
        ((,,,) <$> allVotersF <*> allVotesF <*> allDVotesF <*> allRVotesF)
        opposedFrame
  K.logLE K.Info $ "voters=" <> show totalVoters
  K.logLE K.Info $ "house votes=" <> show totalVotes
  K.logLE K.Info
    $  "D/R/D+R house votes="
    <> show totalDVotes
    <> "/"
    <> show totalRVotes
    <> "/"
    <> show (totalDVotes + totalRVotes)
  K.logLE K.Info
    $  "voters (competitive districts)="
    <> show totalVotersCD
  K.logLE K.Info
    $  "house votes (competitive districts)="
    <> show totalVotesCD
  K.logLE K.Info
    $  "D/R/D+R house votes (competitive districts)="
    <> show totalDVotesCD
    <> "/"
    <> show totalRVotesCD
    <> "/"
    <> show (totalDVotesCD + totalRVotesCD)

totalArrayZipWith :: (A.Ix b, Enum b, Bounded b)
                  => (x -> y -> z)
                  -> A.Array b x
                  -> A.Array b y
                  -> A.Array b z
totalArrayZipWith f xs ys = A.listArray (minBound, maxBound) $ zipWith f (A.elems xs) (A.elems ys)

vlGroupingChart :: Foldable f
                => T.Text
                -> FV.ViewConfig
                -> f (F.Record ['("Group", T.Text)
                               ,'("VotingAgePop", Int)
                               ,'("Turnout",Double)
                               ,'("Voters", Int)
                               ,'("D Voter Preference", Double)
                               ])
                -> GV.VegaLite
vlGroupingChart title vc rows =
  let dat = FV.recordsToVLData id FV.defaultParse rows
      xLabel = "Inferred (%) Likelihood of Voting Democratic"
      estimateXenc = GV.position GV.X [FV.pName @'("D Voter Preference", Double)
                                      ,GV.PmType GV.Quantitative
                                      ,GV.PAxis [GV.AxTitle xLabel]
                                      ]
      estimateYenc = GV.position GV.Y [FV.pName @'("Group",T.Text)
                                      ,GV.PmType GV.Ordinal
                                      ,GV.PAxis [GV.AxTitle "Demographic Group"]
                                      ]
      estimateSizeEnc = GV.size [FV.mName @'("Voters",Int)
                                , GV.MmType GV.Quantitative
                                , GV.MScale [GV.SDomain $ GV.DNumbers [5e6,30e6]]
                                , GV.MLegend [GV.LFormatAsNum]

                                ]
      estimateColorEnc = GV.color [FV.mName @'("Turnout", Double)
                                  , GV.MmType GV.Quantitative
                                  , GV.MScale [GV.SDomain $ GV.DNumbers [0.2,0.8]
                                              ,GV.SScheme "blues" [0.3,1.0]
                                              ]
                                  , GV.MLegend [GV.LGradientLength (vcHeight vc / 3)
                                               , GV.LFormatAsNum
                                               , GV.LFormat "%"
                                               ]
                                  ]
      estEnc = estimateXenc . estimateYenc . estimateSizeEnc . estimateColorEnc
      estSpec = GV.asSpec [(GV.encoding . estEnc) [], GV.mark GV.Point [GV.MFilled True]]
  in
    FV.configuredVegaLite vc [FV.title title, GV.layer [estSpec], dat]

exitCompareChart :: Foldable f
                 => T.Text
                 -> FV.ViewConfig
                 -> f (F.Record ['("Group", T.Text)
                                ,'("Model Dem Pref", Double)
                                ,'("ModelvsExit",Double)
                               ])
                -> GV.VegaLite
exitCompareChart title vc rows =
  let dat = FV.recordsToVLData id FV.defaultParse rows
      xLabel = "Modeled % Likelihood of Voting Democratic"
      xEnc =  GV.position GV.X [FV.pName @'("Model Dem Pref", Double)
                               ,GV.PmType GV.Quantitative
                               ,GV.PAxis [GV.AxTitle xLabel
                                         , GV.AxFormatAsNum
                                         , GV.AxFormat "%"
                                         ]
                               ]
      yEnc = GV.position GV.Y [FV.pName @'("ModelvsExit", Double)
                              ,GV.PmType GV.Quantitative
                              ,GV.PScale [GV.SDomain $ GV.DNumbers [negate 0.15,0.15]]
                              ,GV.PAxis [GV.AxTitle "Model - Exit Poll"
                                        , GV.AxFormatAsNum
                                        , GV.AxFormat "%"
                                        ]
                              ]
      colorEnc = GV.color [FV.mName @'("Group", T.Text)
                          , GV.MmType GV.Nominal
                          ]
      enc = xEnc . yEnc . colorEnc
      spec = GV.asSpec [(GV.encoding . enc) [], GV.mark GV.Point [GV.MFilled True, GV.MSize 100]]
  in
    FV.configuredVegaLite vc [FV.title title, GV.layer [spec], dat]

vlGroupingChartExit :: Foldable f
                    => T.Text
                    -> FV.ViewConfig
                    -> f (F.Record ['("Group", T.Text)
                                   ,'("VotingAgePop", Int)
                                   ,'("Voters", Int)
                                   ,'("D Voter Preference", Double)
                                   ,'("InfMinusExit", Double)
                                   ])
                    -> GV.VegaLite
vlGroupingChartExit title vc rows =
  let dat = FV.recordsToVLData id FV.defaultParse rows
      xLabel = "Inferred Likelihood of Voting Democratic"
      estimateXenc = GV.position GV.X [FV.pName @'("D Voter Preference", Double)
                                      ,GV.PmType GV.Quantitative
                                      ,GV.PAxis [GV.AxTitle xLabel]
                                      ]
      estimateYenc = GV.position GV.Y [FV.pName @'("Group",T.Text)
                                      ,GV.PmType GV.Ordinal
                                      ]
      estimateSizeEnc = GV.size [FV.mName @'("VotingAgePop",Int)
                                , GV.MmType GV.Quantitative]
      estimateColorEnc = GV.color [FV.mName @'("InfMinusExit", Double)
                                  , GV.MmType GV.Quantitative]
      estEnc = estimateXenc . estimateYenc . estimateSizeEnc . estimateColorEnc
      estSpec = GV.asSpec [(GV.encoding . estEnc) [], GV.mark GV.Point []]
  in
    FV.configuredVegaLite vc [FV.title title, GV.layer [estSpec], dat]

max_stars_repo_name: blueripple/preference-model

max_issues_repo_name: blueripple/preference-model

max_forks_repo_name: blueripple/preference-model


text:
-----------------------------------------------------------------------------
-- |
-- Module      :  Data.Random.Distribution.Static.MultivariateNormal
-- Copyright   :  (c) 2016 FP Complete Corporation
-- License     :  MIT (see LICENSE)
-- Maintainer  :  dominic@steinitz.org
--
-- Sample from the multivariate normal distribution with a given
-- vector-valued \(\mu\) and covariance matrix \(\Sigma\). For
-- example, the chart below shows samples from the bivariate normal
-- distribution. The dimension of the mean \(n\) is statically checked
-- to be compatible with the dimension of the covariance matrix \(n \times n\).
--
-- <<diagrams/src_Data_Random_Distribution_Static_MultivariateNormal_diagMS.svg#diagram=diagMS&height=600&width=500>>
--
-- Example code to generate the chart:
--
-- > {-# LANGUAGE DataKinds #-}
-- >
-- > import qualified Graphics.Rendering.Chart as C
-- > import Graphics.Rendering.Chart.Backend.Diagrams
-- >
-- > import Data.Random.Distribution.Static.MultivariateNormal
-- >
-- > import qualified Data.Random as R
-- > import Data.Random.Source.PureMT
-- > import Control.Monad.State
-- > import Numeric.LinearAlgebra.Static
-- >
-- > nSamples :: Int
-- > nSamples = 10000
-- >
-- > sigma1, sigma2, rho :: Double
-- > sigma1 = 3.0
-- > sigma2 = 1.0
-- > rho = 0.5
-- >
-- > singleSample :: R.RVarT (State PureMT) (R 2)
-- > singleSample = R.sample $ Normal (vector [0.0, 0.0])
-- >                (sym $ matrix [ sigma1, rho * sigma1 * sigma2
-- >                              , rho * sigma1 * sigma2, sigma2])
-- >
-- > multiSamples :: [R 2]
-- > multiSamples = evalState (replicateM nSamples $ R.sample singleSample) (pureMT 3)
-- >
-- > pts = map f multiSamples
-- >   where
-- >     f z = (x, y)
-- >       where
-- >         (x, t) = headTail z
-- >         (y, _) = headTail t
-- >
-- > chartPoint pointVals n = C.toRenderable layout
-- >   where
-- >
-- >     fitted = C.plot_points_values .~ pointVals
-- >               $ C.plot_points_style  . C.point_color .~ opaque red
-- >               $ C.plot_points_title .~ "Sample"
-- >               $ def
-- >
-- >     layout = C.layout_title .~ "Sampling Bivariate Normal (" ++ (show n) ++ " samples)"
-- >            $ C.layout_y_axis . C.laxis_generate .~ C.scaledAxis def (-3,3)
-- >            $ C.layout_x_axis . C.laxis_generate .~ C.scaledAxis def (-3,3)
-- >
-- >            $ C.layout_plots .~ [C.toPlot fitted]
-- >            $ def
-- >
-- > diagMS = do
-- >   denv <- defaultEnv C.vectorAlignmentFns 600 500
-- >   return $ fst $ runBackend denv (C.render (chartPoint pts nSamples) (500, 500))
--
-----------------------------------------------------------------------------

{-# OPTIONS_GHC -Wall                     #-}
{-# OPTIONS_GHC -fno-warn-name-shadowing  #-}
{-# OPTIONS_GHC -fno-warn-type-defaults   #-}
{-# OPTIONS_GHC -fno-warn-unused-do-bind  #-}
{-# OPTIONS_GHC -fno-warn-missing-methods #-}
{-# OPTIONS_GHC -fno-warn-orphans         #-}

{-# LANGUAGE MultiParamTypeClasses        #-}
{-# LANGUAGE TypeFamilies                 #-}
{-# LANGUAGE ScopedTypeVariables          #-}
{-# LANGUAGE DataKinds                    #-}

module Data.Random.Distribution.Static.MultivariateNormal
    ( Normal(..)
    ) where

import           Data.Random hiding ( StdNormal, Normal )
import qualified Data.Random as R
import           Control.Monad.State ( replicateM )
import qualified Numeric.LinearAlgebra.HMatrix as H
import           Numeric.LinearAlgebra.Static as S
                 ( R, vector, extract, Sq, Sym, col,
                   tr, linSolve, uncol, chol, (<.>),
                   ℝ, (<>), diag, (#>), eigensystem
                 )
import          GHC.TypeLits ( KnownNat, natVal )
import          Data.Maybe ( fromJust )

normalMultivariate :: KnownNat n =>
                      R n -> Sym n -> RVarT m (R n)
normalMultivariate mu bigSigma = do
  z <- replicateM (fromIntegral $ natVal mu) (rvarT R.StdNormal)
  return $ mu + bigA #> (vector z)
  where
    (vals, bigU) = eigensystem bigSigma
    lSqrt = diag $ mapVector sqrt vals
    bigA = bigU S.<> lSqrt

mapVector :: KnownNat n => (ℝ -> ℝ) -> R n -> R n
mapVector f = vector . H.toList . H.cmap f . extract

sumVector :: KnownNat n => R n -> ℝ
sumVector x = x <.> 1

data family Normal k :: *

data instance Normal (R n) = Normal (R n) (Sym n)

instance KnownNat n => Distribution Normal (R n) where
  rvar (Normal m s) = normalMultivariate m s

normalLogPdf :: KnownNat n =>
                R n -> Sym n -> R n -> Double
normalLogPdf mu bigSigma x = - sumVector (mapVector log (diagonals dec))
                             - 0.5 * (fromIntegral $ natVal mu) * log (2 * pi)
                             - 0.5 * s
  where
    dec = chol bigSigma
    t = uncol $ fromJust $ linSolve (tr dec) (col $ x - mu)
    u = mapVector (\x -> x * x) t
    s = sumVector u

normalPdf :: KnownNat n =>
             R n -> Sym n -> R n -> Double
normalPdf mu sigma x = exp $ normalLogPdf mu sigma x

diagonals :: KnownNat n => Sq n -> R n
diagonals = vector . H.toList . H.takeDiag . extract

instance KnownNat n => PDF Normal (R n) where
  pdf (Normal m s) = normalPdf m s
  logPdf (Normal m s) = normalLogPdf m s

max_stars_repo_name: steinitznavican/random-fu-multivariate

max_issues_repo_name: steinitznavican/random-fu-multivariate

max_forks_repo_name: steinitznavican/random-fu-multivariate


text:
{-# LANGUAGE TemplateHaskell #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE RankNTypes #-}
import           Control.Lens          hiding ((<.>), (|>))
import           Control.Monad.State
import           Numeric.LinearAlgebra
import Data.List (sortBy)
import           System.Random

data RBF = RBF { _position :: Matrix R, _stddev :: Vector R, _weights :: Matrix R, _eta :: R } deriving (Show)

makeLenses ''RBF

type STR a = State (RBF, StdGen) a

rbf :: (Functor f, Field1 s t a b) => (a -> f b) -> s -> f t
rbf = _1

rng :: (Functor f, Field2 s t a b) => (a -> f b) -> s -> f t
rng = _2

distance :: Vector R -> R -> Vector R -> R
distance position stddev x = normalization * exponent
  where
    dims = size position
    diff = x - position
    normalization = 2 * pi * det cov
    exponent = exp (-0.5 * ((diff <# cov) <.> diff))
    cov = ident dims * scalar stddev

relu :: (Num a, Ord a) => a -> a
relu x = if x <= 0 then 0 else x

delta :: Matrix R -> Vector R -> Vector R -> R -> Matrix R
delta w x t eta = scalar eta * (x `outer` (t - out))
  where
    out :: Vector R
    out = tr w #> x

sgn
  :: (Num b, Num a, Ord a, Container c a, Element b) => c a -> c b
sgn = cmap (\x -> if x >= 0 then 1 else -1)

trainpercept :: Vector R -> Vector R -> STR ()
trainpercept x t = do
  r <- use rbf
  rbf.weights += delta (r ^. weights) (addbias x) t (r ^. eta)
  return ()

fwd
  :: (Field1 s s RBF RBF, MonadState s m) => Vector R -> m (Vector R)
fwd x = do
  w <- use $ rbf.weights
  return $ tr w #> addbias x

addbias :: Vector R -> Vector R
addbias x = vjoin [x, scalar 1]

nRandoms :: Int -> STR (Vector R)
nRandoms n = do
  vector <$> replicateM n getone
  where
    getone = do
      r <- use rng
      let (nr, r') = randomR (0.0, 1.0) r
      rng .= r'
      return nr

nullrbf inp hid out eta = RBF { _position = (hid><inp) z
                              , _stddev = hid |> z
                              , _weights = ((hid+1)><out) z
                              , _eta = eta}
  where
    z = repeat 0

addbiasv :: R -> Vector R -> STR (Vector R)
addbiasv hi v = do
  rng <- nRandoms $ fromIntegral $ size v
  return $ (rng*2 - 1)* scalar hi

addbiasm :: R -> Matrix R -> STR (Matrix R)
addbiasm hi m = do
  let r = rows m
      c = cols m
  rng <- nRandoms $ r*c
  return $ reshape c $ (rng*2 - 1)* scalar hi

rngrbf :: R -> STR ()
rngrbf hi = do
  let apply :: (Container c R) => (R -> c R -> STR (c R)) -> Lens' (RBF,StdGen) (c R) -> STR ()
      apply f l = do use l >>= f hi >>= (l .=)
  apply addbiasm (rbf.position)
  apply addbiasv (rbf.stddev)
  apply addbiasm (rbf.weights)
  return ()

xors :: [(Vector R, Vector R)]
xors = [([0,0], [0]), ([0,1], [1]), ([1,0], [1]), ([1,1], [0])] & traverse.both %~ vector & traverse._2 %~ sgn . (subtract 0.5)

ands :: [(Vector R, Vector R)]
ands = [([0,0], [0]), ([0,1], [0]), ([1,0], [0]), ([1,1], [1])] & traverse.both %~ vector & traverse._2 %~ sgn . (subtract 0.5)

shuffle :: [a] -> STR [a]
shuffle x = do
  r <- nRandoms $ length x
  return $ snd <$> sortBy (\x y -> fst x `compare` fst y) (zip (toList r) x)

testa :: IO ()
testa = do
  let rbfn = nullrbf 2 2 1 0.1
  rng <- newStdGen
  let s = flip evalState (rbfn, rng) $ do
        rngrbf 0.01
        replicateM 100 $ do
          a <- shuffle ands
          sequenceA $ uncurry trainpercept <$> a
        ans <- sequenceA (fwd <$> ands ^.. traverse._1)
        w <- use $ rbf.weights
        return (ans,w)

  print s

max_stars_repo_name: fizzoo/kod

max_issues_repo_name: fizzoo/kod

max_forks_repo_name: fizzoo/kod


text:
{-# LANGUAGE OverloadedLists #-}
{-# LANGUAGE DataKinds #-}
{-# OPTIONS_GHC -Wall #-}
{-# OPTIONS_GHC -fno-warn-unused-imports #-}

module Numeric.LinearAlgebra where

import Numeric.LinearAlgebra.Internal

max_stars_repo_name: DataHaskell/numhask-linear-algebra

max_issues_repo_name: DataHaskell/numhask-linear-algebra

max_forks_repo_name: DataHaskell/numhask-linear-algebra


text:
{-# LANGUAGE ApplicativeDo                        #-}
{-# LANGUAGE DataKinds                            #-}
{-# LANGUAGE DeriveFunctor                        #-}
{-# LANGUAGE DeriveGeneric                        #-}
{-# LANGUAGE FlexibleContexts                     #-}
{-# LANGUAGE FlexibleInstances                    #-}
{-# LANGUAGE FunctionalDependencies               #-}
{-# LANGUAGE GADTs                                #-}
{-# LANGUAGE KindSignatures                       #-}
{-# LANGUAGE LambdaCase                           #-}
{-# LANGUAGE MultiParamTypeClasses                #-}
{-# LANGUAGE OverloadedStrings                    #-}
{-# LANGUAGE PartialTypeSignatures                #-}
{-# LANGUAGE PolyKinds                            #-}
{-# LANGUAGE RankNTypes                           #-}
{-# LANGUAGE RecordWildCards                      #-}
{-# LANGUAGE ScopedTypeVariables                  #-}
{-# LANGUAGE StandaloneDeriving                   #-}
{-# LANGUAGE TemplateHaskell                      #-}
{-# LANGUAGE TupleSections                        #-}
{-# LANGUAGE TypeApplications                     #-}
{-# LANGUAGE TypeFamilies                         #-}
{-# LANGUAGE TypeOperators                        #-}
{-# LANGUAGE TypeSynonymInstances                 #-}
{-# LANGUAGE UndecidableInstances                 #-}
{-# LANGUAGE ViewPatterns                         #-}
{-# OPTIONS_GHC -fno-warn-partial-type-signatures #-}

module Egg.Research (
    BonusAmount(..), AsBonusAmount(..)
  , BonusType(..), AsBonusType(..)
  , Bonuses(..), _Bonuses
  , Research(..), HasResearch(..)
  , ResearchTier(..), rtUnlock, rtTechs, SomeResearchTier
  , ResearchData(..), rdCommon, rdEpic, SomeResearchData
  , ResearchStatus(..), rsCommon, rsEpic, SomeResearchStatus
  , ResearchIx(..), _RICommon, _RIEpic
  , ResearchError(..), _REMaxedOut, _RELocked
  , scaleAmount
  , hasEffect
  , bonusEffect
  , bonusEffectFor
  , bonusing
  , bonusingFor
  , emptyBonuses
  , bonuses
  , maxLevel
  , emptyResearchStatus
  , researchBonuses
  , totalBonuses
  , researchCount
  , legalTiers
  , foldResearch
  , purchaseResearch
  , researchIxNum
  , researchIxData
  , researchIxStatus
  , researchIxStatusLegal
  , withSomeResearch
  , researchIxesCommon
  , researchIxesEpic
  , legalResearchIxesCommon
  , legalResearchIxesEpic
  , legalResearchesCommon
  , legalResearchesEpic
  -- * Why?
  , rdrsCommon
  , rdrsEpic
  , researchIx
  ) where

import           Control.Applicative hiding      (some)
import           Control.Lens hiding             ((.=), (:<), Index)
import           Data.Aeson.Encoding
import           Data.Aeson.Types
import           Data.Bifunctor
import           Data.Dependent.Sum
import           Data.Finite
import           Data.Foldable
import           Data.Kind
import           Data.Maybe
import           Data.Monoid                     (First(..))
import           Data.Singletons
import           Data.Singletons.Prelude hiding  (Flip)
import           Data.Singletons.TypeLits
import           Data.Type.Combinator
import           Data.Type.Combinator.Singletons
import           Data.Type.Combinator.Util
import           Data.Type.Conjunction
import           Data.Type.Index
import           Data.Type.Product               as TCP
import           Data.Type.Sum
import           Data.Type.Vector
import           Data.Vector.Sized.Util
import           Egg.Commodity
import           GHC.Generics                    (Generic)
import           Numeric.Lens
import           Numeric.Natural
import           Statistics.LinearRegression
import           Type.Class.Higher
import           Type.Class.Witness
import           Type.Family.Tuple
import qualified Data.Map                        as M
import qualified Data.Text                       as T
import qualified Data.Vector                     as V
import qualified Data.Vector.Sized               as SV

data BonusAmount
    -- | Bonus adds an absolute amount.  Scales by multiplication, and
    -- compounds by addition.
    = BAIncrement Double
    -- | Bonus adds a percentage point multiplier.  Scales by
    -- multiplication (three 10% buffs is a 15% buff) and compounds by
    -- appropriate multiplication.
    | BAPercent Double
    -- | Bonus multiplies by a given ratio.  Scales by exponentiation
    -- (three 2x buffs is a 8x buff) and compounds by multiplication.
    | BAMultiplier Double
  deriving (Show, Eq, Ord, Generic)

makeClassyPrisms ''BonusAmount

data BonusType =
        BTBuildCosts
      | BTDroneRewards
      | BTEggValue
      | BTFarmValue
      | BTFleetSize
      | BTHabCapacity
      | BTHatcheryCapacity
      | BTHatcheryRate
      | BTHoldingHatch
      | BTHoverVehicleCapacity
      | BTInternalHatchery
      | BTInternalHatcheryCalm
      | BTInternalHatcherySharing
      | BTLayingRate
      | BTLongWarpTime
      | BTMaxRunningBonus
      | BTPrestigeEggs
      | BTResearchCosts
      | BTRunningBonus
      | BTSiloQuality
      | BTSoulEggBonus
      | BTVehicleCapacity
      | BTVehicleCosts
      | BTVideoDoublerTime
      | BTVehicleSpeed
  deriving (Show, Eq, Ord, Generic)

makeClassyPrisms ''BonusType

newtype Bonuses = Bonuses { _bMap :: M.Map BonusType [BonusAmount] }
    deriving (Show, Eq, Ord)

makePrisms ''Bonuses
makeWrapped ''Bonuses

data Research a =
    Research { _rName        :: T.Text
             , _rDescription :: T.Text
             , _rBaseBonuses :: Bonuses
             , _rCosts       :: V.Vector (Maybe a)
             }
  deriving (Show, Eq, Ord, Generic, Functor)

makeClassy ''Research

data ResearchTier :: Nat -> Type where
    ResearchTier
        :: { _rtUnlock :: Natural
           , _rtTechs  :: SV.Vector n (Research Bock)
           }
        -> ResearchTier n
  deriving (Show, Eq, Ord, Generic)

makeLenses ''ResearchTier

type SomeResearchTier = DSum Sing ResearchTier

data ResearchData :: [Nat] -> Nat -> Type where
    ResearchData
        :: { _rdCommon :: Prod ResearchTier tiers
           , _rdEpic   :: SV.Vector epic (Research GoldenEgg)
           }
        -> ResearchData tiers epic
  deriving (Show, Eq, Ord)

makeLenses ''ResearchData

type SomeResearchData = DSum Sing (Uncur ResearchData)

data ResearchStatus :: [Nat] -> Nat -> Type where
    ResearchStatus
        :: { _rsCommon :: Prod (Flip SV.Vector Natural) tiers
           , _rsEpic   :: SV.Vector epic Natural
           }
        -> ResearchStatus tiers epic
  deriving (Show, Eq, Ord, Generic)

makeLenses ''ResearchStatus

type SomeResearchStatus = DSum Sing (Uncur ResearchStatus)

-- | A safe index for a given research item, usable with 'ResearchData' and
-- 'ResearchStatus'.
data ResearchIx :: [Nat] -> Nat -> Type -> Type where
    RICommon :: Sum Finite tiers -> ResearchIx tiers epic Bock
    RIEpic   :: Finite epic      -> ResearchIx tiers epic GoldenEgg

deriving instance Show (Sum Finite tiers) => Show (ResearchIx tiers epic a)

_RICommon :: Iso (ResearchIx t1 epic Bock) (ResearchIx t2 epic Bock)
                 (Sum Finite t1          ) (Sum Finite t2          )
_RICommon = iso (\case RICommon i -> i) RICommon

_RIEpic :: Iso (ResearchIx tiers e1 GoldenEgg) (ResearchIx tiers e2 GoldenEgg)
               (Finite e1                    ) (Finite e2                    )
_RIEpic = iso (\case RIEpic i -> i) RIEpic

bonusAmountParseOptions :: Options
bonusAmountParseOptions = defaultOptions
    { sumEncoding = TaggedObject
                      { tagFieldName      = "type"
                      , contentsFieldName = "value"
                      }
    , constructorTagModifier = camelTo2 '-' . drop 2
    }

instance FromJSON BonusAmount where
    parseJSON  = genericParseJSON  bonusAmountParseOptions
instance ToJSON BonusAmount where
    toJSON     = genericToJSON     bonusAmountParseOptions
    toEncoding = genericToEncoding bonusAmountParseOptions

bonusTypeParseOptions :: Options
bonusTypeParseOptions = defaultOptions
    { sumEncoding = UntaggedValue
    , constructorTagModifier = camelTo2 '-' . drop 2
    }

instance FromJSON BonusType where
    parseJSON  = genericParseJSON  bonusTypeParseOptions
instance ToJSON BonusType where
    toJSON     = genericToJSON     bonusTypeParseOptions
    toEncoding = genericToEncoding bonusTypeParseOptions
instance FromJSONKey BonusType where
    fromJSONKey = FromJSONKeyTextParser (parseJSON . String)
instance ToJSONKey BonusType where
    toJSONKey = toJSONKeyText $
        T.pack . constructorTagModifier bonusTypeParseOptions . show

instance Semigroup Bonuses where
    (<>) = mappend

instance Monoid Bonuses where
    mempty = Bonuses M.empty
    mappend (Bonuses x) (Bonuses y) =
      Bonuses (M.unionWith (++) x y)

instance (FromJSON a, RealFrac a) => FromJSON (Research a) where
    parseJSON = withObject "Research" $ \v ->
        Research <$> v .: "name"
                 <*> v .: "description"
                 <*> (mkBase =<< v .: "bonuses")
                 <*> (mkCosts <$> (v .:? "costs") <*> (v .:? "levels"))
      where
        mkBase :: M.Map BonusType Object -> Parser Bonuses
        mkBase = fmap (Bonuses . fmap (:[])) . traverse (.: "base-amount")
        mkCosts :: Maybe [a] -> Maybe Natural -> V.Vector (Maybe a)
        mkCosts Nothing        Nothing  = V.empty
        mkCosts Nothing        (Just l) = V.replicate (fromIntegral l) Nothing
        mkCosts (Just [] )     (Just l) = V.replicate (fromIntegral l) Nothing
        mkCosts (Just v  )     Nothing  = Just <$> V.fromList v
        mkCosts (Just [x])     (Just l) = Just x `V.cons` V.replicate (fromIntegral l - 1) Nothing
        mkCosts (Just v@(_:_)) (Just l) = Just <$> V.fromList v V.++ extras
          where
            lv = length v
            extras = V.generate (fromIntegral l - lv) $ \i -> realToFrac . exp $
              α + β * fromIntegral (i + lv)
            (α, β) = linearRegression (V.fromList $ zipWith const [0..] v)
                                      (V.fromList (log . realToFrac <$> v))

instance ToJSON a => ToJSON (Research a) where
    toJSON Research{..} = object $
        [ "name"        .= _rName
        , "description" .= _rDescription
        , "bonuses"     .= object [ "base-amount" .= _bMap _rBaseBonuses ]
        ] ++ case sequence _rCosts of
            Just v  ->   [ "costs"  .= v ]
            Nothing -> case sequence (V.takeWhile isJust _rCosts) of
              Nothing -> [ "levels" .= V.length _rCosts ]
              Just v  -> [ "costs"  .= v
                         , "levels" .= V.length _rCosts
                         ]
    toEncoding Research{..} = pairs . mconcat $
        [ "name"        .= _rName
        , "description" .= _rDescription
        , pair "bonuses" (pairs ("base-amount" .= _bMap _rBaseBonuses))
        ] ++ case sequence _rCosts of
            Just v  ->   [ "costs"  .= v ]
            Nothing -> case sequence (V.takeWhile isJust _rCosts) of
              Nothing -> [ "levels" .= V.length _rCosts ]
              Just v  -> [ "costs"  .= v
                         , "levels" .= V.length _rCosts
                         ]

researchTierParseOptions :: Options
researchTierParseOptions = defaultOptions
    { fieldLabelModifier = camelTo2 '-' . drop 3
    }

instance KnownNat n => FromJSON (ResearchTier n) where
    parseJSON  = genericParseJSON  researchTierParseOptions
instance KnownNat n => ToJSON (ResearchTier n) where
    toJSON     = genericToJSON     researchTierParseOptions
    toEncoding = genericToEncoding researchTierParseOptions
instance FromJSON SomeResearchTier where
    parseJSON = withObject "ResearchTier" $ \v -> do
      u <- v .: "unlock"
      t <- v .: "techs"
      SV.withSized t $ \tV ->
        return $ sing :=> ResearchTier u tV
instance ToJSON SomeResearchTier where
    toJSON = \case
        SNat :=> r -> toJSON r
    toEncoding = \case
        SNat :=> r -> toEncoding r
instance FromJSON SomeResearchData where
    parseJSON = withObject "ResearchData" $ \v -> do
        res   <- v .: "common"
        epics <- (fmap . fmap) (round @Double) <$> v .: "epic"
        withV res $ \resV ->
          some (withProd (dsumSome . getI) resV) $ \(_ :&: (unzipP->(resS :&: resP))) ->
            SV.withSized epics   $ \sizedV ->
              return (STuple2 (fromTC @(Prod Sing) resS) SNat :=> Uncur (ResearchData resP sizedV))
instance ToJSON SomeResearchData where
    toJSON = \case
        _ :=> Uncur r -> toJSON r
    toEncoding = \case
        _ :=> Uncur r -> toEncoding r
instance (SingI tiers, KnownNat epic) => FromJSON (ResearchData tiers epic) where
    parseJSON = withObject "ResearchData" $ \v -> do
        res   <- v .: "common"
        resV  <- go sing res
        epics <- (fmap . fmap) (round @Double) <$> v .: "epic"
        epicsV <- case SV.toSized epics of
          Nothing -> fail "Bad number of items in list."
          Just eV -> return eV
        return $ ResearchData resV epicsV
      where
        go :: Sing ts -> [Value] -> Parser (Prod ResearchTier ts)
        go = \case
          SNil -> \case
            []  -> return Ø
            _:_ -> fail "Too many items in list"
          SNat `SCons` ss -> \case
            []   -> fail "Too few items in list"
            x:xs -> (:<) <$> parseJSON x <*> go ss xs

instance ToJSON (ResearchData tiers epic) where
    toJSON ResearchData{..} = object
        [ "common" .= TCP.toList (\ResearchTier{..} ->
                          object [ "unlock" .= _rtUnlock
                                 , "techs"  .= SV.fromSized _rtTechs
                                 ]
                        ) _rdCommon
        , "epic"   .= SV.fromSized _rdEpic
        ]
    toEncoding ResearchData{..} = pairs . mconcat $
        [ "common" .= TCP.toList (\ResearchTier{..} ->
                          object [ "unlock" .= _rtUnlock
                                 , "techs"  .= SV.fromSized _rtTechs
                                 ]
                        ) _rdCommon
        , "epic"   .= SV.fromSized _rdEpic
        ]

instance (SingI tiers, KnownNat epic) => FromJSON (ResearchStatus tiers epic) where
    parseJSON = withObject "ResearchStatus" $ \v -> do
        res   <- v .: "common"
        resV  <- go sing res
        epics <- v .: "epic"
        epicsV <- case SV.toSized epics of
          Nothing -> fail "Bad number of items in list."
          Just eV -> return eV
        return $ ResearchStatus resV epicsV
      where
        go :: Sing ts -> [Value] -> Parser (Prod (Flip SV.Vector Natural) ts)
        go = \case
          SNil -> \case
            []  -> return Ø
            _:_ -> fail "Too many items in list"
          SNat `SCons` ss -> \case
            []   -> fail "Too few items in list"
            x:xs -> (:<) <$> parseJSON x <*> go ss xs
instance ToJSON (ResearchStatus tiers epic) where
    toJSON ResearchStatus{..} = object
        [ "common" .= TCP.toList (SV.fromSized . getFlip) _rsCommon
        , "epic"   .= SV.fromSized _rsEpic
        ]
    toEncoding ResearchStatus{..} = pairs . mconcat $
        [ "common" .= TCP.toList (SV.fromSized . getFlip) _rsCommon
        , "epic"   .= SV.fromSized _rsEpic
        ]
instance FromJSON SomeResearchStatus where
    parseJSON = withObject "ResearchStatus" $ \v -> do
        res   <- v .: "common"
        epics <- v .: "epic"
        withV res $ \resV ->
          some (withProd (go . getI) resV) $ \(_ :&: (unzipP->(resS :&: resP))) ->
            SV.withSized epics   $ \sizedV ->
              return (STuple2 (fromTC resS) SNat :=> Uncur (ResearchStatus resP sizedV))
      where
        go :: V.Vector Natural -> Some (Sing :&: Flip SV.Vector Natural)
        go v = SV.withSized v $ \u -> Some (SNat :&: Flip u)
instance ToJSON SomeResearchStatus where
    toJSON = \case
        _ :=> Uncur r -> toJSON r
    toEncoding = \case
        _ :=> Uncur r -> toEncoding r

-- | Scales a bonus amount by a "level".
scaleAmount :: Natural -> BonusAmount -> BonusAmount
scaleAmount n = \case
    BAIncrement  i -> BAIncrement  (fromIntegral n * i)
    BAPercent    p -> BAPercent    (fromIntegral n * p)
    BAMultiplier r -> BAMultiplier (r ^ n)

-- | Tells whether or not a bonus amount creates any actual effect.
hasEffect :: BonusAmount -> Bool
hasEffect = \case
    BAIncrement  i -> i /= 0
    BAPercent    p -> p /= 0
    BAMultiplier r -> r /= 1

-- | A "smart constructor" that clears bonuses with no effect.
bonuses :: M.Map BonusType [BonusAmount] -> Bonuses
bonuses = Bonuses . M.filter (not . null) . fmap (filter hasEffect)

-- | Apply a list of bonuses (from left to right) to a value.
bonusEffect :: Fractional a => [BonusAmount] -> a -> a
bonusEffect = flip . foldl' $ \e -> \case
    BAIncrement  i -> e + realToFrac i
    BAPercent    p -> e * (1 + realToFrac p / 100)
    BAMultiplier r -> e * realToFrac r

-- | Iso on a value under a list of bonuses (from left to right).
bonusing :: (Fractional a, Eq a) => [BonusAmount] -> Iso' a a
bonusing = flip foldl' id $ \e -> \case
             BAIncrement i  -> e . adding (realToFrac i)
             BAPercent   p  -> e . multiplying (1 + realToFrac p / 100)
             BAMultiplier r -> e . multiplying (realToFrac r)

-- | Iso on a value under the bonuses of a given bonus type.
bonusingFor
    :: (Fractional a, Eq a)
    => Bonuses
    -> BonusType
    -> Iso' a a
bonusingFor bs bt = case M.lookup bt (_bMap bs) of
                      Nothing -> id
                      Just bl -> bonusing bl

-- | Apply bonuses for a given type on a value.
bonusEffectFor :: Fractional a => Bonuses -> BonusType -> a -> a
bonusEffectFor bs bt = maybe id bonusEffect $ M.lookup bt (_bMap bs)

-- | No bonuses
emptyBonuses :: Bonuses
emptyBonuses = Bonuses M.empty

-- | Maximum level for a given research.
maxLevel :: Research a -> Natural
maxLevel = fromIntegral . V.length . _rCosts

-- | Empty status (no research).
emptyResearchStatus :: ResearchData tiers epic -> ResearchStatus tiers epic
emptyResearchStatus ResearchData{..} =
    ResearchStatus (map1 clear _rdCommon) (0 <$ _rdEpic)
  where
    clear :: ResearchTier a -> Flip SV.Vector Natural a
    clear = Flip . set mapped 0 . view rtTechs

-- | Zips together all research data and research statuses into an
-- accumulator.
foldResearch
    :: Monoid b
    => (Either (Research Bock) (Research GoldenEgg) -> Natural -> b)
    -> ResearchData tiers epic
    -> ResearchStatus tiers epic
    -> b
foldResearch f ResearchData{..} ResearchStatus{..} = mconcat
    [ foldMap1 (\case d :&: Flip s -> fold $ SV.zipWith (f . Left) (_rtTechs d) s
               )
        (zipP (_rdCommon :&: _rsCommon))
    , fold $ SV.zipWith (f . Right) _rdEpic _rsEpic
    ]

-- | Bonuses from a given 'Research' at a given level.  Assumes no "maximum
-- level".
researchBonuses :: Research a -> Natural -> Bonuses
researchBonuses _ 0 = Bonuses M.empty
researchBonuses r l = _rBaseBonuses r & _Bonuses . mapped . mapped %~ scaleAmount l

-- | Total bonuses from a given 'ResearchStatus'.
totalBonuses :: ResearchData tiers epic -> ResearchStatus tiers epic -> Bonuses
totalBonuses = foldResearch (either researchBonuses researchBonuses)

-- | How many common techs have been researched?
--
-- Used for opening tiers.
researchCount :: ResearchStatus tiers epic -> Natural
researchCount = sumOf $ rsCommon . liftTraversal (_Flip . folded)

legalTiers
    :: forall tiers epic. ()
    => ResearchData tiers epic
    -> ResearchStatus tiers epic
    -> Prod (C Bool) tiers
legalTiers rd rs = rd ^. rdCommon . to (map1 go)
  where
    tot = researchCount rs
    go  :: ResearchTier a
        -> C Bool a
    go = view $ rtUnlock . to (>= tot) . _Unwrapped

data ResearchError = REMaxedOut
                   | RELocked
  deriving (Show, Eq, Ord)

makePrisms ''ResearchError

-- | Purchase research at a given index, incrementing the counter in the
-- 'ResearchStatus'.  Returns 'Nothing' if research is already maxed out.
purchaseResearch
    :: forall tiers epic a. (KnownNat epic, SingI tiers)
    => ResearchData tiers epic
    -> ResearchIx tiers epic a
    -> ResearchStatus tiers epic
    -> Either ResearchError (a, ResearchStatus tiers epic)
purchaseResearch rd i rs0 = pp . getComp . researchIxStatusLegal rd i (Comp . go) $ rs0
  where
    bs = totalBonuses rd rs0
    pp :: Maybe (First a, ResearchStatus tiers epic)
        -> Either ResearchError (a, ResearchStatus tiers epic)
    pp Nothing                     = Left REMaxedOut
    pp (Just (First Nothing , _ )) = Left RELocked
    pp (Just (First (Just c), rs)) = case i of
        RICommon _ -> Right (c ^. bonusingFor bs BTResearchCosts, rs)
        RIEpic   _ -> Right (c, rs)
    go :: Natural -> Maybe (First a, Natural)
    go currLevel = first (First . Just) <$>
      rd ^? researchIxData i
          . rCosts
          . ix (fromIntegral currLevel)
          . to ((, currLevel + 1) . fromMaybe (0 \\ researchIxNum i))

-- | Get a 'Num' instance from a 'ResearchIx'.
researchIxNum :: ResearchIx tiers epic a -> Wit (Num a)
researchIxNum = \case
    RICommon _ -> Wit
    RIEpic   _ -> Wit

-- | A lens into a 'ResearchData', given the appropriate index.
researchIxData
    :: (KnownNat epic, SingI tiers)
    => ResearchIx tiers epic a
    -> Lens' (ResearchData tiers epic) (Research a)
researchIxData = \case
    RICommon i -> \f ->
      let g :: forall a. _ a -> _ (_ a)
          g x@(slot :&: (_ :&: SNat)) = (_2 . _1 . rtTechs . ixSV slot) f x
      in  rdCommon $ \rs -> map1 fanFst . fanSnd
            <$> sumProd g (i :&: zipP (rs :&: toTC sing))
    RIEpic i   -> rdEpic . ixSV i

-- | A lens into a 'ResearchStatus', given the appropriate index.
researchIxStatus
    :: (KnownNat epic, SingI tiers)
    => ResearchIx tiers epic a
    -> Lens' (ResearchStatus tiers epic) Natural
researchIxStatus = \case
    RICommon i -> \f ->
      let g :: forall a. _ a -> _ (_ a)
          g x@(slot :&: (_ :&: SNat)) = (_2 . _1 . _Flip . ixSV slot) f x
      in  rsCommon $ \rs -> map1 fanFst . fanSnd
            <$> sumProd g (i :&: zipP (rs :&: toTC sing))
    RIEpic i   -> rsEpic . ixSV i

-- | Lens into both the research data and research status common slots
-- together.
rdrsCommon
    :: Lens' ((Uncur ResearchData :&: Uncur ResearchStatus) (tiers # epic))
             (Prod (ResearchTier :&: Flip SV.Vector Natural) tiers)
rdrsCommon f (Uncur rd :&: Uncur rs) =
    f (zipP (_rdCommon rd :&: _rsCommon rs)) <&> \(unzipP->(rdc :&: rsc)) ->
      Uncur (rd & rdCommon .~ rdc) :&: Uncur (rs & rsCommon .~ rsc)

-- | Lens into both the research data and research status epic slots
-- together.
rdrsEpic
    :: KnownNat epic
    => Lens' ((Uncur ResearchData :&: Uncur ResearchStatus) (tiers # epic))
             (SV.Vector epic (Research GoldenEgg, Natural))
rdrsEpic f (Uncur rd :&: Uncur rs) =
    f (liftA2 (,) (_rdEpic rd) (_rsEpic rs)) <&> \rdrs ->
      Uncur (rd & rdEpic .~ fmap fst rdrs) :&: Uncur (rs & rsEpic .~ fmap snd rdrs)

-- | Lens into both a ResearchData and ResearchStatus together, from
-- a given index.
researchIx
    :: (SingI tiers, KnownNat epic)
    => ResearchIx tiers epic a
    -> Lens' ((Uncur ResearchData :&: Uncur ResearchStatus) (tiers # epic))
             (Research a, Natural)
researchIx = \case
    RICommon i -> \f ->
      let g :: forall a. _ a
            -> _ (_ a)
          g (slot :&: ((c :&: Flip e) :&: SNat)) =
              f (SV.index (_rtTechs c) slot, SV.index e slot) <&> \(d, s) ->
                let c' = c & rtTechs . ixSV slot .~ d
                    e' = e & ixSV slot .~ s
                in  slot :&: ((c' :&: Flip e') :&: SNat)
      in  rdrsCommon $ \rdrs ->
            map1 fanFst . fanSnd <$> sumProd g (i :&: zipP (rdrs :&: toTC sing))
    RIEpic i -> rdrsEpic . ixSV i

withSomeResearch
    :: DSum Sing (Uncur res)
    -> (forall tiers epic. (KnownNat epic, SingI tiers) => res tiers epic -> r)
    -> r
withSomeResearch = \case
    STuple2 sTs SNat :=> Uncur r -> \f -> withSingI sTs $ f r

-- | Traversal into a given index of a 'ResearchStatus' if the item is in
-- a legal tier.
--
-- Only a legal traversal if the mapping function doesn't change the legal
-- status.
researchIxStatusLegal
    :: (SingI tiers, KnownNat epic)
    => ResearchData tiers epic
    -> ResearchIx tiers epic a
    -> Traversal' (ResearchStatus tiers epic) Natural
researchIxStatusLegal rd = \case
    RICommon i -> \f rs0 -> getUncur . fanSnd <$>
      let totCount = researchCount rs0
          g :: forall a. _ a -> _ (_ a)
          g x@(slot :&: ((rt :&: _) :&: SNat))
            | _rtUnlock rt <= totCount = (_2 . _1 . _2 . _Flip . ixSV slot) f x
            | otherwise                = pure x
      in  (Uncur rd :&: Uncur rs0) & rdrsCommon %%~ \rtrs ->
            map1 fanFst . fanSnd <$> sumProd g (i :&: zipP (rtrs :&: toTC sing))
    RIEpic i   -> rsEpic . ixSV i

-- | All 'ResearchIx' into common researches.
researchIxesCommon
    :: SingI tiers
    => Prod (Flip SV.Vector (ResearchIx tiers epic Bock)) tiers
researchIxesCommon = go sing
  where
    go :: Sing ts -> Prod (Flip SV.Vector (ResearchIx ts epic Bock)) ts
    go = \case
      SNil         -> Ø
      SNat `SCons` ss ->
        let rest = map1 (over (_Flip . mapped . _RICommon) InR) (go ss)
        in  Flip (SV.generate (RICommon . InL)) :< rest

-- | All 'ResearchIx' into epic researches.
researchIxesEpic
    :: KnownNat epic
    => SV.Vector epic (ResearchIx tiers epic GoldenEgg)
researchIxesEpic = SV.generate RIEpic

-- | All legal 'ResearchIx' for common research.
legalResearchIxesCommon
    :: forall tiers epic. SingI tiers
    => ResearchData tiers epic
    -> ResearchStatus tiers epic
    -> Prod (Flip SV.Vector (Either ResearchError (ResearchIx tiers epic Bock))) tiers
legalResearchIxesCommon rd rs =
    imap1 (\i -> Flip . go i) (zipP (zipP (_rdCommon rd :&: _rsCommon rs) :&: toTC sing))
  where
    totCount = researchCount rs
    go  :: forall t. ()
        => Index tiers t
        -> ((ResearchTier :&: Flip SV.Vector Natural) :&: Sing) t
        -> SV.Vector t (Either ResearchError (ResearchIx tiers epic Bock))
    go i ((rt :&: Flip c) :&: SNat)
      | rt ^. rtUnlock > totCount = pure (Left RELocked)
      | otherwise                 =
          let mkIx
                  :: Finite t
                  -> Research Bock
                  -> Natural
                  -> Either ResearchError (ResearchIx tiers epic Bock)
              mkIx j r n
                | n < maxLevel r = Right . RICommon . someSum $ Some (i :&: j)
                | otherwise      = Left REMaxedOut
          in  SV.izipWith mkIx (rt ^. rtTechs) c

-- | All legal 'ResearchIx' for epic research.
--
-- 'Nothing' implies research is maxed-out.
legalResearchIxesEpic
    :: forall tiers epic. KnownNat epic
    => ResearchData tiers epic
    -> ResearchStatus tiers epic
    -> (SV.Vector epic :.: Maybe) (ResearchIx tiers epic GoldenEgg)
legalResearchIxesEpic rd rs = Comp $ SV.izipWith go (_rdEpic rd) (_rsEpic rs)
  where
    go  :: Finite epic
        -> Research GoldenEgg
        -> Natural
        -> Maybe (ResearchIx tiers epic GoldenEgg)
    go i r n
      | n < maxLevel r = Just . RIEpic $ i
      | otherwise      = Nothing

-- | All legal common researches.
legalResearchesCommon
    :: forall tiers epic. SingI tiers
    => ResearchData tiers epic
    -> ResearchStatus tiers epic
    -> Prod (Flip SV.Vector (Either ResearchError Bock)) tiers
legalResearchesCommon rd rs =
    map1 (Flip . go) (zipP (zipP (_rdCommon rd :&: _rsCommon rs) :&: toTC sing))
  where
    bs       = totalBonuses rd rs
    totCount = researchCount rs
    go  :: forall t. ()
        => ((ResearchTier :&: Flip SV.Vector Natural) :&: Sing) t
        -> SV.Vector t (Either ResearchError Bock)
    go ((rt :&: Flip c) :&: SNat)
      | rt ^. rtUnlock > totCount = pure (Left RELocked)
      | otherwise                 = do
          r         <- rt ^. rtTechs
          currLevel <- c
          pure $ case r ^? rCosts . ix (fromIntegral currLevel) of
            Nothing -> Left REMaxedOut
            Just b  -> Right $ fromMaybe 0 b ^. bonusingFor bs BTResearchCosts

-- | All legal epic researches.
--
-- 'Nothing' implies research is maxed-out.
legalResearchesEpic
    :: forall tiers epic. KnownNat epic
    => ResearchData tiers epic
    -> ResearchStatus tiers epic
    -> (SV.Vector epic :.: Maybe) GoldenEgg
legalResearchesEpic rd rs = Comp $ do
    r         <- _rdEpic rd
    currLevel <- _rsEpic rs
    pure $ fromMaybe 0 <$> r ^? rCosts . ix (fromIntegral currLevel)

max_stars_repo_name: mstksg/eggvisor

max_issues_repo_name: mstksg/eggvisor

max_forks_repo_name: mstksg/eggvisor