{-# 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 :: forall (f :: k -> *). Proxy ts -> Rec f '[] -> [CoRec f ts]
rowToColumnAux Proxy ts
_ Rec f '[]
_ = []

instance (r ∈ ts, RowToColumn ts rs) => RowToColumn ts (r ': rs) where
  rowToColumnAux :: forall (f :: a -> *). Proxy ts -> Rec f (r : rs) -> [CoRec f ts]
rowToColumnAux Proxy ts
p (f r
x :& Rec f rs
xs) = forall {k} (a1 :: k) (b :: [k]) (a :: k -> *).
RElem a1 b (RIndex a1 b) =>
a a1 -> CoRec a b
CoRec f r
x forall a. a -> [a] -> [a]
: 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 :: forall {k} (ts :: [k]) (f :: k -> *).
RowToColumn ts ts =>
Rec f ts -> [CoRec f ts]
rowToColumn = forall {k} (ts :: [k]) (rs :: [k]) (f :: k -> *).
RowToColumn ts rs =>
Proxy ts -> Rec f rs -> [CoRec f ts]
rowToColumnAux 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 :: 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 Record ts
r = forall a b. (a -> b) -> [a] -> [b]
map (\CoRec ElField vs
val -> forall (t :: (Symbol, *)). Snd t -> ElField t
Field CoRec ElField vs
val forall {u} (a :: u -> *) (r :: u) (rs :: [u]).
a r -> Rec a rs -> Rec a (r : rs)
:& Record ss
ids) (forall {k} (ts :: [k]) (f :: k -> *).
RowToColumn ts ts =>
Rec f ts -> [CoRec f ts]
rowToColumn Record vs
vals)
  where ids :: Record ss
ids = 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 :: Record vs
vals = 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' :: 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' proxy 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 :: forall (t :: (Symbol, *)) (ts :: [(Symbol, *)]).
Record (t : ts) -> Record (ts ++ '[t])
retroSnoc (ElField r
x :& Rec ElField rs
xs) = forall (ss :: [(Symbol, *)]). Record ss -> Record (ss ++ '[t])
go Rec ElField rs
xs
  where go :: Record ss -> Record (ss ++ '[t])
        go :: forall (ss :: [(Symbol, *)]). Record ss -> Record (ss ++ '[t])
go Rec ElField ss
RNil = ElField r
x forall {u} (a :: u -> *) (r :: u) (rs :: [u]).
a r -> Rec a rs -> Rec a (r : rs)
:& forall {u} (a :: u -> *). Rec a '[]
RNil
        go (ElField r
y :& Rec ElField rs
ys) = ElField r
y forall {u} (a :: u -> *) (r :: u) (rs :: [u]).
a r -> Rec a rs -> Rec a (r : rs)
:& 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 :: 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 = (forall a b. (a -> b) -> [a] -> [b]
map forall (t :: (Symbol, *)) (ts :: [(Symbol, *)]).
Record (t : ts) -> Record (ts ++ '[t])
retroSnoc forall b c a. (b -> c) -> (a -> b) -> a -> c
.) forall b c a. (b -> c) -> (a -> b) -> a -> c
. 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 :: forall (proxy :: [k] -> *). proxy '[] -> Int
hasLength proxy '[]
_ = Int
0
instance forall t ts. HasLength ts => HasLength (t ': ts) where
  hasLength :: forall (proxy :: [k] -> *). proxy (t : ts) -> Int
hasLength proxy (t : ts)
_ = Int
1 forall a. Num a => a -> a -> a
+ forall k (ts :: [k]) (proxy :: [k] -> *).
HasLength ts =>
proxy ts -> Int
hasLength (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 :: forall (vs :: [(Symbol, *)]) (ts :: [(Symbol, *)])
       (ss :: [(Symbol, *)]) (proxy :: [(Symbol, *)] -> *).
(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
p (Frame Int
n Int -> Record ts
v) = forall r. Int -> (Int -> r) -> Frame r
Frame (Int
nforall a. Num a => a -> a -> a
*Int
numVs) Int -> Record (ss ++ '["value" :-> CoRec ElField vs])
go
  where numVs :: Int
numVs = forall k (ts :: [k]) (proxy :: [k] -> *).
HasLength ts =>
proxy ts -> Int
hasLength (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 forall a. Integral a => a -> a -> (a, a)
`quotRem` Int
numVs
               in 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) forall a. [a] -> Int -> a
!! Int
k