{-# OPTIONS_GHC -Wall #-}

module DataSet (DataSet (..), Metric, handleLookups, MetricPrime, transform, fromLists)
where

import Data.Map (Map, (!))
import qualified Data.Map as M

-- import Control.Compose
{-
  So I thought it might be a cute exercise,
  when all is said & done, to make our
  DataSet type be something like
  type DataSet a b = Map a :. Map b
  using Conal's type compose library.
  - wchogg
-}

newtype DataSet a b c = DataSet (Map a (Map b c))
    deriving (Show, Eq)

fromMap :: Map a (Map b c) -> DataSet a b c
fromMap = DataSet

fromLists :: (Ord a, Ord b) => [(a, [(b, c)])] -> DataSet a b c
fromLists = DataSet . M.fromList . (fmap . fmap) M.fromList

addData :: (Ord a, Ord b) => a -> b -> c -> DataSet a b c -> DataSet a b c
addData key subkey val (DataSet set) = DataSet $ addData' key subkey val set

addData' :: (Ord a, Ord b) => a -> b -> c 
                           -> Map a (Map b c) 
                           -> Map a (Map b c)
addData' key subkey val set = 
    case M.lookup key set of
      Just set' -> M.adjust (M.insert subkey val) key set
      Nothing -> M.insert key (M.singleton subkey val) set
                           
getValue :: (Ord a, Ord b) => a -> b -> DataSet a b c -> (Maybe c)
getValue key subkey (DataSet set) = do
  m <- M.lookup key set
  M.lookup subkey m

getValues :: (Ord a, Ord b) => a -> DataSet a b c -> (Maybe (Map b c))
getValues key (DataSet set) = M.lookup key set

type Metric a b c = a -> a -> DataSet a b c -> Maybe c
type MetricPrime b c = Map b c -> Map b c -> Maybe c

handleLookups :: (Ord b, Ord a) =>(Map b c -> Map b c -> Maybe d)
                     -> a -> a -> DataSet a b c -> Maybe d
handleLookups f key key' (DataSet set) = do
  m <- M.lookup key set
  m' <- M.lookup key' set
  f m m'

transform :: (Ord a, Ord b) => DataSet a b c -> DataSet b a c
transform (DataSet set) = DataSet (transform' set)
transform' :: (Ord a, Ord b) => Map a (Map b c) 
                             -> Map b (Map a c)
transform' = M.foldWithKey aux M.empty
 where aux key m set = M.foldWithKey inserter set m
           where inserter key' v set' = addData' key' key v set'