{-# OPTIONS -fglasgow-exts #-}

import Test.QuickCheck hiding ( vector )
import qualified Test.QuickCheck as QC 
      
import Control.Monad

import Data.Maybe          ( fromJust )
import System.IO
import System.IO.Unsafe
import System.Directory    ( getTemporaryDirectory, removeFile )
import System.IO.Error     ( catch )
import Control.Exception   ( finally )
        
import System.IO.MatrixMarket


data Tensor a = DV Int [a]
              | SV Int Int [(Int,a)]
              | DM MatrixType (Int,Int) [a]
              | SM MatrixType (Int,Int) Int [((Int,Int), a)]
    deriving (Eq, Show) 

newtype Comment = C String deriving Show

instance Arbitrary a => Arbitrary (Tensor a) where
    arbitrary = oneof [ vector      >>= \(n,xs)         -> return $ DV n xs
                      , coordVector >>= \(n,nz,ixs)     -> return $ SV n nz ixs
                      , matrix      >>= \(t,mn,xs)      -> return $ DM t mn xs
                      , coordMatrix >>= \(t,mn,nz,ijxs) -> return $ SM t mn nz ijxs
                      ]
                      
    coarbitrary = undefined

instance Arbitrary Comment where
    arbitrary = do
        str <- elements [ ""
                        , "\nleading space"
                        , "trailing space\n"
                        , "more\nthan\none\nline"
                        , "gap\n\n\nin\n\nthe middle"
                        ]
        return (C str)
        
    coarbitrary = undefined
        

instance Arbitrary Field where
    arbitrary   = elements [ Real, Complex, Integer, Pattern ]
    coarbitrary = undefined

instance Arbitrary MatrixType where
    arbitrary   = elements [ General, Symmetric, Hermitian, Skew ]
    coarbitrary = undefined

vector :: (Arbitrary a) => Gen (Int, [a])
vector = sized $ \n -> do
    xs <- QC.vector n
    return (n, xs)
    
coordVector :: (Arbitrary a) => Gen (Int, Int, [(Int,a)])
coordVector = sized $ \n -> do
    nz <- choose (0,n)
    is <- replicateM nz (choose (1,n))
    xs <- QC.vector nz
    let ixs = zip is xs
        
    return (n, nz, ixs)

matrix :: (Arbitrary a) => Gen (MatrixType, (Int,Int), [a])
matrix = sized $ \m -> do
    n  <- choose (0,2*m)
    xs <- QC.vector (m*n)
    t  <- arbitrary
    return $ (t, (m,n), xs)

coordMatrix :: (Arbitrary a) => Gen (MatrixType, (Int,Int), Int, [((Int,Int), a)])
coordMatrix = sized $ \m -> do
    n   <- choose (0,2*m)        
    nz  <- choose (0,m*n)
    is  <- replicateM nz (choose (1,m))
    js  <- replicateM nz (choose (1,n))
    xs  <- QC.vector nz
    let ijs  = zip is js
        ijxs = zip ijs xs        
    attr <- arbitrary
    
    return (attr, (m,n), nz, ijxs)

seqTensor :: Field -> Tensor b -> a -> a
seqTensor Pattern x a =
    case x of
        (DV n xs)           -> n `seq` a
        (SV n nz ixs)       -> n `seq` nz `seq` (length ixs) `seq` a
        (DM t (m,n) _)      -> t `seq` m `seq` n `seq` a
        (SM t (m,n) nz ixs) -> t `seq` m `seq` n `seq` nz `seq` (length ixs) `seq` a

seqTensor _ x a =        
    case x of
        (DV n xs)           -> n `seq` (length xs) `seq` a
        (SV n nz ixs)       -> n `seq` nz `seq` (length ixs) `seq` a
        (DM t (m,n) xs)     -> t `seq` m `seq` n `seq` (length xs) `seq` a
        (SM t (m,n) nz ixs) -> t `seq` m `seq` n `seq` nz `seq` (length ixs) `seq` a

eqTensor :: (Eq a) => Field -> Tensor a -> Tensor a -> Bool
eqTensor Pattern x y =
    case (x,y) of
        ((DV n _),         (DV n' _)) -> 
            n == n'
        ((SV n nz ixs),    (SV n' nz' ixs')) -> 
            n == n' && nz == nz' && fromCoords ixs == fromCoords ixs'
        ((DM t mn _),      (DM t' mn' _)) -> 
            t == t' && mn == mn'
        ((SM t mn nz ixs), (SM t' mn' nz' ixs')) -> 
            t == t' && mn == mn' && nz == nz' && fromCoords ixs == fromCoords ixs'
        _ -> False
        
    where
        fromCoords = fst . unzip
        
eqTensor _ x y = x == y


hPutTensor :: (Show a) => Handle -> Field -> String -> Tensor a -> IO ()
hPutTensor h field desc x = 
    case x of
        (DV n xs)         -> hPutVectorWithDesc      h desc field n xs
        (SV n nz ixs)     -> hPutCoordVectorWithDesc h desc field n nz ixs
        (DM t mn xs)      -> hPutMatrixWithDesc      h desc field t mn xs
        (SM t mn nz ijxs) -> hPutCoordMatrixWithDesc h desc field t mn nz ijxs
        
hGetTensor :: (Read a) => Handle -> Field -> Tensor a -> IO (Tensor a)
hGetTensor h field x =
    case x of
        (DV _ _)     -> 
            hGetVector h field
            >>= \(n,xs)       -> return $ DV n (fromJust xs)
        (SV _ _ _)   -> 
            hGetCoordVector h field 
            >>= \(n,nz,ixs)   -> return $ SV n nz (fromCoords ixs)
        (DM t _ _)   -> 
            hGetMatrix h field t 
            >>= \(mn,xs)      -> return $ DM t mn (fromJust xs)
        (SM t _ _ _) -> 
            hGetCoordMatrix h field t 
            >>= \(mn,nz,ijxs) -> return $ SM t mn nz (fromCoords ijxs)

    where 
        fromCoords ixs = 
            either (\is -> zip is $ repeat undefined) id ixs

writeAndRead :: (Read a, Show a) => Field -> String -> Tensor a -> Tensor a
writeAndRead field desc x = 
    unsafePerformIO $ 
        withTempFile "mytemp.txt" $ \tempfile temph -> do
            hPutTensor temph field desc x
            hClose temph
            withFile tempfile ReadMode $ \h -> do
                x' <- hGetTensor h field x
                -- Because of the unsafePerformIO, we have to make sure we 
                -- are done reading the file before we return.
                seqTensor field x' (return x')
    
-- Taken from _Real World Haskell_ Chapter 6.
withTempFile :: String -> (FilePath -> Handle -> IO a) -> IO a
withTempFile pattern func =
    do tempdir <- catch (getTemporaryDirectory) (\_ -> return ".")
       (tempfile, temph) <- openTempFile tempdir pattern 

       finally (func tempfile temph) 
               (do hClose temph
                   removeFile tempfile)

type T = Tensor Double
 
prop_write_read_id field (x :: T) (C desc) =
    let x' = writeAndRead field desc x
    in eqTensor field x x'
    
main = do
    quickCheck prop_write_read_id