{-# OPTIONS_GHC -O2 -optc-O2 -fglasgow-exts -fbang-patterns #-}

import Control.Monad
import Control.Exception
import Data.List
import Data.Array.IO
import Data.Array.Unboxed
import System
import System.IO.Unsafe
import System.Random
import Data.Array.Base 

matrixMultUnsafe n a b = unsafePerformIO $ 
  do
  c <- newArray_ ((1,1),(n,n)) :: IO (IOUArray (Int,Int) Double)
  let
    f !i !j !k !s | (k==n) = s
    f !i !j !k !s = f i j (k+1) $ s + (a`unsafeAt`(n*i+k))*(b`unsafeAt`(n*k+j))
    jloop !i !j | (j==n) = return()
    jloop !i !j = do unsafeWrite c (i*n+j) (f i j 0 0) ; jloop i (j+1)
    iloop !i | (i==n) = return()
    iloop !i = do jloop i 0; iloop (i+1)
  iloop 0
  unsafeFreeze c

matrixMultSafe n a b = unsafePerformIO $ 
  do
  c <- newArray_ ((1,1),(n,n)) :: IO (IOUArray (Int,Int) Double)
  let
    f !i !j !k !s | (k>n)     = s
    f !i !j !k !s = f i j (k+1) $ s + (a!(i,k))*(b!(k,j))
    jloop !i !j | (j>n) = return()
    jloop !i !j = do writeArray c (i,j) (f i j 1 0) ; jloop i (j+1)
    iloop !i | (i>n) = return()
    iloop !i = do jloop i 1; iloop (i+1)
  iloop 1
  unsafeFreeze c

gaussElimUnsafe matrix =
  do
  ((i1,j1),(m,n)) <- getBounds matrix
  gaussElimUnsafe' matrix (m-i1+1) (n-j1+1)

gaussElimUnsafe' matrix m n = doColumn 0 0
  where
    doColumn !i !j | (i==m||j==n)  = return()
    doColumn !i !j = 
      do 
      (pivotRow,pivotVal) <- findPivot i j
      if nearZero pivotVal
        then doColumn i (j+1)
        else 
          do 
          swapRows i pivotRow
          divideRow i pivotVal
          subtractRows i j
          doColumn (i+1) (j+1)

    findPivot !i !j = f i (i,0)
      where 
        f !i (!maxi,!maxe) | (i==m) = return (maxi,maxe)
        f !i (!maxi,!maxe) = 
          do 
          e <- unsafeRead matrix (i*n+j)
          f (i+1) $ if abs e > abs maxe then (i,e) else (maxi,maxe)

    swapRows !ia !ib = f 0
      where
        f !j | (j==n) = return ()
        f !j =
          do 
          ea <- unsafeRead matrix (ia*n+j)
          eb <- unsafeRead matrix (ib*n+j)
          unsafeWrite matrix (ia*n+j) eb
          unsafeWrite matrix (ib*n+j) ea
          f (j+1)

    divideRow !i !s = f 0
      where
        f !j | (j==n) = return ()
        f !j =
          do
          e <- unsafeRead matrix (i*n+j)
          unsafeWrite matrix (i*n+j) (e/s)
          f (j+1)

    subtractRows !i !j = f 0
      where
        f !u | (u==m) = return ()
        f !u | (u==i) = f (u+1)
        f !u = 
          do 
          s <- unsafeRead matrix (u*n+j)
          g s u 0
          f (u+1)

        g  _  _ !j | (j==n) = return ()
        g !s !u !j = 
          do 
          ei <- unsafeRead matrix (i*n+j)
          eu <- unsafeRead matrix (u*n+j)
          unsafeWrite matrix (u*n+j) (eu - s*ei)
          g s u (j+1)

--------------------------------------------------

gaussElimSafe matrix =
  do
  bnds <- getBounds matrix
  gaussElimSafe' matrix bnds

gaussElimSafe' matrix ((i1,j1),(m,n)) = doColumn i1 j1
  where
    doColumn !i !j | (i>m||j>n) = return()
    doColumn !i !j = 
      do 
      (pivotRow,pivotVal) <- findPivot i j
      if nearZero pivotVal
        then doColumn i (j+1)
        else 
          do 
          swapRows i pivotRow
          divideRow i pivotVal
          subtractRows i j
          doColumn (i+1) (j+1)

    findPivot !i !j = f i (i,0)
      where 
        f !i (!maxi,!maxe) | i>m = return (maxi,maxe)
        f !i (!maxi,!maxe) = 
          do 
          e <- readArray matrix (i,j)
          f (i+1) $ if abs e > abs maxe then (i,e) else (maxi,maxe)

    swapRows !ia !ib = f j1
      where
        f !j | j>n = return ()
        f !j =
          do 
          ea <- readArray matrix (ia,j)
          eb <- readArray matrix (ib,j)
          writeArray matrix (ia,j) eb
          writeArray matrix (ib,j) ea
          f (j+1)

    divideRow !i !s = f j1
      where
        f !j | j>n = return ()
        f !j =
          do 
          e <- readArray matrix (i,j)
          writeArray matrix (i,j) (e/s)
          f (j+1)

    subtractRows !i !j = f i1
      where
        f !u | u>m  = return ()
        f !u | u==i = f (u+1)
        f !u = 
          do 
          s <- readArray matrix (u,j)
          g s u j1
          f (u+1)

        g  _  _ !j | j>n = return ()
        g !s !u !j = 
          do 
          ei <- readArray matrix (i,j)
          eu <- readArray matrix (u,j)
          writeArray matrix (u,j) (eu - s*ei)
          g s u (j+1)

------------------------------------------------------

gaussElim2Unsafe m n matrix =
  do
  _ <- fold1M doColumn [0..n-1]
  return () --matrix
  where
    doColumn i j | i==m = return i
    doColumn i j = 
      do (pivotRow,pivotVal) <- findPivot i j
         if nearZero pivotVal
            then return i
            else do swapRows i pivotRow
                    divideRow i pivotVal
                    mapM_ (\i' -> do e <- unsafeRead matrix (i'*n+j); subtractRow i (e,i')) [0..m-1] 
                    return (i+1)

    findPivot i j =
      do pivotRow <- fold1M
           (\ ia ib ->
             do ea <- unsafeRead matrix (ia*n+j)
                eb <- unsafeRead matrix (ib*n+j)
                if abs ea > abs eb
                    then return ia
                    else return ib
           ) [i..m-1]
         pivotVal <- unsafeRead matrix (pivotRow*n+j)
         return (pivotRow,pivotVal)

    swapRows ia ib = unless (ia == ib) $ mapM_ f [0..n-1]
      where f j = do ea <- unsafeRead matrix (ia*n+j)
                     eb <- unsafeRead matrix (ib*n+j)
                     unsafeWrite matrix (ia*n+j) eb
                     unsafeWrite matrix (ib*n+j) ea

    -- subtract s*row(ia) from row(ib)
    subtractRow ia (s,ib) = unless (ia == ib) $ mapM_ f [0..n-1]
      where f j = do ea <- unsafeRead matrix (ia*n+j)
                     eb <- unsafeRead matrix (ib*n+j)
                     unsafeWrite matrix (ib*n+j) (eb - s*ea)

    --divide row(i) by s
    divideRow i s = mapM_ f [0..n-1]
      where f j = do e <- unsafeRead matrix (i*n+j)
                     unsafeWrite matrix (i*n+j) (e/s)
----------------------------------------------------------------------


gaussElim2Safe matrix ((i1,j1),(m,n)) =
  do
  _ <- fold1M doColumn [j1..n]
  return () --matrix
  where
    doColumn i j | i > m = return i
    doColumn i j = 
      do (pivotRow,pivotVal) <- findPivot i j
         if nearZero pivotVal
            then return i
            else do swapRows i pivotRow
                    divideRow i pivotVal
                    mapM_ (\i' -> do e <- readArray matrix (i',j); subtractRow i (e,i')) [i1..m]
                    return (i+1)

    findPivot i j =
      do pivotRow <- fold1M
           (\ ia ib ->
             do ea <- readArray matrix (ia,j)
                eb <- readArray matrix (ib,j)
                if abs ea > abs eb
                    then return ia
                    else return ib
           ) [i..m]
         pivotVal <- readArray matrix (pivotRow,j)
         return (pivotRow,pivotVal)

    swapRows ia ib = unless (ia == ib) $ mapM_ f [j1..n]
      where f j = do ea <- readArray matrix (ia,j)
                     eb <- readArray matrix (ib,j)
                     writeArray matrix (ia,j) eb
                     writeArray matrix (ib,j) ea

    -- subtract s*row(ia) from row(ib)
    subtractRow ia (s,ib) = unless (ia == ib) $ mapM_ f [j1..n]
      where f j = do ea <- readArray matrix (ia,j)
                     eb <- readArray matrix (ib,j)
                     writeArray matrix (ib,j) (eb - s*ea)

    --divide row(i) by s
    divideRow i s = mapM_ f [j1..n]
      where f j = do e <- readArray matrix (i,j)
                     writeArray matrix (i,j) (e/s)
---------------------------------------------------------------------

fold1M f xs = foldM f (head xs) xs
fold1M_ f xs = fold1M f xs >> return ()





nearZero x = abs x < 1e-5

numItrs = 100000

main =
  do
  rngs <- sequence (replicate numItrs newStdGen)
  
  putStrLn "mulMatrixUnsafe=================="

  forM_ rngs $ \rng ->
    do
    let
      xs = randoms rng
      a = listArray ((1,1),(4,4)) (take 16 xs) :: UArray (Int,Int) Double 
      b = listArray ((1,1),(4,4)) (drop 16 $ take 32 xs) :: UArray (Int,Int) Double 
      c = matrixMultUnsafe 4 a a :: UArray (Int,Int) Double
    evaluate $ c`seq`c

  putStrLn "mulMatrixSafe=================="

  forM_ rngs $ \rng ->
    do
    let
      xs = randoms rng
      a = listArray ((1,1),(4,4)) (take 16 xs) :: UArray (Int,Int) Double 
      b = listArray ((1,1),(4,4)) (drop 16 $ take 32 xs) :: UArray (Int,Int) Double 
      c = matrixMultSafe 4 a a :: UArray (Int,Int) Double
    evaluate $ c`seq`c

  putStrLn "matmult==================="
  forM_ rngs $ \rng ->
    do
    let 
      (a1:a2:a3:a4:a5:a6:a7:a8:a9:a10:a11:a12:a13:a14:a15:a16:_) = randoms rng
      a = [[a1,a2,a3,a4],[a5,a6,a7,a8],[a9,a10,a11,a12],[a13,a14,a15,a16]]
      (b1:b2:b3:b4:b5:b6:b7:b8:b9:b10:b11:b12:b13:b14:b15:b16:_) = randoms rng
      b = [[b1,b2,b3,b4],[b5,b6,b7,b8],[b9,b10,b11,b12],[b13,b14,b15,b16]]
      c = {-# SCC "matmult" #-} matmult a b
    evaluate $ c`seq`c

  putStrLn "gaussElimSafe=================="

  forM_ rngs $ \rng ->
    do
    a <- makeMatrix 4 rng
    gaussElimSafe a

  putStrLn "gaussElimUnsafe=================="

  forM_ rngs $ \rng ->
    do
    a <- makeMatrix 4 rng
    gaussElimUnsafe a
  
  putStrLn "gaussElim2Safe=================="

  forM_ rngs $ \rng ->
    do
    a <- makeMatrix 4 rng
    gaussElim2Safe a ((1,1),(4,8))

  putStrLn "gaussElim2Unsafe=================="

  forM_ rngs $ \rng ->
    do
    a <- makeMatrix 4 rng
    gaussElim2Unsafe 4 8 a

  putStrLn "Mirko's====================="

  forM_ rngs $ \rng -> 
    do
    let 
      (a1:a2:a3:a4:a5:a6:a7:a8:a9:a10:a11:a12:a13:a14:a15:a16:_) = randoms rng
      m = [[a1,a2,a3,a4],[a5,a6,a7,a8],[a9,a10,a11,a12],[a13,a14,a15,a16]]
      mi = {-# SCC "Mirko's" #-} (case inverse m of Just m' -> m' ; Nothing -> error "couldn't invert")
    evaluate $ mi`seq`mi



makeMatrix n rng =
  do
  let 
    m_ = listArray ((1,1),(n,n)) (randoms rng) :: (UArray (Int,Int) Double)
  m <- newArray_ ((1,1),(n,2*n)) :: IO (IOUArray (Int,Int) Double)
  forM_ [1..n] $ \i ->
    do
    forM_ [1..n] $ \j -> writeArray m (i,j) (m_ ! (i,j))
    forM_ [n+1..2*n] $ \j -> writeArray m (i,j) (if j==i+n then 1 else 0)
  return m

printMatrix :: IOUArray (Int,Int) Double -> IO ()
printMatrix m =
  do
  ((r1,c1),(rm,cn)) <- getBounds m
  forM_ [ (i,j) | i<-[r1..rm],j<-[c1..cn] ] $ \(i,j) ->
    do
    a <- readArray m (i,j)
    putStr $ (show a) ++ " " ++ (if j == cn then "\n" else "")
  putStr "\n"







-------------------Mirko's code------------------

type Matrix a = [[a]]

{-# NOINLINE inverse #-}
inverse :: (Monad m, Fractional a) => Matrix a -> m (Matrix a)
inverse m
    | null m       = fail "empty matrix" 
    | any (/=w) ws = fail "matrix not rectangular"
    | h /= w       = fail "matrix not quadractic"
    | otherwise    = liftM (map (drop h)) . clean . gauss $ e
    where h    = length m
	  w:ws = map length m
	  e    = zipWith (++) m $ munit h

munit n = 
    let k = n-1
	rep = flip replicate 0
    in [0..k] >>= \ i -> [rep i ++ [1] ++ rep (k-i)]

gauss m = 
    flip (maybe m) (pivot id m) $ \ (f,row:rows) ->
    map f . (row:) . map (0:) . gauss . map tail $ rows

pivot adjust m
    | all null m = fail "no pivot at all"
    | otherwise  = mplus (pivot1 m >>= return . (,) adjust)
		         (pivot ((0:) . adjust) (map tail m))

pivot1 (row@(0:_):rows) = pivot1 rows >>= \ (r:rs) -> return (r:row:rs)
pivot1 (row@(p:_):rows) = return . (first:) . map add $ rows
    where first         = map (/p) row
          add r@(x:_)   = sub r x first
pivot1 _                = fail "no pivot in first colum"

sub v c = zipWith (-) v . map (*c)

clean (row@(1:rt):rows) = 
    liftM (map (0:)) (clean . map tail $ rows) >>= \ crows ->
    return $ (foldl ( \ v (a,w) -> sub v a w ) row $ zip rt crows) : crows
clean (    (_: _):_   ) = fail "matrix singular"
clean m                 = return m


{-# NOINLINE matmult #-}
matmult a b = map (\v -> map (sum . (zipWith (*) v)) (transpose b)) a