{-# LANGUAGE OverloadedStrings #-}

module Numeric.PHCPack.Parse (
  parseSolutions
  )
  where

import Data.Attoparsec
import Data.Attoparsec.ByteString.Char8 hiding (satisfy)
import Data.Attoparsec.Combinator
import Control.Applicative((<|>))

import Data.ByteString(ByteString)

import Data.ByteString.Char8(pack)
import qualified Data.ByteString as BS

import Numeric.PHCPack.Types
import qualified Data.Map as Map
import Data.Complex

space' :: Parser ()
space' = satisfy isHorizontalSpace >> return ()

spaces :: Parser ()
spaces = skipMany space' >> return ()

spaces1 :: Parser ()
spaces1 = space' >> spaces

parseSolutions :: ByteString -> Maybe [Solution]
parseSolutions = maybeResult . ((flip feed) BS.empty) . (parse solutionsParser)

solutionsParser :: Parser [Solution]
solutionsParser = do
  manyTill anyChar (try (string "START SOLUTIONS : " >> endOfLine))
  skipMany (space' <|> endOfLine)
  solutionCount <- decimal
  spaces1
  dimension <- decimal
  endOfLine

  sepLine
  solutions <- solution `sepBy1` bannerLine
  bannerLine
  skipMany anyChar

  return solutions

sepLine :: Parser ()
sepLine = skipMany (char '=') >> endOfLine

bannerLine :: Parser ()
bannerLine = string "==" >> manyTill anyChar (try endOfLine) >> return ()

solution :: Parser Solution
solution = do
  string "solution "
  solutionNumber <- decimal
  manyTill anyChar (try (string "the solution" >> manyTill anyChar (try endOfLine))) -- skip until "the solution for t :"
  solutions <- variableSolution `sepBy` endOfLine
  endOfLine
  return (Solution $ Map.fromList solutions)

-- variableSolution :: (Num c) => Parser (Unknown, Complex c)
variableSolution = do
  spaces
  name <- manyTill anyChar space
  spaces
  char ':'
  spaces
  real <- double
  spaces
  imag <- double
  return (name, real :+ imag)


{-# LANGUAGE DataKinds           #-}
{-# LANGUAGE PolyKinds           #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeApplications    #-}
{-# LANGUAGE ViewPatterns        #-}

import           Control.Concurrent
import           Control.DeepSeq
import           Control.Monad
import           Control.Monad.Primitive
import           Control.Monad.Trans.Class
import           Control.Monad.Trans.State
import           Data.Bifunctor
import           Data.Char
import           Data.Finite
import           Data.Foldable
import           Data.List
import           Data.Maybe
import           Data.Neural.Activation
import           Data.Neural.HMatrix.Recurrent
import           Data.Neural.HMatrix.Recurrent.Dropout
import           Data.Neural.HMatrix.Recurrent.Generate
import           Data.Neural.HMatrix.Recurrent.Train
import           Data.Ord
import           Data.Proxy
import           Data.Tuple
import           GHC.TypeLits
import           GHC.TypeLits.List
import           Numeric.LinearAlgebra.Static
import           Text.Printf
import qualified Data.Map.Strict                        as M
import qualified Data.Set                               as S
import qualified Data.Vector                            as VB
import qualified Data.Vector.Storable                   as V
import qualified Linear.V                               as L
import qualified System.Random.MWC                      as R
import qualified System.Random.MWC.Distributions        as R

genIx :: forall n. KnownNat n => Finite n -> R n
genIx (fromIntegral->i) = fromJust
                        . create
                        $ V.generate l (\i' -> if i' == i then 1 else 0)
  where
    l = fromInteger $ natVal (Proxy @n)

ixMax :: (Foldable t, Ord a) => t a -> Int
ixMax = fst . maximumBy (comparing snd) . zip [0..] . toList

ixSort :: (Foldable t, Ord a) => t a -> [Int]
ixSort = map fst . sortBy (flip (comparing snd)) . zip [0..] . toList

sanitize :: Char -> Char
sanitize c | isPrint c = c
           | otherwise = '#'

main :: IO ()
main = do
    g <- R.create
    holmes <- readFile "data/holmes.txt"
    let holmesVec :: VB.Vector Char
        holmesVec = VB.fromList holmes
    let allCharsV :: VB.Vector Char
        allCharsV = VB.fromList
                  . S.toList
                  . S.fromList
                  $ holmes

    L.reifyVectorNat allCharsV $ \(allChars :: L.V c Char) -> do
      let charMap :: M.Map Char (Finite c)
          charMap = M.fromList
                  . (`zip` [0..])
                  . toList
                  $ allChars
          c2v     :: Char -> R c
          c2v     = genIx . (charMap M.!)
          holmesSeries :: VB.Vector (L.V 15 Char, Char)
          holmesSeries = processSeries $ VB.zip holmesVec (VB.tail holmesVec)
          holmesSeries' :: VB.Vector ([R c], R c)
          holmesSeries' = bimap (map c2v . toList) c2v <$> holmesSeries
          ins :: VB.Vector [R c]
          ins = fst <$> holmesSeries'
          nextChar :: R c -> IO Int
          nextChar i = R.categorical (V.map ((**8) . max 0) (extract i)) g
          fb :: R c -> IO (R c)
          fb = fmap (genIx . fromIntegral) . nextChar
          cb :: Int -> Int -> Network c hs c -> IO ()
          cb e i n = when (i `mod` 1000 == 0) $ do
            let toCharList = map (sanitize . (allCharsV VB.!))
                           . take 50 . ixSort . V.toList . extract
                insTest   = ins VB.! 100
                insChars  = head . toCharList <$> insTest
                (last->lo, pp) = runNetStream naRLLog n (ins VB.! 100)
            lo' <- fb lo
            testOut   <- (lo:) <$> runNetFeedbackM_ naRLLog fb pp 75 lo'
            let testChars = toCharList <$> testOut
            threadDelay 250000
            printf "%d\t%d\n" e i
            mapM_ putStrLn (take 15 testChars)
            testCharsPick <- map (sanitize . (allCharsV VB.!)) <$> mapM nextChar testOut
            putStrLn $ insChars ++ "|" ++ testCharsPick

      print $ M.keys charMap
      print $ natVal (Proxy @c)
      print $ length holmes

      net0 <- randomNetMWC (-0.1,0.1) g :: IO (Network c '[100,75,50] c)
      net1 <- trainHistory 0.2 0.01 0.005 100 holmesSeries' net0 cb g
      return ()

trainHistory
    :: forall i m hs o. (PrimMonad m, KnownNats hs, KnownNat i, KnownNat o)
    => Double       -- ^ dropout
    -> Double       -- ^ step size (weights)
    -> Double       -- ^ step size (state)
    -> Int          -- ^ number of epochs
    -> VB.Vector ([R i], R o)
    -> Network i hs o
    -> (Int -> Int -> Network i hs o -> m ())  -- ^ callback
    -> R.Gen (PrimState m)
    -> m (Network i hs o)
trainHistory d sw ss e samps net0 cb g =
    flip execStateT net0
  . for_ [1 .. e] $ \i -> do
      samps' <- lift $ R.uniformShuffle samps g
      for_ (zip [1..] (VB.toList samps')) $ \(j, (h, t)) ->
        StateT $ \net -> do
          net' <- trainSeriesDOMWC naRLLog d sw ss t h net g
          cb i j net'
          net' `deepseq` return ((), net')


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

-- |
-- Module    : Statistics.Resampling.Bootstrap
-- Copyright : (c) 2009, 2011 Bryan O'Sullivan
-- License   : BSD3
--
-- Maintainer  : bos@serpentine.com
-- Stability   : experimental
-- Portability : portable
--
-- The bootstrap method for statistical inference.

module Statistics.Resampling.Bootstrap
    (
      Estimate(..)
    , bootstrapBCA
    , scale
    -- * References
    -- $references
    ) where

import Control.Applicative ((<$>), (<*>))
import Control.DeepSeq (NFData)
import Control.Exception (assert)
import Control.Monad.Par (parMap, runPar)
import Data.Aeson (FromJSON, ToJSON)
import Data.Binary (Binary)
import Data.Binary (put, get)
import Data.Data (Data)
import Data.Typeable (Typeable)
import Data.Vector.Unboxed ((!))
import GHC.Generics (Generic)
import Statistics.Distribution (cumulative, quantile)
import Statistics.Distribution.Normal
import Statistics.Resampling (Resample(..), jackknife)
import Statistics.Sample (mean)
import Statistics.Types (Estimator, Sample)
import qualified Data.Vector.Unboxed as U
import qualified Statistics.Resampling as R

-- | A point and interval estimate computed via an 'Estimator'.
data Estimate = Estimate {
      estPoint           :: {-# UNPACK #-} !Double
    -- ^ Point estimate.
    , estLowerBound      :: {-# UNPACK #-} !Double
    -- ^ Lower bound of the estimate interval (i.e. the lower bound of
    -- the confidence interval).
    , estUpperBound      :: {-# UNPACK #-} !Double
    -- ^ Upper bound of the estimate interval (i.e. the upper bound of
    -- the confidence interval).
    , estConfidenceLevel :: {-# UNPACK #-} !Double
    -- ^ Confidence level of the confidence intervals.
    } deriving (Eq, Read, Show, Typeable, Data, Generic)

instance FromJSON Estimate
instance ToJSON Estimate

instance Binary Estimate where
    put (Estimate w x y z) = put w >> put x >> put y >> put z
    get = Estimate <$> get <*> get <*> get <*> get
instance NFData Estimate

-- | Multiply the point, lower bound, and upper bound in an 'Estimate'
-- by the given value.
scale :: Double                 -- ^ Value to multiply by.
      -> Estimate -> Estimate
scale f e@Estimate{..} = e {
                           estPoint = f * estPoint
                         , estLowerBound = f * estLowerBound
                         , estUpperBound = f * estUpperBound
                         }

estimate :: Double -> Double -> Double -> Double -> Estimate
estimate pt lb ub cl =
    assert (lb <= ub) .
    assert (cl > 0 && cl < 1) $
    Estimate { estPoint = pt
             , estLowerBound = lb
             , estUpperBound = ub
             , estConfidenceLevel = cl
             }

data T = {-# UNPACK #-} !Double :< {-# UNPACK #-} !Double
infixl 2 :<

-- | Bias-corrected accelerated (BCA) bootstrap. This adjusts for both
-- bias and skewness in the resampled distribution.
bootstrapBCA :: Double          -- ^ Confidence level
             -> Sample          -- ^ Sample data
             -> [Estimator]     -- ^ Estimators
             -> [Resample]      -- ^ Resampled data
             -> [Estimate]
bootstrapBCA confidenceLevel sample estimators resamples
  | confidenceLevel > 0 && confidenceLevel < 1
      = runPar $ parMap (uncurry e) (zip estimators resamples)
  | otherwise = error "Statistics.Resampling.Bootstrap.bootstrapBCA: confidence level outside (0,1) range"
  where
    e est (Resample resample)
      | U.length sample == 1 || isInfinite bias =
          estimate pt pt pt confidenceLevel
      | otherwise =
          estimate pt (resample ! lo) (resample ! hi) confidenceLevel
      where
        pt    = R.estimate est sample
        lo    = max (cumn a1) 0
          where a1 = bias + b1 / (1 - accel * b1)
                b1 = bias + z1
        hi    = min (cumn a2) (ni - 1)
          where a2 = bias + b2 / (1 - accel * b2)
                b2 = bias - z1
        z1    = quantile standard ((1 - confidenceLevel) / 2)
        cumn  = round . (*n) . cumulative standard
        bias  = quantile standard (probN / n)
          where probN = fromIntegral . U.length . U.filter (<pt) $ resample
        ni    = U.length resample
        n     = fromIntegral ni
        accel = sumCubes / (6 * (sumSquares ** 1.5))
          where (sumSquares :< sumCubes) = U.foldl' f (0 :< 0) jack
                f (s :< c) j = s + d2 :< c + d2 * d
                    where d  = jackMean - j
                          d2 = d * d
                jackMean     = mean jack
        jack  = jackknife est sample

-- $references
--
-- * Davison, A.C; Hinkley, D.V. (1997) Bootstrap methods and their
--   application. <http://statwww.epfl.ch/davison/BMA/>


--
-- CNN : CNN filter generator
--
{-# LANGUAGE BangPatterns #-}

module Main (
  main
) where

import Control.DeepSeq
import Control.Monad
--import Data.List
import Data.Time
import Data.List (foldl')
import Debug.Trace
import Numeric.LinearAlgebra
import System.Environment
import System.IO
import Text.Printf

import CNN.ActLayer
import CNN.Algebra
import CNN.ConvLayer
import CNN.FullConnLayer
import CNN.Image
import CNN.Layer
import CNN.LayerType
import CNN.PoolLayer
import CNN.Propagation

import Pool
import Status

usage :: String
usage = "Usage: train <dir> [<image>]"

-- MAIN

main :: IO ()
main = do
  as <- getArgs
  case length as of
    1 -> train (as !! 0)
    2 -> judge (as !! 0) (as !! 1)
    _ -> error usage

train :: String -> IO ()
train dn = do
  putStrLn "Initializing..."
  st <- loadStatus dn
--  st <- loadStatus ""
  tm0 <- getCurrentTime

  let
    --getTeachers = getImages (poolT st) (batchSz st)
    --getTests    = getImages (poolE st) (testSz st)
    getTeachers = getImagesRandomly (poolT st) (batchSz st)
    getTests    = getImagesRandomly (poolE st) (testSz st)
    putF = putStatus tm0 st getTests
    loopFunc = trainLoop' getTeachers putF st
  putStrLn ("Training the model... (#batch:" ++ (show (batchSz st * nclass st)) ++ ")")
  putF 0 (layers st)
  layers' <- foldM' loopFunc (layers st) [1 .. (repeatCt st)]

  putStrLn "Saving status..."
  saveStatus st layers' (repeatCt st * batchSz st)
  putStrLn "Finished!"

judge :: String -> String -> IO ()
judge dn imf = do
  st <- loadStatus dn
  fh <- openFile imf ReadMode
  lst <- hGetContents fh
  forM_ (lines lst) $ \i -> do
    im <- loadImage i
    putStr (i ++ ",")
    let (y, _) = judgeImage (layers st) im
    forM_ (zip [0..] (toList y)) $ \(c, v) -> do
      let res = show c ++ (printf ":%.2f," (v*100))
      putStr res
    putStrLn ""

--
-- FUNCTIONS
--

{- |
foldM' : monadic version of foldl'

  URL: http://stackoverflow.com/questions/8919026/does-haskell-have-foldlm

-}

foldM' :: (Monad m) => ([a] -> b -> m [a]) -> [a] -> [b] -> m [a]
foldM' _ z [] = return z
foldM' f z (x:xs) = do
  z' <- f z x
  z' `seq` foldM' f z' xs

{- |
trainLoop

  IN : func of getting teachers (including batch size and pool)
       sample of evaluation
       output func
       learning rate
       layers
       epoch numbers

  OUT: updated layers

-}

trainLoop' :: (Int -> IO [Trainer]) -> (Int -> [Layer] -> IO ()) -> Status
           -> [Layer] -> Int -> IO [Layer]
trainLoop' getT putF st ls i = do
  --teachers <- getT ((i-1) * batchSz st)
  teachers <- getT 0
  let ls' = updateLayers (learnR st) teachers ls
  if i `mod` (savePt st) == 0
    then putF i ls'
    else return ()
  return ls'

{- |
updateLayers
-}

updateLayers :: Double -> [Trainer] -> [Layer] -> [Layer]
updateLayers lr ts ls = update lr ls (transpose dls)
  where
    rls = tail $ map reverseLayer $ reverse ls    -- fist element isn't used
    dls = map (trainLayers ls rls) ts             -- dls = diff of layers

{-
putStatus

  IN : start time
       status
       start epoch number
       end epoch number
       step size of status output
       epoch
       layers
       sample of evaluation

-}

putStatus :: UTCTime -> Status -> (Int -> IO [Trainer]) -> Int -> [Layer] -> IO ()
putStatus tm0 st getE i ls = do
  --tests <- getE ((i-1) * testSz st)
  tests <- getE 0
  let
    --(rv, rr) = unzip $ evaluate ls tests
    rr = foldl' (evaluateLayers ls) 0.0 tests
  let
    ite = printf "iter = %5d/%d " i (repeatCt st)
    acc = printf "accuracy = %.10f " (rr / fromIntegral (length tests))
  putStr (ite ++ acc)
  tm <- getCurrentTime
  putStrLn ("time = " ++ show (diffUTCTime tm tm0))
  saveStatus st ls (batchSz st * i)
  --mapM_ putOne rs   -- for debug
  where
    putOne :: ([Double], Double) -> IO ()
    putOne (v, r) = putStrLn ("result:" ++ show v ++ ", ratio:" ++ show r)


module ChartHistogram (
  backendFormats,
  writeHistogram
) where

import           Control.Lens.Setter ((.~))
import           Control.Monad.ST (runST)
import qualified Data.Colour as Colour
import qualified Data.Colour.Names as ColourNames
import           Data.Default.Class
import qualified Data.Vector.Algorithms.Intro as VI
import qualified Data.Vector.Generic as VG
import           Data.Vector.Unboxed ((!))
import qualified Graphics.Rendering.Chart as Chart
import           Statistics.Test.ApproxRand
import           Statistics.Test.Types (TestType(..))
import           System.IO (hPutStrLn, stderr)

import           Histogram
import           ChartBackend

writeHistogram :: TestOptions -> Int -> TestResult -> FilePath -> IO ()
writeHistogram testOptions bins result path =
  case histogram bins result of
    Left err -> hPutStrLn stderr err
    Right  h ->
      writeWithBackend (createHistogram testOptions result h) path

-- Creates a histogram. The histogram is stacked, but the second bar
-- is always empty, except for the bin of the original statistic (if any).
-- There, the first bar is empty and the second bar has the frequency.
-- Yes, this is cheating ;).
createHistogram :: TestOptions -> TestResult -> [(Double, Int)] ->
  Chart.Renderable ()
createHistogram testOptions result his =
  Chart.toRenderable layout
  where
    layout =
        Chart.layout_background    .~ Chart.solidFillStyle opaqueWhite
      $ Chart.layout_y_axis        .  Chart.laxis_title    .~ "Frequency"
      $ Chart.layout_x_axis        .  Chart.laxis_title    .~ "Statistic"
      $ Chart.layout_plots         .~ [ Chart.plotBars randomizationBars,
                                        statisticLine, sigLines ]
      $ Chart.setLayoutForeground   opaqueBlack
      $ def
    randomizationBars =
        Chart.plot_bars_style       .~ Chart.BarsStacked
      $ Chart.plot_bars_spacing     .~ Chart.BarsFixGap 6 2
      -- $ Chart.plot_bars_spacing     ^= Chart.BarsFixGap 0 0
      $ Chart.plot_bars_item_styles .~ [
          (Chart.solidFillStyle $ opaqueGreen, Nothing) ]
      $ Chart.plot_bars_values      .~ map (\(b, f) -> (b, [f])) his
      $ def
    statisticLine =
      Chart.vlinePlot "Statistic for samples" (Chart.solidLine 2 (opaqueRed)) $ trStat result
    sigLines      =
      vlinesPlot "Significance" (Chart.dashedLine 2 [8, 4] opaqueBlack) $
        sigBounds testOptions result

-- Plot vertical lines, adapted from Chart.vlinePlot.
vlinesPlot :: String -> Chart.LineStyle -> [a] -> Chart.Plot a b
vlinesPlot t ls xs = Chart.toPlot vlines
  where
    vlines =
        Chart.plot_lines_title        .~ t
      $ Chart.plot_lines_style        .~ ls
      $ Chart.plot_lines_limit_values .~ minMax
      $ def
    minMax =
      [[(Chart.LValue v, Chart.LMin), (Chart.LValue v, Chart.LMax)] | v <- xs]

-- Calculate the bounds of significance.
sigBounds :: TestOptions -> TestResult -> [Double]
sigBounds (TestOptions testType _ n pTest) (TestResult _ _ stats) =
  case testType of
    TwoTailed -> [sorted ! (nExtreme - 1), sorted ! (n - nExtreme)]
    OneTailed -> [sorted ! (n - nExtreme)]
  where
    sorted           = sortVector stats
    nExtreme         = floor $ (pVal testType pTest) * (fromIntegral n + 1) - 1
    pVal OneTailed p = p
    pVal TwoTailed p = p / 2
    -- XXX: Fix extreme cases: p-value of 0, small n.

sortVector :: (Ord a, VG.Vector v a) => v a -> v a
sortVector v = runST $ do
  s <- VG.thaw v
  VI.sort s
  VG.freeze s

-- Convenience...
opaqueBlack, opaqueGreen, opaqueRed, opaqueWhite :: Colour.AlphaColour Double
opaqueBlack = Colour.opaque ColourNames.black
opaqueGreen = Colour.opaque ColourNames.green
opaqueRed   = Colour.opaque ColourNames.red
opaqueWhite = Colour.opaque ColourNames.white


-- This file is auto-generated.  Do not edit directly.
-- |
--   Stability: Experimental
--
--   Generic interface to Blas using unsafe foreign calls.  Refer to the GHC
--   documentation for more information regarding appropriate use of safe and unsafe
--   foreign calls.
--
--   The functions here are named in a similar fashion to the original Blas interface, with
--   the type-dependent letter(s) removed.  Some functions have been merged with
--   others to allow the interface to work on both real and complex numbers.  If you can't a
--   particular function, try looking for its corresponding complex equivalent (e.g.
--   @symv@ is a special case of 'hemv' applied to real numbers).
--
--   Note: although complex versions of @rot@ and @rotg@ exist in many implementations,
--   they are not part of the official Blas standard and therefore not included here.  If
--   you /really/ need them, submit a ticket so we can try to come up with a solution.
--
--   The documentation here is still incomplete.  Consult the
--   <http://netlib.org/blas/#_blas_routines official documentation> for more
--   information.
--
--   Notation:
--
--     * @⋅@ denotes dot product (without any conjugation).
--     * @*@ denotes complex conjugation.
--     * @⊤@ denotes transpose.
--     * @†@ denotes conjugate transpose (Hermitian conjugate).
--
--   Conventions:
--
--     * All scalars are denoted with lowercase Greek letters
--     * All vectors are denoted with lowercase Latin letters and are
--       assumed to be column vectors (unless transposed).
--     * All matrices are denoted with uppercase Latin letters.
--
--   /Since: 0.1.1/
module Blas.Specialized.ComplexDouble.Unsafe (
  -- * Level 1: vector-vector operations
  -- ** Basic operations
    swap
  , scal
  , copy
  , axpy
  , dotu
  , dotc
  -- ** Norm operations
  , nrm2
  , asum
  , iamax
  -- * Level 2: matrix-vector operations
  -- ** Multiplication
  , gemv
  , gbmv
  , hemv
  , hbmv
  , hpmv
  -- ** Triangular operations
  , trmv
  , tbmv
  , tpmv
  , trsv
  , tbsv
  , tpsv
  -- ** Rank updates
  , geru
  , gerc
  , her
  , hpr
  , her2
  , hpr2
  -- * Level 3: matrix-matrix operations
  -- ** Multiplication
  , gemm
  , symm
  , hemm
  -- ** Rank updates
  , syrk
  , herk
  , syr2k
  , her2k
  -- ** Triangular operations
  , trmm
  , trsm
  ) where
import Prelude (Double, Int, IO)
import Foreign (Ptr)
import Data.Complex (Complex((:+)))
import FFI (getReturnValue)
import Blas.Primitive.Types (Order, Transpose, Uplo, Diag, Side)
import qualified Blas.Primitive.Unsafe as C

-- | Swap two vectors:
--
--   > (x, y) ← (y, x)
swap :: Int
     -> Ptr (Complex Double)
     -> Int
     -> Ptr (Complex Double)
     -> Int
     -> IO ()
swap = C.zswap

-- | Multiply a vector by a scalar.
--
--   > x ← α x
scal :: Int
     -> Complex Double
     -> Ptr (Complex Double)
     -> Int
     -> IO ()
scal n (alpha :+ 0) = C.zdscal n alpha
scal n alpha = C.zscal n alpha

-- | Copy a vector into another vector:
--
--   > y ← x
copy :: Int
     -> Ptr (Complex Double)
     -> Int
     -> Ptr (Complex Double)
     -> Int
     -> IO ()
copy = C.zcopy

-- | Add a scalar-vector product to a vector.
--
--   > y ← α x + y
axpy :: Int
     -> Complex Double
     -> Ptr (Complex Double)
     -> Int
     -> Ptr (Complex Double)
     -> Int
     -> IO ()
axpy = C.zaxpy

-- | Calculate the bilinear dot product of two vectors:
--
--   > x ⋅ y ≡ ∑[i] x[i] y[i]
dotu :: Int
     -> Ptr (Complex Double)
     -> Int
     -> Ptr (Complex Double)
     -> Int
     -> IO (Complex Double)
dotu a b c d e = getReturnValue (C.zdotu_sub a b c d e)

-- | Calculate the sesquilinear dot product of two vectors.
--
--   > x* ⋅ y ≡ ∑[i] x[i]* y[i]
dotc :: Int
     -> Ptr (Complex Double)
     -> Int
     -> Ptr (Complex Double)
     -> Int
     -> IO (Complex Double)
dotc a b c d e = getReturnValue (C.zdotc_sub a b c d e)

-- | Calculate the Euclidean (L²) norm of a vector:
--
--   > ‖x‖₂ ≡ √(∑[i] x[i]²)
nrm2 :: Int
     -> Ptr (Complex Double)
     -> Int
     -> IO Double
nrm2 = C.dznrm2

-- | Calculate the Manhattan (L¹) norm, equal to the sum of the magnitudes of the elements:
--
--   > ‖x‖₁ = ∑[i] |x[i]|
asum :: Int
     -> Ptr (Complex Double)
     -> Int
     -> IO Double
asum = C.dzasum

-- | Calculate the index of the element with the maximum magnitude (absolute value).
iamax :: Int
      -> Ptr (Complex Double)
      -> Int
      -> IO Int
iamax = C.izamax

-- | Perform a general matrix-vector update.
--
--   > y ← α T(A) x + β y
gemv :: Order
     -> Transpose
     -> Int
     -> Int
     -> Complex Double
     -> Ptr (Complex Double)
     -> Int
     -> Ptr (Complex Double)
     -> Int
     -> Complex Double
     -> Ptr (Complex Double)
     -> Int
     -> IO ()
gemv = C.zgemv

-- | Perform a general banded matrix-vector update.
--
--   > y ← α T(A) x + β y
gbmv :: Order
     -> Transpose
     -> Int
     -> Int
     -> Int
     -> Int
     -> Complex Double
     -> Ptr (Complex Double)
     -> Int
     -> Ptr (Complex Double)
     -> Int
     -> Complex Double
     -> Ptr (Complex Double)
     -> Int
     -> IO ()
gbmv = C.zgbmv

-- | Perform a hermitian matrix-vector update.
--
--   > y ← α A x + β y
hemv :: Order
     -> Uplo
     -> Int
     -> Complex Double
     -> Ptr (Complex Double)
     -> Int
     -> Ptr (Complex Double)
     -> Int
     -> Complex Double
     -> Ptr (Complex Double)
     -> Int
     -> IO ()
hemv = C.zhemv

-- | Perform a hermitian banded matrix-vector update.
--
--   > y ← α A x + β y
hbmv :: Order
     -> Uplo
     -> Int
     -> Int
     -> Complex Double
     -> Ptr (Complex Double)
     -> Int
     -> Ptr (Complex Double)
     -> Int
     -> Complex Double
     -> Ptr (Complex Double)
     -> Int
     -> IO ()
hbmv = C.zhbmv

-- | Perform a hermitian packed matrix-vector update.
--
--   > y ← α A x + β y
hpmv :: Order
     -> Uplo
     -> Int
     -> Complex Double
     -> Ptr (Complex Double)
     -> Ptr (Complex Double)
     -> Int
     -> Complex Double
     -> Ptr (Complex Double)
     -> Int
     -> IO ()
hpmv = C.zhpmv

-- | Multiply a triangular matrix by a vector.
--
--   > x ← T(A) x
trmv :: Order
     -> Uplo
     -> Transpose
     -> Diag
     -> Int
     -> Ptr (Complex Double)
     -> Int
     -> Ptr (Complex Double)
     -> Int
     -> IO ()
trmv = C.ztrmv

-- | Multiply a triangular banded matrix by a vector.
--
--   > x ← T(A) x
tbmv :: Order
     -> Uplo
     -> Transpose
     -> Diag
     -> Int
     -> Int
     -> Ptr (Complex Double)
     -> Int
     -> Ptr (Complex Double)
     -> Int
     -> IO ()
tbmv = C.ztbmv

-- | Multiply a triangular packed matrix by a vector.
--
--   > x ← T(A) x
tpmv :: Order
     -> Uplo
     -> Transpose
     -> Diag
     -> Int
     -> Ptr (Complex Double)
     -> Ptr (Complex Double)
     -> Int
     -> IO ()
tpmv = C.ztpmv

-- | Multiply an inverse triangular matrix by a vector.
--
--   > x ← T(A⁻¹) x
trsv :: Order
     -> Uplo
     -> Transpose
     -> Diag
     -> Int
     -> Ptr (Complex Double)
     -> Int
     -> Ptr (Complex Double)
     -> Int
     -> IO ()
trsv = C.ztrsv

-- | Multiply an inverse triangular banded matrix by a vector.
--
--   > x ← T(A⁻¹) x
tbsv :: Order
     -> Uplo
     -> Transpose
     -> Diag
     -> Int
     -> Int
     -> Ptr (Complex Double)
     -> Int
     -> Ptr (Complex Double)
     -> Int
     -> IO ()
tbsv = C.ztbsv

-- | Multiply an inverse triangular packed matrix by a vector.
--
--   > x ← T(A⁻¹) x
tpsv :: Order
     -> Uplo
     -> Transpose
     -> Diag
     -> Int
     -> Ptr (Complex Double)
     -> Ptr (Complex Double)
     -> Int
     -> IO ()
tpsv = C.ztpsv

-- | Perform an unconjugated rank-1 update of a general matrix.
--
--   > A ← α x y⊤ + A
geru :: Order
     -> Int
     -> Int
     -> Complex Double
     -> Ptr (Complex Double)
     -> Int
     -> Ptr (Complex Double)
     -> Int
     -> Ptr (Complex Double)
     -> Int
     -> IO ()
geru = C.zgeru

-- | Perform a conjugated rank-1 update of a general matrix.
--
--   > A ← α x y† + A
gerc :: Order
     -> Int
     -> Int
     -> Complex Double
     -> Ptr (Complex Double)
     -> Int
     -> Ptr (Complex Double)
     -> Int
     -> Ptr (Complex Double)
     -> Int
     -> IO ()
gerc = C.zgerc

-- | Perform a rank-1 update of a Hermitian matrix.
--
--   > A ← α x y† + A
her :: Order
    -> Uplo
    -> Int
    -> Double
    -> Ptr (Complex Double)
    -> Int
    -> Ptr (Complex Double)
    -> Int
    -> IO ()
her = C.zher

-- | Perform a rank-1 update of a Hermitian packed matrix.
--
--   > A ← α x y† + A
hpr :: Order
    -> Uplo
    -> Int
    -> Double
    -> Ptr (Complex Double)
    -> Int
    -> Ptr (Complex Double)
    -> IO ()
hpr = C.zhpr

-- | Perform a rank-2 update of a Hermitian matrix.
--
--   > A ← α x y† + y (α x)† + A
her2 :: Order
     -> Uplo
     -> Int
     -> Complex Double
     -> Ptr (Complex Double)
     -> Int
     -> Ptr (Complex Double)
     -> Int
     -> Ptr (Complex Double)
     -> Int
     -> IO ()
her2 = C.zher2

-- | Perform a rank-2 update of a Hermitian packed matrix.
--
--   > A ← α x y† + y (α x)† + A
hpr2 :: Order
     -> Uplo
     -> Int
     -> Complex Double
     -> Ptr (Complex Double)
     -> Int
     -> Ptr (Complex Double)
     -> Int
     -> Ptr (Complex Double)
     -> IO ()
hpr2 = C.zhpr2

-- | Perform a general matrix-matrix update.
--
--   > C ← α T(A) U(B) + β C
gemm :: Order -- ^ Layout of all the matrices.
     -> Transpose -- ^ The operation @T@ to be applied to @A@.
     -> Transpose -- ^ The operation @U@ to be applied to @B@.
     -> Int -- ^ Number of rows of @T(A)@ and @C@.
     -> Int -- ^ Number of columns of @U(B)@ and @C@.
     -> Int -- ^ Number of columns of @T(A)@ and number of rows of @U(B)@.
     -> Complex Double -- ^ Scaling factor @α@ of the product.
     -> Ptr (Complex Double) -- ^ Pointer to a matrix @A@.
     -> Int -- ^ Stride of the major dimension of @A@.
     -> Ptr (Complex Double) -- ^ Pointer to a matrix @B@.
     -> Int -- ^ Stride of the major dimension of @B@.
     -> Complex Double -- ^ Scaling factor @β@ of the original @C@.
     -> Ptr (Complex Double) -- ^ Pointer to a mutable matrix @C@.
     -> Int -- ^ Stride of the major dimension of @C@.
     -> IO ()
gemm = C.zgemm

-- | Perform a symmetric matrix-matrix update.
--
--   > C ← α A B + β C    or    C ← α B A + β C
--
--   where @A@ is symmetric.  The matrix @A@ must be in an unpacked format, although the
--   routine will only access half of it as specified by the @'Uplo'@ argument.
symm :: Order -- ^ Layout of all the matrices.
     -> Side -- ^ Side that @A@ appears in the product.
     -> Uplo -- ^ The part of @A@ that is used.
     -> Int -- ^ Number of rows of @C@.
     -> Int -- ^ Number of columns of @C@.
     -> Complex Double -- ^ Scaling factor @α@ of the product.
     -> Ptr (Complex Double) -- ^ Pointer to a symmetric matrix @A@.
     -> Int -- ^ Stride of the major dimension of @A@.
     -> Ptr (Complex Double) -- ^ Pointer to a matrix @B@.
     -> Int -- ^ Stride of the major dimension of @B@.
     -> Complex Double -- ^ Scaling factor @α@ of the original @C@.
     -> Ptr (Complex Double) -- ^ Pointer to a mutable matrix @C@.
     -> Int -- ^ Stride of the major dimension of @C@.
     -> IO ()
symm = C.zsymm

-- | Perform a Hermitian matrix-matrix update.
--
--   > C ← α A B + β C    or    C ← α B A + β C
--
--   where @A@ is Hermitian.  The matrix @A@ must be in an unpacked format, although the
--   routine will only access half of it as specified by the @'Uplo'@ argument.
hemm :: Order -- ^ Layout of all the matrices.
     -> Side -- ^ Side that @A@ appears in the product.
     -> Uplo -- ^ The part of @A@ that is used.
     -> Int -- ^ Number of rows of @C@.
     -> Int -- ^ Number of columns of @C@.
     -> Complex Double -- ^ Scaling factor @α@ of the product.
     -> Ptr (Complex Double) -- ^ Pointer to a Hermitian matrix @A@.
     -> Int -- ^ Stride of the major dimension of @A@.
     -> Ptr (Complex Double) -- ^ Pointer to a matrix @B@.
     -> Int -- ^ Stride of the major dimension of @B@.
     -> Complex Double -- ^ Scaling factor @α@ of the original @C@.
     -> Ptr (Complex Double) -- ^ Pointer to a mutable matrix @C@.
     -> Int -- ^ Stride of the major dimension of @C@.
     -> IO ()
hemm = C.zhemm

-- | Perform a symmetric rank-k update.
--
--   > C ← α A A⊤ + β C    or    C ← α A⊤ A + β C
syrk :: Order
     -> Uplo
     -> Transpose
     -> Int
     -> Int
     -> Complex Double
     -> Ptr (Complex Double)
     -> Int
     -> Complex Double
     -> Ptr (Complex Double)
     -> Int
     -> IO ()
syrk = C.zsyrk

-- | Perform a Hermitian rank-k update.
--
--   > C ← α A A† + β C    or    C ← α A† A + β C
herk :: Order
     -> Uplo
     -> Transpose
     -> Int
     -> Int
     -> Double
     -> Ptr (Complex Double)
     -> Int
     -> Double
     -> Ptr (Complex Double)
     -> Int
     -> IO ()
herk = C.zherk

-- | Perform a symmetric rank-2k update.
--
--   > C ← α A B⊤ + α* B A⊤ + β C    or    C ← α A⊤ B + α* B⊤ A + β C
syr2k :: Order
      -> Uplo
      -> Transpose
      -> Int
      -> Int
      -> Complex Double
      -> Ptr (Complex Double)
      -> Int
      -> Ptr (Complex Double)
      -> Int
      -> Complex Double
      -> Ptr (Complex Double)
      -> Int
      -> IO ()
syr2k = C.zsyr2k

-- | Perform a Hermitian rank-2k update.
--
--   > C ← α A B† + α* B A† + β C    or    C ← α A† B + α* B† A + β C
her2k :: Order
      -> Uplo
      -> Transpose
      -> Int
      -> Int
      -> Complex Double
      -> Ptr (Complex Double)
      -> Int
      -> Ptr (Complex Double)
      -> Int
      -> Double
      -> Ptr (Complex Double)
      -> Int
      -> IO ()
her2k = C.zher2k

-- | Perform a triangular matrix-matrix multiplication.
--
--   > B ← α T(A) B    or    B ← α B T(A)
--
--   where @A@ is triangular.
trmm :: Order
     -> Side
     -> Uplo
     -> Transpose
     -> Diag
     -> Int
     -> Int
     -> Complex Double
     -> Ptr (Complex Double)
     -> Int
     -> Ptr (Complex Double)
     -> Int
     -> IO ()
trmm = C.ztrmm

-- | Perform an inverse triangular matrix-matrix multiplication.
--
--   > B ← α T(A⁻¹) B    or    B ← α B T(A⁻¹)
--
--   where @A@ is triangular.
trsm :: Order
     -> Side
     -> Uplo
     -> Transpose
     -> Diag
     -> Int
     -> Int
     -> Complex Double
     -> Ptr (Complex Double)
     -> Int
     -> Ptr (Complex Double)
     -> Int
     -> IO ()
trsm = C.ztrsm


{-# LANGUAGE CPP  #-}
#if defined(__GLASGOW_HASKELL__) && __GLASGOW_HASKELL__ >= 702
{-# LANGUAGE Safe #-}
#endif
-------------------------------------------------------------------------------
-- |
-- Module      : OGLS.Engine.Math.Vectors
-- description : 3D and 4D vectors arithmetic
-- Copyright   : (c) Adam, 2017
-- License     : MIT
-- Maintainer  : awkure@protonmail.ch
-- Stability   : unstable
-- Portability : POSIX
-------------------------------------------------------------------------------

module OGLS.Engine.Math.Instances where

import Data.Hashable              ( Hashable (..)         )
import Control.Monad.Fix          ( MonadFix (..)         )
import Control.Monad.Zip          ( MonadZip (..)         )
import Control.Applicative        ( liftA2                )
import Data.Functor.Bind          ( Apply (..), Bind (..) )
import Data.Semigroup.Foldable    ( Foldable1 (..)        )
import Data.Semigroup.Traversable ( Traversable1 (..)     )
import Data.Complex               ( Complex (..)          )
import Data.Semigroup             ( (<>)                  )
import Data.HashMap.Lazy as HashMap

instance (Hashable k, Eq k) => Apply (HashMap k) where
  (<.>) = HashMap.intersectionWith id

instance (Hashable k, Eq k) => Bind (HashMap k) where
  m >>- f = HashMap.fromList $ do
    (k, a) <- HashMap.toList m
    case HashMap.lookup k (f a) of
      Just b -> [(k,b)]
      Nothing -> []

instance Apply Complex where
  (a :+ b) <.> (c :+ d) = a c :+ b d

instance Bind Complex where
  (a :+ b) >>- f = a' :+ b' where
    a' :+ _  = f a
    _  :+ b' = f b
  {-# INLINE (>>-) #-}

instance MonadZip Complex where
  mzipWith = liftA2

instance MonadFix Complex where
  mfix f = (let a :+ _ = f a in a) :+ (let _ :+ a = f a in a)

instance Foldable1 Complex where
  foldMap1 f (a :+ b) = f a <> f b
  {-# INLINE foldMap1 #-}

instance Traversable1 Complex where
  traverse1 f (a :+ b) = (:+) <$> f a <.> f b
  {-# INLINE traverse1 #-}


module LispVal
  ( LispVal(..)
  ) where

import Data.Complex

data LispVal
  = Atom String
  | List [LispVal]
  | DottedList [LispVal] LispVal
  | Number Integer
  | String String
  | Bool Bool
  | Character Char
  | Float Double
  | Ratio Rational
  | Complex (Complex Double)


{-# LANGUAGE FlexibleContexts #-}

module Numeric.Trainee.Data.Common (
	CsvM, runCsvM, col_, read_, single_, enum_, class_, sample_
	) where

import Prelude.Unicode

import Control.DeepSeq
import Control.Monad.Except
import Control.Monad.Reader
import Data.List
import qualified Data.Vector as V
import qualified Data.Text as T
import Numeric.LinearAlgebra
import Text.Read

import Numeric.Trainee.Types

type CsvM a = ReaderT (V.Vector T.Text) (Either String) a

runCsvM ∷ CsvM a → V.Vector T.Text → Either String a
runCsvM = runReaderT

col_ ∷ Int → CsvM String
col_ idx = ReaderT $ \r → maybe (throwError $ "invalid column index: " ++ show idx) (return ∘ T.unpack) (r V.!? idx)

read_ ∷ Read a ⇒ String → CsvM a
read_ s = maybe (throwError $ "can't parse: " ++ s) return ∘ readMaybe $ s

single_ ∷ a → CsvM [a]
single_ = return ∘ return

enum_ ∷ Fractional a ⇒ [String] → String → CsvM a
enum_ names name = case findIndex (≡ name) names of
	Nothing → throwError $ "invalid enum value: " ++ name ++ ", expected " ++ intercalate ", " names
	Just idx → return (fromIntegral idx / fromIntegral (length names - 1))

class_ ∷ Num a ⇒ [String] → String → CsvM [a]
class_ names name = case findIndex (≡ name) names of
	Nothing → throwError $ "invalid class value: " ++ name ++ ", expected " ++ intercalate ", " names
	Just idx → return $ replicate idx 0 ++ [1] ++ replicate (length names - idx - 1) 0

sample_ ∷ Container Vector a ⇒ [a] → [a] → CsvM (Sample (Vector a) (Vector a))
sample_ is os = return $!! Sample (fromList is) (fromList os)


{-# LANGUAGE GADTs               #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeApplications    #-}
-- |
-- Module      : Data.Array.Accelerate.Numeric.LinearAlgebra.LLVM.PTX.Level2
-- Copyright   : [2017..2020] Trevor L. McDonell
-- License     : BSD3
--
-- Maintainer  : Trevor L. McDonell <trevor.mcdonell@gmail.com>
-- Stability   : experimental
-- Portability : non-portable (GHC extensions)
--

module Data.Array.Accelerate.Numeric.LinearAlgebra.LLVM.PTX.Level2
  where

import Data.Complex
import Data.Array.Accelerate.Representation.Array
import Data.Array.Accelerate.Representation.Shape
import Data.Array.Accelerate.Sugar.Elt

import Data.Array.Accelerate.LLVM.PTX.Foreign
import Data.Array.Accelerate.Numeric.LinearAlgebra.LLVM.PTX.Base
import Data.Array.Accelerate.Numeric.LinearAlgebra.LLVM.PTX.Context
import Data.Array.Accelerate.Numeric.LinearAlgebra.LLVM.PTX.Level3
import Data.Array.Accelerate.Numeric.LinearAlgebra.Type

import Foreign.Marshal                                              ( with )
import qualified Foreign.CUDA.Ptr                                   as CUDA
import qualified Foreign.CUDA.BLAS                                  as BLAS

import Control.Monad.Reader

-- NOTE: cuBLAS requires matrices to be stored in column-major order
-- (Fortran-style), but Accelerate uses C-style arrays in row-major order.
--
-- If the operation is N or T, we can just swap the operation. For
-- conjugate-transpose (H) operations (on complex valued arguments), since there
-- is no conjugate-no-transpose operation, we implement that via 'gemm', which
-- I assume is more efficient than ?geam followed by ?gemv.
--
gemv :: NumericR s e
     -> Transpose
     -> ForeignAcc ((((((), Scalar e), Matrix e), Vector e)) -> Vector e)
gemv eR opA = ForeignAcc "ptx.gemv" (gemv' eR opA)

gemv' :: NumericR s e
      -> Transpose
      -> ((((), Scalar e), Matrix e), Vector e)
      -> Par PTX (Future (Vector e))
gemv' NumericRcomplex32 H = as_gemm NumericRcomplex32 H
gemv' NumericRcomplex64 H = as_gemm NumericRcomplex64 H
gemv' nR                t = as_gemv nR t

as_gemm
    :: NumericR s e
    -> Transpose
    -> ((((), Scalar e), Matrix e), Vector e)
    -> Par PTX (Future (Vector e))
as_gemm nR opA ((((), alpha), matA), Array sh adata) = do
  let matB = Array (sh,1) adata
  --
  future <- new
  result <- gemm' nR opA N ((((), alpha), matA), matB)
  fork $ do Array (sh',_) vecy <- get result
            put future (Array sh' vecy)
  return future

as_gemv
    :: NumericR s e
    -> Transpose
    -> ((((), Scalar e), Matrix e), Vector e)
    -> Par PTX (Future (Vector e))
as_gemv nR opA ((((), alpha), matA), vecx) = do
  let
      (((), rowsA), colsA) = shape matA

      sizeY   = case opA of
                  N -> rowsA
                  _ -> colsA

      opA'    = encodeTranspose
              $ case opA of
                  N -> T
                  _ -> N

      aR      = ArrayR dim1 eR
      eR      = case nR of
                  NumericRfloat32   -> eltR @Float
                  NumericRfloat64   -> eltR @Double
                  NumericRcomplex32 -> eltR @(Complex Float)
                  NumericRcomplex64 -> eltR @(Complex Double)
  --
  future  <- new
  stream  <- asks ptxStream
  vecy    <- allocateRemote aR ((), sizeY)
  alpha'  <- indexRemote eR alpha 0
  ()      <- liftPar $ do
    withArray nR matA stream   $ \ptr_A -> do
     withArray nR vecx stream  $ \ptr_x -> do
      withArray nR vecy stream $ \ptr_y -> do
       withBLAS                $ \hdl   -> do
         case nR of
           NumericRfloat32 -> liftIO $
            with alpha' $ \ptr_alpha ->
             with 0     $ \ptr_beta  ->
               BLAS.sgemv hdl opA' colsA rowsA ptr_alpha ptr_A colsA ptr_x 1 ptr_beta ptr_y 1

           NumericRfloat64 -> liftIO $
            with alpha' $ \ptr_alpha ->
             with 0     $ \ptr_beta  ->
               BLAS.dgemv hdl opA' colsA rowsA ptr_alpha ptr_A colsA ptr_x 1 ptr_beta ptr_y 1

           NumericRcomplex32 -> liftIO $
            withV2 nR alpha' $ \ptr_alpha ->
             with 0          $ \ptr_beta  ->
               BLAS.cgemv hdl opA' colsA rowsA ptr_alpha (CUDA.castDevPtr ptr_A) colsA (CUDA.castDevPtr ptr_x) 1 ptr_beta (CUDA.castDevPtr ptr_y)  1

           NumericRcomplex64 -> liftIO $
            withV2 nR alpha' $ \ptr_alpha ->
             with 0          $ \ptr_beta  ->
               BLAS.zgemv hdl opA' colsA rowsA ptr_alpha (CUDA.castDevPtr ptr_A) colsA (CUDA.castDevPtr ptr_x) 1 ptr_beta (CUDA.castDevPtr ptr_y)  1
  --
  put future vecy
  return future


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

{-# OPTIONS_GHC  -O0 #-}

module BlueRipple.Data.UsefulDataJoins where

import qualified Control.Foldl                 as FL
import           Control.Monad                  ( join )
import qualified Data.Array                    as A
import           Data.Function                  ( on )
import qualified Data.List                     as L
import qualified Data.Set                      as S
import qualified Data.Map                      as M
import           Data.Maybe                     ( isJust )

import qualified Data.Text                     as T
import qualified Data.Serialize                as SE
import           Data.Serialize.Text            ( )
import qualified Frames                        as F
import qualified Frames.Melt                   as F
import qualified Frames.InCore                 as FI
import qualified Data.Vinyl                    as V
import qualified Data.Vinyl.TypeLevel          as V
import qualified Data.Discrimination.Grouping  as G

import qualified Control.MapReduce             as MR
import qualified Frames.Transform              as FT
import qualified Frames.Folds                  as FF
import qualified Frames.MapReduce              as FMR
import qualified Frames.Enumerations           as FE
import qualified Frames.SimpleJoins            as FJ
import qualified Frames.Serialize              as FS

import qualified Graphics.Vega.VegaLite        as GV
import qualified Knit.Report                   as K

import           BlueRipple.Utilities.KnitUtils

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

import qualified BlueRipple.Data.DataFrames    as BR
import qualified BlueRipple.Data.DemographicTypes
                                               as BR

import qualified BlueRipple.Data.ElectionTypes
                                               as BR

--import qualified BlueRipple.Data.PrefModel     as BR
import qualified BlueRipple.Model.TurnoutAdjustment
                                               as BR

import qualified BlueRipple.Utilities.KnitUtils
                                               as BR
import qualified BlueRipple.Data.Keyed as Keyed

rollupF :: forall as bs ds. (Ord (F.Record bs)
                           , bs F.⊆ (as V.++ bs V.++ ds)
                           , ds F.⊆ (as V.++ bs V.++ ds)
                           , FI.RecVec (bs V.++ ds)
                           ) => FL.Fold (F.Record ds) (F.Record ds) -> FL.Fold (F.Record (as V.++ bs V.++ ds)) (F.FrameRec (bs V.++ ds))
rollupF dataF = FMR.concatFold $ FMR.mapReduceFold
         FMR.noUnpack
         (FMR.assignKeysAndData @bs)
         (FMR.foldAndAddKey dataF)

rollupSumF :: forall as bs ds. (Ord (F.Record bs)
                               , bs F.⊆ (as V.++ bs V.++ ds)
                               , ds F.⊆ (as V.++ bs V.++ ds)
                               , FI.RecVec (bs V.++ ds)
                               , FF.ConstrainedFoldable Num ds
                               )
        => FL.Fold (F.Record (as V.++ bs V.++ ds)) (F.FrameRec (bs V.++ ds))
rollupSumF = rollupF @as @bs @ds (FF.foldAllConstrained @Num FL.sum)

addElectoralWeight ::
  BR.ElectoralWeightSourceT
  -> BR.ElectoralWeightOfT
  -> (F.Record rs -> Double)
  -> F.Record rs
  -> F.Record (rs V.++ [BR.ElectoralWeightSource, BR.ElectoralWeightOf, BR.ElectoralWeight])
addElectoralWeight ews ewof calcEW r =
  let ewcs :: F.Record [BR.ElectoralWeightSource, BR.ElectoralWeightOf, BR.ElectoralWeight] = ews F.&: ewof F.&: calcEW r F.&: V.RNil
  in r `V.rappend` ewcs

type PCols p = [BR.PopCountOf, p]
type PEWCols p = PCols p V.++ BR.EWCols --[BR.PopCountOf, p, BR.ElectoralWeightSource, BR.ElectoralWeightOf, BR.ElectoralWeight]

joinDemoAndWeights
  :: forall js ks p
  . ( FJ.CanLeftJoinM  js (ks V.++ PCols p) (js V.++ BR.EWCols)
    , ((ks V.++ PCols p) V.++ F.RDeleteAll js (js V.++ BR.EWCols)) ~ ((ks V.++ PCols p) V.++ BR.EWCols)
    , js F.⊆ ks
    , FI.RecVec (ks V.++ PCols p V.++ BR.EWCols)
    , (ks V.++ PCols p V.++ BR.EWCols) F.⊆ ((ks V.++ PCols p) V.++ F.RDeleteAll js (js V.++ BR.EWCols))
    , js F.⊆  (ks V.++ PCols p)
    , js F.⊆  (js V.++ BR.EWCols)
    , (ks V.++ PCols p) F.⊆ ((ks V.++ PCols p) V.++ F.RDeleteAll js (js V.++ BR.EWCols))
    , F.RDeleteAll js (js V.++ BR.EWCols) F.⊆  (js V.++ BR.EWCols)
    , V.RMap (ks V.++ PCols p)
    , V.RMap  ((ks V.++ PCols p) V.++ F.RDeleteAll js (js V.++ BR.EWCols))
    , V.RecApplicative  (F.RDeleteAll js (js V.++ BR.EWCols))
    , G.Grouping (F.Record js)
    , FI.RecVec  (ks V.++ PCols p)
    , FI.RecVec (F.RDeleteAll js (js V.++ BR.EWCols))
    , FI.RecVec ((ks V.++ PCols p) V.++ F.RDeleteAll js (js V.++ BR.EWCols))
    )
  => F.FrameRec (ks V.++ (PCols p))
  -> F.FrameRec (js V.++ BR.EWCols)
  -> Either [F.Record js] (F.FrameRec (ks V.++ PCols p V.++ BR.EWCols))
joinDemoAndWeights d w = FJ.leftJoinE @js d w
{-F.toFrame
                         $ fmap F.rcast
                         $ catMaybes
                         $ fmap F.recMaybe
                         $ F.leftJoin @js d w
-}

{-
--defaultPopRec :: BR.PopCountOfT -> V.Snd p -> F.Record (PCols p)
--defaultPopRec pco x = pco F.&: x F.&: V.RNil

joinDemoAndWeightsWithDefaults
  :: forall js ks p
  . (js F.⊆ ks
    , Keyed.FiniteSet (F.Record ks)
    , FI.RecVec (ks V.++ PCols p V.++ BR.EWCols)
    , (ks V.++ PCols p V.++ BR.EWCols) F.⊆ ((ks V.++ PCols p) V.++ F.RDeleteAll js (js V.++ BR.EWCols))
    , js F.⊆  (ks V.++ PCols p)
    , js F.⊆  (js V.++ BR.EWCols)
    , (ks V.++ PCols p) F.⊆ ((ks V.++ PCols p) V.++ F.RDeleteAll js (js V.++ BR.EWCols))
    , (F.RDeleteAll js (js V.++ BR.EWCols)) F.⊆  (js V.++ BR.EWCols)
    , V.RMap (ks V.++ PCols p)
    , V.RMap  ((ks V.++ PCols p) V.++ F.RDeleteAll js (js V.++ BR.EWCols))
    , V.RecApplicative  (F.RDeleteAll js (js V.++ BR.EWCols))
    , G.Grouping (F.Record js)
    , FI.RecVec  (ks V.++ PCols p)
    , FI.RecVec (F.RDeleteAll js (js V.++ BR.EWCols))
    , FI.RecVec ((ks V.++ PCols p) V.++ F.RDeleteAll js (js V.++ BR.EWCols))
    , V.KnownField p
    , F.ElemOf (ks V.++ (PCols p)) p
    , F.ElemOf (ks V.++ (PCols p)) BR.PopCountOf
    , ks F.⊆ (ks V.++ (PCols p))
    , (ks V.++ PCols p V.++ BR.EWCols) F.⊆ ((js V.++ BR.EWCols) V.++ F.RDeleteAll js (ks V.++ PCols p))
    , (js V.++ BR.EWCols) F.⊆ ((js V.++ BR.EWCols) V.++ F.RDeleteAll js (ks V.++ PCols p))
    , (F.RDeleteAll js (ks V.++ PCols p)) F.⊆ (ks V.++ PCols p)
    , V.RMap (js V.++ BR.EWCols)
    , V.RMap ((js V.++ BR.EWCols) V.++ F.RDeleteAll js (ks V.++ PCols p))
    , V.RecApplicative (F.RDeleteAll js (ks V.++ PCols p))
    , FI.RecVec (js V.++ BR.EWCols)
    , FI.RecVec (F.RDeleteAll js (ks V.++ PCols p))
    , FI.RecVec ((js V.++ BR.EWCols) V.++ F.RDeleteAll js (ks V.++ PCols p))
    )
  => BR.PopCountOfT
  -> V.Snd p
  -> F.FrameRec (ks V.++ (PCols p))
  -> F.FrameRec (js V.++ BR.EWCols)
  -> F.FrameRec (ks V.++ PCols p V.++ BR.EWCols)
joinDemoAndWeightsWithDefaults pco x d w =
  let demoWithDefaults :: F.FrameRec (ks V.++ (PCols p))
      demoWithDefaults =  F.toFrame $ FL.fold (Keyed.addDefaultRec @ks @(PCols p) (pco F.&: x F.&: V.RNil)) d
  in F.toFrame
     $ fmap F.rcast
     $ catMaybes
     $ fmap F.recMaybe
     $ F.leftJoin @js w demoWithDefaults
-}

adjustWeightsForStateTotals
  :: forall ks p r
  . (K.KnitEffects r
    , V.KnownField p
    , V.Snd p ~ Int
    , F.ElemOf (BR.WithYS ks) BR.Year
    , F.ElemOf (BR.WithYS ks) BR.StateAbbreviation
    , (ks V.++ PCols p V.++ BR.EWCols) F.⊆ BR.WithYS (ks V.++ PCols p V.++ BR.EWCols)
    , F.ElemOf (ks V.++ PCols p V.++ BR.EWCols) p
    , F.ElemOf (ks V.++ PCols p V.++ BR.EWCols) BR.ElectoralWeight
    , FI.RecVec (ks V.++ PCols p V.++ BR.EWCols)
    , Show (F.Record ((ks V.++ PCols '(V.Fst p, Int)) V.++ BR.EWCols))
    )
  => F.Frame BR.StateTurnout
  -> F.FrameRec (BR.WithYS ks V.++ PCols p V.++ BR.EWCols)
  -> K.Sem r (F.FrameRec (BR.WithYS ks V.++ PCols p  V.++ BR.EWCols))
adjustWeightsForStateTotals stateTurnout unadj =
  FL.foldM
    (BR.adjTurnoutFold @p @BR.ElectoralWeight stateTurnout)
    unadj

demographicsWithAdjTurnoutByState
  :: forall catCols p ks js effs
  . ( K.KnitEffects effs
    , FJ.CanLeftJoinM (js V.++ catCols) ((BR.WithYS ks) V.++ catCols V.++ (PCols p)) (js V.++ catCols V.++ BR.EWCols)
    , (((ks V.++ catCols) V.++ PCols p) V.++ F.RDeleteAll (js V.++ catCols) ((js V.++ catCols) V.++ BR.EWCols)) ~ (((ks V.++ catCols) V.++ PCols p) V.++ BR.EWCols
                                                                                )

    , V.RMap (js V.++ catCols)
    , V.ReifyConstraint Show V.ElField (js V.++ catCols)
    , V.RecordToList (js V.++ catCols)
    , V.KnownField p
    , V.Snd p ~ Int
    , F.ElemOf (ks V.++ catCols V.++ PEWCols p) p
    , F.ElemOf (ks V.++ catCols V.++ PEWCols p) BR.ElectoralWeight
    , (ks V.++ catCols V.++ PEWCols p) F.⊆ BR.WithYS (ks V.++ catCols V.++ PEWCols p)
    , FI.RecVec (ks V.++ catCols V.++ PEWCols p)
    , (ks V.++ catCols V.++ PEWCols p) F.⊆ (BR.WithYS (ks V.++ catCols V.++ PCols p) V.++ F.RDeleteAll (js V.++ catCols) (js V.++ catCols V.++ BR.EWCols))
    , (js V.++ catCols) F.⊆ ((js V.++ catCols V.++ BR.EWCols))
    , (js V.++ catCols) F.⊆ BR.WithYS (ks V.++ catCols)
    , (js V.++ catCols) F.⊆ BR.WithYS (ks V.++ catCols V.++ (PCols p))
    , (ks V.++ catCols V.++ (PCols p)) F.⊆ BR.WithYS (ks V.++ catCols V.++ (PCols p) V.++ F.RDeleteAll (js V.++ catCols) ((js V.++ catCols) V.++ BR.EWCols))
    , (F.RDeleteAll (js V.++ catCols) ((js V.++ catCols) V.++ BR.EWCols)) F.⊆ ((js V.++ catCols) V.++ BR.EWCols)
    , V.RMap ((ks V.++ catCols) V.++ PCols p)
    , V.RMap ((((ks V.++ catCols) V.++ PCols p) V.++ F.RDeleteAll (js V.++ catCols) ((js V.++ catCols) V.++ BR.EWCols)))
    , V.RecApplicative (F.RDeleteAll (js V.++ catCols) ((js V.++ catCols) V.++ BR.EWCols))
    , G.Grouping (F.Record (js V.++ catCols))
    , FI.RecVec ((ks V.++ catCols) V.++ PCols p)
    , FI.RecVec (F.RDeleteAll (js V.++ catCols) ((js V.++ catCols) V.++ BR.EWCols))
    , FI.RecVec (((ks V.++ catCols) V.++ PCols p) V.++ F.RDeleteAll (js V.++ catCols) ((js V.++ catCols) V.++ BR.EWCols))
    , (ks V.++ catCols V.++ PCols p V.++ BR.EWCols) ~ (ks V.++ catCols V.++ (PCols p V.++ BR.EWCols))
    , Show (F.Record ((ks V.++ catCols)
                           V.++ '[ '("PopCountOf", BR.PopCountOfT), p,
                                   BR.ElectoralWeightSource, BR.ElectoralWeightOf,
                                   BR.ElectoralWeight]))
    )
  => F.Frame BR.StateTurnout
  -> F.FrameRec ((BR.WithYS ks) V.++ catCols V.++ (PCols p))
  -> F.FrameRec (js V.++ catCols V.++ BR.EWCols)
  -> K.Sem effs (F.FrameRec ((BR.WithYS ks) V.++ catCols V.++ (PEWCols p)))
demographicsWithAdjTurnoutByState stateTurnout demos ews = do
  let joinedE :: Either [F.Record (js V.++ catCols)] (F.FrameRec ((BR.WithYS ks) V.++ catCols V.++ (PEWCols p)))
      joinedE = joinDemoAndWeights @(js V.++ catCols) @((BR.WithYS ks) V.++ catCols) @p demos ews
  joined <- K.knitEither
            $ either (\missingKeys -> Left $ "missing keys in electoral weights in demographicsWithAdjTurnoutByState: " <> show missingKeys) Right
            $ joinedE
  adjustWeightsForStateTotals @(ks V.++ catCols) @p stateTurnout joined

{-
demographicsWithDefaultsWithAdjTurnoutByState
  :: forall catCols p ks js effs
  . ( K.KnitEffects effs
    , V.KnownField p
    , V.Snd p ~ Int
    , F.ElemOf (ks V.++ catCols V.++ PEWCols p) p
    , F.ElemOf (ks V.++ catCols V.++ PEWCols p) BR.ElectoralWeight
    , (ks V.++ catCols V.++ PEWCols p) F.⊆ BR.WithYS (ks V.++ catCols V.++ PEWCols p)
    , FI.RecVec (ks V.++ catCols V.++ PEWCols p)
    , (ks V.++ catCols V.++ PEWCols p) F.⊆ (BR.WithYS (ks V.++ catCols V.++ PCols p) V.++ F.RDeleteAll (js V.++ catCols) (js V.++ catCols V.++ BR.EWCols))
    , (js V.++ catCols) F.⊆ ((js V.++ catCols V.++ BR.EWCols))
    , (js V.++ catCols) F.⊆ BR.WithYS (ks V.++ catCols)
    , (js V.++ catCols) F.⊆ BR.WithYS (ks V.++ catCols V.++ (PCols p))
    , (ks V.++ catCols V.++ (PCols p)) F.⊆ BR.WithYS (ks V.++ catCols V.++ (PCols p) V.++ F.RDeleteAll (js V.++ catCols) ((js V.++ catCols) V.++ BR.EWCols))
    , (F.RDeleteAll (js V.++ catCols) ((js V.++ catCols) V.++ BR.EWCols)) F.⊆ ((js V.++ catCols) V.++ BR.EWCols)
    , V.RMap ((ks V.++ catCols) V.++ PCols p)
    , V.RMap ((((ks V.++ catCols) V.++ PCols p) V.++ F.RDeleteAll (js V.++ catCols) ((js V.++ catCols) V.++ BR.EWCols)))
    , V.RecApplicative (F.RDeleteAll (js V.++ catCols) ((js V.++ catCols) V.++ BR.EWCols))
    , G.Grouping (F.Record (js V.++ catCols))
    , FI.RecVec ((ks V.++ catCols) V.++ PCols p)
    , FI.RecVec (F.RDeleteAll (js V.++ catCols) ((js V.++ catCols) V.++ BR.EWCols))
    , FI.RecVec (((ks V.++ catCols) V.++ PCols p) V.++ F.RDeleteAll (js V.++ catCols) ((js V.++ catCols) V.++ BR.EWCols))
    , (ks V.++ catCols V.++ PCols p V.++ BR.EWCols) ~ (ks V.++ catCols V.++ (PCols p V.++ BR.EWCols))
    )
  => BR.PopCountOfT
  -> V.Snd p
  -> F.Frame BR.StateTurnout
  -> F.FrameRec ((BR.WithYS ks) V.++ catCols V.++ (PCols p))
  -> F.FrameRec (js V.++ catCols V.++ BR.EWCols)
  -> K.Sem effs (F.FrameRec ((BR.WithYS ks) V.++ catCols V.++ (PEWCols p)))
demographicsWithDefaultsWithAdjTurnoutByState pco x stateTurnout demos ews = do
  let joined :: F.FrameRec ((BR.WithYS ks) V.++ catCols V.++ (PEWCols p))
      joined = joinDemoAndWeightsWithDefaults @(js V.++ catCols) @((BR.WithYS ks) V.++ catCols) @p pco x demos ews
  adjustWeightsForStateTotals @(ks V.++ catCols) @p stateTurnout joined

-}

-- This is monstrous
rollupAdjustAndJoin
  :: forall as catCols p ks js effs
  .(K.KnitEffects effs
   , FJ.CanLeftJoinM (js V.++ catCols) ((BR.WithYS ks) V.++ catCols V.++ (PCols p)) (js V.++ catCols V.++ BR.EWCols)
   , (((ks V.++ catCols) V.++ PCols p) V.++ F.RDeleteAll (js V.++ catCols) ((js V.++ catCols) V.++ BR.EWCols)) ~ (((ks V.++ catCols) V.++ PCols p) V.++ BR.EWCols)
   , V.RMap (js V.++ catCols)
   , V.ReifyConstraint Show V.ElField (js V.++ catCols)
   , V.RecordToList (js V.++ catCols)
   , V.KnownField p
   , V.Snd p ~ Int
   , F.ElemOf (ks V.++ catCols V.++ (PEWCols p)) p
   , F.ElemOf (ks V.++ catCols V.++ (PEWCols p)) BR.ElectoralWeight
   , (ks V.++ catCols V.++ (PEWCols p)) F.⊆ BR.WithYS (ks V.++ catCols V.++ (PEWCols p))
   , FI.RecVec (ks V.++ catCols V.++ (PEWCols p))
   , (ks V.++ catCols V.++ (PEWCols p)) F.⊆ BR.WithYS (ks V.++ catCols V.++ (PCols p) V.++ F.RDeleteAll (js V.++ catCols) ((js V.++ catCols V.++ BR.EWCols)))
   , (js V.++ catCols) F.⊆ ((js V.++ catCols V.++ BR.EWCols))
   , (js V.++ catCols) F.⊆ BR.WithYS (ks V.++ catCols)
   , (js V.++ catCols) F.⊆ BR.WithYS (ks V.++ catCols V.++ (PCols p))
   , (ks V.++ catCols V.++ (PCols p)) F.⊆ BR.WithYS (ks V.++ catCols V.++ (PCols p) V.++ F.RDeleteAll (js V.++ catCols) ((js V.++ catCols) V.++ BR.EWCols))
   , (F.RDeleteAll (js V.++ catCols) ((js V.++ catCols) V.++ BR.EWCols)) F.⊆ ((js V.++ catCols) V.++ BR.EWCols)
   , V.RMap ((ks V.++ catCols) V.++ PCols p)
   , V.RMap ((((ks V.++ catCols) V.++ PCols p) V.++ F.RDeleteAll (js V.++ catCols) ((js V.++ catCols) V.++ BR.EWCols)))
   , V.RecApplicative (F.RDeleteAll (js V.++ catCols) ((js V.++ catCols) V.++ BR.EWCols))
   , G.Grouping (F.Record (js V.++ catCols))
   , FI.RecVec ((ks V.++ catCols) V.++ PCols p)
   , FI.RecVec (F.RDeleteAll (js V.++ catCols) ((js V.++ catCols) V.++ BR.EWCols))
   , FI.RecVec (((ks V.++ catCols) V.++ PCols p) V.++ F.RDeleteAll (js V.++ catCols) ((js V.++ catCols) V.++ BR.EWCols))
   , (as V.++ BR.WithYS ks V.++ catCols V.++ PCols p) ~ (as V.++ BR.WithYS (ks V.++ catCols V.++ '[BR.PopCountOf]) V.++ '[p])
   , (((ks V.++ catCols) V.++ '[BR.PopCountOf]) V.++ '[p]) ~ ((ks V.++ catCols) V.++ PCols p)
   , Ord (F.Record ((ks V.++ catCols) V.++ '[BR.PopCountOf]))
   , F.ElemOf  (as V.++ BR.WithYS (ks V.++ catCols V.++ '[BR.PopCountOf]) V.++ '[p]) p
   , F.ElemOf  (as V.++ BR.WithYS (ks V.++ catCols V.++ '[BR.PopCountOf]) V.++ '[p]) BR.StateAbbreviation
   , F.ElemOf  (as V.++ BR.WithYS (ks V.++ catCols V.++ '[BR.PopCountOf]) V.++ '[p]) BR.Year
   , ((ks V.++ catCols) V.++ '[BR.PopCountOf]) F.⊆ (as V.++ BR.WithYS (ks V.++ catCols V.++ '[BR.PopCountOf]) V.++ '[p])
   , FI.RecVec (as V.++ BR.WithYS (ks V.++ catCols V.++ '[BR.PopCountOf]) V.++ '[p] V.++  F.RDeleteAll (ks V.++ catCols) ((ks V.++ catCols) V.++ BR.EWCols))
   , V.AllConstrained G.Grouping (F.UnColumn (ks V.++ catCols))
   , (ks V.++ catCols) F.⊆ (as V.++ BR.WithYS (ks V.++ catCols V.++ '[BR.PopCountOf]) V.++ '[p])
   , (as V.++ BR.WithYS (ks V.++ catCols V.++ '[BR.PopCountOf]) V.++ '[p]) F.⊆ (as V.++ BR.WithYS (ks V.++ catCols V.++ '[BR.PopCountOf]) V.++ '[p] V.++  F.RDeleteAll (ks V.++ catCols) ((ks V.++ catCols) V.++ BR.EWCols))
   , V.RMap (as V.++ BR.WithYS (ks V.++ catCols V.++ '[BR.PopCountOf]) V.++ '[p])
   , V.RMap (as V.++ BR.WithYS (ks V.++ catCols V.++ '[BR.PopCountOf]) V.++ '[p] V.++  F.RDeleteAll (ks V.++ catCols) ((ks V.++ catCols) V.++ BR.EWCols))
   , G.Grouping (F.Record (ks V.++ catCols))
   , FI.RecVec (as V.++ BR.WithYS (ks V.++ catCols V.++ '[BR.PopCountOf]) V.++ '[p])
   , (ks V.++ catCols V.++ PCols p V.++ BR.EWCols) ~ (ks V.++ catCols V.++ (PCols p V.++ BR.EWCols))
   , (as V.++ BR.WithYS (ks V.++ catCols V.++ '[BR.PopCountOf]) V.++ '[p] V.++ BR.EWCols)
    ~ ((as V.++ BR.WithYS (ks V.++ catCols V.++ '[BR.PopCountOf]) V.++ '[p]) V.++ F.RDeleteAll (ks V.++ catCols) (ks V.++ catCols V.++ BR.EWCols))
   , (ks V.++ catCols V.++ BR.EWCols) F.⊆ BR.WithYS (ks V.++ catCols V.++ PEWCols p)
   , (ks V.++ catCols) F.⊆ BR.WithYS (ks V.++ catCols V.++ BR.EWCols)
   , (F.RDeleteAll (ks V.++ catCols) ((ks V.++ catCols) V.++ BR.EWCols)) F.⊆ BR.WithYS (ks V.++ catCols V.++ BR.EWCols)
   , V.RecApplicative (F.RDeleteAll (ks V.++ catCols) ((ks V.++ catCols) V.++ BR.EWCols))
   , FI.RecVec (F.RDeleteAll (ks V.++ catCols) ((ks V.++ catCols) V.++ BR.EWCols))
   , Show (F.Record ((ks V.++ catCols)
                           V.++ '[ '("PopCountOf", BR.PopCountOfT), p,
                                   BR.ElectoralWeightSource, BR.ElectoralWeightOf,
                                   BR.ElectoralWeight]))
   )
  => F.Frame BR.StateTurnout
  -> F.FrameRec (as V.++ (BR.WithYS ks) V.++ catCols V.++ (PCols p))
  -> F.FrameRec (js V.++ catCols V.++ BR.EWCols)
  -> K.Sem effs (F.FrameRec (as V.++ (BR.WithYS ks) V.++ catCols V.++ PCols p V.++ BR.EWCols))
rollupAdjustAndJoin stateTurnout demos ews = do
  let rolledUpDemos :: F.FrameRec ((BR.WithYS ks) V.++ catCols V.++ (PCols p))
      rolledUpDemos = FL.fold (rollupSumF @as @((BR.WithYS ks) V.++ catCols V.++ '[BR.PopCountOf]) @'[p]) demos
  adjusted :: F.FrameRec (BR.WithYS ks V.++ catCols V.++ PEWCols p) <- demographicsWithAdjTurnoutByState @catCols @p @ks @js  stateTurnout rolledUpDemos ews
  let adjustedWithoutDemo = fmap (F.rcast @(BR.WithYS ks V.++ catCols V.++ BR.EWCols)) adjusted
  return
      $ F.toFrame
      $ catMaybes
      $ fmap F.recMaybe
      $ F.leftJoin @(BR.WithYS ks V.++ catCols) demos adjustedWithoutDemo

type TurnoutCols = [BR.Population, BR.Citizen, BR.Registered, BR.Voted]
type ACSColsCD = [BR.CongressionalDistrict, BR.Year, BR.StateAbbreviation, BR.StateFIPS, BR.StateName]
type ACSCols = [BR.Year, BR.StateAbbreviation, BR.StateFIPS, BR.StateName]

-- This is also monstrous.  Which is surprising??
acsDemographicsWithAdjCensusTurnoutByCD
  :: forall catCols r
  . (K.KnitEffects r
    , BR.CacheEffects r
    , V.RMap catCols
    , V.ReifyConstraint Show V.ElField catCols
    , V.RecordToList catCols
    , ((catCols V.++ PCols BR.ACSCount) V.++ BR.EWCols) ~ (catCols V.++ PEWCols BR.ACSCount)
    , ((catCols V.++ '[BR.PopCountOf]) V.++ '[BR.ACSCount]) ~ (catCols V.++ [BR.PopCountOf, BR.ACSCount])
    , ((catCols V.++ PCols BR.ACSCount) V.++ F.RDeleteAll catCols (catCols V.++ BR.EWCols)) ~ (catCols V.++ PEWCols BR.ACSCount)
    , BR.RecSerializerC (BR.ACSKeys V.++ catCols V.++ '[BR.ACSCount, BR.VotedPctOfAll])
    , V.RMap (catCols V.++ '[BR.ACSCount, BR.VotedPctOfAll])
    , FI.RecVec (catCols V.++ '[BR.ACSCount, BR.VotedPctOfAll])
    , (catCols V.++ PCols BR.ACSCount) F.⊆ (BR.ACSKeys V.++ (catCols V.++ '[BR.ACSCount]) V.++ '[BR.PopCountOf])
    , (catCols V.++ BR.EWCols) F.⊆ (BR.Year : (catCols V.++ TurnoutCols V.++ BR.EWCols))
    , F.ElemOf (catCols V.++ TurnoutCols) BR.Population
    , F.ElemOf (catCols V.++ TurnoutCols) BR.Voted
    , V.AllConstrained G.Grouping (F.UnColumn catCols)
    , V.RMap (catCols V.++ PCols BR.ACSCount)
    , V.RMap (catCols V.++ PEWCols BR.ACSCount)
    , V.RecApplicative (F.RDeleteAll catCols (catCols V.++ BR.EWCols))
    , G.Grouping (F.Record catCols)
    , FI.RecVec (catCols V.++ PEWCols BR.ACSCount)
    , FI.RecVec (catCols V.++ PCols BR.ACSCount)
    , FI.RecVec (F.RDeleteAll catCols (catCols V.++ BR.EWCols))
    , Ord (F.Record (catCols V.++ '[BR.PopCountOf]))
    , F.ElemOf (catCols V.++ PEWCols BR.ACSCount) BR.ACSCount
    , F.ElemOf (catCols V.++ PCols BR.ACSCount) BR.ACSCount
    , F.ElemOf (catCols V.++ PEWCols BR.ACSCount) BR.ElectoralWeight
    , (catCols V.++ PEWCols BR.ACSCount) F.⊆ (ACSCols V.++ catCols V.++ PEWCols BR.ACSCount)
    , (catCols V.++ '[BR.PopCountOf]) F.⊆ (ACSColsCD V.++ catCols V.++ PCols BR.ACSCount)
    , (catCols V.++ PCols BR.ACSCount) F.⊆ (ACSCols V.++ catCols V.++ PEWCols BR.ACSCount)
    , (catCols V.++ PCols BR.ACSCount) F.⊆ (ACSColsCD V.++ catCols V.++ PEWCols BR.ACSCount)
    , (catCols V.++ BR.EWCols) F.⊆ (ACSCols V.++ catCols V.++ PEWCols BR.ACSCount)
    , catCols F.⊆ (ACSColsCD V.++ catCols V.++ PCols BR.ACSCount)
    , catCols F.⊆ (BR.Year ': (catCols V.++ BR.EWCols))
    , catCols F.⊆ (ACSCols V.++ catCols)
    , catCols F.⊆ (ACSCols V.++ catCols V.++ PCols BR.ACSCount)
    , catCols F.⊆ (ACSCols V.++ catCols V.++ BR.EWCols)
    , (F.RDeleteAll catCols (catCols V.++ BR.EWCols))  F.⊆ (BR.Year ': (catCols V.++ BR.EWCols))
    , (F.RDeleteAll catCols (catCols V.++ BR.EWCols))  F.⊆ (ACSCols V.++ (catCols V.++ BR.EWCols))
    , (catCols V.++ '[BR.ACSCount, BR.VotedPctOfAll]) F.⊆ (ACSColsCD V.++ (F.RDelete BR.ElectoralWeight (catCols V.++ PEWCols BR.ACSCount)) V.++ '[BR.VotedPctOfAll])
    , (F.RDelete BR.ElectoralWeight (catCols V.++ PEWCols BR.ACSCount)) F.⊆ (ACSColsCD V.++ catCols V.++ PEWCols BR.ACSCount)
    , V.ReifyConstraint Show V.ElField (catCols
                                        V.++ '[ '("PopCountOf", BR.PopCountOfT), '("ACSCount", Int),
                                                BR.ElectoralWeightSource, BR.ElectoralWeightOf,
                                                BR.ElectoralWeight])
    , V.RecordToList (catCols
                       V.++ '[ '("PopCountOf", BR.PopCountOfT), '("ACSCount", Int),
                               BR.ElectoralWeightSource, BR.ElectoralWeightOf,
                               BR.ElectoralWeight])
    )
  => T.Text
  -> K.ActionWithCacheTime r (F.FrameRec (BR.ACSKeys V.++ catCols V.++ '[BR.ACSCount]))
  -> K.ActionWithCacheTime r (F.FrameRec ('[BR.Year] V.++ catCols V.++ '[BR.Population, BR.Citizen, BR.Registered, BR.Voted]))
  -> K.ActionWithCacheTime r (F.Frame BR.StateTurnout)
  -> K.Sem r (K.ActionWithCacheTime r (F.FrameRec (BR.ACSKeys V.++ catCols V.++ '[BR.ACSCount, BR.VotedPctOfAll])))
acsDemographicsWithAdjCensusTurnoutByCD cacheKey cachedDemo cachedTurnout cachedStateTurnout = do
  let cachedDeps = (,,) <$> cachedDemo <*> cachedTurnout <*> cachedStateTurnout
  BR.retrieveOrMakeFrame cacheKey cachedDeps $ \(demoF, turnoutF, stateTurnoutF) -> do
    let demo' = fmap (FT.mutate $ const $ FT.recordSingleton @BR.PopCountOf BR.PC_All) demoF
        demo'' = fmap (F.rcast @([BR.CongressionalDistrict, BR.Year, BR.StateAbbreviation, BR.StateFIPS, BR.StateName] V.++ catCols V.++ (PCols BR.ACSCount))) demo'
        vpa r = realToFrac (F.rgetField @BR.Voted r) / realToFrac (F.rgetField @BR.Population r)
    let turnout' = fmap (F.rcast @('[BR.Year] V.++ catCols V.++ BR.EWCols) . addElectoralWeight BR.EW_Census BR.EW_All vpa) turnoutF
    result <- rollupAdjustAndJoin @'[BR.CongressionalDistrict] @catCols @BR.ACSCount @[BR.StateFIPS, BR.StateName] @'[BR.Year] stateTurnoutF demo'' turnout'
    return $ F.toFrame $ fmap (F.rcast . FT.retypeColumn @BR.ElectoralWeight @BR.VotedPctOfAll) result

cachedASEDemographicsWithAdjTurnoutByCD
  :: (K.KnitEffects r
     , BR.CacheEffects r
     )
  => K.ActionWithCacheTime r (F.FrameRec (BR.ACSKeys V.++ BR.CatColsASE V.++ '[BR.ACSCount]))
  -> K.ActionWithCacheTime r (F.FrameRec ( '[BR.Year] V.++ BR.CatColsASE V.++ '[BR.Population, BR.Citizen, BR.Registered, BR.Voted]))
  -> K.ActionWithCacheTime r (F.Frame BR.StateTurnout)
  -> K.Sem r (K.ActionWithCacheTime r (F.FrameRec (BR.ACSKeys V.++ BR.CatColsASE V.++'[BR.ACSCount, BR.VotedPctOfAll])))
cachedASEDemographicsWithAdjTurnoutByCD = acsDemographicsWithAdjCensusTurnoutByCD @BR.CatColsASE "turnout/aseDemoWithStateAdjTurnoutByCD.bin"

cachedASRDemographicsWithAdjTurnoutByCD
  :: (K.KnitEffects r
     , BR.CacheEffects r
     )
  => K.ActionWithCacheTime r (F.FrameRec (BR.ACSKeys V.++ BR.CatColsASR V.++ '[BR.ACSCount]))
  -> K.ActionWithCacheTime r (F.FrameRec ( '[BR.Year] V.++ BR.CatColsASR V.++ '[BR.Population, BR.Citizen, BR.Registered, BR.Voted]))
  -> K.ActionWithCacheTime r (F.Frame BR.StateTurnout)
  -> K.Sem r (K.ActionWithCacheTime r (F.FrameRec (BR.ACSKeys V.++ BR.CatColsASR V.++'[BR.ACSCount, BR.VotedPctOfAll])))
cachedASRDemographicsWithAdjTurnoutByCD = acsDemographicsWithAdjCensusTurnoutByCD @BR.CatColsASR "turnout/asrDemoWithStateAdjTurnoutByCD.bin"


{-# 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]