#!/usr/bin/env nix-script
#!>haskell
#! haskell | vector

{- Solve a linear system by Gauss elimination. -}

import Data.Vector (Vector, cons, snoc, empty)
import qualified Data.Vector as V

type Matrix a = Vector (Vector a)

decons :: Vector a -> (a, Vector a)
decons x = (V.head x, V.tail x)

asList :: ([a] -> [b]) -> (Vector a -> Vector b)
asList f = V.fromList . f . V.toList

fromLists :: [[a]] -> Matrix a
fromLists = V.fromList . map V.fromList


gauss :: (Eq a, Fractional a) => Matrix a -> Vector a -> Vector a
gauss a b = x
  where b' = fmap pure b
        a' = V.zipWith (V.++) a b'
        x  = resubstitute (triangular a')


triangular :: (Eq a, Fractional a) => Matrix a -> Matrix a
triangular m
  | null m    = empty
  | otherwise = cons r (triangular rs')
  where (r, rs) = decons (rotatePivot m)
        rs' = fmap f rs
        f bs
          | (V.head bs) == 0 = V.drop 1 bs
          | otherwise        = V.drop 1 $ V.zipWith (-) (fmap (*c) bs) r
          where  c = (V.head r)/(V.head bs)


rotatePivot :: (Eq a, Fractional a) => Matrix a -> Matrix a
rotatePivot m
  | (V.head r) /= 0 = m
  | otherwise       = rotatePivot (snoc rs r)
  where (r,rs) = decons m


resubstitute :: (Num a, Fractional a) => Matrix a -> Vector a
resubstitute = V.reverse . f . V.reverse . fmap V.reverse
  where f m
          | null m     = empty
          | otherwise  = cons x (f rs')
          where (r,rs) = decons m
                x      = (V.head r)/(V.last r)
                rs'    = fmap (asList subUnknown) rs
                subUnknown (a1:(a2:as')) = ((a1-x*a2):as')