{-# LANGUAGE DataKinds #-}
{-# LANGUAGE Arrows, FlexibleContexts #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE DeriveGeneric #-}

module Examples.MultiCamera where

import Control.Arrow hiding ((|||))
import Control.Monad (replicateM, forM)
import Control.Monad.Trans (MonadIO, liftIO, lift)
import Control.Monad.Bayes.Class (MonadSample, MonadCond, logCategorical)
import Control.Monad.Bayes.Sampler (SamplerIO, sampleIO)

import Data.Aeson hiding (Result)
import Data.Maybe (catMaybes)

import GHC.Generics (Generic)

import qualified Data.ByteString.Lazy.Char8 as BS (putStrLn)

import Numeric.LinearAlgebra.Static

import Inference (zdsparticles, zunheap)
import DelayedSampling (MDistr (..), children, state, State, MarginalT (..), SMarginalT(..), distr, typeOfDSDistr, DelayedInfer)
import qualified SymbolicDistr as DS
import DSProg (DeepForce (..), Result (..), Expr' (..), Expr, marginal, zdeepForce)
import Distributions (normal, poisson, sample, observe, bernoulli, replicateIID, passert, factor, uniformIntRange)
import MVDistributions (shuffleList, mvNormal, uniformNSphere)
import Util.ZStream (ZStream)
import qualified Util.ZStream as ZS
import Util.Ref (MonadState, Heap, emptyHeap, readRef)

import Numeric.SpecFunctions (logGamma)
import Numeric.Log (Log (Exp), ln)

import qualified Data.Vector.Storable as V (fromList)

import qualified Metaprob as MP

numCameras :: Int
numCameras = 2

numPoses :: Int
numPoses = 10

type CameraID = Int
type Pose = Int

data TrackG pv = Track
  { posWH :: pv
  , startTime :: Double
  , trackID :: Int
  , identity :: Int
  , camera :: CameraID
  }
  deriving (Generic, Show)

instance ToJSON pv => ToJSON (TrackG pv)

instance DeepForce pv => DeepForce (TrackG pv) where
  type Forced (TrackG pv) = TrackG (Forced pv)
  deepForce (Track pv h t id c) = Track <$> deepForce pv <*> deepForce h <*> pure t <*> pure id <*> pure c
  deepConst (Track pv h t id c) = Track (deepConst pv) (deepConst h) t id c

type STrack = TrackG (Expr (R 4))
type Track = TrackG (R 4)
type MarginalTrack = TrackG (Result (R 4))

type Appearance = Expr (R 10)

tdiff :: Double
tdiff = 1

posWHCovBlocks :: Sym 2 -> Sym 2 -> Sym 4
posWHCovBlocks pcov whcov = sym $
  ((unSym pcov ||| (konst 0 :: Sq 2))
        ===
  ((konst 0 :: Sq 2) ||| unSym whcov))

newTrack'' :: MonadState Heap m => MonadSample m => Double -> Int -> Int -> Int -> m STrack
newTrack'' t tid id c = do
  pv <- DS.sample (DS.mvNormal (Const mu) cov)
  pure (Track pv t tid id c)
  where
  mu = (konst 0 :: R 2) # (konst 1 :: R 2)
  cov = posWHCovBlocks (sym eye :: Sym 2) (0.2 * sym eye :: Sym 2)

newTrack :: MonadState Heap m => MonadSample m => [Appearance] -> Double -> Int -> m (STrack, [Appearance])
newTrack appearances t tid = do
  camera <- sample (uniformIntRange numCameras)
  newAppearance <- sample (bernoulli (1 / (fromIntegral (length appearances) + 1)))
  (appearanceID, appearances') <- if newAppearance
    then do
         app <- DS.sample (DS.mvNormal (Const 0) (sym eye))
         pure (length appearances, appearances ++ [app])
    else do
         appID <- sample (uniformIntRange (length appearances))
         pure (appID, appearances)
  newTrack <- newTrack'' t tid appearanceID camera
  pure (newTrack, appearances')

trackMotion :: MonadState Heap m => MonadSample m => Double -> STrack -> m STrack
trackMotion tdiff track = do
  pv' <- DS.sample (DS.mvNormal (MVMul (Const motionMatrix) (posWH track)) motionCov)
  pure $ track { posWH = pv' }
  where
  motionMatrix = eye :: Sq 4
  motionCov :: Sym 4
  motionCov = sym (konst tdiff) * posWHCovBlocks (sym eye) (sym (konst 0.1 * eye))

trackSurvivalMotion :: MonadState Heap m => MonadSample m => Double -> STrack -> m (Maybe STrack)
trackSurvivalMotion tdiff track = do
  survived <- sample (bernoulli (exp (- tdiff * deathRate)))
  if survived
    then Just <$> trackMotion tdiff track
    else pure Nothing
  where
  deathRate = 0.02

trackMeasurement :: [Appearance] -> STrack -> DS.Distr (Expr (R 4), (R 10, Expr (R 1)))
trackMeasurement appearances track =
  DS.indep (DS.mvNormal (posWH track) (sym eye)) $
  DS.bind uniformNSphere $ \pose ->
    DS.mvNormal (MVMul (Const (row pose)) (appearances !! identity track)) (sym eye)

-- Naive association
associationWithClutter :: forall a b. DeepForce b => DS.Distr [b] -> (a -> DS.Distr b) -> [a] -> DS.Distr [b]
associationWithClutter cluttersD obsD tracks = DS.Distr is iobs undefined where
  is :: MonadState Heap m => MonadSample m => m [b]
  is = do
    clutter <- DS.sample cluttersD
    observations <- mapM (DS.sample . obsD) tracks
    shuffleList (observations ++ clutter)
  defaultObserve :: DelayedInfer m => [Forced b] -> m ()
  defaultObserve observations = do
    passert (length tracks <= length observations)
    observations' <- shuffleList observations
    let (observations'', clutter) = splitAt (length tracks) observations'
    mapM_ (\(t, o) -> DS.observe (obsD t) o) (zip tracks observations'')
    DS.observe cluttersD clutter
  iobs :: DelayedInfer m => [Forced b] -> m ()
  iobs = if True
    then assocWithClutterCustomProposal cluttersD obsD tracks
    else defaultObserve

assocWithClutterCustomProposal :: DelayedInfer m => DeepForce b => DS.Distr [b] -> (a -> DS.Distr b) -> [a] -> [Forced b] -> m ()
assocWithClutterCustomProposal cluttersD obsD allTracks allObservations = go allTracks allObservations where
  numTracks = fromIntegral (length (allTracks))
  numObservations = fromIntegral (length (allObservations))
  go [] clutter = do
    let numClutter = fromIntegral (length clutter)
    factor (- (logGamma (numTracks + numClutter + 1) - logGamma (numClutter + 1)))
    DS.observe cluttersD clutter
  go (track : tracks) [] =
    passert False
  go (track : tracks) observations = do
    let obsDistr = obsD track
    lls <- map Exp <$> mapM (DS.score obsDistr) observations
    let totalLL = sum lls
    let adjLLs = map (/ totalLL) lls
    i <- logCategorical (V.fromList adjLLs)
    let (os1, obs : os2) = splitAt i observations
    DS.observe obsDistr obs
    factor (ln (recip (adjLLs !! i))) -- proposal correction
    go tracks (os1 ++ os2)

tracksMeasurement :: [Appearance] -> [STrack] -> DS.Distr [(Expr (R 4), (R 10, Expr (R 1)))]
tracksMeasurement appearances = associationWithClutter cluttersDistr (trackMeasurement appearances)
  where
  cluttersDistr = DS.replicateIID (poisson amtClutter) clutterDistr
  amtClutter :: Double
  amtClutter = 1
  clutterDistr :: DS.Distr (Expr (R 4), (R 10, Expr (R 1)))
  clutterDistr = DS.indep (DS.mvNormal (Const ((konst 0 :: R 2) # (konst 1 :: R 2))) (sym eye))
    $ DS.bind (mvNormal 0 (sym eye)) (\_ -> DS.mvNormal (Const 0) (sym eye))

-- NOTE: Currently, it is possible that we will generate a new track
-- with the same identity and camera as an existing track.
-- Should consider how to change the model so that this doesn't happen.
tracksMotion :: MonadState Heap m => MonadSample m => Double -> Double -> [STrack] -> Int -> [Appearance] -> m ([STrack], Int, [Appearance])
tracksMotion t tdiff tracks numTracks appearances = do
  liveTracks <- catMaybes <$> mapM (trackSurvivalMotion tdiff) tracks
  numNewTracks <- sample (poisson (birthRate * tdiff))
  (newTracks, appearances') <- generateNewTracks appearances (numTracks + 1) numNewTracks
  pure (liveTracks ++ newTracks, numTracks + numNewTracks, appearances')
  where
  generateNewTracks apps nt 0 = pure ([], apps)
  generateNewTracks apps nt n = do
    (newT, apps') <- newTrack apps t nt
    (\(x, y) -> (newT : x, y)) <$> generateNewTracks apps' (nt + 1) (n - 1)
  birthRate = 0.1

-- Note how extraneous state about the number of tracks and the time is kept hidden
zstepGen :: MonadState Heap m => MonadSample m => ZStream (MP.Gen m) () ([STrack], [(CameraID, R 10, Expr (R 4), Expr (R 1))])
zstepGen = ZS.fromStep stepf initState
  where
  initState :: (Double, [TrackG pv], Int, [Appearance])
  initState = (0, [], 0, [])
  stepf (t, tracks, numTracks, appearances) () = do
    (tracks', newNumTracks, appearances') <- lift (tracksMotion t tdiff tracks numTracks appearances)
    observations <- forM [0 .. numCameras - 1] $ \camID -> do
      obs <- ("obs." ++ show camID) MP.~~ MP.dsPrim (tracksMeasurement appearances' (filter (\t -> camera t == camID) tracks'))
      return [ (camID, pose, pvwh, obsAppear) | (pvwh, (pose, obsAppear)) <- obs ]
    return ((t + tdiff, tracks', newNumTracks, appearances'), (tracks', concat observations))

processObservationsStream :: DelayedInfer m => ZStream m [(CameraID, R 10, R 4, R 1)] [MarginalTrack]
processObservationsStream = proc observations -> do
  (tracks', _) <- MP.zobserving zstepGen -< ((), "obs" MP.|-> MP.obs observations)
  ZS.run -< mapM trackf tracks'
  where
  trackf :: DelayedInfer m => STrack -> m MarginalTrack
  trackf track = do
    Just pv <- marginal (posWH track)
    pure $ track { posWH = pv }

generateGroundTruth :: MonadState Heap m => MonadSample m => ZStream m () ([STrack], [(CameraID, R 10, Expr (R 4), Expr (R 1))])
generateGroundTruth = ZS.liftM MP.sim zstepGen

runMTTPF :: Int -> ZStream SamplerIO () ([Track], [[MarginalTrack]], [(CameraID, R 10, R 4, R 1)])
runMTTPF numParticles = proc () -> do
  (groundTruth, obs) <- zdeepForce generateGroundTruth -< ()
  particles <- zdsparticles numParticles processObservationsStream -< obs
  returnA -< (groundTruth, particles, obs)

runExample :: IO ()
runExample = sampleIO $ ZS.runStream (liftIO . BS.putStrLn . encode) (runMTTPF 100)


{-# LANGUAGE OverloadedStrings #-}
module Analysis.FFT (auralDistance, peakList) where

-- need to create frequency frames for peaks after the FFT has been applied
import Data.Complex
import Data.List
import Data.List.Split
import Data.Ord
import Data.Audio
import Numeric.Extra
import Data.Int
import Data.Array.IArray
import Numeric.Transform.Fourier.FFT

import Control.Parallel.Strategies

import Analysis.FFT_Py

import Types.Common
import qualified Settings as S
import Utils

import Shelly (shelly, silently, run)
import Data.String

auralDistance :: (FilePath, AudioFormat) -> (FilePath, AudioFormat) -> IO Double
auralDistance a1 a2 = do
  ps1 <- peakList a1
  ps2 <- peakList a2
  -- since I am not looking at time shifting filter for now, I can zipWith over time to remove the time domain
  -- NB this means the times need to line up 'exactly', ie this will not work well for real-world examples
  --    but only for contrsucted examples where I apply the filter to the sample myself (eg in audacity)
  --mapM_ plot $ take 5 $ zip (map sort ps1) (map sort ps2)
  let
    distances = zipWith (comparePeaks) (map sort ps1) (map sort ps2)
    average ls = sum ls / genericLength ls
  return $ average distances

plot :: ([Peak], [Peak]) -> IO()
plot (ps1, ps2) = do
  writeFile "tmp1.csv" (listToCSV ps1)
  writeFile "tmp2.csv" (listToCSV ps2)
  results <- shelly $ silently $ run (fromString "gnuplot")
               [ "-p"
               , "plotter.gnuplot" ]
  return ()

comparePeaks ps1 ps2 =
  sum $ zipWith comparePeak ps1 ps2

-- TODO rethink how i do distances...
-- | for a single time slice, do euclid distance treating (freq,amp) as (x,y)
--   should always have the same number of points, but they might not lineup exactly...
--   eg (3,0.1),(4,0.5),(5,0.5) vs (4,0.5),(5,0.5),(6,0.1) should be rated as fairly close
comparePeak :: Peak -> Peak -> Double
comparePeak peak1 peak2 = let
  freq1 = intToDouble $ getFreq peak1
  freq2 = intToDouble $ getFreq peak2
  amp1 = getAmp peak1
  amp2 = getAmp peak2
 in
    {-if amp2 == 0
    then 1000 --penalize zero files
    else -}
    --traceShow ((show freq1)++" "++(show amp1)++", "++(show freq2)++" "++(show amp2) )
    (euclidDistance (freq1,amp1) (freq2,amp2)) / 1000-- + 1/(amp2+0.00001)

-- | break an audio file into time slices and i
--   find the freq peaks that are most predominate for each time slice
peakList :: (FilePath, AudioFormat) -> IO (OverTime (OverFreq Peak))
peakList (fp,a) = do
  peaks <- peakListPython fp
  let normPeaks = normalize peaks
  return $ getMainPeaks normPeaks

-- | Scale so max over all time slices is 1
--   note that this means some time slices will not have a max of 1
--   another option is to normlize so that the integral of the fft is 1
normalize :: [[Peak]] -> [[Peak]]
normalize ps = let
  maxPeak = maximum $ concatMap (map getAmp) ps
  scaleFactor = 1 / maxPeak
 in
  map (map (\(frq,amp) -> (frq, amp* scaleFactor))) ps

{-
--takes wave file and turns it's values into list of Complex Doubles
wavList :: AudioFormat -> OverTime Double
wavList wav = let
    l1 = sampleData wav
  --takes at most a certain time frame
  in take (16384*10) $ elems $ amap (toSample) l1

-- mkFrames takes assocs of audio file and breaks it into 4096 sample (.09s) frames (overlapped by 50%)
mkFrames :: OverTime Double -> OverTime (OverTime Double)
mkFrames list1 =
  if (length list1) < S.frameRes --Settings.framerate
    then []
    else (take S.frameRes list1):(mkFrames (drop S.overlap list1))

-- | performs FFT on a list of samples, and conversts each sample to a list of peaks as the triple (freq,amp,phase)
--   returns a list (in the time domain) of peaks (freq domain)
constellateAll :: OverTime (OverTime Double) -> OverTime (OverFreq Peak)
constellateAll timeSlices = let
    -- what the heck does this do?
    sparsifying n= foldr (.) id (replicate n (remove_every_nth 2))
    ars1 = --sparsifying S.resolution $
               --why not length sample here?
               map (\samples -> listArray (0,(S.frameRes - 1)) samples) timeSlices
    in map ((take (S.frameRes `div` 2)).assocs.constellate.rfft) ars1

-- constellate takes results of FFT and turns it into (amp,phase) at each frequency point
constellate :: Array Int (Complex Double) -> Array Int (Double,Double)
constellate arr1 =
  let list1 = assocs arr1
      polar1 = map (polar.snd) list1 --TODO use amap for array
      in listArray (0,((length list1) - 1)) polar1
-}

-- | with the full set of peaks for a list of frames,
--   we want to essentially apply a filter to each time slice
--   so that we only have the main peaks for that time slice
getMainPeaks :: OverTime (OverFreq Peak) -> OverTime (OverFreq Peak)
getMainPeaks ts =
  map (findBiggestPeaks. freqBins) ts

-- | partition a sample slice's peaks by freq into bins
--   this allows for a bit of flexibilty in freq analysis
--   ie dont need the exact same freqs to be the same
--   sameness of frequency is determined by bin size
--   NB the OverFreq interprtation seems to break down here
freqBins :: OverFreq Peak -> [[Peak]]
freqBins = chunksOf S.binSize

-- | for a given sample slices' peaks in each freq bin,
--   get the biggest (by amp) in each bin
--   then take the biggest numPeaks of those
findBiggestPeaks :: [[Peak]] -> [Peak]
findBiggestPeaks ts = let
   -- first sort each bin, and only take the largest
   -- this eliminates the 'noisy' freqs nearby the peaks
   largestPeaksPerBin = map (last.(sortBy (comparing getAmp))) ts
   -- then we want to get the biggest peaks our of all the loudest bins (or 'freq areas')
   -- so we first sort by loudest 'freq areas'
   largestBinPerSampleSlice = reverse $ sortBy (comparing getAmp) largestPeaksPerBin
  in
   -- to find the S.numPeaks biggest peaks
   take S.numPeaks $ largestBinPerSampleSlice

-- listTriple turns tuple of int and tuple and makes it into a triple
listTriple :: [(a,(b,b))] -> [(a,b,b)]
listTriple xs =
   map (\(x,(y,z)) -> (x,y,z)) xs


{-# LANGUAGE GADTs, NoImplicitPrelude, TypeFamilies,
InstanceSigs, StandaloneDeriving, UndecidableInstances,
ScopedTypeVariables, FlexibleInstances, DataKinds,
FunctionalDependencies, PolyKinds,
TypeOperators, RankNTypes, ImpredicativeTypes, MultiParamTypeClasses,
AllowAmbiguousTypes, TypeApplications, FlexibleContexts #-}
---
module Main where

import qualified Data.Map.Strict as Map
import Prelude hiding ((.), id)
import Linear.Vector
import Data.Tuple
import Data.Complex
import Linear.Epsilon (nearZero)
import Control.Category
import Control.Monad ((<=<))
import GHC.TypeNats
import Data.Proxy
import Diag (mymain)
-- import qualified Text.PrettyPrint.ANSI.Leijen as PP
import qualified Text.PrettyPrint.Boxes as B
-- import Control.Arrow
{-
type MapVec r k = Map.Map k r

vadd :: (Ord k, Num r) => MapVec r k -> MapVec r k -> MapVec r k
vadd = Map.unionWith (+)

smul :: (Ord k, Num r) => k -> MapVec r k -> MapVec r k
smul

-- kmett alread did this stuff
-}

-- http://blog.sigfpe.com/2007/03/monads-vector-spaces-and-quantum.html

data W b a = W { runW :: [(a,b)] } deriving (Eq,Show,Ord)
instance Semigroup (W b a) where
  (W x) <> (W y) = W (x <> y)
instance Monoid (W b a) where
  mempty = W mempty
mapW f (W l) = W $ map (\(a,b) -> (a,f b)) l
instance Functor (W b) where
    fmap f (W a) = W $ map (\(a,p) -> (f a,p)) a

instance Num b => Applicative (W b) where
  pure x = W [(x,1)]
  (W fs) <*> (W xs) = W [(f x, a * b) | (f, a) <- fs, (x, b) <- xs]
  -- the fs is a a vector of rows kind of.
    -- error "I don't feel like it"

instance Num b => Monad (W b) where
   return x = W [(x,1)]
   l >>= f = W $ concatMap (\(W d,p) -> map (\(x,q)->(x,p*q)) d) (runW $ fmap f l)

a .* b = mapW (a*) b

instance (Eq a,Show a,Num b) => Num (W b a) where
     W a + W b = W $ (a ++ b)
     a - b = a + (-1) .* b
     _ * _ = error "Num is annoying"
     abs _ = error "Num is annoying"
     signum _ = error "Num is annoying"
     fromInteger a = if a==0 then W [] else error "fromInteger can only take zero argument"

collect :: (Ord a, Num b) => W b a -> W b a
collect = W . Map.toList . Map.fromListWith (+) . runW

trimZero = W . filter (\(k,v) -> not $ nearZero v) . runW
simplify :: Ord a => Q a -> Q a
simplify = trimZero . collect
-- filter (not . nearZero . snd)

type P a = W Double a

type Q a = W (Complex Double) a

v1 :: Map.Map Bool Double
v1 = zero

boolbase :: [Map.Map Bool Double]
boolbase = basis

e1 = Map.singleton True 1.0
e2 = Map.singleton False 1.0

star = mapW conjugate

{-
There are two types of particles

-}
data FibAnyon = Id | Tau
-- I started using DataKinds and it ended up being a problem.

data A a

data Tau
data Id
data FibTree root leaves where
   TTT :: FibTree Tau l -> FibTree Tau r -> FibTree Tau (l,r)
   ITT :: FibTree Tau l -> FibTree Tau r -> FibTree Id (l,r) -- we would like to enforce that isingany can split into a b, but thisis tough
   TIT :: FibTree Id l -> FibTree Tau r -> FibTree Tau (l,r)
   TTI :: FibTree Tau l -> FibTree Id r -> FibTree Tau (l,r)
   III :: FibTree Id l -> FibTree Id r -> FibTree Id (l,r)
   -- III :: FibTree 'Id l -> FibTree 'Id l' -> FibTree 'Id (l,l') -- maybe shouldn't exist. uneccessary. Does really help for monoidal instance
   TLeaf :: FibTree Tau Tau
   ILeaf :: FibTree Id Id

-- pretty printing would be hella nice
--deriving instance Show (FibTree a b)

instance Show (FibTree a b) where
  show = drawTree
deriving instance Eq (FibTree a b)
--deriving instance Ord (FibTree a b)
{-
r = PP.text "\\"
l = PP.text "/"
ls = PP.group $ (PP.text "  /" <> PP.line <> PP.text " /" <> PP.line <>PP.text "/")
rs = PP.group $ (PP.text "\\" <> PP.line <> PP.text " \\" <> PP.line <>PP.text "  \\")
-}

sp = B.char ' '
l = B.char '/'
r = B.char '\\'
h = B.char '-'
v = B.char '|'
t = B.char 'T'
i = B.char 'I'
rs x =  (r  B.<> sp B.<> sp) B.//
      (sp B.<> r  B.<> x) B.//
      (sp B.<> sp  B.<> r)
ls x =  (sp  B.<> sp B.<> l) B.//
      (x B.<> l  B.<> sp) B.//
      (l B.<> sp  B.<> sp)

{-
drawty = v B.// v B.// v B.// (h B.<> h B.<> h B.<> h)
drawb :: FibTree a b -> B.Box
drawb ILeaf = i
drawb TLeaf = t
drawb (TTT l r) = (drawb l) B.vcat   B.// ((drawb r)  B.<> v )
-}
-- PP.nest 2 l <> PP.line <> PP.nest 1 l <> PP.line <> l
-- text "Tau \\"
-- text "Id \\"

-- From Data.Tree

draw :: FibTree a b -> [String]
draw (ITT l r) = "I" : drawSubTrees (l,r)
draw (TIT l r) = "T" : drawSubTrees (l,r)
draw (TTT l r) = "T" : drawSubTrees (l,r)
draw (TTI l r) = "T" : drawSubTrees (l,r)
draw (III l r) = "I" : drawSubTrees (l,r)
draw  ILeaf    = "I" : []
draw  TLeaf    = "T" : []

-- drawSubTrees' [] = []
drawSubTrees' t =
    "|" : shift "`- " "   " (draw t)
drawSubTrees (t,ts) =
    "|" : shift "+- " "|  " (draw t) ++ drawSubTrees' ts

shift first other = zipWith (++) (first : repeat other)

{-
draw' ILeaf = ["I"]
draw' TLeaf = ["T"]
draw' (TTT l r) =  " -" ++ ((fmap . fmap) (\s -> s ++ "  |") (draw' l)) ++ (["-+"] ++ draw' r))
-}
-- drawex = drawb (TTT (TTT TLeaf TLeaf) TLeaf)

-- | Neat 2-dimensional drawing of a tree.
drawTree :: FibTree a b -> String
drawTree  = unlines . draw

-- | Neat 2-dimensional drawing of a forest.
-- drawForest :: Forest String -> String
-- drawForest  = unlines . map drawTree

instance Ord (FibTree a b) where
  compare (ITT l r) (ITT l' r') | l < l' = LT
                                | l > l' = GT
                                | otherwise = compare r r'
  compare (ITT _ _) _ = LT
  compare _ (ITT _ _) = GT
  compare (TTI l r) (TTI l' r') | l < l' = LT
                                | l > l' = GT
                                | otherwise = compare r r'
  compare (TIT l r) (TIT l' r') | l < l' = LT
                                | l > l' = GT
                                | otherwise = compare r r'
  compare (TTT l r) (TTT l' r') | l < l' = LT
                                | l > l' = GT
                                | otherwise = compare r r'
  compare (III l r) (III l' r') | l < l' = LT
                                | l > l' = GT
                                | otherwise = compare r r'

  compare (TTI _ _) _ = LT
  compare _ (TTI _ _) = GT
  compare (TIT _ _) _ = LT
  compare _ (TIT _ _) = GT
  compare (TTT _ _) _ = LT
  compare _ (TTT _ _) = GT
  compare (III _ _) _ = LT
  compare _ (III _ _) = GT
  compare TLeaf TLeaf = EQ
  compare ILeaf ILeaf = EQ
  -- whoa. GHC is smart enough to realize some of these patterns can't happen
  -- because of the types.

{-
instance Enum (FibTree 'Tau (A 'Tau)) where
  toEnum _ = TLeaf
  fromEnum _ = 0
-}
{-
instance Enum (FibTree 'Id 'Id) where
  toEnum _ = TLeaf
  fromEnum _ = 0
  -}
-- and no inhabitants types?
--instance Enum (FibTree d b), Enum (FibTree e c) => Enum (FibTree a (b,c))
{-
instance (Enum (FibTree 'Tau b),
  Enum (FibTree 'Tau c),
  Enum (FibTree ' b),
  Enum (FibTree 'Id c)) => Enum (FibTree 'Tau (b,c))
  toEnum 0 = TTI
  toEnum
  fromEnum
-}

--instance Enum (FibTree a (b,c)) Enum (FibTree a (b,c))=> Enum (FibTree 'Id (b,c))

-- (FibTree a (b,c))

type TreeFun a b c d = FibTree a b -> FibTree c d
{-
lmap :: (forall a. FibTree a b -> FibTree (a) c) -> (forall e. FibTree e (b,d) -> FibTree e (c,d)) --  forall (a :: FibAnyon).
lmap f (ITT l r) = (ITT (f l) r)
lmap f (TTI l r) = (TTI (f l) r)
lmap f (TIT l r) = (TIT (f l) r)
lmap f (TTT l r) = (TTT (f l) r)
-}
lmap :: (forall a. FibTree a b -> Q (FibTree a c)) -> (FibTree e (b,d) -> Q (FibTree e (c,d)))
lmap f (ITT l r) = fmap (\l' -> ITT l' r) (f l)
lmap f (TTI l r) = fmap (\l' -> TTI l' r) (f l)
lmap f (TIT l r) = fmap (\l' -> TIT l' r) (f l)
lmap f (TTT l r) = fmap (\l' -> TTT l' r) (f l)
lmap f (III l r) = fmap (\l' -> III l' r) (f l)
{-
rmap :: (forall a. FibTree a b -> FibTree a c) -> (forall e. FibTree e (d,b) -> FibTree e (d,c)) --  forall (a :: FibAnyon).
rmap f (ITT l r) = ITT l (f r)
rmap f (TTI l r) = TTI l (f r)
rmap f (TIT l r) = TIT l (f r)
rmap f (TTT l r) = TTT l (f r)
-}
rmap :: (forall a. FibTree a b -> Q (FibTree a c)) -> (FibTree e (d,b) -> Q (FibTree e (d,c)))
rmap f (ITT l r) = fmap (\r' -> ITT l r') (f r)
rmap f (TTI l r) = fmap (\r' -> TTI l r') (f r)
rmap f (TIT l r) = fmap (\r' -> TIT l r') (f r)
rmap f (TTT l r) = fmap (\r' -> TTT l r') (f r)
rmap f (III l r) = fmap (\r' -> III l r') (f r)

fibswap :: FibTree a (l,l') -> FibTree a (l',l)
fibswap (ITT l r) = (ITT r l)
fibswap (TTI l r) = (TIT r l)
fibswap (TIT l r) = (TTI r l)
fibswap (TTT l r) = (TTT r l)

eye = 0 :+ 1

braid :: FibTree a (l,r) -> Q (FibTree a (r,l))
braid (ITT l r) = W [(ITT r l,  cis $ 4 * pi / 5)]  -- different scalar factors for trivial and non trivial fusion
braid (TTT l r) = W [(TTT r l,  (cis $ - 3 * pi / 5))]
braid (TTI l r) = pure $ TIT r l-- exchange with trivial means nothing
braid (TIT l r) = pure $ TTI r l
braid (III l r) = pure $ III r l

-- The inverse of braid
braid' :: FibTree a (l,r) -> Q (FibTree a (r,l))
braid' = star . braid

-- property
-- braid (braid' v) == v
-- braid' (braid v) == v

{-
fibassoc :: FibTree a ((c,d),e) -> FibTree a (c,(d,e))
fibassoc (ITT l r) = (ITT r l)
fibassoc (TTI l r) = (TIT r l)
fibassoc (TIT l r) = (TTI r l)
fibassoc (TTT l r) = (TTT r l)
-}
-- Looks like we need more constructors for the tree
{-
fmove :: FibTree a (c,(d,e)) -> FibTree a ((c,d),e)
-- fmove (ITT  a  (TTI b c)) = ITI ( TTT  a b) c
fmove (ITT  a  (TIT b c)) = ITT ( TTI  a b) c
fmove (ITT  a  (TTT b c)) = ITT ( TTT  a b) c

fmove (TTT  a  (TTI b c)) = TTI ( TTT  a b) c
fmove (TTT  a  (TIT b c)) = TTT ( TTI  a b) c
fmove (TTT  a  (TTT b c)) = TTT ( TTT  a b) c

fmove (TIT  a  (TTI b c)) = TTI ( TIT  a b) c
-- fmove (TIT  a  (TIT b c)) = TTT ( III  a b) c
fmove (TIT  a  (TTT b c)) = TTT ( TIT  a b) c
-}

fmove :: FibTree a (c,(d,e)) -> Q (FibTree a ((c,d),e))
-- fmove (ITT  a  (TTI b c)) = pure $ ITI ( TTT  a b) c -- pure (auto (auto a b) c) -- no maybe not. The internal one isn't auto
fmove (ITT  a  (TIT b c)) = pure $ ITT ( TTI  a b) c
fmove (ITT  a  (TTT b c)) = pure $ ITT ( TTT  a b) c
fmove (ITT  a  (TTI b c)) = pure $ III ( ITT  a b) c

fmove (TIT  a  (TTT b c)) = pure $ TTT ( TIT  a b) c
fmove (TIT  a  (TTI b c)) = pure $ TTI ( TIT  a b) c
fmove (TIT  a  (TIT b c)) = pure $ TIT ( III  a b) c

-- fmove (TIT  a  (TIT b c)) = TTT ( III  a b) c
-- the nontrivial ones have all tau on the leafs and root
-- internal I
fmove (TTI  a  (III b c)) = pure $ TTI ( TTI  a b) c
fmove (TTI  a  (ITT b c)) = W [(TIT ( ITT  a b) c, tau)         , (TTT ( TTT  a b) c, sqrt tau)]
-- internal T
fmove (TTT  a  (TTT b c)) = W [(TIT ( ITT  a b) c, sqrt tau)  ,   (TTT ( TTT  a b) c, - tau   )]
fmove (TTT  a  (TTI b c)) = pure $ TTI ( TTT  a b) c
fmove (TTT  a  (TIT b c)) = pure $ TTT ( TTI  a b) c

fmove (III  a  (ITT b c)) = pure $ ITT ( TIT  a b) c
fmove (III  a  (III b c)) = pure $ III ( III  a b) c

-- largely just a tranpose of the above case.
fmove' :: FibTree a ((c,d),e) -> Q (FibTree a (c,(d,e)))
fmove' (ITT ( TTI  a b) c) = pure $ (ITT  a  (TIT b c))
fmove' (ITT ( TTT  a b) c) = pure $  (ITT  a  (TTT b c))
fmove' (ITT ( TIT  a b) c) = pure $  (III  a  (ITT b c))

--fmoveq (ITT  a  (TTT b c)) = pure $

fmove' (TTI ( TTT  a b) c) = pure $ (TTT  a  (TTI b c))
fmove' (TTI ( TTI  a b) c) = pure $ (TTI  a  (III b c))
fmove' (TTI ( TIT  a b) c) = pure $ TIT  a  (TTI b c)
--fmoveq (TTT  a  (TTI b c)) = pure $ TTI ( TTT  a b) c

--fmoveq (TTT  a  (TIT b c)) = pure $ TTT ( TTI  a b) c
fmove' (TIT ( ITT  a b) c) = W [(TTI  a  (ITT b c), tau)         , (TTT  a  (TTT b c) , sqrt tau)]
fmove' (TIT ( III  a b) c ) = pure $ TIT  a  (TIT b c)

fmove' (TTT ( TTI  a b) c ) = pure $ TTT  a  (TIT b c)
fmove' (TTT ( TIT  a b) c ) = pure $ TIT  a  (TTT b c)
fmove' (TTT ( TTT  a b) c) = W [(TTI  a  (ITT b c), sqrt tau)  , (TTT  a  (TTT b c),   - tau  )]

fmove' (III ( III  a b) c ) = pure $ III  a  (III b c)
fmove' (III ( ITT  a b) c ) = pure $ ITT  a  (TTI b c)

--fmoveq (TIT  a  (TTI b c)) = pure $ TTI ( TIT  a b) c
-- fmove (TIT  a  (TIT b c)) = TTT ( III  a b) c
-- the nontrivial ones have all tau on the leafs and root

-- internal I

--fmoveq (TTI  a  (ITT b c)) = W [(TIT ( ITT  a b) c, recip tau)         , (TTT ( TTT  a b) c, recip $ sqrt tau)]
-- internal T

--fmoveq (TTT  a  (TTT b c)) = W [(TIT ( ITT  a b) c, recip $ sqrt tau)  , (TTT ( TTT  a b) c,   - recip tau  )]

checkf :: Q (FibTree Tau ((Tau, Tau),Tau))
checkf = simplify $ fmove' (TIT (ITT TLeaf TLeaf) TLeaf) >>= fmove
checkf' = simplify $ fmove' (TTT (TTT TLeaf TLeaf) TLeaf) >>= fmove
-- checkf''' = fmove' (TIT (ITT TLeaf TLeaf) TLeaf)
-- checkf'' = fmove' (ITT (TTT TLeaf TLeaf) TLeaf) >>= fmove

-- tau**2 + tau == 1
tau :: Complex Double
tau =  ((sqrt 5) - 1) / 2 :+ 0 -- 0.618 :+ 0

--test1 = lmap fibswap
-- test2 = lmap $ lmap fibswap
-- test3 = rmap fmove

type Vec1 b r = [(b, r)]

smul :: Num r => r -> Vec1 b r -> Vec1 b r
smul s = map (\(b,s') -> (b, s' * s))

linapply :: Num r => (a -> Vec1 b r) -> (Vec1 a r -> Vec1 b r) -- bind
linapply f [] = []
linapply f ((b,s) : xs) = smul s (f b) ++ linapply f xs

dTTT :: FibTree a (Tau, Tau) -> Q (FibTree a Tau)
dTTT (TTT _ _) = pure TLeaf
dTTT _ = mempty
{-
dITT :: FibTree a (Tau, Tau) -> Q (FibTree a Id)
dTTT (TTT _ _) = pure TLeaf
dTTT _ = mempty
-}
dTTI :: FibTree a (Tau, Id) -> Q (FibTree a Tau) -- A dual tree of type  'tau (tau,'tau)
dTTI (TTI _ _) =  pure TLeaf
dTTI _ = mempty

dTTT'''' = dot (TTT TLeaf TLeaf)
-- I could do this as a typeclass. NOPE. don't need that garbage

-- I think that incompatible FibTree a' a is empty by default vs unconstructible
-- which is better?
dot :: FibTree a (b, c) -> FibTree a' (b, c) -> Q (FibTree a' a)
dot x@(TTI _ _) y@(TTI _ _) | x == y = pure TLeaf
                             | otherwise = mempty
dot x@(TIT _ _) y@(TIT _ _) | x == y = pure TLeaf
                             | otherwise = mempty
dot x@(TTT _ _) y@(TTT _ _) | x == y = pure TLeaf
                            | otherwise = mempty
dot x@(III _ _) y@(III _ _) | x == y = pure ILeaf
                            | otherwise = mempty
dot x@(ITT _ _) y@(ITT _ _) | x == y = pure ILeaf
                            | otherwise = mempty
dot _ _ = mempty

test5 = (pure $ TTT TLeaf TLeaf) >>= dTTT
test6 = (pure $ ITT TLeaf TLeaf) >>= dTTT
test7 v = do y <- rmap dTTI v  -- wait. What the hell am I doing?
             lmap dTTT y
test8 v = rmap dTTT v  >>= dTTT
-- Building from the leaves up...?
-- you  hve to apply the deeper stuff before the upper stuff.
--- So this is a dual tree with TTT in the right branch. Then TTT above that.
-- In total (T,(T,T)) -> Tleaf.
-- maybe we need to cps it.
dTTT' :: (Q (FibTree a Tau) -> (Complex Double)) -> (FibTree a (Tau, Tau)) -> Complex Double
dTTT' f = f . dTTT

dTTT'' :: (FibTree a Tau -> Complex Double) -> FibTree a (Tau, Tau) -> Complex Double
dTTT'' f = undefined

-- sort of similar to bind in a way. lift function to a Q function
-- (Q FibTree a ) -> (FibTree a -> Q FibTree b)  -> Q FibTree b
dot' :: (FibTree a b -> Complex Double) -> Q (FibTree a b) -> Complex Double
dot' f (W v) = sum $ map (\(e,a) -> a * (f e)) v

dTTT''' = dot' . dTTT''

pentagon1 ::  FibTree a (e,(d,(c,b))) -> Q (FibTree a (((e,d),c),b))
pentagon1 v =  do
                 v1 <- fmove v
                 fmove v1

-- type annotations not necessary. For clarity.
pentagon2 :: FibTree a (b,(c,(d,e))) -> Q (FibTree a (((b,c),d),e))
pentagon2 v = do
                v1 :: FibTree a (b,((c,d),e)) <- rmap fmove v
                v2 :: FibTree a ((b,(c,d)),e) <- fmove v1
                lmap fmove v2

ex3 = TTT TLeaf (TTI TLeaf (ITT TLeaf TLeaf))
ex1 = TTT TLeaf (TTT TLeaf (TTT TLeaf TLeaf))
pentagon =  simplify $ ((pentagon1 ex1) - (pentagon2 ex1))
pentagon' =  simplify $ ((pentagon1 ex3) - (pentagon2 ex3))

hexagon1 :: FibTree a (b,(c,d)) -> Q (FibTree a ((d,b),c))
hexagon1 v = do
             v1 :: FibTree a ((b,c),d) <- fmove v
             v2 :: FibTree a (d,(b,c)) <- braid v1
             fmove v2

hexagon2 :: FibTree a (b,(c,d)) -> Q (FibTree a ((d,b),c))
hexagon2 v = do
             v1 :: FibTree a (b,(d,c)) <- rmap braid v
             v2 :: FibTree a ((b,d),c) <- fmove v1
             lmap braid v2

ex4 = (TTI TLeaf (ITT TLeaf TLeaf))
ex2 = (TTT TLeaf (TTT TLeaf TLeaf))
hexagon =  simplify $ ((hexagon1 ex2) - (hexagon2 ex2))
hexagon' =  simplify $ ((hexagon1 ex4) - (hexagon2 ex4))

-- hexagon2 :: FibTree a (b,(c,d)) -> Q (FibTree a (c,(d,b)))

-- quickcheck forall v, pentagon1 v == pentagon2 v

-- test9 = test8 $ pure $ TTT (TLeaf) ()
-- Vec1 (FibTree 'Tau l) r -> Vec1 (FibTree 'Tau l') r -> Vec1 (FibTree 'Id (l,l')) r

-- linearized tensor products
--linITT :: Num r => Vec1 (FibTree 'Tau l) r -> Vec1 (FibTree 'Tau l') r -> Vec1 (FibTree 'Id (l,l')) r -- we would like to enforce that isingany can split into a b, but thisis tough
--linITT = [ (ITT l r, s * s')  |  (l, s) <- v1 , (r, s') <- v2   ]

-- class AutoNode c a b where
      -- prod ::  FibTree a l -> FibTree b l' -> FibTree  c (l,l')
{-
 instance AutoProd 'Tau 'Tau 'Tau
    prod = TTT
 instance AutoProd 'Id 'Tau 'Tau
    prod = ITT
-- etc.

-- I can't make a fst. It extracts a skolemized object.
-- So I do need to do complete pattern matching? No. I CAN do this... ? Not always clear I'd want to?
fibfst :: AutoNode a e _ => FibTree a (b,c) -> FibTree e b

-}
class TProd a b c where
  tprod :: Vec1 (FibTree a l) r -> Vec1 (FibTree b l') r -> Vec1 (FibTree c (l,l')) r

-- tprod :: AutoNode a b c => Vec1 (FibTree a l) r -> Vec1 (FibTree b l') r -> Vec1 (FibTree c (l,l')) r
--class Index tree n  -- reference leaf by number?

-- class AutoFMove tree n

{-
linTTI :: FibTree 'Tau l -> FibTree 'Id l' -> FibTree 'Tau (l,l')
linTIT :: FibTree 'Id l -> FibTree 'Tau l' -> FibTree 'Tau (l,l')
libTTT :: FibTree 'Tau l -> FibTree 'Tau l' -> FibTree 'Tau (l,l')
-}
(~*) :: r -> b -> (b,r) -- **, `smul`
(~*) = flip (,)

(~+) :: a -> [a] -> [a] -- ++ instead?
(~+) = (:)

test4 = [1 ~* 'a', 2 ~* 'b'] -- eh. What's the point.
-- test3 = (lmap . lmap) fibswap
--lmap (ITT l r)
-- ^+^
-- zero
--

-- TODOS

-- Categorical Interface
-- lmap, rmap vs second first.

{-
pullLeft (Tau,x) = (Tau,x)
pullLeft (Id,x) = (Id,x)
pullLeft (a,b) = fmove ((a',c),b) where (a',c) = pullLeft a
-- ((a,b),c) = (, (b,c)) where (a',d) = pullLeft a
-}
{-
type family FMove' a where
  FMove' ((a,b),c)  = (a,(b,c))

type family PullLeft a where
  PullLeft Tau = Tau
  PullLeft Id = Id
  PullLeft (Tau,b) = (Tau,b)
  PullLeft (Id,b) =  (Id,b)
  PullLeft ((a,b),c)  = FMove' (PullLeft (a,b), c)

-- Auto F Moves. / auto braid.

pullLeft :: (PullLeft (b,c) ~ (b',c')) => FibTree a (b,c) -> Q (FibTree a (b',c'))
pullLeft x@(TTT TLeaf _) = pure x
pullLeft (TTT l@(TTT _ _ ) r) =  do
                      l' <- pullLeft l
                      fmove' (TTT l' r)
-}
class PullLeft a b | a -> b where -- | a -> b functional dependency causes errors?
  pullLeft :: FibTree c a -> Q (FibTree c b)

instance PullLeft (Tau,c) (Tau,c) where
  pullLeft = pure

instance PullLeft (Id,c) (Id,c) where
  pullLeft = pure

instance PullLeft Tau Tau where
  pullLeft = pure

instance PullLeft Id Id where
  pullLeft = pure

instance (PullLeft (a,b) (a',b'), r ~ (a',(b',c))) => PullLeft ((a, b),c) r where
  pullLeft t = do
           t' <- lmap pullLeft t
           fmove' t'
{-
instance (PullLeft a a', r ~ (a',(b,c))) => PullLeft (a, (b,c)) r where
  pullLeft t = lmap pullLeft t
-}
{-
instance (PullLeft a a', PullLeft (a',b) (a'',b''), ) => PullLeft (a,b) r where
  pullLeft = lmap pullLeft t
-}

{-  pullLeft (TTT l r) =  do
                        l' <- pullLeft l
                        fmove' (TTT l' r)
-}

class PullRight a b | a -> b where -- | a -> b functional dependency causes errors?
  pullRight :: FibTree c a -> Q (FibTree c b)
{-
instance (PullRight a a', r ~ ((b,c),a')) => PullRight ((b,c),a) r where
  pullRight t = rmap pullRight t
-}
instance PullRight Tau Tau where
  pullRight = pure

instance PullRight Id Id where
  pullRight = pure

instance PullRight (c,Tau) (c,Tau) where
  pullRight = pure

instance PullRight (c,Id) (c,Id) where
  pullRight = pure

instance (PullRight (a,b) (a',b'), r ~ ((c,a'),b')) => PullRight (c,(a, b)) r where
  pullRight t = do
           t' <- rmap pullRight t
           fmove t'

type family Count a where
  Count Tau = 1
  Count Id = 1
  Count (a,b) = (Count a) + (Count b)

type family LeftCount a where
  LeftCount (a,b) = Count a

-- The version without the explicit ordering supplied.
class LCA n a b c d | n a c -> b d where
  lcamap :: (forall r. FibTree r b -> Q (FibTree r c)) -> (FibTree e a) -> Q (FibTree e d)
{-
lcamapP :: LCA n a b c d => Proxy n ->  (forall r. FibTree r b -> Q (FibTree r c)) -> (FibTree e a) -> Q (FibTree e d)
lcamapP _ f x = lcamap f x
-}
instance (lc ~ (LeftCount a),
          gte ~ (CmpNat lc n),
         LCA' n gte a b c d) => LCA n a b c d where
  lcamap f x = lcamap' @n @gte f x

class LCA' n gte a b c d | n gte a c -> b d where
  lcamap' :: (forall r. FibTree r b -> Q (FibTree r c)) -> (FibTree e a) -> Q (FibTree e d)

-- we find b at the lca and pass it back up. c gets passed all the way down, d gets computed by rebuilding out of c.
-- a drives the search.
instance (n' ~ (n - Count l), -- we're searching in the right subtree. Subtract the leaf number in the left subtree
        lc ~ (LeftCount r), -- dip one left down to order which way we have to go next
        gte ~ (CmpNat lc n'), -- Do we go left, right or havce we arrived in the next layer?
        LCA' n' gte r b c d',  -- recurive call
        d ~ (l,d') -- reconstruct total return type from recurive return type. left tree is unaffected by lcamapping
        ) => LCA' n 'LT (l,r) b c d where
    lcamap' f x = rmap (lcamap' @n' @gte f) x

instance (lc ~ (LeftCount l),
          gte ~ (CmpNat lc n),
          LCA' n gte l b c d',
          d ~ (d',r)
          ) => LCA' n 'GT (l,r) b c d where
    lcamap' f x = lmap (lcamap' @n @gte f) x

instance (b ~ a, d ~ c) => LCA' n 'EQ a b c d where
  lcamap' f x = f x

class LeafMap n gte a b c d | n gte a c -> b d where
  leafmap :: (forall r. FibTree r b -> Q (FibTree r c)) -> (FibTree e a) -> Q (FibTree e d)

instance (n' ~ (n - Count l), -- we're searching in the right subtree. Subtract the leaf number in the left subtree
        lc ~ (LeftCount r), -- dip one left down to order which way we have to go next
        gte ~ (CmpNat lc n'), -- Do we go left, right or havce we arrived in the next layer?
        LeafMap n' gte r b c d',  -- recurive call
        d ~ (l,d') -- reconstruct total return type from recurive return type. left tree is unaffected by lcamapping
        ) => LeafMap n 'LT (l,r) b c d where
    leafmap f x = rmap (leafmap @n' @gte f) x

instance (lc ~ (LeftCount l),
          gte ~ (CmpNat lc n),
          LeafMap n gte l b c d',
          d ~ (d',r)
          ) => LeafMap n 'GT (l,r) b c d where
    leafmap f x = lmap (leafmap @n @gte f) x

-- In the equals case, we now continue onward.
instance (lc ~ (LeftCount l),
          gte ~ (CmpNat lc n),
          LeafMap n gte l b c d',
          d ~ (d',r)
          ) => LeafMap n 'EQ (l,r) b c d where
    leafmap f x = lmap (leafmap @n @gte f) x

-- base case
instance (b ~ Tau, d ~ c) => LeafMap 1 'EQ Tau b c d where
  leafmap f x = f x

instance (b ~ Id, d ~ c) => LeafMap 1 'EQ Id b c d where
  leafmap f x = f x

-- split
-- leafmap @3 (const (TTT TLeaf TLeaf))
-- need one last arbitrary left or right fmove to put the two on the same stalk
-- neighbormap :: (LCA n a (l,r) (l',r') d, PullRight l (l',x), PullLeft r (y,r') ) => (FibTree.  -> Q FibTree) -> FibTree

-- I need this to also work for (Tau, Tau)
-- (Tau, yada yada)
-- I.e. I need to do different things depending on

-- Jesus. What a shitshow. But I don't see how to do better.
-- a is starting tree
-- b is extracted part
-- c is transformed b
-- d is recoustrcted a replacing b with c and rearranging.
class NeighborMap a b c d | a c -> b d where
  nmap :: (forall r. FibTree r b -> Q (FibTree r c)) -> FibTree e a -> Q (FibTree e d)

instance NeighborMap ((l,x),(y,r)) (x,y) c ((l,c),r) where
   nmap f x = do
              x'  <- fmove x -- (((l',x),y),r')
              x'' <- lmap fmove' x' -- ((l',(x,y)),r')
              lmap (rmap f) x''
instance NeighborMap (Tau, (y,r)) (Tau, y)  c  (c,r) where
   nmap f x = fmove x >>= lmap f
instance NeighborMap (Id, (y,r)) (Id, y) c (c,r) where
   nmap f x = fmove x >>= lmap f
instance NeighborMap ((l,x), Tau) (x,Tau) c (l,c) where
   nmap f x = fmove' x >>= rmap f
instance NeighborMap ((l,x), Id) (x,Id) c  (l,c) where
   nmap f x = fmove' x >>= rmap f
instance NeighborMap (Tau, Tau) (Tau,Tau) c  c where
   nmap f x = f x
instance NeighborMap (Id, Id) (Id,Id) c  c where
   nmap f x = f x
instance NeighborMap (Tau, Id) (Tau,Id) c  c where
   nmap f x = f x
instance NeighborMap (Id, Tau) (Id,Tau) c  c where
   nmap f x = f x

neighbormap :: forall n a b c d l l' r' x y e r z. (LCA n a b c d,
   b ~ (l,r),
   c ~ ((l',z),r'),
   PullRight l (l',x),
   PullLeft r (y,r')) =>
   (forall r. FibTree r (x,y) -> Q (FibTree r z)) -> FibTree e a -> Q (FibTree e d)
neighbormap f z = lcamap @n @a @b @c @d (helper f) z

helper ::  (c ~ ((l',z),r'),
   PullRight l (l',x),
   PullLeft r (y,r')) => (forall r. FibTree r (x,y) -> Q (FibTree r z)) -> FibTree e (l,r) -> Q (FibTree e c)
helper f x = do
            x' <- rootneighbor x
            lmap (rmap f) x'

rootneighbor :: (PullRight l (l',x), PullLeft r (y,r')) => FibTree e (l,r) -> Q (FibTree e ((l',(x,y)),r'))
rootneighbor x = do
                x' <- lmap pullRight x
                x'' <- rmap pullLeft x' -- ((l',x),(y,r'))
                x''' <- fmove x'' -- (((l',x),y),r')
                lmap fmove' x''' -- ((l',(x,y)),r')

neighbormap' :: forall n a b c d l l' r r' x y e z b'. (LCA n a b c d,
   b ~ (l,r),
   -- c ~ c',
   -- c ~ ((l',z),r'),
   PullRight l l',
   PullLeft r r',
   NeighborMap (l',r') b' z c) =>
   (forall r. FibTree r b' -> Q (FibTree r z)) -> FibTree e a -> Q (FibTree e d)
neighbormap' f z = lcamap @n @a @b @c @d (\x -> do
                                          x'  <- lmap pullRight x
                                          x'' <- rmap pullLeft x'
                                          nmap f x'') z

neighbormap'' :: forall n a b' z d c l l' r r' x y e b. (LCA n a b c d,
   b ~ (l,r),
   -- c ~ c',
   -- c ~ ((l',z),r'),
   PullRight l l',
   PullLeft r r',
   NeighborMap (l',r') b' z c) =>
   (forall r. FibTree r b' -> Q (FibTree r z)) -> FibTree e a -> Q (FibTree e d)
neighbormap'' f z = lcamap @n @a @b @c @d (\x -> do
                                          x'  <- lmap pullRight x
                                          x'' <- rmap pullLeft x'
                                          nmap f x'') z

t1 = neighbormap' @2 braid (TTT (TTI TLeaf ILeaf) (TTT TLeaf TLeaf))
t2 = neighbormap' @1 braid (TTT (TTI TLeaf ILeaf) (TTT TLeaf TLeaf))
t3 = neighbormap' @2 braid (TTT (TTT (TTT TLeaf TLeaf) TLeaf) (TTT TLeaf TLeaf))

                                      {-
neighbormap p f z = let helper (x :: FibTree _ (l,r)) = do
                                x' <- rootneighbor x
                                lmap (rmap f) x'
                    in
                    lcamap helper z
                    -}
{-  lcamap
((\x -> do
                                x' <- rootneighbor x
                                lmap (rmap f) x'
                                ) :: (forall s. FibTree s (l,r) -> Q (FibTree s c))) z
-}

{-(\x -> do
                            x' <- lmap pullRight x
                            x'' <- rmap pullLeft x' -- ((l',x),(y,r'))
                            x''' <- fmove x'' -- (((l',x),y),r')
                            x4 <- lmap fmove' x''' -- ((l',(x,y)),r')
                            lmap (lmap f) x4 -- ((l',z),r')
                            ) z
-}
{-
-- neighbormap f z = lcamap @n (\x -> do
                            x' <- lmap pullRight x
                            x'' <- rmap pullLeft x'
                            x''' <- fmove x''
                            lmap? f x'''   ) z
-}
{-
abraid :: forall n a b c d l l' r' x y e r z b'. (LCA n a b c d,
   b ~ (l,r),
   -- c ~ c',
   -- c ~ ((l',z),r'),
   PullRight l l',
   PullLeft r r',
   NeighborMap (l',r') (x,y) (y,x) c) => Q (FibTree e a) -> Q (FibTree e d)

   -}

abraid :: forall n a b c d l l' r r' x y e z b'. (LCA n a b c d,
   b ~ (l,r),
   -- c ~ c',
   -- c ~ ((l',z),r'),
   PullRight l l',
   PullLeft r r',
   NeighborMap (l',r') (x,y) (y,x) c) =>
   Q (FibTree e a) -> Q (FibTree e d)
abraid x = x >>= (neighbormap' @n @a @b @c @d @l @l' @r @r' @x @y @e @(y,x) @(x,y)) braid

abraid' :: forall n a b c d l l' r r' x y e z b'. (LCA n a b c d,
   b ~ (l,r),
   -- c ~ c',
   -- c ~ ((l',z),r'),
   PullRight l l',
   PullLeft r r',
   NeighborMap (l',r') (x,y) (y,x) c) =>
   Q (FibTree e a) -> Q (FibTree e d)
abraid' x = x >>= (neighbormap' @n @a @b @c @d @l @l' @r @r' @x @y @e @(y,x) @(x,y)) braid'

{-
abraid
  :: forall n a b c d l l' r r' x y e z b'. (LCA n a (l, r) c d, PullRight l l', PullLeft r r',
      NeighborMap (l', r') (x, y) (y, x) c) =>
     (FibTree e a) -> Q (FibTree e d)
abraid x = neighbormap' @n @a @(l,r) @c @d @l @l' @r @r' braid x
-}
{-
t4 = abraid @2 $
     abraid @4 $
     abraid @3 $
     pure (TTT (TTT (TTT TLeaf TLeaf) TLeaf) (TTT TLeaf TLeaf))
-}
{-
ndot' :: forall n a b c d l l' r r' x y e z b' a'. (LCA n a b c d,
   b ~ (l,r),
   -- c ~ c',
   -- c ~ ((l',z),r'),
   PullRight l l',
   PullLeft r r',
   NeighborMap (l',r') (x,y) a' c) =>
   FibTree a' (x,y) -> FibTree e a -> Q (FibTree e d)
   -}
-- ndot' x = (neighbormap' @n @a @b @c @d @l @l' @r @r' @x @y @e @a' @(x,y)) (dot' x)

ndot :: forall n a l r c d l' r' x y a' e. (LCA n a (l, r) c d, PullRight l l', PullLeft r r',
      NeighborMap (l', r') (x, y) a' c) =>
     FibTree a' (x, y) -> FibTree e a -> Q (FibTree e d)
ndot x = neighbormap' @n @a @(l,r) (dot x)

-- This a b c d pattern is something like a lens. We are replacing
-- piece of b with c which turns a into d. b is piece of a.
-- s t a b

-- quantified constraints
--  (forall b' c'. (g a b' c' d , f b' b c c') => Compose f g a b c d

-- (forall r. FibTree r b' -> Q (FibTree r z)) ->
{-
abraid' :: forall n a b c d l l' r r' x y e b'. (LCA n a b c d,
   b ~ (l,r),
   -- c ~ c',
   -- c ~ ((l',z),r'),
   PullRight l l',
   PullLeft r r',
   NeighborMap (l',r') (x,y) (y,x) c) =>
   FibTree e a -> Q (FibTree e d)
abraid' = (neighbormap'' @n @a @(x,y) @(y,x) @d) braid
-}

{-
abraid' :: forall n a b c d l l' r r' x y e z b'. (LCA n a b c d,
   b ~ (l,r),
   -- c ~ c',
   -- c ~ ((l',z),r'),
   PullRight l l',
   PullLeft r r',
   NeighborMap (l',r') (x,y) (y,x) c) =>
   FibTree e a -> Q (FibTree e d)
abraid' = (neighbormap' @n @a @b @c @d @l @l' @r @r' @x @y @e @(y,x) @(x,y)) braid'
-}

{-
abraid :: (LCA n a (l1, r) c d, PullRight l1 l'1, PullLeft r r',
      NeighborMap (l'1, r') (l2, l'2) (l'2, l2) c) =>
     FibTree e a -> Q (FibTree e d)
abraid = neighbormap' braid
-}
-- autobraid = neighbormap braid
-- autodot x = neighbormap (dot x)

-- Build state monad to carry along the vector.
-- should be able to do in applicative style, since data does not change structre of computation

-- auto braid
-- find least common ancestor
-- pullLeft, and pullRight
-- then braid

-- stateful
-- Pretty Print tree

-- Operators.
newtype FibOp a b = FibOp (forall c. FibTree c a -> Q (FibTree c b))
type FibOp' c a b = FibTree c a -> Q (FibTree c b)

instance Category FibOp where
  id = FibOp pure
  (FibOp f) . (FibOp g) = FibOp (f <=< g)
{-
instance Arrow FibOp where
  arr = error "No. No arr."
  (***) =
-}
{-
class Monoidal k where
  (***) :: k a b -> k c d -> k (a,c) (b,d)
-}
tensor :: Num r => (b -> c -> a) -> W r b -> W r c -> W r a
tensor f (W xs) (W ys) =  W [ ( f x y , xc * yc ) | (x, xc) <- xs,  (y,yc) <- ys]

-- tensor :: Num r => (b -> c -> a) -> b -> c -> W r a
-- fg f g (TTT l r) = tensor TTT (f l) (g r)

{-
instance Monoidal FibOp where
  (***) :: FibOp a b -> FibOp c d -> FibOp (a,c) (b,d)
  (FibOp f) *** (FibOp g) = FibOp fg where
                               -- fg ::  (a,c) (b,d)
                               fg :: FibTree e (a,c) -> Q (FibTree e (b,d))
                               fg (TTT l r) = tensor TTT (f l) (g r)
                               -- fg (ITT l r) = FibOp $ tensor ITT (f l) (g r)
                               -- fg (TIT l r) = FibOp $ tensor TIT (f l) (g r)
-}
-- try to remove the issue of being in a instance. Which is kind of silly anyhow
{-
(****) :: FibOp' e a b -> FibOp' e c d -> FibOp' e (a,c) (b,d)
f **** g = fg where
                   -- fg ::  (a,c) (b,d)
                 fg :: FibTree e (a,c) -> Q (FibTree e (b,d))
                 fg (TTT l r) = tensor TTT (f l) (g r)
                 fg (ITT l r) = tensor ITT (f l) (g r)
-}
(***) :: (forall e. FibOp' e a b) -> (forall e'. FibOp' e' c d)  -> FibOp' e'' (a,c) (b,d)
(***) f g (TTT l r) = tensor TTT (f l) (g r)
(***) f g (ITT l r) = tensor ITT (f l) (g r)
(***) f g (TIT l r) = tensor TIT (f l) (g r)

first = lmap
second = rmap
-- braid
-- fmove

-- cup, cap? dot'

-- Is the problem that I'd trying to pattern match on kind of the output?
 -- densification

-- Quickcheck props

-- implement qubit abtraction

-- Did I fall into a way to make end, co-end work as einstein notation?6
-- The bounded meaning of forall and exists in the GADT context

main :: IO ()
main = mymain
--return ()


{-# LANGUAGE BangPatterns,
             ScopedTypeVariables,
             RecordWildCards,
             FlexibleContexts,
             TypeFamilies,
             GeneralizedNewtypeDeriving #-}

module AI.HsANN.ComplexNetwork
    (
    -- * Types
      ComplexNetwork(..)
    , ActivationFunction
    , ActivationFunctionDerivative
    , Sample
    , Samples
    , (-->)

    -- * Creating a neural network
    , createComplexNetwork
    , fromWeightMatrices

    -- * Computing a neural network's output
    , output
    , sigmoid
    , sigmoid'
    , complexSigmoid
    , complexSigmoid'
    , randComplex

    -- * Training a neural network
    , trainUntil
    , trainNTimes
    , trainUntilErrorBelow
    , quadError

    -- * Loading and saving a neural network
    , loadComplexNetwork
    , saveComplexNetwork
    ) where

import Codec.Compression.Zlib     (compress, decompress)
import Data.Binary                (Binary(..), encode, decode)
import Data.List                  (foldl')
import Foreign.Storable           (Storable)
import qualified Data.ByteString.Lazy  as B
import qualified Data.Vector           as V
import qualified Data.Complex as C
import Debug.Trace

import System.Random.MWC
import Numeric.LinearAlgebra
import Data.Functor ((<$>))

-- | Our feed-forward neural network type. Note the 'Binary' instance, which means you can use
--   'encode' and 'decode' in case you need to serialize your neural nets somewhere else than
--   in a file (e.g over the network)
newtype ComplexNetwork a = ComplexNetwork
                 { matrices   :: V.Vector (Matrix a) -- ^ the weight matrices
                 } deriving Show

instance (Element a, Binary a) => Binary (ComplexNetwork a) where
  put (ComplexNetwork ms) = put . V.toList $ ms
  get = (ComplexNetwork . V.fromList) `fmap` get

--instance (Complex a, Variate a) => Variate (Complex a) where
--	uniform g = uniform g :+ uniform g
--	uniformR (x,y) g = uniformR (x,y) g :+ uniformR (x,y) g

-- | The type of an activation function, mostly used for clarity in signatures
type ActivationFunction a = a -> a

-- | The type of an activation function's derivative, mostly used for clarity in signatures
type ActivationFunctionDerivative a = a -> a

--

randComplex :: (Variate a, Num a) => IO(Complex a)
randComplex = withSystemRandom . asGenST $ \gen -> do
  r <- uniformR (-1,1) gen
  i <- uniformR (-1,1) gen
  return (r :+ i)

randComplexList :: (Variate a, Num a) => Int -> IO([Complex a])
randComplexList 0 = return []
randComplexList n = do
  c <- randComplex
  lst <- randComplexList (n-1)
  return (c : lst)

randComplexMatrix :: (Variate a, Storable a, Num a) => (Int,Int) -> IO(Matrix (Complex a))
randComplexMatrix (rows,cols) = do
  lst <- randComplexList (rows*cols)
  return (reshape cols $ Numeric.LinearAlgebra.fromList lst)

-- | The following creates a neural network with 'n' inputs and if 'l' is [n1, n2, ...]
--   the net will have n1 neurons on the first layer, n2 neurons on the second, and so on
--   ending with k neurons on the output layer, with random weight matrices as a courtesy of
-- 'System.Random.MWC.uniformR'.
-- > createComplexNetwork n l k
createComplexNetwork :: (Variate a, Storable a, Num a) => Int -> [Int] -> Int -> IO (ComplexNetwork (Complex a))
createComplexNetwork nInputs hiddens nOutputs =
  fmap ComplexNetwork $ go dimensions V.empty
  where
        go [] !ms         = return ms
        go ((!n,!m):ds) ms = do
          !mat <- randComplexMatrix (n,m)
          go ds (ms `V.snoc` mat)
        dimensions      = zip (hiddens ++ [nOutputs]) $
                              (nInputs : hiddens)
{-# INLINE createComplexNetwork #-}

-- | Creates a neural network with exactly the weight matrices given as input here.
--   We don't check that the numbers of rows/columns are compatible, etc.
fromWeightMatrices :: Storable a => V.Vector (Matrix (Complex a)) -> ComplexNetwork (Complex a)
fromWeightMatrices ws = ComplexNetwork ws
{-# INLINE fromWeightMatrices #-}

-- The `join [input, 1]' trick  below is a courtesy of Alberto Ruiz
-- <http://dis.um.es/~alberto/>. Per his words:
--
-- "The idea is that the constant input in the first layer can be automatically transferred to the following layers
-- by the learning algorithm (by setting the weights of a neuron to 1,0,0,0,...). This allows for a simpler
-- implementation and in my experiments those networks are able to easily solve non linearly separable problems."

--without the trick (of constant input)

output :: (Floating (Vector a), Product a, Storable a, Num (Vector a)) => ComplexNetwork a -> ActivationFunction a -> Vector a -> Vector a
output (ComplexNetwork{..}) act input = V.foldl' f input matrices
  where f !inp m = mapVector act $ m <> inp
{-# INLINE output #-}

-- | Computes and keeps the output of all the layers of the neural network with the given activation function
outputs :: (Floating (Vector a), Product a, Storable a, Num (Vector a)) => ComplexNetwork a -> ActivationFunction a -> Vector a -> V.Vector (Vector a)
outputs (ComplexNetwork{..}) act input = V.scanl f input matrices
  where f !inp m = mapVector act $ m <> inp
{-# INLINE outputs #-}

conju (x:+y) = x:+(-y)

deltas :: (Floating b, Floating (Vector a), Floating a, Product a, Storable a, Num (Vector a), Container Vector a, a ~ Complex b) => ComplexNetwork a -> ActivationFunctionDerivative a -> V.Vector (Vector a) -> Vector a -> V.Vector (Matrix a)
deltas (ComplexNetwork{..}) act' os expected = V.zipWith outer (V.tail ds) (V.init (V.map (mapVector conju) os))
  where !dl = (V.last os - expected) * (deriv $ mapVector conju (V.last os)) -- = (out_last - target)*(f'(netin_last)) = dev_last * deriv(out_last) = partial_last
        !ds = V.scanr f dl (V.zip os matrices) -- generates dev_i
        f (!o, m) !del = deriv o * ((ctrans m) <> del) -- dev_k = deriv(out_k)*dev_k = deriv(out_k)*(WEIGHTS <> partial_k+1)
        deriv = mapVector act'
{-# INLINE deltas #-}

updateComplexNetwork :: (Floating b, Floating (Vector a), Floating a, Product a, Storable a, Num (Vector a), Container Vector a, a ~ Complex b) => a -> ActivationFunction a -> ActivationFunctionDerivative a -> ComplexNetwork a -> Sample a -> ComplexNetwork a
updateComplexNetwork alpha act act' n@(ComplexNetwork{..}) (input, expectedOutput) = ComplexNetwork $ V.zipWith (+) matrices corr
    where !xs = outputs n act input
          !ds = deltas n act' xs expectedOutput
          !corr = V.map (scale (-alpha)) ds
{-# INLINE updateComplexNetwork #-}

-- | Input vector and expected output vector
type Sample a = (Vector a, Vector a)

-- | List of 'Sample's
type Samples a = [Sample a]

-- | Handy operator to describe your learning set, avoiding unnecessary parentheses. It's just a synonym for '(,)'.
--   Generally you'll load your learning set from a file, a database or something like that, but it can be nice for
--   quickly playing with hnn or for simple problems where you manually specify your learning set.
--   That is, instead of writing:
--
-- > samples :: Samples Double
-- > samples = [ (fromList [0, 0], fromList [0])
-- >           , (fromList [0, 1], fromList [1])
-- >           , (fromList [1, 0], fromList [1])
-- >           , (fromList [1, 1], fromList [0])
-- >           ]
--
--   You can write:
--
-- > samples :: Samples Double
-- > samples = [ fromList [0, 0] --> fromList [0]
-- >           , fromList [0, 1] --> fromList [1]
-- >           , fromList [1, 0] --> fromList [1]
-- >           , fromList [1, 1] --> fromList [0]
-- >           ]
(-->) :: Vector a -> Vector a -> Sample a
(-->) = (,)

backpropOnce :: (Floating (Vector a), Floating b, Floating a, Product a, Num (Vector a), Container Vector a, a ~ Complex b) => a -> ActivationFunction a -> ActivationFunctionDerivative a -> ComplexNetwork a -> Samples a -> ComplexNetwork a
backpropOnce rate act act' n samples = foldl' (updateComplexNetwork rate act act') n samples
{-# INLINE backpropOnce #-}

-- | Generic training function.
--
-- The first argument is a predicate that will tell the backpropagation algorithm when to stop.
-- The first argument to the predicate is the epoch, i.e the number of times the backprop has been
-- executed on the samples. The second argument is /the current network/, and the third is the list of samples.
-- You can thus combine these arguments to create your own criterion.
--
-- For example, if you want to stop learning either when the network's quadratic error on the samples,
-- using the tanh function, is below 0.01, or after 1000 epochs, whichever comes first, you could
-- use the following predicate:
--
-- > pred epochs net samples = if epochs == 1000 then True else quadError tanh net samples < 0.01
--
-- You could even use 'Debug.Trace.trace' to print the error, to see how the error evolves while it's learning,
-- or redirect this to a file from your shell in order to generate a pretty graphics and what not.
--
-- The second argument (after the predicate) is the learning rate. Then come the activation function you want,
-- its derivative, the initial neural network, and your training set.
-- Note that we provide 'trainNTimes' and 'trainUntilErrorBelow' for common use cases.
trainUntil :: (Floating (Vector a), Floating b, Floating a, Product a, Num (Vector a), Container Vector a, a ~ Complex b) => (Int -> ComplexNetwork a -> Samples a -> Bool) -> a -> ActivationFunction a -> ActivationFunctionDerivative a -> ComplexNetwork a -> Samples a -> ComplexNetwork a
trainUntil pr learningRate act act' net samples = go net 0
  where go n !k | pr k n samples = n
                | otherwise      = case backpropOnce learningRate act act' n samples of
                                    n' -> go n' (k+1)
{-# INLINE trainUntil #-}

-- | Trains the neural network with backpropagation the number of times specified by the 'Int' argument,
-- using the given learning rate (second argument).
trainNTimes :: (Floating (Vector a), Floating b, Floating a, Product a, Num (Vector a), Container Vector a, a ~ Complex b) => Int -> a -> ActivationFunction a -> ActivationFunctionDerivative a -> ComplexNetwork a -> Samples a -> ComplexNetwork a
trainNTimes n = trainUntil (\k _ _ -> k > n)
{-# INLINE trainNTimes #-}

-- | Quadratic error on the given training set using the given activation function. Useful to create
-- your own predicates for 'trainUntil'.
quadError :: (Floating (Vector a), Floating b, Floating a, Num (Vector a), Num (RealOf a), Product a, a ~ Complex b) => ActivationFunction a -> ComplexNetwork a -> Samples a -> RealOf a
quadError act net samples = foldl' (\err (inp, out) -> err + (norm2 $ output net act inp - out)) 0 samples
{-# INLINE quadError #-}

-- | Trains the neural network until the quadratic error ('quadError') comes below the given value (first argument),
-- using the given learning rate (second argument).
--
-- /Note/: this can loop pretty much forever when you're using a bad architecture for the problem, or unappropriate activation
-- functions.
trainUntilErrorBelow :: (Floating (Vector a), Floating b, Floating a, Product a, Num (Vector a), Ord a, Container Vector a, Num (RealOf a), a ~ RealOf a, Show a, a ~ Complex b) => a -> a -> ActivationFunction a -> ActivationFunctionDerivative a -> ComplexNetwork a -> Samples a -> ComplexNetwork a
trainUntilErrorBelow x rate act = trainUntil (\_ n s -> quadError act n s < x) rate act
{-# INLINE trainUntilErrorBelow #-}

-- | The sigmoid function:  1 / (1 + exp (-x))
sigmoid :: Floating a => a -> a
sigmoid !x = 1 / (1 + exp (-x))
{-# INLINE sigmoid #-}

-- | Derivative of the sigmoid function: sigmoid x * (1 - sigmoid x)
sigmoid' :: Floating a => a -> a
sigmoid' !x = case sigmoid x of
  s -> s * (1 - s)
{-# INLINE sigmoid' #-}

complexSigmoid :: Floating a => Complex a -> Complex a
complexSigmoid (x :+ y) = (sigmoid x) :+ (sigmoid y)

complexSigmoid' :: Floating a => Complex a -> Complex a
complexSigmoid' (x :+ y) = (sigmoid' x) :+ (sigmoid' y)

-- | Loading a neural network from a file (uses zlib compression on top of serialization using the binary package).
--   Will throw an exception if the file isn't there.
loadComplexNetwork :: (Storable a, Element a, Binary a) => FilePath -> IO (ComplexNetwork a)
loadComplexNetwork fp = decode . decompress <$> B.readFile fp
{-# INLINE loadComplexNetwork #-}

-- | Saving a neural network to a file (uses zlib compression on top of serialization using the binary package).
saveComplexNetwork :: (Storable a, Element a, Binary a) => FilePath -> ComplexNetwork a -> IO ()
saveComplexNetwork fp = B.writeFile fp . compress . encode
{-# INLINE saveComplexNetwork #-}


module Deconvolution where

import           Control.Monad             as M
import           Control.Monad.Parallel    as MP
import           Data.Array.Repa           as R
import           Data.Binary
import           Data.Complex
import           Data.List                 as L
import           DFT.Plan
import           Filter.Utils
import           FokkerPlanck.DomainChange
import           FokkerPlanck.MonteCarlo
import           FokkerPlanck.Pinwheel
import           Graphics.Gnuplot.Simple
import           Image.IO
import           STC
import           System.Directory
import           System.Environment
import           System.FilePath
import           Text.Printf
import           Types
import           Utils.Array

main = do
  args@(numPointStr:numOrientationStr:numScaleStr:thetaSigmaStr:scaleSigmaStr:maxScaleStr:taoStr:numTrailStr:maxTrailStr:theta0FreqsStr:thetaFreqsStr:scale0FreqsStr:scaleFreqsStr:initDistStr:histFilePath:alphaStr:pinwheelFlagStr:numThreadStr:_) <-
    getArgs
  print args
  let numPoint = read numPointStr :: Int
      numOrientation = read numOrientationStr :: Int
      numScale = read numScaleStr :: Int
      thetaSigma = read thetaSigmaStr :: Double
      scaleSigma = read scaleSigmaStr :: Double
      maxScale = read maxScaleStr :: Double
      tao = read taoStr :: Double
      numTrail = read numTrailStr :: Int
      maxTrail = read maxTrailStr :: Int
      theta0Freq = read theta0FreqsStr :: Double
      theta0Freqs = [-theta0Freq .. theta0Freq]
      thetaFreq = read thetaFreqsStr :: Double
      thetaFreqs = [-thetaFreq .. thetaFreq]
      initDist = read initDistStr :: [R2S1RPPoint]
      scale0Freq = read scale0FreqsStr :: Double
      scaleFreq = read scaleFreqsStr :: Double
      scale0Freqs = [-scale0Freq .. scale0Freq]
      scaleFreqs = [-scaleFreq .. scaleFreq]
      alpha = read alphaStr :: Double
      pinwheelFlag = read pinwheelFlagStr :: Bool
      numThread = read numThreadStr :: Int
      folderPath = "output/test/Deconvolution"
      (R2S1RPPoint (_, _, _, s)) = L.head initDist
  createDirectoryIfMissing True folderPath
  flag <- doesFileExist histFilePath
  radialArr <-
    if flag
      then R.map magnitude . getNormalizedHistogramArr <$>
           decodeFile histFilePath
      else do
        putStrLn "Couldn't find a Green's function data. Start simulation..."
        solveMonteCarloR2Z2T0S0Radial
          numThread
          numTrail
          maxTrail
          numPoint
          numPoint
          thetaSigma
          scaleSigma
          maxScale
          tao
          theta0Freqs
          thetaFreqs
          scale0Freqs
          scaleFreqs
          histFilePath
          (emptyHistogram
             [ (round . sqrt . fromIntegral $ 2 * (div numPoint 2) ^ 2)
             , L.length scale0Freqs
             , L.length theta0Freqs
             , L.length scaleFreqs
             , L.length thetaFreqs
             ]
             0)
  radialArrSink <-
    R.map magnitude . getNormalizedHistogramArr <$> decodeFile histFilePath
  arrR2Z2T0S0 <-
    computeUnboxedP $
    computeR2Z2T0S0ArrayRadial
      radialArr
      numPoint
      numPoint
      1
      maxScale
      thetaFreqs
      scaleFreqs
      theta0Freqs
      scale0Freqs
  arrR2Z2T0S0Sink <-
    computeUnboxedP $
    computeR2Z2T0S0ArrayRadial
      (cutoff 24 radialArrSink)
      numPoint
      numPoint
      1
      maxScale
      thetaFreqs
      scaleFreqs
      theta0Freqs
      scale0Freqs
  plan <- makeR2Z2T0S0Plan emptyPlan arrR2Z2T0S0
  sourceDistArr <-
    computeInitialDistributionR2T0S0
      plan
      numPoint
      numPoint
      theta0Freqs
      scale0Freqs
      maxScale
      initDist
  sourceDistArrR2S1RP <-
    r2z2Tor2s1rpP numOrientation thetaFreqs numScale scaleFreqs $ sourceDistArr
  let xIndex =
        [ (fromIntegral i) * 360 / (fromIntegral numOrientation)
        | i <- [0 .. numOrientation - 1]
        ] :: [Double]
      inputStyle =
        L.zipWith
          (\(R2S1RPPoint (x', y', _, _)) i ->
             let x = x' + center numPoint
                 y = y' + center numPoint
              in ( defaultStyle
                     { plotType = LinesPoints
                     , lineSpec =
                         CustomStyle
                           [LineTitle (printf "(%d,%d)" x y), PointType i]
                     }
                 , L.zip xIndex .
                   R.toList . R.slice (R.map magnitude sourceDistArrR2S1RP) $
                   (Z :. All :. (0 :: Int) :. x :. y)))
          initDist
          [1 ..]
  plotPathsStyle [PNG (folderPath </> "Input.png"), Title "Input"] inputStyle
  arrR2Z2T0S0F <-
    dftR2Z2T0S0 plan .
    computeS .
    makeFilter2D .
    R.traverse arrR2Z2T0S0
      -- (R.traverse (arrR2Z2T0S0) id $ \f idx@(Z :. _ :. _ :. _ :. _ :. a :. b) ->
      --    if a == center numPoint && b == center numPoint
      --      then 0
      --      else f idx)
      id $ \f (Z :. tf' :. sf' :. t0f :. s0f :. i :. j) ->
      let idx = (Z :. tf' :. sf' :. t0f :. s0f :. i :. j)
       in if i == center numPoint && j == center numPoint &&
             tf' == t0f
                    -- tf' == div (L.length thetaFreqs) 2 &&
                    -- sf' == div (L.length scaleFreqs) 2 &&
                    -- t0f == div (L.length theta0Freqs) 2 &&
                    -- s0f == div (L.length scale0Freqs) 2
            then f idx + (0.0 :+ 0)
            else f idx
  -- let arr' =
  --       r2z2t0s0Tor2s1rps1rp
  --         numOrientation
  --         thetaFreqs
  --         theta0Freqs
  --         numScale
  --         scaleFreqs
  --         scale0Freqs
  --         maxScale $
  --       arrR2Z2T0S0Sink
  --     arr'' =
  --       r2s1rps1rpTor2z2t0s0
  --         numOrientation
  --         thetaFreqs
  --         theta0Freqs
  --         numScale
  --         scaleFreqs
  --         scale0Freqs
  --         maxScale .
  --       computeUnboxedS .
  --       R.map
  --         (\x ->
  --            let (m, p) = polar x
  --             in if m == 0
  --                  then 0
  --                  else mkPolar (1 / m) p) $
  --       arr'
  arrR2Z2T0S0SinkF <-
    dftR2Z2T0S0 plan .
    computeS .
    makeFilter2D . computeSinkFromSourceR2Z2T0S0 thetaFreqs theta0Freqs
    $
    arrR2Z2T0S0Sink
  --   R.traverse
  --     (-- computeSinkFromSourceR2Z2T0S0 thetaFreqs theta0Freqs $
  --      arrR2Z2T0S0)
  --     id $ \f idx@(Z :. a :. b :. c :. d :. i :. j) ->
  --     let (m, p) =
  --           polar $
  --           conjugate
  --             (f (Z :. (L.length thetaFreqs - 1 - a) :.
  --                 (L.length scaleFreqs - 1 - b) :.
  --                 (L.length theta0Freqs - 1 - c) :.
  --                 (L.length scale0Freqs - 1 - d) :.
  --                 i :.
  --                 j))
  --      in if a == div (L.length thetaFreqs) 2 &&
  --            b == div (L.length scaleFreqs) 2 &&
  --            c == div (L.length theta0Freqs) 2 &&
  --            d == div (L.length scale0Freqs) 2
  --           then 1 + conjugate
  --                       (f (Z :. (L.length thetaFreqs - 1 - a) :.
  --                           (L.length scaleFreqs - 1 - b) :.
  --                           (L.length theta0Freqs - 1 - c) :.
  --                           (L.length scale0Freqs - 1 - d) :.
  --                           i :.
  --                           j))
  --           else conjugate
  --                  (f (Z :. (L.length thetaFreqs - 1 - a) :.
  --                      (L.length scaleFreqs - 1 - b) :.
  --                      (L.length theta0Freqs - 1 - c) :.
  --                      (L.length scale0Freqs - 1 - d) :.
  --                      i :.
  --                      j))
  -- Source field
  sourceArr <- convolveR2T0S0 plan arrR2Z2T0S0F sourceDistArr
  sourceR2Z2' <- R.sumP . R.sumS . rotateR2Z2T0S0Array $ sourceArr
  sourceR2S1RP <-
    r2z2Tor2s1rpP numOrientation thetaFreqs numScale scaleFreqs $ sourceR2Z2'
  let sourceR2Z2 -- = sourceR2Z2'
       =
        R.traverse2 sourceR2Z2' sourceDistArr const $ \f1 f2 idx@(Z :. a :. b :. c :. d) ->
          let x = div (L.length theta0Freqs) 2
              y = div (L.length scale0Freqs) 2
           in if f2 (Z :. x :. y :. c :. d) == 0
                then 0
                else f1 idx
      sourceStyle =
        L.zipWith
          (\(R2S1RPPoint (x', y', _, _)) i ->
             let x = x' + center numPoint
                 y = y' + center numPoint
              in ( defaultStyle
                     { plotType = LinesPoints
                     , lineSpec =
                         CustomStyle
                           [LineTitle (printf "(%d,%d)" x y), PointType i]
                     }
                 , L.zip xIndex .
                   R.toList . R.slice (R.map magnitude sourceR2S1RP) $
                   (Z :. All :. (0 :: Int) :. x :. y)))
          initDist
          [1 ..]
  plotPathsStyle [PNG (folderPath </> "Conv.png"), Title "Conv"] sourceStyle
  sourceField <-
    fmap (computeS . R.extend (Z :. (1 :: Int) :. All :. All)) .
    R.sumP .
    R.sumS .
    rotate4D .
    rotate4D . r2z2Tor2s1rp numOrientation thetaFreqs numScale scaleFreqs $
    sourceR2Z2
  plotImageRepaComplex (folderPath </> "Source.png") . ImageRepa 8 $ sourceField
  -- Deconvolution
  deconvSourceArr <- convolveR2T0S0 plan arrR2Z2T0S0SinkF . computeUnboxedS $ sourceR2Z2
  deconvSourceR2Z2 <-
    R.sumP . R.sumS . rotateR2Z2T0S0Array $ deconvSourceArr
  deconvSourceField <-
    fmap (computeS . R.extend (Z :. (1 :: Int) :. All :. All)) .
    R.sumP .
    R.sumS .
    rotate4D .
    rotate4D . r2z2Tor2s1rp numOrientation thetaFreqs numScale scaleFreqs $
    deconvSourceR2Z2
  plotImageRepaComplex (folderPath </> "DeconvSource.png") . ImageRepa 8 $
    deconvSourceField
  deconvSourceR2S1RP <-
    r2z2Tor2s1rpP numOrientation thetaFreqs numScale scaleFreqs $
    deconvSourceR2Z2
  let outputStyle =
        L.zipWith
          (\(R2S1RPPoint (x', y', _, _)) i ->
             let x = x' + center numPoint
                 y = y' + center numPoint
              in ( defaultStyle
                     { plotType = LinesPoints
                     , lineSpec =
                         CustomStyle
                           [LineTitle (printf "(%d,%d)" x y), PointType i]
                     }
                 , L.zip xIndex .
                   R.toList . R.slice (R.map magnitude deconvSourceR2S1RP) $
                   (Z :. All :. (0 :: Int) :. x :. y)
                     -- (R.sumS .
                   --    R.backpermute
                   --      (Z :. (L.length theta0Freqs) :. numPoint :. numPoint :.
                   --       (L.length scale0Freqs))
                   --      (\(Z :. a :. b :. c :. d) -> (Z :. a :. d :. b :. c)) .
                   --    R.map magnitude $
                   --    deconvSourceR2S1RP) $
                   -- (Z :. All :. x :. y)
                  ))
          initDist
          [1 ..]
  plotPathsStyle [PNG (folderPath </> "Output.png"), Title "Output"] outputStyle


{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
-----------------------------------------------------------------------------
-- |
-- Module      :  Numeric.Signal.Internal
-- Copyright   :  (c) Alexander Vivian Hugh McPhail 2010, 2014, 2015, 2016
-- License     :  BSD3
--
-- Maintainer  :  haskell.vivian.mcphail <at> gmail <dot> com
-- Stability   :  provisional
-- Portability :  uses FFI
--
-- low-level interface
--
-----------------------------------------------------------------------------

module Numeric.Signal.Internal (
                Convolvable(..),
                Filterable(..),
                freqz,
                pwelch,
                hilbert
                ) where

import Numeric.LinearAlgebra
import Numeric.LinearAlgebra.Devel

--import Numeric.LinearAlgebra.Linear

import qualified Numeric.GSL.Fourier as F
import Foreign

--import Data.Complex
import Foreign.C.Types

import Prelude hiding(filter)

import System.IO.Unsafe(unsafePerformIO)

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

infixr 1 #
a # b = applyRaw a b
{-# INLINE (#) #-}

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

type PD = Ptr Double
type PC = Ptr (Complex Double)
type PF = Ptr Float

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

class Convolvable a where
    -- | convolve two containers, output is the size of the second argument, no zero padding
    convolve :: a -> a -> a

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

class (Storable a, Container Vector a, Num (Vector a)
      , Convert a, Floating (Vector a), RealElement a
      , Num a)
       => Filterable a where
    -- | convert from Vector Double
    fromDouble :: Vector Double -> Vector a
--       b ~ ComplexOf a, Container Vector b, Convert b) => Filterable a where
    -- | filter a signal
    filter_ :: Vector a -- ^ zero coefficients
            -> Vector a -- ^ pole coefficients
            -> Vector a -- ^ input signal
            -> Vector a -- ^ output signal
    -- | coefficients of a Hamming window
    hamming_ :: Int           -- ^ length
           -> Vector a -- ^ the Hamming coeffficents
    -- | the complex power : real $ v * (conj v)
    complex_power_ :: Vector (Complex Double) -- ^ input
                  -> Vector a                       -- ^ output
     -- | resample, take one sample every n samples in the original
    downsample_ :: Int -> Vector a -> Vector a
    -- | the difference between consecutive elements of a vector
    deriv_ :: Vector a -> Vector a
    -- | unwrap the phase of signal (input expected to be within (-pi,pi)
    unwrap_ :: Vector a -> Vector a
    -- | evaluate a real coefficient polynomial for complex arguments
    polyEval_ :: Vector a           -- ^ the real coefficients
         -> Vector (Complex Double) -- ^ the points at which to be evaluated
         -> Vector (Complex Double) -- ^ the values
    -- | the cross covariance of two signals
    cross_covariance_ :: Int            -- ^ maximum delay
                     -> Vector a       -- ^ time series
                     -> Vector a       -- ^ time series
                     -> (a,a,Vector a) -- ^ (sd_x,sd_y,cross_covariance)
    -- | the cumulative sum of a signal
    cumulative_sum_ :: Vector a    -- ^ time series
                    -> Vector a    -- ^ result

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

instance Convolvable (Vector Double) where
    convolve x y = fst $ fromComplex $ F.ifft $ (F.fft (complex x) * F.fft (complex y))
--    convolve = convolve_vector_double

convolve_vector_double c a = unsafePerformIO $ do
                             r <- createVector (size a)
                             (c # a # r # id) signal_vector_double_convolve #| "signalDoubleConvolve"
                             return r

foreign import ccall "signal-aux.h vector_double_convolve" signal_vector_double_convolve :: CInt -> PD -> CInt -> PD -> CInt -> PD -> IO CInt

instance Convolvable (Vector Float) where
    convolve x y = single $ fst $ fromComplex $ F.ifft $ (F.fft (complex $ double x) * F.fft (complex $ double y))
--    convolve = convolve_vector_double

convolve_vector_float c a = unsafePerformIO $ do
                             r <- createVector (size a)
                             (c # a # r # id ) signal_vector_float_convolve #| "signalFloatConvolve"
                             return r

foreign import ccall "signal-aux.h vector_float_convolve" signal_vector_float_convolve :: CInt -> PF -> CInt -> PF -> CInt -> PF -> IO CInt

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

instance Convolvable (Vector (Complex Double)) where
    convolve x y = F.ifft $ (F.fft x * F.fft y)
--    convolve = convolve_vector_complex

convolve_vector_complex c a = unsafePerformIO $ do
                              r <- createVector (size a)
                              (c # a # r # id) signal_vector_complex_convolve #| "signalComplexConvolve"
                              return r

foreign import ccall "signal-aux.h vector_complex_convolve" signal_vector_complex_convolve :: CInt -> PC -> CInt -> PC -> CInt -> PC -> IO CInt

instance Convolvable (Vector (Complex Float)) where
    convolve x y = single $ F.ifft $ (F.fft (double x) * F.fft (double y))

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

instance Filterable Double where
    fromDouble = id
    filter_ = filterD
    hamming_ = hammingD
    complex_power_ = complex_powerD
    downsample_ = downsampleD
    deriv_ = derivD
    unwrap_ = unwrapD
    polyEval_ = polyEval
    cross_covariance_ = crossCovarianceD
    cumulative_sum_ = cumSumD

instance Filterable Float where
    fromDouble = single
    filter_ = filterF
    hamming_ = hammingF
    complex_power_ = complex_powerF
    downsample_ = downsampleF
    deriv_ = derivF
    unwrap_ = unwrapF
    polyEval_ c = polyEval (double c)
    cross_covariance_ = crossCovarianceF
    cumulative_sum_ = cumSumF

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

-- | filters the signal
filterD :: Vector Double -- ^ zero coefficients
       -> Vector Double -- ^ pole coefficients
       -> Vector Double -- ^ input signal
       -> Vector Double -- ^ output signal
filterD l k v = unsafePerformIO $ do
               r <- createVector (size v)
               (l # k # v # r # id) signal_filter_double #| "signalFilter"
               return r

foreign import ccall "signal-aux.h filter_double" signal_filter_double :: CInt -> PD -> CInt -> PD -> CInt -> PD -> CInt -> PD -> IO CInt

-- | filters the signal
filterF :: Vector Float -- ^ zero coefficients
       -> Vector Float -- ^ pole coefficients
       -> Vector Float -- ^ input signal
       -> Vector Float -- ^ output signal
filterF l k v = unsafePerformIO $ do
               r <- createVector (size v)
               (l # k # v # r # id) signal_filter_float #| "signalFilter"
               return r

foreign import ccall "signal-aux.h filter_float" signal_filter_float :: CInt -> PF -> CInt -> PF -> CInt -> PF -> CInt -> PF -> IO CInt

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

-- | Hilbert transform with original vector as real value, transformed as imaginary
hilbert :: Vector Double -> Vector (Complex Double)
hilbert v = unsafePerformIO $ do
            let r = complex v
            -- could use (complex v) to make a complex vector in haskell rather than C
            (r # id) signal_hilbert #| "hilbert"
            return r

foreign import ccall "signal-aux.h hilbert" signal_hilbert :: CInt -> PC -> IO CInt

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

-- | Welch (1967) power spectrum density using periodogram/FFT method
pwelch :: Int            -- ^ window size (multiple of 2)
       -> Vector Double  -- ^ input signal
       -> Vector Double  -- ^ power density
pwelch w v = unsafePerformIO $ do
             let r = konst 0.0 ((w `div` 2) + 1)
             (complex v # r # id) (signal_pwelch $ fromIntegral w) #| "pwelch"
             return r

foreign import ccall "signal-aux.h pwelch" signal_pwelch :: CInt -> CInt -> PC -> CInt -> PD -> IO CInt

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

-- | coefficients of a Hamming window
hammingD :: Int           -- ^ length
        -> Vector Double -- ^ the Hamming coeffficents
hammingD l
    | l == 1          = konst 1.0 1
    | otherwise       = unsafePerformIO $ do
                        r <- createVector l
                        (r # id) signal_hamming_double #| "Hamming"
                        return r

foreign import ccall "signal-aux.h hamming_double" signal_hamming_double :: CInt -> PD -> IO CInt

-- | coefficients of a Hamming window
hammingF :: Int           -- ^ length
        -> Vector Float -- ^ the Hamming coeffficents
hammingF l
    | l == 1          = konst 1.0 1
    | otherwise       = unsafePerformIO $ do
                        r <- createVector l
                        (r # id) signal_hamming_float #| "Hamming"
                        return r

foreign import ccall "signal-aux.h hamming_float" signal_hamming_float :: CInt -> PF -> IO CInt

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

-- | determine the frequency response of a filter
{-freqz :: (Filterable a, Storable a, Container Vector a, Convert a, RealElement a,
         DoubleOf a ~ DoubleOf (RealOf b), RealElement c, c ~ DoubleOf a, c ~ DoubleOf (RealOf b),
         b ~ Complex a, b ~ ComplexOf a, Convert b, Container Vector b,
         Container Vector c, Convert c, b ~ ComplexOf a, b ~ ComplexOf c)
        ⇒ Vector a     -- ^ zero coefficients
      -> Vector a       -- ^ pole coefficients
      -> Vector a       -- ^ points (between 0 and 2*pi)
      -> Vector a       -- ^ response
-}
freqz :: (Filterable a, Complex Double ~ ComplexOf (DoubleOf a)
        ,Filterable (DoubleOf a)) =>
        Vector a       -- ^ zero coefficients
      -> Vector a       -- ^ pole coefficients
      -> Vector a       -- ^ points (between 0 and 2*pi)
      -> Vector a       -- ^ response
freqz b a w = let k = max (size b) (size a)
                  hb = polyEval_ (postpad b k) (exp (scale (0 :+ 1) ((complex $ double w))))
                  ha = polyEval_ (postpad a k) (exp (scale (0 :+ 1) ((complex $ double w))))
              in complex_power_ (hb / ha)

postpad v n = let d = size v
              in if d < n then vjoin [v,(konst 0.0 (n-d))]
              else v

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

-- | evaluate a real coefficient polynomial for complex arguments
polyEval :: Vector Double           -- ^ the real coefficients
         -> Vector (Complex Double) -- ^ the points at which to be evaluated
         -> Vector (Complex Double) -- ^ the values
polyEval c z = unsafePerformIO $ do
               r <- createVector (size z)
               (c # z # r # id) signal_real_poly_complex_eval #| "polyEval"
               return r

foreign import ccall "signal-aux.h real_poly_complex_eval" signal_real_poly_complex_eval :: CInt -> PD -> CInt -> PC -> CInt -> PC -> IO CInt

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

-- | the complex power : real $ v * (conj v)
complex_powerD :: Vector (Complex Double) -- ^ input
              -> Vector Double           -- ^ output
complex_powerD v = unsafePerformIO $ do
                  r <- createVector (size v)
                  (v # r # id) signal_complex_power_double #| "complex_power"
                  return r

foreign import ccall "signal-aux.h complex_power_double" signal_complex_power_double :: CInt -> PC -> CInt -> PD -> IO CInt

-- | the complex power : real $ v * (conj v)
complex_powerF :: Vector (Complex Double) -- ^ input
              -> Vector Float             -- ^ output
complex_powerF v = unsafePerformIO $ do
                  r <- createVector (size v)
                  (v # r # id) signal_complex_power_float #| "complex_power"
                  return r

foreign import ccall "signal-aux.h complex_power_float" signal_complex_power_float :: CInt -> PC -> CInt -> PF -> IO CInt

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

-- | resample, take one sample every n samples in the original
downsampleD :: Int -> Vector Double -> Vector Double
downsampleD n v = unsafePerformIO $ do
               r <- createVector (size v `div` n)
               (v # r # id) (signal_downsample_double $ fromIntegral n) #| "downsample"
               return r

foreign import ccall "signal-aux.h downsample_double" signal_downsample_double :: CInt -> CInt -> PD -> CInt -> PD -> IO CInt

-- | resample, take one sample every n samples in the original
downsampleF :: Int -> Vector Float -> Vector Float
downsampleF n v = unsafePerformIO $ do
               r <- createVector (size v `div` n)
               (v # r # id) (signal_downsample_float $ fromIntegral n) #| "downsample"
               return r

foreign import ccall "signal-aux.h downsample_float" signal_downsample_float :: CInt -> CInt -> PF -> CInt -> PF -> IO CInt

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

-- | the difference between consecutive elements of a vector
derivD :: Vector Double -> Vector Double
derivD v = unsafePerformIO $ do
          r <- createVector (size v - 1)
          (v # r # id) (signal_diff_double) #| "diff"
          return r

foreign import ccall "signal-aux.h vector_diff_double" signal_diff_double :: CInt -> PD -> CInt -> PD -> IO CInt

-- | the difference between consecutive elements of a vector
derivF :: Vector Float -> Vector Float
derivF v = unsafePerformIO $ do
          r <- createVector (size v - 1)
          (v # r # id) (signal_diff_float) #| "diff"
          return r

foreign import ccall "signal-aux.h vector_diff_float" signal_diff_float :: CInt -> PF -> CInt -> PF -> IO CInt

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

-- | unwrap the phase of signal (input expected to be within (-pi,pi)
unwrapD :: Vector Double -> Vector Double
unwrapD v = unsafePerformIO $ do
           r <- createVector $ size v
           (v # r # id) signal_unwrap_double #| "unwrap"
           return r

foreign import ccall "signal-aux.h unwrap_double" signal_unwrap_double :: CInt -> PD -> CInt -> PD -> IO CInt

-- | unwrap the phase of signal (input expected to be within (-pi,pi)
unwrapF :: Vector Float -> Vector Float
unwrapF v = unsafePerformIO $ do
           r <- createVector $ size v
           (v # r # id) signal_unwrap_float #| "unwrap"
           return r

foreign import ccall "signal-aux.h unwrap_float" signal_unwrap_float :: CInt -> PF -> CInt -> PF -> IO CInt

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

-- | compute the cross covariance of two signals
crossCovarianceD :: Int -> Vector Double -> Vector Double -> (Double,Double,Vector Double)
crossCovarianceD l x y = unsafePerformIO $ do
                           r <- createVector (2*l)
                           alloca $ \sx ->
                               alloca $ \sy -> do
                                 (x # y # r # id) (signal_cross_covariance_double (fromIntegral l) sx sy) #| "cross_covariance"
                                 sx' <- peek sx
                                 sy' <- peek sy
                                 return (sx',sy',r)

foreign import ccall "signal-aux.h cross_covariance_double"
        signal_cross_covariance_double :: CInt -> PD -> PD -> CInt -> PD -> CInt
                                       -> PD -> CInt -> PD -> IO CInt

-- | compute the cross covariance of two signals
crossCovarianceF :: Int -> Vector Float -> Vector Float -> (Float,Float,Vector Float)
crossCovarianceF l x y = unsafePerformIO $ do
                           r <- createVector (2*l)
                           alloca $ \sx ->
                               alloca $ \sy -> do
                                 (x # y # r # id) (signal_cross_covariance_float (fromIntegral l) sx sy) #| "cross_covariance"
                                 sx' <- peek sx
                                 sy' <- peek sy
                                 return (sx',sy',r)

foreign import ccall "signal-aux.h cross_covariance_float"
        signal_cross_covariance_float :: CInt -> PF -> PF -> CInt -> PF -> CInt
                                       -> PF -> CInt -> PF -> IO CInt

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

cumSumD :: Vector Double -> Vector Double
cumSumD v = unsafePerformIO $ do
              r <- createVector (size v)
              (v # r # id) signal_cum_sum_double #| "cumSumD"
              return r

cumSumF :: Vector Float -> Vector Float
cumSumF v = unsafePerformIO $ do
              r <- createVector (size v)
              (v # r # id) signal_cum_sum_float #| "cumSumF"
              return r

foreign import ccall "signal-aux.h cum_sum_double"
        signal_cum_sum_double :: CInt -> PD -> CInt -> PD -> IO CInt
foreign import ccall "signal-aux.h cum_sum_float"
        signal_cum_sum_float :: CInt -> PF -> CInt -> PF -> IO CInt

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


{-# LANGUAGE DataKinds             #-}
{-# LANGUAGE DeriveAnyClass        #-}
{-# LANGUAGE DeriveGeneric         #-}
{-# LANGUAGE FlexibleInstances     #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE RankNTypes            #-}
{-# LANGUAGE ScopedTypeVariables   #-}
{-# LANGUAGE TypeFamilies          #-}
{-# LANGUAGE TypeOperators         #-}
{-|
Module      : Grenade.Core.Softmax
Description : Softmax loss layer
Copyright   : (c) Huw Campbell, 2016-2017
License     : BSD2
Stability   : experimental
-}
module Grenade.Layers.Softmax (
    Softmax (..)
  , softmax
  , softmax'
  , SpecSoftmax (..)
  , specSoftmax
  , softmaxLayer
  ) where

import           Data.Serialize

import           Control.DeepSeq                (NFData (..))
import           Data.Reflection                (reifyNat)
import           Data.Singletons
import           GHC.Generics                   (Generic)
import           GHC.TypeLits
import           Numeric.LinearAlgebra.Static   as LAS

import           Grenade.Core
import           Grenade.Dynamic
import           Grenade.Dynamic.Internal.Build

-- | A Softmax layer
--
--   This layer is like a logit layer, but normalises
--   a set of matricies to be probabilities.
--
--   One can use this layer as the last layer in a network
--   if they need normalised probabilities.
data Softmax = Softmax
  deriving (Show, Generic, NFData)

instance UpdateLayer Softmax where
  type Gradient Softmax = ()
  runUpdate _ _ _ = Softmax

instance RandomLayer Softmax where
  createRandomWith _ _ = return Softmax

instance ( KnownNat i ) => Layer Softmax ('D1 i) ('D1 i) where
  type Tape Softmax ('D1 i) ('D1 i) = S ('D1 i)

  runForwards _ (S1D y) = (S1D y, S1D (softmax y))
  runBackwards _ (S1D y) (S1D dEdy) = ((), S1D (softmax' y dEdy))

instance Serialize Softmax where
  put _ = return ()
  get = return Softmax

softmax :: KnownNat i => LAS.R i -> LAS.R i
softmax xs =
  let xs' = LAS.dvmap exp xs
      s   = LAS.dot xs' 1
  in  LAS.dvmap (/ s) xs'

softmax' :: KnownNat i => LAS.R i -> LAS.R i -> LAS.R i
softmax' x grad =
  let yTy = outer sm sm
      d   = diag sm
      g   = d - yTy
  in  g #> grad
    where
  sm = softmax x

-------------------- DynamicNetwork instance --------------------

instance FromDynamicLayer Softmax where
  fromDynamicLayer inp _ Softmax = case tripleFromSomeShape inp of
    (rows, 1, 1) -> SpecNetLayer $ SpecSoftmax rows
    _ -> error "Error in specification: The layer Softmax may only be used with 1D input!"

instance ToDynamicLayer SpecSoftmax where
  toDynamicLayer _ _ (SpecSoftmax rows) =
    reifyNat rows $ \(_ :: (KnownNat i) => Proxy i) ->
    return $ SpecLayer Softmax (sing :: Sing ('D1 i)) (sing :: Sing ('D1 i))

-- | Create a specification for a elu layer.
specSoftmax :: Integer -> SpecNet
specSoftmax = SpecNetLayer . SpecSoftmax

-- | Add a Softmax layer to your build.
softmaxLayer :: BuildM ()
softmaxLayer = buildRequireLastLayerOut Is1D >>= buildAddSpec . SpecNetLayer . SpecSoftmax . fst3
  where
    fst3 (x, _, _) = x

-------------------- GNum instances --------------------

instance GNum Softmax where
  _ |* Softmax = Softmax
  _ |+ Softmax = Softmax


module Statistics.GModeling.Models.HMM
  (
  ) where

import Statistics.GModeling.DSL
import Statistics.GModeling.Gibbs
import qualified Data.Vector as V
import qualified Data.Vector.Unboxed as U
import qualified Data.Vector.Generic as G
import qualified Data.Vector.Unboxed.Mutable as MU

data HMM = Alpha | Beta | Transition
         | Topic | Symbols | Symbol

hmm :: Network HMM
hmm =
  [
    Only Alpha            :-> Transition
  , Only Beta             :-> Symbols
  , (Transition :@ Topic) :-> Topic
  , (Symbols :@ Topic)    :-> Symbol
  ]

type D = V.Vector (U.Vector (Int,Int))

reader :: D -> Reader HMM Int
reader v = Reader
  {
    size = G.sum (G.map G.length v)
  , readn = \idx k ->
      let (i,j) = indices G.! idx
          (topic,symbol) = v G.! i G.! j
          (prev_topic,_) = v G.! i G.! (j-1)
      in case k of
           Topic -> topic
           Symbol -> symbol
           Symbols -> topic
           Transition -> if j==0
                         -- reserve Topic 0 for initial distribution
                         then 0
                         else prev_topic
  , copy = do
      m <- G.mapM U.thaw v
      return Writer
        {
          writen = \idx k v -> do
            let (i,j) = indices G.! idx
            case k of
              Topic -> do
                (_,s) <- MU.unsafeRead (m G.! i) j
                MU.unsafeWrite (m G.! i) j (v,s)
        , readOnly = G.mapM U.freeze m >>= return . reader
        }
  }
  where indices = U.fromList $ do
          i <- [0..G.length v-1]
          j <- [0..G.length (v G.! i)-1]
          return (i,j)


{-# LANGUAGE EmptyDataDecls #-}
{-# LANGUAGE ExplicitForAll #-}
--{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE FunctionalDependencies #-}
{-# LANGUAGE MultiParamTypeClasses #-}
--{-# LANGUAGE PartialTypeSignatures #-}
{-# LANGUAGE RankNTypes #-}
--{-# LANGUAGE RebindableSyntax #-}
--{-# LANGUAGE ScopedTypeVariables #-}

{-# LANGUAGE OverloadedLists #-}
--{-# LANGUAGE NamedFieldPuns #-}
{-# OPTIONS_GHC -XBangPatterns #-}

module Test.Random ( testRandom ) where

import Prelude.Extended
-- import Control.Monad.Eff (Eff)
-- import Control.Monad.Eff.Console ( CONSOLE, log )
-- import Control.Monad.Eff.Random ( RANDOM )
-- import Data.List.Lazy ( replicateM )
import Data.Vector.Unboxed as A
  ( zipWith, fromList )
-- import Data.Traversable ( for )
-- import Data.Tuple ( Tuple (..) )
import Statistics.Sample ( meanVariance )
import System.Random ( RandomGen, newStdGen )
import Data.Time.Clock
-- import Text.Printf
-- import Control.Parallel
-- import Control.Parallel.Strategies

import Data.Cov ( chol, fromArray, toArray, (*.), Vec5 )
import FV.Fit ( fit )
import FV.Types ( VHMeas(..), HMeas(..), MMeas(..)
  , invMass, fromQMeas, fitMomenta )

{-- import qualified Graphics.Gnuplot.Frame.OptionSet as Opts --}
{-- import Graphics.Histogram --}

-- | Randomizable TypeClass to provide randomize method
-- | for MC smearing of a measurement
class Randomizable a where
  randomize :: RandomGen g => a -> g -> (a, g)
-- | randomize a single helix parameters measurement, based on the cov matrix
-- | return randomized helix
instance Randomizable HMeas where
  randomize (HMeas h hh w0) g = (HMeas h' hh w0, g') where
    (rs, g') = normals 5 g
    r5 :: Vec5
    r5 = fromArray rs
    h' = fromArray $ A.zipWith (+) (toArray h) (toArray (chol hh *. r5))

-- | randomize a vertex measurement by randomizing each helix parameter measurement
-- | leaving the initial vertex untouched
instance Randomizable VHMeas where
  randomize (VHMeas { vertex= v, helices= hl}) g =
    (VHMeas { vertex= v, helices= hl' }, g') where
      doit :: RandomGen g => HMeas -> (List HMeas, g) -> (List HMeas, g)
      doit h (hs, g) = (hs', g') where
        (h', g') = randomize h g
        hs' = h' : hs
      (hl', g') = foldr doit ([], g) hl

-- calc fitted invariant mass of VHMeas
fitm :: VHMeas -> Number
fitm vm = m where
  MMeas {m= m} = invMass <<< map fromQMeas <<< fitMomenta $ fit vm

testRandom ::Int
             -> VHMeas
             -> IO Text
testRandom cnt vm = do
  g <- newStdGen
  t0 <- getCurrentTime
  let
      mf = invMass <<< map fromQMeas <<< fitMomenta <<< fit $ vm
      !mr = MMeas {m=m, dm=sqrt dm2}
      (m, dm2) = meanVariance $ A.fromList ms
      ls :: [Int]
      ls = [0 .. (cnt-1)]
      (ms, _) = foldl doit ([], g) ls where
        doit :: RandomGen g => (List Number, g) -> Int -> (List Number, g)
        doit (ms, g) _ = (ms', g') where
          (vm', g') = randomize vm g
          m = fitm vm'
          ms' = m : ms
  t1 <- getCurrentTime

  -- ms = map func 1..cnt
  -- func vm g = (m, g') where
  --   (vm', g') = randomize vm g
  --   m' = fitm vm'

  pure $ "Fit Mass  " <> tshow mf
       <> "\nMean Mass " <> tshow mr
       <> "\ntime: "
            <> to2fix (realToFrac (diffUTCTime t1 t0) :: Number)
            <> " seconds"
  {-- let hist = histogram binSturges (V.toList hf) --}
  {-- _ <- plot "invMass.png" hist --}


{-# LANGUAGE UndecidableInstances #-}
{-# LANGUAGE QuantifiedConstraints #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE MultiParamTypeClasses #-}
module Q.Interpolation where
import qualified Q.SortedVector as SV
import Numeric.GSL.Interpolation
import qualified Numeric.LinearAlgebra as V (Vector, fromList)
import Foreign (Storable)
import Data.List
class (Ord k, Storable k, Storable v) => Interpolator a k v where
  interpolate :: a -> [(k, v)] -> k -> v

class (Ord k, Storable k, Storable v) => InterpolatorV a k v where
  interpolateV :: a -> SV.SortedVector k -> V.Vector v -> k -> v

instance (Ord k, Storable k, Storable v, InterpolatorV a k v) => Interpolator a k v where
  interpolate a pts = interpolateV a xs' ys' where
    (xs, ys) = (unzip . sortOn fst) pts
    xs'      = SV.fromSortedList xs
    ys'      = V.fromList ys