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

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

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

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

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

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

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

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

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

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

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

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


{-# LANGUAGE TemplateHaskell
           , TupleSections
           , DeriveGeneric
           , DeriveAnyClass
           , OverloadedLists #-}
module Physics.Collision where

import Numeric.LinearAlgebra (Vector, Matrix, (<>), scale, dot, cross, (#>))
import Control.Lens
import Numeric.LinearAlgebra.Data
import Data.Vector.Storable (init)
import GHC.Generics
import qualified Data.Vector as V
import qualified Data.Vector.Storable as VS
import Data.Function (on)

import Control.DeepSeq
import Utils
import Prelude hiding (init)

data Contact = Contact { _contactPoint :: Vector Float
                       , _contactNormal :: Vector Float
                       , _penetration :: Float
                       , _sign :: Float } deriving Show

makeLenses ''Contact

makeContactBasis cont =
  let ys = [0,1,0]
      y2 = [0,0,1]
      x  = cont^.contactNormal
      zs = normalize $ cross x ys
      z  = if VS.any isNaN zs
           then normalize $ cross x y2
           else zs
      y  = normalize $ cross z x
  in fromRows [x,y,z]

data CollisionData = CollisionData { _contacts :: [Contact]
                                   , _left :: Int }

makeLenses ''CollisionData

data Primitive = Sphere { _radius     :: Float
                        , _primoffset :: Matrix Float }
               | Plane  { _normal     :: Vector Float
                        , _planeoff   :: Float }
               | Box    { _points     :: V.Vector (Vector Float)
                        , _faces      :: V.Vector Squad }
                 deriving Show

data PrimitiveComputation = CSphere { _r    :: Float
                                    , _poff :: Vector Float }
                          | CPlane  { _cnormal   :: Vector Float
                                    , _cplaneoff :: Float}
                          | CBox    { _cpoints :: V.Vector (Vector Float)
                                    , _quads :: V.Vector Quad } deriving Show

data Quad = Quad { _qnorm   :: Vector Float
                 , _qoff    :: Float} deriving Show

quadToCPlane (Quad n o) = CPlane n o

data Squad = Squad { _snorm :: Vector Float
                   , _spoint:: Vector Float } deriving Show

sToQ t (Squad n p) = Quad newN o
  where
    newN = init . (t#>) . flip VS.snoc 0 $ n
    o    = dot newN . init . (t#>) . flip VS.snoc 1 $ p

preparePrimitive :: Matrix Float -> Primitive -> PrimitiveComputation
preparePrimitive t (Sphere r p1) = CSphere r ppos
  where ppos = init . (!!3) . toColumns $ t <> p1
preparePrimitive t (Box p q) = CBox (V.map adjust p) $ V.map (sToQ t) q
  where adjust = init . (t #>) . flip VS.snoc 1
preparePrimitive _ (Plane n o) = CPlane n o

tolerance = 0.05

quadDepth (Quad n o) p =
  (dot n p - o, n)

boxPtcoll (CBox _ qs) p cdata =
  let dots = V.map (flip quadDepth p) qs
      (d,n) = V.minimumBy (compare `on` fst) dots
  in if V.any ((>=0) . fst) dots
     then cdata
     else over contacts (Contact p n d 1:) cdata

quadPenetration (Quad n o) = planePenetration (CPlane n o)

planePenetration (CPlane n o) p cdata =
  let vertDist = dot p n
      cPoint   = p -- scale (vertDist - o) n + p
      contact  = Contact cPoint n (o - vertDist) 1
  in if cdata^.left > 0 && vertDist <= o + tolerance
     then over contacts (contact:) . over left (subtract 1) $ cdata
     else cdata

-- checkAxe :: Vector Float -> Vector Float -> V.Vector
checkAxe pos r ps ix =
  let vals = V.map (VS.!ix) ps
  in (pos VS.! ix + r) >= minimum vals &&
     (pos VS.! ix - r) <= maximum vals

checkConvexAxe ps1 ps2 ix =
  let vals1 = V.map (VS.! ix) ps1
      vals2 = V.map (VS.! ix) ps2
  in maximum vals1 >= minimum vals2 &&
     minimum vals1 <= maximum vals2

axeList :: [Int]
axeList = [0,1,2]

separatingAxes :: PrimitiveComputation -> PrimitiveComputation -> Bool
separatingAxes (CSphere r pos) (CBox ps _) = all (checkAxe pos r ps)     axeList
separatingAxes (CBox ps _) (CSphere r pos) = all (checkAxe pos r ps)     axeList
separatingAxes (CBox ps _) (CBox ps2 _)    = all (checkConvexAxe ps ps2) axeList

detectCollision prim1 prim2 =
  collisionDetector prim1 prim2 (CollisionData [] 8)

quadSphereCollision (CSphere r p) q@(Quad n o) =
  (dot n p - r - o, n)

quadContact (CSphere r p) (dist, n) =
  let contactPos = p - scale (dist + r) n
  in Contact contactPos n (-dist) 1

collisionDetector (CSphere r1 pos1) (CSphere r2 pos2) cdata =
  let mid     = pos1 - pos2
      midnorm = norm mid
  in if cdata^.left <= 0 || midnorm <= 0.0 || midnorm >= r1 + r2
     then cdata
     else let normal     = scale (1/midnorm) mid
              contactPos = pos1 + scale 0.5 mid
              contactPen = r1 + r2 - midnorm
              contact    = Contact contactPos normal contactPen 1
          in over contacts (contact:) . over left (subtract 1) $ cdata
collisionDetector p@(CPlane _ _) s@(CSphere _ _) cdata =
  set (contacts.traverse.sign) (-1) $
  collisionDetector s p cdata
collisionDetector (CSphere r pos1) (CPlane n p2) cdata =
  if cdata^.left <= 0 then cdata
  else let dist = dot n pos1 - r - p2
       in if dist >= 0
          then cdata
          else let contactPos = pos1 - scale (dist + r) n
                   contact    = Contact contactPos n (-dist) 1
               in over contacts (contact:) cdata
collisionDetector (CBox ps _) p@(CPlane _ _) cdata =
  V.foldr (planePenetration p) cdata ps
collisionDetector s@(CSphere _ _) b@(CBox _ _) cdata =
  set (contacts.traverse.sign) (-1) $
  collisionDetector b s cdata
collisionDetector (CBox _ q) s@(CSphere _ _) cdata =
  let cols = V.map (quadSphereCollision s) q
      best = V.maximumBy (compare `on` fst) cols
  in if V.any ((>=0) . fst) cols
     then cdata
     else over contacts (quadContact s best:) cdata
collisionDetector p@(CPlane _ _) b@(CBox _ _) cdata =
  set (contacts.traverse.sign) (-1) $
  collisionDetector b p cdata
collisionDetector b1@(CBox ps qs) b2@(CBox ps2 qs2) cdata =
  V.foldr (boxPtcoll b2) (V.foldr (boxPtcoll b1) cdata ps2) ps
collisionDetector _ _ cdata = cdata


{-# language Rank2Types #-}
{-# language ScopedTypeVariables #-}
{-# language DataKinds #-}
{-# language TypeOperators #-}
{-# language TypeFamilies #-}
{-# language FlexibleInstances #-}
{-# language FlexibleContexts #-}
{-# language MultiParamTypeClasses #-}
{-# language PolyKinds #-}
{-# language GADTs #-}
{-# language ConstraintKinds #-}
{-# language GeneralizedNewtypeDeriving #-}

module Feldspar.Software.Verify.Primitive where

import Feldspar.Sugar
import Feldspar.Representation
import Feldspar.Software.Primitive
import Feldspar.Software.Expression
import Feldspar.Software.Representation hiding (Nil)
import Feldspar.Software.Verify.Command
import Feldspar.Verify.Arithmetic

import Feldspar.Verify.Monad (Verify)
import qualified Feldspar.Verify.FirstOrder as FO
import qualified Feldspar.Verify.Monad as V
import qualified Feldspar.Verify.SMT as SMT
import qualified Feldspar.Verify.Abstract as A

import Data.Struct
import qualified Data.Map.Strict as Map

import qualified Control.Monad.RWS.Strict as S

import qualified SimpleSMT as SMT hiding (not, declareFun)

import qualified Language.Embedded.Expression as Imp
import qualified Language.Embedded.Imperative.CMD as Imp

import qualified Data.Bits as Bits
import qualified Data.Complex as Complex
import Data.Constraint hiding (Sub)
import Data.Int
import Data.Word
import Data.Typeable

import Language.Syntactic

import GHC.Stack

--------------------------------------------------------------------------------
-- *
--------------------------------------------------------------------------------

newtype Symbolic a = Symbolic { unSymbolic :: Rat }
  deriving (Eq, Ord, Show, V.TypedSExpr, V.SMTOrd)

instance V.Fresh (Symbolic a)
  where
    fresh = V.freshSExpr

symbCast :: Symbolic a -> Symbolic b
symbCast = Symbolic . unSymbolic

symbFun :: SymbParam a => String -> [Symbolic a] -> Symbolic a
symbFun name (args :: [Symbolic a]) = V.fromSMT $
  SMT.fun (symbType (undefined :: a) ++ "-" ++ name) (map V.toSMT args)

fromComplexConstant :: (RealFrac a, SymbParam b) => Complex.Complex a -> Symbolic b
fromComplexConstant c = symbFun "complex" [real, imag]
  where
    real = Symbolic $ fromRational $ toRational $ Complex.realPart c
    imag = Symbolic $ fromRational $ toRational $ Complex.imagPart c

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

data SymbFloat
data SymbDouble
data SymbComplexFloat
data SymbComplexDouble

class SymbParam a
  where
    symbType :: a -> String

instance SymbParam SymbFloat  where symbType _ = "float"
instance SymbParam SymbDouble where symbType _ = "double"
instance SymbParam SymbComplexFloat  where symbType _ = "cfloat"
instance SymbParam SymbComplexDouble where symbType _ = "cdouble"

instance SymbParam a => Num (Symbolic a)
  where
    fromInteger = Symbolic . fromInteger
    x + y  = symbFun "+" [x, y]
    x - y  = symbFun "-" [x, y]
    x * y  = symbFun "*" [x, y]
    abs    = smtAbs
    signum = smtSignum

instance SymbParam a => Fractional (Symbolic a) where
  fromRational = Symbolic . fromRational
  x / y = symbFun "/" [x, y]

instance SymbParam a => Floating (Symbolic a) where
  pi      = fromRational (toRational pi)
  exp x   = symbFun "exp" [x]
  log x   = symbFun "log" [x]
  sqrt x  = symbFun "sqrt" [x]
  x ** y  = symbFun "pow" [x, y]
  sin x   = symbFun "sin" [x]
  cos x   = symbFun "cos" [x]
  tan x   = symbFun "tan" [x]
  asin x  = symbFun "asin" [x]
  acos x  = symbFun "acos" [x]
  atan x  = symbFun "atan" [x]
  sinh x  = symbFun "sinh" [x]
  cosh x  = symbFun "cosh" [x]
  tanh x  = symbFun "tanh" [x]
  asinh x = symbFun "asinh" [x]
  acosh x = symbFun "acosh" [x]
  atanh x = symbFun "atanh" [x]

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

class Floating a => Complex a
  where
    type RealPart a
    complex   :: RealPart a -> RealPart a -> a
    polar     :: RealPart a -> RealPart a -> a
    real      :: a -> RealPart a
    imag      :: a -> RealPart a
    magnitude :: a -> RealPart a
    phase     :: a -> RealPart a
    conjugate :: a -> a

instance Complex (V.SMTExpr Prim (Complex.Complex Float))
  where
    type RealPart (V.SMTExpr Prim (Complex.Complex Float)) = V.SMTExpr Prim Float
    complex   (Float x) (Float y) = ComplexFloat (complex x y)
    polar     (Float x) (Float y) = ComplexFloat (polar x y)
    real      (ComplexFloat x)    = Float (real x)
    imag      (ComplexFloat x)    = Float (imag x)
    magnitude (ComplexFloat x)    = Float (magnitude x)
    phase     (ComplexFloat x)    = Float (phase x)
    conjugate (ComplexFloat x)    = ComplexFloat (conjugate x)

instance Complex (V.SMTExpr Prim (Complex.Complex Double))
  where
    type RealPart (V.SMTExpr Prim (Complex.Complex Double)) = V.SMTExpr Prim Double
    complex   (Double x) (Double y) = ComplexDouble (complex x y)
    polar     (Double x) (Double y) = ComplexDouble (polar x y)
    real      (ComplexDouble x)     = Double (real x)
    imag      (ComplexDouble x)     = Double (imag x)
    magnitude (ComplexDouble x)     = Double (magnitude x)
    phase     (ComplexDouble x)     = Double (phase x)
    conjugate (ComplexDouble x)     = ComplexDouble (conjugate x)

witnessComplex :: (SoftwarePrimType a, SoftwarePrimType (Complex.Complex a)) =>
  Prim (Complex.Complex a) ->
  Dict ( Floating (V.SMTExpr Prim a)
       , Complex  (V.SMTExpr Prim (Complex.Complex a))
       , RealPart (V.SMTExpr Prim (Complex.Complex a)) ~ V.SMTExpr Prim a)
witnessComplex (_ :: Prim (Complex.Complex a)) =
  case softwareRep :: SoftwarePrimTypeRep (Complex.Complex a) of
    ComplexFloatST  -> Dict
    ComplexDoubleST -> Dict

witnessFractional :: (SoftwarePrimType a, Fractional a) =>
  Prim a ->
  Dict (Floating (V.SMTExpr Prim a))
witnessFractional (_ :: Prim a) = case softwareRep :: SoftwarePrimTypeRep a of
    FloatST         -> Dict
    DoubleST        -> Dict
    ComplexFloatST  -> Dict
    ComplexDoubleST -> Dict

witnessIntegral :: (SoftwarePrimType a, Integral a) =>
  Prim a ->
  Dict (Integral (V.SMTExpr Prim a), Bits (V.SMTExpr Prim a))
witnessIntegral (_ :: Prim a) = case softwareRep :: SoftwarePrimTypeRep a of
  Int8ST   -> Dict
  Int16ST  -> Dict
  Int32ST  -> Dict
  Int64ST  -> Dict
  Word8ST  -> Dict
  Word16ST -> Dict
  Word32ST -> Dict
  Word64ST -> Dict

witnessBits :: (SoftwarePrimType a, Bits.Bits a) =>
  Prim a ->
  Dict ( Num a
       , Integral (V.SMTExpr Prim a)
       , Bits (V.SMTExpr Prim a))
witnessBits (_ :: Prim a) = case softwareRep :: SoftwarePrimTypeRep a of
  Int8ST   -> Dict
  Int16ST  -> Dict
  Int32ST  -> Dict
  Int64ST  -> Dict
  Word8ST  -> Dict
  Word16ST -> Dict
  Word32ST -> Dict
  Word64ST -> Dict

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

toRat :: (SoftwarePrimType a, Fractional a) => V.SMTExpr Prim a -> Rat
toRat (x :: V.SMTExpr Prim a) = case softwareRep :: SoftwarePrimTypeRep a of
  FloatST         -> let Float         (Symbolic y) = x in y
  DoubleST        -> let Double        (Symbolic y) = x in y
  ComplexFloatST  -> let ComplexFloat  (Symbolic y) = x in y
  ComplexDoubleST -> let ComplexDouble (Symbolic y) = x in y

fromRat :: forall a. (SoftwarePrimType a, Num a) => Rat -> V.SMTExpr Prim a
fromRat x = case softwareRep :: SoftwarePrimTypeRep a of Int8ST -> Int8 (f2i x)
  where
    f2i :: forall s w. (Sign s, Width w) => Rat -> BV s w
    f2i (Rat x) = BV $ SMT.List [SMT.fam "int2bv" [width (undefined :: w)], SMT.fun "to_int" [x]]

i2n :: forall a b.
  ( V.SMTEval Prim a, SoftwarePrimType a, Integral a
  , V.SMTEval Prim b, SoftwarePrimType b, Num b) =>
  V.SMTExpr Prim a ->
  V.SMTExpr Prim b
i2n x = toBV x $ case softwareRep :: SoftwarePrimTypeRep b of
  Int8ST   -> Int8   . i2i
  Int16ST  -> Int16  . i2i
  Int32ST  -> Int32  . i2i
  Int64ST  -> Int64  . i2i
  Word8ST  -> Word8  . i2i
  Word16ST -> Word16 . i2i
  Word32ST -> Word32 . i2i
  Word64ST -> Word64 . i2i
  FloatST  -> Float  . Symbolic . i2f
  DoubleST -> Double . Symbolic . i2f
  ComplexFloatST  -> ComplexFloat  . Symbolic . i2f
  ComplexDoubleST -> ComplexDouble . Symbolic . i2f
  where
    toBV :: forall a b. (V.SMTEval Prim a, SoftwarePrimType a, Integral a) =>
      V.SMTExpr Prim a -> (forall s w. (Sign s, Width w) => BV s w -> b) -> b
    toBV (x :: V.SMTExpr Prim a) k = case softwareRep :: SoftwarePrimTypeRep a of
      Int8ST   -> let Int8   y = x in k y
      Int16ST  -> let Int16  y = x in k y
      Int32ST  -> let Int32  y = x in k y
      Int64ST  -> let Int64  y = x in k y
      Word8ST  -> let Word8  y = x in k y
      Word16ST -> let Word16 y = x in k y
      Word32ST -> let Word32 y = x in k y
      Word64ST -> let Word64 y = x in k y

    i2f :: (Sign s, Width w) => BV s w -> Rat
    i2f (BV x) = Rat (SMT.fun "to_real" [SMT.fun "bv2int" [x]])

    i2i :: forall s1 w1 s2 w2. (Sign s1, Width w1, Sign s2, Width w2) => BV s1 w1 -> BV s2 w2
    i2i x = case compare m n of
      LT | isSigned x -> V.fromSMT (SMT.signExtend (n-m) (V.toSMT x))
         | otherwise  -> V.fromSMT (SMT.zeroExtend (n-m) (V.toSMT x))
      EQ -> V.fromSMT (V.toSMT x)
      GT -> V.fromSMT (SMT.extract (V.toSMT x) (n-1) 0)
      where
        m = width (undefined :: w1)
        n = width (undefined :: w2)

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

class SymbParam a => SymbComplex a
  where
    type SymbRealPart a

instance SymbComplex SymbComplexFloat
  where
    type SymbRealPart SymbComplexFloat  = SymbFloat

instance SymbComplex SymbComplexDouble
  where
    type SymbRealPart SymbComplexDouble = SymbDouble

instance SymbComplex a => Complex (Symbolic a)
  where
    type RealPart (Symbolic a) = Symbolic (SymbRealPart a)
    complex x y = symbFun "complex" [symbCast x, symbCast y]
    polar x y   = symbFun "polar" [symbCast x, symbCast y]
    real x      = symbCast (symbFun "real" [x])
    imag x      = symbCast (symbFun "imag" [x])
    magnitude x = symbCast (symbFun "magnitude" [x])
    phase x     = symbCast (symbFun "phase" [x])
    conjugate x = symbFun "conjugate" [x]

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

declareSymbFun :: SymbParam a => String -> a ->
  [SMT.SExpr] -> SMT.SExpr -> SMT.SMT ()
declareSymbFun name (_ :: a) args res =
  S.void $ SMT.declareFun (symbType (undefined :: a) ++ "-" ++ name) args res

declareSymbArith :: SymbParam a => a -> SMT.SMT ()
declareSymbArith x = do
  declareSymbFun "+" x     [SMT.tReal, SMT.tReal] SMT.tReal
  declareSymbFun "-" x     [SMT.tReal, SMT.tReal] SMT.tReal
  declareSymbFun "*" x     [SMT.tReal, SMT.tReal] SMT.tReal
  declareSymbFun "/" x     [SMT.tReal, SMT.tReal] SMT.tReal
  declareSymbFun "exp" x   [SMT.tReal] SMT.tReal
  declareSymbFun "log" x   [SMT.tReal] SMT.tReal
  declareSymbFun "sqrt" x  [SMT.tReal] SMT.tReal
  declareSymbFun "pow" x   [SMT.tReal, SMT.tReal] SMT.tReal
  declareSymbFun "sin" x   [SMT.tReal] SMT.tReal
  declareSymbFun "cos" x   [SMT.tReal] SMT.tReal
  declareSymbFun "tan" x   [SMT.tReal] SMT.tReal
  declareSymbFun "asin" x  [SMT.tReal] SMT.tReal
  declareSymbFun "acos" x  [SMT.tReal] SMT.tReal
  declareSymbFun "atan" x  [SMT.tReal] SMT.tReal
  declareSymbFun "sinh" x  [SMT.tReal] SMT.tReal
  declareSymbFun "cosh" x  [SMT.tReal] SMT.tReal
  declareSymbFun "tanh" x  [SMT.tReal] SMT.tReal
  declareSymbFun "asinh" x [SMT.tReal] SMT.tReal
  declareSymbFun "acosh" x [SMT.tReal] SMT.tReal
  declareSymbFun "atanh" x [SMT.tReal] SMT.tReal

declareSymbComplex :: SymbParam a => a -> SMT.SMT ()
declareSymbComplex x = do
  declareSymbArith x
  declareSymbFun "complex" x   [SMT.tReal, SMT.tReal] SMT.tReal
  declareSymbFun "polar" x     [SMT.tReal, SMT.tReal] SMT.tReal
  declareSymbFun "real" x      [SMT.tReal] SMT.tReal
  declareSymbFun "imag" x      [SMT.tReal] SMT.tReal
  declareSymbFun "magnitude" x [SMT.tReal] SMT.tReal
  declareSymbFun "phase" x     [SMT.tReal] SMT.tReal
  declareSymbFun "conjugate" x [SMT.tReal] SMT.tReal

declareFeldsparGlobals :: SMT.SMT ()
declareFeldsparGlobals = do
  declareSymbArith   (undefined :: SymbFloat)
  declareSymbArith   (undefined :: SymbDouble)
  declareSymbComplex (undefined :: SymbComplexFloat)
  declareSymbComplex (undefined :: SymbComplexDouble)
  SMT.declareFun "skolem-int8"   [] (SMT.tBits 8)
  SMT.declareFun "skolem-int16"  [] (SMT.tBits 16)
  SMT.declareFun "skolem-int32"  [] (SMT.tBits 32)
  SMT.declareFun "skolem-int64"  [] (SMT.tBits 64)
  SMT.declareFun "skolem-word8"  [] (SMT.tBits 8)
  SMT.declareFun "skolem-word16" [] (SMT.tBits 16)
  SMT.declareFun "skolem-word32" [] (SMT.tBits 32)
  SMT.declareFun "skolem-word64" [] (SMT.tBits 64)
  return ()

instance FO.Substitute Prim
  where
    type SubstPred Prim = SoftwarePrimType
    subst sub (Prim exp) = Prim (everywhereUp f exp)
      where
        f :: ASTF SoftwarePrimDomain a -> ASTF SoftwarePrimDomain a
        f (Sym (FreeVar x :&: ty)) = case FO.lookupSubst sub (Imp.ValComp x) of
          Imp.ValComp y -> Sym (FreeVar y :&: ty)
          Imp.ValRun  z -> Sym (Lit z :&: ty)
        f (Sym (ArrIx iarr :&: ty) :$ i) = Sym (ArrIx iarr' :&: ty) :$ i
          where iarr' = FO.lookupSubst sub iarr
        f x = x

instance FO.TypeablePred SoftwarePrimType
  where
    witnessTypeable Dict = Dict

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

instance V.SMTEval Prim Bool where
  fromConstant = Bool . SMT.bool
  witnessOrd _ = Dict

instance V.SMTEval Prim Int8 where
  fromConstant = Int8 . fromIntegral
  witnessNum _ = Dict
  witnessOrd _ = Dict
  skolemIndex  = V.fromSMT (SMT.fun "skolem-int8" [])

instance V.SMTEval Prim Int16 where
  fromConstant = Int16 . fromIntegral
  witnessNum _ = Dict
  witnessOrd _ = Dict
  skolemIndex  = V.fromSMT (SMT.fun "skolem-int16" [])

instance V.SMTEval Prim Int32 where
  fromConstant = Int32 . fromIntegral
  witnessNum _ = Dict
  witnessOrd _ = Dict
  skolemIndex  = V.fromSMT (SMT.fun "skolem-int32" [])

instance V.SMTEval Prim Int64 where
  fromConstant = Int64 . fromIntegral
  witnessNum _ = Dict
  witnessOrd _ = Dict
  skolemIndex  = V.fromSMT (SMT.fun "skolem-int64" [])

instance V.SMTEval Prim Word8 where
  fromConstant = Word8 . fromIntegral
  witnessNum _ = Dict
  witnessOrd _ = Dict
  skolemIndex  = V.fromSMT (SMT.fun "skolem-word8" [])

instance V.SMTEval Prim Word16 where
  fromConstant = Word16 . fromIntegral
  witnessNum _ = Dict
  witnessOrd _ = Dict
  skolemIndex  = V.fromSMT (SMT.fun "skolem-word16" [])

instance V.SMTEval Prim Word32 where
  fromConstant = Word32 . fromIntegral
  witnessNum _ = Dict
  witnessOrd _ = Dict
  skolemIndex  = V.fromSMT (SMT.fun "skolem-word32" [])

instance V.SMTEval Prim Word64 where
  fromConstant = Word64 . fromIntegral
  witnessNum _ = Dict
  witnessOrd _ = Dict
  skolemIndex  = V.fromSMT (SMT.fun "skolem-word64" [])

instance V.SMTEval Prim Float where
  fromConstant = Float . fromRational . toRational
  witnessOrd _ = Dict
  witnessNum _ = Dict

instance V.SMTEval Prim Double where
  fromConstant = Double . fromRational . toRational
  witnessOrd _ = Dict
  witnessNum _ = Dict

instance V.SMTEval Prim (Complex.Complex Float) where
  fromConstant = ComplexFloat . fromComplexConstant
  witnessNum _ = Dict

instance V.SMTEval Prim (Complex.Complex Double) where
  fromConstant = ComplexDouble . fromComplexConstant
  witnessNum _ = Dict

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

instance V.SMTEval1 Prim where
  type Pred Prim = SoftwarePrimType
  newtype SMTExpr Prim Bool   = Bool SMT.SExpr
    deriving (Typeable)
  newtype SMTExpr Prim Float  = Float  (Symbolic SymbFloat)
    deriving (Typeable, Num, Fractional, Floating, V.SMTOrd, V.TypedSExpr)
  newtype SMTExpr Prim Double = Double (Symbolic SymbDouble)
    deriving (Typeable, Num, Fractional, Floating, V.SMTOrd, V.TypedSExpr)
  newtype SMTExpr Prim (Complex.Complex Float) =
      ComplexFloat (Symbolic SymbComplexFloat)
    deriving (Typeable, Num, Fractional, Floating, V.TypedSExpr)
  newtype SMTExpr Prim (Complex.Complex Double) =
      ComplexDouble (Symbolic SymbComplexDouble)
    deriving (Typeable, Num, Fractional, Floating, V.TypedSExpr)
  newtype SMTExpr Prim Int8   = Int8   (BV Signed W8)
    deriving (Typeable, Num, Real, Enum, Integral, Bits, V.SMTOrd, V.TypedSExpr)
  newtype SMTExpr Prim Int16  = Int16  (BV Signed W16)
    deriving (Typeable, Num, Real, Enum, Integral, Bits, V.SMTOrd, V.TypedSExpr)
  newtype SMTExpr Prim Int32  = Int32  (BV Signed W32)
    deriving (Typeable, Num, Real, Enum, Integral, Bits, V.SMTOrd, V.TypedSExpr)
  newtype SMTExpr Prim Int64  = Int64  (BV Signed W64)
    deriving (Typeable, Num, Real, Enum, Integral, Bits, V.SMTOrd, V.TypedSExpr)
  newtype SMTExpr Prim Word8  = Word8  (BV Unsigned W8)
    deriving (Typeable, Num, Real, Enum, Integral, Bits, V.SMTOrd, V.TypedSExpr)
  newtype SMTExpr Prim Word16 = Word16 (BV Unsigned W16)
    deriving (Typeable, Num, Real, Enum, Integral, Bits, V.SMTOrd, V.TypedSExpr)
  newtype SMTExpr Prim Word32 = Word32 (BV Unsigned W32)
    deriving (Typeable, Num, Real, Enum, Integral, Bits, V.SMTOrd, V.TypedSExpr)
  newtype SMTExpr Prim Word64 = Word64 (BV Unsigned W64)
    deriving (Typeable, Num, Real, Enum, Integral, Bits, V.SMTOrd, V.TypedSExpr)

  eval (Prim exp :: Prim a) =
    simpleMatch (\(exp :&: ty) -> case softwarePrimWitType ty of
      Dict -> case V.witnessPred (undefined :: Prim a) of
        Dict -> verifyPrim exp) exp

  witnessPred (_ :: Prim a) = case softwareRep :: SoftwarePrimTypeRep a of
    BoolST   -> Dict
    Int8ST   -> Dict
    Int16ST  -> Dict
    Int32ST  -> Dict
    Int64ST  -> Dict
    Word8ST  -> Dict
    Word16ST -> Dict
    Word32ST -> Dict
    Word64ST -> Dict
    FloatST  -> Dict
    DoubleST -> Dict
    ComplexFloatST  -> Dict
    ComplexDoubleST -> Dict

instance V.SMTOrd (V.SMTExpr Prim Bool) where
  Bool x .<.  Bool y = SMT.not x SMT..&&. y
  Bool x .>.  Bool y = x SMT..&&. SMT.not y
  Bool x .<=. Bool y = SMT.not x SMT..||. y
  Bool x .>=. Bool y = x SMT..||. SMT.not y

instance V.TypedSExpr (V.SMTExpr Prim Bool) where
  smtType _ = SMT.tBool
  toSMT (Bool x) = x
  fromSMT x = Bool x

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

verifyPrim :: forall a .
  (SoftwarePrimType (DenResult a), V.SMTEval Prim (DenResult a), HasCallStack) =>
  SoftwarePrim a ->
  Args (AST SoftwarePrimDomain) a ->
  V.Verify (V.SMTExpr Prim (DenResult a))
verifyPrim (FreeVar x) _ = peekValue x
verifyPrim (Lit x) _     = return (V.fromConstant x)
verifyPrim Add (x :* y :* Nil)
  | Dict <- V.witnessNum (undefined :: Prim (DenResult a)) =
    S.liftM2 (+) (V.eval (Prim x)) (V.eval (Prim y))
verifyPrim Sub (x :* y :* Nil)
  | Dict <- V.witnessNum (undefined :: Prim (DenResult a)) =
    S.liftM2 (-) (V.eval (Prim x)) (V.eval (Prim y))
verifyPrim Mul (x :* y :* Nil)
  | Dict <- V.witnessNum (undefined :: Prim (DenResult a)) =
    S.liftM2 (*) (V.eval (Prim x)) (V.eval (Prim y))
verifyPrim Neg (x :* Nil)
  | Dict <- V.witnessNum (undefined :: Prim (DenResult a)) =
    fmap negate (V.eval (Prim x))
verifyPrim Abs (x :* Nil)
  | Dict <- V.witnessNum (undefined :: Prim (DenResult a)) =
    fmap abs (V.eval (Prim x))
verifyPrim Sign (x :* Nil)
  | Dict <- V.witnessNum (undefined :: Prim (DenResult a)) =
    fmap signum (V.eval (Prim x))
verifyPrim Div (x :* y :* Nil)
  | Dict <- witnessIntegral (undefined :: Prim (DenResult a)) =
    S.liftM2 div (V.eval (Prim x)) (V.eval (Prim y))
verifyPrim Mod (x :* y :* Nil)
  | Dict <- witnessIntegral (undefined :: Prim (DenResult a)) =
    S.liftM2 mod (V.eval (Prim x)) (V.eval (Prim y))
verifyPrim Quot (x :* y :* Nil)
  | Dict <- witnessIntegral (undefined :: Prim (DenResult a)) =
    S.liftM2 quot (V.eval (Prim x)) (V.eval (Prim y))
verifyPrim Rem (x :* y :* Nil)
  | Dict <- witnessIntegral (undefined :: Prim (DenResult a)) =
    S.liftM2 rem (V.eval (Prim x)) (V.eval (Prim y))
verifyPrim FDiv (x :* y :* Nil)
  | Dict <- witnessFractional (undefined :: Prim (DenResult a)) =
    S.liftM2 (/) (V.eval (Prim x)) (V.eval (Prim y))
verifyPrim Pi Nil
  | Dict <- witnessFractional (undefined :: Prim (DenResult a)) =
    return pi
verifyPrim Exp (x :* Nil)
  | Dict <- witnessFractional (undefined :: Prim (DenResult a)) =
    fmap exp (V.eval (Prim x))
verifyPrim Log (x :* Nil)
  | Dict <- witnessFractional (undefined :: Prim (DenResult a)) =
    fmap log (V.eval (Prim x))
verifyPrim Sqrt (x :* Nil)
  | Dict <- witnessFractional (undefined :: Prim (DenResult a)) =
    fmap sqrt (V.eval (Prim x))
verifyPrim Pow (x :* y :* Nil)
  | Dict <- witnessFractional (undefined :: Prim (DenResult a)) =
    S.liftM2 (**) (V.eval (Prim x)) (V.eval (Prim y))
verifyPrim Sin (x :* Nil)
  | Dict <- witnessFractional (undefined :: Prim (DenResult a)) =
    fmap sin (V.eval (Prim x))
verifyPrim Cos (x :* Nil)
  | Dict <- witnessFractional (undefined :: Prim (DenResult a)) =
    fmap cos (V.eval (Prim x))
verifyPrim Tan (x :* Nil)
  | Dict <- witnessFractional (undefined :: Prim (DenResult a)) =
    fmap tan (V.eval (Prim x))
verifyPrim Asin (x :* Nil)
  | Dict <- witnessFractional (undefined :: Prim (DenResult a)) =
    fmap asin (V.eval (Prim x))
verifyPrim Acos (x :* Nil)
  | Dict <- witnessFractional (undefined :: Prim (DenResult a)) =
    fmap acos (V.eval (Prim x))
verifyPrim Atan (x :* Nil)
  | Dict <- witnessFractional (undefined :: Prim (DenResult a)) =
    fmap atan (V.eval (Prim x))
verifyPrim Sinh (x :* Nil)
  | Dict <- witnessFractional (undefined :: Prim (DenResult a)) =
    fmap sinh (V.eval (Prim x))
verifyPrim Cosh (x :* Nil)
  | Dict <- witnessFractional (undefined :: Prim (DenResult a)) =
    fmap cosh (V.eval (Prim x))
verifyPrim Tanh (x :* Nil)
  | Dict <- witnessFractional (undefined :: Prim (DenResult a)) =
    fmap tanh (V.eval (Prim x))
verifyPrim Asinh (x :* Nil)
  | Dict <- witnessFractional (undefined :: Prim (DenResult a)) =
    fmap asinh (V.eval (Prim x))
verifyPrim Acosh (x :* Nil)
  | Dict <- witnessFractional (undefined :: Prim (DenResult a)) =
    fmap acosh (V.eval (Prim x))
verifyPrim Atanh (x :* Nil)
  | Dict <- witnessFractional (undefined :: Prim (DenResult a)) =
    fmap atanh (V.eval (Prim x))
verifyPrim Complex (x :* y :* Nil)
  | Dict <- witnessComplex (undefined :: Prim (DenResult a)) =
    S.liftM2 complex (V.eval (Prim x)) (V.eval (Prim y))
verifyPrim Polar (x :* y :* Nil)
  | Dict <- witnessComplex (undefined :: Prim (DenResult a)) =
    S.liftM2 polar (V.eval (Prim x)) (V.eval (Prim y))
verifyPrim Real ((x :: ASTF SoftwarePrimDomain b) :* Nil)
  | Dict <- witnessComplex (undefined :: Prim b) =
    fmap real (V.eval (Prim x))
verifyPrim Imag ((x :: ASTF SoftwarePrimDomain b) :* Nil)
  | Dict <- witnessComplex (undefined :: Prim b) =
    fmap imag (V.eval (Prim x))
verifyPrim Magnitude ((x :: ASTF SoftwarePrimDomain b) :* Nil)
  | Dict <- witnessComplex (undefined :: Prim b) =
    fmap magnitude (V.eval (Prim x))
verifyPrim Phase ((x :: ASTF SoftwarePrimDomain b) :* Nil)
  | Dict <- witnessComplex (undefined :: Prim b) =
    fmap phase (V.eval (Prim x))
verifyPrim Conjugate ((x :: ASTF SoftwarePrimDomain b) :* Nil)
  | Dict <- witnessComplex (undefined :: Prim b) =
    fmap conjugate (V.eval (Prim x))
verifyPrim I2N ((x :: ASTF SoftwarePrimDomain b) :* Nil)
  | Dict <- V.witnessPred (undefined :: Prim b),
    Dict <- V.witnessNum (undefined :: Prim b) = do
    fmap i2n (V.eval (Prim x))
verifyPrim I2B ((x :: ASTF SoftwarePrimDomain b) :* Nil)
  | Dict <- V.witnessPred (undefined :: Prim b),
    Dict <- V.witnessNum (undefined :: Prim b) = do
    x <- V.eval (Prim x)
    return (V.fromSMT (SMT.not (x V..==. 0)))
verifyPrim B2I (x :* Nil)
  | Dict <- V.witnessNum (undefined :: Prim (DenResult a)) = do
    x <- V.eval (Prim x)
    return (V.smtIte (V.toSMT x) 1 0)
verifyPrim Round ((x :: ASTF SoftwarePrimDomain b) :* Nil) = do
  x <- V.eval (Prim x)
  return (fromRat (toRat x))
verifyPrim Not (x :* Nil) =
  fmap (V.fromSMT . SMT.not . V.toSMT) (V.eval (Prim x))
verifyPrim And (x :* y :* Nil) = do
  x <- V.eval (Prim x)
  y <- V.eval (Prim y)
  return (V.fromSMT (V.toSMT x SMT..&&. V.toSMT y))
verifyPrim Or (x :* y :* Nil) = do
  x <- V.eval (Prim x)
  y <- V.eval (Prim y)
  return (V.fromSMT (V.toSMT x SMT..||. V.toSMT y))
verifyPrim Eq ((x :: ASTF SoftwarePrimDomain b) :* y :* Nil)
  | Dict <- V.witnessPred (undefined :: Prim b) =
    fmap V.fromSMT (S.liftM2 (V..==.) (V.eval (Prim x)) (V.eval (Prim y)))
verifyPrim Neq ((x :: ASTF SoftwarePrimDomain b) :* y :* Nil)
  | Dict <- V.witnessPred (undefined :: Prim b) =
    fmap V.fromSMT (S.liftM2 (./=.) (V.eval (Prim x)) (V.eval (Prim y)))
  where
    x ./=. y = SMT.not (x V..==. y)
verifyPrim Lt ((x :: ASTF SoftwarePrimDomain b) :* y :* Nil)
  | Dict <- V.witnessPred (undefined :: Prim b),
    Dict <- V.witnessOrd  (undefined :: Prim b) =
    fmap V.fromSMT (S.liftM2 (V..<.) (V.eval (Prim x)) (V.eval (Prim y)))
verifyPrim Gt ((x :: ASTF SoftwarePrimDomain b) :* y :* Nil)
  | Dict <- V.witnessPred (undefined :: Prim b),
    Dict <- V.witnessOrd  (undefined :: Prim b) =
    fmap V.fromSMT (S.liftM2 (V..>.) (V.eval (Prim x)) (V.eval (Prim y)))
verifyPrim Lte ((x :: ASTF SoftwarePrimDomain b) :* y :* Nil)
  | Dict <- V.witnessPred (undefined :: Prim b),
    Dict <- V.witnessOrd  (undefined :: Prim b) =
    fmap V.fromSMT (S.liftM2 (V..<=.) (V.eval (Prim x)) (V.eval (Prim y)))
verifyPrim Gte ((x :: ASTF SoftwarePrimDomain b) :* y :* Nil)
  | Dict <- V.witnessPred (undefined :: Prim b),
    Dict <- V.witnessOrd  (undefined :: Prim b) =
    fmap V.fromSMT (S.liftM2 (V..>=.) (V.eval (Prim x)) (V.eval (Prim y)))
verifyPrim BitAnd (x :* y :* Nil)
  | Dict <- witnessBits (undefined :: Prim (DenResult a)) =
    S.liftM2 (.&.) (V.eval (Prim x)) (V.eval (Prim y))
verifyPrim BitOr (x :* y :* Nil)
  | Dict <- witnessBits (undefined :: Prim (DenResult a)) =
    S.liftM2 (.|.) (V.eval (Prim x)) (V.eval (Prim y))
verifyPrim BitXor (x :* y :* Nil)
  | Dict <- witnessBits (undefined :: Prim (DenResult a)) =
    S.liftM2 xor (V.eval (Prim x)) (V.eval (Prim y))
verifyPrim BitCompl (x :* Nil)
  | Dict <- witnessBits (undefined :: Prim (DenResult a)) =
    fmap complement (V.eval (Prim x))
verifyPrim ShiftL (x :* (y :: ASTF SoftwarePrimDomain b) :* Nil)
  | Dict <- witnessBits (undefined :: Prim (DenResult a)),
    Dict <- V.witnessPred (undefined :: Prim b),
    Dict <- witnessIntegral (undefined :: Prim b) = do
    -- todo: should check for undefined behaviour
    x <- V.eval (Prim x)
    y <- V.eval (Prim y)
    return (shiftL x (i2n y))
verifyPrim ShiftR (x :* (y :: ASTF SoftwarePrimDomain b) :* Nil)
  | Dict <- witnessBits (undefined :: Prim (DenResult a)),
    Dict <- V.witnessPred (undefined :: Prim b),
    Dict <- witnessIntegral (undefined :: Prim b) = do
    -- todo: should check for undefined behaviour
    x <- V.eval (Prim x)
    y <- V.eval (Prim y)
    return (shiftR x (i2n y))
verifyPrim (ArrIx (Imp.IArrComp name :: Imp.IArr Index b)) (i :* Nil) = do
  i <- V.eval (Prim i)
  readArray name i
verifyPrim Cond (cond :* x :* y :* Nil) =
  S.liftM3 V.smtIte
    (fmap V.toSMT (V.eval (Prim cond)))
    (V.eval (Prim x))
    (V.eval (Prim y))
verifyPrim exp _ = error ("Unimplemented: " ++ show exp)

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


{-# LANGUAGE ConstraintKinds      #-}
{-# LANGUAGE DataKinds            #-}
{-# LANGUAGE DeriveGeneric        #-}
{-# LANGUAGE FlexibleContexts     #-}
{-# LANGUAGE FlexibleInstances         #-}
{-# LANGUAGE GADTs                     #-}
{-# LANGUAGE OverloadedStrings         #-}
{-# LANGUAGE PolyKinds                 #-}
{-# LANGUAGE QuasiQuotes               #-}
{-# LANGUAGE RankNTypes                #-}
{-# LANGUAGE ScopedTypeVariables       #-}
{-# LANGUAGE TemplateHaskell           #-}
{-# LANGUAGE TypeApplications          #-}
{-# LANGUAGE TypeFamilies              #-}
{-# LANGUAGE TypeOperators             #-}
{-# LANGUAGE UndecidableInstances      #-}
{-# LANGUAGE AllowAmbiguousTypes       #-}
{-# LANGUAGE NoMonomorphismRestriction #-}
{-# LANGUAGE PartialTypeSignatures     #-}
--{-# LANGUAGE MonomorphismRestriction #-}
module Main where

import           DataSources
import qualified Frames.Utils               as FU
import qualified Frames.Folds               as FF
import qualified Frames.KMeans              as KM
import qualified Frames.Regression          as FR
import qualified Frames.MaybeUtils          as FM
import qualified Frames.VegaLite            as FV
import qualified Frames.Transform           as FT
import qualified Frames.Table               as Table
import qualified Math.Rescale               as MR
import qualified Frames.MapReduce           as MR

import qualified Knit.Report                as K
import qualified Knit.Effect.RandomFu       as KR
import qualified Knit.Effect.Pandoc         as K (newPandoc, NamedDoc (..))
import qualified Knit.Report.Other.Lucid    as KL
{-
import qualified Polysemy             as PS
import qualified Knit.Effects.Logger      as Log
import qualified Knit.Effects.PandocMonad as PM
import qualified Knit.Effects.Pandoc      as PE
import           Knit.Effects.RandomFu      (Random, runRandomIOPureMT)
import           Knit.Effects.Docs        (toNamedDocListWithM)
--import qualified Knit.Report.Blaze            as RB
import qualified Knit.Report.Pandoc              as RP
import qualified Knit.Report.Lucid            as RL
-}

import qualified Control.Foldl              as FL
import           Control.Monad.IO.Class     (MonadIO, liftIO)
import           Control.Monad              (sequence)
import           Control.Lens ((%~),(^.))
import           Data.Bool                  (bool)
import qualified Data.Map                   as M
import           Data.Maybe                 (catMaybes, fromJust, fromMaybe)
import           Data.Monoid                ((<>))
import qualified Data.Monoid                as MO
import           Data.Proxy                 (Proxy (..))
import qualified Data.Text                  as T
import qualified Data.Text.IO               as T
import qualified Data.Text.Lazy             as TL
import qualified Data.Vector                as V
import qualified Data.Vinyl                 as V
import qualified Data.Vinyl.Functor         as V
import qualified Data.Vinyl.Class.Method    as V
import qualified Data.Vinyl.TypeLevel       as V
import           Data.Vinyl.Lens            (type (∈))
import           Frames                     ((:.), (&:))
import qualified Frames                     as F
import qualified Frames.CSV                 as F
import qualified Frames.Melt                 as F
import qualified Frames.InCore              as FI
import qualified Frames.TH                  as F
import qualified Graphics.Vega.VegaLite     as GV
import qualified Pipes                      as P
import qualified Pipes.Prelude              as P
import qualified Lucid                      as HL
import qualified Text.Blaze.Html.Renderer.Text as BH
import           Data.Random.Source.PureMT as R
import           Data.Random as R
import qualified System.Clock as C
import qualified Statistics.Types as ST
-- stage restriction means this all has to be up top
F.tableTypes' (F.rowGen fipsByCountyFP) { F.rowTypeName = "FIPSByCountyRenamed", F.columnNames = ["fips","County","State"]}

type CO_AnalysisVERA_Cols = [Year, State, TotalPop, TotalJailAdm, TotalJailPop, TotalPrisonAdm, TotalPrisonPop]

justsFromRec :: V.RMap fs => F.Record fs -> F.Rec (Maybe :. F.ElField) fs
justsFromRec = V.rmap (V.Compose . Just)

rDiv :: (Real a, Real b, Fractional c) => a -> b -> c
rDiv a b = realToFrac a/realToFrac b

type instance FI.VectorFor (Maybe a) = V.Vector

type DblX = "X" F.:-> Double
type DblY = "Y" F.:-> Double

templateVars = M.fromList
  [
    ("lang", "English")
  , ("author", "Adam Conner-Sax")
  , ("pagetitle", "Colorado Incarceration Data Analysis")
  , ("tufte","True")
  ]

main :: IO ()
main = do
  -- create streams which are filtered to CO
  let writeNamedHtml (K.NamedDoc n lt) = T.writeFile (T.unpack $ "analysis/" <> n <> ".html") $ TL.toStrict lt
      writeAllHtml = fmap (const ()) . traverse writeNamedHtml
      pandocWriterConfig = K.PandocWriterConfig (Just "pandoc-templates/minWithVega-pandoc.html")  templateVars K.mindocOptionsF
  startReal <- C.getTime C.Monotonic
  eitherDocs <- K.knitHtmls (Just "colorado-analysis.Main") K.logAll pandocWriterConfig $ KR.runRandomIOPureMT (pureMT 1) $ do
    K.logLE K.Info "Creating data producers from CSV files"
    let parserOptions = F.defaultParser { F.quotingMode =  F.RFC4180Quoting ' ' }
        veraData :: F.MonadSafe m => P.Producer (FM.MaybeRow IncarcerationTrends)  m ()
        veraData = F.readTableMaybeOpt F.defaultParser veraTrendsFP  P.>-> P.filter (FU.filterOnMaybeField @State (=="CO"))
        povertyData :: F.MonadSafe m => P.Producer SAIPE m ()
        povertyData = F.readTableOpt parserOptions censusSAIPE_FP P.>-> P.filter (FU.filterOnField @Abbreviation (== "CO"))
        fipsByCountyData :: F.MonadSafe m => P.Producer FIPSByCountyRenamed m ()
        fipsByCountyData = F.readTableOpt parserOptions fipsByCountyFP  P.>-> P.filter (FU.filterOnField @State (== "CO"))
        -- NB: This data has 2 rows per county, one for misdemeanors, one for felonies
        countyBondCO_Data :: F.MonadSafe m => P.Producer (FM.MaybeRow CountyBondCO) m ()
        countyBondCO_Data = F.readTableMaybeOpt parserOptions countyBondCO_FP
        countyDistrictCO_Data :: F.MonadSafe m => P.Producer CountyDistrictCO m ()
        countyDistrictCO_Data = F.readTableOpt parserOptions countyDistrictCrosswalkCO_FP
        -- This data has 3 rows per county and year, one for each type of crime (against persons, against property, against society)
        crimeStatsCO_Data :: F.MonadSafe m => P.Producer (FM.MaybeRow CrimeStatsCO) m ()
        crimeStatsCO_Data = F.readTableMaybeOpt parserOptions crimeStatsCO_FP
    -- load streams into memory for joins, subsetting as we go
    K.logLE K.Info "loading producers into memory for joining"
    fipsByCountyFrame <- liftIO $ F.inCoreAoS $ fipsByCountyData P.>-> P.map (F.rcast @[Fips,County]) -- get rid of state col
    povertyFrame <- liftIO $ F.inCoreAoS $ povertyData P.>-> P.map (F.rcast @[Fips, Year, MedianHI,MedianHIMOE,PovertyR])
    countyBondFrameM <- liftIO $ fmap F.boxedFrame $ F.runSafeEffect $ P.toListM countyBondCO_Data
    veraFrameM <- liftIO $ fmap F.boxedFrame $ F.runSafeEffect $ P.toListM $ veraData P.>-> P.map (F.rcast @[Fips,Year,TotalPop,Urbanicity,IndexCrime])
    countyDistrictFrame <- liftIO $F.inCoreAoS countyDistrictCO_Data
    K.logLE K.Diagnostic $ (T.pack $ show $ FL.fold FL.length fipsByCountyFrame) <> " rows in fipsByCountyFrame."
    K.logLE K.Diagnostic $ (T.pack $ show $ FL.fold FL.length povertyFrame) <> " rows in povertyFrame."
    K.logLE K.Diagnostic $ (T.pack $ show $ FL.fold FL.length countyBondFrameM) <> " rows in countyBondFrameM."
    K.logLE K.Diagnostic $ (T.pack $ show $ FL.fold FL.length veraFrameM) <> " rows in veraFrameM."
    K.logLE K.Diagnostic $ (T.pack $ show $ FL.fold FL.length countyDistrictFrame) <> " rows in countyDistrictFrame."
    -- do joins
    K.logLE K.Info $ "Doing initial the joins..."
    let countyBondPlusFIPS = FM.leftJoinMaybe (Proxy @'[County]) countyBondFrameM (justsFromRec <$> fipsByCountyFrame)
        countyBondPlusFIPSAndDistrict = FM.leftJoinMaybe (Proxy @'[County]) (F.boxedFrame countyBondPlusFIPS) (justsFromRec <$> countyDistrictFrame)
        countyBondPlusFIPSAndSAIPE = FM.leftJoinMaybe (Proxy @[Fips, Year]) (F.boxedFrame countyBondPlusFIPSAndDistrict) (justsFromRec <$> povertyFrame)
        countyBondPlusFIPSAndSAIPEAndVera = FM.leftJoinMaybe (Proxy @[Fips, Year]) (F.boxedFrame countyBondPlusFIPSAndSAIPE) veraFrameM
    K.newPandoc "moneyBondRateAndPovertyRate" $ kmMoneyBondPctAnalysis countyBondPlusFIPSAndSAIPEAndVera
    K.newPandoc "moneyBondRateAndCrimeRate" $ bondVsCrimeAnalysis countyBondCO_Data crimeStatsCO_Data
    endReal <- liftIO $ C.getTime C.Monotonic
    let realTime = C.diffTimeSpec endReal startReal
        printTime (C.TimeSpec s ns) = (T.pack $ show $ realToFrac s + realToFrac ns/(10^9))
    K.logLE K.Info $ "Time (real): " <> printTime realTime <> "s"
    return ()
  case eitherDocs of
    Right namedDocs -> writeAllHtml namedDocs
    Left err -> putStrLn $ "pandoc error: " ++ show err

-- extra columns we will need
-- CrimeRate is defined in DataSources since we use it in more places
type MoneyBondRate = "money_bond_rate" F.:-> Double
type PostedBondRate = "posted_bond_rate" F.:-> Double
type PostedBondPerCapita = "posted_bond_per_capita" F.:-> Double
type ReoffenseRate = "reoffense_rate" F.:-> Double
type PostedBondFreq = "posted_bond_freq" F.:-> Int
type CrimeRateError = "crime_rate_error" F.:-> Double
type CrimeRateFitted = "crime_rate_fitted" F.:-> Double
type CrimeRateFittedErr = "crime_rate_fitted_err" F.:-> Double

-- functions to populate some of those columns from other columns
moneyBondRate r = let t = r ^. totalBondFreq in FT.recordSingleton @MoneyBondRate $ bool ((r ^. moneyBondFreq) `rDiv` t) 0 (t == 0)
postedBondRate r = let t = r ^. totalBondFreq in FT.recordSingleton @PostedBondRate $ bool ((r ^. moneyPosted + r ^. prPosted) `rDiv` t) 0 (t == 0)
postedBondPerCapita r = FT.recordSingleton @PostedBondPerCapita $ (r ^. moneyPosted + r ^. prPosted) `rDiv` (r ^. estPopulation)
cRate r = FT.recordSingleton @CrimeRate $ (r ^. crimes) `rDiv` (r ^. estPopulation)
reoffenseRate r = FT.recordSingleton @ReoffenseRate $ (r ^. moneyNewYes + r ^. prNewYes) `rDiv` (r ^. moneyPosted + r ^. prPosted)
postedBondFreq r = FT.recordSingleton @PostedBondFreq $ (r ^. moneyPosted + r ^. prPosted)

bondVsCrimeAnalysis :: forall effs. ( MonadIO (K.Semantic effs)
                                    , K.PandocEffects effs
                                    , K.Member K.ToPandoc effs
                                    , K.Member KR.Random effs)
                    => P.Producer (FM.MaybeRow CountyBondCO) (F.SafeT IO) ()
                    -> P.Producer (FM.MaybeRow CrimeStatsCO) (F.SafeT IO) ()
                    -> K.Semantic effs ()
bondVsCrimeAnalysis bondDataMaybeProducer crimeDataMaybeProducer = K.wrapPrefix "BondRateVsCrimeRate" $ do
  K.logLE K.Info "Doing bond rate vs crime rate analysis..."
  countyBondFrameM <- liftIO $ fmap F.boxedFrame $ F.runSafeT $ P.toListM bondDataMaybeProducer
  let blanksToZeroes :: F.Rec (Maybe :. F.ElField) '[Crimes,Offenses] -> F.Rec (Maybe :. F.ElField) '[Crimes,Offenses]
      blanksToZeroes = FM.fromMaybeMono 0
  -- turn Nothing into 0, sequence the Maybes out, use concat to "fold" over those maybes, removing any nothings
  crimeStatsList <- liftIO $ F.runSafeT $ P.toListM $ crimeDataMaybeProducer P.>-> P.map (F.recMaybe . (F.rsubset %~ blanksToZeroes)) P.>-> P.concat
  let sumCrimesFold = FF.sequenceEndoFolds $ FF.FoldEndo FL.sum V.:& FF.FoldEndo FL.sum V.:& FF.FoldEndo (fmap (fromMaybe 0) FL.last) V.:& V.RNil
      foldAllCrimes = MR.concatFold $ MR.hashableMapReduceFold MR.noUnpack (MR.assignKeysAndData @[County, Year] @[Crimes, Offenses, EstPopulation]) (MR.foldAndAddKey sumCrimesFold)
      mergedCrimeStatsFrame = FL.fold foldAllCrimes crimeStatsList
      unmergedCrimeStatsFrame = F.boxedFrame crimeStatsList
      foldAllBonds = MR.concatFold $ MR.hashableMapReduceFold
        (MR.unpackGoodRows @[County,Year,MoneyBondFreq,PrBondFreq,TotalBondFreq,MoneyPosted,PrPosted,MoneyNewYes,PrNewYes])
        (MR.splitOnKeys @[County,Year])
        (MR.foldAndAddKey $ FF.foldAllMonoid @MO.Sum)
      mergedBondDataFrame = FL.fold foldAllBonds countyBondFrameM
      countyBondAndCrimeMerged = F.leftJoin @[County,Year] mergedBondDataFrame mergedCrimeStatsFrame
      countyBondAndCrimeUnmerged = F.leftJoin @[County,Year] mergedBondDataFrame unmergedCrimeStatsFrame
  K.logLE K.Info "Joined crime data and bond data"
  K.logLE K.Diagnostic $ (T.pack $ show $ FL.fold FL.length crimeStatsList) <> " rows in crimeStatsList (unmerged)."
  K.logLE K.Diagnostic $ (T.pack $ show $ FL.fold FL.length mergedCrimeStatsFrame) <> " rows in crimeStatsFrame(merged)."
  let initialCentroidsF n = MR.functionToFoldM (KM.kMeansPPCentroids @DblX @DblY @EstPopulation KM.euclidSq n)
      kmReduce f k rows = sequence $ M.singleton k $ f 10 10 initialCentroidsF (KM.weighted2DRecord @DblX @DblY @EstPopulation) KM.euclidSq rows
      sunCrimeRateF = FL.premap (F.rgetField @CrimeRate) $ MR.scaleAndUnscale (MR.RescaleNormalize 1) (MR.RescaleNone) id
      sunMoneyBondRateF = FL.premap (F.rgetField @MoneyBondRate) $ MR.scaleAndUnscale (MR.RescaleNormalize 1) (MR.RescaleNone) id
      sunPostedBondRateF = FL.premap (F.rgetField @PostedBondRate) $ MR.scaleAndUnscale (MR.RescaleNormalize 1) MR.RescaleNone id
      sunReoffenseRateF = FL.premap (F.rgetField @ReoffenseRate) $ MR.scaleAndUnscale (MR.RescaleNormalize 1) MR.RescaleNone id
      mbrAndCr r = moneyBondRate r F.<+> cRate r
      pbrAndCr r = postedBondRate r F.<+> cRate r
      rorAndMbr r = reoffenseRate r F.<+> moneyBondRate r
  kmMergedCrimeRateVsMoneyBondRateByYear  <- do
    K.logLE K.Info "Doing weighted-KMeans on crime rate vs. money-bond rate (merged crime types)."
    let unpack = fmap (FT.mutate mbrAndCr) (MR.unpackGoodRows @[Year,County,MoneyBondFreq,TotalBondFreq,Crimes,Offenses,EstPopulation])
        reduce :: _
        reduce = MR.ReduceM $ kmReduce (KM.kMeansOneWithClusters @MoneyBondRate @CrimeRate @EstPopulation sunMoneyBondRateF sunCrimeRateF)
        kmFoldM = MR.concatFoldM $ MR.hashableMapReduceFoldM (MR.generalizeUnpack unpack) (MR.generalizeAssign $ MR.assignKeys @'[Year]) reduce
    flip FL.foldM countyBondAndCrimeMerged $
      fmap (KM.clusteredRows @MoneyBondRate @CrimeRate @EstPopulation (F.rgetField @County)) $ kmFoldM
  kmCrimeRateVsMoneyBondRateByYearAndType <- do
    K.logLE K.Info "Doing weighted-KMeans on crime rate vs. money-bond rate (separate crime types)."
    let unpack = fmap (FT.mutate mbrAndCr) (MR.unpackGoodRows @[Year,County,CrimeAgainst,MoneyBondFreq,TotalBondFreq,Crimes,Offenses,EstPopulation])
        reduce :: _
        reduce =  MR.ReduceM $ kmReduce (KM.kMeansOneWithClusters @MoneyBondRate @CrimeRate @EstPopulation sunMoneyBondRateF sunCrimeRateF)
        kmFoldM = MR.concatFoldM $ MR.hashableMapReduceFoldM (MR.generalizeUnpack unpack) (MR.generalizeAssign $ MR.assignKeys @'[Year, CrimeAgainst]) reduce
    flip FL.foldM countyBondAndCrimeUnmerged $
      fmap (KM.clusteredRows @MoneyBondRate @CrimeRate @EstPopulation (F.rgetField @County)) $ kmFoldM
  kmMergedCrimeRateVsPostedBondRateByYear <- do
    K.logLE K.Info "Doing weighted-KMeans on crime rate vs. posted-bond rate (merged)."
    let unpack = fmap (FT.mutate pbrAndCr) (MR.unpackGoodRows @[Year,County,TotalBondFreq,MoneyPosted,PrPosted,Crimes,Offenses,EstPopulation])
        reduce :: _
        reduce =  MR.ReduceM $ kmReduce (KM.kMeansOneWithClusters @PostedBondRate @CrimeRate @EstPopulation sunPostedBondRateF sunCrimeRateF)
        kmFoldM = MR.concatFoldM $ MR.hashableMapReduceFoldM (MR.generalizeUnpack unpack) (MR.generalizeAssign $ MR.assignKeys @'[Year]) reduce
    flip FL.foldM countyBondAndCrimeMerged $
      fmap (KM.clusteredRows @PostedBondRate @CrimeRate @EstPopulation (F.rgetField @County)) $ kmFoldM
  kmCrimeRateVsPostedBondRateByYearAndType <- do
    K.logLE K.Info "Doing weighted-KMeans on crime rate vs. posted-bond rate (unmerged)."
    let unpack = fmap (FT.mutate pbrAndCr) (MR.unpackGoodRows @[Year,County,CrimeAgainst,TotalBondFreq,MoneyPosted,PrPosted,Crimes,Offenses,EstPopulation])
        reduce :: _
        reduce = MR.ReduceM $ kmReduce (KM.kMeansOneWithClusters @PostedBondRate @CrimeRate @EstPopulation sunPostedBondRateF sunCrimeRateF)
        kmFoldM = MR.concatFoldM $ MR.hashableMapReduceFoldM (MR.generalizeUnpack unpack) (MR.generalizeAssign $ MR.assignKeys @'[Year,CrimeAgainst]) reduce
    flip FL.foldM countyBondAndCrimeUnmerged $
      fmap (KM.clusteredRows @PostedBondRate @CrimeRate @EstPopulation (F.rgetField @County)) $ kmFoldM
  kmReoffenseRateVsMergedMoneyBondRateByYear <- do
    K.logLE K.Info "Doing weighted-KMeans on re-offense rate vs. money-bond rate (merged)."
    let unpack = fmap (FT.mutate rorAndMbr) (MR.unpackGoodRows @[Year, County, MoneyNewYes, PrNewYes, MoneyPosted, PrPosted, TotalBondFreq, MoneyBondFreq, EstPopulation])
        reduce :: _
        reduce = MR.ReduceM $ kmReduce (KM.kMeansOneWithClusters @MoneyBondRate @ReoffenseRate @EstPopulation sunMoneyBondRateF sunReoffenseRateF)
        kmFoldM = MR.concatFoldM $ MR.hashableMapReduceFoldM (MR.generalizeUnpack unpack) (MR.generalizeAssign $ MR.assignKeys @'[Year]) reduce
    flip FL.foldM countyBondAndCrimeMerged $
      fmap (KM.clusteredRows @MoneyBondRate @ReoffenseRate @EstPopulation (F.rgetField @County)) $ kmFoldM
  -- regressions
  let rMut r = mbrAndCr r F.<+> postedBondFreq r F.<+> postedBondRate r F.<+> postedBondPerCapita r
  (rData, regressionRes, regressionResMany) <- do
    let regUnpack = fmap (FT.mutate rMut) $ (MR.unpackGoodRows @[Year,MoneyBondFreq,PrBondFreq,PrPosted,MoneyPosted,TotalBondFreq,Crimes,EstPopulation])
--        regReduce :: Foldable f => (forall g. Foldable g => g (F.Record rs) -> PS.Semantic effs b) -> k -> f (F.Record rs) -> PS.Semantic effs (M.Map k b)

        regMR r = MR.concatFoldM $ MR.hashableMapReduceFoldM (MR.generalizeUnpack regUnpack) (MR.generalizeAssign $ MR.assignKeys @'[Year]) r --(MR.ReduceM $ regReduce r)
        dataMR = MR.unpackOnlyFoldM (MR.generalizeUnpack regUnpack)
        guess = [0,0] -- guess has one extra dimension for constant
        regressOneBM = MR.functionToFoldM $ return . FR.leastSquaresByMinimization @Crimes @'[EstPopulation, MoneyBondFreq] False guess
        regressOneOLS = MR.functionToFoldM $ FR.ordinaryLeastSquares @effs @Crimes @False @'[EstPopulation]
        regressOneWLS = MR.functionToFoldM $ FR.popWeightedLeastSquares @effs @CrimeRate @True @'[] @EstPopulation
        regressOneWLS2 = MR.functionToFoldM $ FR.varWeightedLeastSquares @effs @Crimes @False @'[EstPopulation, PostedBondFreq] @EstPopulation
        regressOneWTLS = MR.functionToFoldM $ FR.varWeightedTLS @effs @Crimes @False @'[EstPopulation, PostedBondFreq] @EstPopulation
        regReduceF r k = fmap (M.singleton k) r
        allMR  = (,,,,,) <$> dataMR
                 <*> (regMR $ MR.ReduceFoldM $ regReduceF regressOneBM)
                 <*> (regMR $ MR.ReduceFoldM $ regReduceF regressOneOLS)
                 <*> (regMR $ MR.ReduceFoldM $ regReduceF regressOneWLS)
                 <*> (regMR $ MR.ReduceFoldM $ regReduceF regressOneWLS2)
                 <*> (regMR $ MR.ReduceFoldM $ regReduceF regressOneWTLS)
    K.logLE K.Info "Regressing Crime Rate on Money Bond Rate"
    (rData, r1, r2,r3,r4,r5) <- FL.foldM allMR countyBondAndCrimeMerged
--    K.logLE K.Info $ "regression (by minimization) results: " <> (T.pack $ show $ fmap (flip FR.prettyPrintRegressionResult 0.95) r1)
    K.logLE K.Info $ "regression (by OLS) results: " <> FR.prettyPrintRegressionResults FR.keyRecordText (M.toList r2) ST.cl95 FR.prettyPrintRegressionResult "\n"
    K.logLE K.Info $ "regression (rates by pop weighted LS) results: " <> FR.prettyPrintRegressionResults FR.keyRecordText (M.toList r3) ST.cl95 FR.prettyPrintRegressionResult "\n"
    K.logLE K.Info $ "regression (counts by inv sqrt pop weighted LS) results: " <> FR.prettyPrintRegressionResults FR.keyRecordText (M.toList r4) ST.cl95 FR.prettyPrintRegressionResult "\n"
    K.logLE K.Info $ "regression (counts by TLS) results: " <> FR.prettyPrintRegressionResults FR.keyRecordText  (M.toList r5) ST.cl95 FR.prettyPrintRegressionResult "\n"
    let rData2016 = F.filterFrame ((== 2016) . F.rgetField @Year) $ F.toFrame rData
        regressionR = fromJust $ M.lookup (FT.recordSingleton @Year 2016) r4
    return $ (rData2016, regressionR, r4)
--    K.log K.Info $ "data=\n" <> Table.textTable rData

  K.logLE K.Info "Creating Doc"
  K.addLucid $ do
    HL.h1_ "Colorado Money Bond Rate vs Crime rate"
    KL.placeTextSection $ do
      HL.h2_ "Colorado Bond Rates and Crime Rates (preliminary)"
      HL.p_ [HL.class_ "subtitle"] "Adam Conner-Sax"
      HL.p_ "Each county in Colorado issues money bonds and personal recognizance bonds.  For each county I look at the % of money bonds out of all bonds issued and the crime rate.  We have 3 years of data and there are 64 counties in Colorado (each with vastly different populations).  So I've used a population-weighted k-means clustering technique (see notes below) to reduce the number of points to at most 10 per year. Each circle in the plot below represents one cluster of counties with similar money bond and poverty rates.  The size of the circle represents the total population in the cluster."
    KL.placeVisualization "crimeRateVsMoneyBondRateMerged" $ cRVsMBRVL True kmMergedCrimeRateVsMoneyBondRateByYear
    HL.p_ "Broken down by Colorado's crime categories:"
    KL.placeVisualization "crimeRateVsMoneyBondRateUnMerged" $ cRVsMBRVL False kmCrimeRateVsMoneyBondRateByYearAndType
    KL.placeTextSection $ do
      HL.p_ "Notes:"
      HL.ul_ $ do
        HL.li_ $ do
          HL.span_ "Money Bond Rate is (number of money bonds/total number of bonds). That data comes from "
          HL.a_ [HL.href_ "https://github.com/Data4Democracy/incarceration-trends/blob/master/Colorado_ACLU/4-money-bail-analysis/complete-county-bond.csv"] "complete-county-bond.csv"
          HL.span_ ". NB: Bond rates are not broken down by crime type because they are not broken down in the data we have."
        HL.li_ $ do
          HL.span_ "Crime Rate is crimes/estimated_population. Those numbers come from the Colorado crime statistics "
          HL.a_ [HL.href_ "https://coloradocrimestats.state.co.us/"] "web-site"
          HL.span_ ".  We also have some of that data in the "
          HL.a_ [HL.href_ "https://github.com/Data4Democracy/incarceration-trends/blob/master/Colorado_ACLU/4-money-bail-analysis/crime-rate-bycounty.csv"] "repo"
          HL.span_ ", but only for 2016, so I re-downloaded it for this. Also, there are some crimes in that data that don't roll up to a particular county but instead are attributed to the Colorado Bureau of Investigation or the Colorado State Patrol.  I've ignored those."
    KL.placeTextSection $ do
      HL.p_ "Below I look at the percentage of all bonds which are \"posted\" (posting bond means paying the money bond or agreeing to the terms of the personal recognizance bond ?) rather than the % of money bonds. I use the same clustering technique."
    KL.placeVisualization "crimeRateVsPostedBondRate" $ cRVsPBRVL True kmMergedCrimeRateVsPostedBondRateByYear
    HL.p_ "Broken down by Colorado's crime categories:"
    KL.placeVisualization "crimeRateVsPostedBondRateUnMerged" $ cRVsPBRVL False kmCrimeRateVsPostedBondRateByYearAndType
    KL.placeTextSection $ do
      HL.p_ "Notes:"
      HL.ul_ $ do
        HL.li_ $ do
          HL.span_ "Posted Bond Rate is [(number of posted money bonds + number of posted PR bonds)/total number of bonds] where that data comes from "
          HL.a_ [HL.href_ "https://github.com/Data4Democracy/incarceration-trends/blob/master/Colorado_ACLU/4-money-bail-analysis/complete-county-bond.csv"] "complete-county-bond.csv."
        HL.li_ $ do
          HL.span_ "Crime Rate, as above, is crimes/estimated_population where those numbers come from the Colorado crime statistics "
          HL.a_ [HL.href_ "https://coloradocrimestats.state.co.us/"] "web-site."
    HL.p_ "Below I look at the rate of re-offending among people out on bond, here considered in each county along with money-bond rate and clustered as above."
    KL.placeVisualization "reoffenseRateVsMoneyBondRate" $ rRVsMBRVL True kmReoffenseRateVsMergedMoneyBondRateByYear
    KL.placeTextSection $ do
      HL.p_ "Notes:"
      HL.ul_ $ do
        HL.li_ $ do
          HL.span_ "Reoffense Rate is (number of new offenses for all bonds/total number of posted bonds) where that data comes from "
          HL.a_ [HL.href_ "https://github.com/Data4Democracy/incarceration-trends/blob/master/Colorado_ACLU/4-money-bail-analysis/complete-county-bond.csv"] "complete-county-bond.csv."

    KL.placeTextSection $ do
      HL.h2_ "Regressions"
      HL.p_ "We can use linear regression to investigate the relationship between Crime Rate and the use of money and personal recognizance bonds. We begin by finding a best fit (using population-weighted least squares) to the model "
      KL.latex_ "$c_i = cr (p_{i}) + A (pb_{i}) + e_{i}$"
      HL.p_ "where, for each county (denoted by the subscipt i), c is the number of crimes, p is the population and pb is the number of posted bonds and e is an error term we seek to minimize. We look at the result for each of the years in the data:"
    KL.placeVisualization "regresssionCoeffs" $ FV.regressionCoefficientPlotMany (T.pack . show . F.rgetField @Year) "Regression Results (by year)" [" cr","A"] (M.toList (fmap FR.regressionResult regressionResMany)) ST.cl95
    HL.p_ "We look at these regressions directly by overlaying this model on the data itself:"
--    H.placeVisualization "regresssionScatter"  $ FV.scatterWithFit @PostedBondPerCapita @CrimeRate errF (FV.FitToPlot "WLS regression" fitF) "test scatter with fit" rData
--    H.placeVisualization "regresssionScatter"  $ FV.scatterWithFit @PostedBondPerCapita @CrimeRate errF (FV.FitToPlot "WLS regression" fitF) "test scatter with fit" rData
    KL.placeVisualization "regresssionScatter"  $ FV.frameScatterWithFit "test scatter with fit" (Just "WLS") regressionRes ST.cl95 rData
    kMeansNotes
--  liftIO $ T.writeFile "analysis/moneyBondRateAndCrimeRate.html" $ TL.toStrict $ htmlAsText

cRVsMBRVL mergedOffenseAgainst dataRecords =
  let dat = FV.recordsToVLData (transformF @MoneyBondRate (*100) . transformF @CrimeRate (*100)) dataRecords
      ptEnc = GV.color [FV.mName @Year, GV.MmType GV.Nominal]
              . GV.size [FV.mName @EstPopulation, GV.MmType GV.Quantitative, GV.MLegend [GV.LTitle "population"]]
  in case mergedOffenseAgainst of
    True ->
      let vlF = FV.clustersWithClickIntoVL @MoneyBondRate @CrimeRate @KM.IsCentroid @KM.ClusterId @KM.MarkLabel @'[Year]
      in vlF "Money Bond Rate (%, All Crimes)" "Crime Rate (%, All Crimes)" "Crime Rate vs Money Bond Rate" ptEnc id id id dat
    False ->
      let ptEncFaceted = ptEnc
                         . GV.row [FV.fName @CrimeAgainst, GV.FmType GV.Nominal,GV.FHeader [GV.HTitle "Crime Against"]]
          vlF = FV.clustersWithClickIntoVL @MoneyBondRate @CrimeRate @KM.IsCentroid @KM.ClusterId @KM.MarkLabel @[Year,CrimeAgainst]
      in vlF "Money Bond Rate (%, All Crimes)" "Crime Rate (%)" "Crime Rate vs Money Bond Rate" ptEncFaceted id id id dat

cRVsPBRVL mergedOffenseAgainst dataRecords =
  let dat = FV.recordsToVLData (transformF @PostedBondRate (*100) . transformF @CrimeRate (*100)) dataRecords
      ptEnc = GV.color [FV.mName @Year, GV.MmType GV.Nominal]
              . GV.size [FV.mName @EstPopulation, GV.MmType GV.Quantitative, GV.MLegend [GV.LTitle "population"]]
  in case mergedOffenseAgainst of
    True ->
      let vlF = FV.clustersWithClickIntoVL @PostedBondRate @CrimeRate @KM.IsCentroid @KM.ClusterId @KM.MarkLabel @'[Year]
      in vlF "Posted Bond Rate (%, All Crimes)" "Crime Rate (%, All Crimes)" "Crime Rate vs Posted Bond Rate" ptEnc id id id dat
    False ->
      let ptEncFaceted = ptEnc
                         . GV.row [FV.fName @CrimeAgainst, GV.FmType GV.Nominal,GV.FHeader [GV.HTitle "Crime Against"]]
          vlF = FV.clustersWithClickIntoVL @PostedBondRate @CrimeRate @KM.IsCentroid @KM.ClusterId @KM.MarkLabel @[Year,CrimeAgainst]
      in vlF "Posted Bond Rate (%, All Crimes)" "Crime Rate (%)" "Crime Rate vs Posted Bond Rate" ptEncFaceted id id id dat

rRVsMBRVL mergedOffenseAgainst dataRecords =
  let dat = FV.recordsToVLData (transformF @MoneyBondRate (*100) . transformF @ReoffenseRate (*100)) dataRecords
      ptEnc = GV.color [FV.mName @Year, GV.MmType GV.Nominal]
              . GV.size [FV.mName @EstPopulation, GV.MmType GV.Quantitative, GV.MLegend [GV.LTitle "population"]]
  in case mergedOffenseAgainst of
    True ->
      let vlF = FV.clustersWithClickIntoVL @MoneyBondRate @ReoffenseRate @KM.IsCentroid @KM.ClusterId @KM.MarkLabel @'[Year]
      in vlF "Money Bond Rate (%, All Crimes)" "Reoffense Rate (%, All Crimes)" "Reoffense Rate vs Money Bond Rate" ptEnc id id id dat
    False ->
      let ptEncFaceted = ptEnc
                         . GV.row [FV.fName @CrimeAgainst, GV.FmType GV.Nominal,GV.FHeader [GV.HTitle "Crime Against"]]
          vlF = FV.clustersWithClickIntoVL @MoneyBondRate @CrimeRate @KM.IsCentroid @KM.ClusterId @KM.MarkLabel @[Year,CrimeAgainst]
      in vlF "Money Bond Rate (%, All Crimes)" "Reoffense Rate (%)" "Reoffense Rate vs Money Bond Rate" ptEncFaceted id id id dat

-- NB: The data has two rows for each county and year, one for misdemeanors and one for felonies.
--type MoneyPct = "moneyPct" F.:-> Double --F.declareColumn "moneyPct" ''Double
kmMoneyBondPctAnalysis :: _
kmMoneyBondPctAnalysis joinedData = K.wrapPrefix "MoneyBondVsPoverty" $ do
  K.logLE K.Info "Doing money bond % vs poverty rate analysis..."
  let sunPovertyRF = FL.premap (F.rgetField @PovertyR) $ MR.scaleAndUnscale (MR.RescaleNormalize 1) MR.RescaleNone id
      sunMoneyBondRateF = FL.premap (F.rgetField @MoneyBondRate) $ MR.scaleAndUnscale (MR.RescaleNormalize 1) MR.RescaleNone id
      initialCentroids = KM.kMeansPPCentroids @DblX @DblY @TotalPop KM.euclidSq
      unpack = fmap (FT.mutate moneyBondRate) $ MR.unpackGoodRows @[Year,County,OffType,Urbanicity,PovertyR, MoneyBondFreq,TotalBondFreq,TotalPop]
      reduce :: _
      reduce = MR.ReduceM $ \_ -> KM.kMeansOne @PovertyR @MoneyBondRate @TotalPop sunPovertyRF sunMoneyBondRateF 5 initialCentroids (KM.weighted2DRecord @DblX @DblY @TotalPop) KM.euclidSq
      toRec (x, y, z) = (x F.&: y F.&: z F.&: V.RNil) :: F.Record [PovertyR, MoneyBondRate, TotalPop]
      kmByYearF = MR.concatFoldM $ MR.hashableMapReduceFoldM
        (MR.generalizeUnpack unpack)
        (MR.generalizeAssign $ MR.assignKeysAndData @[Year, OffType] @[PovertyR, MoneyBondRate, TotalPop])
        (MR.makeRecsWithKeyM toRec reduce)
      kmByYearUrbF = MR.concatFoldM $ MR.hashableMapReduceFoldM
        (MR.generalizeUnpack unpack)
        (MR.generalizeAssign $ MR.assignKeysAndData @'[Year, OffType, Urbanicity] @[PovertyR, MoneyBondRate, TotalPop])
        (MR.makeRecsWithKeyM toRec reduce)
  (kmByYear, kmByYearUrb) <- FL.foldM ((,) <$> kmByYearF <*> kmByYearUrbF) joinedData
{-  (kmByYear, kmByYearUrb) <- FL.foldM ((,)
                                        <$> KM.kMeans @[Year, OffType] @PovertyR @MoneyBondRate @TotalPop sunPovertyRF sunMoneyBondRateF 5 (KM.kMeansPPCentroids KM.euclidSq) KM.euclidSq
                                        <*> KM.kMeans @[Year, OffType, Urbanicity] @PovertyR @MoneyBondRate @TotalPop sunPovertyRF sunMoneyBondRateF 5 (KM.kMeansPPCentroids KM.euclidSq) KM.euclidSq) kmData -}
  K.logLE K.Info "Creating Doc"
  K.addLucid $ do
    HL.h1_ "Colorado Money Bonds and Poverty"
    KL.placeTextSection $ do
      HL.h2_ "Colorado Money Bond Rate and Poverty (preliminary)"
      HL.p_ [HL.class_ "subtitle"] "Adam Conner-Sax"
      HL.p_ "Colorado issues two types of bonds when releasing people from jail before trial. Sometimes people are released on a \"money bond\" and sometimes on a personal recognizance bond. We have county-level data of all bonds (?) issued in 2014, 2015 and 2016.  Plotting it all is very noisy (since there are 64 counties in CO) so we use a population-weighted k-means clustering technique to combine similar counties into clusters. In the plots below, each circle represents one cluster of counties and the circle's size represents the total population of all the counties included in the cluster.  We consider felonies and misdemeanors separately."
    KL.placeVisualization "mBondsVspRateByYrFelonies" $ moneyBondPctVsPovertyRateVL False kmByYear
    KL.placeTextSection $ HL.h3_ "Broken down by \"urbanicity\":"
    KL.placeVisualization "mBondsVspRateByYrUrbFelonies" $ moneyBondPctVsPovertyRateVL True kmByYearUrb
    KL.placeTextSection $ do
      HL.p_ "Notes:"
      HL.ul_ $ do
        HL.li_ "Denver, the only urban county in CO, didn't report misdemeanors in this data, so the last plot is blank."
        HL.li_ $ do
          HL.span_ "Money Bond Rate is (number of money bonds/total number of bonds). That data comes from "
          HL.a_ [HL.href_ "https://github.com/Data4Democracy/incarceration-trends/blob/master/Colorado_ACLU/4-money-bail-analysis/complete-county-bond.csv"] "complete-county-bond.csv."
        HL.li_ $ do
          HL.span_ "Poverty Rate comes from the census and we have that data in "
          HL.a_ [HL.href_ "https://github.com/Data4Democracy/incarceration-trends/blob/master/Colorado_ACLU/4-money-bail-analysis/complete-county-bond-SAIPE.csv"] "complete-county-bond-SAIPE.csv"
          HL.span_ ". That data originates from the "
          HL.a_ [HL.href_ "https://www.census.gov/programs-surveys/saipe.html"] "SAIPE website"
          HL.span_ " at the census bureau."
        HL.li_ $ do
          HL.span_ "Urbanicity classification comes from the "
          HL.a_ [HL.href_ "https://www.vera.org/projects/incarceration-trends"] "VERA"
          HL.span_ " data. In the repo "
          HL.a_ [HL.href_ "https://github.com/Data4Democracy/incarceration-trends/blob/master/Colorado_ACLU/4-money-bail-analysis/county_bond_saipe_vera_data.csv"] "here."
    kMeansNotes
--  liftIO $ T.writeFile "analysis/moneyBondRateAndPovertyRate.html" $ TL.toStrict $ htmlAsText

moneyBondPctVsPovertyRateVL facetByUrb dataRecords =
  let dat = FV.recordsToVLData (transformF @MoneyBondRate (*100)) dataRecords
      enc = GV.encoding
        . GV.position GV.X [FV.pName @PovertyR, GV.PmType GV.Quantitative, GV.PAxis [GV.AxTitle "Poverty Rate (%)"]]
        . GV.position GV.Y [FV.pName @MoneyBondRate, GV.PmType GV.Quantitative, GV.PAxis [GV.AxTitle "% Money Bonds"]]
        . GV.color [FV.mName @Year, GV.MmType GV.Nominal]
        . GV.size [FV.mName @TotalPop, GV.MmType GV.Quantitative, GV.MLegend [GV.LTitle "population"]]
        . GV.row [FV.fName @OffType, GV.FmType GV.Nominal, GV.FHeader [GV.HTitle "Type of Offense"]]
        . if facetByUrb then GV.column [FV.fName @Urbanicity, GV.FmType GV.Nominal] else id
      sizing = if facetByUrb
               then [GV.autosize [GV.AFit, GV.AResize]]
               else [GV.autosize [GV.AFit], GV.height 300, GV.width 700]
      vl = GV.toVegaLite $
        [ GV.description "Vega-Lite Attempt"
        , GV.title ("% of money bonds (out of money and personal recognizance bonds) vs poverty rate in CO")
        , GV.background "white"
        , GV.mark GV.Point []
        , enc []
        , dat] <> sizing
  in vl

kMeansNotes =  KL.placeTextSection $ do
  HL.h3_ "Some notes on weighted k-means"
  HL.ul_ $ do
    HL.li_ $ do
      HL.a_ [HL.href_ "https://en.wikipedia.org/wiki/K-means_clustering"] "k-means"
      HL.span_ " works by choosing random starting locations for cluster centers, assigning each data-point to the nearest cluster center, moving the center of each cluster to the weighted average of the data-points assigned to the same cluster and then repeating until no data-points move to a new cluster."
    HL.li_ "How do you define distance between points?  Each axis has a different sort of data and it's not clear how to combine differences in each into a meanignful overall distance.  Here, before we hand the data off to the k-means algorithm, we shift and rescale the data so that the set of points has mean 0 and std-deviation 1 in each variable.  Then we use ordinary Euclidean distance, that is the sum of the squares of the differences in each coordinate.  Before plotting, we reverse that scaling so that we can visualize the data in the original. This attempts to give the two variables equal weight in determining what \"nearby\" means for these data points."
    HL.li_ "This clustering happens for each combination plotted."
    HL.li_ $ do
      HL.span_ "k-means is sensitive to the choice of starting points.  Here we use the \"k-means++\" method.  This chooses one of the data points as a starting point.  Then chooses among the remaining points in a way designed to make it likely that the initial centers are widely dispersed.  See "
      HL.a_ [HL.href_ "https://en.wikipedia.org/wiki/K-means%2B%2B"] "here"
      HL.span_ " for more information.  We repeat the k-means clustering with a few different starting centers chosen this way and choose the best clustering, in the sense of minimizing total weighted distance of all points from their cluster centers."

transformF :: forall x rs. (V.KnownField x, x ∈ rs) => (V.Snd x -> V.Snd x) -> F.Record rs -> F.Record rs
transformF f r = F.rputField @x (f $ F.rgetField @x r) r

--  kmBondRatevsCrimeRateAnalysis countyBondPlusFIPSAndSAIPEAndVera
type IndexCrimeRate = "index_crime_rate" F.:-> Double
type TotalBondRate = "total_bond_rate" F.:-> Double
indexCrimeRate r = FT.recordSingleton @IndexCrimeRate $ (r ^. indexCrime) `rDiv` (r ^. totalPop)
totalBondRate r = FT.recordSingleton @TotalBondRate $ (r ^. totalBondFreq) `rDiv` (r ^. totalPop)
kmBondRatevsCrimeRateAnalysis joinedData = do
  let select = F.rcast @[Year,County,Urbanicity,TotalBondFreq,IndexCrime,TotalPop]
      mutation r = totalBondRate r F.<+> indexCrimeRate r
      mData = fmap (FT.mutate mutation) . catMaybes $ fmap (F.recMaybe . select) joinedData
  F.writeCSV "data/raw_COCrimeRatevsBondRate.csv" mData
{-  let dataCols = Proxy @[TotalBondRate, IndexCrimeRate, TotalPop]
      xScaling = FL.premap (F.rgetField @TotalBondRate &&& F.rgetField @TotalPop) $ MR.weightedScaleAndUnscale (MR.RescaleNormalize 1) MR.RescaleNone id
      yScaling = FL.premap (F.rgetField @IndexCrimeRate &&& F.rgetField @TotalPop) $ MR.weightedScaleAndUnscale (MR.RescaleNormalize 1) MR.RescaleNone id
-}

--  FM.writeCSV_Maybe "data/countyBondPlusFIPS.csv" countyBondPlusFIPS
--  FM.writeCSV_Maybe "data/countyBondPlusSAIPE.csv" countyBondPlusFIPSAndSAIPE
--  FM.writeCSV_Maybe "data/countyBondPlus.csv" countyBondPlusFIPSAndSAIPEAndVera
--  coloradoRowCheck <- F.runSafeEffect $ FL.purely P.fold (goodDataByKey  [F.pr1|Year|]) coloradoTrendsData
--  putStrLn $ "(CO rows, CO rows with all fields) = " ++ show coloradoRowCheck


{-# LANGUAGE DataKinds             #-}
{-# LANGUAGE ScopedTypeVariables   #-}
{-# LANGUAGE GADTs                 #-}
{-# LANGUAGE TypeOperators         #-}
{-# LANGUAGE TypeFamilies          #-}
{-# LANGUAGE FlexibleContexts      #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-|
Module      : Grenade.Layers.Crop
Description : Cropping layer
Copyright   : (c) Huw Campbell, 2016-2017
License     : BSD2
Stability   : experimental
-}
module Grenade.Layers.Crop (
    Crop (..)
  ) where

import           Data.Maybe
import           Data.Proxy
import           Data.Singletons.TypeLits
import           GHC.TypeLits

import           Grenade.Core
import           Grenade.Layers.Internal.Pad

import           Numeric.LinearAlgebra (konst, subMatrix, diagBlock)
import           Numeric.LinearAlgebra.Static (extract, create)

-- | A cropping layer for a neural network.
data Crop :: Nat
          -> Nat
          -> Nat
          -> Nat -> * where
  Crop :: Crop cropLeft cropTop cropRight cropBottom

instance Show (Crop cropLeft cropTop cropRight cropBottom) where
  show Crop = "Crop"

instance UpdateLayer (Crop l t r b) where
  type Gradient (Crop l t r b) = ()
  runUpdate _ x _ = x
  createRandom = return Crop

-- | A two dimentional image can be cropped.
instance ( KnownNat cropLeft
         , KnownNat cropTop
         , KnownNat cropRight
         , KnownNat cropBottom
         , KnownNat inputRows
         , KnownNat inputColumns
         , KnownNat outputRows
         , KnownNat outputColumns
         , (outputRows + cropTop + cropBottom) ~ inputRows
         , (outputColumns + cropLeft + cropRight) ~ inputColumns
         ) => Layer (Crop cropLeft cropTop cropRight cropBottom) ('D2 inputRows inputColumns) ('D2 outputRows outputColumns) where
  type Tape (Crop cropLeft cropTop cropRight cropBottom) ('D2 inputRows inputColumns) ('D2 outputRows outputColumns) = ()
  runForwards Crop (S2D input) =
    let cropl = fromIntegral $ natVal (Proxy :: Proxy cropLeft)
        cropt = fromIntegral $ natVal (Proxy :: Proxy cropTop)
        nrows = fromIntegral $ natVal (Proxy :: Proxy outputRows)
        ncols = fromIntegral $ natVal (Proxy :: Proxy outputColumns)
        m  = extract input
        r  = subMatrix (cropt, cropl) (nrows, ncols) m
    in  ((), S2D . fromJust . create $ r)
  runBackwards _ _ (S2D dEdy) =
    let cropl = fromIntegral $ natVal (Proxy :: Proxy cropLeft)
        cropt = fromIntegral $ natVal (Proxy :: Proxy cropTop)
        cropr = fromIntegral $ natVal (Proxy :: Proxy cropRight)
        cropb = fromIntegral $ natVal (Proxy :: Proxy cropBottom)
        eo    = extract dEdy
        vs    = diagBlock [konst 0 (cropt,cropl), eo, konst 0 (cropb,cropr)]
    in  ((), S2D . fromJust . create $ vs)

-- | A two dimentional image can be cropped.
instance ( KnownNat cropLeft
         , KnownNat cropTop
         , KnownNat cropRight
         , KnownNat cropBottom
         , KnownNat inputRows
         , KnownNat inputColumns
         , KnownNat outputRows
         , KnownNat outputColumns
         , KnownNat channels
         , KnownNat (inputRows * channels)
         , KnownNat (outputRows * channels)
         , (outputRows + cropTop + cropBottom) ~ inputRows
         , (outputColumns + cropLeft + cropRight) ~ inputColumns
         ) => Layer (Crop cropLeft cropTop cropRight cropBottom) ('D3 inputRows inputColumns channels) ('D3 outputRows outputColumns channels) where
  type Tape (Crop cropLeft cropTop cropRight cropBottom) ('D3 inputRows inputColumns channels) ('D3 outputRows outputColumns channels)  = ()
  runForwards Crop (S3D input) =
    let padl  = fromIntegral $ natVal (Proxy :: Proxy cropLeft)
        padt  = fromIntegral $ natVal (Proxy :: Proxy cropTop)
        padr  = fromIntegral $ natVal (Proxy :: Proxy cropRight)
        padb  = fromIntegral $ natVal (Proxy :: Proxy cropBottom)
        inr   = fromIntegral $ natVal (Proxy :: Proxy inputRows)
        inc   = fromIntegral $ natVal (Proxy :: Proxy inputColumns)
        outr  = fromIntegral $ natVal (Proxy :: Proxy outputRows)
        outc  = fromIntegral $ natVal (Proxy :: Proxy outputColumns)
        ch    = fromIntegral $ natVal (Proxy :: Proxy channels)
        m     = extract input
        cropped = crop ch padl padt padr padb outr outc inr inc m
    in  ((), S3D . fromJust . create $ cropped)

  runBackwards Crop () (S3D gradient) =
    let padl  = fromIntegral $ natVal (Proxy :: Proxy cropLeft)
        padt  = fromIntegral $ natVal (Proxy :: Proxy cropTop)
        padr  = fromIntegral $ natVal (Proxy :: Proxy cropRight)
        padb  = fromIntegral $ natVal (Proxy :: Proxy cropBottom)
        inr   = fromIntegral $ natVal (Proxy :: Proxy inputRows)
        inc   = fromIntegral $ natVal (Proxy :: Proxy inputColumns)
        outr  = fromIntegral $ natVal (Proxy :: Proxy outputRows)
        outc  = fromIntegral $ natVal (Proxy :: Proxy outputColumns)
        ch    = fromIntegral $ natVal (Proxy :: Proxy channels)
        m     = extract gradient
        padded = pad ch padl padt padr padb outr outc inr inc m
    in  ((), S3D . fromJust . create $ padded)