{-# LANGUAGE DeriveGeneric #-}
{-# LANGUAGE DeriveAnyClass #-}

module Joins (
    readFileToTable,
    selectFromTable,
    sortTableOnCol,
    nestedLoopJoin,
    sortMergeJoin,
    partitionTableOnId,
    ) where


import Data.List.Split (splitOn, chunksOf)
import Data.List (elemIndex, sortOn, groupBy)
import Data.Maybe (fromJust)
import Data.Char (toUpper, digitToInt)
import GHC.Generics
import Control.DeepSeq (NFData)

type Column = String

type Record = [String]


data Table = Table {
    name :: String,
    columns :: [Column],
    rows :: [Record]
} deriving (Show, Generic, NFData)


instance Semigroup Table where
    (Table name1 cols1 rows1) <> (Table name2 cols2 rows2) = 
        Table name1 cols1 (rows1 ++ rows2)
    

-- Takes a Table, a Column to partition on, and a number of partitions
--  Partitions on the Column
--  Assumes that values are Strings and begin with "1" through "9"
partitionTableOnId :: Table -> Column -> Int -> [Table]
partitionTableOnId t col numPartitions = 
    [Table {name = name t, columns = columns t, rows = partRow} | partRow <- partedRows]
    where
        partitions = chunksOf numPartitions [1..9]
        partedRows = [partitionHelper (rows t) colIndex part | part <- partitions]
        colIndex = fromJust $ getColIndex t col
        partitionVal = (!! colIndex)


-- Helper for partitioning
--  Takes a record, a column index to partition on, and a list of integer values 
--  representing a group of values to partition on
--  For example, a group of [1..3] means that any id starting with "1", "2", or "3" 
--  will be in this partition
partitionHelper :: [Record] -> Int -> [Int] -> [Record]
partitionHelper rows colIndex pred = 
    filter ((`elem` pred) . digitToInt . head . partitionVal) rows
    where
        partitionVal = (!! colIndex)


-- Takes two Tables and returns True if the first has fewer rows than the second 
-- Not using the Ord typeclass here because defining Eq doesn't really make sense
firstTableSmaller :: Table -> Table -> Bool
firstTableSmaller t1 t2 = card1 < card2
    where
        card1 = length $ rows t1
        card2 = length $ rows t2


-- Takes a FilePath of the form "/path/to/file_name.ext" and returns <file_name>
parseNameFromFile :: FilePath -> String
parseNameFromFile name = nameWithoutExtension
    where
        nameWithoutPath = last (splitOn "/" name)
        nameWithoutExtension = head $ splitOn "." nameWithoutPath


-- Takes a "header" row from the file and returns it as [Columns]
parseColumnsFromHeader :: String -> String -> [Column]
parseColumnsFromHeader header delim = splitOn delim $ filterStopChars header
    where
        filterStopChars = filter (`notElem` ['\r', '\n', '"'])


-- Takes the body (file without header) and returns a [Record]
parsePayloadsFromBody :: [String] -> String -> [Record]
parsePayloadsFromBody rows delim = map (splitOn ";" . filterStopChars) rows
    where
        filterStopChars = filter (`notElem` ['\r', '\n', '"'])


-- Takes a FilePath and a delimiter, and returns a Table
readFileToTable :: FilePath -> String -> IO Table
readFileToTable fileName delim = do
    file <- readFile fileName
    let tableName = parseNameFromFile fileName
        tableLines = lines file
        columns = parseColumnsFromHeader (head tableLines) delim
        rows = parsePayloadsFromBody (tail tableLines) delim
    return $ Table{name = tableName, columns = columns, rows = rows}


-- Takes a Table and a [Column], returns True if all columns in Table
colInTable :: Table -> Column -> Bool
colInTable t col = col `elem` columns t


-- Takes a Table and a Column, returns the index of that Column in Table{columns}
getColIndex :: Table -> Column -> Maybe Int
getColIndex t col = elemIndex col (columns t)


-- Takes a Table and [Column], and returns the Table with only those columns
selectFromTable :: Table -> [Column] -> Table
selectFromTable t cols = Table {name = selectName, columns = cols, rows = selectRows}
    where
        selectName = "select" ++ [toUpper $ head $ name t] ++ tail (name t)
        indices = [fromJust $ getColIndex t col | col <- cols, colInTable t col]
        selectRows = [[(!! i) r | i <- indices] | r <- rows t]


-- Takes a Table, a Column, and a value, and returns a Table where rows have col == comp
filterTableOnCol :: Table -> Column -> String -> Table
filterTableOnCol t col comp =
    if colInTable t col then t {rows = filtered} else error errorMessage

    where
        errorMessage = "No column(s) '" ++ col ++ "' in table " ++ name t
        colIx = elemIndex col (columns t)
        filtered = [p | p <- rows t, (!! fromJust colIx) p == comp]


-- Takes a Table and a Column, and returns the Table sorted on that column
sortTableOnCol :: Table -> Column -> Table
sortTableOnCol t col =
    if colInTable t col then t {rows = sorted} else error errorMessage
    where
        errorMessage = "No column(s) '" ++ col ++ "' in table " ++ name t
        colIx = elemIndex col (columns t)
        sorted = sortOn (!! fromJust colIx) (rows t)


nestedLoopJoin :: Table -> Table -> Column -> Table
nestedLoopJoin t1 t2 col = Table {name = joinName, columns = joinCols, rows = joinRows}
    where
        (outer, inner) = if firstTableSmaller t1 t2 then (t1, t2) else (t2, t1)
        joinName = name outer ++ " X " ++ name inner
        joinCols = columns outer ++ columns inner
        outerRows = rows outer
        innerRows = rows inner
        outerColIx = fromJust $ elemIndex col $ columns outer
        innerColIx = fromJust $ elemIndex col $ columns inner
        joinRows = [x ++ y | x <- outerRows, y <- innerRows, x !! outerColIx == y !! innerColIx]


sortMergeJoin :: Table -> Table -> Column -> Table
sortMergeJoin t1 t2 col = Table {name = joinName, columns = joinCols, rows = joinRows}
    where
        (outer, inner) = if firstTableSmaller t1 t2 then (t1, t2) else (t2, t1)
        joinName = name outer ++ " X " ++ name inner
        joinCols = columns outer ++ columns inner
        outerInd = fromJust $ getColIndex outer col
        innerInd = fromJust $ getColIndex inner col
        sortedOuter = sortTableOnCol outer col
        sortedInner = sortTableOnCol inner col
        outerRows = rows sortedOuter
        innerRows = rows sortedInner
        joinRows = sortMergeJoinHelper outerRows outerInd innerRows innerInd col


-- Helper for sortMergeJoin that operates on rows
-- Smaller table will be the left (outer) table
sortMergeJoinHelper :: [Record] -> Int -> [Record] -> Int -> Column -> [Record]
sortMergeJoinHelper [] _ ys _ _ = []
sortMergeJoinHelper xs _ [] _ _ = []
-- If the outer (smaller) list is a singleton, advance inner list until there 
--  is no match
sortMergeJoinHelper [x] outerInd (y:ys) innerInd col = if getOuterVal == getInnerVal then
    (x ++ y) : continue else continue
    where
        getOuterVal = (!! outerInd) x
        getInnerVal = (!! innerInd) y
        continue = sortMergeJoinHelper [x] outerInd ys innerInd col
-- If there is a match, match the current outer row with all rows in inner rows that match
--  Then advance outer row
sortMergeJoinHelper rowsT1@(x:xs) outerInd rowsT2@(y:ys) innerInd col
    | getOuterVal == getInnerVal =
        [x ++ inner_y | inner_y <- rowsT2, (!! innerInd) inner_y == getOuterVal] ++ advanceOuter
    | getOuterVal < getInnerVal = advanceOuter
    | otherwise = advanceInner
    where
        getOuterVal = (!! outerInd) x
        getInnerVal = (!! innerInd) y
        peekAhead = (!! outerInd) (head xs)
        advanceOuter = sortMergeJoinHelper xs outerInd rowsT2 innerInd col
        advanceInner = sortMergeJoinHelper rowsT1 outerInd ys innerInd col
