--module MI where

import HUnit
import Control.Arrow
import Monad

data Frac = Frac Integer Integer
--            Frac Frac
--          | FracInt Int
--          | FracSqrt Int

instance Show Frac where
    show (Frac a 1) = show a
    show (Frac 0 _) = "0"
    show (Frac a b) = show a ++ "/" ++ show b
--    show (FracInt 0) = "0"
--    show (FracInt a) = show a
--    show (FracSqrt a) = "[a]"

instance Eq Frac where
    (==) (Frac a b) (Frac a' b') = a == a' && b == b'
--    (==) (FracInt a) (FracInt b) = a == b
--    (==) (FracInt _) (FracSqrt _) = False
--    (==) a@(FracSqrt _) b@(FracInt _) = b == a
--    (==) (FracSqrt a) (FracSqrt b) = a == b

instance Num Frac where
    (+) (Frac a b) (Frac a' b') = (a*b'+b*a')%(b'*b)
--    (+) (FracInt a) (FracInt b) = FracInt (a + b)
--    (+) (FracSqrt a) (FracSqrt b) = error "undefined"
--    (+) (FracInt a) (FracSqrt b) = error "undefined"

    (*) (Frac a b) (Frac a' b') = (a*a')%(b*b')
--    (*) (FracInt a) (FracInt b) = FracInt (a * b)
--    (*) (FracInt a) b@(FracSqrt _) = (FracSqrt (a*a))*b
--    (*) a@(FracSqrt _) b@(FracInt _) = b * a
--    (*) (FracSqrt a) (FracSqrt b) = FracSqrt (a * b)

    abs (Frac a b) = (abs a)%(abs b)
--    abs (FracInt a) = FracInt $ abs a
--    abs a@(FracSqrt _) = a

    fromInteger = (%1)

    signum (Frac a b) = fromInteger ((signum a)*(signum b))
--    signum (FracInt a) = FracInt $ signum a
--    signum (FracSqrt _) = 1

    negate (Frac a b) = (-a)%b
--    negate (FracInt a) = FracInt $ negate a
--    negate a@(FracSqrt _) = Frac a (FracInt (-1))

data Matrix = Matrix [[Frac]]
            deriving Show

--instance Show Matrix where
--    show (Matrix rows) = concatMap (\cells -> concatMap show cells) rows

instance Eq Matrix where
    (==) (Matrix rows) (Matrix rows') =
        foldl (&&) True (zipWith (==) rows rows')

dim :: Matrix -> (Int, Int)
dim (Matrix []) = (0,0)
dim (Matrix rows) = ((length rows), (length (rows!!0)))

dimRow :: Matrix -> Int
dimRow = length . mUnpack

dimCol :: Matrix -> Int
dimCol (Matrix []) = 0
dimCol m = (length . (!!0) . mUnpack) m

sameDim :: Matrix -> Matrix -> Bool
sameDim a b = (dim a) == (dim b)

transpose :: Matrix -> Matrix
transpose = Matrix . (foldr (zipWith (:)) (repeat [])) . mUnpack

instance Num Matrix where
    (+) a b | not $ sameDim a b =
                error ("+: Matrices must have the same dimension, given "
                       ++ show a ++ show b)
    (+) (Matrix rows) (Matrix rows') =
        Matrix $ map (uncurry $ zipWith (+))
                   (zip rows rows')
    (-) a b | not $ sameDim a b =
                error ("-: Matrices must have the same dimension, given "
                       ++ show a ++ show b)
    (-) a b = a + (negate b)

    (*) m1 m2 =
        if (fst $ dim m1) == 1 && (snd $ dim m2) == 1 then
            Matrix [[sum $ zipWith (*)
                       ((mUnpack m1)!!0)
                       ((mUnpack $ transpose m2)!!0)]]
        else
            mRowMap (\row -> mColMap ((mRowHead row) *) m2) m1
    abs = error "undefined operation"
    fromInteger i = Matrix [[fromInteger i]]
    signum = error "undefined operation"
    negate (Matrix rows) = Matrix (map (\cells -> map negate cells) rows)

(%) :: Integer -> Integer -> Frac
(%) _ 0 = error "Divisor can't be 0"
(%) a b =
    Frac (sig*((abs a) `div` g)) ((abs b) `div` g)
    where g = gcd a b
          sig = (signum a)*(signum b)

--matrix' :: [[Frac]]
--matrix' = [[1,2,3],
--          [4,5,6],
--          [7,8,9]]
--
--rh :: [[Frac]]
--rh = [[1,0,0],
--      [0,1,0],
--      [0,0,1]]
--
scale :: Frac -> Matrix -> Matrix
scale n m = Matrix $ map (map (*n)) (mUnpack m)

cDiv :: Frac -> Frac -> Frac
cDiv (Frac a b) (Frac a' b') = (a*b')%(b*a')

-- Gets a row matrix out of a matrix
mRow :: Int -> Matrix -> Matrix
mRow i (Matrix rows) = Matrix [rows!!i]

mCell :: (Int,Int) -> Matrix -> Frac
mCell (row,col) (Matrix rows) = rows!!row!!col

emptyMatrix :: Matrix
emptyMatrix = Matrix [[]]

mColAppend  :: Matrix -> Matrix -> Matrix
mColAppend m1 m2 | m2 == emptyMatrix = m1
mColAppend m1 m2 | m1 == emptyMatrix = m2
mColAppend (Matrix r1) (Matrix r2) = Matrix $ zipWith (++) r1 r2

mRowAppend :: Matrix -> Matrix -> Matrix
mRowAppend m1 m2 | m2 == emptyMatrix = m1
mRowAppend m1 m2 | m1 == emptyMatrix = m2
mRowAppend (Matrix r1) (Matrix r2) = Matrix (r1++r2)

mUnpack :: Matrix -> [[Frac]]
mUnpack (Matrix rows) = rows

mRowHead :: Matrix -> Matrix
mRowHead = mRow 0

mRowTail :: Matrix -> Matrix
mRowTail = Matrix . tail . mUnpack

mColHead :: Matrix -> Matrix
mColHead = Matrix . (map (take 1)) . mUnpack

mColTail ::  Matrix -> Matrix
mColTail = Matrix . (map tail) . mUnpack

mRowMap :: (Matrix -> Matrix) -> Matrix -> Matrix
mRowMap f (Matrix rows) =
    foldl mRowAppend emptyMatrix
              (map (\row -> f (Matrix [row])) rows)

mColMap :: (Matrix -> Matrix) -> Matrix -> Matrix
mColMap = onTranspose . mRowMap . onTranspose
    where onTranspose f = transpose . f . transpose
--    where onTranspose = (transpose .) . (. transpose)
-- Alternative:
--mColMap f m = transpose $ mRowMap (transpose.f.transpose) (transpose m)
-- Alternative:
--    if mColHead m == emptyMatrix then
--        emptyMatrix
--    else
--        (f $ mColHead m) `mColAppend` (mColMap f (mColTail m))

(-|-) :: Matrix -> Matrix -> Frac
(-|-) m1 m2 | not $ sameDim m1 m2 = error "Matrices must both be mx1"
(-|-) m1 _ | (snd $ dim m1) /= 1 = error "Matrices must be mx1"
(-|-) m1 m2 =
    sum $ zipWith (*) (f m1) (f m2)
    where f m = (mUnpack $ transpose m) !! 0

identityMatrix :: Int -> Matrix
identityMatrix 0 = Matrix [[]]
identityMatrix n =
    mRowAppend
     (Matrix [1 : replicate (n-1) 0])
     (mColAppend
      (Matrix $ replicate (n-1) [0])
      (identityMatrix (n-1)))

--identityMatrix n =
--    foldl
--     (\m n ->
--      (mRowAppend
--       (Matrix [1 : replicate (n-1) [0]])
--       (mColAppend
--        (Matrix $ replicate (n-1) [0])
--        m)))
--     Matrix [[]]
--     (reverse [0..n])

flipMatrix :: Matrix -> Matrix
flipMatrix m = Matrix $ reverse (map reverse (mUnpack m))

simpleGauss :: Matrix -> Matrix
simpleGauss m@(Matrix [_]) = m
simpleGauss m =
    let m' = mRowMap
             (\r -> (mColTail r) -
                    (scale ((mCell (0,0) r) `cDiv` (mCell (0,0) m))
                               (mColTail $ mRowHead m)))
             (mRowTail m) in
    (mRowAppend
     (mRowHead m)
     (mColAppend
      (Matrix $ replicate (dimCol m) [0])
      (simpleGauss m')))


splitCol :: Matrix -> (Matrix,Matrix)
splitCol m@(Matrix rows) =
    let dim' = ((dimCol m) `div` 2) in
    join (***) Matrix
             (map (take dim') rows,
              map (drop dim') rows)

scaleToOne :: Matrix -> Matrix
scaleToOne m = mRowMap (\row -> scale (1 `cDiv` (firstNonZero row)) row) m
    where firstNonZero = head . (filter (/=0)) . head . mUnpack


invertMatrix :: Matrix -> Matrix
invertMatrix m =
    let m' = mColAppend m (identityMatrix $ dimRow m)
        aux :: Matrix -> Matrix
        aux = snd . splitCol . scaleToOne . flipAppend . splitCol . simpleGauss . flipAppend . splitCol . simpleGauss
        flipAppend :: (Matrix, Matrix) -> Matrix
        flipAppend m'' = uncurry mColAppend (join (***) flipMatrix m'') in
    aux m'

tests :: Test
tests = test [
         "/" ~: "1/4"
                 ~=? (show (1%4)),
         "+" ~: "1"
                 ~=? (show ((1%2) + (1%2))),
         "sig" ~: "-1/3"
                   ~=? (show ((-1)%3)),
         "sig2" ~: "-1/3"
                    ~=? (show (1%(-3))),
         "div" ~: "1/3"
               ~=? (show (5%15)),
         "dim" ~: (show (2::Int,3::Int))
               ~=? (show $ dim $ Matrix [[1, 2, 3],[1, 2, 3]]),
         "dim0" ~: (show (0::Int,0::Int))
               ~=? (show $ dim $ Matrix []),
         "matrix +" ~: (show $ Matrix [[5,7,9],[7,7,7]])
                        ~=? (show $ (Matrix [[1,2,3],[3,2,1]]) +
                                      (Matrix [[4,5,6], [4,5,6]])),
         "matrix negate" ~: (show $ Matrix [[1,-2],[-3,4]])
                         ~=? (show $ negate (Matrix [[-1,2],[3,-4]])),
         "scale" ~: (show $ Matrix [[2,4],[6,8]])
                     ~=? (show $ scale 2 (Matrix [[1,2],[3,4]])),
         "cDiv" ~: (show (3%1))
                        ~=? (show ((3%1) `cDiv` (1%1))),
         "mColTail" ~: (show $ Matrix [[2,3],[5,6]])
                        ~=? (show $ mColTail $ Matrix [[1,2,3],[4,5,6]]),
         "*" ~: (show $ Matrix [[19,22],[43,50]])
                 ~=? (show ((Matrix [[1,2],[3,4]]) * (Matrix [[5,6],[7,8]]))),
         "transpose" ~: (show $ Matrix [[1,2],[3,4]])
                     ~=? (show $ transpose $ Matrix [[1,3],[2,4]]),
         --
         -- / 1 2 \ + / 1 1 \ = / 2 3 \
         -- \ 3 4 / + \ 2 2 / = \ 5 6 /
         "mColMap" ~: (show $ Matrix [[2,3],[5,6]])
                   ~=? (show $ mColMap (+(Matrix [[1],[2]]))
                                 (Matrix [[1,2],[3,4]])),

         "scalar" ~: "11"
                  ~=? (show $ (Matrix [[1],[2]]) -|- (Matrix [[3],[4]])),

         "identityMatrix" ~: (show $ Matrix [[1,0,0],[0,1,0],[0,0,1]])
                              ~=? (show $ identityMatrix 3),
         "invertMatrix" ~: (show $ Matrix [[-1,1%3,1%3],[0,1%3,-2%3],[2%3,-1%3,1%3]])
                        ~=? (show $ invertMatrix $ Matrix
                                      [[1,2,3],[4,5,6],[2,1,3]]),

         "flipMatrix" ~: (show $ Matrix [[6,5,4],[3,2,1]])
                      ~=? (show $ flipMatrix $ Matrix [[1,2,3],[4,5,6]]),
         "scaleToOne" ~: (show $ Matrix [[0,1,2%3],[1,3,8]])
                          ~=? (show $ scaleToOne $ Matrix [[0,3,2], [1,3,8]])
        ]

--m = Matrix [[1,2,3],[4,5,6],[2,1,3]]


main :: IO Counts
main =
    runTestTT tests;