content: {-# LANGUAGE OverloadedStrings #-}
  
  module DataFrame.Operations.Join where
  
  import qualified Data.Map.Strict as M
  import qualified Data.Text as T
  import qualified Data.Vector as VB
  import qualified Data.Vector.Unboxed as VU
  import DataFrame.Internal.Column as D
  import DataFrame.Internal.DataFrame as D
  import DataFrame.Operations.Aggregation as D
  import DataFrame.Operations.Core as D
  
  -- | Equivalent to SQL join types.
  data JoinType
      = INNER
      | LEFT
      | RIGHT
      | FULL_OUTER
  
  {- | Join two dataframes using SQL join semantics.
  
  Only inner join is implemented for now.
  -}
  join ::
      JoinType ->
      [T.Text] ->
      DataFrame -> -- Right hand side
      DataFrame -> -- Left hand side
      DataFrame
  join INNER xs right = innerJoin xs right
  join LEFT xs right = error "UNIMPLEMENTED"
  join RIGHT xs right = error "UNIMPLEMENTED"
  join FULL_OUTER xs right = error "UNIMPLEMENTED"
  
  {- | Inner join of two dataframes. Note: for chaining, the left dataframe is actually
  on the right side.
  -}
  innerJoin :: [T.Text] -> DataFrame -> DataFrame -> DataFrame
  innerJoin cs right left =
      let
          leftIndicesToGroup = M.elems $ M.filterWithKey (\k _ -> k `elem` cs) (D.columnIndices left)
          leftRowRepresentations = VU.generate (fst (D.dimensions left)) (D.mkRowRep leftIndicesToGroup left)
          -- key -> [index0, index1]
          leftKeyCountsAndIndices =
              VU.foldr
                  (\(i, v) acc -> M.insertWith (++) v [i] acc)
                  M.empty
                  (VU.indexed leftRowRepresentations)
          -- key -> [index0, index1]
          rightIndicesToGroup = M.elems $ M.filterWithKey (\k _ -> k `elem` cs) (D.columnIndices right)
          rightRowRepresentations = VU.generate (fst (D.dimensions right)) (D.mkRowRep rightIndicesToGroup right)
          rightKeyCountsAndIndices =
              VU.foldr
                  (\(i, v) acc -> M.insertWith (++) v [i] acc)
                  M.empty
                  (VU.indexed rightRowRepresentations)
          -- key -> [(left_indexes0, right_indexes1)]
          mergedKeyCountsAndIndices =
              M.foldrWithKey
                  ( \k v m ->
                      if k `M.member` rightKeyCountsAndIndices
                          then M.insert k (VU.fromList v, VU.fromList (rightKeyCountsAndIndices M.! k)) m
                          else m
                  )
                  M.empty
                  leftKeyCountsAndIndices
          -- [(ints, ints)]
          leftAndRightIndicies = M.elems mergedKeyCountsAndIndices
          -- [(ints, ints)] (expanded to n * m)
          expandedIndices =
              map
                  ( \(l, r) -> (mconcat (replicate (VU.length r) l), mconcat (replicate (VU.length l) r))
                  )
                  leftAndRightIndicies
          expandedLeftIndicies = mconcat (map fst expandedIndices)
          expandedRightIndicies = mconcat (map snd expandedIndices)
          -- df
          expandedLeft =
              left
                  { columns = VB.map (D.atIndicesStable expandedLeftIndicies) (D.columns left)
                  , dataframeDimensions =
                      (VU.length expandedLeftIndicies, snd (D.dataframeDimensions left))
                  }
          -- df
          expandedRight =
              right
                  { columns = VB.map (D.atIndicesStable expandedRightIndicies) (D.columns right)
                  , dataframeDimensions =
                      (VU.length expandedRightIndicies, snd (D.dataframeDimensions right))
                  }
          -- [string]
          leftColumns = D.columnNames left
          rightColumns = D.columnNames right
          initDf = expandedLeft
          insertIfPresent _ Nothing df = df
          insertIfPresent name (Just c) df = D.insertColumn name c df
       in
          D.fold
              ( \name df ->
                  if name `elem` cs
                      then df
                      else
                          ( if name `elem` leftColumns
                              then insertIfPresent ("Right_" <> name) (D.getColumn name expandedRight) df
                              else insertIfPresent name (D.getColumn name expandedRight) df
                          )
              )
              rightColumns
              initDf
  