import Neuronet
import System.Random(randomRIO)
import Numeric.LinearAlgebra (Matrix,Vector,tr,scale,cmap,(#>),randn,toList,fromList,toLists,fromLists,Container)
import Data.List
import Data.List.Split
import Data.Tuple.Extra
import System.Random.Shuffle
import qualified Data.ByteString as BS
import System.IO

-- a single data-sample with input and expected output
type Sample = (Vector Double,Vector Double)

-- a list of samples
type Samples =[Sample]

-- split samples into random batches of n elements
batches :: Int -> Samples -> IO [Samples]
batches n x = do s<-shuffleM x
                 return . filter ((==n).length) $ chunksOf n s

-- get multiple epochs, with the given number of samples each.
epochs :: Int -> Int -> Samples -> IO [[Samples]]
epochs 0 _ _ = return []
epochs e n s = (:) <$> batches n s <*> epochs (e-1) n s

-- train for multiple epochs
training :: Double -> Neuronet -> [[Samples]] -> Neuronet
training r net s = foldl' f net (concat s)
    where f a v = trainBatch r a (fst.unzip $ v) (snd.unzip $ v)

-- test with given samples and return number of correct answers
testing :: (Vector Double -> Vector Double -> Bool) -> Neuronet -> Samples -> Int
testing f net s = length . filter id $ map (\(x,y)->f y (asknet net x)) s

-- finally some learning and testing with MNIST
main = mainMNIST

-- Paths to the unpacked files from http://yann.lecun.com/exdb/mnist/
train_img="train-images-idx3-ubyte"
train_lbl="train-labels-idx1-ubyte"
test_img="t10k-images-idx3-ubyte"
test_lbl="t10k-labels-idx1-ubyte"

mainMNIST =do
        nt <- neuronet [28*28,10,10]
        h1 <- openFile train_img ReadMode
        h2 <- openFile train_lbl ReadMode
        xs <- BS.hGetContents h1 >>= return. map (fromList.map fromIntegral) . chunksOf (28*28) . BS.unpack.BS.drop 16
        ys <- BS.hGetContents h2 >>= return. map (fromList.val.fromIntegral) . BS.unpack. BS.drop 8
        tr <- epochs 1 100 (zip xs ys) >>= return . training (1/100) nt

        print "bye"
    where  zrs= take 9 $ repeat 0
           val x= take x zrs ++ [1::Double] ++ drop x zrs

main2 = do

        easynet<-neuronet [20,10,10]
        samples<-(map ( \(a,b) -> (fromList (val a++val b),fromList $ val $ a+b)) . filter ((<10).uncurry (+))) <$> gensamp 4000
        tstsmpl<-(map ( \(a,b) -> (fromList (val a++val b),fromList $ val $ a+b)) . filter ((<10).uncurry (+))) <$> gensamp 1000
        trained <- epochs 200 100 samples >>= return . training (3/100) easynet
        print $ testing tst trained tstsmpl

    where  zrs= take 9 $ repeat 0.01
           val x= take x zrs ++ [0.99] ++ drop x zrs
           getVal x=snd . last . sort $ zip x [0..9]
           tst y1 y2 = getVal (toList y1) == getVal (toList y2)

gensamp :: Int -> IO [(Int,Int)]
gensamp  0 = return []
gensamp n = (\a b r->(a,b):r) <$> randomRIO (0,9) <*> randomRIO (0,9) <*> gensamp (n-1)

{-
main3 = do

    let numtrain=300000

    easynet<-neuronet [20,15,10]

    samples<-filter ((<10).uncurry (+)) <$> gensamp numtrain
    let testsmp=filter ((<10).uncurry (+))[(a,b)|a<-[0..10],b<-[0..10]]

    let trained=foldl' training easynet samples
    let res = map (\(a,b)->test trained (a,b) == a+b) testsmp
    let right=length $ filter (id) res
    let all = length res

    print $ "trained adding up to 10 with "++ show numtrain ++" random training samples"
    print $ "correct answers during testing: " ++ show right ++ "/" ++ show all ++ " (" ++ show (right*100`div`all) ++ "%)"
    print "what can I add for you?"
    str<-lines<$>getContents
    mapM_ (addit trained) str

    where  zrs= take 9 $ repeat 0.01
           val x= take x zrs ++ [0.99] ++ drop x zrs
           getVal x=snd . last . sort $ zip x [0..9]

           training net (s1,s2) = train 0.25 net (fromList $ val s1++val s2) (fromList $ val $ s1+s2)
           test net (s1,s2) = getVal . toList $ asknet net (fromList $ val s1++val s2)


           addit n x=   do  let w=words x
                            let a=read $ w!!0 :: Int
                            let b=read $ w!!1 :: Int
                            putStrLn $ "I try to add " ++ show a ++ " and " ++ show b
                            putStrLn $ "I think this is = " ++ show (test n (a,b) )
                            putStrLn "what two numbers do you want me to add master?"

                            





showImg a | BS.null a = print "--" 
          | otherwise = do mapM_ putVal $ BS.unpack $ BS.take 28 a
                           putStr "\n"
                           showImg $ BS.drop 28 a

putVal v | v==0 = putStr " "
         | v<50 = putStr "."
         | v<128 = putStr "o"
         | v<160 = putStr "O"
         | otherwise = putStr "0"


main2 = do 
    --              0  1   2   3   4   5   6   7   8   9 
    let m=(4><10)[  0, 100,0,  100,0,  100,0,  100,0,  100
                   ,0, 0,  100,100,0,  0,  100,100,0,  0
                   ,0, 0,  0,  0,  100,100,100,100,0,  0
                   ,0, 0,  0,  0,  0,  0,  0,  0,  100,100]

    let b=vector . take 4 . repeat $ (-50)
                
    let z= repeat 0.01
    let vv x= vector $ take x z ++ [0.99] ++ take (9-x) z
    let v = vv 6
    let z= m #> v + b
    let a= cmap sigmoid z
    print m
    print v
    print z 
    print a
-}