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

module BlueRipple.Model.Preference where

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

import qualified Text.Pandoc.Error             as PA

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

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

import qualified Numeric.LinearAlgebra         as LA

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

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

import qualified Frames.Visualization.VegaLite.Correlation
                                               as FV

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

    totalVoteDiagnostics @b resultsWithPopulationsAndGGAdjFrame opposedFrame

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

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

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

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

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

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

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

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

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

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

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


{-# LANGUAGE CPP #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE DeriveTraversable #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE UndecidableInstances #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE RankNTypes #-}
#if defined(__GLASGOW_HASKELL__) && __GLASGOW_HASKELL__ >= 702
{-# LANGUAGE Trustworthy #-}
{-# LANGUAGE DeriveGeneric #-}
#endif
{-# LANGUAGE DeriveDataTypeable #-}
#if defined(__GLASGOW_HASKELL__) && __GLASGOW_HASKELL__ < 708
{-# LANGUAGE StandaloneDeriving #-}
{-# LANGUAGE ScopedTypeVariables #-}
#endif
-----------------------------------------------------------------------------
-- |
-- License     :  BSD-style (see the file LICENSE)
-- Maintainer  :  Edward Kmett <ekmett@gmail.com>
-- Stability   :  provisional
-- Portability :  portable
--
-- Operations on affine spaces.
-----------------------------------------------------------------------------
module Linear.Affine where

import Control.Applicative
import Control.Lens
import Data.Complex (Complex)
import Data.Data
import Data.Distributive
import Data.Foldable as Foldable
import Data.Functor.Bind
import Data.Functor.Rep as Rep
import Data.HashMap.Lazy (HashMap)
import Data.Hashable
import Data.IntMap (IntMap)
import Data.Ix
import Data.Map (Map)
import Data.Vector (Vector)
import Foreign.Storable
#if defined(__GLASGOW_HASKELL__) && __GLASGOW_HASKELL__ >= 702
import GHC.Generics (Generic)
#endif
#if defined(__GLASGOW_HASKELL__) && __GLASGOW_HASKELL__ >= 706
import GHC.Generics (Generic1)
#endif
import Linear.Epsilon
import Linear.Metric
import Linear.Plucker
import Linear.Quaternion
import Linear.V
import Linear.V0
import Linear.V1
import Linear.V2
import Linear.V3
import Linear.V4
import Linear.Vector

#ifdef HLINT
{-# ANN module "HLint: ignore Unused LANGUAGE pragma" #-}
#endif

-- | An affine space is roughly a vector space in which we have
-- forgotten or at least pretend to have forgotten the origin.
--
-- > a .+^ (b .-. a)  =  b@
-- > (a .+^ u) .+^ v  =  a .+^ (u ^+^ v)@
-- > (a .-. b) ^+^ v  =  (a .+^ v) .-. q@
class Additive (Diff p) => Affine p where
  type Diff p :: * -> *

  infixl 6 .-.
  -- | Get the difference between two points as a vector offset.
  (.-.) :: Num a => p a -> p a -> Diff p a

  infixl 6 .+^
  -- | Add a vector offset to a point.
  (.+^) :: Num a => p a -> Diff p a -> p a

  infixl 6 .-^
  -- | Subtract a vector offset from a point.
  (.-^) :: Num a => p a -> Diff p a -> p a
  p .-^ v = p .+^ negated v
  {-# INLINE (.-^) #-}

-- | Compute the quadrance of the difference (the square of the distance)
qdA :: (Affine p, Foldable (Diff p), Num a) => p a -> p a -> a
qdA a b = Foldable.sum (fmap (join (*)) (a .-. b))
{-# INLINE qdA #-}

-- | Distance between two points in an affine space
distanceA :: (Floating a, Foldable (Diff p), Affine p) => p a -> p a -> a
distanceA a b = sqrt (qdA a b)
{-# INLINE distanceA #-}

#define ADDITIVEC(CTX,T) instance CTX => Affine T where type Diff T = T ; \
  (.-.) = (^-^) ; {-# INLINE (.-.) #-} ; (.+^) = (^+^) ; {-# INLINE (.+^) #-} ; \
  (.-^) = (^-^) ; {-# INLINE (.-^) #-}
#define ADDITIVE(T) ADDITIVEC((), T)

ADDITIVE([])
ADDITIVE(Complex)
ADDITIVE(ZipList)
ADDITIVE(Maybe)
ADDITIVE(IntMap)
ADDITIVE(Identity)
ADDITIVE(Vector)
ADDITIVE(V0)
ADDITIVE(V1)
ADDITIVE(V2)
ADDITIVE(V3)
ADDITIVE(V4)
ADDITIVE(Plucker)
ADDITIVE(Quaternion)
ADDITIVE(((->) b))
ADDITIVEC(Ord k, (Map k))
ADDITIVEC((Eq k, Hashable k), (HashMap k))
ADDITIVEC(Dim n, (V n))

-- | A handy wrapper to help distinguish points from vectors at the
-- type level
newtype Point f a = P (f a)
  deriving ( Eq, Ord, Show, Read, Monad, Functor, Applicative, Foldable
           , Traversable, Apply, Additive, Metric
           , Fractional , Num, Ix, Storable, Epsilon
           , Hashable
#if defined(__GLASGOW_HASKELL__) && __GLASGOW_HASKELL__ >= 702
           , Generic
#endif
#if defined(__GLASGOW_HASKELL__) && __GLASGOW_HASKELL__ >= 706
           , Generic1
#endif
#if defined(__GLASGOW_HASKELL__) && __GLASGOW_HASKELL__ >= 708
           , Typeable, Data
#endif
           )

#if defined(__GLASGOW_HASKELL__) && __GLASGOW_HASKELL__ < 708
instance forall f. Typeable1 f => Typeable1 (Point f) where
  typeOf1 _ = mkTyConApp (mkTyCon3 "linear" "Linear.Affine" "Point") [] `mkAppTy`
              typeOf1 (undefined :: f a)

deriving instance (Data (f a), Typeable1 f, Typeable a) => Data (Point f a)
#endif

lensP :: Lens' (Point g a) (g a)
lensP afb (P a) = P <$> afb a
{-# INLINE lensP #-}

_Point :: Iso' (Point f a) (f a)
_Point = iso (\(P a) -> a) P
{-# INLINE _Point #-}

instance (t ~ Point g b) => Rewrapped (Point f a) t
instance Wrapped (Point f a) where
  type Unwrapped (Point f a) = f a
  _Wrapped' = _Point
  {-# INLINE _Wrapped' #-}

instance Bind f => Bind (Point f) where
  join (P m) = P $ join $ fmap (\(P m')->m') m

instance Distributive f => Distributive (Point f) where
  distribute = P . collect (\(P p) -> p)

instance Representable f => Representable (Point f) where
  type Rep (Point f) = Rep f
  tabulate f = P (tabulate f)
  {-# INLINE tabulate #-}
  index (P xs) = Rep.index xs
  {-# INLINE index #-}

type instance Index (Point f a) = Index (f a)
type instance IxValue (Point f a) = IxValue (f a)

instance Ixed (f a) => Ixed (Point f a) where
  ix l = lensP . ix l
  {-# INLINE ix #-}

instance Traversable f => Each (Point f a) (Point f b) a b where
  each = traverse
  {-# INLINE each #-}

instance R1 f => R1 (Point f) where
  _x = lensP . _x
  {-# INLINE _x #-}

instance R2 f => R2 (Point f) where
  _y = lensP . _y
  {-# INLINE _y #-}
  _xy = lensP . _xy
  {-# INLINE _xy #-}

instance R3 f => R3 (Point f) where
  _z = lensP . _z
  {-# INLINE _z #-}
  _xyz = lensP . _xyz
  {-# INLINE _xyz #-}

instance R4 f => R4 (Point f) where
  _w = lensP . _w
  {-# INLINE _w #-}
  _xyzw = lensP . _xyzw
  {-# INLINE _xyzw #-}

instance Additive f => Affine (Point f) where
  type Diff (Point f) = f
  P x .-. P y = x ^-^ y
  {-# INLINE (.-.) #-}
  P x .+^ v = P (x ^+^ v)
  {-# INLINE (.+^) #-}
  P x .-^ v = P (x ^-^ v)
  {-# INLINE (.-^) #-}

-- | Vector spaces have origins.
origin :: (Additive f, Num a) => Point f a
origin = P zero

-- | An isomorphism between points and vectors, given a reference
--   point.
relative :: (Additive f, Num a) => Point f a -> Iso' (Point f a) (f a)
relative p0 = iso (.-. p0) (p0 .+^)
{-# INLINE relative #-}


{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE ConstraintKinds #-}
{-# LANGUAGE CPP #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE KindSignatures #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE Rank2Types #-}
{-# LANGUAGE ScopedTypeVariables #-}
#if __GLASGOW_HASKELL__ >= 805
{-# LANGUAGE ExplicitNamespaces #-}
{-# LANGUAGE NoStarIsType #-}
#endif
{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE TypeOperators #-}

module Eigen.Matrix
  (
    -- * Types
    Matrix(..)
  , Vec(..)
  , MatrixXf
  , MatrixXd
  , MatrixXcf
  , MatrixXcd

    -- * Common API
  , Elem
  , C
  , natToInt
  , Row(..)
  , Col(..)

    -- * Encode/Decode a Matrix
  , encode
  , decode

    -- * Querying a Matrix
  , null
  , square
  , rows
  , cols
  , dims

    -- * Constructing a Matrix
  , empty
  , constant
  , zero
  , ones
  , identity
  , random
  , diagonal

  , (!)
  , coeff
  , generate
  , sum
  , prod
  , mean
  , trace
  , all
  , any
  , count
  , norm
  , squaredNorm
  , blueNorm
  , hypotNorm
  , determinant
  , add
  , sub
  , mul
  , map
  , imap
  , TriangularMode(..)
  , triangularView
  , filter
  , ifilter
  , length
  , foldl
  , foldl'
  , inverse
  , adjoint
  , transpose
  , conjugate
  , normalize
  , modify
  , block
  , unsafeFreeze
  , unsafeWith
  , fromList
  , toList
  ) where

import Control.Monad (when)
import Control.Monad.ST (ST)
import Prelude hiding
  (map, null, filter, length, foldl, any, all, sum)
import Control.Monad (forM_)
import Control.Monad.Primitive (PrimMonad(..))
import Data.Binary (Binary(..))
import qualified Data.Binary as Binary
import qualified Data.ByteString.Lazy as BSL
import Data.Complex (Complex)
import Data.Constraint.Nat
import Eigen.Internal
  ( Elem
  , Cast(..)
  , natToInt
  , Row(..)
  , Col(..)
  )
import qualified Eigen.Internal as Internal
import qualified Eigen.Matrix.Mutable as M
import qualified Data.List as List
import Data.Kind (Type)
import GHC.TypeLits (Nat, type (*), type (<=), KnownNat)
import Foreign.C.Types (CInt)
import Foreign.C.String (CString)
import Foreign.Marshal.Alloc (alloca)
import Foreign.Ptr (Ptr)
import Foreign.Storable (peek)

import qualified Data.Vector.Storable as VS
import qualified Data.Vector.Storable.Mutable as VSM

-- | Matrix to be used in pure computations.
--
--   * Uses column majour memory layout.
--
--   * Has a copy-free FFI using the <http://eigen.tuxfamily.org Eigen> library.
--
newtype Matrix :: Nat -> Nat -> Type -> Type where
  Matrix :: Vec (n * m) a -> Matrix n m a

-- | Used internally to track the size and corresponding C type of the matrix.
newtype Vec :: Nat -> Type -> Type where
  Vec :: VS.Vector (C a) -> Vec n a

instance forall n m a. (Elem a, Show a, KnownNat n, KnownNat m) => Show (Matrix n m a) where
  show m = List.concat
    [ "Matrix ", show (rows m), "x", show (cols m)
    , "\n", List.intercalate "\n" $ List.map (List.intercalate "\t" . List.map show) $ toList m, "\n"
    ]

instance forall n m a. (KnownNat n, KnownNat m, Elem a) => Binary (Matrix n m a) where
  put (Matrix (Vec vals)) = do
    put $ Internal.magicCode (undefined :: C a)
    put $ natToInt @n
    put $ natToInt @m
    put vals

  get = do
    get >>= (`when` fail "wrong matrix type") . (/= Internal.magicCode (undefined :: C a))
    Matrix . Vec <$> get

-- | Encode the sparse matrix as a lazy bytestring
encode :: (Elem a, KnownNat n, KnownNat m) => Matrix n m a -> BSL.ByteString
encode = Binary.encode

-- | Decode the sparse matrix from a lazy bytestring
decode :: (Elem a, KnownNat n, KnownNat m) => BSL.ByteString -> Matrix n m a
decode = Binary.decode

-- | Alias for single precision matrix
type MatrixXf n m = Matrix n m Float
-- | Alias for double precision matrix
type MatrixXd n m = Matrix n m Double
-- | Alias for single precision matrix of complex numbers
type MatrixXcf n m = Matrix n m (Complex Float)
-- | Alias for double precision matrix of complex numbers
type MatrixXcd n m = Matrix n m (Complex Double)

-- | Construct an empty 0x0 matrix
empty :: Elem a => Matrix 0 0 a
{-# INLINE empty #-}
empty = Matrix (Vec (VS.empty))

-- | Is matrix empty?
null :: (Elem a, KnownNat n, KnownNat m) => Matrix n m a -> Bool
{-# INLINE null #-}
null m = cols m == 0 && rows m == 0

-- | Is matrix square?
--
square :: forall n m a. (Elem a, KnownNat n, KnownNat m) => Matrix n m a -> Bool
{-# INLINE square #-}
square _ = natToInt @n == natToInt @m

-- | Matrix where all coeffs are filled with the given value
constant :: forall n m a. (Elem a, KnownNat n, KnownNat m) => a -> Matrix n m a
{-# INLINE constant #-}
constant !val =
  let !cval = toC val
  in withDims $ \rs cs -> VS.replicate (rs * cs) cval

-- | Matrix where all coeffs are filled with 0
zero :: (Elem a, KnownNat n, KnownNat m) => Matrix n m a
{-# INLINE zero #-}
zero = constant 0

-- | Matrix where all coeffs are filled with 1
ones :: (Elem a, KnownNat n, KnownNat m) => Matrix n m a
{-# INLINE ones #-}
ones = constant 1

-- | The identity matrix (not necessarily square)
identity :: forall n m a. (Elem a, KnownNat n, KnownNat m) => Matrix n m a
identity =
  Internal.performIO $ do
     m :: M.IOMatrix n m a <- M.new
     Internal.call $ M.unsafeWith m Internal.identity
     unsafeFreeze m

-- | The random matrix of a given size
random :: forall n m a. (Elem a, KnownNat n, KnownNat m) => IO (Matrix n m a)
random = do
  m :: M.IOMatrix n m a <- M.new
  Internal.call $ M.unsafeWith m Internal.random
  unsafeFreeze m

withDims :: forall n m a. (Elem a, KnownNat n, KnownNat m) => (Int -> Int -> VS.Vector (C a)) -> Matrix n m a
{-# INLINE withDims #-}
withDims f =
  let !r = natToInt @n
      !c = natToInt @m
  in Matrix $ Vec $ f r c

-- | The number of rows in the matrix
rows :: forall n m a. KnownNat n => Matrix n m a -> Int
{-# INLINE rows #-}
rows _ = natToInt @n

-- | The number of colums in the matrix
cols :: forall n m a. KnownNat m => Matrix n m a -> Int
{-# INLINE cols #-}
cols _ = natToInt @m

-- | Return Matrix size as a pair of (rows, cols)
dims :: forall n m a. (Elem a, KnownNat n, KnownNat m) => Matrix n m a -> (Int, Int)
{-# INLINE dims #-}
dims _ = (natToInt @n, natToInt @m)

-- | Return the value at the given position.
(!) :: forall n m a r c. (Elem a, KnownNat n, KnownNat r, KnownNat c, r <= n, c <= m) => Row r -> Col c -> Matrix n m a -> a
{-# INLINE (!) #-}
(!) = coeff

-- | Return the value at the given position.
coeff :: forall n m a r c. (Elem a, KnownNat n, KnownNat r, KnownNat c, r <= n, c <= m) => Row r -> Col c -> Matrix n m a -> a
{-# INLINE coeff #-}
coeff _ _ m@(Matrix (Vec vals)) =
  let !row  = natToInt @r
      !col  = natToInt @c
  in fromC $! VS.unsafeIndex vals $! col * rows m + row

unsafeCoeff :: (Elem a, KnownNat n) => Int -> Int -> Matrix n m a -> a
{-# INLINE unsafeCoeff #-}
unsafeCoeff row col m@(Matrix (Vec vals)) = fromC $! VS.unsafeIndex vals $! col * rows m + row

-- | Given a generation function `f :: Int -> Int -> a`, construct a Matrix of known size
--   using points in the matrix as inputs.
generate :: forall n m a. (Elem a, KnownNat n, KnownNat m) => (Int -> Int -> a) -> Matrix n m a
generate f = withDims $ \rs cs -> VS.create $ do
  vals :: VSM.MVector s (C a) <- VSM.new (rs * cs)
  forM_ [0 .. pred rs] $ \r ->
    forM_ [0 .. pred cs] $ \c ->
      VSM.write vals (c * rs + r) (toC $! f r c)
  pure vals

-- | The sum of all coefficients in the matrix
sum :: (Elem a, KnownNat n, KnownNat m) => Matrix n m a -> a
sum = _prop Internal.sum

-- | The product of all coefficients in the matrix
prod :: (Elem a, KnownNat n, KnownNat m) => Matrix n m a -> a
prod = _prop Internal.prod

-- | The arithmetic mean of all coefficients in the matrix
mean :: (Elem a, KnownNat n, KnownNat m) => Matrix n m a -> a
mean = _prop Internal.mean

-- | The trace of a matrix is the sum of the diagonal coefficients.
--
--   'trace' m == 'sum' ('diagonal' m)
trace :: (Elem a, KnownNat n, KnownNat m) => Matrix n m a -> a
trace = _prop Internal.trace

-- | Given a predicate p, determine if all values in the Matrix satisfy p.
all :: (Elem a, KnownNat n, KnownNat m) => (a -> Bool) -> Matrix n m a -> Bool
all f (Matrix (Vec vals)) = VS.all (f . fromC) vals

-- | Given a predicate p, determine if any values in the Matrix satisfy p.
any :: (Elem a, KnownNat n, KnownNat m) => (a -> Bool) -> Matrix n m a -> Bool
any f (Matrix (Vec vals)) = VS.any (f . fromC) vals

-- | Given a predicate p, determine how many values in the Matrix satisfy p.
count :: (Elem a, KnownNat n, KnownNat m) => (a -> Bool) -> Matrix n m a -> Int
count f (Matrix (Vec vals)) = VS.foldl' (\n x-> if f (fromC x) then (n + 1) else n) 0 vals

norm, squaredNorm, blueNorm, hypotNorm :: (Elem a, KnownNat n, KnownNat m) => Matrix n m a -> a

{-| For vectors, the l2 norm, and for matrices the Frobenius norm.
    In both cases, it consists in the square root of the sum of the square of all the matrix entries.
    For vectors, this is also equals to the square root of the dot product of this with itself.
-}
norm = _prop Internal.norm

-- | For vectors, the squared l2 norm, and for matrices the Frobenius norm. In both cases, it consists in the sum of the square of all the matrix entries. For vectors, this is also equals to the dot product of this with itself.
squaredNorm = _prop Internal.squaredNorm

-- | The l2 norm of the matrix using the Blue's algorithm. A Portable Fortran Program to Find the Euclidean Norm of a Vector, ACM TOMS, Vol 4, Issue 1, 1978.
blueNorm = _prop Internal.blueNorm

-- | The l2 norm of the matrix avoiding undeflow and overflow. This version use a concatenation of hypot calls, and it is very slow.
hypotNorm = _prop Internal.hypotNorm

-- | The determinant of the matrix
determinant :: forall n a. (Elem a, KnownNat n) => Matrix n n a -> a
determinant m = _prop Internal.determinant m

-- | Add two matrices.
add :: (Elem a, KnownNat n, KnownNat m) => Matrix n m a -> Matrix n m a -> Matrix n m a
add m1 m2 = _binop Internal.add m1 m2

-- | Subtract two matrices.
sub :: (Elem a, KnownNat n, KnownNat m) => Matrix n m a -> Matrix n m a -> Matrix n m a
sub m1 m2 = _binop Internal.sub m1 m2

-- | Multiply two matrices.
mul :: (Elem a, KnownNat p, KnownNat q, KnownNat r) => Matrix p q a -> Matrix q r a -> Matrix p r a
mul m1 m2 = _binop Internal.mul m1 m2

{- | Apply a given function to each element of the matrix.
Here is an example how to implement scalar matrix multiplication:

>>> let a = fromList [[1,2],[3,4]] :: MatrixXf 2 2
>>> a
Matrix 2x2
1.0 2.0
3.0 4.0
>>> map (*10) a
Matrix 2x2
10.0    20.0
30.0    40.0
-}
map :: Elem a => (a -> a) -> Matrix n m a -> Matrix n m a
map f (Matrix (Vec vals)) = Matrix $ Vec $ VS.map (toC . f . fromC) vals

{- | Apply a given function to each element of the matrix.
Here is an example how upper triangular matrix can be implemented:
>>> let a = fromList [[1,2,3],[4,5,6],[7,8,9]] :: MatrixXf
>>> a
Matrix 3x3
1.0 2.0 3.0
4.0 5.0 6.0
7.0 8.0 9.0
>>> imap (\row col val -> if row <= col then val else 0) a
Matrix 3x3
1.0 2.0 3.0
0.0 5.0 6.0
0.0 0.0 9.0
-}
imap :: (Elem a, KnownNat n, KnownNat m) => (Int -> Int -> a -> a) -> Matrix n m a -> Matrix n m a
imap f (Matrix (Vec vals)) =
  withDims $ \rs _ ->
    VS.imap (\n ->
      let (c,r) = divMod n rs
      in toC . f r c . fromC) vals

-- | Provide a view of the matrix for extraction of a subset.
data TriangularMode
  -- | View matrix as a lower triangular matrix.
  = Lower
  -- | View matrix as an upper triangular matrix.
  | Upper
  -- | View matrix as a lower triangular matrix with zeros on the diagonal.
  | StrictlyLower
  -- | View matrix as an upper triangular matrix with zeros on the diagonal.
  | StrictlyUpper
  -- | View matrix as a lower triangular matrix with ones on the diagonal.
  | UnitLower
  -- | View matrix as an upper triangular matrix with ones on the diagonal.
  | UnitUpper
  deriving (Eq, Enum, Show, Read)

-- | Triangular view extracted from the current matrix
triangularView :: (Elem a, KnownNat n, KnownNat m) => TriangularMode -> Matrix n m a -> Matrix n m a
triangularView = \case
  Lower         -> imap $ \row col val -> case compare row col of { LT -> 0; _ -> val }
  Upper         -> imap $ \row col val -> case compare row col of { GT -> 0; _ -> val }
  StrictlyLower -> imap $ \row col val -> case compare row col of { GT -> val; _ -> 0 }
  StrictlyUpper -> imap $ \row col val -> case compare row col of { LT -> val; _ -> 0 }
  UnitLower     -> imap $ \row col val -> case compare row col of { GT -> val; LT -> 0; EQ -> 1 }
  UnitUpper     -> imap $ \row col val -> case compare row col of { LT -> val; GT -> 0; EQ -> 1 }

-- | Filter elements in the matrix. Filtered elements will be replaced by 0.
filter :: Elem a => (a -> Bool) -> Matrix n m a -> Matrix n m a
filter f = map (\x -> if f x then x else 0)

-- | Filter elements in the matrix with an indexed predicate. Filtered elements will be replaces by 0.
ifilter :: (Elem a, KnownNat n, KnownNat m) => (Int -> Int -> a -> Bool) -> Matrix n m a -> Matrix n m a
ifilter f = imap (\r c x -> if f r c x then x else 0)

-- | The length of the matrix.
length :: forall n m a r. (Elem a, KnownNat n, KnownNat m, r ~ (n * m), KnownNat r) => Matrix n m a -> Int
length _ = natToInt @r

-- | Left fold of a matrix, where accumulation is lazy.
foldl :: (Elem a, KnownNat n, KnownNat m) => (b -> a -> b) -> b -> Matrix n m a -> b
foldl f b (Matrix (Vec vals)) = VS.foldl (\a x -> f a (fromC x)) b vals

-- | Right fold of a matrix, where accumulation is strict.
foldl' :: Elem a => (b -> a -> b) -> b -> Matrix n m a -> b
foldl' f b (Matrix (Vec vals)) = VS.foldl' (\ !a x -> f a (fromC x)) b vals

-- | Return the diagonal of a matrix.
diagonal :: (Elem a, KnownNat n, KnownNat m, r ~ Min n m, KnownNat r) => Matrix n m a -> Matrix r 1 a
diagonal = _unop Internal.diagonal

{- | Inverse of the matrix
For small fixed sizes up to 4x4, this method uses cofactors. In the general case, this method uses PartialPivLU decomposition
-}
inverse :: forall n a. (Elem a, KnownNat n) => Matrix n n a -> Matrix n n a
inverse = _unop Internal.inverse

-- | Adjoint of the matrix
adjoint :: (Elem a, KnownNat n, KnownNat m) => Matrix n m a -> Matrix m n a
adjoint = _unop Internal.adjoint

-- | Transpose of the matrix
transpose :: (Elem a, KnownNat n, KnownNat m) => Matrix n m a -> Matrix m n a
transpose = _unop Internal.transpose

-- | Conjugate of the matrix
conjugate :: (Elem a, KnownNat n, KnownNat m) => Matrix n m a -> Matrix n m a
conjugate = _unop Internal.conjugate

-- | Normalise the matrix by dividing it on its 'norm'
normalize :: forall n m a. (Elem a, KnownNat n, KnownNat m) => Matrix n m a -> Matrix n m a
normalize (Matrix (Vec vals)) = Internal.performIO $ do
  vals' <- VS.thaw vals
  VSM.unsafeWith vals' $ \p ->
    let !rs = natToInt @n
        !cs = natToInt @m
    in Internal.call $ Internal.normalize p (toC rs) (toC cs)
  Matrix . Vec <$> VS.unsafeFreeze vals'

-- | Apply a destructive operation to a matrix. The operation will be performed in-place, if it is safe
--   to do so - otherwise, it will create a copy of the matrix.
modify :: (Elem a, KnownNat n, KnownNat m) => (forall s. M.MMatrix n m s a -> ST s ()) -> Matrix n m a -> Matrix n m a
modify f (Matrix (Vec vals)) = Matrix $ Vec $ VS.modify (f . M.fromVector ) vals

-- | Extract rectangular block from matrix defined by startRow startCol blockRows blockCols
block :: forall sr sc br bc n m a.
     (Elem a, KnownNat sr, KnownNat sc, KnownNat br, KnownNat bc, KnownNat n, KnownNat m)
  => (sr <= n, sc <= m, br <= n, bc <= m)
  => Row sr -- ^ starting row
  -> Col sc -- ^ starting col
  -> Row br -- ^ block of rows
  -> Col bc -- ^ block of cols
  -> Matrix n m a -- ^ extract from this
  -> Matrix br bc a -- ^ extraction
block _ _ _ _ m =
  let !startRow = natToInt @sr
      !startCol = natToInt @sc
  in generate $ \row col -> unsafeCoeff (startRow + row) (startCol + col) m

-- | Turn a mutable matrix into an immutable matrix without copying.
--   The mutable matrix should not be modified after this conversion.
unsafeFreeze :: (Elem a, KnownNat n, KnownNat m, PrimMonad p) => M.MMatrix n m (PrimState p) a -> p (Matrix n m a)
unsafeFreeze m = VS.unsafeFreeze (M.vals m) >>= pure . Matrix . Vec

-- | Pass a pointer to the matrix's data to the IO action. The data may not be modified through the pointer.
unsafeWith  :: (Elem a, KnownNat n, KnownNat m) => Matrix n m a -> (Ptr (C a) -> CInt -> CInt -> IO b) -> IO b
unsafeWith m@(Matrix (Vec (vals))) f =
  VS.unsafeWith vals $ \p ->
    let !rs = toC $! rows m
        !cs = toC $! cols m
    in f p rs cs

_prop :: (Elem a, KnownNat n, KnownNat m) => (Ptr (C a) -> Ptr (C a) -> CInt -> CInt -> IO CString) -> Matrix n m a -> a
{-# INLINE _prop #-}
_prop f m = fromC $ Internal.performIO $ alloca $ \p -> do
   Internal.call $ unsafeWith m (f p)
   peek p

_binop :: forall n m n1 m1 n2 m2 a. (Elem a, KnownNat n, KnownNat m, KnownNat n1, KnownNat m1, KnownNat n2, KnownNat m2)
  => (Ptr (C a) -> CInt -> CInt -> Ptr (C a) -> CInt -> CInt -> Ptr (C a) -> CInt -> CInt -> IO CString)
  -> Matrix n m a
  -> Matrix n1 m1 a
  -> Matrix n2 m2 a
{-# INLINE _binop #-}
_binop g m1 m2 = Internal.performIO $ do
  m0 :: M.IOMatrix n2 m2 a <- M.new
  M.unsafeWith m0 $ \vals0 rows0 cols0 ->
      unsafeWith m1 $ \vals1 rows1 cols1 ->
          unsafeWith m2 $ \vals2 rows2 cols2 ->
              Internal.call $ g
                vals0 rows0 cols0
                vals1 rows1 cols1
                vals2 rows2 cols2
  unsafeFreeze m0

_unop :: forall n m n1 m1 a. (Elem a, KnownNat n, KnownNat m, KnownNat n1, KnownNat m1)
  => (Ptr (C a) -> CInt -> CInt -> Ptr (C a) -> CInt -> CInt -> IO CString)
  -> Matrix n m a
  -> Matrix n1 m1 a
{-# INLINE _unop #-}
_unop g m1 = Internal.performIO $ do
  m0 :: M.IOMatrix n1 m1 a <- M.new
  M.unsafeWith m0 $ \vals0 rows0 cols0 ->
      unsafeWith m1 $ \vals1 rows1 cols1 ->
          Internal.call $ g
              vals0 rows0 cols0
              vals1 rows1 cols1
  unsafeFreeze m0

-- | Convert a matrix to a list.
toList :: (Elem a, KnownNat n, KnownNat m) => Matrix n m a -> [[a]]
{-# INLINE toList #-}
toList m@(Matrix (Vec vals))
  | null m = []
  | otherwise = [[fromC $ vals `VS.unsafeIndex` (col * _rows + row) | col <- [0..pred _cols]] | row <- [0..pred _rows]]
  where
    !_rows = rows m
    !_cols = cols m

-- | Convert a list to a matrix. Returns 'Nothing' if the dimensions of the list do not match that
--   of the matrix.
fromList :: forall n m a. (Elem a, KnownNat n, KnownNat m) => [[a]] -> Maybe (Matrix n m a)
fromList list = do
  let myRows = natToInt @n
  let myCols = natToInt @m
  let _rows  = List.length list
  let _cols  = List.foldl' max 0 (List.map List.length list)
  if ((myRows /= _rows) || (myCols /= _cols))
    then Nothing
    else (Just . Matrix . Vec) $ VS.create $ do
      vm <- VSM.replicate (_rows * _cols) (toC (0 :: a))
      forM_ (zip [0..] list) $ \(row,vals) ->
        forM_ (zip [0..] vals) $ \(col, val) ->
          VSM.write vm (col * _rows + row) (toC val)
      pure vm


{-# OPTIONS_GHC -fplugin GHC.TypeLits.KnownNat.Solver #-}

{-# LANGUAGE CPP                  #-}
{-# LANGUAGE DataKinds            #-}
{-# LANGUAGE FlexibleContexts     #-}
{-# LANGUAGE GADTs                #-}
{-# LANGUAGE RankNTypes           #-}
{-# LANGUAGE ScopedTypeVariables  #-}
{-# LANGUAGE StandaloneDeriving   #-}
{-# LANGUAGE TypeFamilies         #-}
{-# LANGUAGE TypeOperators        #-}
{-# LANGUAGE UndecidableInstances #-}
{-|
Module      : Grenade.Core.Shape
Description : Dependently typed shapes of data which are passed between layers of a network
Copyright   : (c) Huw Campbell, 2016-2017
License     : BSD2
Stability   : experimental

-}
module Grenade.Core.Shape (
    S (..)
  , Shape (..)
#if MIN_VERSION_singletons(2,6,0)
  , SShape (..)
#else
  , Sing (..)
#endif

  , randomOfShape
  , fromStorable
  , fromStorableMatrix
  , nk
  , visualise2D
  , splitChannels
  , combineChannels
  ) where

#if MIN_VERSION_singletons(2,6,0)
import           Data.Kind                    (Type)
#endif

import           Control.DeepSeq              (NFData (..))
import           Data.List.Split              (chunksOf)
import           Data.Maybe                   (fromJust)
import           Data.Proxy
import           Data.Serialize
import           Data.Singletons
import           Data.Singletons.TypeLits
import           Data.Vector.Storable         (Vector)
import qualified Data.Vector.Storable         as V
import           GHC.TypeLits                 hiding (natVal)
import           Numeric.LinearAlgebra        (Matrix)
import qualified Numeric.LinearAlgebra        as NLA
import qualified Numeric.LinearAlgebra.Data   as NLAD
import           Numeric.LinearAlgebra.Static
import qualified Numeric.LinearAlgebra.Static as H

import           System.Random.MWC

import           Grenade.Types

-- | The current shapes we accept.
--   at the moment this is just one, two, and three dimensional
--   Vectors/Matricies.
--
--   These are only used with DataKinds, as Kind `Shape`, with Types 'D1, 'D2, 'D3.
data Shape
  = D1 Nat
  -- ^ One dimensional vector
  | D2 Nat Nat
  -- ^ Two dimensional matrix. Row, Column.
  | D3 Nat Nat Nat
  -- ^ Three dimensional matrix. Row, Column, Channels.
  | D4 Nat Nat Nat Nat
  -- ^ Four dimensional matrix. Depth, Channels, Row, Column

-- | Concrete data structures for a Shape.
--
--   All shapes are held in contiguous memory.
--   3D is held in a matrix (usually row oriented) which has height depth * rows.
data S (n :: Shape) where
  S1D :: ( KnownNat len )
      => R len
      -> S ('D1 len)

  S2D :: ( KnownNat rows, KnownNat columns )
      => L rows columns
      -> S ('D2 rows columns)

  S3D :: ( KnownNat rows
         , KnownNat columns
         , KnownNat channels
         , KnownNat (rows * channels) )
      => L (rows * channels) columns
      -> S ('D3 rows columns channels)

  S4D :: ( KnownNat rows
         , KnownNat columns
         , KnownNat channels
         , KnownNat depth)
      => L (depth * channels * rows) columns
      -> S ('D4 depth channels rows columns)

deriving instance Show (S n)

-- Singleton instances.
--
-- These could probably be derived with template haskell, but this seems
-- clear and makes adding the KnownNat constraints simple.
-- We can also keep our code TH free, which is great.
#if MIN_VERSION_singletons(2,6,0)
-- In singletons 2.6 Sing switched from a data family to a type family.
type instance Sing = SShape

-- | Datatype to allow "pattern matching" on shapes at the type level through
--   the use the sing function from Singletons.
data SShape :: Shape -> Type where
  D1Sing :: Sing a -> SShape ('D1 a)
  D2Sing :: Sing a -> Sing b -> SShape ('D2 a b)
  D3Sing :: KnownNat (a * c) => Sing a -> Sing b -> Sing c -> SShape ('D3 a b c)
  D4Sing :: KnownNat (a * c * d) => Sing a -> Sing b -> Sing c -> Sing d -> SShape ('D4 a b c d)
#else
data instance Sing (n :: Shape) where
  D1Sing :: Sing a -> Sing ('D1 a)
  D2Sing :: Sing a -> Sing b -> Sing ('D2 a b)
  D3Sing :: KnownNat (a * c) => Sing a -> Sing b -> Sing c -> Sing ('D3 a b c)
  D4Sing :: KnownNat (a * c * d) => Sing a -> Sing b -> Sing c -> Sing d -> Sing ('D4 a b c d)
#endif

instance KnownNat a => SingI ('D1 a) where
  sing = D1Sing sing
instance (KnownNat a, KnownNat b) => SingI ('D2 a b) where
  sing = D2Sing sing sing
instance (KnownNat a, KnownNat b, KnownNat c, KnownNat (a * c)) => SingI ('D3 a b c) where
  sing = D3Sing sing sing sing
instance (KnownNat a, KnownNat b, KnownNat c, KnownNat d) => SingI ('D4 a b c d) where
  sing = D4Sing sing sing sing sing

instance SingI x => Num (S x) where
  (+) = n2 (+)
  (-) = n2 (-)
  (*) = n2 (*)
  abs = n1 abs
  signum = n1 signum
  fromInteger x = nk (fromInteger x)

instance SingI x => Fractional (S x) where
  (/) = n2 (/)
  recip = n1 recip
  fromRational x = nk (fromRational x)

instance SingI x => Floating (S x) where
  pi = nk pi
  exp = n1 exp
  log = n1 log
  sqrt = n1 sqrt
  (**) = n2 (**)
  logBase = n2 logBase
  sin = n1 sin
  cos = n1 cos
  tan = n1 tan
  asin = n1 asin
  acos = n1 acos
  atan = n1 atan
  sinh = n1 sinh
  cosh = n1 cosh
  tanh = n1 tanh
  asinh = n1 asinh
  acosh = n1 acosh
  atanh = n1 atanh

--
-- I haven't made shapes strict, as sometimes they're not needed
-- (the last input gradient back for instance)
--
instance NFData (S x) where
  rnf (S1D x) = rnf x
  rnf (S2D x) = rnf x
  rnf (S3D x) = rnf x
  rnf (S4D x) = rnf x

-- | Generate random data of the desired shape
randomOfShape :: forall x . (SingI x) => IO (S x)
randomOfShape = do
  seed :: Int <- withSystemRandom . asGenST $ \gen -> uniform gen
  return $ case (sing :: Sing x) of
    D1Sing SNat ->
        S1D (H.randomVector seed H.Uniform * 2 - 1)

    D2Sing SNat SNat ->
        S2D (H.uniformSample seed (-1) 1)

    D3Sing SNat SNat SNat ->
        S3D (H.uniformSample seed (-1) 1)

    D4Sing SNat SNat SNat SNat ->
        S4D (H.uniformSample seed (-1) 1)

-- | Generate a shape from a Storable Vector.
--
--   Returns Nothing if the vector is of the wrong size.
fromStorable :: forall x. SingI x => Vector RealNum -> Maybe (S x)
fromStorable xs = case sing :: Sing x of
    D1Sing SNat ->
      S1D <$> H.create xs

    D2Sing SNat SNat ->
      S2D <$> mkL xs

    D3Sing SNat SNat SNat ->
      S3D <$> mkL xs

    D4Sing SNat SNat SNat SNat ->
      S4D <$> mkL xs
  where
    mkL :: forall rows columns. (KnownNat rows, KnownNat columns)
        => Vector RealNum -> Maybe (L rows columns)
    mkL v =
      let rows    = fromIntegral $ natVal (Proxy :: Proxy rows)
          columns = fromIntegral $ natVal (Proxy :: Proxy columns)
      in  if rows * columns == V.length v
             then H.create $ NLA.reshape columns v
             else Nothing

-- | Generate a shape from a Storable Matrix.
--
--   Returns Nothing if the matrix is of the wrong size (or one dimensional)
fromStorableMatrix :: forall x. SingI x => Matrix RealNum -> Maybe (S x)
fromStorableMatrix xs = case sing :: Sing x of
    D1Sing SNat ->
      Nothing

    D2Sing SNat SNat ->
      S2D <$> mkL xs

    D3Sing SNat SNat SNat ->
      S3D <$> mkL xs

    D4Sing SNat SNat SNat SNat ->
      S4D <$> mkL xs
  where
    mkL :: forall rows columns. (KnownNat rows, KnownNat columns)
        => Matrix RealNum -> Maybe (L rows columns)
    mkL m =
      let rows    = fromIntegral $ natVal (Proxy :: Proxy rows)
          columns = fromIntegral $ natVal (Proxy :: Proxy columns)
          (h, w)  = NLA.size m
      in  if rows == h && columns == w
             then H.create m
             else Nothing

instance SingI x => Serialize (S x) where
  put i = (case i of
            (S1D x) -> putListOf put . NLA.toList . H.extract $ x
            (S2D x) -> putListOf put . NLA.toList . NLA.flatten . H.extract $ x
            (S3D x) -> putListOf put . NLA.toList . NLA.flatten . H.extract $ x
            (S4D x) -> putListOf put . NLA.toList . NLA.flatten . H.extract $ x
          ) :: PutM ()

  get = do
    Just i <- fromStorable . V.fromList <$> getListOf get
    return i

-- Helper function for creating the number instances
n1 :: ( forall a. Floating a => a -> a ) -> S x -> S x
n1 f (S1D x) = S1D (f x)
n1 f (S2D x) = S2D (f x)
n1 f (S3D x) = S3D (f x)
n1 f (S4D x) = S4D (f x)

-- Helper function for creating the number instances
n2 :: ( forall a. Floating a => a -> a -> a ) -> S x -> S x -> S x
n2 f (S1D x) (S1D y) = S1D (f x y)
n2 f (S2D x) (S2D y) = S2D (f x y)
n2 f (S3D x) (S3D y) = S3D (f x y)
n2 f (S4D x) (S4D y) = S4D (f x y)

-- | Helper function for creating the number instances
nk :: forall x. SingI x => RealNum -> S x
nk x = case (sing :: Sing x) of
  D1Sing SNat ->
    S1D (H.konst x)

  D2Sing SNat SNat ->
    S2D (H.konst x)

  D3Sing SNat SNat SNat ->
    S3D (H.konst x)

  D4Sing SNat SNat SNat SNat ->
    S4D (H.konst x)

-- | Prints out the contents of a matrix with entries between zero and the max value
visualise2D :: S ('D2 a b)  -- ^ input matrix
            -> RealNum      -- ^ maximum element value
            -> String       -- ^ pretty printed matrix
visualise2D (S2D mm) max =
  let m  = H.extract mm
      ms = NLAD.toLists m
      render n' | n' <= 0.2 * max  = ' '
                | n' <= 0.4 * max  = '.'
                | n' <= 0.6 * max  = '-'
                | n' <= 0.8 * max  = '='
                | otherwise =  '#'
      px = (fmap . fmap) render ms
  in unlines px

-- | TODO Theo
splitChannels :: forall rows columns channels.
                 (KnownNat rows, KnownNat columns)
                 => S ('D3 rows columns channels) -> [S ('D2 rows columns)]
splitChannels (S3D x)
  = let r   = fromIntegral $ natVal (Proxy :: Proxy rows)
        rs  = NLA.toRows $ H.extract x
        rs' = chunksOf r rs
        ms  = map (S2D . fromJust . H.create . NLA.fromRows) rs' :: [S ('D2 rows columns)]
    in ms

-- | TODO Theo
combineChannels :: forall rows columns channels.
                 (KnownNat rows, KnownNat columns, KnownNat channels)
                 => [S ('D2 rows columns)] -> S ('D3 rows columns channels)
combineChannels xs
  = let xs' = map (\(S2D x) -> NLA.toRows $ H.extract x) xs :: [[Vector RealNum]]
        xs'' = concat xs'
    in  S3D $ fromJust . H.create . NLA.fromRows $ xs''


-- | See Hoffman, Gelman (2011) The No U-Turn Sampler: Adaptively Setting Path
--   Lengths in Hamiltonian Monte Carlo.

module Numeric.MCMC.NUTS where

import Control.Monad
import Control.Monad.Loops
import Control.Monad.Primitive
import System.Random.MWC
import System.Random.MWC.Distributions hiding (gamma)
import Statistics.Distribution.Normal

type Parameters = [Double]
type Density    = Parameters -> Double
type Gradient   = Parameters -> Parameters
type Particle   = (Parameters, Parameters)
type StateSpecs = ([Double], [Double], [Double], [Double], [Double], Int, Int)

newtype StateTree = StateTree {
    getStateTree :: StateSpecs
  }

instance Show StateTree where
  show (StateTree (tm, rm, tp, rp, t', n, s)) =
       "\n" ++ "tm: " ++ show tm
    ++ "\n" ++ "tp: " ++ show tp
    ++ "\n" ++ "t': " ++ show t'
    ++ "\n" ++ "n : " ++ show n
    ++ "\n" ++ "s : " ++ show s

data DualAveragingParameters = DualAveragingParameters {
    mAdapt :: Int
  , delta  :: Double
  , mu     :: Double
  , gamma  :: Double
  , tau0   :: Double
  , kappa  :: Double
  } deriving Show

-- | The NUTS sampler.
nuts
  :: PrimMonad m
  => Density
  -> Gradient
  -> Int
  -> Double
  -> Parameters
  -> Gen (PrimState m)
  -> m [Parameters]
nuts lTarget glTarget n e t g = go t 0 []
  where go position j acc
          | j >= n    = return acc
          | otherwise = do
              nextPosition <- nutsKernel lTarget glTarget e position g
              go nextPosition (succ j) (nextPosition : acc)

-- | The NUTS sampler with dual averaging.
nutsDualAveraging
  :: PrimMonad m
  => Density
  -> Gradient
  -> Int
  -> Int
  -> Parameters
  -> Gen (PrimState m)
  -> m [Parameters]
nutsDualAveraging lTarget glTarget n nAdapt t g = do
    e0 <- findReasonableEpsilon lTarget glTarget t g
    let daParams = basicDualAveragingParameters e0 nAdapt
    chain <- unfoldrM (kernel daParams) (1, e0, 1, 0, t)
    return $ drop nAdapt chain
  where
    kernel params (m, e, eAvg, h, t0) = do
      (eNext, eAvgNext, hNext, tNext) <-
        nutsKernelDualAvg lTarget glTarget e eAvg h m params t0 g
      return $ if   m > n + nAdapt
               then Nothing
               else Just (t0, (succ m, eNext, eAvgNext, hNext, tNext))

-- | Default DA parameters, given a base step size and burn in period.
basicDualAveragingParameters :: Double -> Int -> DualAveragingParameters
basicDualAveragingParameters step burnInPeriod = DualAveragingParameters {
    mu     = log (10 * step)
  , delta  = 0.5
  , mAdapt = burnInPeriod
  , gamma  = 0.05
  , tau0   = 10
  , kappa  = 0.75
  }

-- | A single iteration of dual-averaging NUTS.
nutsKernelDualAvg
  :: PrimMonad m
  => Density
  -> Gradient
  -> Double
  -> Double
  -> Double
  -> Int
  -> DualAveragingParameters
  -> [Double]
  -> Gen (PrimState m)
  -> m (Double, Double, Double, Parameters)
nutsKernelDualAvg lTarget glTarget e eAvg h m daParams t g = do
  r0 <- replicateM (length t) (normal 0 1 g)
  z0 <- exponential 1 g
  let logu = log (auxilliaryTarget lTarget t r0) - z0

  let go (tn, tp, rn, rp, tm, j, n, s, a, na) g
        | s == 1 = do
            vj <- symmetricCategorical [-1, 1] g
            z  <- uniform g

            (tnn, rnn, tpp, rpp, t1, n1, s1, a1, na1) <-
              if   vj == -1
              then do
                (tnn', rnn', _, _, t1', n1', s1', a1', na1') <-
                  buildTreeDualAvg lTarget glTarget g tn rn logu vj j e t r0
                return (tnn', rnn', tp, rp, t1', n1', s1', a1', na1')
              else do
                (_, _, tpp', rpp', t1', n1', s1', a1', na1') <-
                  buildTreeDualAvg lTarget glTarget g tp rp logu vj j e t r0
                return (tn, rn, tpp', rpp', t1', n1', s1', a1', na1')

            let accept = s1 == 1 && (min 1 (fi n1 / fi n :: Double)) > z

                n2 = n + n1
                s2 = s1 * stopCriterion tnn tpp rnn rpp
                j1 = succ j
                t2 | accept    = t1
                   | otherwise = tm

            go (tnn, tpp, rnn, rpp, t2, j1, n2, s2, a1, na1) g

        | otherwise = return (tm, a, na)

  (nextPosition, alpha, nalpha) <- go (t, t, r0, r0, t, 0, 1, 1, 0, 0) g

  let (hNext, eNext, eAvgNext) =
          if   m <= mAdapt daParams
          then (hm, exp logEm, exp logEbarM)
          else (h, eAvg, eAvg)
        where
          eta = 1 / (fromIntegral m + tau0 daParams)
          hm  = (1 - eta) * h
              + eta * (delta daParams - alpha / fromIntegral nalpha)

          zeta = fromIntegral m ** (- (kappa daParams))

          logEm    = mu daParams - sqrt (fromIntegral m) / gamma daParams * hm
          logEbarM = (1 - zeta) * log eAvg + zeta * logEm

  return (eNext, eAvgNext, hNext, nextPosition)

-- | A single iteration of NUTS.
nutsKernel
  :: PrimMonad m
  => Density
  -> Gradient
  -> Double
  -> Parameters
  -> Gen (PrimState m)
  -> m Parameters
nutsKernel lTarget glTarget e t g = do
  r0   <- replicateM (length t) (normal 0 1 g)
  z0   <- exponential 1 g
  let logu = log (auxilliaryTarget lTarget t r0) - z0

  let go (tn, tp, rn, rp, tm, j, n, s) g
        | s == 1 = do
            vj <- symmetricCategorical [-1, 1] g
            z  <- uniform g

            (tnn, rnn, tpp, rpp, t1, n1, s1) <-
              if   vj == -1
              then do
                (tnn', rnn', _, _, t1', n1', s1') <-
                  buildTree lTarget glTarget g tn rn logu vj j e
                return (tnn', rnn', tp, rp, t1', n1', s1')
              else do
                (_, _, tpp', rpp', t1', n1', s1') <-
                  buildTree lTarget glTarget g tp rp logu vj j e
                return (tn, rn, tpp', rpp', t1', n1', s1')

            let accept = s1 == 1 && (min 1 (fi n1 / fi n :: Double)) > z

                n2 = n + n1
                s2 = s1 * stopCriterion tnn tpp rnn rpp
                j1 = succ j
                t2 | accept    = t1
                   | otherwise = tm

            go (tnn, tpp, rnn, rpp, t2, j1, n2, s2) g

        | otherwise = return tm

  go (t, t, r0, r0, t, 0, 1, 1) g

-- | Build the 'tree' of candidate states.
buildTree
  :: PrimMonad m
  => Density
  -> Gradient
  -> Gen (PrimState m)
  -> Parameters
  -> Parameters
  -> Double
  -> Double
  -> Int
  -> Double
  -> m StateSpecs
buildTree lTarget glTarget g t r logu v 0 e = do
  let (t0, r0) = leapfrog glTarget (t, r) (v * e)
      joint    = log $ auxilliaryTarget lTarget t0 r0
      n        = indicate (logu < joint)
      s        = indicate (logu - 1000 < joint)
  return (t0, r0, t0, r0, t0, n, s)

buildTree lTarget glTarget g t r logu v j e = do
  z <- uniform g
  (tn, rn, tp, rp, t0, n0, s0) <-
    buildTree lTarget glTarget g t r logu v (pred j) e

  if   s0 == 1
  then do
    (tnn, rnn, tpp, rpp, t1, n1, s1) <-
      if   v == -1
      then do
        (tnn', rnn', _, _, t1', n1', s1') <-
          buildTree lTarget glTarget g tn rn logu v (pred j) e
        return (tnn', rnn', tp, rp, t1', n1', s1')
      else do
        (_, _, tpp', rpp', t1', n1', s1') <-
          buildTree lTarget glTarget g tp rp logu v (pred j) e
        return (tn, rn, tpp', rpp', t1', n1', s1')

    let accept = (fi n1 / max (fi (n0 + n1)) 1) > (z :: Double)
        n2     = n0 + n1
        s2     = s0 * s1 * stopCriterion tnn tpp rnn rpp
        t2     | accept    = t1
               | otherwise = t0

    return (tnn, rnn, tpp, rpp, t2, n2, s2)
  else return (tn, rn, tp, rp, t0, n0, s0)

-- | Determine whether or not to stop doubling the tree of candidate states.
stopCriterion :: (Integral a, Num b, Ord b) => [b] -> [b] -> [b] -> [b] -> a
stopCriterion tn tp rn rp =
      indicate (positionDifference `innerProduct` rn >= 0)
    * indicate (positionDifference `innerProduct` rp >= 0)
  where
    positionDifference = tp .- tn

-- | Build the tree of candidate states under dual averaging.
buildTreeDualAvg
  :: PrimMonad m
  => Density
  -> Gradient
  -> Gen (PrimState m)
  -> Parameters
  -> Parameters
  -> Double
  -> Double
  -> Int
  -> Double
  -> Parameters
  -> Parameters
  -> m ([Double], [Double], [Double], [Double], [Double], Int, Int, Double, Int)
buildTreeDualAvg lTarget glTarget g t r logu v 0 e t0 r0 = do
  let (t1, r1) = leapfrog glTarget (t, r) (v * e)
      joint    = log $ auxilliaryTarget lTarget t1 r1
      n        = indicate (logu < joint)
      s        = indicate (logu - 1000 <  joint)
      a        = min 1 (acceptanceRatio lTarget t0 t1 r0 r1)
  return (t1, r1, t1, r1, t1, n, s, a, 1)

buildTreeDualAvg lTarget glTarget g t r logu v j e t0 r0 = do
  z <- uniform g
  (tn, rn, tp, rp, t1, n1, s1, a1, na1) <-
    buildTreeDualAvg lTarget glTarget g t r logu v (pred j) e t0 r0

  if   s1 == 1
  then do
    (tnn, rnn, tpp, rpp, t2, n2, s2, a2, na2) <-
      if   v == -1
      then do
        (tnn', rnn', _, _, t1', n1', s1', a1', na1') <-
          buildTreeDualAvg lTarget glTarget g tn rn logu v (pred j) e t0 r0
        return (tnn', rnn', tp, rp, t1', n1', s1', a1', na1')
      else do
        (_, _, tpp', rpp', t1', n1', s1', a1', na1') <-
          buildTreeDualAvg lTarget glTarget g tp rp logu v (pred j) e t0 r0
        return (tn, rn, tpp', rpp', t1', n1', s1', a1', na1')

    let p      = fi n2 / max (fi (n1 + n2)) 1
        accept = p > (z :: Double)
        n3     = n1 + n2
        a3     = a1 + a2
        na3    = na1 + na2
        s3     = s1 * s2 * stopCriterion tnn tpp rnn rpp

        t3  | accept    = t2
            | otherwise = t1

    return (tnn, rnn, tpp, rpp, t3, n3, s3, a3, na3)
  else return (tn, rn, tp, rp, t1, n1, s1, a1, na1)

-- | Heuristic for initializing step size.
findReasonableEpsilon
  :: PrimMonad m
  => Density
  -> Gradient
  -> Parameters
  -> Gen (PrimState m)
  -> m Double
findReasonableEpsilon lTarget glTarget t0 g = do
  r0 <- replicateM (length t0) (normal 0 1 g)
  let (t1, r1) = leapfrog glTarget (t0, r0) 1.0
      a        = 2 * indicate (acceptanceRatio lTarget t0 t1 r0 r1 > 0.5) - 1

      go j e t r
        | j <= 0 = e -- no need to shrink this excessively
        | (acceptanceRatio lTarget t0 t r0 r) ^^ a > 2 ^^ (-a) =
            let (tn, rn) = leapfrog glTarget (t, r) e
            in  go (pred j) (2 ^^ a * e) tn rn
        | otherwise = e

  return $ go 10 1.0 t1 r1

-- | Simulate a single step of Hamiltonian dynamics.
leapfrog :: Gradient -> Particle -> Double -> Particle
leapfrog glTarget (t, r) e = (tf, rf)
  where
    rm = adjustMomentum glTarget e t r
    tf = adjustPosition e rm t
    rf = adjustMomentum glTarget e tf rm

-- | Adjust momentum.
adjustMomentum :: Fractional c => (t -> [c]) -> c -> t -> [c] -> [c]
adjustMomentum glTarget e t r = r .+ ((e / 2) .* glTarget t)

-- | Adjust position.
adjustPosition :: Num c => c -> [c] -> [c] -> [c]
adjustPosition e r t = t .+ (e .* r)

-- | The MH acceptance ratio for a given proposal.
acceptanceRatio :: Floating a => (t -> a) -> t -> t -> [a] -> [a] -> a
acceptanceRatio lTarget t0 t1 r0 r1 = auxilliaryTarget lTarget t1 r1
                                    / auxilliaryTarget lTarget t0 r0

-- | The negative potential.
auxilliaryTarget :: Floating a => (t -> a) -> t -> [a] -> a
auxilliaryTarget lTarget t r = exp (lTarget t - 0.5 * innerProduct r r)

-- | Simple inner product.
innerProduct :: Num a => [a] -> [a] -> a
innerProduct xs ys = sum $ zipWith (*) xs ys

-- | Vectorized multiplication.
(.*) :: Num b => b -> [b] -> [b]
z .* xs = map (* z) xs

-- | Vectorized subtraction.
(.-) :: Num a => [a] -> [a] -> [a]
xs .- ys = zipWith (-) xs ys

-- | Vectorized addition.
(.+) :: Num a => [a] -> [a] -> [a]
xs .+ ys = zipWith (+) xs ys

-- | Indicator function.
indicate :: Integral a => Bool -> a
indicate True  = 1
indicate False = 0

-- | A symmetric categorical (discrete uniform) distribution.
symmetricCategorical :: PrimMonad m => [a] -> Gen (PrimState m) -> m a
symmetricCategorical [] _ = error "symmetricCategorical: no candidates"
symmetricCategorical zs g = do
  j <- uniformR (0, length zs - 1) g
  return $ zs !! j

-- | Alias for fromIntegral.
fi :: (Integral a, Num b) => a -> b
fi = fromIntegral


{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE TypeOperators    #-}
module FokkerPlanck.Pinwheel
  ( module Filter.Pinwheel
  , module FokkerPlanck.Pinwheel
  ) where

import           Data.Array.Repa               as R
import           Data.Complex
import           Data.List                     as L
import           Data.Vector.Unboxed           as VU
import           Filter.Pinwheel
import           FokkerPlanck.Interpolation
import           Numeric.LinearAlgebra.Data    as NL
import           Numeric.LinearAlgebra.HMatrix
import           Types
import           Utils.Array
import           Utils.Coordinates
import           Utils.Parallel

{-# INLINE computeR2Z1T0ArrayRadial #-}
computeR2Z1T0ArrayRadial ::
     (R.Source r Double)
  => R.Array r DIM3 Double
  -> Int
  -> Int
  -> Double
  -> [Double]
  -> [Double]
  -> R.Array D DIM4 (Complex Double)
computeR2Z1T0ArrayRadial radialArr xLen yLen scaleFactor thetaFreqs theta0Freqs =
  let pinwheelArr =
        traverse2
          (fromListUnboxed (Z :. L.length thetaFreqs) thetaFreqs)
          (fromListUnboxed (Z :. L.length theta0Freqs) theta0Freqs)
          (\(Z :. numThetaFreq) (Z :. numTheta0Freq) ->
             (Z :. numThetaFreq :. numTheta0Freq :. xLen :. yLen)) $ \f f0 (Z :. k :. l :. i :. j) ->
          pinwheel
            (f (Z :. k) - f0 (Z :. l))
            0
            (exp 1)
            0
            (i - center xLen)
            (j - center yLen)
   in radialCubicInterpolation radialArr scaleFactor pinwheelArr

{-# INLINE computeR2Z2T0S0ArrayRadial #-}
computeR2Z2T0S0ArrayRadial ::
     (R.Source r Double)
  => PinwheelType
  -> R.Array r DIM5 Double
  -> Int
  -> Int
  -> Double
  -> Double
  -> [Double]
  -> [Double]
  -> [Double]
  -> [Double]
  -> R.Array D DIM6 (Complex Double)
computeR2Z2T0S0ArrayRadial pinwheelType radialArr xLen yLen scaleFactor rMax thetaFreqs scaleFreqs theta0Freqs scale0Freqs =
  let pinwheelArr =
        traverse4
          (fromListUnboxed (Z :. L.length thetaFreqs) thetaFreqs)
          (fromListUnboxed (Z :. L.length scaleFreqs) scaleFreqs)
          (fromListUnboxed (Z :. L.length theta0Freqs) theta0Freqs)
          (fromListUnboxed (Z :. L.length scale0Freqs) scale0Freqs)
          (\(Z :. numThetaFreq) (Z :. numScaleFreq) (Z :. numTheta0Freq) (Z :. numScale0Freq) ->
             (Z :. numThetaFreq :. numScaleFreq :. numTheta0Freq :.
              numScale0Freq :.
              xLen :.
              yLen)) $ \ft fs ft0 fs0 (Z :. t :. s :. t0 :. s0 :. i :. j) ->
          pinwheelFunc
            pinwheelType
            (ft (Z :. t) - ft0 (Z :. t0))
            (fs (Z :. s) - fs0 (Z :. s0))
            rMax
            0
            (i - center xLen)
            (j - center yLen)
  in radialCubicInterpolation radialArr scaleFactor pinwheelArr

{-# INLINE computeR2Z2T0S0ArrayRadial' #-}
computeR2Z2T0S0ArrayRadial' ::
     Int
  -> Int
  -> Double
  -> Double
  -> [Double]
  -> [Double]
  -> [Double]
  -> [Double]
  -> R.Array D DIM6 (Complex Double)
computeR2Z2T0S0ArrayRadial' xLen yLen scaleFactor rMax thetaFreqs scaleFreqs theta0Freqs scale0Freqs =
  let pinwheelArr =
        traverse4
          (fromListUnboxed (Z :. L.length thetaFreqs) thetaFreqs)
          (fromListUnboxed (Z :. L.length scaleFreqs) scaleFreqs)
          (fromListUnboxed (Z :. L.length theta0Freqs) theta0Freqs)
          (fromListUnboxed (Z :. L.length scale0Freqs) scale0Freqs)
          (\(Z :. numThetaFreq) (Z :. numScaleFreq) (Z :. numTheta0Freq) (Z :. numScale0Freq) ->
             (Z :. numThetaFreq :. numScaleFreq :. numTheta0Freq :.
              numScale0Freq :.
              xLen :.
              yLen)) $ \ft fs ft0 fs0 (Z :. t :. s :. t0 :. s0 :. i :. j) ->
          pinwheel
            (ft (Z :. t) - ft0 (Z :. t0))
            (fs (Z :. s) - fs0 (Z :. s0))
            rMax
            0
            (i - center xLen)
            (j - center yLen)
   in pinwheelArr

{-# INLINE cutoff #-}
cutoff :: (R.Source r Double) => Int -> R.Array r DIM5 Double -> R.Array D DIM5 Double
cutoff r arr =
  R.traverse arr id $ \f idx@(Z :. _ :. _ :. _ :. _ :. e) ->
    if e > r
      then 0
      else f idx

{-# INLINE computeLocalEigenVector #-}
computeLocalEigenVector ::
     (R.Source r Double)
  => ParallelParams
  -> (Double -> Double -> Double -> Double -> Int -> Int -> Complex Double)
  -> R.Array r DIM5 Double
  -> Int
  -> Int
  -> Double
  -> [Double]
  -> [Double]
  -> R.Array U DIM4 (Complex Double)
computeLocalEigenVector parallelParams pinwheelFunc radialArr xLen yLen rMax thetaFreqs scaleFreqs =
  let (Z :. numThetaFreq :. numScaleFreq :. _ :. _ :. _) = extent radialArr
      len = numThetaFreq * numScaleFreq
   in computeS .
      rotate4D .
      rotate4D .
      fromUnboxed (Z :. xLen :. yLen :. numThetaFreq :. numScaleFreq) .
      VU.concat .
      parMapChunk
        parallelParams
        rdeepseq
        (\(x, y) ->
           let r =
                 sqrt $
                 (fromIntegral $ x - center xLen) ^ 2 +
                 (fromIntegral $ y - center yLen) ^ 2
               interpolatedArray = cubicInterpolation radialArr r
               interpolatedPinwheel =
                 R.traverse3
                   interpolatedArray
                   (fromListUnboxed (Z :. numThetaFreq) thetaFreqs)
                   (fromListUnboxed (Z :. numScaleFreq) scaleFreqs)
                   (\sh _ _ -> sh) $ \f ft fs idx@(Z :. t :. s :. t0 :. s0) ->
                   (f idx :+ 0) *
                   pinwheelFunc
                     (ft (Z :. t) - ft (Z :. t0))
                     (fs (Z :. s) + fs (Z :. s0))
                     rMax
                     0
                     (x - center xLen)
                     (y - center yLen)
               (eigVal, eigVec) =
                 eig . (len >< len) . R.toList $ interpolatedPinwheel
               (val, vec) =
                 L.head . L.reverse . L.sortOn (magnitude . fst) $
                 L.zip (NL.toList eigVal) (toColumns $ eigVec)
               dominantVec = VU.fromList . L.map (* val) . NL.toList $ vec
            in -- VU.zipWith
               --   (*)
               --   (VU.fromList
               --      [ exp
               --        (0 :+
               --         (tf) *
               --         (angleFunctionRad
               --            (fromIntegral $ x - center xLen)
               --            (fromIntegral $ y - center yLen)))
               --      | tf <- thetaFreqs
               --      , sf <- scaleFreqs
               --      ])
               --   dominantVec
               dominantVec
        ) $
      [(x, y) | x <- [0 .. xLen - 1], y <- [0 .. yLen - 1]]

{-# INLINE computeLocalEigenVectorSink #-}
computeLocalEigenVectorSink ::
     (R.Source r Double)
  => ParallelParams
  -> (Double -> Double -> Double -> Double -> Int -> Int -> Complex Double)
  -> R.Array r DIM5 Double
  -> Int
  -> Int
  -> Double
  -> [Double]
  -> [Double]
  -> R.Array U DIM4 (Complex Double)
computeLocalEigenVectorSink parallelParams pinwheelFunc radialArr xLen yLen rMax thetaFreqs scaleFreqs =
  let (Z :. numThetaFreq :. numScaleFreq :. _ :. _ :. _) = extent radialArr
      len = numThetaFreq * numScaleFreq
   in computeS .
      rotate4D .
      rotate4D .
      fromUnboxed (Z :. xLen :. yLen :. numThetaFreq :. numScaleFreq) .
      VU.concat .
      parMapChunk
        parallelParams
        rdeepseq
        (\(x, y) ->
           let r =
                 sqrt $
                 (fromIntegral $ x - center xLen) ^ 2 +
                 (fromIntegral $ y - center yLen) ^ 2
               interpolatedArray = cubicInterpolation radialArr r
               interpolatedPinwheel =
                 R.traverse3
                   interpolatedArray
                   (fromListUnboxed (Z :. numThetaFreq) thetaFreqs)
                   (fromListUnboxed (Z :. numScaleFreq) scaleFreqs)
                   (\sh _ _ -> sh) $ \f ft fs idx@(Z :. t :. s :. t0 :. s0) ->
                   (f idx :+ 0) * (exp (0 :+ (ft (Z :. t) - ft (Z :. t0)) * pi)) *
                   pinwheelFunc
                     (ft (Z :. t) - ft (Z :. t0))
                     (fs (Z :. s) + fs (Z :. s0))
                     rMax
                     0
                     (x - center xLen)
                     (y - center yLen)
               (eigVal, eigVec) =
                 eig . (len >< len) . R.toList $ interpolatedPinwheel
               (val, vec) =
                 L.head . L.reverse . L.sortOn (realPart . fst) $
                 L.zip (NL.toList eigVal) (toColumns $ eigVec)
            in VU.fromList . L.map (* val) . NL.toList $ vec) $
      [(x, y) | x <- [0 .. xLen - 1], y <- [0 .. yLen - 1]]


{-# LANGUAGE RecordWildCards #-}
module DFA where

import qualified Data.Set as Set
import qualified Data.Map.Strict as Map
import Control.Monad.State
import Data.Maybe
import Data.Semigroup (stimes)
import Numeric.LinearAlgebra hiding (remap, (<>))
import NFA
import Utils
import qualified Partition as P
import LruCache

-- | The DFA type is parameterized on the state type, so that we can use
-- Sets during the subset construction and then map them to integers for
-- better efficiency.
data DFA s c = DFA
  { dfaStart :: s
  , dfaFinal :: Set.Set s
  , dfaStates :: Set.Set s
  , dfaTransitions :: Set.Set (c, s, s)
  }
  deriving Show

-- | At first, the DFA states are sets of NFA states
type DfaState1 = Set.Set NfaState

-- | Then, the sets are mapped into integers
type DfaState2 = Int

----------------------------------------------------------------------
--                      The subset construction
----------------------------------------------------------------------

-- | Return all epsilon-edges from an NFA state
epsilonEdges :: Ord c => NFA (Maybe c) -> NfaState -> Set.Set NfaState
epsilonEdges NFA{nfaTransitions} s =
  fromMaybe mempty . Map.lookup Nothing . fromMaybe mempty . Map.lookup s $
    transitionMap nfaTransitions

-- | Compute an epsilon-closure
epsilonClosure :: Ord c => NFA (Maybe c) -> Set.Set NfaState -> Set.Set NfaState
epsilonClosure nfa = go mempty
  where
    go !seen0 !new0
      | Set.null new0 = seen0
      | otherwise =
      let
        seen1 = Set.union seen0 new0
        new1 = Set.difference (sconcatMap (epsilonEdges nfa) new0) seen0
      in go seen1 new1

-- | @move nfa s a@ is a set of NFA states where we can get from one of the
-- states in @s@ after receiving @a@ as an input
move :: Ord c => NFA (Maybe c) -> DfaState1 -> c -> DfaState1
move NFA{nfaTransitions} ss c = sconcatMap
  (\s -> fromMaybe mempty . Map.lookup (Just c) .
    fromMaybe mempty . Map.lookup s $
    transitionMap nfaTransitions)
  ss

-- | Build a DFA from an NFA
nfaToDfa :: forall c . Ord c => [c] -> NFA (Maybe c) -> DFA DfaState1 c
nfaToDfa alphabet nfa =
  let
    start :: DfaState1
    start = epsilonClosure nfa (nfaStart nfa)
    transitions = go Set.empty (Set.singleton $ start) Set.empty
    all_states = concat [ [s1,s2] | (_, s1, s2) <- Set.toList transitions ]
    contains_final :: DfaState1 -> Bool
    contains_final dfa_st = not . Set.null $ Set.intersection (nfaFinal nfa) dfa_st
  in
    DFA
      { dfaStart = start
      , dfaFinal = Set.fromList $ filter contains_final all_states
      , dfaStates = Set.fromList all_states
      , dfaTransitions = transitions
      }
  where
    go
      :: Set.Set DfaState1
      -> Set.Set DfaState1
      -> Set.Set (c, DfaState1, DfaState1)
      -> Set.Set (c, DfaState1, DfaState1)
    go !seen0 !new0 !transitions0
      | Set.null new0 = transitions0
      | otherwise =
      let
        seen1 = Set.union seen0 new0
        -- where can we get from the DFA states in new1?
        (transitions1, new1) =
          flip foldMap new0 $ \s ->
          flip foldMap alphabet $ \c ->
          let
            next :: DfaState1
            next = epsilonClosure nfa (move nfa s c)
          in (Set.singleton (c, s, next), Set.singleton next)

        transitions2 = transitions0 `Set.union` transitions1
        new2 = new1 `Set.difference` seen1
      in go seen1 new2 transitions2

----------------------------------------------------------------------
--                         DFA optimization
----------------------------------------------------------------------

-- | Map all DFA states
mapStates
  :: (Ord s1, Ord s2, Ord c)
  => (s1 -> s2)
  -> DFA s1 c
  -> DFA s2 c
mapStates remap DFA{..} =
  DFA
  { dfaStart = remap dfaStart
  , dfaFinal = Set.map remap dfaFinal
  , dfaStates = Set.map remap dfaStates
  , dfaTransitions = Set.map (\(c,s0,s1) -> (c, remap s0, remap s1)) dfaTransitions
  }

-- | Convert all states of the DFA to consecutive integers starting from 0
mapStatesToInt
  :: (Ord s, Ord c)
  => DFA s c
  -> DFA DfaState2 c
mapStatesToInt dfa =
  let
    stateMap = Map.fromList $ zip (Set.toList . dfaStates $ dfa) [0..]
    remap = (stateMap Map.!)
  in
    mapStates remap dfa

minimizeDfa
  :: forall s c . (Ord c, Ord s)
  => DFA s c
  -> DFA Int c
minimizeDfa dfa@DFA{..} = flip evalState 0 $ do
  initialPartition <- liftM2 (<>)
    (P.singleton dfaFinal)
    (P.singleton $ dfaStates `Set.difference` dfaFinal)

  let
    transitionMap :: Map.Map s (Map.Map c s)
    transitionMap = Map.fromListWith Map.union
      [ (s0, Map.singleton c s1)
      | (c, s0, s1) <- Set.toList dfaTransitions
      ]

    classify :: P.Partition s -> s -> Map.Map c Int
    classify pt s = fmap (P.lookupState pt) $ transitionMap Map.! s

  finalPartition <- P.subpartitionFixpoint classify initialPartition

  return $ mapStates (P.lookupState finalPartition) dfa

----------------------------------------------------------------------
--                Transfer matrices and probabilities
----------------------------------------------------------------------

data TransferMatrix = TransferMatrix
  { tmMatrix :: Matrix Double
    -- ^ transfer matrix itself
  , tmStart :: Vector Double
    -- ^ the indicator vector of the start state
  , tmFinal :: Vector Double
    -- ^ the indicator vector of the final states
  }
  deriving (Show)

-- | Construct the DFA transfer matrix assuming the uniform distribution
-- over nucleotides.
--
-- Return the matrix and the 0-based indices of the start and end states.
dfaTransferMatrix
  :: (Ord c, Ord s)
  => Map.Map c Double -- ^ character probabilities
  -> DFA s c
  -> TransferMatrix
dfaTransferMatrix freqs (mapStatesToInt -> DFA{..}) =
  -- we remap the DFA once again to make sure the numbers are consecutive
  let
    n = Set.size dfaStates
    elts =
      [ ((s0, s1), freqs Map.! c)
      | (c, s0, s1) <- Set.toList dfaTransitions
      ]
  in
    TransferMatrix
    { tmMatrix = (assoc (n,n) 0 . assocSum) elts
    , tmStart = assoc n 0 [ (dfaStart, 1) ]
    , tmFinal = assoc n 0 [ (st, 1) | st <- Set.toList dfaFinal ]
    }
  where
    assocSum = Map.toList . Map.fromListWith (+)

dfaProbability
  :: TransferMatrix
  -> Int -- ^ length of the random string where they motif may occur
  -> Double
dfaProbability TransferMatrix{..} len =
  (tmStart <# stimes len tmMatrix) <.> tmFinal

dfaProbabilityCached
  :: MonadLru Int (Matrix Double) m
  => TransferMatrix
  -> Int -- ^ length of the random string where they motif may occur
  -> m Double
dfaProbabilityCached TransferMatrix{..} len = do
  pw <- lruPow len tmMatrix
  return $ (tmStart <# pw) <.> tmFinal