{-# LANGUAGE AllowAmbiguousTypes  #-}
{-# LANGUAGE ConstraintKinds      #-}
{-# LANGUAGE DataKinds            #-}
{-# LANGUAGE PolyKinds            #-}
{-# LANGUAGE GADTs                #-}
{-# LANGUAGE FlexibleContexts     #-}
{-# LANGUAGE FlexibleInstances    #-}
{-# LANGUAGE OverloadedStrings    #-}
{-# LANGUAGE TypeFamilies         #-}
{-# LANGUAGE TypeOperators        #-}
{-# LANGUAGE TypeApplications     #-}
{-# LANGUAGE Rank2Types           #-}
{-# LANGUAGE ScopedTypeVariables  #-}
{-# LANGUAGE UndecidableInstances       #-}
{-# OPTIONS_GHC -O0 -freduction-depth=0 #-}
module BlueRipple.Data.CPSVoterPUMS
  (
    cpsVoterPUMSLoader
  , cpsVoterPUMSWithCDLoader
  , CPSVoterPUMS
  , CPSVoterPUMSWeight
  , cpsVoterPUMSRollup
  , cpsVoterPUMSRollupWeightedCounts
  , cpsVoterPUMSElectoralWeights
  , cpsVoterPUMSElectoralWeightsByCD
  , cpsVoterPUMSElectoralWeightsByState
  , cpsVoterPUMSNationalElectoralWeights
  , cpsCountVotersByStateF
--  , cpsKeysToASE
--  , cpsKeysToASR
  , cpsKeysToASER
--  , cpsKeysToASER4
  , cpsKeysToASER5
  , cpsKeysToASER4H
--  , cpsKeysToIdentity
  , cpsPossibleVoter
  , cpsVoted

  ) where

import qualified BlueRipple.Data.CPSVoterPUMS.CPSVoterPUMS_Frame as BR
import qualified BlueRipple.Data.DemographicTypes as BR
import qualified BlueRipple.Data.ElectionTypes as BR
import qualified BlueRipple.Data.DataFrames as BR
import qualified BlueRipple.Data.LoadersCore as BR
import qualified BlueRipple.Data.Loaders as BR
import qualified BlueRipple.Data.Keyed as BR
import qualified BlueRipple.Utilities.KnitUtils as BR
import qualified BlueRipple.Data.CountFolds as BRCF

import qualified Control.Foldl                 as FL
import           Control.Lens                   ((%~))
import qualified Control.Monad.Except          as X
import qualified Control.Monad.State           as ST
import qualified Data.Array                    as A
import qualified Data.Serialize                as S
import qualified Data.Serialize.Text           as S
import qualified Data.List                     as L
import qualified Data.Map                      as M
import           Data.Maybe                     ( fromMaybe, catMaybes)
import qualified Data.Text                     as T
import           Data.Text                      ( Text )
import           Text.Read                      (readMaybe)
import qualified Data.Vinyl                    as V
import           Data.Vinyl.TypeLevel                     (type (++))
import qualified Data.Vinyl.TypeLevel          as V
import qualified Data.Vinyl.Functor            as V
import qualified Frames                        as F
import           Data.Vinyl.Lens               (type (⊆))
import           Frames                         ( (:.)(..) )
import qualified Frames.CSV                    as F
import qualified Frames.InCore                 as FI
import qualified Frames.TH                     as F
import qualified Frames.Melt                   as F
import qualified Text.Read                     as TR

import qualified Frames.Folds                  as FF
import qualified Frames.MapReduce              as FMR
import qualified Frames.ParseableTypes         as FP
import qualified Frames.Transform              as FT
import qualified Frames.MaybeUtils             as FM
import qualified Frames.Utils                  as FU
import qualified Frames.MapReduce              as MR
import qualified Frames.Enumerations           as FE
import qualified Frames.Serialize              as FS
import qualified Frames.SimpleJoins            as FJ
import qualified Frames.Visualization.VegaLite.Data
                                               as FV
--import qualified Graphics.Vega.VegaLite        as GV

{-
import qualified Data.IndexedSet               as IS
import qualified Numeric.GLM.ProblemTypes      as GLM
import qualified Numeric.GLM.ModelTypes      as GLM
import qualified Numeric.GLM.Predict            as GLM
import qualified Numeric.LinearAlgebra         as LA
-}

import           Data.Hashable                  ( Hashable )
import qualified Data.Vector                   as V
--import qualified Data.Vector.Boxed             as VB
import           GHC.Generics                   ( Generic, Rep )

import qualified Knit.Report as K
import qualified Polysemy.Error                as P (mapError, Error)
import qualified Polysemy                as P (raise)

import GHC.TypeLits (Symbol)
import Data.Kind (Type)

cpsVoterPUMSLoader :: (K.KnitEffects r, BR.CacheEffects r)
                   => K.Sem r (K.ActionWithCacheTime r (F.FrameRec CPSVoterPUMS))
cpsVoterPUMSLoader = do
  let cpsPUMSDataPath = BR.LocalData $ T.pack BR.cpsVoterPUMSCSV
  cachedStateAbbrCrosswalk <- BR.stateAbbrCrosswalkLoader
  cachedCPSDataPath <- K.liftKnit $ BR.dataPathWithCacheTime cpsPUMSDataPath
  let cachedDeps = (,) <$> cachedStateAbbrCrosswalk <*> cachedCPSDataPath
  BR.retrieveOrMakeFrame "data/cpsVoterPUMSWithAbbrs.bin" cachedDeps $ \(stateAbbrCrosswalk, dataPath) -> do
    let filter r = (F.rgetField @BR.CPSAGE r >= 18) && (F.rgetField @BR.CPSCITIZEN r /= 5)
    withoutAbbr <- K.ignoreCacheTimeM
                   $ BR.cachedFrameLoader @(F.RecordColumns BR.CPSVoterPUMS_Raw) @CPSVoterPUMS'
                   dataPath
                   Nothing
                   (Just filter)
                   transformCPSVoterPUMSRow
                   Nothing
                   "cpsVoterPUMS.bin"
    fmap (fmap F.rcast) (K.knitMaybe "missing state abbreviation in state abbreviation crosswalk"
                          $ FJ.leftJoinM @'[BR.StateFIPS] withoutAbbr stateAbbrCrosswalk)

-- NB: This should not be used for state-level rollup since some rows will be duplicated if the county is in more than one CD.
cpsVoterPUMSWithCDLoader :: (K.KnitEffects r, BR.CacheEffects r)
                          => K.Sem r (K.ActionWithCacheTime r (F.FrameRec (CPSVoterPUMS V.++ [BR.CongressionalDistrict, BR.CountyWeight])))
cpsVoterPUMSWithCDLoader = do
  cachedCPSVoterPUMS <- cpsVoterPUMSLoader
  cachedCountyToCD <- BR.county2010ToCD116Loader
  let cachedDeps = (,) <$> cachedCPSVoterPUMS <*> cachedCountyToCD
  BR.retrieveOrMakeFrame "data/cpsVoterPUMSWithAbbrsAndCDs.bin" cachedDeps $ \(cpsVoterPUMS, countyToCD) -> do
    K.knitEither . either (Left . show) Right
      $ FJ.leftJoinE
      @[BR.StateFIPS, BR.CountyFIPS]
      (F.filterFrame ((> 0) . F.rgetField @BR.CountyFIPS) cpsVoterPUMS)
      (fmap (F.rcast @[BR.CountyFIPS, BR.StateFIPS, BR.CongressionalDistrict, BR.CountyWeight]) countyToCD)

type CPSVoterPUMSWeight = "CPSVoterPUMSWeight" F.:-> Double

type CPSVoterPUMS = '[ BR.Year
                     , BR.StateFIPS
                     , BR.StateAbbreviation
                     , BR.CountyFIPS
                     , BR.Age4C
                     , BR.SexC
                     , BR.RaceAlone4C
                     , BR.HispC
                     , BR.IsCitizen
                     , BR.CollegeGradC
                     , BR.InCollege
                     , BR.VoteWhyNotC
                     , BR.RegWhyNotC
                     , BR.VoteHowC
                     , BR.VoteWhenC
                     , BR.VotedYNC
                     , BR.RegisteredYNC
                     , CPSVoterPUMSWeight
                     ]

type CPSVoterPUMS' = V.RDelete BR.StateAbbreviation CPSVoterPUMS

cpsVoterPUMSRollup
  :: forall cs rs ks ds.
  ( cs F.⊆ rs
  , ks F.⊆ (ks V.++ cs)
  , cs F.⊆ (ks V.++ cs)
  , FI.RecVec (ks V.++ ds)
  , Ord (F.Record ks)
  )
  => (F.Record rs -> F.Record ks)
  -> FL.Fold (F.Record cs) (F.Record ds)
  -> FL.Fold (F.Record rs) (F.FrameRec (ks V.++ ds))
cpsVoterPUMSRollup getKeys foldData =
  FMR.concatFold
  $ FMR.mapReduceFold
  (FMR.Unpack $ \r -> [getKeys r `V.rappend` F.rcast @cs r])
  (FMR.assignKeysAndData @ks @cs)
  (FMR.foldAndAddKey foldData)

cpsVoterPUMSRollupWithDefault
  ::  forall cs rs ls ks ds.
  ( cs F.⊆ rs
  , (ls V.++ ks) F.⊆ (ls V.++ ks V.++ cs)
  , cs F.⊆ (ls V.++ ks V.++ cs)
  , FI.RecVec (ls V.++ ks V.++ ds)
  , Ord (F.Record (ls V.++ ks))
  , BR.FiniteSet (F.Record ks)
  , (ls V.++ (ks V.++ ds)) ~ ((ls V.++ ks) V.++ ds)
  , Ord (F.Record ls)
  , ls F.⊆ ((ls V.++ ks) V.++ ds)
  , (ks V.++ ds) F.⊆ ((ls V.++ ks) V.++ ds)
  , ks F.⊆ (ks V.++ ds)
  , ds F.⊆ (ks V.++ ds)
  )
  => (F.Record rs -> F.Record ls)
  -> (F.Record rs -> F.Record ks)
  -> FL.Fold (F.Record cs) (F.Record ds)
  -> F.Record ds -- default value if there are no records for a given value of F.Record ks
  -> FL.Fold (F.Record rs) (F.FrameRec (ls V.++ ks V.++ ds))
cpsVoterPUMSRollupWithDefault getLoc getKey foldData defD =
  let addDefaultF = FMR.concatFold
                    $ FMR.mapReduceFold
                    FMR.noUnpack
                    (FMR.assignKeysAndData @ls @(ks V.++ ds))
                    (FMR.makeRecsWithKey id
                     $ FMR.ReduceFold
                     $ const
                     $ BR.addDefaultRec @ks defD
                    )
  in FL.fold addDefaultF <$> cpsVoterPUMSRollup (\r -> getLoc r `V.rappend` getKey r) foldData

cpsVoterPUMSRollupWeightedCounts
  :: forall cs rs ls ks ds.
  ( cs F.⊆ rs
  , (ls V.++ ks) F.⊆ (ls V.++ ks V.++ cs)
  , cs F.⊆ (ls V.++ ks V.++ cs)
  , FI.RecVec (ls V.++ ks V.++ ds)
  , Ord (F.Record (ls V.++ ks))
  , (ls V.++ (ks V.++ ds)) ~ ((ls V.++ ks) V.++ ds)
  , BR.FiniteSet (F.Record ks)
  , Ord (F.Record ls)
  , ls F.⊆ ((ls V.++ ks) V.++ ds)
  , (ks V.++ ds) F.⊆ ((ls V.++ ks) V.++ ds)
  , ks F.⊆ (ks V.++ ds)
  , ds F.⊆ (ks V.++ ds)
  )
  => (F.Record rs -> F.Record ls)
  -> (F.Record rs -> F.Record ks)
  -> (F.Record cs -> Bool) -- which to include
  -> (F.Record cs -> Bool) -- which to count
  -> (F.Record cs -> Double) -- get weight
  -> (Double -> F.Record ds) -- turn the weighted counts into a record
  -> F.Record ds -- default
  -> FL.Fold (F.Record rs) (F.FrameRec (ls V.++ ks V.++ ds))
cpsVoterPUMSRollupWeightedCounts getLoc getKey filterData countIf wgt countToRec defD =
  let countF :: FL.Fold (F.Record cs) (F.Record ds)
      countF =
        let wgtdAllF = FL.premap wgt FL.sum
            wgtdCountF = FL.prefilter countIf $ FL.premap wgt FL.sum
            safeDiv n d = if d > 0 then n/d else 0
            wF = safeDiv <$> wgtdCountF <*> wgtdAllF -- this is acceptable here since (d == 0) iff (n == 0)
        in fmap countToRec wF
      ewFold :: FL.Fold (F.Record cs) (F.Record ds)
      ewFold = FL.prefilter filterData countF
  in cpsVoterPUMSRollupWithDefault getLoc getKey ewFold defD

cpsVoterPUMSElectoralWeights
  :: forall cs rs ls ks.
  ((ls V.++ ks) F.⊆ (ls V.++ ks V.++ cs)
  , cs F.⊆ (ls V.++ ks V.++ cs)
  , cs F.⊆ rs
  , F.ElemOf cs BR.IsCitizen
  , F.ElemOf cs BR.VotedYNC
  , FI.RecVec (ls V.++ ks V.++ BR.EWCols)
  , (ls V.++ (ks V.++ BR.EWCols)) ~ ((ls V.++ ks) V.++ BR.EWCols)
  , Ord (F.Record (ls V.++ ks))
  , Ord (F.Record ls)
  , BR.FiniteSet (F.Record ks)
  , ls F.⊆ ((ls V.++ ks) V.++ BR.EWCols)
  , (ks V.++ BR.EWCols) F.⊆ ((ls V.++ ks) V.++ BR.EWCols)
  , ks F.⊆ (ks V.++ BR.EWCols)
  , BR.EWCols F.⊆ (ks V.++ BR.EWCols)
  )
  => (F.Record rs -> F.Record ls)
  -> (F.Record rs -> F.Record ks)
  -> (F.Record cs -> Double)
  -> FL.Fold (F.Record rs) (F.FrameRec (ls V.++ ks V.++ BR.EWCols))
cpsVoterPUMSElectoralWeights getLoc getKey getWgt =
  let toRec :: Double -> F.Record BR.EWCols
      toRec w =  BR.EW_Census F.&: BR.EW_Citizen F.&: w F.&: V.RNil
      citizen r = F.rgetField @BR.IsCitizen r
      possibleVoter r = cpsPossibleVoter $ F.rgetField @BR.VotedYNC r
      voted r = cpsVoted $ F.rgetField @BR.VotedYNC r
  in cpsVoterPUMSRollupWeightedCounts @cs
     getLoc
     getKey
     (\r -> citizen r &&  possibleVoter r)
     voted
     getWgt
     toRec
     (toRec 0) -- this is a reasonable default to use for inference but not great to use this data directly

cpsVoterPUMSElectoralWeightsByState
  :: forall ks.
  (F.ElemOf (ks ++ '[BR.IsCitizen, BR.VotedYNC, CPSVoterPUMSWeight]) CPSVoterPUMSWeight
  , F.ElemOf (ks ++ '[BR.IsCitizen, BR.VotedYNC, CPSVoterPUMSWeight]) BR.IsCitizen
  , F.ElemOf (ks ++ '[BR.IsCitizen, BR.VotedYNC, CPSVoterPUMSWeight]) BR.VotedYNC
  , FI.RecVec (ks ++ BR.EWCols)
  , Ord (F.Record ks)
  , ks F.⊆ ('[BR.Year, BR.StateAbbreviation, BR.StateFIPS] V.++ ks V.++ '[BR.IsCitizen, BR.VotedYNC, CPSVoterPUMSWeight])
  , BR.FiniteSet (F.Record ks)
  , F.ElemOf (ks V.++ BR.EWCols) BR.ElectoralWeight
  , F.ElemOf (ks V.++ BR.EWCols) BR.ElectoralWeightOf
  , F.ElemOf (ks V.++ BR.EWCols) BR.ElectoralWeightSource
  , (ks V.++ BR.EWCols) F.⊆ ('[BR.Year, BR.StateAbbreviation, BR.StateFIPS] V.++ ks V.++ BR.EWCols)
  , ks F.⊆ (ks V.++ BR.EWCols)
  )
  => (F.Record CPSVoterPUMS -> F.Record ks)
  -> FL.Fold (F.Record CPSVoterPUMS) (F.FrameRec ('[BR.Year, BR.StateAbbreviation, BR.StateFIPS] V.++ ks V.++ BR.EWCols))
cpsVoterPUMSElectoralWeightsByState getCatKey =
  cpsVoterPUMSElectoralWeights @[BR.IsCitizen, BR.VotedYNC, CPSVoterPUMSWeight]
  (F.rcast @[BR.Year, BR.StateAbbreviation, BR.StateFIPS]) getCatKey (F.rgetField @CPSVoterPUMSWeight)

cpsVoterPUMSNationalElectoralWeights
  :: forall ks.
  (F.ElemOf (ks ++ '[BR.IsCitizen, BR.VotedYNC, CPSVoterPUMSWeight]) CPSVoterPUMSWeight
  , F.ElemOf (ks ++ '[BR.IsCitizen, BR.VotedYNC, CPSVoterPUMSWeight]) BR.IsCitizen
  , F.ElemOf (ks ++ '[BR.IsCitizen, BR.VotedYNC, CPSVoterPUMSWeight]) BR.VotedYNC
  , FI.RecVec (ks ++ BR.EWCols)
  , Ord (F.Record ks)
  , ks F.⊆ ('[BR.Year] V.++ ks V.++ '[BR.IsCitizen, BR.VotedYNC, CPSVoterPUMSWeight])
  , BR.FiniteSet (F.Record ks)
  , F.ElemOf (ks V.++ BR.EWCols) BR.ElectoralWeight
  , F.ElemOf (ks V.++ BR.EWCols) BR.ElectoralWeightOf
  , F.ElemOf (ks V.++ BR.EWCols) BR.ElectoralWeightSource
  , (ks V.++ BR.EWCols) F.⊆ ('[BR.Year] V.++ ks V.++ BR.EWCols)
  , ks F.⊆ (ks V.++ BR.EWCols)
  )
  => (F.Record CPSVoterPUMS -> F.Record ks)
  -> FL.Fold (F.Record CPSVoterPUMS) (F.FrameRec ('[BR.Year] V.++ ks V.++ BR.EWCols))
cpsVoterPUMSNationalElectoralWeights getCatKey =
  cpsVoterPUMSElectoralWeights @[BR.IsCitizen, BR.VotedYNC, CPSVoterPUMSWeight] (F.rcast @'[BR.Year]) getCatKey (F.rgetField @CPSVoterPUMSWeight)

cpsVoterPUMSElectoralWeightsByCD
  :: forall ks.
  (F.ElemOf (ks ++ '[BR.IsCitizen, BR.VotedYNC, CPSVoterPUMSWeight, BR.CountyWeight]) CPSVoterPUMSWeight
  , F.ElemOf (ks ++ '[BR.IsCitizen, BR.VotedYNC, CPSVoterPUMSWeight, BR.CountyWeight]) BR.CountyWeight
  , F.ElemOf (ks ++ '[BR.IsCitizen, BR.VotedYNC, CPSVoterPUMSWeight, BR.CountyWeight]) BR.IsCitizen
  , F.ElemOf (ks ++ '[BR.IsCitizen, BR.VotedYNC, CPSVoterPUMSWeight, BR.CountyWeight]) BR.VotedYNC
  , FI.RecVec (ks ++ BR.EWCols)
  , Ord (F.Record ks)
  , ks F.⊆ ('[BR.Year, BR.StateAbbreviation, BR.StateFIPS, BR.CongressionalDistrict]
             V.++ ks
             V.++ '[BR.IsCitizen, BR.VotedYNC, CPSVoterPUMSWeight, BR.CountyWeight])
  , BR.FiniteSet (F.Record ks)
  , F.ElemOf (ks V.++ BR.EWCols) BR.ElectoralWeight
  , F.ElemOf (ks V.++ BR.EWCols) BR.ElectoralWeightOf
  , F.ElemOf (ks V.++ BR.EWCols) BR.ElectoralWeightSource
  , (ks V.++ BR.EWCols) F.⊆ ('[BR.Year, BR.StateAbbreviation, BR.StateFIPS, BR.CongressionalDistrict] V.++ ks V.++ BR.EWCols)
  , ks F.⊆ (ks V.++ BR.EWCols)
  )
  => (F.Record (CPSVoterPUMS V.++ [BR.CongressionalDistrict, BR.CountyWeight]) -> F.Record ks)
  -> FL.Fold (F.Record (CPSVoterPUMS V.++ [BR.CongressionalDistrict, BR.CountyWeight]))
  (F.FrameRec ('[BR.Year, BR.StateAbbreviation, BR.StateFIPS, BR.CongressionalDistrict] V.++ ks V.++ BR.EWCols))
cpsVoterPUMSElectoralWeightsByCD getCatKey =
  let wgt r = F.rgetField @CPSVoterPUMSWeight r * F.rgetField @BR.CountyWeight r
  in cpsVoterPUMSElectoralWeights @[BR.IsCitizen, BR.VotedYNC, CPSVoterPUMSWeight, BR.CountyWeight]
     (F.rcast @[BR.Year, BR.StateAbbreviation, BR.StateFIPS, BR.CongressionalDistrict]) getCatKey wgt

-- rollup for MRP
cpsCountVotersByStateF
  :: forall ks.
  (Ord (F.Record ks)
  , FI.RecVec (ks V.++ BRCF.CountCols)
  , ks F.⊆ CPSVoterPUMS
  )
  => (F.Record CPSVoterPUMS -> F.Record ks)
  -> Int -- year
  -> FMR.Fold
  (F.Record CPSVoterPUMS)
  (F.FrameRec ('[BR.StateAbbreviation] V.++ ks V.++ BRCF.CountCols))
cpsCountVotersByStateF getCatKey year =
  let isYear y r = F.rgetField @BR.Year r == y
      possible r = cpsPossibleVoter $ F.rgetField @BR.VotedYNC r
      citizen r = F.rgetField @BR.IsCitizen r
      includeRow r = isYear year r &&  possible r && citizen r
      voted r = cpsVoted $ F.rgetField @BR.VotedYNC r
  in BRCF.weightedCountFold
     (F.rcast @(BR.StateAbbreviation ': ks))
     (F.rcast @[BR.VotedYNC, CPSVoterPUMSWeight])
     includeRow
     voted
     (F.rgetField @CPSVoterPUMSWeight)

cpsCountVotersByCDF
  :: forall ks.
  (Ord (F.Record ks)
  , FI.RecVec (ks V.++ BRCF.CountCols)
  , ks F.⊆ (CPSVoterPUMS V.++ [BR.CongressionalDistrict, BR.CountyWeight])
  )
  => (F.Record (CPSVoterPUMS V.++ [BR.CongressionalDistrict, BR.CountyWeight]) -> F.Record ks)
  -> Int -- year
  -> FMR.Fold
  (F.Record (CPSVoterPUMS V.++ [BR.CongressionalDistrict, BR.CountyWeight]))
  (F.FrameRec ('[BR.StateAbbreviation, BR.CongressionalDistrict] V.++ ks V.++ BRCF.CountCols))
cpsCountVotersByCDF getCatKey year =
  let isYear y r = F.rgetField @BR.Year r == y
      possible r = cpsPossibleVoter $ F.rgetField @BR.VotedYNC r
      citizen r = F.rgetField @BR.IsCitizen r
      includeRow r = isYear year r &&  possible r && citizen r
      voted r = cpsVoted $ F.rgetField @BR.VotedYNC r
      wgt r = F.rgetField @CPSVoterPUMSWeight r * F.rgetField @BR.CountyWeight r
  in BRCF.weightedCountFold
     (F.rcast @([BR.StateAbbreviation, BR.CongressionalDistrict] V.++ ks))
     (F.rcast  @[BR.VotedYNC, CPSVoterPUMSWeight, BR.CountyWeight])
     includeRow
     voted
     wgt

-- We give the option of counting "In College" as "College Grad". This is different from what the census summary tables do.
-- NB: This needs to be done consistently with the demographics.
-- We don't have this information for the preferences, at least not from CCES, so doing this amounts to assigning
-- "in college" people to either Grad or NonGrad buckets in terms of voting pref.

cpsKeysToASER5 :: Bool
               -> F.Record '[BR.Age4C, BR.SexC, BR.CollegeGradC, BR.InCollege, BR.RaceAlone4C, BR.HispC]
               -> F.Record BR.CatColsASER5
cpsKeysToASER5 addInCollegeToGrads r =
  let cg = F.rgetField @BR.CollegeGradC r
      ic = addInCollegeToGrads && F.rgetField @BR.InCollege r
      ra4 =  F.rgetField @BR.RaceAlone4C r
      h = F.rgetField @BR.HispC r
  in (BR.age4ToSimple $ F.rgetField @BR.Age4C r)
     F.&: (F.rgetField @BR.SexC r)
     F.&: (if cg == BR.Grad || ic then BR.Grad else BR.NonGrad)
     F.&: BR.race5FromRaceAlone4AndHisp True ra4 h
     F.&: V.RNil

cpsKeysToASER4H :: Bool
                -> F.Record [BR.Age4C, BR.SexC, BR.CollegeGradC, BR.InCollege, BR.RaceAlone4C, BR.HispC]
                -> F.Record [BR.SimpleAgeC, BR.SexC, BR.CollegeGradC, BR.RaceAlone4C, BR.HispC]
cpsKeysToASER4H addInCollegeToGrads r =
  let cg = F.rgetField @BR.CollegeGradC r
      ic = addInCollegeToGrads && F.rgetField @BR.InCollege r
  in (BR.age4ToSimple $ F.rgetField @BR.Age4C r)
     F.&: (F.rgetField @BR.SexC r)
     F.&: (if cg == BR.Grad || ic then BR.Grad else BR.NonGrad)
     F.&: F.rgetField @BR.RaceAlone4C r
     F.&: F.rgetField @BR.HispC r
     F.&: V.RNil

cpsKeysToASER :: Bool -> F.Record '[BR.Age4C, BR.SexC, BR.CollegeGradC, BR.InCollege, BR.RaceAlone4C, BR.HispC] -> F.Record BR.CatColsASER
cpsKeysToASER addInCollegeToGrads r =
  let cg = F.rgetField @BR.CollegeGradC r
      ic = addInCollegeToGrads && F.rgetField @BR.InCollege r
      ra4 =  F.rgetField @BR.RaceAlone4C r
      h = F.rgetField @BR.HispC r
  in (BR.age4ToSimple $ F.rgetField @BR.Age4C r)
     F.&: (F.rgetField @BR.SexC r)
     F.&: (if cg == BR.Grad || ic then BR.Grad else BR.NonGrad)
     F.&: (BR.simpleRaceFromRaceAlone4AndHisp True ra4 h)
     F.&: V.RNil

{-
cpsKeysToASER4 :: Bool -> F.Record '[BR.Age4C, BR.SexC, BR.CollegeGradC, BR.InCollege, BR.Race5C] -> F.Record BR.CatColsASER4
cpsKeysToASER4 addInCollegeToGrads r =
  let cg = F.rgetField @BR.CollegeGradC r
      ic = addInCollegeToGrads && F.rgetField @BR.InCollege r
  in (BR.age4ToSimple $ F.rgetField @BR.Age4C r)
     F.&: (F.rgetField @BR.SexC r)
     F.&: (if cg == BR.Grad || ic then BR.Grad else BR.NonGrad)
     F.&: (BR.race4FromRace5 $ F.rgetField @BR.Race5C r)
     F.&: V.RNil

cpsKeysToASER :: Bool -> F.Record '[BR.Age4C, BR.SexC, BR.CollegeGradC, BR.InCollege, BR.Race5C] -> F.Record BR.CatColsASER
cpsKeysToASER addInCollegeToGrads r =
  let cg = F.rgetField @BR.CollegeGradC r
      ic = addInCollegeToGrads && F.rgetField @BR.InCollege r
  in (BR.age4ToSimple $ F.rgetField @BR.Age4C r)
     F.&: (F.rgetField @BR.SexC r)
     F.&: (if cg == BR.Grad || ic then BR.Grad else BR.NonGrad)
     F.&: (BR.simpleRaceFromRace5 $ F.rgetField @BR.Race5C r)
     F.&: V.RNil

cpsKeysToASE :: Bool -> F.Record '[BR.Age4C, BR.SexC, BR.CollegeGradC, BR.InCollege, BR.Race5C] -> F.Record BR.CatColsASE
cpsKeysToASE addInCollegeToGrads r =
  let cg = F.rgetField @BR.CollegeGradC r
      ic = addInCollegeToGrads && F.rgetField @BR.InCollege r
  in (BR.age4ToSimple $ F.rgetField @BR.Age4C r)
     F.&: (F.rgetField @BR.SexC r)
     F.&: (if cg == BR.Grad || ic then BR.Grad else BR.NonGrad)
     F.&: V.RNil

cpsKeysToASR :: F.Record '[BR.Age4C, BR.SexC, BR.CollegeGradC, BR.InCollege, BR.Race5C] -> F.Record BR.CatColsASR
cpsKeysToASR r =
  (BR.age4ToSimple $ F.rgetField @BR.Age4C r)
  F.&: (F.rgetField @BR.SexC r)
  F.&: (BR.simpleRaceFromRace5 $ F.rgetField @BR.Race5C r)
  F.&: V.RNil

cpsKeysToIdentity :: F.Record '[BR.Age4C, BR.SexC, BR.CollegeGradC, BR.InCollege, BR.Race5C] -> F.Record '[]
cpsKeysToIdentity = const V.RNil
-}
-- we have to drop all records with age < 18
intToAge4 :: Int -> BR.Age4
intToAge4 n
  | n < 18 = error "CPS Voter record with age < 18"
  | n < 24 = BR.A4_18To24
  | n < 45 = BR.A4_25To44
  | n < 65 = BR.A4_45To64
  | otherwise = BR.A4_65AndOver

intToCollegeGrad :: Int -> BR.CollegeGrad
intToCollegeGrad n = if n < 111 then BR.NonGrad else BR.Grad

intToInCollege :: Int -> Bool
intToInCollege n = n == 4

intToSex :: Int -> BR.Sex
intToSex n = if n == 1 then BR.Male else BR.Female

intToRaceAlone4 :: Int -> BR.RaceAlone4
intToRaceAlone4 rN
  | rN == 100 = BR.RA4_White
  | rN == 200 = BR.RA4_Black
  | rN == 651 = BR.RA4_Asian
  | otherwise = BR.RA4_Other

intToHisp :: Int -> BR.Hisp
intToHisp hN
  | (hN >= 100) && (hN <= 901) = BR.Hispanic
  | otherwise = BR.NonHispanic

intsToRace5 :: Int -> Int -> BR.Race5
intsToRace5 hN rN
  | (hN >= 100) && (hN <= 901) = BR.R5_Hispanic
  | rN == 100 = BR.R5_WhiteNonHispanic
  | rN == 200 = BR.R5_Black
  | rN == 651 = BR.R5_Asian
  | otherwise = BR.R5_Other

intToIsCitizen :: Int -> Bool
intToIsCitizen n = n /= 5

intToVoteWhyNot :: Int -> BR.VoteWhyNot
intToVoteWhyNot n
  | n == 1 = BR.VWN_PhysicallyUnable
  | n == 2 = BR.VWN_Away
  | n == 3 = BR.VWN_Forgot
  | n == 4 = BR.VWN_NotInterested
  | n == 5 = BR.VWN_Busy
  | n == 6 = BR.VWN_Transport
  | n == 7 = BR.VWN_DislikeChoices
  | n == 8 = BR.VWN_RegIssue
  | n == 9 = BR.VWN_Weather
  | n == 10 = BR.VWN_BadPollingPlace
  | otherwise = BR.VWN_Other

intToRegWhyNot :: Int -> BR.RegWhyNot
intToRegWhyNot n
  | n == 1 = BR.RWN_MissedDeadline
  | n == 2 = BR.RWN_DidntKnowHow
  | n == 3 = BR.RWN_Residency
  | n == 4 = BR.RWN_PhysicallyUnable
  | n == 5 = BR.RWN_Language
  | n == 6 = BR.RWN_NotInterested
  | n == 7 = BR.RWN_MyVoteIrrelevant
  | otherwise = BR.RWN_Other

intToVoteHow :: Int -> BR.VoteHow
intToVoteHow n
  | n == 1 = BR.VH_InPerson
  | n == 2 = BR.VH_ByMail
  | otherwise = BR.VH_Other

intToVoteWhen :: Int -> BR.VoteWhen
intToVoteWhen n
  | n == 1 = BR.VW_ElectionDay
  | n == 2 = BR.VW_BeforeElectionDay
  | otherwise = BR.VW_Other

intToVotedYN :: Int -> BR.VotedYN
intToVotedYN n
  | n == 1 = BR.VYN_DidNotVote
  | n == 2 = BR.VYN_Voted
  | n == 96 = BR.VYN_Refused
  | n == 97 = BR.VYN_DontKnow
  | n == 98 = BR.VYN_NoResponse
  | otherwise = BR.VYN_NotInUniverse

cpsPossibleVoter :: BR.VotedYN -> Bool
cpsPossibleVoter BR.VYN_DidNotVote = True
cpsPossibleVoter BR.VYN_Voted = True
cpsPossibleVoter BR.VYN_Refused = False
cpsPossibleVoter BR.VYN_DontKnow = False
cpsPossibleVoter BR.VYN_NoResponse = False
cpsPossibleVoter BR.VYN_NotInUniverse = False

cpsVoted :: BR.VotedYN -> Bool
cpsVoted BR.VYN_Voted = True
cpsVoted _ = False

intToRegisteredYN :: Int -> BR.RegisteredYN
intToRegisteredYN n
  | n == 1 = BR.RYN_NotRegistered
  | n == 2 = BR.RYN_Registered
  | otherwise = BR.RYN_Other

transformCPSVoterPUMSRow :: BR.CPSVoterPUMS_Raw -> F.Record CPSVoterPUMS'
transformCPSVoterPUMSRow = F.rcast . addCols  where
  addCols = (FT.addOneFromOne @BR.CPSRACE @BR.RaceAlone4C intToRaceAlone4)
            . (FT.addOneFromOne @BR.CPSHISPAN @BR.HispC intToHisp)
            . (FT.addName @BR.CPSVOSUPPWT @CPSVoterPUMSWeight)
            . (FT.addName @BR.CPSCOUNTY @BR.CountyFIPS)
            . (FT.addName @BR.CPSSTATEFIP @BR.StateFIPS)
            . (FT.addName @BR.CPSYEAR @BR.Year)
            . (FT.addOneFromOne @BR.CPSVOREG @BR.RegisteredYNC intToRegisteredYN)
            . (FT.addOneFromOne @BR.CPSVOTED @BR.VotedYNC intToVotedYN)
            . (FT.addOneFromOne @BR.CPSVOTEWHEN @BR.VoteWhenC intToVoteWhen)
            . (FT.addOneFromOne @BR.CPSVOTEHOW @BR.VoteHowC intToVoteHow)
            . (FT.addOneFromOne @BR.CPSVOYNOTREG @BR.RegWhyNotC intToRegWhyNot)
            . (FT.addOneFromOne @BR.CPSVOWHYNOT @BR.VoteWhyNotC intToVoteWhyNot)
            . (FT.addOneFromOne @BR.CPSSCHLCOLL @BR.InCollege intToInCollege)
            . (FT.addOneFromOne @BR.CPSEDUC @BR.CollegeGradC intToCollegeGrad)
            . (FT.addOneFromOne @BR.CPSCITIZEN @BR.IsCitizen intToIsCitizen)
            . (FT.addOneFromOne @BR.CPSSEX @BR.SexC intToSex)
            . (FT.addOneFromOne @BR.CPSAGE @BR.Age4C intToAge4)

{-
transformCPSVoterPUMSRow :: BR.CPSVoterPUMS_Raw -> F.Record CPSVoterPUMS'
transformCPSVoterPUMSRow r = F.rcast @CPSVoterPUMS' (mutate r) where
  addAge = FT.recordSingleton @BR.Age4C . intToAge4 . F.rgetField @BR.CPSAGE
  addSex =  FT.recordSingleton @BR.SexC . intToSex . F.rgetField @BR.CPSSEX
  hN = F.rgetField @BR.CPSHISPAN
  rN = F.rgetField @BR.CPSRACE
  addRace r = FT.recordSingleton @BR.Race5C $ intsToRace5 (hN r) (rN r)
  addIsCitizen = FT.recordSingleton @BR.IsCitizen . intToIsCitizen . F.rgetField @BR.CPSCITIZEN
  addCollegeGrad = FT.recordSingleton @BR.CollegeGradC . intToCollegeGrad . F.rgetField @BR.CPSEDUC
  addInCollege = FT.recordSingleton @BR.InCollege . intToInCollege . F.rgetField @BR.CPSSCHLCOLL
  addVoteWhyNot = FT.recordSingleton @BR.VoteWhyNotC . intToVoteWhyNot . F.rgetField @BR.CPSVOWHYNOT
  addRegWhyNot = FT.recordSingleton @BR.RegWhyNotC . intToRegWhyNot . F.rgetField @BR.CPSVOYNOTREG
  addVoteHow = FT.recordSingleton @BR.VoteHowC . intToVoteHow . F.rgetField @BR.CPSVOTEHOW
  addVoteWhen = FT.recordSingleton @BR.VoteWhenC . intToVoteWhen . F.rgetField @BR.CPSVOTEWHEN
  addVotedYN = FT.recordSingleton @BR.VotedYNC . intToVotedYN . F.rgetField @BR.CPSVOTED
  addRegisteredYN = FT.recordSingleton @BR.RegisteredYNC . intToRegisteredYN . F.rgetField @BR.CPSVOREG
  mutate = FT.retypeColumn @BR.CPSYEAR @BR.Year
           . FT.retypeColumn @BR.CPSSTATEFIP @BR.StateFIPS
           . FT.retypeColumn @BR.CPSCOUNTY @BR.CountyFIPS
           . FT.mutate addAge
           . FT.mutate addSex
           . FT.mutate addRace
           . FT.mutate addIsCitizen
           . FT.mutate addCollegeGrad
           . FT.mutate addInCollege
           . FT.mutate addVoteWhyNot
           . FT.mutate addRegWhyNot
           . FT.mutate addVoteHow
           . FT.mutate addVoteWhen
           . FT.mutate addVotedYN
           . FT.mutate addRegisteredYN
           . FT.retypeColumn @BR.CPSVOSUPPWT @CPSVoterPUMSWeight
-}

-- to use in maybeRecsToFrame
-- if SCHG indicates not in school we map to 0 so we will interpret as "Not In College"
{-
fixPUMSRow :: F.Rec (Maybe F.:. F.ElField) PUMS_Raw -> F.Rec (Maybe F.:. F.ElField) PUMS_Raw
fixPUMSRow r = (F.rsubset %~ missingInCollegeTo0)
               $ r where
  missingInCollegeTo0 :: F.Rec (Maybe :. F.ElField) '[BR.PUMSSCHG] -> F.Rec (Maybe :. F.ElField) '[BR.PUMSSCHG]
  missingInCollegeTo0 = FM.fromMaybeMono 0
-}


{-# LANGUAGE GADTs             #-}
{-# LANGUAGE OverloadedStrings #-}

module MatrixMarket (Matrix(..), readMatrix) where

import Control.Applicative                      hiding ( many )

import Data.Int
import Data.Complex
import Data.Attoparsec.ByteString.Char8
import Data.ByteString.Lex.Fractional
import qualified Data.Attoparsec.Lazy           as L
import qualified Data.ByteString.Lazy           as L

-- | Specifies the element type.  Pattern matrices do not have any elements,
-- only indices, and only make sense for coordinate matrices and vectors.
--
data Field  = Real | Complex | Integer | Pattern
    deriving (Eq, Show)

-- | Specifies either sparse or dense storage.  In sparse (\"coordinate\")
-- storage, elements are given in (i,j,x) triplets for matrices (or (i,x) for
-- vectors).  Indices are 1-based, so that A(1,1) is the first element of a
-- matrix, and x(1) is the first element of a vector.
--
-- In dense (\"array\") storage, elements are given in column-major order.
--
-- In both cases, each element is given on a separate line.
--
data Format = Coordinate | Array
    deriving (Eq, Show)

-- | Specifies any special structure in the matrix.  For symmetric and hermition
-- matrices, only the lower-triangular part of the matrix is given. For skew
-- matrices, only the entries below the diagonal are stored.
--
data Structure = General | Symmetric | Hermitian | Skew
    deriving (Eq, Show)

-- We really want a type parameter to Matrix, but I think that requires some
-- kind of dynamic typing so that we can determine (a ~ Integral) or (a ~
-- RealFloat), and so forth, depending on the file being read. This will do for
-- our purposes...
--
-- Format is: (rows,columns) nnz [(row,column,value)]
--
data Matrix
  = PatternMatrix (Int,Int) Int [(Int32,Int32)]
  | IntMatrix     (Int,Int) Int [(Int32,Int32,Int)]
  | RealMatrix    (Int,Int) Int [(Int32,Int32,Float)]
  | ComplexMatrix (Int,Int) Int [(Int32,Int32,Complex Float)]
  deriving Show

--------------------------------------------------------------------------------
-- Combinators
--------------------------------------------------------------------------------

comment :: Parser ()
comment = char '%' *> skipWhile (not . eol) *> endOfLine
  where
    eol w = w `elem` ("\n\r" :: String)

floating :: Fractional a => Parser a
floating = do
  mv <- readDecimal <$> (skipSpace *> takeTill isSpace)  -- readDecimal does the fancy stuff
  case mv of
       Just (v,_) -> return v
       Nothing    -> fail "floating-point number"

integral :: Integral a => Parser a
integral = skipSpace *> decimal

format :: Parser Format
format =  string "coordinate" *> pure Coordinate
      <|> string "array"      *> pure Array
      <?> "matrix format"

field :: Parser Field
field =  string "real"    *> pure Real
     <|> string "complex" *> pure Complex
     <|> string "integer" *> pure Integer
     <|> string "pattern" *> pure Pattern
     <?> "matrix field"

structure :: Parser Structure
structure =  string "general"        *> pure General
         <|> string "symmetric"      *> pure Symmetric
         <|> string "hermitian"      *> pure Hermitian
         <|> string "skew-symmetric" *> pure Skew
         <?> "matrix structure"

header :: Parser (Format,Field,Structure)
header =  string "%%MatrixMarket matrix"
       >> (,,) <$> (skipSpace *> format)
               <*> (skipSpace *> field)
               <*> (skipSpace *> structure)
               <*  endOfLine
               <?> "MatrixMarket header"

extent :: Parser (Int,Int,Int)
extent = do
  [m,n,l] <- skipWhile isSpace *> count 3 integral <* endOfLine
  return (m,n,l)

line :: Integral i => Parser a -> Parser (i,i,a)
line f = (,,) <$> integral
              <*> integral
              <*> f
              <*  endOfLine

--------------------------------------------------------------------------------
-- Matrix Market
--------------------------------------------------------------------------------

matrix :: Parser Matrix
matrix = do
  (_,t,_) <- header
  (m,n,l) <- skipMany comment *> extent
  case t of
    Real    -> RealMatrix    (m,n) l `fmap` many1 (line floating)
    Complex -> ComplexMatrix (m,n) l `fmap` many1 (line ((:+) <$> floating <*> floating))
    Integer -> IntMatrix     (m,n) l `fmap` many1 (line integral)
    Pattern -> PatternMatrix (m,n) l `fmap` many1 ((,) <$> integral <*> integral)

readMatrix :: FilePath -> IO Matrix
readMatrix file = do
  chunks <- L.readFile file
  case L.parse matrix chunks of
    L.Fail _ _ msg      -> error $ file ++ ": " ++ msg
    L.Done _ mtx        -> return mtx


-- |
-- Module      : Main
-- Description :
-- Copyright   : (c) Jonatan H Sundqvist, 2015
-- License     : MIT
-- Maintainer  : Jonatan H Sundqvist
-- Stability   : experimental|stable
-- Portability : POSIX (not sure)
--

-- Created September 20 2015

-- TODO | -
--        -

-- SPEC | -
--        -

--------------------------------------------------------------------------------------------------------------------------------------------
-- GHC Pragmas
--------------------------------------------------------------------------------------------------------------------------------------------

--------------------------------------------------------------------------------------------------------------------------------------------
-- API
--------------------------------------------------------------------------------------------------------------------------------------------
module Main where

--------------------------------------------------------------------------------------------------------------------------------------------
-- We'll need these
--------------------------------------------------------------------------------------------------------------------------------------------
import Data.IORef
import Data.Complex
import qualified Data.Set as S
import qualified Data.Map as M

import Graphics.UI.Gtk
import qualified Graphics.Rendering.Cairo as Cairo

import Southpaw.Math.Constants
import qualified Southpaw.Picasso.Palette as Palette

import Occlusion.Types
import Occlusion.Vector
import qualified Occlusion.Window     as Window
import qualified Occlusion.Events     as Events
import qualified Occlusion.Behaviours as Behaviours

--------------------------------------------------------------------------------------------------------------------------------------------
-- Entry point
--------------------------------------------------------------------------------------------------------------------------------------------
-- |
main :: IO ()
main = do
  (window, canvas) <- Window.create size
  assets'          <- loadAssets
  stateref         <- newIORef $ AppState { _animation = AnimationData { _fps=30.0, _frame=0 },
                                            _input     = InputData { _keyboard=S.empty, _mouse=0:+0, _click=Nothing },
                                            _scene     = Scene { _player=Character { _position=0:+0, _velocity=0:+0, _health=100, _name="Democritus", _colour=Palette.peru,  _behaviour=Behaviours.player },
                                                                 _npcs=[Character { _position=0:+0, _velocity=20:+20, _health=58, _name="Agamemnon", _colour=Palette.orchid, _behaviour=Behaviours.npc patrol }],
                                                                 _obstacles=theobstacles },
                                            _gui       = GUI { _window=window, _canvas=canvas },
                                            _assets    = assets' }
  Events.attach window canvas stateref
  mainGUI
  where
    size = (720:+480) * 1.2
    patrol = map (+(200:+150)) [20:+20, 300:+32, 200:+10, 20:+20]

-- |
loadAssets :: IO Assets
loadAssets = do
  tree <- Cairo.imageSurfaceCreateFromPNG "assets/images/TreeAtDuskSmall.png"
  return $ Assets { _images=M.fromList [("tree", tree)] }

-- |
theobstacles :: [Polygon Double]
theobstacles = [[origin + let θ = α + n*2*π/ns in mkPolar mag θ | n <- [0..(ns-1)]] | (α, ns, mag, origin) <- [(0, 5, 78, 522:+150), (0, 8, 92, 192:+342)] ]


{-# LANGUAGE OverloadedLists #-}
{-# LANGUAGE Strict #-}
{-# LANGUAGE BangPatterns#-}
{-# LANGUAGE TemplateHaskell #-}
{-# LANGUAGE RankNTypes#-}
{-# LANGUAGE DeriveGeneric, DeriveAnyClass #-}
module Tasks.Swarm where

import           Numeric.LinearAlgebra (Vector, scale)
import qualified Data.Vector.Storable as VS
import           Control.Lens
import qualified Data.Vector as V
import qualified Data.Vector.Mutable as MV
import           Data.Vector.Mutable (STVector)
import           Control.Monad.ST
import           Data.List (nub)
import           Control.Monad ((<=<), foldM, when)
import Debug.Trace
import qualified Data.HashTable.ST.Cuckoo as HT
import qualified Data.HashMap as HM
import Data.Vector.Strategies
import GHC.Generics
import Control.Arrow

import Data.GameObject
import Physics.Mechanics
import Physics.Collision
import Data.PhysicsData

import Utils (norm,normalize)

type CollisionGrid s = HT.HashTable s (Int,Int,Int) [Int]

integrityWeight = 0.5
separationWeight = 2
gravityWeight = 0.1

cellSize = 4

vehSize = 2
minDist = 4

off = map ((.) (round.(/cellSize))) [(+ vehSize),(+ negate vehSize)]

addInPl :: forall s. Vector Float -> STVector s Float -> ST s ()
addInPl v mv = VS.zipWithM_ (\i a -> MV.modify mv (+a) i) [0,1,2] v

data BoidComputation = BoidComputation { _object :: GameObject
                                       , _sumSp :: Vector Float
                                       , _separ :: Vector Float
                                       , _nrNeighb :: {-# UNPACK #-} Int} deriving Show

makeLenses ''BoidComputation

ins :: Int -> Maybe [Int] -> (Maybe [Int], ())
ins ind (Just l) = (Just (ind:l),())
ins ind Nothing  = (Just [ind],())

markObject t ind=
  mapM_ (\v -> HT.mutate t (VS.unsafeIndex v 0,VS.unsafeIndex v 1,VS.unsafeIndex v 2) $ ins ind) .
  VS.mapM (\x -> ($ x) <$> off) .
  view location

createGrid :: V.Vector GameObject ->
               ST s (CollisionGrid s, V.Vector PrimitiveComputation)
createGrid v = do
  ht <- HT.newSized 6000
  vec <- V.imapM (process ht) v
  return (ht, vec)
  where
    process ht ind ob = do
      markObject ht ind ob
      return $ prepareCPrim ob

getCollisionData :: CollisionGrid s ->
                     V.Vector PrimitiveComputation ->
                     ST s (HM.Map (Int, Int) [Contact])
getCollisionData g prims =
  HT.foldM (\a p -> return $ collisionsInCell prims a (second nub p)) HM.empty g

getContacts :: V.Vector GameObject -> HM.Map (Int, Int) [Contact]
getContacts v =
  traceShowId $ runST $ createGrid v >>= uncurry getCollisionData

resolveContacts :: V.Vector GameObject -> V.Vector GameObject
resolveContacts =
  getContacts >>= HM.foldWithKey folder id
  where folder k = (.) . resolveCollision k

collisionsInCell :: V.Vector PrimitiveComputation ->
                      HM.Map (Int, Int) [Contact] ->
                      (a, [Int]) ->
                      HM.Map (Int, Int) [Contact]
collisionsInCell prims m (x, ix1:ix2:l) =
  let newM = insertCollisions prims m ix1 ix2
  in collisionsInCell prims newM (x, ix2:l)
collisionsInCell _ m _ = m

insertCollisions :: V.Vector PrimitiveComputation ->
                      HM.Map (Int, Int) [Contact] ->
                      Int -> Int ->
                      HM.Map (Int, Int) [Contact]
insertCollisions prims m ix1 ix2 =
  let key = (ix1, ix2)
      check = V.unsafeIndex prims
  in if HM.member key m
     then m
     else let val = view contacts $
                    detectCollision (check ix1) (check ix2)
          in if null val
             then m
             else HM.insert key val m

fillGrid :: V.Vector GameObject -> ST s (CollisionGrid s, STVector s BoidComputation)
fillGrid v = do
  ht <- HT.newSized 6000
  vec <- V.thaw =<< V.imapM (process ht) v
  return (ht,vec)
  where
    process g ind ob = do
      markObject g ind ob
      return $ initBoidComputation ob

computeForces :: (GameObject -> GameObject) -> V.Vector GameObject -> V.Vector GameObject
computeForces = computeForcesGrid fillGrid

incCount = over nrNeighb (+1)

initBoidComputation v = BoidComputation v [0,0,0] [0,0,0] 0

computeBoid (BoidComputation o v1 v2 n) =
  if n /= 0 then
    let Just p = o^.physics
        integ  = scale (integrityWeight / fromIntegral n) $ v1 - p^.speed
        sep    = scale separationWeight v2
    in putForceOnMid (integ+sep) o
  else o

computeForcesGrid :: (forall s. V.Vector GameObject -> ST s (CollisionGrid s, STVector s BoidComputation)) -> (GameObject -> GameObject) -> V.Vector GameObject -> V.Vector GameObject
computeForcesGrid s f v =
  V.map (f . computeBoid) $ runST $
             V.freeze <=< uncurry computeForcesMap <=< s $ v

computeForcesMap :: CollisionGrid s -> STVector s BoidComputation -> ST s (STVector s BoidComputation)
computeForcesMap m v = HT.foldM (\x -> computeForcesST x . nub . snd) v m

computeForcesST :: STVector s BoidComputation -> [Int] -> ST s (STVector s BoidComputation)
computeForcesST v (i:l@(_:_)) =
  foldM (computeForcesPair i) v l >>= flip computeForcesST l
computeForcesST v _ = pure v

computeForcesPair
  :: Int
  -> STVector s BoidComputation
  -> Int
  -> ST s (STVector s BoidComputation)
computeForcesPair i1 v i2 = do
  obj1 <- view object <$> MV.read v i1
  obj2 <- view object <$> MV.read v i2
  when (norm (obj1^.location - obj2^.location) < minDist) $ do
    let sepForce = sep obj1 obj2
    MV.modify v (over separ (+ sepForce)) i1
    MV.modify v (over sumSp (+ integ obj2)) i1
    MV.modify v (over separ (subtract sepForce)) i2
    MV.modify v (over sumSp (+ integ obj1)) i2
    MV.modify v incCount i1
    MV.modify v incCount i2
  pure v

sep :: GameObject -> GameObject -> Vector Float
sep obj1 obj2 = obj1^.location - obj2^.location
  -- let v = obj1^.location - obj2^.location
  --     m = norm v
  -- in scale (1/m/(m-1)) v

integ :: GameObject -> Vector Float
integ obj =
  case obj^.physics of
    Just x  -> x^.speed
    Nothing -> [0,0,0]

clamp :: Float -> Vector Float -> Vector Float
clamp m v =
  let n = norm v
  in if n > m then scale (m/n) v else v

addGravityForce :: Vector Float -> GameObject -> GameObject
addGravityForce gravCent obj =
  let direction = gravCent - obj^.location
--       dnorm = norm direction
      force = scale gravityWeight
              -- . normalize $
              direction
  in putForceOnMid force obj
    -- over (physics._Just.acc) (+force) obj


-- |
-- Module    : GRN.Uncertainty
-- Copyright : (c) 2011 Jason Knight
-- License   : BSD3
--
-- Maintainer  : jason@jasonknight.us
-- Stability   : experimental
-- Portability : portable
--
-- Print out a set of boolean networks that represent the uncertainty class.
--

module GRN.Uncertainty where

import GRN.StateTransition
import GRN.Render
import GRN.Sparse
import GRN.Types
import GRN.Utils
import GRN.EM

import qualified Data.Map as M
import Data.List
import Data.Char
import Data.Bits
import Control.Monad
import Text.Printf
import System.Console.ParseArgs

import qualified Data.Vector as V
import qualified Data.Vector.Unboxed as U
import qualified Data.Vector.Generic as G

import Control.Parallel.Strategies
import Data.Vector.Strategies

import Data.Graph.Inductive
import Data.Graph.Analysis.Algorithms.Common (componentsOf)

import Statistics.Distribution
import Statistics.Distribution.Normal

import Graphics.Gnuplot.Simple
import Text.PrettyPrint

-- | Collapses SSD onto space of non-input nodes
collapseSSD :: ParseData -> SSD -> SSD
collapseSSD pdata orig = snd $ foldl' splitAndSum (pdata,orig) inputSet where
  inputSet = inputGeneNames pdata

-- | Partition an SSD based on the value of a gene and sum the SSD (reducing
-- its size by 2 in the process).
splitAndSum :: (ParseData,SSD) -> Gene -> (ParseData,SSD)
splitAndSum (pd, sd) gene = (M.delete gene pd, U.zipWith (filtSum) sp1 sp2) where
  filtSum (_,x1) (_,x2) = x1+x2 -- Have to get rid of the indices
  (sp1,sp2) = U.partition evalGene (U.indexed sd) -- Wish there was a U.ipartition
  evalGene (ind,val) = testBit ind (numGenes-geneIndex-1)
  geneIndex = M.findIndex gene pd
  numGenes = M.size pd

-- | Find the names of the input genes
inputGeneNames :: ParseData -> [Gene]
inputGeneNames pdata = filter (inputGene pdata) (M.keys pdata)

-- | Use this to find input genes by the parse data.
inputGene :: ParseData -> Gene -> Bool
inputGene pdata g = null $ depends gi
  where
    Just gi = M.lookup g pdata

computeSSD :: IO ()
computeSSD = do
    let r file = liftM (U.fromList . head . read) $ readFile file :: IO (U.Vector Double)
    con1a <- r "moham_100.dat"
    con1b <- r "moham_111.dat"

    con2a <- r "moham_111.dat"
    con2b <- r "moham_101.dat"

    con3a <- r "moham_x1x.dat"
    con3b <- r "moham_1xx.dat"

    let ssd1 = U.sum . U.zipWith (\x y->(min x y)*0.5) con1a $ con1b
        ssd2 = U.sum . U.zipWith (\x y->(min x y)*0.5) con2a $ con2b
        ssd3 = U.sum . U.zipWith (\x y->(min x y)*0.5) con3a $ con3b

    print ssd1
    print ssd2
    print ssd3

uncertaintyPrint :: Args String -> ParseData -> IO ()
uncertaintyPrint args p = do
  let
    {-gen = gotArg args "generate"-}
    ks = buildKmaps p
    ms = measureMap p
    n1 = getRequiredArg args "n1" :: Int
    uvals = uncertainVals ks

    ng = M.size ks
    n = length uvals

    {-hypercorners = map (U.fromList.map fromIntegral.dec2bin n) [0..2^n-1]-}

    xs = clamped [0,0.5,1]
    center = replicate n 0.5
    clamped xs = map (\x-> if (x>1.0) then 1.0 else (if (x<0.0) then 0.0 else x)) xs
    grid = [[]] >>= foldr (>=>) return (replicate n (\x-> [x++[y] | y <- xs]))
    specialgrid = center : (grid \\ [center])
    hypergrid = map (U.fromList) specialgrid

    kmaps = map (fillKentries uvals ks) hypergrid
    ssds = parMap rdeepseq (collapseSSD p . simulateDOKUnif args . flip kmapToDOK 0) kmaps

  putStr (pprint ssds)
  {-plotList [] (U.toList . head $ ssds)-}

  {-print $ map U.length ssds-}
  {-putStrLn ""-}
  {-print $ length ssds-}
  {-print $ map (fst . normalizeSSD) ssds-}
  {-print $ "Optimum Theta: " ++ show optimumTheta-}

  return ()

class Pretty a where
  prettyPrint :: a -> String
  prettyPrint = renderStyle (Style LeftMode 80 0) . prettyShow
  prettyShow :: a -> Doc

pprint :: Pretty a => a -> String
pprint = prettyPrint

instance Pretty Int where
  prettyShow = text . show
instance Pretty Double where
  prettyShow = text . printf "%f"
instance Pretty a => Pretty [a] where
  prettyShow = brackets . hcat . punctuate comma . map prettyShow
instance (Pretty a, U.Unbox a) => Pretty (U.Vector a) where
  prettyShow = brackets . hcat . punctuate comma . map prettyShow . U.toList
instance (Pretty a, Pretty b) => Pretty (a,b) where
  prettyShow (a,b) = prettyShow a <> colon <> prettyShow b


module School.FileIO.MatrixSink
( MatrixSink
, matrixDoubleSink
, matrixIntSink
) where

import Conduit ((.|), ConduitM, mapC, sinkFileBS, yield)
import Data.Void (Void)
import School.FileIO.AppIO (AppIO)
import School.FileIO.FileType (FileType)
import School.FileIO.FileHeader (FileHeader(..), headerBuilder)
import School.Types.Encoding (matrixDoubleToBin, matrixIntToBin)
import Numeric.LinearAlgebra.Data (I, Matrix, R)

type MatrixSink a = ConduitM (Matrix a)
                             Void
                             AppIO
                             ()

matrixDoubleSink :: FileType
                 -> FileHeader
                 -> FilePath
                 -> MatrixSink R
matrixDoubleSink fType header path =
  byteBuilder .| sinkFileBS path where
    dType = dataType header
    byteBuilder = do
      yield $ headerBuilder fType header
      mapC (matrixDoubleToBin dType)

matrixIntSink :: FileType
              -> FileHeader
              -> FilePath
              -> MatrixSink I
matrixIntSink fType header path =
  byteBuilder .| sinkFileBS path where
    dType = dataType header
    byteBuilder = do
      yield $ headerBuilder fType header
      mapC (matrixIntToBin dType)


{-# LANGUAGE TemplateHaskell, OverloadedLists #-}
module Myelin.PyNN.Types.Node where

import Data.Aeson.TH
import Control.Lens
import Numeric.LinearAlgebra

-- | pyNN population node
data Node =
    IFCurrentAlpha {
        _tau_m :: Vector Float, -- ^ membrane time constant
        _tau_refrac :: Vector Float, -- ^ refractory time
        _v_thresh :: Vector Float, -- ^ threshhold potential
        _tau_syn_E :: Vector Float, -- ^ excitatory synaptic time constant
        _v_rest :: Vector Float, -- ^ resting potential
        _cm :: Vector Float, -- ^ membrane capactitance
        _v_reset :: Vector Float, -- ^ reset potential
        _tau_syn_I :: Vector Float, -- ^ inhibitory synaptic time constant
        _i_offset :: Vector Float -- ^ offset current
    }
    | IFCondAlpha {
        _v_rest :: Vector Float, -- ^ resting membrane potential
        _cm :: Vector Float, -- ^ capacity of the membrane
        _tau_m :: Vector Float, -- ^ membrane time constant
        _tau_refrac :: Vector Float, -- ^ duration of refractory period
        _tau_syn_E :: Vector Float, -- ^ rise time of the excitatory synaptic alpha function
        _tau_syn_I :: Vector Float, -- ^ rise time of the inhibitory synaptic alpha function
        _e_rev_E :: Vector Float, -- ^ reversal potential for excitatory input
        _e_rev_I :: Vector Float, -- ^ reversal potential for inhibitory input
        _v_thresh :: Vector Float, -- ^ spike threshold
        _v_reset :: Vector Float, -- ^ reset potential after a spike
        _i_offset :: Vector Float -- ^ offset current
    }
    | IFSpikey {
        _e_rev_I :: Vector Float, -- ^ excitatory reversal current
        _g_leak :: Vector Float, -- ^ leak conductance
        _tau_refrac :: Vector Float, -- ^ refractory time
        _v_reset :: Vector Float, -- ^ reset potential
        _v_rest :: Vector Float, -- ^ resting potential
        _v_thresh :: Vector Float -- ^ threshhold potential
    }
    | IFCurrExp {
        _cm :: Vector Float, -- ^ membrane capacitance
        _tau_m :: Vector Float, -- ^ membrane time constant
        _tau_syn_E :: Vector Float, -- ^ excitatory synaptic time constant
        _tau_syn_I :: Vector Float, -- ^ inhibitory synaptic time constant
        _tau_refrac :: Vector Float, -- ^ refractory time
        _v_thresh :: Vector Float, -- ^ threshhold voltage
        _v_rest :: Vector Float, -- ^ resting potential
        _v_reset :: Vector Float, -- ^ reset potential
        _i_offset :: Vector Float -- ^ offset current
    }
    | IFCondExp {
        _v_rest :: Vector Float, -- ^ resting potential
        _cm :: Vector Float, -- ^ membrane capacitance
        _tau_m :: Vector Float, -- ^ membrane time constant
        _tau_refrac :: Vector Float, -- ^ refractory time
        _tau_syn_E :: Vector Float, -- ^ excitatory synaptic time constant
        _tau_syn_I :: Vector Float, -- ^ inhibitory synaptic time constant
        _e_rev_E :: Vector Float, -- ^ excitatory reversal potential
        _e_rev_I :: Vector Float, -- ^ inhibitory reversal potential
        _v_thresh :: Vector Float, -- ^ spike initiation threshold
        _v_reset :: Vector Float, -- ^ reset value for membrane potential after a spike
        _i_offset :: Vector Float -- ^ offset current
    }
    | Izhikevich {
        _a :: Vector Float,
        _i_offset :: Vector Float, -- ^ offset current
        _c :: Vector Float,
        _d :: Vector Float,
        _b :: Vector Float
    }
    | EIFCondExp {
        _v_reset :: Vector Float, -- ^ reset potential
        _i_offset :: Vector Float, -- ^ offset current
        _tau_w :: Vector Float,
        _tau_syn_I :: Vector Float, -- ^ inhibitory synaptic time constant
        _e_rev_E :: Vector Float,
        _v_rest :: Vector Float, -- ^ resting potential
        _cm :: Vector Float, -- ^ capactivie memory
        _tau_syn_E :: Vector Float, -- ^ excitatory synaptic time constant
        _tau_m :: Vector Float, -- ^ membrane timeconstant
        _a :: Vector Float,
        _delta_T :: Vector Float,
        _v_thresh :: Vector Float, -- ^ threshhold potential
        _b :: Vector Float,
        _v_spike :: Vector Float,
        _e_rev_I :: Vector Float,
        _tau_refrac :: Vector Float -- ^ refractory time constant
    }
    | EIFCondAlpha {
        _cm :: Vector Float, -- ^ capacity of the membrane
        _tau_refrac :: Vector Float, -- ^ duration of the refractory period
        _v_spike :: Vector Float, -- ^ spike detection threshold
        _v_reset :: Vector Float, -- ^ reset value for membrane potential after a spike
        _v_rest :: Vector Float, -- ^ resting membrane potential (Leak reversal potential)
        _tau_m :: Vector Float, -- ^ membrane time constant
        _i_offset :: Vector Float, -- ^ offset current
        _a :: Vector Float, -- ^ subthreshold adaptation conductance
        _b :: Vector Float, -- ^ spike-triggered adaptation
        _delta_T :: Vector Float, -- ^ slope factor
        _tau_w :: Vector Float, -- ^ adaptation time constant
        _v_thresh :: Vector Float, -- ^ spike initiation threshold
        _e_rev_E :: Vector Float, -- ^ excitatory reversal potential
        _tau_syn_E :: Vector Float, -- ^ rise time of excitatory synaptic conductance (alpha function)
        _e_rev_I :: Vector Float, -- ^ inhibitory reversal potential
        _tau_syn_I :: Vector Float -- ^ rise time of the inhibitory synaptic conductance (alpha function)
    }
    | HHCondExp {
        _cm :: Vector Float, -- ^ capacity of the membrane
        _e_rev_E :: Vector Float, -- ^ excitatory reversal potential
        _e_rev_I :: Vector Float, -- ^ inhibitory reversal potential
        _e_rev_K :: Vector Float,
        _e_rev_Na :: Vector Float,
        _e_rev_leak :: Vector Float,
        _g_leak :: Vector Float,
        _gbar_K :: Vector Float,
        _gbar_Na :: Vector Float,
        _i_offset :: Vector Float,
        _tau_syn_E :: Vector Float,
        _tau_syn_I :: Vector Float,
        _v_offset :: Vector Float
    }
    | SpikeSourcePoisson {
        _start :: Float,
        _duration :: Float,
        _rate :: Float
    }
    | SpikeSourceArray {
        _spike_times :: Vector Float
    }
    | SpikeSourceInhGamma {
        _a :: Vector Float,
        _tbins :: Vector Float,
        _duration :: Float,
        _b :: Vector Float,
        _start :: Float
    }
    deriving (Eq, Show, Ord, Read)

-- |  Defaults taken from http://neuralensemble.org/docs/PyNN/standardmodels.html

-- | default parameters for the integrate and fire conductance based neuron with exponential term
if_cond_exp :: Node
if_cond_exp = IFCondExp {
    _v_rest = [-65.0],
    _cm = [1.0],
    _tau_m = [20.0],
    _tau_refrac = [0.0],
    _tau_syn_E = [5.0],
    _tau_syn_I = [5.0],
    _e_rev_E = [0.0],
    _e_rev_I = [-70.0],
    _v_thresh = [-50.0],
    _v_reset = [-65.0],
    _i_offset = [0.0]
}

if_current_alpha :: Node
if_current_alpha = IFCurrentAlpha {
    _tau_m = [20.0],
    _tau_refrac = [0.1],
    _v_thresh = [-50.0],
    _tau_syn_E = [ 0.5],
    _tau_syn_I = [0.5],
    _v_rest = [-65.0],
    _cm = [1.0],
    _v_reset = [-65.0],
    _i_offset = [0.0]
}

if_cond_alpha :: Node
if_cond_alpha = IFCondAlpha {
    _v_rest = [-65.0],
    _cm = [1.0],
    _tau_m = [20.0],
    _tau_refrac = [0.0],
    _tau_syn_E = [5.0],
    _tau_syn_I = [5.0],
    _e_rev_E = [0.0],
    _e_rev_I = [-70.0],
    _v_thresh = [-50.0],
    _v_reset = [-65.0],
    _i_offset = [0.0]
}

if_spikey :: Node
if_spikey = IFSpikey {
    _e_rev_I = [-80.0],
    _g_leak = [20.0],
    _tau_refrac = [1.0],
    _v_reset = [-80.0],
    _v_rest = [-75.0],
    _v_thresh = [-55.0]
}

if_current_exponential :: Node
if_current_exponential = IFCurrExp {
    _cm = [1.0],
    _tau_m = [20.0],
    _tau_syn_E = [5.0],
    _tau_syn_I = [5.0],
    _tau_refrac = [0.1],
    _v_thresh = [-50.0],
    _v_rest = [-65.0],
    _v_reset = [-65.0],
    _i_offset = [0.0]
}

izhikevich :: Node
izhikevich = Izhikevich {
    _a = [0.02],
    _i_offset = [0.0],
    _c = [-65.0],
    _d = [2.0],
    _b = [0.2]
}

eif_cond_exp :: Node
eif_cond_exp = EIFCondExp {
    _v_reset = [-70.6],
    _i_offset = [0.0],
    _tau_w = [144.0],
    _tau_syn_I = [5.0],
    _e_rev_E = [0.0],
    _v_rest = [-70.6],
    _cm = [0.281],
    _tau_syn_E = [5.0],
    _tau_m = [9.3667],
    _a = [4.0],
    _delta_T = [2.0],
    _v_thresh = [-50.4],
    _b = [0.0805],
    _v_spike = [-40.0],
    _e_rev_I = [-80.0],
    _tau_refrac = [0.1]
}

eif_cond_alpha :: Node
eif_cond_alpha = EIFCondAlpha {
    _v_reset = [-70.6],
    _i_offset = [0.0],
    _tau_w = [144.0],
    _tau_syn_I = [5.0],
    _e_rev_E = [0.0],
    _v_rest = [-70.6],
    _cm = [0.281],
    _tau_syn_E = [5.0],
    _tau_m = [9.3667],
    _a = [4.0],
    _delta_T = [2.0],
    _v_thresh = [-50.4],
    _b = [0.0805],
    _v_spike = [-40.0],
    _e_rev_I = [-80.0],
    _tau_refrac = [0.1]
}

spike_source_poisson = SpikeSourcePoisson {
    _start = 0.0,
    _duration = 100.0,
    _rate = 0.1
}

spike_source_array = SpikeSourceArray {
    _spike_times = []
}

deriveJSON defaultOptions ''Node
makeLenses ''Node
makePrisms ''Node


{-# LANGUAGE FlexibleContexts,
             UndecidableInstances #-}

module Network.Network.FeedForwardNetwork
( FeedForwardNetwork(..)

, emptyFeedForwardNetwork
, isEmptyFeedForwardNetwork
, addFeedForwardNetworks

, loadFeedForwardNetwork
, saveFeedForwardNetwork

, apply
) where

import Network.Neuron
import Network.Layer
import Network.Network

import System.Random
import Numeric.LinearAlgebra
import qualified Data.ByteString.Lazy as B
import System.IO
import Data.Monoid (Monoid(..))
import Data.Binary (encode, decode, Binary(..))

-- | Networks are constructed front to back. Start by adding an input layer,
--   then each hidden layer, and finally an output layer.
data FeedForwardNetwork = FeedForwardNetwork { layers :: [Layer] } deriving Show

-- | We gain the ability to combine two networks of the same proportions
--   by abstracting a network as a monoid. This is useful in backpropagation
--   for batch training
instance Monoid (FeedForwardNetwork) where
  mempty = emptyFeedForwardNetwork
  mappend = addFeedForwardNetworks

-- | A tuple of (input, expected output)
type TrainingData = (Vector Double, Vector Double)

-- | Our Unit, an empty network with no layers
emptyFeedForwardNetwork :: FeedForwardNetwork
emptyFeedForwardNetwork = FeedForwardNetwork []

-- | A boolean to check if the network is the unit network or not
isEmptyFeedForwardNetwork :: FeedForwardNetwork -> Bool
isEmptyFeedForwardNetwork n = length (layers n) == 0

-- | A function to combine two networks
addFeedForwardNetworks :: FeedForwardNetwork -> FeedForwardNetwork -> FeedForwardNetwork
addFeedForwardNetworks n1 n2 = if isEmptyFeedForwardNetwork n1 then n2 else
  if isEmptyFeedForwardNetwork n2 then n1 else
    FeedForwardNetwork $ zipWith combineLayers (layers n1) (layers n2)
  where combineLayers l1 l2 =
          Layer ((weightMatrix l1) + (weightMatrix l2))
          ((biasVector l1) + (biasVector l2)) (neuron l1)

instance Network (FeedForwardNetwork) where
  -- | Predict folds over each layer of the network using the input vector as the
  --   first value of the accumulator. It operates on whatever network you pass in.
  predict input network = foldl apply input (layers network)

  -- | The createNetwork function takes in a random transform used for weight
  --   initialization, a source of entropy, and a list of layer definitions,
  --   and returns a network with the weights initialized per the random transform.
  createNetwork t g [] = emptyFeedForwardNetwork
  createNetwork t g (layerDef : []) = emptyFeedForwardNetwork
  createNetwork t g (layerDef : layerDef' : otherLayerDefs) =
    FeedForwardNetwork (layer : layers restOfNetwork)
    where layer = createLayer t g' layerDef layerDef'
          restOfNetwork = createNetwork t g'' (layerDef' : otherLayerDefs)
          (g', g'') = split g

-- | A function used in the fold in predict that applies the activation
--   function and pushes the input through a layer of the network.
apply :: Vector Double -> Layer -> Vector Double
apply vector layer = mapVector sigma (weights <> vector + bias)
  where sigma = activation (neuron layer)
        weights = weightMatrix layer
        bias = biasVector layer

-- | Given a filename and a network, we want to save the weights and biases
--   of the network to the file for later use.
saveFeedForwardNetwork :: Binary (Layer) => FilePath -> FeedForwardNetwork -> IO ()
saveFeedForwardNetwork file n = B.writeFile file (encode (layers n))

-- | Given a filename, and a list of layer definitions, we want to reexpand
--   the data back into a network.
loadFeedForwardNetwork :: Binary (Layer) => FilePath -> [LayerDefinition] -> IO (FeedForwardNetwork)
loadFeedForwardNetwork file defs = B.readFile file >>= \sls ->
  return $ FeedForwardNetwork (map showableToLayer (zip (decode sls) defs))


module Main (main) where

import           Debug.Trace

import           Criterion
import           Criterion.IO
import           Criterion.Types
import           Data.Foldable
import           Data.List
import           Data.List.NonEmpty (NonEmpty(..))
import qualified Data.List.NonEmpty as NonEmpty
import           Data.List.Split
import           Data.Map.Strict (Map)
import qualified Data.Map.Strict as Map
import           Graphics.Rendering.Chart.Backend.Cairo
import           Graphics.Rendering.Chart.Easy
import           Statistics.Types
import           System.Environment
import           System.IO
import           Text.Read

main :: IO ()
main = do
  [file] <- getArgs
  res <- readJSONReports file
  case res of
    Left err -> putStrLn err
    Right (_, _, reports) ->
      case traverse process reports of
        Nothing -> putStrLn "Failed to process result"
        Just rs -> do
          let plotData = aggregate rs
          for_ plotData $ \p -> do
            let fileName = plotName p <> ".svg"
            putStrLn ("Generating " <> fileName)
            toFile (def & fo_format .~ SVG) fileName (generatePlot p)

-- We assume that

data Result = Result
  { operation :: String
  , implementation :: String
  , size :: Double
  , iters :: Estimate ConfInt Double
  } deriving Show

process :: Report -> Maybe Result
process r = Result op impl <$> readMaybe sizeStr <*> getIters r
  where
    (sizeStr:impl:rest) = reverse (splitOn "/" (reportName r))
    op = intercalate "/" (reverse rest)

getIters :: Report -> Maybe (Estimate ConfInt Double)
getIters r =
  Map.lookup "iters" . regCoeffs =<<
  find (\r -> regResponder r == "time") (anRegress (reportAnalysis r))

data PlotData = PlotData
  { plotName :: String
  , lines :: [LineData]
  } deriving Show

data LineData = LineData
  { lineName :: String
  , linePoints :: [(Double, Estimate ConfInt Double)]
  } deriving Show

aggregate :: [Result] -> [PlotData]
aggregate rs =
  (map (uncurry PlotData) .
   Map.toList . Map.map (map (uncurry toLineData) . Map.toList))
    byImpl
  where
    byOp :: Map String (NonEmpty Result)
    byOp = aggregateBy operation rs
    byImpl :: Map String (Map String (NonEmpty Result))
    byImpl = Map.map (aggregateBy implementation . toList) byOp
    toLineData :: String -> NonEmpty Result -> LineData
    toLineData k rs = LineData k (map (\r -> (size r, iters r)) (toList rs))

aggregateBy :: Ord k => (a -> k) -> [a] -> Map k (NonEmpty a)
aggregateBy f = Map.fromListWith (<>) . map (\x -> (f x, x :| []))

generatePlot :: PlotData -> EC (Layout Double Double) ()
generatePlot (PlotData n ls) = do
  layout_title .= n
  for_ ls $ \(LineData n' ps) -> do
    let ps' = map (\(x, e) -> (LogValue x, estPoint e)) ps
    plot (line n' [ps'])
    plot (points n' ps')


{-# LANGUAGE CPP                   #-}
{-# LANGUAGE DataKinds             #-}
{-# LANGUAGE FlexibleContexts      #-}
{-# LANGUAGE GADTs                 #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE PolyKinds             #-}
{-# LANGUAGE ScopedTypeVariables   #-}
{-# LANGUAGE TypeFamilies          #-}
{-# LANGUAGE TypeOperators         #-}
{-# LANGUAGE UndecidableInstances  #-}

module Grenade.Recurrent.Layers.BasicRecurrent (
    BasicRecurrent (..)
  ) where

import           Data.Proxy
import           Data.Singletons.TypeLits     hiding (natVal)

#if MIN_VERSION_base(4,9,0)
import           Data.Kind                    (Type)
#endif

import           Numeric.LinearAlgebra.Static

import           GHC.TypeLits

import           Grenade.Core
import           Grenade.Recurrent.Core
import           Grenade.Utils.LinearAlgebra

data BasicRecurrent :: Nat -- Input layer size
                    -> Nat -- Output layer size
                    -> Type where
  BasicRecurrent :: ( KnownNat input
                    , KnownNat output
                    , KnownNat matrixCols
                    , matrixCols ~ (input + output))
                 => !(R output)   -- Bias neuron weights
                 -> !(R output)   -- Bias neuron momentum
                 -> !(L output matrixCols) -- Activation
                 -> !(L output matrixCols) -- Momentum
                 -> BasicRecurrent input output

data BasicRecurrent' :: Nat -- Input layer size
                     -> Nat -- Output layer size
                     -> Type where
  BasicRecurrent' :: ( KnownNat input
                     , KnownNat output
                     , KnownNat matrixCols
                     , matrixCols ~ (input + output))
                  => !(R output)   -- Bias neuron gradients
                  -> !(L output matrixCols)
                  -> BasicRecurrent' input output

instance Show (BasicRecurrent i o) where
  show BasicRecurrent {} = "BasicRecurrent"

instance FoldableGradient (BasicRecurrent' input output) where
  mapGradient f (BasicRecurrent' bias act) = BasicRecurrent' (dvmap f bias) (dmmap f act)
  squaredSums (BasicRecurrent' bias act) = [sumV . squareV $ bias, sumM . squareM $ act]

instance (KnownNat i, KnownNat o, KnownNat (i + o)) => UpdateLayer (BasicRecurrent i o) where
  type Gradient (BasicRecurrent i o) = (BasicRecurrent' i o)

  runUpdate (OptSGD lRate lMomentum lRegulariser) (BasicRecurrent oldBias oldBiasMomentum oldActivations oldMomentum) (BasicRecurrent' biasGradient activationGradient) =
    let newBiasMomentum = konst lMomentum * oldBiasMomentum - konst lRate * biasGradient
        newBias         = oldBias + newBiasMomentum
        newMomentum     = konst lMomentum * oldMomentum - konst lRate * activationGradient
        regulariser     = konst (lRegulariser * lRate) * oldActivations
        newActivations  = oldActivations + newMomentum - regulariser
    in BasicRecurrent newBias newBiasMomentum newActivations newMomentum
  runUpdate _ l d = runUpdate defOptimizer l d
  reduceGradient = error "Attempt to reduce gradient of batch in BasicRecurrent instance"

instance (KnownNat i, KnownNat o, KnownNat x, KnownNat (x*o), x ~ (i+o)) => RandomLayer (BasicRecurrent i o) where
  createRandomWith m gen = do
    wB <- getRandomVector i o m gen
    wN <- getRandomMatrix i o m gen
    let bm = konst 0
        mm = konst 0
    return $ BasicRecurrent wB bm wN mm
      where i = natVal (Proxy :: Proxy i)
            o = natVal (Proxy :: Proxy o)

instance (KnownNat i, KnownNat o, KnownNat (i + o), i <= (i + o), o ~ ((i + o) - i)) => RecurrentUpdateLayer (BasicRecurrent i o) where
  type RecurrentShape (BasicRecurrent i o) = S ('D1 o)

instance (KnownNat i, KnownNat o, KnownNat (i + o), i <= (i + o), o ~ ((i + o) - i)) => RecurrentLayer (BasicRecurrent i o) ('D1 i) ('D1 o) where

  type RecTape (BasicRecurrent i o) ('D1 i) ('D1 o) = (S ('D1 o), S ('D1 i))
  -- Do a matrix vector multiplication and return the result.
  runRecurrentForwards (BasicRecurrent wB _ wN _) (S1D lastOutput) (S1D thisInput) =
    let thisOutput = S1D $ wB + wN #> (thisInput # lastOutput)
    in ((S1D lastOutput, S1D thisInput), thisOutput, thisOutput)

  -- Run a backpropogation step for a full connected layer.
  runRecurrentBackwards (BasicRecurrent _ _ wN _) (S1D lastOutput, S1D thisInput) (S1D dRec) (S1D dEdy) =
    let biasGradient        = (dRec + dEdy)
        layerGrad           = (dRec + dEdy) `outer` (thisInput # lastOutput)
        -- calcluate derivatives for next step
        (backGrad, recGrad) = split $ tr wN #> (dRec + dEdy)
    in  (BasicRecurrent' biasGradient layerGrad, S1D recGrad, S1D backGrad)


-- |
-- Module      : RenderUtils
-- Description : Cairo rendering utilities
-- Copyright   : (c) Jonatan H Sundqvist, 2015
-- License     : MIT
-- Maintainer  : Jonatan H Sundqvist
-- Stability   : experimental|stable
-- Portability : POSIX (not sure)
--

-- Created July 12 2015

-- TODO | -
--        -

-- SPEC | -
--        -

module Southpaw.Picasso.RenderUtils where

--------------------------------------------------------------------------------------------------------------------------------------------
-- We'll need these
--------------------------------------------------------------------------------------------------------------------------------------------
import Data.Complex
import Control.Monad.IO.Class (MonadIO)
import Control.Monad          (liftM, liftM2)

import qualified Graphics.Rendering.Cairo as Cairo

import qualified Southpaw.Picasso.Palette as Palette
import           Southpaw.Cartesian.Plane.Types
import           Southpaw.Cartesian.Plane.Utilities

--------------------------------------------------------------------------------------------------------------------------------------------
-- Functions
--------------------------------------------------------------------------------------------------------------------------------------------
-- General rendering utilities -------------------------------------------------------------------------------------------------------------

-- | Choose a colour
choose :: Palette.Colour Double -> Cairo.Render ()
choose (r, g, b, a) = Cairo.setSourceRGBA r g b a

-- Vector utilities ------------------------------------------------------------------------------------------------------------------------

-- TODO: Move to another module (?)
-- |
vectorise :: (Double -> Double -> a) -> Complex Double -> a
vectorise f (x:+y) = f x y

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

-- |
imageSurfaceSize :: MonadIO m => Cairo.Surface -> m (Complex Double)
imageSurfaceSize im = liftM (dotmap fromIntegral) $ liftM2 (:+) (Cairo.imageSurfaceGetWidth im) (Cairo.imageSurfaceGetHeight im)

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

-- |
textsize :: String -> Cairo.Render (Complex Double)
textsize text = do
  extents <- Cairo.textExtents text
  return $ Cairo.textExtentsWidth extents :+ Cairo.textExtentsHeight extents

-- |
textbounds :: String -> Complex Double -> Cairo.Render (BoundingBox Double)
textbounds text centre = do
  size <- textsize text
  return $ BoundingBox { _centre=centre, _size=size }

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

-- |
-- TODO: Unsafe, use Maybe (?)
closePath :: [Complex Double] -> [Complex Double]
closepath [p]  = [p]
closePath path = path ++ take 1 path

-- Control structures ----------------------------------------------------------------------------------------------------------------------

-- |
grid :: (Integral n, Enum n) => n -> n -> (n -> n -> a) -> [a]
grid cols rows f = [ f cl rw | cl <- [0..(cols-1)], rw <- [0..(rows-1)] ] -- Tiles

-- |
-- TODO: Rename (eg. something pertaining to 2D loops)
-- TODO: Underscores, grrr
gridM_ :: (Integral n, Enum n, Monad m) => n -> n -> (n -> n -> m a) -> m ()
gridM_ cols rows f = sequence_ $ grid cols rows f


{-
 - SPDX-FileCopyrightText: 2020-2021 UnrelatedString <https://github.com/UnrelatedString> and other Perhaps contributors
 -
 - SPDX-License-Identifier: BSD-3-Clause
 -}

{-# LANGUAGE ExistentialQuantification #-}
{-# LANGUAGE FlexibleInstances #-}
-- {-# LANGUAGE StandaloneDeriving #-}
{-# LANGUAGE QuantifiedConstraints #-}

module Perhaps.Data
    ( Token (LiteralT, PrimitiveT, OperatorT),
      Value (Number, Char, List),
      Expression (PrimitiveE, LiteralE, DerivedE),
      Function (PrimitiveF, LiteralF, DerivedF, TrainF),
      Primitive,
      Operator (Operator),
      Derived (Derived),
      derive,
      Number,
      integerMaybe
    ) where

import Data.Ratio (Rational, numerator, denominator)
--import Data.Complex (Complex, realPart, imagPart)

data Token = LiteralT Value
           | PrimitiveT Primitive
           | OperatorT Operator deriving (Show)

data Expression = PrimitiveE Primitive
               | LiteralE Value
               | DerivedE (Derived (Maybe Expression))
               | LeftBindE (Maybe Expression) (Maybe Expression)
               | RightBindE (Maybe Expression) (Maybe Expression) deriving (Show)

data Function = PrimitiveF Primitive
              | LiteralF Value
              | DerivedF (Derived Function)
              | TrainF [Function]
              | LeftBindF Function Function
              | RightBindF Function Function deriving (Show)

-- Syntactic adicity, not semantic adicity
data Adicity = Niladic | Monadic | Dyadic | Variadic (Adicity -> Adicity)

class Adic a where
    adicity :: a -> Adicity

-- TODO: flagged lists -- just using lists for convenience at moment
-- replace before even implementing choice
data Value = Number Number | Char Char | List [Value] deriving (Show)

-- TODO: replace with symbolic math lmao, don't really want to approximate radicals/pi
type Number = Rational

integerMaybe :: Number -> Maybe Integer
integerMaybe x
    | denominator x == 1 = Just $ numerator x
    | otherwise = Nothing

-- placeholder
type Primitive = String

data Derived a = forall f. (Traversable f, forall b. Show b => Show (f b)) =>
                     Derived String (f Function -> String) (f a)

instance Functor Derived where
    fmap f (Derived name op args) = Derived name op $ fmap f args
instance Foldable Derived where
    foldr f x (Derived _ _ args) = foldr f x args
instance Traversable Derived where
    sequenceA (Derived name op args) = Derived name op <$> sequenceA args

instance Show a => Show (Derived a) where
    show (Derived name _ args) = name ++ " " ++ show args

derive :: Derived Function -> String -- bad name
derive (Derived _ op args) = op args

data Operator = Operator String Bool ([Maybe Expression] -> (Derived (Maybe Expression), [Maybe Expression]))

instance Show Operator where
    show (Operator name _ _) = name