{-# LANGUAGE TupleSections #-}
module Diagonal where

import Colonnade
import Data.Profunctor
import Data.Bifunctor

-- Row numbers, as integers with two extra values at the end
data RowNumber = RN Int | Ellipsis | Omega deriving Eq
instance Show RowNumber where
  show (RN n) = show n
  show Ellipsis = "..."
  show Omega = "ω"

-- Rendered cell data, along with its row number and an optional value to show instead
data RenderCell a = RC {
    cellRow :: Maybe RowNumber,
    cellForceShow :: Maybe String,
    cellData :: a
  } deriving Functor

-- If a forced value is availaible, use it; otherwise, show the inside
instance Show a => Show (RenderCell a) where
  show (RC _ Nothing x)  = show x
  show (RC _ (Just s) _) = s

-- Lift a RenderCell over Either to an Either over RenderCells
-- If the force-string is set, then a `Left` version with an undefined value is returned
liftRC :: RenderCell (Either a b) -> Either (RenderCell a) (RenderCell b)
liftRC (RC rn (Just x) _) = Left $ RC rn (Just x) undefined
liftRC (RC rn Nothing (Left v)) = Left $ RC rn Nothing v
liftRC (RC rn Nothing (Right v)) = Right $ RC rn Nothing v

-- Single column with a fixed header and data coming from a row number
newColumn header f = columns f (\_ -> Headed $ RC Nothing (Just header) undefined) [()]

-- Numbered rows of data, as opposed to data for rendering
data DiagRow a = DR {
    rowNumber :: RowNumber,
    rowData :: [a]
  } deriving Functor

-- Same as `newColumn`, but reads into a DiagRow's row numberRows
-- The display function `f` is used to convert the value to a String
newColumnR header f = newColumn header (\_ c -> RC (Just $ rowNumber c) (Just $ f $ rowNumber c) undefined)

-- Number rows with the title "Sequence"
numberColumn = newColumnR "Sequence" show
-- Column of ellipses
ellipsisColumn = newColumnR "..." (const "...")

-- Create a Colonnade over DiagRows by mapping them into lists
-- The ellipsis row is mapped to a cells which always display as ellipses
-- The remaining rows are `Right` if they lie on the main diagonal, or `Left` otherwise
diagBox' n = columns cell headCell [0..n]
  where cell _ (DR Ellipsis  _) = RC (Just Ellipsis) (Just "...") undefined
        cell i (DR rn       rd) = RC (Just rn) Nothing $ (if RN i == rn then Right else Left) (rd !! i)
        headCell i = Headed $ RC Nothing (Just $ show i) undefined

diagBox n = numberColumn <> diagBox' n <> ellipsisColumn

displayRows f = (displayTable .) . tagRows'
  where -- Tag rows by (predicate function, extra value) tuples
        tagRows' [] r@(RC rn _ _) = Left r  -- Tag missing things as left
        tagRows' ((p, f):ps) r@(RC rn _ _) = if p == rn then Right (f, r) else tagRows' ps r -- Match found, tag right
        displayTable row = case bimap liftRC (fmap liftRC) row of
          (Left (Left r)) -> Left r
          (Left (Right r)) -> Right $ f r
          (Right (g, r)) -> Right $ g $ either id id r

rowsOmega :: [[a]] -> [a] -> [DiagRow a]
rowsOmega xs y = zipWith (\x n -> DR (RN n) x) xs [0..] ++ [DR Ellipsis [], DR Omega y]


displayDiagTable ps = rmap style . diagBox
  where style    = either show id . displayRows (("Red " ++) . show) (map (, ("Green " ++). show) ps)