-- Info:

-- This program solves linear equation of any dimensions using LU
-- decomposition

-- To solve:
-- 3x - 4y = 0
-- 9x - 8y = 12
-- do in the repl:
-- solve [[3,-4],[9,-8]] [0,12]

-- You can use whatever size matrix


----------------------------------------------
-- These are the helper functions and types --

type Mat a = [Row a]

type Row a = [a]

p_transpose :: Mat a -> Mat a
p_transpose ([]:_) = []
p_transpose x = map head x : p_transpose (map tail x)

-- Matrix sun
smult :: Fractional a => [a] -> [a] -> a
smult xs = sum . zipWith (*) xs

-- Matrix multiplication
mult :: (Fractional a) => [[a]] -> [[a]] -> [[a]]
mult ma mb = [map (smult row) mbt | row <- ma]
  where mbt = p_transpose mb

---------------- LU decomposition-----------------

lu_mat :: Fractional a => [Row a] -> ([Row a], [[a]])
lu_mat [a] = ([a],[[1]])
lu_mat (xs:xss)
  = (com [mat, subm], (1:(map negate facts)):(map (0:) u))
  where facts = map negate $ map (/head xs) (map head xss)
        mat = zipWith (\x -> zipWith (+) (map (x*) xs)) facts xss
        (subm, u) = lu_mat $ map tail mat

lu :: Fractional a => [Row a] -> (Mat a, Mat a)
lu m = (p_transpose u, head m : l)
  where (l,u) = lu_mat m

com :: [Mat a] -> Mat a
com [a] = a
com xs = head a : (zipWith (:) (map head (tail a)) b)
  where a = head xs
        b = com (tail xs)

--------------Solving system equation----------------

-- solve :: Fractional a => Mat a -> Mat a -> Mat a
-- LZ = C
-- 
solve a c = let (l,u) = lu a
                z = solve_lower l c
                x = reverse $
                      solve_lower (reverse $ p_transpose $ reverse $ p_transpose u)
                 (reverse z)
            in x

-- AX = Z
solve_lower :: (Foldable t, Fractional a) => t [a] -> [a] -> [a]
solve_lower ma z = fst $ foldl
 (\(mx,(zf:ze)) row ->
   let (sum, (num:_)) = dropProdSum mx row
   in (mx ++ [(zf-sum)/num], ze))
 ([],z) ma
                  
dropProdSum :: Num a => [a] -> [a] -> (a, [a])
dropProdSum [] ys = (0, ys)
dropProdSum (x:xs) (y:ys) = (x*y + a, b)
  where (a,b) = dropProdSum xs ys