{-# LANGUAGE GADTs,MultiParamTypeClasses #-}
import Data.Graph.Inductive.Graph
import Data.Graph.Inductive.Basic
import Data.Graph.Inductive.PatriciaTree
--Showing
import Data.Graph.Inductive.Query.DFS (topsort)
import Text.PrettyPrint.Boxes -- (render,text,(<+>),vcat)
import Data.List (transpose)
import Data.Maybe (fromJust)
--Showing End
--Test Comonad2 - the double version
import Comonad2
--fgl boxes split

data NN a b = NN (Con a b) (Gr a b)

newtype Con a b = Con (Context a b)

liftC :: (Context a b -> Context c d) -> Con a b -> Con c d
liftC f (Con c) = Con (f c)
dropC :: (Con a b -> Con c d) -> Context a b -> Context c d
dropC f c = c2 where Con c2 = f (Con c)
liftCA :: (a->a) -> Con a b -> Con a b
liftCA f (Con (i,n,a,o)) = Con (i,n,f a,o)
pullC :: Con a b -> Context a b
pullC (Con c)=c
       
contexts :: Graph gr => gr a b -> [Context a b]
contexts = gsel (const True) 
   
dextractN :: NN a b -> Node
dextractN (NN (Con (_,n,_,_)) _)= n

dextractA :: NN a b -> a
dextractA (NN (Con (_,_,a,_)) _)= a

unzipPre :: NN a b -> ([a],[b])    
unzipPre nn= (as,bs) where
    (ns,bs) = unzip $ lpre nn (dextractN nn)
    as = map (fromJust.lab nn) ns
  
instance Graph NN where
    empty = NN undefined empty --TODO carefull!
    isEmpty (NN _ net) = isEmpty net
    match n (NN c net) = (mc,NN c g)
        where (mc,g)=match n net
    mkGraph ns es=NN (Con $ fst $ matchAny newNet) newNet --first edge
        where newNet=mkGraph ns es
    labNodes (NN _ net) = labNodes net
    nodeRange (NN _ net) = nodeRange net
instance DynGraph NN where
    nc & (NN c net)= NN (Con nc) (nc & net)
    
instance DFunctor NN Con where
    dmap = gmap.dropC
    
instance DCopointed NN Con where
    dextract (NN c _) = c
    
instance DComonad NN Con where
    dextend f g@(NN _ nn)= dmap convert g where
        convert c'= f $ NN c' nn
 
main = putStrLn $ show $ dextendX step 10000 t

--example
step :: PseudounipolarNN a b -> Con (Axon a b) (Dendrite a b)
step nn = (liftCA $ (stepAxon.unzipPre) nn) (dextract nn)
    
type Neuron a b = ([(Dendrite a b,Node)]
                  ,Node
                  ,Axon a b
                  ,[(Dendrite a b,Node)])

data Dendrite a b = Dendrite {
    weightFunc :: a -> a
    ,weight :: b
    }
   
data Axon a b where 
     Axon :: ([a] -> a) -> a -> Axon a b
activityFunc (Axon a _)=a
activity (Axon _ a)=a

stepAxon :: ([Axon a b],[Dendrite a b]) -> Axon a b -> Axon a b
stepAxon (as,bs) a= Axon aFunc (aFunc $ zipWith stepAxon' as bs) where
    aFunc = activityFunc a
    stepAxon' = \x y -> (weightFunc y) (activity x)

type PseudounipolarNN a b=NN (Axon a b) (Dendrite a b)

t = mkGraph nz ez :: PseudounipolarNN Double Double

sig :: Double -> Double
sig a =1.0 / (1 + exp (-1*a)) -0.5

nz = zip ns nns
    where   ns = [1..23]
            nns1 = take 10 $ repeat (Axon (const 0.25) 0.25)
            nns= nns1 ++ ( map (Axon $ sig.sum) is )
            is= [-1,-0.9..0] ++ repeat 0
ez= let 
        a2layer=[(x,y) | x<-[1..10],y<-[11..20]
                    , or [x+10==y,x*2==y,x*3==y] ]
        a3layer=a2layer++[(x,y) | x<-[11..20],y<-[21,22]
                    , or [ and [y==22,x>15] , and [y==21,x<=15] ]
                    ]
        a4layer=a3layer++[(21,23),(22,23)]
        zipit (x,y) z=(x,y,z)
        ss = zipWith Dendrite (map (*) ds) ds
            where ds=[-2,-1.8..]
    in
        zipWith zipit a4layer ss
    
--Printing Stuff
instance Show b =>Show (Dendrite a b) 
    where show s= (take 4 $ show $ weight s)

instance Show a => Show (Axon a b) 
    where show (Axon _ activity)= (take 4 $ show $ activity)

instance (Show a, Show b) => Show (NN a b) where 
    show (NN (Con (_,node,_,_)) net) 
       |isEmpty net = "empty"
       |otherwise = render $ hcat top $ map (vcat left) cList
        where cList=transpose $ map (showC.context net) (topsort net)
              showC (i,n,a,o)=(case n of n | n==node -> [text "*"]
                                           |  otherwise -> [text " "])
                   ++[alignHoriz right 5 $ text(show n++": ")
                  ,alignHoriz right 4 $ text $ show a , text " --> "
                  ,text $ show es ,text " with weights: " ,text $ show ws]
                where  (ws,es)=unzip o