-- Vicotoria Lam
-- COMP 554 Spring 2018
-- Interesting Haskell solution to this problem
import Data.Set (fromList,union,empty,isSubsetOf,member,Set)
import Data.List (sortOn,delete)

main :: IO ()
main = return ()
-- ghcid gets mad if this isn't here.

graph2 :: [[Double]]
graph2 = [ [ 0, 2, 1]
         , [ 2, 0, 7]
         , [ 1, 7, 0] ]

graph4 :: [[Double]]
graph4 = [ [ 0,10, 5, 7]
         , [10, 0, 9, 6]
         , [ 5, 9, 0, 8]
         , [ 7, 6, 8, 0] ]

graph6 :: [[Double]]
graph6 = [ [ 0,10,12,11, 3, 5]
         , [10, 0, 9,12, 2, 4]
         , [12, 9, 0, 8, 0, 0]
         , [11,12, 8, 0, 7, 0]
         , [ 3, 2, 0, 7, 0, 6]
         , [ 5, 4, 0, 0, 6, 0] ]

-- Part I.

v :: [[Double]] -> Set Int
v g = fromList [0..(length (g)-1) ]

scan :: [[Double]] -> [(Double, Int, Int)]
scan g = helpList
  where
    helpList = [ foo i j | j <- [0..(n-1)], i <- [j..(n-1)], keep i j ]
    n = length g
    keep :: Int -> Int -> Bool
    keep i j = g !! i !! j /= 0
    foo :: Int -> Int -> (Double, Int, Int)
    foo i j = (g !! i !! j, i, j)

sort_ :: [[Double]] -> [(Int, Int)]
sort_ g = fmap getEdge (sortOn getCost (scan g))
  where
    getEdge :: (Double, Int, Int) -> (Int, Int)
    getEdge (cost, v1, v2) = (v1, v2)
    getCost :: (Double, Int, Int) -> Double
    getCost (cost, v1, v2) = cost
-- note: sortOn function list


-- Part II.

type Tree = Set (Int, Int)

extract :: (Int, Int) -> Set Int
extract (a, b) =
  fromList [a,b]

loop :: Set Int -> (Tree, Set Int, [(Int, Int)] ) -> Tree
loop v (tspan, included, sorted) = if included == v
  then tspan
  else loop v (step (tspan, included, sorted))

step :: (Tree, Set Int, [(Int, Int)] )
     -> (Tree, Set Int, [(Int, Int)] )
step (tspan, included, sorted) = (tspan', included', sorted')
  where
    check (i, j) = (member i included || member j included) && not (isSubsetOf (fromList [i,j]) included)
    viable = filter check sorted
    tspan' = union tspan (fromList [head viable])
    included' = union included (extract (head viable))
    sorted' =  delete (head viable) sorted

function :: [[Double]] -> Tree
function g = loop (v g) (tspan, included, sorted)
  where
    tspan = empty
    included = fromList [0]
    sorted = sort_ g