{-# LANGUAGE ConstraintKinds, DataKinds, FlexibleContexts, FlexibleInstances,
             KindSignatures, MultiParamTypeClasses, PolyKinds,
             ScopedTypeVariables, TypeFamilies, TypeOperators,
             UndecidableInstances #-}
module Frames.Melt where
import Data.Proxy
import Data.Vinyl
import Data.Vinyl.CoRec (CoRec(..))
import Data.Vinyl.TypeLevel
import Frames.Col
import Frames.Frame (Frame(..), FrameRec)
import Frames.Rec
import Frames.RecF (ColumnHeaders(..))

type family Elem t ts :: Bool where
  Elem t '[] = 'False
  Elem t (t ': ts) = 'True
  Elem t (s ': ts) = Elem t ts

type family Or (a :: Bool) (b :: Bool) :: Bool where
  Or 'True b = 'True
  Or a b = b

type family Not a :: Bool where
  Not 'True = 'False
  Not 'False = 'True

type family Disjoint ss ts :: Bool where
  Disjoint '[] ts = 'True
  Disjoint (s ': ss) ts = Or (Not (Elem s ts)) (Disjoint ss ts)

type ElemOf ts r = RElem r ts (RIndex r ts)

class RowToColumn ts rs where
  rowToColumnAux :: Proxy ts -> Rec f rs -> [CoRec f ts]

instance RowToColumn ts '[] where
  rowToColumnAux :: Proxy ts -> Rec f '[] -> [CoRec f ts]
rowToColumnAux Proxy ts
_ Rec f '[]
_ = []

instance (r ∈ ts, RowToColumn ts rs) => RowToColumn ts (r ': rs) where
  rowToColumnAux :: Proxy ts -> Rec f (r : rs) -> [CoRec f ts]
rowToColumnAux Proxy ts
p (f r
x :& Rec f rs
xs) = f r -> CoRec f ts
forall k (a1 :: k) (b :: [k]) (a :: k -> *).
RElem a1 b (RIndex a1 b) =>
a a1 -> CoRec a b
CoRec f r
x CoRec f ts -> [CoRec f ts] -> [CoRec f ts]
forall a. a -> [a] -> [a]
: Proxy ts -> Rec f rs -> [CoRec f ts]
forall k (ts :: [k]) (rs :: [k]) (f :: k -> *).
RowToColumn ts rs =>
Proxy ts -> Rec f rs -> [CoRec f ts]
rowToColumnAux Proxy ts
p Rec f rs
xs

-- | Transform a record into a list of its fields, retaining proof
-- that each field is part of the whole.
rowToColumn :: RowToColumn ts ts => Rec f ts -> [CoRec f ts]
rowToColumn :: Rec f ts -> [CoRec f ts]
rowToColumn = Proxy ts -> Rec f ts -> [CoRec f ts]
forall k (ts :: [k]) (rs :: [k]) (f :: k -> *).
RowToColumn ts rs =>
Proxy ts -> Rec f rs -> [CoRec f ts]
rowToColumnAux Proxy ts
forall k (t :: k). Proxy t
Proxy

meltAux :: forall vs ss ts.
           (vs ⊆ ts, ss ⊆ ts, Disjoint ss ts ~ 'True, ts ≅ (vs ++ ss),
           ColumnHeaders vs, RowToColumn vs vs)
        => Record ts
        -> [Record ("value" :-> CoRec ElField vs ': ss)]
meltAux :: Record ts -> [Record (("value" :-> CoRec ElField vs) : ss)]
meltAux Record ts
r = (CoRec ElField vs -> Record (("value" :-> CoRec ElField vs) : ss))
-> [CoRec ElField vs]
-> [Record (("value" :-> CoRec ElField vs) : ss)]
forall a b. (a -> b) -> [a] -> [b]
map (\CoRec ElField vs
val -> CoRec ElField vs -> ElField ("value" :-> CoRec ElField vs)
forall (s :: Symbol) t. KnownSymbol s => t -> ElField '(s, t)
Field CoRec ElField vs
val ElField ("value" :-> CoRec ElField vs)
-> Rec ElField ss -> Record (("value" :-> CoRec ElField vs) : ss)
forall u (a :: u -> *) (r :: u) (rs :: [u]).
a r -> Rec a rs -> Rec a (r : rs)
:& Rec ElField ss
ids) (Rec ElField vs -> [CoRec ElField vs]
forall k (ts :: [k]) (f :: k -> *).
RowToColumn ts ts =>
Rec f ts -> [CoRec f ts]
rowToColumn Rec ElField vs
vals)
  where ids :: Rec ElField ss
ids = Record ts -> Rec ElField ss
forall k1 k2 (rs :: [k1]) (ss :: [k1]) (f :: k2 -> *)
       (record :: (k2 -> *) -> [k1] -> *) (is :: [Nat]).
(RecSubset record rs ss is, RecSubsetFCtx record f) =>
record f ss -> record f rs
rcast Record ts
r :: Record ss
        vals :: Rec ElField vs
vals = Record ts -> Rec ElField vs
forall k1 k2 (rs :: [k1]) (ss :: [k1]) (f :: k2 -> *)
       (record :: (k2 -> *) -> [k1] -> *) (is :: [Nat]).
(RecSubset record rs ss is, RecSubsetFCtx record f) =>
record f ss -> record f rs
rcast Record ts
r :: Record vs

type family RDeleteAll ss ts where
  RDeleteAll '[] ts = ts
  RDeleteAll (s ': ss) ts = RDeleteAll ss (RDelete s ts)

-- | This is 'melt', but the variables are at the front of the record,
-- which reads a bit odd.
meltRow' :: forall proxy vs ts ss. (vs ⊆ ts, ss ⊆ ts, vs ~ RDeleteAll ss ts,
            Disjoint ss ts ~ 'True, ts ≅ (vs ++ ss),
            ColumnHeaders vs, RowToColumn vs vs)
         => proxy ss
         -> Record ts
         -> [Record ("value" :-> CoRec ElField vs ': ss)]
meltRow' :: proxy ss
-> Record ts -> [Record (("value" :-> CoRec ElField vs) : ss)]
meltRow' proxy ss
_ = Record ts -> [Record (("value" :-> CoRec ElField vs) : ss)]
forall (vs :: [(Symbol, *)]) (ss :: [(Symbol, *)])
       (ts :: [(Symbol, *)]).
(vs ⊆ ts, ss ⊆ ts, Disjoint ss ts ~ 'True, ts ≅ (vs ++ ss),
 ColumnHeaders vs, RowToColumn vs vs) =>
Record ts -> [Record (("value" :-> CoRec ElField vs) : ss)]
meltAux

-- | Turn a cons into a snoc after the fact.
retroSnoc :: forall t ts. Record (t ': ts) -> Record (ts ++ '[t])
retroSnoc :: Record (t : ts) -> Record (ts ++ '[t])
retroSnoc (ElField r
x :& Rec ElField rs
xs) = Rec ElField rs -> Record (rs ++ '[t])
forall (ss :: [(Symbol, *)]). Record ss -> Record (ss ++ '[t])
go Rec ElField rs
xs
  where go :: Record ss -> Record (ss ++ '[t])
        go :: Record ss -> Record (ss ++ '[t])
go Record ss
RNil = ElField r
x ElField r -> Rec ElField '[] -> Rec ElField '[r]
forall u (a :: u -> *) (r :: u) (rs :: [u]).
a r -> Rec a rs -> Rec a (r : rs)
:& Rec ElField '[]
forall u (a :: u -> *). Rec a '[]
RNil
        go (ElField r
y :& Rec ElField rs
ys) = ElField r
y ElField r
-> Rec ElField (rs ++ '[t]) -> Rec ElField (r : (rs ++ '[t]))
forall u (a :: u -> *) (r :: u) (rs :: [u]).
a r -> Rec a rs -> Rec a (r : rs)
:& Rec ElField rs -> Rec ElField (rs ++ '[t])
forall (ss :: [(Symbol, *)]). Record ss -> Record (ss ++ '[t])
go Rec ElField rs
ys

-- | Like @melt@ in the @reshape2@ package for the @R@ language. It
-- stacks multiple columns into a single column over multiple
-- rows. Takes a specification of the id columns that remain
-- unchanged. The remaining columns will be stacked.
--
-- Suppose we have a record, @r :: Record [Name,Age,Weight]@. If we
-- apply @melt [pr1|Name|] r@, we get two values with type @Record
-- [Name, "value" :-> CoRec Identity [Age,Weight]]@. The first will
-- contain @Age@ in the @value@ column, and the second will contain
-- @Weight@ in the @value@ column.
meltRow :: (vs ⊆ ts, ss ⊆ ts, vs ~ RDeleteAll ss ts,
            Disjoint ss ts ~ 'True, ts ≅ (vs ++ ss),
            ColumnHeaders vs, RowToColumn vs vs)
        => proxy ss
        -> Record ts
        -> [Record (ss ++ '["value" :-> CoRec ElField vs])]
meltRow :: proxy ss
-> Record ts -> [Record (ss ++ '["value" :-> CoRec ElField vs])]
meltRow = ((Record (("value" :-> CoRec ElField vs) : ss)
 -> Record (ss ++ '["value" :-> CoRec ElField vs]))
-> [Record (("value" :-> CoRec ElField vs) : ss)]
-> [Record (ss ++ '["value" :-> CoRec ElField vs])]
forall a b. (a -> b) -> [a] -> [b]
map Record (("value" :-> CoRec ElField vs) : ss)
-> Record (ss ++ '["value" :-> CoRec ElField vs])
forall (t :: (Symbol, *)) (ts :: [(Symbol, *)]).
Record (t : ts) -> Record (ts ++ '[t])
retroSnoc ([Record (("value" :-> CoRec ElField vs) : ss)]
 -> [Record (ss ++ '["value" :-> CoRec ElField vs])])
-> (Record ts -> [Record (("value" :-> CoRec ElField vs) : ss)])
-> Record ts
-> [Record (ss ++ '["value" :-> CoRec ElField vs])]
forall b c a. (b -> c) -> (a -> b) -> a -> c
.) ((Record ts -> [Record (("value" :-> CoRec ElField vs) : ss)])
 -> Record ts -> [Record (ss ++ '["value" :-> CoRec ElField vs])])
-> (proxy ss
    -> Record ts -> [Record (("value" :-> CoRec ElField vs) : ss)])
-> proxy ss
-> Record ts
-> [Record (ss ++ '["value" :-> CoRec ElField vs])]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. proxy ss
-> Record ts -> [Record (("value" :-> CoRec ElField vs) : ss)]
forall (proxy :: [(Symbol, *)] -> *) (vs :: [(Symbol, *)])
       (ts :: [(Symbol, *)]) (ss :: [(Symbol, *)]).
(vs ⊆ ts, ss ⊆ ts, vs ~ RDeleteAll ss ts, Disjoint ss ts ~ 'True,
 ts ≅ (vs ++ ss), ColumnHeaders vs, RowToColumn vs vs) =>
proxy ss
-> Record ts -> [Record (("value" :-> CoRec ElField vs) : ss)]
meltRow'

class HasLength (ts :: [k]) where
  hasLength :: proxy ts -> Int

instance HasLength '[] where hasLength :: proxy '[] -> Int
hasLength proxy '[]
_ = Int
0
instance forall t ts. HasLength ts => HasLength (t ': ts) where
  hasLength :: proxy (t : ts) -> Int
hasLength proxy (t : ts)
_ = Int
1 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Proxy ts -> Int
forall k (ts :: [k]) (proxy :: [k] -> *).
HasLength ts =>
proxy ts -> Int
hasLength (Proxy ts
forall k (t :: k). Proxy t
Proxy :: Proxy ts)

-- | Applies 'meltRow' to each row of a 'FrameRec'.
melt :: forall vs ts ss proxy.
        (vs ⊆ ts, ss ⊆ ts, vs ~ RDeleteAll ss ts, HasLength vs,
         Disjoint ss ts ~ 'True, ts ≅ (vs ++ ss),
         ColumnHeaders vs, RowToColumn vs vs)
     => proxy ss
     -> FrameRec ts
     -> FrameRec (ss ++ '["value" :-> CoRec ElField vs])
melt :: proxy ss
-> FrameRec ts -> FrameRec (ss ++ '["value" :-> CoRec ElField vs])
melt proxy ss
p (Frame Int
n Int -> Record ts
v) = Int
-> (Int -> Record (ss ++ '["value" :-> CoRec ElField vs]))
-> FrameRec (ss ++ '["value" :-> CoRec ElField vs])
forall r. Int -> (Int -> r) -> Frame r
Frame (Int
nInt -> Int -> Int
forall a. Num a => a -> a -> a
*Int
numVs) Int -> Record (ss ++ '["value" :-> CoRec ElField vs])
go
  where numVs :: Int
numVs = Proxy vs -> Int
forall k (ts :: [k]) (proxy :: [k] -> *).
HasLength ts =>
proxy ts -> Int
hasLength (Proxy vs
forall k (t :: k). Proxy t
Proxy :: Proxy vs)
        go :: Int -> Record (ss ++ '["value" :-> CoRec ElField vs])
go Int
i = let (Int
j,Int
k) = Int
i Int -> Int -> (Int, Int)
forall a. Integral a => a -> a -> (a, a)
`quotRem` Int
numVs
               in proxy ss
-> Record ts -> [Record (ss ++ '["value" :-> CoRec ElField vs])]
forall (vs :: [(Symbol, *)]) (ts :: [(Symbol, *)])
       (ss :: [(Symbol, *)]) (proxy :: [(Symbol, *)] -> *).
(vs ⊆ ts, ss ⊆ ts, vs ~ RDeleteAll ss ts, Disjoint ss ts ~ 'True,
 ts ≅ (vs ++ ss), ColumnHeaders vs, RowToColumn vs vs) =>
proxy ss
-> Record ts -> [Record (ss ++ '["value" :-> CoRec ElField vs])]
meltRow proxy ss
p (Int -> Record ts
v Int
j) [Record (ss ++ '["value" :-> CoRec ElField vs])]
-> Int -> Record (ss ++ '["value" :-> CoRec ElField vs])
forall a. [a] -> Int -> a
!! Int
k