content: {-|
  Module: Squeal.PostgreSQL.Expression.Aggregate
  Description: aggregate functions and arguments
  Copyright: (c) Eitan Chatav, 2019
  Maintainer: eitan@morphism.tech
  Stability: experimental
  
  aggregate functions and arguments
  -}
  
  {-# LANGUAGE
      DataKinds
    , DeriveGeneric
    , FlexibleContexts
    , FlexibleInstances
    , FunctionalDependencies
    , LambdaCase
    , MultiParamTypeClasses
    , OverloadedStrings
    , PatternSynonyms
    , PolyKinds
    , ScopedTypeVariables
    , StandaloneDeriving
    , TypeApplications
    , TypeFamilies
    , TypeOperators
    , UndecidableInstances
  #-}
  
  module Squeal.PostgreSQL.Expression.Aggregate
    ( -- * Aggregate
      Aggregate (..)
      -- * Aggregate Arguments
    , AggregateArg (..)
    , pattern All
    , pattern Alls
    , allNotNull
    , pattern Distinct
    , pattern Distincts
    , distinctNotNull
    , FilterWhere (..)
      -- * Aggregate Types
    , PGSum
    , PGAvg
    ) where
  
  import Data.ByteString (ByteString)
  import GHC.TypeLits
  
  import qualified Generics.SOP as SOP
  
  import Squeal.PostgreSQL.Type.Alias
  import Squeal.PostgreSQL.Expression
  import Squeal.PostgreSQL.Expression.Logic
  import Squeal.PostgreSQL.Expression.Null
  import Squeal.PostgreSQL.Expression.Sort
  import Squeal.PostgreSQL.Type.List
  import Squeal.PostgreSQL.Render
  import Squeal.PostgreSQL.Type.Schema
  
  -- $setup
  -- >>> import Squeal.PostgreSQL
  
  {- |
  `Aggregate` functions compute a single result from a set of input values.
  `Aggregate` functions can be used as `Grouped` `Expression`s as well
  as `Squeal.PostgreSQL.Expression.Window.WindowFunction`s.
  -}
  class Aggregate arg expr | expr -> arg where
  
    -- | A special aggregation that does not require an input
    --
    -- >>> :{
    -- let
    --   expression :: Expression ('Grouped bys) '[] with db params from ('NotNull 'PGint8)
    --   expression = countStar
    -- in printSQL expression
    -- :}
    -- count(*)
    countStar :: expr lat with db params from ('NotNull 'PGint8)
  
    -- | >>> :{
    -- let
    --   expression :: Expression ('Grouped bys) '[] with db params '[tab ::: '["col" ::: null ty]] ('NotNull 'PGint8)
    --   expression = count (All #col)
    -- in printSQL expression
    -- :}
    -- count(ALL "col")
    count
      :: arg '[ty] lat with db params from
      -- ^ argument
      -> expr lat with db params from ('NotNull 'PGint8)
  
    -- | >>> :{
    -- let
    --   expression :: Expression ('Grouped bys) '[] with db params '[tab ::: '["col" ::: 'Null 'PGnumeric]] ('Null 'PGnumeric)
    --   expression = sum_ (Distinct #col & filterWhere (#col .< 100))
    -- in printSQL expression
    -- :}
    -- sum(DISTINCT "col") FILTER (WHERE ("col" < (100.0 :: numeric)))
    sum_
      :: arg '[null ty] lat with db params from
      -- ^ argument
      -> expr lat with db params from ('Null (PGSum ty))
  
    -- | input values, including nulls, concatenated into an array
    --
    -- >>> :{
    -- let
    --   expression :: Expression ('Grouped bys) '[] with db params '[tab ::: '["col" ::: 'Null 'PGnumeric]] ('Null ('PGvararray ('Null 'PGnumeric)))
    --   expression = arrayAgg (All #col & orderBy [AscNullsFirst #col] & filterWhere (#col .< 100))
    -- in printSQL expression
    -- :}
    -- array_agg(ALL "col" ORDER BY "col" ASC NULLS FIRST) FILTER (WHERE ("col" < (100.0 :: numeric)))
    arrayAgg
      :: arg '[ty] lat with db params from
      -- ^ argument
      -> expr lat with db params from ('Null ('PGvararray ty))
  
    -- | aggregates values as a JSON array
    jsonAgg
      :: arg '[ty] lat with db params from
      -- ^ argument
      -> expr lat with db params from ('Null 'PGjson)
  
    -- | aggregates values as a JSON array
    jsonbAgg
      :: arg '[ty] lat with db params from
      -- ^ argument
      -> expr lat with db params from ('Null 'PGjsonb)
  
    {- |
    the bitwise AND of all non-null input values, or null if none
  
    >>> :{
    let
      expression :: Expression ('Grouped bys) '[] with db params '[tab ::: '["col" ::: null 'PGint4]] ('Null 'PGint4)
      expression = bitAnd (Distinct #col)
    in printSQL expression
    :}
    bit_and(DISTINCT "col")
    -}
    bitAnd
      :: int `In` PGIntegral
      => arg '[null int] lat with db params from
      -- ^ argument
      -> expr lat with db params from ('Null int)
  
    {- |
    the bitwise OR of all non-null input values, or null if none
  
    >>> :{
    let
      expression :: Expression ('Grouped bys) '[] with db params '[tab ::: '["col" ::: null 'PGint4]] ('Null 'PGint4)
      expression = bitOr (All #col)
    in printSQL expression
    :}
    bit_or(ALL "col")
    -}
    bitOr
      :: int `In` PGIntegral
      => arg '[null int] lat with db params from
      -- ^ argument
      -> expr lat with db params from ('Null int)
  
    {- |
    true if all input values are true, otherwise false
  
    >>> :{
    let
      winFun :: WindowFunction  'Ungrouped '[] with db params '[tab ::: '["col" ::: null 'PGbool]] ('Null 'PGbool)
      winFun = boolAnd (Window #col)
    in printSQL winFun
    :}
    bool_and("col")
    -}
    boolAnd
      :: arg '[null 'PGbool] lat with db params from
      -- ^ argument
      -> expr lat with db params from ('Null 'PGbool)
  
    {- |
    true if at least one input value is true, otherwise false
  
    >>> :{
    let
      expression :: Expression ('Grouped bys) '[] with db params '[tab ::: '["col" ::: null 'PGbool]] ('Null 'PGbool)
      expression = boolOr (All #col)
    in printSQL expression
    :}
    bool_or(ALL "col")
    -}
    boolOr
      :: arg '[null 'PGbool] lat with db params from
      -- ^ argument
      -> expr lat with db params from ('Null 'PGbool)
  
    {- |
    equivalent to `boolAnd`
  
    >>> :{
    let
      expression :: Expression ('Grouped bys) '[] with db params '[tab ::: '["col" ::: null 'PGbool]] ('Null 'PGbool)
      expression = every (Distinct #col)
    in printSQL expression
    :}
    every(DISTINCT "col")
    -}
    every
      :: arg '[null 'PGbool] lat with db params from
      -- ^ argument
      -> expr lat with db params from ('Null 'PGbool)
  
    {- |maximum value of expression across all input values-}
    max_
      :: arg '[null ty] lat with db params from
      -- ^ argument
      -> expr lat with db params from ('Null ty)
  
    -- | minimum value of expression across all input values
    min_
      :: arg '[null ty] lat with db params from
      -- ^ argument
      -> expr lat with db params from ('Null ty)
  
    -- | the average (arithmetic mean) of all input values
    avg
      :: arg '[null ty] lat with db params from
      -- ^ argument
      -> expr lat with db params from ('Null (PGAvg ty))
  
    {- | correlation coefficient
  
    >>> :{
    let
      expression :: Expression ('Grouped g) '[] c s p '[t ::: '["x" ::: 'NotNull 'PGfloat8, "y" ::: 'NotNull 'PGfloat8]] ('Null 'PGfloat8)
      expression = corr (Alls (#y *: #x))
    in printSQL expression
    :}
    corr(ALL "y", "x")
    -}
    corr
      :: arg '[null 'PGfloat8, null 'PGfloat8] lat with db params from
      -- ^ arguments
      -> expr lat with db params from ('Null 'PGfloat8)
  
    {- | population covariance
  
    >>> :{
    let
      expression :: Expression ('Grouped g) '[] c s p '[t ::: '["x" ::: 'NotNull 'PGfloat8, "y" ::: 'NotNull 'PGfloat8]] ('Null 'PGfloat8)
      expression = covarPop (Alls (#y *: #x))
    in printSQL expression
    :}
    covar_pop(ALL "y", "x")
    -}
    covarPop
      :: arg '[null 'PGfloat8, null 'PGfloat8] lat with db params from
      -- ^ arguments
      -> expr lat with db params from ('Null 'PGfloat8)
  
    {- | sample covariance
  
    >>> :{
    let
      winFun :: WindowFunction  'Ungrouped '[] c s p '[t ::: '["x" ::: 'NotNull 'PGfloat8, "y" ::: 'NotNull 'PGfloat8]] ('Null 'PGfloat8)
      winFun = covarSamp (Windows (#y *: #x))
    in printSQL winFun
    :}
    covar_samp("y", "x")
    -}
    covarSamp
      :: arg '[null 'PGfloat8, null 'PGfloat8] lat with db params from
      -- ^ arguments
      -> expr lat with db params from ('Null 'PGfloat8)
  
    {- | average of the independent variable (sum(X)/N)
  
    >>> :{
    let
      expression :: Expression ('Grouped g) '[] c s p '[t ::: '["x" ::: 'NotNull 'PGfloat8, "y" ::: 'NotNull 'PGfloat8]] ('Null 'PGfloat8)
      expression = regrAvgX (Alls (#y *: #x))
    in printSQL expression
    :}
    regr_avgx(ALL "y", "x")
    -}
    regrAvgX
      :: arg '[null 'PGfloat8, null 'PGfloat8] lat with db params from
      -- ^ arguments
      -> expr lat with db params from ('Null 'PGfloat8)
  
    {- | average of the dependent variable (sum(Y)/N)
  
    >>> :{
    let
      winFun :: WindowFunction  'Ungrouped '[] c s p '[t ::: '["x" ::: 'NotNull 'PGfloat8, "y" ::: 'NotNull 'PGfloat8]] ('Null 'PGfloat8)
      winFun = regrAvgY (Windows (#y *: #x))
    in printSQL winFun
    :}
    regr_avgy("y", "x")
    -}
    regrAvgY
      :: arg '[null 'PGfloat8, null 'PGfloat8] lat with db params from
      -- ^ arguments
      -> expr lat with db params from ('Null 'PGfloat8)
  
    {- | number of input rows in which both expressions are nonnull
  
    >>> :{
    let
      winFun :: WindowFunction  'Ungrouped '[] c s p '[t ::: '["x" ::: 'NotNull 'PGfloat8, "y" ::: 'NotNull 'PGfloat8]] ('Null 'PGint8)
      winFun = regrCount (Windows (#y *: #x))
    in printSQL winFun
    :}
    regr_count("y", "x")
    -}
    regrCount
      :: arg '[null 'PGfloat8, null 'PGfloat8] lat with db params from
      -- ^ arguments
      -> expr lat with db params from ('Null 'PGint8)
  
    {- | y-intercept of the least-squares-fit linear equation determined by the (X, Y) pairs
  
    >>> :{
    let
      expression :: Expression ('Grouped g) '[] c s p '[t ::: '["x" ::: 'NotNull 'PGfloat8, "y" ::: 'NotNull 'PGfloat8]] ('Null 'PGfloat8)
      expression = regrIntercept (Alls (#y *: #x))
    in printSQL expression
    :}
    regr_intercept(ALL "y", "x")
    -}
    regrIntercept
      :: arg '[null 'PGfloat8, null 'PGfloat8] lat with db params from
      -- ^ arguments
      -> expr lat with db params from ('Null 'PGfloat8)
  
    -- | @regr_r2(Y, X)@, square of the correlation coefficient
    regrR2
      :: arg '[null 'PGfloat8, null 'PGfloat8] lat with db params from
      -- ^ arguments
      -> expr lat with db params from ('Null 'PGfloat8)
  
    -- | @regr_slope(Y, X)@, slope of the least-squares-fit linear equation
    -- determined by the (X, Y) pairs
    regrSlope
      :: arg '[null 'PGfloat8, null 'PGfloat8] lat with db params from
      -- ^ arguments
      -> expr lat with db params from ('Null 'PGfloat8)
  
    -- | @regr_sxx(Y, X)@, sum(X^2) - sum(X)^2/N
    -- (“sum of squares” of the independent variable)
    regrSxx
      :: arg '[null 'PGfloat8, null 'PGfloat8] lat with db params from
      -- ^ arguments
      -> expr lat with db params from ('Null 'PGfloat8)
  
    -- | @regr_sxy(Y, X)@, sum(X*Y) - sum(X) * sum(Y)/N
    -- (“sum of products” of independent times dependent variable)
    regrSxy
      :: arg '[null 'PGfloat8, null 'PGfloat8] lat with db params from
      -- ^ arguments
      -> expr lat with db params from ('Null 'PGfloat8)
  
    -- | @regr_syy(Y, X)@, sum(Y^2) - sum(Y)^2/N
    -- (“sum of squares” of the dependent variable)
    regrSyy
      :: arg '[null 'PGfloat8, null 'PGfloat8] lat with db params from
      -- ^ arguments
      -> expr lat with db params from ('Null 'PGfloat8)
  
    -- | historical alias for `stddevSamp`
    stddev
      :: arg '[null ty] lat with db params from
      -- ^ argument
      -> expr lat with db params from ('Null (PGAvg ty))
  
    -- | population standard deviation of the input values
    stddevPop
      :: arg '[null ty] lat with db params from
      -- ^ argument
      -> expr lat with db params from ('Null (PGAvg ty))
  
    -- | sample standard deviation of the input values
    stddevSamp
      :: arg '[null ty] lat with db params from
      -- ^ argument
      -> expr lat with db params from ('Null (PGAvg ty))
  
    -- | historical alias for `varSamp`
    variance
      :: arg '[null ty] lat with db params from
      -- ^ argument
      -> expr lat with db params from ('Null (PGAvg ty))
  
    -- | population variance of the input values
    -- (square of the population standard deviation)
    varPop
      :: arg '[null ty] lat with db params from
      -- ^ argument
      -> expr lat with db params from ('Null (PGAvg ty))
  
    -- | sample variance of the input values
    -- (square of the sample standard deviation)
    varSamp
      :: arg '[null ty] lat with db params from
      -- ^ argument
      -> expr lat with db params from ('Null (PGAvg ty))
  
  {- |
  `AggregateArg`s are used for the input of `Aggregate` `Expression`s.
  -}
  data AggregateArg
    (xs :: [NullType])
    (lat :: FromType)
    (with :: FromType)
    (db :: SchemasType)
    (params :: [NullType])
    (from :: FromType)
    = AggregateAll
    { aggregateArgs :: NP (Expression 'Ungrouped lat with db params from) xs
    , aggregateOrder :: [SortExpression 'Ungrouped lat with db params from]
      -- ^ `orderBy`
    , aggregateFilter :: [Condition 'Ungrouped lat with db params from]
      -- ^ `filterWhere`
    }
    | AggregateDistinct
    { aggregateArgs :: NP (Expression 'Ungrouped lat with db params from) xs
    , aggregateOrder :: [SortExpression 'Ungrouped lat with db params from]
      -- ^ `orderBy`
    , aggregateFilter :: [Condition 'Ungrouped lat with db params from]
      -- ^ `filterWhere`
    }
  
  instance (HasUnique tab (Join from lat) row, Has col row ty)
    => IsLabel col (AggregateArg '[ty] lat with db params from) where
      fromLabel = All (fromLabel @col)
  instance (Has tab (Join from lat) row, Has col row ty)
    => IsQualified tab col (AggregateArg '[ty] lat with db params from) where
      tab ! col = All (tab ! col)
  
  instance SOP.SListI xs => RenderSQL (AggregateArg xs lat with db params from) where
    renderSQL = \case
      AggregateAll args sorts filters ->
        parenthesized
        ("ALL" <+> renderCommaSeparated renderSQL args<> renderSQL sorts)
        <> renderFilters filters
      AggregateDistinct args sorts filters ->
        parenthesized
        ("DISTINCT" <+> renderCommaSeparated renderSQL args <> renderSQL sorts)
        <> renderFilters filters
      where
        renderFilter wh = "FILTER" <+> parenthesized ("WHERE" <+> wh)
        renderFilters = \case
          [] -> ""
          wh:whs -> " " <> renderFilter (renderSQL (foldr (.&&) wh whs))
  
  instance OrderBy (AggregateArg xs) 'Ungrouped where
    orderBy sorts1 = \case
      AggregateAll xs sorts0 whs -> AggregateAll xs (sorts0 ++ sorts1) whs
      AggregateDistinct xs sorts0 whs -> AggregateDistinct xs (sorts0 ++ sorts1) whs
  
  -- | `All` invokes the aggregate on a single
  -- argument once for each input row.
  pattern All
    :: Expression 'Ungrouped lat with db params from x
    -- ^ argument
    -> AggregateArg '[x] lat with db params from
  pattern All x = Alls (x :* Nil)
  
  -- | `All` invokes the aggregate on multiple
  -- arguments once for each input row.
  pattern Alls
    :: NP (Expression 'Ungrouped lat with db params from) xs
    -- ^ arguments
    -> AggregateArg xs lat with db params from
  pattern Alls xs = AggregateAll xs [] []
  
  -- | `allNotNull` invokes the aggregate on a single
  -- argument once for each input row where the argument
  -- is not null
  allNotNull
    :: Expression 'Ungrouped lat with db params from ('Null x)
    -- ^ argument
    -> AggregateArg '[ 'NotNull x] lat with db params from
  allNotNull x = All (unsafeNotNull x) & filterWhere (not_ (isNull x))
  
  {- |
  `Distinct` invokes the aggregate once for each
  distinct value of the expression found in the input.
  -}
  pattern Distinct
    :: Expression 'Ungrouped lat with db params from x
    -- ^ argument
    -> AggregateArg '[x] lat with db params from
  pattern Distinct x = Distincts (x :* Nil)
  
  {- |
  `Distincts` invokes the aggregate once for each
  distinct set of values, for multiple expressions, found in the input.
  -}
  pattern Distincts
    :: NP (Expression 'Ungrouped lat with db params from) xs
    -- ^ arguments
    -> AggregateArg xs lat with db params from
  pattern Distincts xs = AggregateDistinct xs [] []
  
  {- |
  `distinctNotNull` invokes the aggregate once for each
  distinct, not null value of the expression found in the input.
  -}
  distinctNotNull
    :: Expression 'Ungrouped lat with db params from ('Null x)
    -- ^ argument
    -> AggregateArg '[ 'NotNull x] lat with db params from
  distinctNotNull x = Distinct (unsafeNotNull x) & filterWhere (not_ (isNull x))
  
  -- | Permits filtering
  -- `Squeal.PostgreSQL.Expression.Window.WindowArg`s and `AggregateArg`s
  class FilterWhere arg grp | arg -> grp where
    {- |
    If `filterWhere` is specified, then only the input rows for which
    the `Condition` evaluates to true are fed to the aggregate function;
    other rows are discarded.
    -}
    filterWhere
      :: Condition grp lat with db params from
      -- ^ include rows which evaluate to true
      -> arg xs lat with db params from
      -> arg xs lat with db params from
  instance FilterWhere AggregateArg 'Ungrouped where
    filterWhere wh = \case
      AggregateAll xs sorts whs -> AggregateAll xs sorts (wh : whs)
      AggregateDistinct xs sorts whs -> AggregateDistinct xs sorts (wh : whs)
  
  instance Aggregate AggregateArg (Expression ('Grouped bys)) where
    countStar = UnsafeExpression "count(*)"
    count = unsafeAggregate "count"
    sum_ = unsafeAggregate "sum"
    arrayAgg = unsafeAggregate "array_agg"
    jsonAgg = unsafeAggregate "json_agg"
    jsonbAgg = unsafeAggregate "jsonb_agg"
    bitAnd = unsafeAggregate "bit_and"
    bitOr = unsafeAggregate "bit_or"
    boolAnd = unsafeAggregate "bool_and"
    boolOr = unsafeAggregate "bool_or"
    every = unsafeAggregate "every"
    max_ = unsafeAggregate "max"
    min_ = unsafeAggregate "min"
    avg = unsafeAggregate "avg"
    corr = unsafeAggregate "corr"
    covarPop = unsafeAggregate "covar_pop"
    covarSamp = unsafeAggregate "covar_samp"
    regrAvgX = unsafeAggregate "regr_avgx"
    regrAvgY = unsafeAggregate "regr_avgy"
    regrCount = unsafeAggregate "regr_count"
    regrIntercept = unsafeAggregate "regr_intercept"
    regrR2 = unsafeAggregate "regr_r2"
    regrSlope = unsafeAggregate "regr_slope"
    regrSxx = unsafeAggregate "regr_sxx"
    regrSxy = unsafeAggregate "regr_sxy"
    regrSyy = unsafeAggregate "regr_syy"
    stddev = unsafeAggregate "stddev"
    stddevPop = unsafeAggregate "stddev_pop"
    stddevSamp = unsafeAggregate "stddev_samp"
    variance = unsafeAggregate "variance"
    varPop = unsafeAggregate "var_pop"
    varSamp = unsafeAggregate "var_samp"
  
  -- provides a nicer type error when we forget to group by
  -- note that we need to make our 'a' polymorphic so that we can still match when it's ambiguous
  instance ( TypeError ('Text "Cannot use aggregate functions to construct an Ungrouped Expression. Add a 'groupBy' to your TableExpression. If you want to aggregate across the entire result set, use 'groupBy Nil'.")
           , a ~ AggregateArg
           ) => Aggregate a (Expression 'Ungrouped) where
    countStar = impossibleAggregateError
    count = impossibleAggregateError
    sum_ = impossibleAggregateError
    arrayAgg = impossibleAggregateError
    jsonAgg = impossibleAggregateError
    jsonbAgg = impossibleAggregateError
    bitAnd = impossibleAggregateError
    bitOr = impossibleAggregateError
    boolAnd = impossibleAggregateError
    boolOr = impossibleAggregateError
    every = impossibleAggregateError
    max_ = impossibleAggregateError
    min_ = impossibleAggregateError
    avg = impossibleAggregateError
    corr = impossibleAggregateError
    covarPop = impossibleAggregateError
    covarSamp = impossibleAggregateError
    regrAvgX = impossibleAggregateError
    regrAvgY = impossibleAggregateError
    regrCount = impossibleAggregateError
    regrIntercept = impossibleAggregateError
    regrR2 = impossibleAggregateError
    regrSlope = impossibleAggregateError
    regrSxx = impossibleAggregateError
    regrSxy = impossibleAggregateError
    regrSyy = impossibleAggregateError
    stddev = impossibleAggregateError
    stddevPop = impossibleAggregateError
    stddevSamp = impossibleAggregateError
    variance = impossibleAggregateError
    varPop = impossibleAggregateError
    varSamp = impossibleAggregateError
  
  -- | helper function for our errors above
  impossibleAggregateError :: a
  impossibleAggregateError = error "impossible; called aggregate function for Ungrouped even though the Aggregate instance has a type error constraint."
  
  -- | escape hatch to define aggregate functions
  unsafeAggregate
    :: SOP.SListI xs
    => ByteString -- ^ function
    -> AggregateArg xs lat with db params from
    -> Expression ('Grouped bys) lat with db params from y
  unsafeAggregate fun xs = UnsafeExpression $ fun <> renderSQL xs
  
  -- | A type family that calculates `PGSum` `PGType` of
  -- a given argument `PGType`.
  type family PGSum ty where
    PGSum 'PGint2 = 'PGint8
    PGSum 'PGint4 = 'PGint8
    PGSum 'PGint8 = 'PGnumeric
    PGSum 'PGfloat4 = 'PGfloat4
    PGSum 'PGfloat8 = 'PGfloat8
    PGSum 'PGnumeric = 'PGnumeric
    PGSum 'PGinterval = 'PGinterval
    PGSum 'PGmoney = 'PGmoney
    PGSum pg = TypeError
      ( 'Text "Squeal type error: Cannot sum with argument type "
        ':<>: 'ShowType pg )
  
  -- | A type family that calculates `PGAvg` type of a `PGType`.
  type family PGAvg ty where
    PGAvg 'PGint2 = 'PGnumeric
    PGAvg 'PGint4 = 'PGnumeric
    PGAvg 'PGint8 = 'PGnumeric
    PGAvg 'PGnumeric = 'PGnumeric
    PGAvg 'PGfloat4 = 'PGfloat8
    PGAvg 'PGfloat8 = 'PGfloat8
    PGAvg 'PGinterval = 'PGinterval
    PGAvg pg = TypeError
      ('Text "Squeal type error: No average for " ':<>: 'ShowType pg)
  