{-# LANGUAGE OverloadedStrings #-}

-- | Using statistics-linreg to find a given bug date
module Main
  ( main,
  )
where

import Control.Exception (catch)
import Data.Maybe (fromJust, isJust)
import Data.Time.Clock (UTCTime)
import Data.Time.Clock.POSIX (POSIXTime, posixSecondsToUTCTime, utcTimeToPOSIXSeconds)
import qualified Data.Vector as V
import Statistics.LinearRegression (linearRegression)
import System.Environment (getArgs)
import Web.Bugzilla hiding (bugId)

getTime :: BugzillaSession -> BugId -> IO (Maybe UTCTime)
getTime session bugId = catch go ignoreException
  where
    go :: IO (Maybe UTCTime)
    go = do
      bug <- getBug session bugId
      return $ case bug of
        Just bug' -> Just $ bugCreationTime bug'
        Nothing -> Nothing
    ignoreException :: BugzillaException -> IO (Maybe UTCTime)
    ignoreException _ = return Nothing

getTimes :: [BugId] -> IO [Maybe UTCTime]
getTimes bugs = do
  session <- anonymousSession <$> newBugzillaContext "bugzilla.redhat.com"
  mapM (getTime session) bugs

filterBadBugs :: [BugId] -> [Maybe UTCTime] -> ([BugId], [UTCTime])
filterBadBugs bugs times = (goodBugs, map fromJust goodTimes)
  where
    (goodBugs, goodTimes) = unzip $ filter (isJust . snd) $ zip bugs times

-- >>> let bugs = [0, 100]
-- >>> now <- getCurrentTime
-- >>> let yesterday = addUTCTime (-3600*24) now
-- >>> let dates = [yesterday, now]
-- >>> predictTime' bugs dates 200
-- tomorrow -- in utc time
predictTime' :: [BugId] -> [UTCTime] -> BugId -> UTCTime
predictTime' bugs times target = posixSecondsToUTCTime $ predict $ linearRegression xs ys
  where
    predict :: (Double, Double) -> POSIXTime
    predict (alpha, beta) = realToFrac ((fromIntegral target - alpha) / beta)
    ys = V.fromList $ map fromIntegral bugs
    xs = V.fromList $ map (realToFrac . utcTimeToPOSIXSeconds) times

predictTime :: BugId -> Int -> BugId -> IO UTCTime
predictTime highestBug sampleSize target =
  do
    sampleTimes <- getTimes sampleBugs
    let (bugs, times) = filterBadBugs sampleBugs sampleTimes
    print $ "Sampling with " <> show (length bugs) <> " bugs from " <> show (last bugs) <> " to " <> show (head bugs)
    return $ predictTime' bugs times target
  where
    sampleBugs = take sampleSize $ getRange highestBug 1
    -- Look for bugId spread apart
    getRange cur step
      | cur < 0 = []
      | otherwise = cur : getRange (cur - step) (step + 20)

main :: IO ()
main = do
  args <- getArgs
  case fmap read args of
    [highestBug, sampleSize, target] -> do
      prediction <- predictTime highestBug sampleSize target
      print prediction
    _ -> putStrLn "usage: bugzilla-demo highest-bug sample-size target-bug"


{-# LANGUAGE BangPatterns #-}

-- |
-- Module    : Statistics.Matrix.Fast.Algorithms
-- Copyright : 2019 Magalame
-- License   : BSD3
--
-- Useful matrix functions.

module Statistics.Matrix.Fast.Algorithms
    (
      qr
    ) where

import Control.Applicative ((<$>), (<*>))
import Control.Monad.ST (ST, runST)
import Prelude hiding (replicate)
import Statistics.Matrix (Matrix (..),dimension, for)
import qualified Statistics.Matrix.Mutable as M
import qualified Data.Vector.Unboxed as U

-- | /O(r*c)/ Compute the QR decomposition of a matrix.
-- The result returned is the matrices (/q/,/r/).
qr :: Matrix -> (Matrix, Matrix)
qr mat = runST $ do
  let (m,n) = dimension mat

  r <- M.replicate n n 0
  a <- M.thaw mat
  for 0 n $ \j -> do
    cn <- M.immutably a $ \aa -> sqrt $ normCol j aa
    M.unsafeWrite r j j cn
    for 0 m $ \i -> M.unsafeModify a i j (/ cn)
    for (j+1) n $ \jj -> do
      p <- innerProduct a j jj
      M.unsafeWrite r j jj p
      for 0 m $ \i -> do
        aij <- M.unsafeRead a i j
        M.unsafeModify a i jj $ subtract (p * aij)
  (,) <$> M.unsafeFreeze a <*> M.unsafeFreeze r

normCol :: Int -> Matrix -> Double
normCol jthcol (Matrix r c v) = sub 0 0
  where sub !acc !ij | ij == r = acc
                     | otherwise = sub ( valCol*valCol + acc ) (ij+1)
                                   where
                                    valCol = U.unsafeIndex v (ij*c+jthcol)

innerProduct :: M.MMatrix s -> Int -> Int -> ST s Double
innerProduct mmat j k = M.immutably mmat $ \mat ->
  dotCol j mat k mat

dotCol :: Int -> Matrix -> Int -> Matrix -> Double
dotCol jthcol (Matrix r1 c1 v1) kthcol (Matrix _ c2 v2) = sub 0 0
  where sub !acc !ij | ij == r1 = acc
                     | otherwise = sub ( valColj*valColk + acc ) (ij+1)
                                   where
                                    valColk = U.unsafeIndex v2 (ij*c2+kthcol)
                                    valColj = U.unsafeIndex v1 (ij*c1+jthcol)


module LogisticRegression where

import Numeric.LinearAlgebra

sigmoid :: Matrix Double -> Matrix Double
sigmoid x = fromRows $ map (cmap sigmoid') (toRows x)

-- sigmoid applied to a single real number.
sigmoid' :: Double -> Double
sigmoid' x = 1 / (1 + (exp (-1 * x)))


{-# LANGUAGE TemplateHaskell #-}
{-# LANGUAGE DataKinds #-}
module Lib
    (
    ) where
import Control.Lens
import Numeric.LinearAlgebra.Static as NLS
import Numeric.Backprop
import Debug.SimpleReflect as SR
data Network = Net { _weight1 :: NLS.L 20 100
                   , _bias1   :: NLS.R 20
                   , _weight2 :: NLS.L  5  20
                   , _bias2   :: NLS.R  5
                   }
Control.Lens.makeLenses ''Network
runNet :: Network -> NLS.R 100 -> NLS.R 5
runNet net x = z
  where
    y = logistic $ (net ^. weight1) #> x + (net ^. bias1)
    z = logistic $ (net ^. weight2) #> y + (net ^. bias2)
logistic :: Floating a => a -> a
logistic x = 1 / (1 + exp (-x))
myFun :: Floating a => a -> a -> a
myFun x y = x + y
myFunc :: Floating a => a -> a
myFunc x = sqrt (x*4)
{-
data Point = Point { _x :: Double, _y :: Double }deriving (Show)
data Atom = Atom { _element :: String, _point :: Point } deriving (Show)
data Molecule = Molecule { _atoms :: [Atom] } deriving (Show)
Control.Lens.makeLenses ''Atom
Control.Lens.makeLenses ''Point
Control.Lens.makeLenses ''Molecule
shiftAtomX :: Atom -> Atom
shiftAtomX = over (point . x) (+ 1)
shiftMoleculeX :: Molecule -> Molecule
shiftMoleculeX = over (atoms . traverse . point . x) (+ 1)
atom1 = Atom { _element = "C", _point = Point { _x = 1.0, _y = 2.0 } }
atom2 = Atom { _element = "O", _point = Point { _x = 3.0, _y = 4.0 } }
molecule = Molecule { _atoms = [atom1, atom2] }
newMolecule = shiftMoleculeX molecule
t = atom1 ^. (point . x)
-}


module Tests.Parametric (tests) where

import Data.Maybe (fromJust)
import Statistics.Test.StudentT
import Statistics.Types
import qualified Data.Vector.Unboxed as U
import Test.Tasty (testGroup)
import Tests.Helpers  (testEquality)
import qualified Test.Tasty as Tst

tests :: Tst.TestTree
tests = testGroup "Parametric tests" studentTTests

-- 2 samples x 20 obs data
--
-- Both samples are samples from normal distributions with the same variance (= 1.0),
-- but their means are different (0.0 and 0.5, respectively).
--
-- You can reproduce the data with R (3.1.0) as follows:
--   set.seed(0)
--   sample1 = rnorm(20)
--   sample2 = rnorm(20, 0.5)
--   student = t.test(sample1, sample2, var.equal=T)
--   welch = t.test(sample1, sample2)
--   paired = t.test(sample1, sample2, paired=T)
sample1, sample2 :: U.Vector Double
sample1 = U.fromList [
  1.262954284880793e+00,
 -3.262333607056494e-01,
  1.329799262922501e+00,
  1.272429321429405e+00,
  4.146414344564082e-01,
 -1.539950041903710e+00,
 -9.285670347135381e-01,
 -2.947204467905602e-01,
 -5.767172747536955e-03,
  2.404653388857951e+00,
  7.635934611404596e-01,
 -7.990092489893682e-01,
 -1.147657009236351e+00,
 -2.894615736882233e-01,
 -2.992151178973161e-01,
 -4.115108327950670e-01,
  2.522234481561323e-01,
 -8.919211272845686e-01,
  4.356832993557186e-01,
 -1.237538421929958e+00]
sample2 = U.fromList [
  2.757321147216907e-01,
  8.773956459817011e-01,
  6.333363608148415e-01,
  1.304189509744908e+00,
  4.428932256161913e-01,
  1.003607972233726e+00,
  1.585769362145687e+00,
 -1.909538396968303e-01,
 -7.845993538721883e-01,
  5.467261721883520e-01,
  2.642934435604988e-01,
 -4.288825501025439e-02,
  6.668968254321778e-02,
 -1.494716467962331e-01,
  1.226750747385451e+00,
  1.651911754087200e+00,
  1.492160365445798e+00,
  7.048689050811874e-02,
  1.738304100853380e+00,
  2.206537181457307e-01]

testTTest :: String
          -> PValue Double
          -> Test d
          -> [Tst.TestTree]
testTTest name pVal test =
  [ testEquality name (isSignificant pVal test) NotSignificant
  , testEquality name (isSignificant (mkPValue $ pValue pVal + 1e-5) test)
    Significant
  ]

studentTTests :: [Tst.TestTree]
studentTTests = concat
  [ -- R: t.test(sample1, sample2, alt="two.sided", var.equal=T)
    testTTest "two-sample t-test SamplesDiffer Student"
      (mkPValue 0.03410) (fromJust $ studentTTest SamplesDiffer sample1 sample2)
    -- R: t.test(sample1, sample2, alt="two.sided", var.equal=F)
  , testTTest "two-sample t-test SamplesDiffer Welch"
      (mkPValue 0.03483) (fromJust $ welchTTest SamplesDiffer sample1 sample2)
    -- R: t.test(sample1, sample2, alt="two.sided", paired=T)
  , testTTest "two-sample t-test SamplesDiffer Paired"
      (mkPValue 0.03411) (fromJust $ pairedTTest SamplesDiffer sample12)
    -- R: t.test(sample1, sample2, alt="less", var.equal=T)
  , testTTest "two-sample t-test BGreater Student"
      (mkPValue 0.01705) (fromJust $ studentTTest BGreater sample1 sample2)
    -- R: t.test(sample1, sample2, alt="less", var.equal=F)
  , testTTest "two-sample t-test BGreater Welch"
      (mkPValue 0.01741) (fromJust $ welchTTest BGreater sample1 sample2)
    -- R: t.test(sample1, sample2, alt="less", paired=F)
  , testTTest "two-sample t-test BGreater Paired"
      (mkPValue 0.01705) (fromJust $ pairedTTest BGreater sample12)
  ]
  where sample12 = U.zip sample1 sample2


module FinancialTimeseries.Algorithm.MovingAverage where

import qualified Data.Vector as Vec

import qualified Statistics.Sample as Sample

import FinancialTimeseries.Type.Labeled (Labeled(..))
import FinancialTimeseries.Type.Segment (Segment(..), HalfSegment(..))
import FinancialTimeseries.Type.Strategy (Strategy(..))
import FinancialTimeseries.Type.Timeseries (TimeseriesRaw(..), Timeseries(..))
import FinancialTimeseries.Type.Types (StripPrice, stripPrice)

import FinancialTimeseries.Util.Pretty (Pretty, pretty)
import FinancialTimeseries.Util.ToFileString (ToFileString, toFileString)

newtype Window = Window {
  unWindow :: Int
  } deriving (Show)

instance ToFileString Window where
  toFileString (Window w) = "Window-" ++ show w

instance Pretty Window where
  pretty = show

data UpDown =
  Up Int
  | Down Int
  deriving (Show)

movingAverage ::
  (StripPrice price, Real a, Fractional a) =>
  Window -> TimeseriesRaw price a -> Timeseries price a
movingAverage (Window m) ts@(TimeseriesRaw _ pvs) =
  let vs = stripPrice pvs
      idx = Vec.fromList [0 .. Vec.length vs - m]

      slc i =
        let ws = Vec.slice i m vs
            (tn, xn) = Vec.last ws
        in (tn, realToFrac (Sample.mean (Vec.map (realToFrac . snd) ws)), xn)
      us = Vec.map slc idx

      f i ((_, m0, x0), (_, m1, x1)) acc
        | m0 <= x0 && m1 > x1 = Up (i+m) : acc
        | m0 > x0 && m1 <= x1 = Down (i+m) : acc
        | otherwise = acc

      xs = Vec.ifoldr' f [] (Vec.zip us (Vec.tail us))

      g [] = (Nothing, [])
      g [Down i] = (Just (HalfSegment i), [])
      g (Down i:Up j:zs) = fmap (Segment i j :) (g zs)
      g (Up _:zs) = g zs
      g (Down _:zs) = g zs

      h (t, s, _) = (t, s)

      (hs, isegs) = g xs

      res = Timeseries {
        timeseriesRaw = ts
        , investedSegments = isegs
        , lastSegment = hs
        , additionalSeries = [Labeled "Moving Average" (Vec.map h us)]
        }
  in res

movingAverageStrategy ::
  (StripPrice price, Real a, Fractional a) =>
  Window -> Strategy Window price a
movingAverageStrategy w = Strategy "MovingAverage" w movingAverage


{-# LANGUAGE FlexibleContexts #-}
module PlotGreensFunctionR2Z2T0S0 where

import           Control.Monad.Parallel  as MP
import           Data.Array.Repa         as R
import           Data.Binary             (decodeFile)
import           Data.Complex
import           Data.List               as L
import           FokkerPlanck.MonteCarlo
import           FokkerPlanck.Pinwheel
import           Image.IO
import           System.Directory
import           System.Environment
import           System.FilePath
import           Types
import           Utils.Array
import Data.Vector.Unboxed as VU

main = do
  args@(numPointStr:numOrientationStr:thetaSigmaStr:scaleSigmaStr:maxScaleStr:taoStr:lenStr:initStr:numTrailStr:maxTrailStr:theta0FreqsStr:thetaFreqsStr:scale0FreqsStr:scaleFreqsStr:histFileName:alphaStr:numThreadStr:_) <-
    getArgs
  print args
  let numPoint = read numPointStr :: Int
      numOrientation = read numOrientationStr :: Int
      thetaSigma = read thetaSigmaStr :: Double
      scaleSigma = read scaleSigmaStr :: Double
      maxScale = read maxScaleStr :: Double
      tao = read taoStr :: Double
      len = read lenStr :: Int
      init = read initStr :: (Double, Double, Double, Double, Double, Double)
      numTrail = read numTrailStr :: Int
      maxTrail = read maxTrailStr :: Int
      theta0Freq = read theta0FreqsStr :: Double
      thetaFreq = read thetaFreqsStr :: Double
      scale0Freq = read scale0FreqsStr :: Double
      scaleFreq = read scaleFreqsStr :: Double
      theta0Freqs = [-theta0Freq .. -theta0Freq]
      thetaFreqs = [thetaFreq .. thetaFreq]
      scale0Freqs = [(scale0Freq - 1) .. scale0Freq]
      scaleFreqs = [-scaleFreq .. scaleFreq]
      alpha = read alphaStr :: Double
      numThread = read numThreadStr :: Int
      folderPath = "output/app/PlotGreensFunctionR2Z2T0S0"
      histPath = folderPath </> histFileName
      radialHistPath = folderPath </> ("Radial_" L.++ histFileName)
  createDirectoryIfMissing True (folderPath </> "GreensR2Z2T0S0")
  flag <- doesFileExist histPath
  hist <-
    if flag
      then decodeFile histPath
      else return $
           emptyHistogram
             [ numPoint
             , numPoint
             , L.length scale0Freqs
             , L.length theta0Freqs
             , L.length scaleFreqs
             , L.length thetaFreqs
             ]
             0
  arr <-
    solveMonteCarloR2Z2T0S0
      numThread
      numTrail
      maxTrail
      numPoint
      numPoint
      thetaSigma
      scaleSigma
      maxScale
      tao
      theta0Freqs
      thetaFreqs
      scale0Freqs
      scaleFreqs
      histPath
      hist
  let arr4d =
        R.slice
          arr
          (Z :. (L.length theta0Freqs - 1) :. (L.length scale0Freqs - 1) :. All :.
           All :.
           All :.
           All)
      arr4dMax = VU.maximum . VU.map magnitude . toUnboxed . computeS $ arr4d
      -- arr4dNormalized = R.map (\x -> x / (arr4dMax :+ 0) * (255 :+ 0))
  -- MP.mapM_
  --   (\(i, j) ->
  --      plotImageRepaComplex
  --        (folderPath </> "GreensR2Z2T0S0" </> show (i + 1) L.++ "_" L.++
  --         show (j + 1) L.++
  --         ".png") .
  --      ImageRepa 8 .
  --      computeS . R.extend (Z :. (1 :: Int) :. All :. All) . R.slice arr4dNormalized $
  --      (Z :. i :. j :. All :. All))
  --   [ (i, j)
  --   | i <- [0 .. (L.length thetaFreqs) - 1]
  --   , j <- [0 .. (L.length scaleFreqs) - 1]
  --   ]
  flag <- doesFileExist radialHistPath
  radialArr <-
    if flag
      then R.map magnitude . getNormalizedHistogramArr <$>
           decodeFile radialHistPath
      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
          radialHistPath
          (emptyHistogram
             [ if maxScale >= 2
                 then round maxScale + 1
                 else (round . sqrt . fromIntegral $ 2 * (div numPoint 2) ^ 2)
             , L.length scale0Freqs
             , L.length theta0Freqs
             , L.length scaleFreqs
             , L.length thetaFreqs
             ]
             0)
  arr' <-
    computeUnboxedP $
    computeR2Z2T0S0ArrayRadial
      (PinwheelHollow0 1)
      radialArr
      numPoint
      numPoint
      1
      maxScale
      thetaFreqs
      scaleFreqs
      theta0Freqs
      scale0Freqs
  let arr4d' =
        R.slice
          arr'
          (Z :. All :. All :. (L.length thetaFreqs - 1) :.
           (L.length scaleFreqs - 1) :.
           All :.
           All)
  MP.mapM_
    (\(i, j) ->
       plotImageRepaComplex
         (folderPath </> "GreensR2Z2T0S0" </> "Radial_" L.++ show (i + 1) L.++
          "_" L.++
          show (j + 1) L.++
          ".png") .
       ImageRepa 8 .
       computeS . R.extend (Z :. (1 :: Int) :. All :. All) . R.slice arr4d' $
       (Z :. i :. j :. All :. All))
    [ (i, j)
    | i <- [0 .. (L.length thetaFreqs) - 1]
    , j <- [0 .. (L.length scaleFreqs) - 1]
    ]


{-# OPTIONS_GHC -Wall #-}
{-# LANGUAGE CPP #-}

{- |
Module      :  Physics.Learn.BlochSphere
Copyright   :  (c) Scott N. Walck 2016
License     :  BSD3 (see LICENSE)
Maintainer  :  Scott N. Walck <walck@lvc.edu>
Stability   :  experimental

This module contains functions for displaying the
state of a spin-1/2 particle or other quantum two-level
system as a point on the Bloch Sphere.
-}

module Physics.Learn.BlochSphere
    ( VisObj
    , toPos
    , ketToPos
    , staticBlochSphere
    , displayStaticState
    , animatedBlochSphere
    , simulateBlochSphere
    , simulateBlochSphereK
    , stateProp
    , statePropK
    , evolutionBlochSphere
    , evolutionBlochSphereK
    , hamRabi
    )
    where

import qualified Physics.Learn.QuantumMat as M
import qualified Physics.Learn.Ket as K
import Physics.Learn.Ket
    ( Ket
    , Operator
    , (<>)
    , dagger
    )
import Numeric.LinearAlgebra
    ( Vector
    , Matrix
    , C
    , iC
--    , (<>)  -- matrix multiplication
--    , (|>)  -- vector definition
    , (!)   -- vector element access
    , (><)  -- matrix definition
    , scale
    , size
    )
import Data.Complex
    ( Complex(..)
    , conjugate
    , realPart
    , imagPart
    )
import Physics.Learn
    ( Position
    , v3FromPos
    , cart
    )
import SpatialMath
    ( Euler(..)
    )
import Vis
    ( VisObject(..)
    , Flavour(..)
    , Options(..)
    , Camera0(..)
    , defaultOpts
    , display
    , simulate
    , blue
    , red
    )
#if MIN_VERSION_base(4,11,0)
import Prelude hiding ((<>))
#endif

{-
3 ways to specify the state of a spin-1/2 particle:
Vector C
Ket
Position  (Bloch vector)

2 ways to specify a Hamiltonian:
Matrix C
Operator

3 choices for Vis' world:
(Float, Vector C)
(Float, Ket)
(Float, Position)
-}

-- | A Vis object.
type VisObj = VisObject Double

-- | Convert a 2x1 complex state vector for a qubit
--   into Bloch (x,y,z) coordinates.
toPos :: Vector C -> Position
toPos v
    = if size v /= 2
      then error "toPos only for size 2 vectors"
      else let z1 = v ! 0
               z2 = v ! 1
           in cart (2 * realPart (conjugate z1 * z2))
                   (2 * imagPart (conjugate z1 * z2))
                   (realPart (conjugate z1 * z1 - conjugate z2 * z2))

-- | Convert a qubit ket
--   into Bloch (x,y,z) coordinates.
ketToPos :: Ket -> Position
ketToPos psi
    = if K.dim psi /= 2
      then error "ketToPos only for qubit kets"
      else let z1 = dagger K.zp <> psi
               z2 = dagger K.zm <> psi
           in cart (2 * realPart (conjugate z1 * z2))
                   (2 * imagPart (conjugate z1 * z2))
                   (realPart (conjugate z1 * z1 - conjugate z2 * z2))

-- | A static 'VisObj' for the state of a qubit.
staticBlochSphere :: Position -> VisObj
staticBlochSphere r
    = RotEulerDeg (Euler 270 0 0) $ RotEulerDeg (Euler 0 180 0) $
      VisObjects [ Sphere 1 Wireframe blue
                 , Trans (v3FromPos r) (Sphere 0.05 Solid red)
                 ]

displayStaticBlochSphere :: Position -> IO ()
displayStaticBlochSphere r
    = display myOptions (staticBlochSphere r)

-- | Display a qubit state vector as a point on the Bloch Sphere.
displayStaticState :: Vector C -> IO ()
displayStaticState = displayStaticBlochSphere . toPos

-- | Given a Bloch vector as a function of time,
--   return a 'VisObj' as a function of time.
animatedBlochSphere :: (Double -> Position) -> (Float -> VisObj)
animatedBlochSphere f
    = staticBlochSphere . f . realToFrac

-- | Given a sample rate, initial qubit state vector, and
--   state propagation function, produce a simulation.
--   The 'Float' in the state propagation function is the time
--   since the beginning of the simulation.
simulateBlochSphere :: Double -> Vector C -> (Float -> (Float,Vector C) -> (Float,Vector C)) -> IO ()
simulateBlochSphere sampleRate initial statePropFunc
    = simulate myOptions sampleRate (0,initial) (staticBlochSphere . toPos . snd) statePropFunc

-- | Given a sample rate, initial qubit state ket, and
--   state propagation function, produce a simulation.
--   The 'Float' in the state propagation function is the time
--   since the beginning of the simulation.
simulateBlochSphereK :: Double -> Ket -> (Float -> (Float,Ket) -> (Float,Ket)) -> IO ()
simulateBlochSphereK sampleRate initial statePropFuncK
    = simulate myOptions sampleRate (0,initial) (staticBlochSphere . ketToPos . snd) statePropFuncK

{-
-- | Given a sample rate, initial qubit state vector, and
--   state propagation function, produce a simulation.
--   The 'Float' in the state propagation function is the time
--   since the beginning of the simulation.
playBlochSphere :: Double -> Vector C -> (Float -> (Float,Vector C) -> (Float,Vector C)) -> IO ()
playBlochSphere sampleRate initial statePropFunc
    = play myOptions sampleRate (0,initial) (staticBlochSphere . toPos . snd) statePropFunc
-}

-- | Produce a state propagation function from a time-dependent Hamiltonian.
stateProp :: (Double -> Matrix C) -> Float -> (Float,Vector C) -> (Float,Vector C)
stateProp ham tNew (tOld,v)
    = (tNew, M.timeEv (realToFrac dt) (ham tMid) v)
      where
        dt = tNew - tOld
        tMid = realToFrac $ (tNew + tOld) / 2

-- | Produce a state propagation function from a time-dependent Hamiltonian.
statePropK :: (Double -> Operator) -> Float -> (Float,Ket) -> (Float,Ket)
statePropK ham tNew (tOld,psi)
    = (tNew, K.timeEv (realToFrac dt) (ham tMid) psi)
      where
        dt = tNew - tOld
        tMid = realToFrac $ (tNew + tOld) / 2

-- | Given an initial qubit state and a time-dependent Hamiltonian,
--   produce a visualization.
evolutionBlochSphere :: Vector C -> (Double -> Matrix C) -> IO ()
evolutionBlochSphere psi0 ham
    = simulateBlochSphere 0.01 psi0 (stateProp ham)

-- | Given an initial qubit ket and a time-dependent Hamiltonian,
--   produce a visualization.
evolutionBlochSphereK :: Ket -> (Double -> Operator) -> IO ()
evolutionBlochSphereK psi0 ham
    = simulateBlochSphereK 0.01 psi0 (statePropK ham)

myOptions :: Options
myOptions = defaultOpts {optWindowName = "Bloch Sphere"
                        ,optInitialCamera = Just (Camera0 75 20 4)}

{-
staticBz1 :: IO ()
staticBz1 = evolutionBlochSphere M.xp (const (scale 0.9 M.sz))

staticBz2 :: IO ()
staticBz2 = evolutionBlochSphere ((2|>) [(cos (pi / 8)), (sin (pi / 8))]) (const M.sz)

staticBy1 :: IO ()
staticBy1 = evolutionBlochSphere M.xp (const M.sy)
-}

-- | Hamiltonian for nuclear magnetic resonance.
--   Explain omega0, omegaR, omega.
hamRabi :: Double ->  Double ->  Double ->  Double -> Matrix C
hamRabi omega0 omegaR omega t
    = let h11 = omega0 :+ 0
          h12 = (omegaR :+ 0) * exp (-iC * ((omega * t) :+ 0))
      in scale (1/2) $ (2><2) [h11, h12, (conjugate h12), (-h11)]

-- need to scale time

-- a pi pulse


{-# OPTIONS_GHC -Wall #-}
{-# LANGUAGE Safe #-}
{-# LANGUAGE MultiParamTypeClasses, FunctionalDependencies #-}
{-# LANGUAGE TypeSynonymInstances, FlexibleInstances #-}
{-# LANGUAGE CPP #-}

{- |
Module      :  Physics.Learn.Ket
Copyright   :  (c) Scott N. Walck 2016-2018
License     :  BSD3 (see LICENSE)
Maintainer  :  Scott N. Walck <walck@lvc.edu>
Stability   :  experimental

This module contains ket vectors, bra vectors,
and operators for quantum mechanics.
-}

-- a Ket layer on top of QuantumMat

module Physics.Learn.Ket
    (
    -- * Basic data types
      C
    , i
    , magnitude
    , Ket
    , Bra
    , Operator
    -- * Kets for spin-1/2 particles
    , xp
    , xm
    , yp
    , ym
    , zp
    , zm
    , np
    , nm
    -- * Operators for spin-1/2 particles
    , sx
    , sy
    , sz
    , sn
    , sn'
    -- * Quantum Dynamics
    , timeEvOp
    , timeEv
    -- * Composition
    , Kron(..)
    -- * Measurement
    , possibleOutcomes
    , outcomesProjectors
    , outcomesProbabilities
--    , prob
--    , probs
    -- * Generic multiplication
    , Mult(..)
    -- * Adjoint operation
    , Dagger(..)
    -- * Normalization
    , HasNorm(..)
    -- * Representation
    , Representable(..)
    -- * Orthonormal bases
    , OrthonormalBasis
    , makeOB
    , listBasis
    , size
    -- * Orthonormal bases for spin-1/2 particles
    , xBasis
    , yBasis
    , zBasis
    , nBasis
    -- , angularMomentumXMatrix
    -- , angularMomentumYMatrix
    -- , angularMomentumZMatrix
    -- , angularMomentumPlusMatrix
    -- , angularMomentumMinusMatrix
    -- , jXMatrix
    -- , jYMatrix
    -- , jZMatrix
    -- , matrixCommutator
    -- , rotationMatrix
    -- , jmColumn
    )
    where

-- We try to import only from QuantumMat
-- and not from Numeric.LinearAlgebra

import qualified Data.Complex as C
import Data.Complex
    ( Complex(..)
    , conjugate
    )
import qualified Physics.Learn.QuantumMat as M
import Physics.Learn.QuantumMat
    ( C
    , Vector
    , Matrix
    , (#>)
    , (<#)
    , conjugateTranspose
    , scaleV
    , scaleM
    , conjV
    , fromList
    , toList
    , fromLists
    )
#if MIN_VERSION_base(4,11,0)
import Prelude hiding ((<>))
#endif

infixl 7 <>

-- | A ket vector describes the state of a quantum system.
data Ket = Ket (Vector C)

instance Show Ket where
    show k =
        let message = "Use 'rep <basis name> <ket name>'."
        in if dim k == 2
           then "Representation in zBasis:\n" ++
                show (rep zBasis k) ++ "\n" ++ message
           else message

-- | An operator describes an observable (a Hermitian operator)
--   or an action (a unitary operator).
data Operator = Operator (Matrix C)

instance Show Operator where
    show op =
        let message = "Use 'rep <basis name> <operator name>'."
        in if dim op == 2
           then "Representation in zBasis:\n" ++
                show (rep zBasis op) ++ "\n" ++ message
           else message

-- | A bra vector describes the state of a quantum system.
data Bra = Bra (Vector C)

instance Show Bra where
    show _ = "<bra>\nTry 'rep zBasis <bra name>'"

magnitude :: C -> Double
magnitude = C.magnitude

i :: C
i = 0 :+ 1

-- | Generic multiplication including inner product,
--   outer product, operator product, and whatever else makes sense.
--   No conjugation takes place in this operation.
class Mult a b c | a b -> c where
    (<>) :: a -> b -> c

instance Mult C C C where
    z1 <> z2 = z1 * z2

instance Mult C Ket Ket where
    c <> Ket matrixKet = Ket (scaleV c matrixKet)

instance Mult C Bra Bra where
    c <> Bra matrixBra = Bra (scaleV c matrixBra)

instance Mult C Operator Operator where
    c <> Operator m = Operator (scaleM c m)

instance Mult Ket C Ket where
    Ket matrixKet <> c = Ket (scaleV c matrixKet)

instance Mult Bra C Bra where
    Bra matrixBra <> c = Bra (scaleV c matrixBra)

instance Mult Operator C Operator where
    Operator m <> c = Operator (scaleM c m)

instance Mult Bra Ket C where
    Bra matrixBra <> Ket matrixKet
        = sum $ zipWith (*) (toList matrixBra) (toList matrixKet)

instance Mult Bra Operator Bra where
    Bra matrixBra <> Operator matrixOp
        = Bra (matrixBra <# matrixOp)

instance Mult Operator Ket Ket where
    Operator matrixOp <> Ket matrixKet
        = Ket (matrixOp #> matrixKet)

instance Mult Ket Bra Operator where
    Ket k <> Bra b =
        Operator
        (fromLists [[ x*y | y <- toList b] | x <- toList k])

instance Mult Operator Operator Operator where
    Operator m1 <> Operator m2 = Operator (m1 M.<> m2)

instance Num Ket where
    Ket v1 + Ket v2 = Ket (v1 + v2)
    Ket v1 - Ket v2 = Ket (v1 - v2)
    (*)         = error "Multiplication is not defined for kets"
    negate (Ket v) = Ket (negate v)
    abs         = error "abs is not defined for kets"
    signum      = error "signum is not defined for kets"
    fromInteger = error "fromInteger is not defined for kets"

instance Num Bra where
    Bra v1 + Bra v2 = Bra (v1 + v2)
    Bra v1 - Bra v2 = Bra (v1 - v2)
    (*)         = error "Multiplication is not defined for bra vectors"
    negate (Bra v) = Bra (negate v)
    abs         = error "abs is not defined for bra vectors"
    signum      = error "signum is not defined for bra vectors"
    fromInteger = error "fromInteger is not defined for bra vectors"

instance Num Operator where
    Operator v1 + Operator v2 = Operator (v1 + v2)
    Operator v1 - Operator v2 = Operator (v1 - v2)
    Operator v1 * Operator v2 = Operator (v1 M.<> v2)
    negate (Operator v) = Operator (negate v)
    abs         = error "abs is not defined for operators"
    signum      = error "signum is not defined for operators"
    fromInteger = error "fromInteger is not defined for operators"

-- | The adjoint operation on complex numbers, kets,
--   bras, and operators.
class Dagger a b | a -> b where
    dagger :: a -> b

instance Dagger Ket Bra where
    dagger (Ket v) = Bra (conjV v)

instance Dagger Bra Ket where
    dagger (Bra v) = Ket (conjV v)

instance Dagger Operator Operator where
    dagger (Operator m) = Operator (conjugateTranspose m)

instance Dagger C C where
    dagger c = conjugate c

class HasNorm a where
    norm      :: a -> Double
    normalize :: a -> a

instance HasNorm Ket where
    norm (Ket v) = M.norm v
    normalize k  = (1 / norm k :+ 0) <> k

instance HasNorm Bra where
    norm (Bra v) = M.norm v
    normalize b  = (1 / norm b :+ 0) <> b

-- | An orthonormal basis of kets.
newtype OrthonormalBasis = OB [Ket]
    deriving (Show)

-- | Make an orthonormal basis from a list of linearly independent kets.
makeOB :: [Ket] -> OrthonormalBasis
makeOB = OB . gramSchmidt

size :: OrthonormalBasis -> Int
size (OB ks) = length ks

listBasis :: OrthonormalBasis -> [Ket]
listBasis (OB ks) = ks

{-
newOrthonormalBasis :: Int -> OrthonormalBasis
newOrthonormalBasis = undefined
-}

class Representable a b | a -> b where
    rep :: OrthonormalBasis -> a -> b
    dim :: a -> Int

instance Representable Ket (Vector C) where
    rep (OB ks) k = fromList (map (\bk -> dagger bk <> k) ks)
    dim (Ket v) = M.dim v

instance Representable Bra (Vector C) where
    rep (OB ks) b = fromList (map (\bk -> b <> bk) ks)
    dim (Bra v) = M.dim v

instance Representable Operator (Matrix C) where
    rep (OB ks) op = fromLists [[ dagger k1 <> op <> k2 | k2 <- ks ] | k1 <- ks ]
    dim (Operator m) = let (p,q) = M.size m
                       in if p == q then p else error "dim: non-square operator"

--------------
-- Spin 1/2 --
--------------

-- | State of a spin-1/2 particle if measurement
--   in the x-direction would give angular momentum +hbar/2.
xp :: Ket
xp = Ket M.xp

-- | State of a spin-1/2 particle if measurement
--   in the x-direction would give angular momentum -hbar/2.
xm :: Ket
xm = Ket M.xm

-- | State of a spin-1/2 particle if measurement
--   in the y-direction would give angular momentum +hbar/2.
yp :: Ket
yp = Ket M.yp

-- | State of a spin-1/2 particle if measurement
--   in the y-direction would give angular momentum -hbar/2.
ym :: Ket
ym = Ket M.ym

-- | State of a spin-1/2 particle if measurement
--   in the z-direction would give angular momentum +hbar/2.
zp :: Ket
zp = Ket M.zp

-- | State of a spin-1/2 particle if measurement
--   in the z-direction would give angular momentum -hbar/2.
zm :: Ket
zm = Ket M.zm

-- | State of a spin-1/2 particle if measurement
--   in the n-direction, described by spherical polar angle theta
--   and azimuthal angle phi, would give angular momentum +hbar/2.
np :: Double -> Double -> Ket
np theta phi
    = (cos (theta / 2) :+ 0) <> zp
      + (sin (theta / 2) :+ 0) * (cos phi :+ sin phi) <> zm

-- | State of a spin-1/2 particle if measurement
--   in the n-direction, described by spherical polar angle theta
--   and azimuthal angle phi, would give angular momentum -hbar/2.
nm :: Double -> Double -> Ket
nm theta phi
    = (sin (theta / 2) :+ 0) <> zp
      - (cos (theta / 2) :+ 0) * (cos phi :+ sin phi) <> zm

-- | The orthonormal basis composed of 'xp' and 'xm'.
xBasis :: OrthonormalBasis
xBasis = makeOB [xp,xm]

-- | The orthonormal basis composed of 'yp' and 'ym'.
yBasis :: OrthonormalBasis
yBasis = makeOB [yp,ym]

-- | The orthonormal basis composed of 'zp' and 'zm'.
zBasis :: OrthonormalBasis
zBasis = makeOB [zp,zm]

-- | Given spherical polar angle theta and azimuthal angle phi,
--   the orthonormal basis composed of 'np' theta phi and 'nm' theta phi.
nBasis :: Double -> Double -> OrthonormalBasis
nBasis theta phi = makeOB [np theta phi,nm theta phi]

-- | The Pauli X operator.
sx :: Operator
sx = xp <> dagger xp - xm <> dagger xm

-- | The Pauli Y operator.
sy :: Operator
sy = yp <> dagger yp - ym <> dagger ym

-- | The Pauli Z operator.
sz :: Operator
sz = zp <> dagger zp - zm <> dagger zm

-- | Pauli operator for an arbitrary direction given
--   by spherical coordinates theta and phi.
sn :: Double -> Double -> Operator
sn theta phi
    = (sin theta * cos phi :+ 0) <> sx +
      (sin theta * sin phi :+ 0) <> sy +
      (cos theta           :+ 0) <> sz

-- | Alternative definition
--   of Pauli operator for an arbitrary direction.
sn' :: Double -> Double -> Operator
sn' theta phi
    = np theta phi <> dagger (np theta phi) -
      nm theta phi <> dagger (nm theta phi)

----------------------
-- Quantum Dynamics --
----------------------

-- | Given a time step and a Hamiltonian operator,
--   produce a unitary time evolution operator.
--   Unless you really need the time evolution operator,
--   it is better to use 'timeEv', which gives the
--   same numerical results without doing an explicit
--   matrix inversion.  The function assumes hbar = 1.
timeEvOp :: Double -> Operator -> Operator
timeEvOp dt (Operator m) = Operator (M.timeEvMat dt m)

-- | Given a time step and a Hamiltonian operator,
--   advance the state ket using the Schrodinger equation.
--   This method should be faster than using 'timeEvOp'
--   since it solves a linear system rather than calculating
--   an inverse matrix.  The function assumes hbar = 1.
timeEv :: Double -> Operator -> Ket -> Ket
timeEv dt (Operator m) (Ket k) = Ket $ M.timeEv dt m k

-----------------
-- Composition --
-----------------

class Kron a where
    kron :: a -> a -> a

instance Kron Ket where
    kron (Ket v1) (Ket v2) = Ket (M.kron v1 v2)

instance Kron Bra where
    kron (Bra v1) (Bra v2) = Bra (M.kron v1 v2)

instance Kron Operator where
    kron (Operator m1) (Operator m2) = Operator (M.kron m1 m2)

-----------------
-- Measurement --
-----------------

-- | The possible outcomes of a measurement
--   of an observable.
--   These are the eigenvalues of the operator
--   of the observable.
possibleOutcomes :: Operator -> [Double]
possibleOutcomes (Operator observable) = M.possibleOutcomes observable

-- | Given an obervable, return a list of pairs
--   of possible outcomes and projectors
--   for each outcome.
outcomesProjectors :: Operator -> [(Double,Operator)]
outcomesProjectors (Operator m)
    = [(val,Operator p) | (val,p) <- M.outcomesProjectors m]

-- | Given an observable and a state ket, return a list of pairs
--   of possible outcomes and probabilites
--   for each outcome.
outcomesProbabilities :: Operator -> Ket -> [(Double,Double)]
outcomesProbabilities (Operator m) (Ket v)
    = M.outcomesProbabilities m v

{-
prob :: Ket -> Ket -> Double
prob k1 k2 = magnitude c ** 2
    where
      c = dagger k1 <> k2

probs :: OrthonormalBasis -> Ket -> [Double]
probs (OB ks) k = map (\bk -> let c = dagger bk <> k in magnitude c ** 2) ks
-}

{-
----------------------------------------
-- Angular Momentum of arbitrary size --
----------------------------------------

angularMomentumZMatrix :: Rational -> Matrix Cyclotomic
angularMomentumZMatrix j
    = case twoJPlusOne j of
        Nothing -> error "j must be a nonnegative integer or half-integer"
        Just d  -> matrix d d (\(r,c) -> if r == c then fromRational (j + 1 - fromIntegral r) else 0)

twoJPlusOne :: Rational -> Maybe Int
twoJPlusOne j
    | j >= 0 && (denominator j == 1 || denominator j == 2)  = Just $ fromIntegral $ numerator (2 * j + 1)
    | otherwise                                             = Nothing

angularMomentumPlusMatrix :: Rational -> Matrix Cyclotomic
angularMomentumPlusMatrix j
    = case twoJPlusOne j of
        Nothing -> error "j must be a nonnegative integer or half-integer"
        Just d  -> matrix d d (\(r,c) -> let mr = j + 1 - fromIntegral r
                                             mc = j + 1 - fromIntegral c
                                         in if mr == mc + 1
                                            then sqrtRat (j*(j+1) - mc*mr)
                                            else 0
                              )

angularMomentumMinusMatrix :: Rational -> Matrix Cyclotomic
angularMomentumMinusMatrix j
    = case twoJPlusOne j of
        Nothing -> error "j must be a nonnegative integer or half-integer"
        Just d  -> matrix d d (\(r,c) -> let mr = j + 1 - fromIntegral r
                                             mc = j + 1 - fromIntegral c
                                         in if mr + 1 == mc
                                            then sqrtRat (j*(j+1) - mc*mr)
                                            else 0
                              )

angularMomentumXMatrix :: Rational -> Matrix Cyclotomic
angularMomentumXMatrix j
    = scaleMatrix (1/2) (angularMomentumPlusMatrix j + angularMomentumMinusMatrix j)

angularMomentumYMatrix :: Rational -> Matrix Cyclotomic
angularMomentumYMatrix j
    = scaleMatrix (1/(2*i)) (angularMomentumPlusMatrix j - angularMomentumMinusMatrix j)

jXMatrix :: Rational -> Matrix Cyclotomic
jXMatrix = angularMomentumXMatrix

jYMatrix :: Rational -> Matrix Cyclotomic
jYMatrix = angularMomentumYMatrix

jZMatrix :: Rational -> Matrix Cyclotomic
jZMatrix = angularMomentumZMatrix

matrixCommutator :: Matrix Cyclotomic -> Matrix Cyclotomic -> Matrix Cyclotomic
matrixCommutator m1 m2 = m1 * m2 - m2 * m1

-----------------------
-- Rotation matrices --
-----------------------

r2i :: Rational -> Integer
r2i r
    | denominator r == 1  = numerator r
    | otherwise           = error "r2i:  not an integer"

-- from Sakurai, revised, (3.8.33)
-- beta in degrees
smallDRotElement :: Rational -> Rational -> Rational -> Rational -> Cyclotomic
smallDRotElement j m' m beta
    = sum [parity(k-m+m') * sqrtRat (fact(j+m) * fact(j-m) * fact(j+m') * fact(j-m'))
                   / fromRational (fact(j+m-k) * fact(k) * fact(j-k-m') * fact(k-m+m'))
                         * cosDeg (beta/2) ^ r2i(2*j-2*k+m-m')
                         * sinDeg (beta/2) ^ r2i(2*k-m+m') | k <- [max 0 (m-m') .. min (j+m) (j-m')]]

parity :: Rational -> Cyclotomic
parity = fromIntegral . parityInteger . r2i

-- | (-1)^n, where n might be negative
parityInteger :: Integer -> Integer
parityInteger n
    | odd n      = -1
    | otherwise  =  1

factInteger :: Integer -> Integer
factInteger n
    | n == 0     = 1
    | n >  0     = n * factInteger (n-1)
    | otherwise  = error "factInteger called on negative argument"

fact :: Rational -> Rational
fact = fromIntegral . factInteger . r2i

-- | Rotation matrix elements.
--   From Sakurai, Revised Edition, (3.5.50).
--   The matrix desribes a rotation by gamma about the z axis,
--   followed by a rotation by beta about the y axis,
--   followed by a rotation by alpha about the z axis.
bigDRotElement :: Rational  -- ^ j, a nonnegative integer or half-integer
               -> Rational  -- ^ m', an integer or half-integer index for the row
               -> Rational  -- ^ m, an integer or half-integer index for the column
               -> Rational  -- ^ alpha, in degrees
               -> Rational  -- ^ beta, in degrees
               -> Rational  -- ^ gamma, in degrees
               -> Cyclotomic  -- ^ rotation matrix element
bigDRotElement j m' m alpha beta gamma
    = polarRat 1 (-(m' * alpha + m * gamma) / 360) * smallDRotElement j m' m beta

-- | Rotation matrix for a spin-j particle.
--   The matrix desribes a rotation by gamma about the z axis,
--   followed by a rotation by beta about the y axis,
--   followed by a rotation by alpha about the z axis.
rotationMatrix :: Rational  -- ^ j, a nonnegative integer or half-integer
               -> Rational  -- ^ alpha, in degrees
               -> Rational  -- ^ beta, in degrees
               -> Rational  -- ^ gamma, in degrees
               -> Matrix Cyclotomic  -- ^ rotation matrix
rotationMatrix j alpha beta gamma
    = case twoJPlusOne j of
        Nothing -> error "bigDRotMatrix: j must be a nonnegative integer or half-integer"
        Just d  -> matrix d d (\(r,c) -> let m' = j + 1 - fromIntegral r
                                             m  = j + 1 - fromIntegral c
                                         in bigDRotElement j m' m alpha beta gamma
                              )

----------------------------------
-- Angular Momentum eigenstates --
----------------------------------

jmColumn :: Rational -> Rational -> Matrix Cyclotomic
jmColumn j m
    = case twoJPlusOne j of
        Nothing -> error "bigDRotMatrix: j must be a nonnegative integer or half-integer"
        Just d  -> matrix d 1 (\(r,_) -> let m' = j + 1 - fromIntegral r
                                         in if m == m'
                                            then 1
                                            else 0
                              )
-}

------------------
-- Gram-Schmidt --
------------------

-- | Form an orthonormal list of kets from
--   a list of linearly independent kets.
gramSchmidt :: [Ket] -> [Ket]
gramSchmidt [] = []
gramSchmidt [k] = [normalize k]
gramSchmidt (k:ks) = let nks = gramSchmidt ks
                         nk = normalize (foldl (-) k [w <> dagger w <> k | w <- nks])
                     in nk:nks

-- todo:  Clebsch-Gordon coeffs


module Spec_3_4_3 where

import Test.Hspec
import Numeric.LinearAlgebra

import Section_3_4_3 (initNetwork, forward)

spec :: Spec
spec = do
  describe "forward" $ do
    it "(1.0, 0.5) -> (0.326, 0.696)" $ do
        let x = vector [1.0, 0.5]
            answer = vector [0.31682708, 0.69627909]
            prediction = forward initNetwork x
            err = 0.0001
        norm_2 (answer - prediction) < err `shouldBe` True


module RC.Helpers
  ( addBiases
  , flatten'
  , unflatten'
  , randList
  , randMatrix
  , randSparse
  , hsigmoid
  ) where

import           System.Random
import           Data.List ( unfoldr )
import           Numeric.LinearAlgebra as LA

-- | Hard sigmoid
hsigmoid :: (Fractional a, Ord a)
         => (a, a, a)
         -- ^ Vertical scaling, width, offset
         -> a
         -> a
hsigmoid (β, width, offset) x = f
  where
    f | x < offset = 0.0
      | x < width = β * (x - offset)
      | otherwise = β * (width - offset)
{-# SPECIALISE hsigmoid :: (Double, Double, Double) -> Double -> Double #-}
{-# SPECIALISE hsigmoid :: (Float, Float, Float) -> Float -> Float #-}

-- | Prepend a row of ones
--
-- >>> addBiases $ (2><3) [20..26]
-- (3><3)
--  [  1.0,  1.0,  1.0
--  , 20.0, 21.0, 22.0
--  , 23.0, 24.0, 25.0 ]
addBiases :: LA.Matrix Double -> LA.Matrix Double
addBiases m = let no = LA.cols m
                  m' = LA.konst 1.0 (1, no)
              in m' LA.=== m

-- | Random matrix with elements in the range of [minVal; maxVal]
randMatrix
  :: StdGen
     -> (Int, Int)
     -- ^ Number of rows and columns
     -> (Double, Double)
     -- ^ Minimal and maximal values
     -> LA.Matrix Double
randMatrix seed (rows', cols') (minVal, maxVal) = (LA.reshape cols'. LA.vector) xs
  where
    xs = f <$> randList (rows' * cols') seed
    f x = (maxVal - minVal) * x + minVal

-- | Random sparse matrix
--
-- NB: at the moment, the matrix is stored in memory as
-- an ordinary (dense) matrix.
randSparse g (rows', cols') (minVal, maxVal) connectivity =
    LA.reshape cols' $ LA.vector xs
  where
    (g1, g2) = split g
    rlist = randList (rows' * cols')
    xs = zipWith f (rlist g1) (rlist g2)
    f lv rv | lv < connectivity = (maxVal - minVal) * rv + minVal
            | otherwise = 0.0

randList :: (Random a, Floating a) => Int -> StdGen -> [a]
randList n = take n. unfoldr (Just. random)
{-# SPECIALISE randList :: Int -> StdGen -> [Float] #-}
{-# SPECIALISE randList :: Int -> StdGen -> [Double] #-}

-- | Flatten by columns
flatten' :: Matrix Double -> Vector Double
flatten' = flatten. tr

-- | Reshape using columns instead of rows
unflatten' :: Int -> Vector Double -> Matrix Double
unflatten' nodes = tr. reshape nodes


{-# LANGUAGE DataKinds           #-}
{-# LANGUAGE FlexibleContexts    #-}
{-# LANGUAGE RankNTypes          #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeApplications    #-}
{-# LANGUAGE TypeOperators       #-}

module Backprop.Learn.Loss (
  -- * Loss functions
    Loss
  , crossEntropy
  , squaredError, absError, totalSquaredError, squaredErrorV
  -- , totalCov
  -- ** Manipulate loss functions
  , scaleLoss
  , sumLoss
  , sumLossDecay
  , lastLoss
  , zipLoss
  , t2Loss
  , t3Loss
  -- * Regularization
  , Regularizer
  , l2Reg
  , l1Reg
  , noReg
  -- ** Manipulate regularizers
  , addReg
  , scaleReg
  ) where

import           Control.Applicative
import           Data.Finite
import           Data.Type.Tuple hiding                (T2(..), T3(..))
import           GHC.TypeNats
import           Numeric.Backprop
import           Numeric.LinearAlgebra.Static.Backprop
import           Numeric.Opto.Update hiding            ((<.>))
import qualified Data.Type.Tuple                       as T
import qualified Data.Vector.Sized                     as SV
import qualified Prelude.Backprop                      as B

type Loss a = forall s. Reifies s W => a -> BVar s a -> BVar s Double

crossEntropy :: KnownNat n => Loss (R n)
crossEntropy targ res = -(log res <.> constVar targ)

squaredErrorV :: KnownNat n => Loss (R n)
squaredErrorV targ res = e <.> e
  where
    e = res - constVar targ

totalSquaredError
  :: (Backprop (t Double), Num (t Double), Foldable t, Functor t)
    => Loss (t Double)
totalSquaredError targ res = B.sum (e * e)
  where
    e = constVar targ - res

squaredError :: Loss Double
squaredError targ res = (res - constVar targ) ** 2

absError :: Loss Double
absError targ res = abs (res - constVar targ)

-- -- | Sum of covariances between matching components.  Not sure if anyone
-- -- uses this.
-- totalCov :: (Backprop (t Double), Foldable t, Functor t) => Loss (t Double)
-- totalCov targ res = -(xy / fromIntegral n - (x * y) / fromIntegral (n * n))
--   where
--     x = constVar $ sum targ
--     y = B.sum res
--     xy = B.sum (constVar targ * res)
--     n = length targ

-- klDivergence :: Loss Double
-- klDivergence =

sumLoss
    :: (Traversable t, Applicative t, Backprop a)
    => Loss a
    -> Loss (t a)
sumLoss l targ = sum . liftA2 l targ . sequenceVar

zipLoss
    :: (Traversable t, Applicative t, Backprop a)
    => t Double
    -> Loss a
    -> Loss (t a)
zipLoss xs l targ = sum
                  . liftA3 (\α t -> (* constVar α) . l t) xs targ
                  . sequenceVar

sumLossDecay
    :: forall n a. (KnownNat n, Backprop a)
    => Double
    -> Loss a
    -> Loss (SV.Vector n a)
sumLossDecay β = zipLoss $ SV.generate (\i -> β ** (fromIntegral i - n))
  where
    n = fromIntegral $ maxBound @(Finite n)

lastLoss
    :: (KnownNat (n + 1), Backprop a)
    => Loss a
    -> Loss (SV.Vector (n + 1) a)
lastLoss l targ = l (SV.last targ) . viewVar (SV.ix maxBound)

-- | Scale the result of a loss function.
scaleLoss :: Double -> Loss a -> Loss a
scaleLoss β l x = (* constVar β) . l x

-- | Lift and sum a loss function over the components of a 'T.T2'.
t2Loss
    :: (Backprop a, Backprop b)
    => Loss a                   -- ^ loss on first component
    -> Loss b                   -- ^ loss on second component
    -> Loss (T.T2 a b)
t2Loss f g (T.T2 xT yT) (T2B xR yR) = f xT xR + g yT yR

-- | Lift and sum a loss function over the components of a 'T.T3'.
t3Loss
    :: (Backprop a, Backprop b, Backprop c)
    => Loss a                   -- ^ loss on first component
    -> Loss b                   -- ^ loss on second component
    -> Loss c                   -- ^ loss on third component
    -> Loss (T.T3 a b c)
t3Loss f g h (T.T3 xT yT zT) xRyRzR
        = f xT (xRyRzR ^^. t3_1)
        + g yT (xRyRzR ^^. t3_2)
        + h zT (xRyRzR ^^. t3_3)

-- | A regularizer on parameters
type Regularizer p = forall s. Reifies s W => BVar s p -> BVar s Double

-- | L2 regularization
--
-- \[
-- \sum_w \frac{1}{2} w^2
-- \]
l2Reg
    :: (Metric Double p, Backprop p)
    => Double                   -- ^ scaling factor (often 0.5)
    -> Regularizer p
l2Reg λ = liftOp1 . op1 $ \x ->
            ( λ * quadrance x / 2, (.* x) . (* λ))

-- | L1 regularization
--
-- \[
-- \sum_w \lvert w \rvert
-- \]
l1Reg
    :: (Num p, Metric Double p, Backprop p)
    => Double                   -- ^ scaling factor (often 0.5)
    -> Regularizer p
l1Reg λ = liftOp1 . op1 $ \x ->
            ( λ * norm_1 x, (.* signum x) . (* λ)
            )

-- | No regularization
noReg :: Regularizer p
noReg _ = constVar 0

-- | Add together two regularizers
addReg :: Regularizer p -> Regularizer p -> Regularizer p
addReg = liftA2 (+)

-- | Scale a regularizer's influence
scaleReg :: Double -> Regularizer p -> Regularizer p
scaleReg λ reg = (* constVar λ) . reg


-- |
-- Module      : Control.Monad.Bayes.Class
-- Description : Types for probabilistic modelling
-- Copyright   : (c) Adam Scibior, 2015-2020
-- License     : MIT
-- Maintainer  : leonhard.markert@tweag.io
-- Stability   : experimental
-- Portability : GHC
--
-- This module defines 'MonadInfer', which can be used to represent a simple model
-- like the following:
--
-- @
-- import Control.Monad (when)
-- import Control.Monad.Bayes.Class
--
-- model :: MonadInfer m => m Bool
-- model = do
--   rain <- bernoulli 0.3
--   sprinkler <-
--     bernoulli $
--     if rain
--       then 0.1
--       else 0.4
--   let wetProb =
--     case (rain, sprinkler) of
--       (True,  True)  -> 0.98
--       (True,  False) -> 0.80
--       (False, True)  -> 0.90
--       (False, False) -> 0.00
--   score wetProb
--   return rain
-- @
module Control.Monad.Bayes.Class
  ( MonadSample,
    random,
    uniform,
    normal,
    gamma,
    beta,
    bernoulli,
    categorical,
    logCategorical,
    uniformD,
    geometric,
    poisson,
    dirichlet,
    MonadCond,
    score,
    factor,
    condition,
    MonadInfer,
    normalPdf,
  )
where

import Control.Monad (when)
import Control.Monad.Trans.Class
import Control.Monad.Trans.Cont
import Control.Monad.Trans.Identity
import Control.Monad.Trans.List
import Control.Monad.Trans.Maybe
import Control.Monad.Trans.RWS hiding (tell)
import Control.Monad.Trans.Reader
import Control.Monad.Trans.State
import Control.Monad.Trans.Writer
import qualified Data.Vector as V
import Data.Vector.Generic as VG
import Numeric.Log
import Statistics.Distribution
import Statistics.Distribution.Beta (betaDistr)
import Statistics.Distribution.Gamma (gammaDistr)
import Statistics.Distribution.Geometric (geometric0)
import Statistics.Distribution.Normal (normalDistr)
import qualified Statistics.Distribution.Poisson as Poisson
import Statistics.Distribution.Uniform (uniformDistr)

-- | Monads that can draw random variables.
class Monad m => MonadSample m where
  -- | Draw from a uniform distribution.
  random ::
    -- | \(\sim \mathcal{U}(0, 1)\)
    m Double

  -- | Draw from a uniform distribution.
  uniform ::
    -- | lower bound a
    Double ->
    -- | upper bound b
    Double ->
    -- | \(\sim \mathcal{U}(a, b)\).
    m Double
  uniform a b = draw (uniformDistr a b)

  -- | Draw from a normal distribution.
  normal ::
    -- | mean μ
    Double ->
    -- | standard deviation σ
    Double ->
    -- | \(\sim \mathcal{N}(\mu, \sigma^2)\)
    m Double
  normal m s = draw (normalDistr m s)

  -- | Draw from a gamma distribution.
  gamma ::
    -- | shape k
    Double ->
    -- | scale θ
    Double ->
    -- | \(\sim \Gamma(k, \theta)\)
    m Double
  gamma shape scale = draw (gammaDistr shape scale)

  -- | Draw from a beta distribution.
  beta ::
    -- | shape α
    Double ->
    -- | shape β
    Double ->
    -- | \(\sim \mathrm{Beta}(\alpha, \beta)\)
    m Double
  beta a b = draw (betaDistr a b)

  -- | Draw from a Bernoulli distribution.
  bernoulli ::
    -- | probability p
    Double ->
    -- | \(\sim \mathrm{B}(1, p)\)
    m Bool
  bernoulli p = fmap (< p) random

  -- | Draw from a categorical distribution.
  categorical ::
    Vector v Double =>
    -- | event probabilities
    v Double ->
    -- | outcome category
    m Int
  categorical ps = fromPMF (ps !)

  -- | Draw from a categorical distribution in the log domain.
  logCategorical ::
    (Vector v (Log Double), Vector v Double) =>
    -- | event probabilities
    v (Log Double) ->
    -- | outcome category
    m Int
  logCategorical = categorical . VG.map (exp . ln)

  -- | Draw from a discrete uniform distribution.
  uniformD ::
    -- | observable outcomes @xs@
    [a] ->
    -- | \(\sim \mathcal{U}\{\mathrm{xs}\}\)
    m a
  uniformD xs = do
    let n = Prelude.length xs
    i <- categorical $ V.replicate n (1 / fromIntegral n)
    return (xs !! i)

  -- | Draw from a geometric distribution.
  geometric ::
    -- | success rate p
    Double ->
    -- | \(\sim\) number of failed Bernoulli trials with success probability p before first success
    m Int
  geometric = discrete . geometric0

  -- | Draw from a Poisson distribution.
  poisson ::
    -- | parameter λ
    Double ->
    -- | \(\sim \mathrm{Pois}(\lambda)\)
    m Int
  poisson = discrete . Poisson.poisson

  -- | Draw from a Dirichlet distribution.
  dirichlet ::
    Vector v Double =>
    -- | concentration parameters @as@
    v Double ->
    -- | \(\sim \mathrm{Dir}(\mathrm{as})\)
    m (v Double)
  dirichlet as = do
    xs <- VG.mapM (`gamma` 1) as
    let s = VG.sum xs
    let ys = VG.map (/ s) xs
    return ys

-- | Draw from a continuous distribution using the inverse cumulative density
-- function.
draw :: (ContDistr d, MonadSample m) => d -> m Double
draw d = fmap (quantile d) random

-- | Draw from a discrete distribution using a sequence of draws from
-- Bernoulli.
fromPMF :: MonadSample m => (Int -> Double) -> m Int
fromPMF p = f 0 1
  where
    f i r = do
      when (r < 0) $ error "fromPMF: total PMF above 1"
      let q = p i
      when (q < 0 || q > 1) $ error "fromPMF: invalid probability value"
      b <- bernoulli (q / r)
      if b then pure i else f (i + 1) (r - q)

-- | Draw from a discrete distributions using the probability mass function.
discrete :: (DiscreteDistr d, MonadSample m) => d -> m Int
discrete = fromPMF . probability

-- | Monads that can score different execution paths.
class Monad m => MonadCond m where
  -- | Record a likelihood.
  score ::
    -- | likelihood of the execution path
    Log Double ->
    m ()

-- | Synonym for 'score'.
factor ::
  MonadCond m =>
  -- | likelihood of the execution path
  Log Double ->
  m ()
factor = score

-- | Hard conditioning.
condition :: MonadCond m => Bool -> m ()
condition b = score $ if b then 1 else 0

-- | Monads that support both sampling and scoring.
class (MonadSample m, MonadCond m) => MonadInfer m

-- | Probability density function of the normal distribution.
normalPdf ::
  -- | mean μ
  Double ->
  -- | standard deviation σ
  Double ->
  -- | sample x
  Double ->
  -- | relative likelihood of observing sample x in \(\mathcal{N}(\mu, \sigma^2)\)
  Log Double
normalPdf mu sigma x = Exp $ logDensity (normalDistr mu sigma) x

----------------------------------------------------------------------------
-- Instances that lift probabilistic effects to standard tranformers.

instance MonadSample m => MonadSample (IdentityT m) where
  random = lift random
  bernoulli = lift . bernoulli

instance MonadCond m => MonadCond (IdentityT m) where
  score = lift . score

instance MonadInfer m => MonadInfer (IdentityT m)

instance MonadSample m => MonadSample (MaybeT m) where
  random = lift random

instance MonadCond m => MonadCond (MaybeT m) where
  score = lift . score

instance MonadInfer m => MonadInfer (MaybeT m)

instance MonadSample m => MonadSample (ReaderT r m) where
  random = lift random
  bernoulli = lift . bernoulli

instance MonadCond m => MonadCond (ReaderT r m) where
  score = lift . score

instance MonadInfer m => MonadInfer (ReaderT r m)

instance (Monoid w, MonadSample m) => MonadSample (WriterT w m) where
  random = lift random
  bernoulli = lift . bernoulli
  categorical = lift . categorical

instance (Monoid w, MonadCond m) => MonadCond (WriterT w m) where
  score = lift . score

instance (Monoid w, MonadInfer m) => MonadInfer (WriterT w m)

instance MonadSample m => MonadSample (StateT s m) where
  random = lift random
  bernoulli = lift . bernoulli
  categorical = lift . categorical

instance MonadCond m => MonadCond (StateT s m) where
  score = lift . score

instance MonadInfer m => MonadInfer (StateT s m)

instance (MonadSample m, Monoid w) => MonadSample (RWST r w s m) where
  random = lift random

instance (MonadCond m, Monoid w) => MonadCond (RWST r w s m) where
  score = lift . score

instance (MonadInfer m, Monoid w) => MonadInfer (RWST r w s m)

instance MonadSample m => MonadSample (ListT m) where
  random = lift random
  bernoulli = lift . bernoulli
  categorical = lift . categorical

instance MonadCond m => MonadCond (ListT m) where
  score = lift . score

instance MonadInfer m => MonadInfer (ListT m)

instance MonadSample m => MonadSample (ContT r m) where
  random = lift random

instance MonadCond m => MonadCond (ContT r m) where
  score = lift . score

instance MonadInfer m => MonadInfer (ContT r m)


module Legacy.Base {-# DEPRECATED "Use Vector instead" #-} where

import Data.Complex
import Data.Either
import Data.List
import Control.Applicative

newtype StateVector = StateVector [Complex Double]
    deriving (Eq, Show)

newtype TransformTensor = TransformTensor [[Complex Double]]
    deriving (Eq, Show)

{-|
    integerLog2 determines if an integer is a power of 2
    If it is, the log base 2 of the integer will be wrapped
    by Right constructor; if not, Left constructor with
    corresponding error will be returned.
    The function takes 1 argument of type 'Int'
-}
integerLog2 :: Int -> Either String Int
integerLog2 x
    | x < 1 = Left "Input less than 0"
    | x == 1 = Right 0
    | y * 2 == x = (+1) <$> integerLog2 y
    | otherwise = Left "Input is not a power of 2"
    where
        y = quot x 2

initQubit :: Complex Double -> Complex Double -> Either String StateVector
initQubit a b
    | magnitude a ** 2 + magnitude b ** 2 == 1 = Right $ StateVector [a, b]
    | otherwise = Left "Magnitude Error"

initQubit0 :: StateVector
initQubit0 = StateVector [1:+0, 0:+0]

initNumQubits :: Int -> Either String StateVector
initNumQubits n
    | n < 1 = Left "Number of Qubits Must be Greater than 1"
    | otherwise = Right $ foldr mergeState (StateVector [1 :+ 0]) $ replicate n initQubit0

initState :: [Complex Double] -> Either String StateVector
initState q
    | isLeft numQubits = Left "Input State Dimension Invalid"
    | otherwise = Right $ StateVector q
    where
        numQubits = integerLog2 $ length q

mergeState :: StateVector -> StateVector -> StateVector
mergeState (StateVector x) (StateVector y) = StateVector $ liftA2 (*) x y

mergeTensor :: TransformTensor -> TransformTensor -> TransformTensor
mergeTensor (TransformTensor x) (TransformTensor y) = TransformTensor $ fmap ((<*>) . fmap (*)) x <*> y

tensorAdd :: TransformTensor -> TransformTensor -> Either String TransformTensor
tensorAdd (TransformTensor i) (TransformTensor j)
    | m == x && n == y = Right $ TransformTensor $ zipWith (zipWith (+)) i j
    | otherwise = Left "Adding Two Tensors with Different Dimensions"
    where
        m = length i
        n = length $ head i
        x = length j
        y = length $ head j

doubleBindEither :: Either String TransformTensor -> Either String TransformTensor -> (TransformTensor -> TransformTensor -> Either String TransformTensor) -> Either String TransformTensor
doubleBindEither x y f
    | isLeft x && isRight y = Left $ fromLeft [] x ++ "\n" ++ fromLeft [] y
    | isLeft x = x
    | isLeft y = y
    | otherwise = f (fromRight (TransformTensor [[]]) x) (fromRight (TransformTensor [[]]) y)

applyTensorToState :: TransformTensor -> StateVector -> Either String StateVector
applyTensorToState (TransformTensor g) (StateVector q)
    | length f == length q && length q == length g =
        Right $ StateVector $ foldr (zipWith (+)) acc $ zipWith fmap (fmap (*) q) f
    | otherwise = Left "Transformation Matrix Dimension Mismatch"
    where
        acc = replicate (length q) $0:+0
        f = transpose g

-- The qubits are 1-indexed
applyGateToQubit :: TransformTensor -> Int -> StateVector -> Either String StateVector
applyGateToQubit g n (StateVector q) = t >>= (`applyTensorToState` StateVector q)
    where
        numQubits = integerLog2 $ length q
        t = numQubits >>= flip (genTensor g) n

-- Gate to insert -> Number of Qubits -> Index of Qubit
genTensor :: TransformTensor -> Int -> Int -> Either String TransformTensor
genTensor g n i
    | n < i || i < 1 = Left "Index Out of Range"
    | otherwise = Right $ foldr mergeTensor (TransformTensor [[1:+0]]) expGate
    where
        expGate = fstId ++ [g] ++ lstId
        fstId = replicate (i - 1) pauliId
        lstId = replicate (n - i) pauliId

getQubitState :: Int -> StateVector -> Either String (Double, Double)
getQubitState x (StateVector q)
    | x < 1 = Left "Index Out of Bounds"
    | otherwise = Right (zeroP, oneP)
    where
        numRow = 2 ^ (x - 1)
        q2 = fmap (** 2) q
        zeroP = magnitude $ sum $ takeHalf q2 numRow
        oneP = magnitude $ sum $ takeHalf (drop numRow q2) numRow
        takeHalf :: [Complex Double] -> Int -> [Complex Double]
        takeHalf xs n
            | length xs < n = []
            | otherwise = take n xs ++ takeHalf (drop (2 * n) xs) n

-- Gate to Apply -> Control Bit Index -> Application Bit Index -> Circuit To Apply
applyControlToQubits :: TransformTensor -> Int -> Int -> StateVector -> Either String StateVector
applyControlToQubits g m n (StateVector q)
    | m < n = t m n mask1 g >>= (`applyTensorToState` StateVector q)
    | m > n = t n m g mask1 >>= (`applyTensorToState` StateVector q)
    | otherwise = Left "Control and Application Qubit Cannot be Identical"
    where
        zeroMask = numQubits >>= flip (genTensor mask0) m
        numQubits = integerLog2 $ length q
        t :: Int -> Int -> TransformTensor -> TransformTensor -> Either String TransformTensor
        t i j fstGate sndGate = doubleBindEither (fmap mergeTensor fstTensor <*> sndTensor) zeroMask tensorAdd
            where
                fstTensor = genTensor fstGate i i
                sndTensor = numQubits >>= flip (genTensor sndGate) (j - i) . flip (-) i

pauliX :: TransformTensor
pauliX = TransformTensor [[0 :+ 0, 1 :+ 0],
                      [1 :+ 0, 0 :+ 0]]

pauliY :: TransformTensor
pauliY = TransformTensor [[0 :+ 0, 0 :+ (-1)],
                      [0 :+ 1, 0 :+ 0]]

pauliZ :: TransformTensor
pauliZ = TransformTensor [[1 :+ 0, 0 :+ 0],
                      [0 :+ 0, (-1) :+ 0]]

pauliId :: TransformTensor
pauliId = TransformTensor [[1 :+ 0, 0 :+ 0],
                      [0 :+ 0, 1 :+ 0]]

hadamard :: TransformTensor
hadamard = TransformTensor [[1 / sqrt 2 :+ 0, 1 / sqrt 2 :+ 0],
                        [1 / sqrt 2 :+ 0, (-1) / sqrt 2 :+ 0]]

mask0 :: TransformTensor
mask0 = TransformTensor [[1 :+ 0, 0 :+ 0],
                    [0 :+ 0, 0 :+ 0]]

mask1 :: TransformTensor
mask1 = TransformTensor [[0 :+ 0, 0 :+ 0],
                    [0 :+ 0, 1 :+ 0]]

-- main :: IO()
-- main = print result
--     where q0 = initQubit0
--           crt = mergeState q0 q0
--           flip1 = applyGateToQubit hadamard 1 crt
--           flip2 = flip1 >>= applyControlToQubits pauliY 1 2
--           --result = flip2 >>= (getQubitState 2)
--           result = TransformTensor [[0 :+ 0, 0 :+ 0],
--                                     [0 :+ 0, 1 :+ 0]]

-- Linear algebra versions

-- safeMultiply :: Num a => T.Tensor a -> T.Tensor a -> Either String (T.Tensor a)
-- safeMultiply (T.Tensor x) (T.Tensor y)
--     | T.isSquare (T.Tensor x)
--         && T.isSquare (T.Tensor y)
--         && length x == length y = Right $ T.Tensor x * T.Tensor y
--     | otherwise = Left "Dimension mismatch"
-- safeMultiply x y = Right $ x * y

-- -- combine two qubit circuits
-- addQubit :: Num a => T.Tensor a -- ^ first half statevector
--   -> T.Tensor a -- ^ second half statevector
--   -> T.Tensor a
-- addQubit (T.Tensor x) (T.Tensor y)
--     | T.rank (T.Tensor x) /= 1 || T.rank (T.Tensor y) /= 1 = error "One of the elements is not a valid statevector"
--     | otherwise = T.Tensor $ liftA2 (*) x y
-- addQubit _ _ = error "One of the elements is not a valid statevector"

-- Applies input gate to qubit of index

-- -- Builds a transformation tensor to apply input gate to selected qubit
-- build :: RealFloat a => Int -- ^ Number of qubits this tensor applies to
--   -> [(Int, T.Tensor (Complex a))]-- ^ Index / Gate pair (Pre-sorted)
--   -> T.Tensor (Complex a)-- ^ Default Gate
--   -> T.Tensor (Complex a)
-- build len lst dft = T.flatten2d $ foldr T.product (T.fromList [1]) (go len lst 0)
--     where
--         go l (x:xs) offset
--             | fst x >= l = go l [] offset
--             | otherwise = replicate (fst x - offset) dft ++ [snd x] ++ go l xs (fst x + 1)
--         go l [] offset = replicate (l - offset) dft

-- applyGate :: RealFloat a => T.Tensor (Complex a) -- ^ Gate tensor
--   -> Int -- ^ Index
--   -> T.Tensor (Complex a) -- ^ State vector
--   -> Either String (T.Tensor (Complex a))
-- applyGate (T.Tensor xs) i (T.Tensor ys)
--     | T.rank x /= 2 || not (T.isSquare x) || numT == 0 = Left "Invalid Gate Tensor"
--     | T.rank y /= 1 || numV == 0 = Left "Invalid Statevector"
--     | i < 0 || i >= numV = Left "Index out of bounds"
--     | otherwise = Right $ build numV [(i, x)] pauliId * y
--     where
--         x = T.Tensor xs
--         y = T.Tensor ys
--         numT = fromRight 0 $ integralLog 2 (length xs)
--         numV = fromRight 0 $ integralLog 2 (length ys)
-- applyGate _ _ _ = Left "Input is not a Tensor"

-- applyGateAll :: RealFloat a => T.Tensor (Complex a) -> T.Tensor (Complex a) -> Either String (T.Tensor (Complex a))
-- applyGateAll (T.Tensor xs) (T.Tensor ys)
--     | T.rank x /= 2 || not (T.isSquare x) || numT == 0 = Left "Invalid Gate Tensor"
--     | T.rank y /= 1 || numV == 0 = Left "Invalid Statevector"
--     | otherwise = Right $ build numV [] x * y
--     where
--         x = T.Tensor xs
--         y = T.Tensor ys
--         numT = fromRight 0 $ integralLog 2 (length xs)
--         numV = fromRight 0 $ integralLog 2 (length ys)
-- applyGateAll _ _ = Left "Input is not a Tensor"

-- applyControl :: RealFloat a => T.Tensor (Complex a) -- ^ Gate tensor
--   -> Int -- ^ Control qubit index
--   -> Int -- ^ Apply qubit index
--   -> T.Tensor (Complex a) -- ^ State vector
--   -> Either String (T.Tensor (Complex a))
-- applyControl (T.Tensor xs) ctl i (T.Tensor ys)
--     | T.rank x /= 2 || not (T.isSquare x) || numT == 0 = Left "Invalid Gate Tensor"
--     | T.rank y /= 1 || numV == 0 = Left "Invalid Statevector"
--     | i < 0 || i >= numV = Left "Application qubit index out of bounds"
--     | ctl < 0 || ctl >= numV = Left "Control qubit index out of bounds"
--     | ctl == i = Left "Control and application qubit cannot be the same"
--     | otherwise = Right $ (build numV [(ctl, mask0)] pauliId
--         + build numV (sortBy (compare `on` fst) [(ctl, mask1), (i, x)]) pauliId) * y
--     where
--         x = T.Tensor xs
--         y = T.Tensor ys
--         numT = fromRight 0 $ integralLog 2 (length xs)
--         numV = fromRight 0 $ integralLog 2 (length ys)

-- applyControl _ _ _ _ = Left "Input is not a Tensor"


module Statistics.Information.Utils.KNN where

import qualified Data.KdTree.Static as KdTree
import Data.Matrix
import Data.Metric.Class
import Data.Metric.Vector.Real
import qualified Data.Vector as Vector
import Numeric.SpecFunctions
import Statistics.Information.Utils.Matrix

chebyshevDistance u v =
  distance (Chebyshev $ Vector.fromList u) (Chebyshev $ Vector.fromList v)

kdTree metric xs = KdTree.buildWithDist id metric xs

kdDistances tree metric xs k =
  map
    (\row -> let neighbor = KdTree.kNearest tree (k+1) row !! k in
      metric neighbor row)
    xs

knnDistances :: Real a => ([a] -> [a] -> a) -> [[a]] -> Int -> [a]
knnDistances metric xs k = kdDistances tree metric xs k where
  tree = kdTree metric xs

avgdigamma metric xs distances =
  mean . map (digamma . fromIntegral . length) $ radiusPoints
  where
    tree = kdTree metric (toLists xs)
    mean xs = sum xs / (fromIntegral $ length xs)
    radiusPoints = [KdTree.inRadius tree (d - 1e-15) x | (x, d) <- pds]
    pds = (map Vector.toList $ rows xs) `zip` distances


{-|
Module: MachineLearning.NeuralNetwork.MultiSvmLoss
Description: Multi SVM Loss.
Copyright: (c) Alexander Ignatyev, 2017
License: BSD-3
Stability: experimental
Portability: POSIX

Multi SVM Loss.
-}

module MachineLearning.NeuralNetwork.MultiSvmLoss
(
  scores
  , gradient
  , loss
)

where

import qualified Data.Vector.Storable as V
import qualified Numeric.LinearAlgebra as LA
import MachineLearning.Types (R, Matrix)
import MachineLearning.Utils (sumByRows, reduceByRows)

-- | SVM Delta
svmD :: R
svmD = 1.0

scores = id

gradient :: Matrix -> Matrix -> Matrix
gradient scores y =
    let correct_scores = sumByRows $ scores*(LA.step y)
        margins = scores - (correct_scores - (LA.scalar svmD))
        margins' = (1-y)*(LA.step margins)
        k = sumByRows margins'
    in margins' - (y * k)

loss :: Matrix -> Matrix -> R
loss scores y =
  let nSamples = fromIntegral $ LA.rows scores
      correct_scores = sumByRows $ scores*(LA.step y)
      margins = scores - (correct_scores - (LA.scalar svmD))
      margins' = margins * (LA.step margins)
      loss = LA.sumElements margins' - nSamples * svmD
  in loss / nSamples


module VB where

import Numeric.SpecFunctions (digamma)
import Numeric.LinearAlgebra.HMatrix
       (Matrix, size, (!), fromRows, scalar, sumElements, toRows, reshape,
        flatten, tr, dot, (<>), ( #> ), flatten, asRow, matrix, fromList, Vector, cmap)
import Numeric.LinearAlgebra.Devel (zipVectorWith)

data LogType
    = Log  -- ^ log x
    | InvLog  -- ^ log (1 - x)

type ExpectationQZ = Matrix Double
type ExpectationQmu = (Matrix Double, Matrix Double)  -- ^ Matrix Log and Matrix InvLog
type QZ = Matrix Double
type Qmu = (Matrix Double, Matrix Double)  -- ^ Matrix A and matrix B

-- | Moments of logarithmically transformed random variables
expectationBeta :: Double  -- ^ alpha
                -> Double -- ^ beta
                -> LogType  -- ^ log x or log (1 - x)
                -> Double
expectationBeta alpha beta (Log) = digamma alpha -
    digamma (alpha + beta)
expectationBeta alpha beta (InvLog) = digamma beta -
    digamma (alpha + beta)

em :: Matrix Double
   -> Vector Double
   -> Vector Double
   -> Double
   -> Double
   -> Vector Double
em = go 0 (-1)
  where
    go lQthetaNew lQthetaOld matrixX vectorPi vectorAlpha a b
        | abs (lQthetaNew - lQthetaOld) <
              0.1 =
            vectorPi
        | otherwise =
            let (eQz,eQmu) = eStep matrixX initQZ initQmu
                vectorPiNew = mStep vectorPi vectorAlpha eQz
            in go
                   (lQTheta matrixX vectorPiNew vectorAlpha eQz eQmu)
                   lQthetaNew
                   matrixX
                   vectorPiNew
                   vectorAlpha
                   a
                   b
        where initQZ = fromRows (replicate nRow vectorPi)
              initQmu = (matrixAB a, matrixAB b)
              matrixAB v = matrix
                      nColumn
                      [v | _ <-
                              [1 .. nMixtures * nColumn]]
              (nRow,nColumn) = size matrixX

eStep :: Matrix Double -> QZ -> Qmu -> (ExpectationQZ, ExpectationQmu)
eStep matrixX qZ qMu = let eQmu = calculateEQmu qMu
                           eQz = calculateEQZ qZ
                           qZ' = calculateQZ matrixX eQmu qZ
                           eQz' = calculateEQZ qZ'
                           qMu' = calculateQmu matrixX eQz' qMu
                           eQmu' = calculateEQmu qMu'
    in go
           (lQ matrixX eQz' eQmu')
           (lQ matrixX eQz eQmu)
           qZ'
           eQz'
           qMu'
           eQmu'
    where go lQnew lQold qZn eQzn qMun eQmun
              | abs (lQnew - lQold) <
                    0.1 =
                  (eQzn, eQmun)
              | otherwise =
                  let qZn' = calculateQZ matrixX eQmun qZn
                      eQzn' = calculateEQZ qZn'
                      qMun' = calculateQmu matrixX eQzn' qMun
                      eQmun' = calculateEQmu qMun'
                  in go (lQ matrixX eQzn' eQmun') lQnew qZn' eQzn' qMun' eQmun'

lQ :: Matrix Double -> ExpectationQZ -> ExpectationQmu -> Double
lQ matrixX eQz (eQmuM,eQmuMi) = sum
        (map
             (\n ->
                   (eQz ! n) `dot`
                   flatten
                       (asRow (matrixX ! n) <>
                        teQmuM +
                        asRow
                             (1 -
                              (matrixX ! n)) <>
                        teQmuMi))
             [0 .. nRow - 1])
    where teQmuM = tr eQmuM
          teQmuMi = tr eQmuMi
          (nRow,_) = size matrixX

calculateEQZ :: QZ -> ExpectationQZ
calculateEQZ qZ = fromRows
        (map
             (\row ->
                   row /
                   scalar (sumElements row))
             (toRows qZ))

calculateQZ :: Matrix Double -> ExpectationQmu -> QZ -> QZ
calculateQZ matrixX (eQmuM,eQmuMi) qZ = qZ + matrixX <> tr eQmuM +
    (1 - matrixX) <>
    tr eQmuMi

calculateEQmu :: Qmu -> ExpectationQmu
calculateEQmu (aM, bM) = (reshape d qMuM, reshape d qMuMi)
  where
    aV = flatten aM
    bV = flatten bM
    qMuM = zipVectorWith (\a b -> expectationBeta a b Log) aV bV
    qMuMi = zipVectorWith (\a b -> expectationBeta a b InvLog) aV bV
    (k, d) = size aM

calculateQmu :: Matrix Double -> ExpectationQZ -> Qmu -> Qmu
calculateQmu matrixX eQz (aOld,bOld) = ( aOld - 1 + tr eQz <> matrixX
                                       , bOld - 1 +
                                         tr eQz <>
                                         (1 - matrixX))

lQTheta :: Matrix Double
        -> Vector Double
        -> Vector Double
        -> ExpectationQZ
        -> ExpectationQmu
        -> Double
lQTheta matrixX vectorPi vectorAlpha eQz eQmu = lQ matrixX eQz eQmu +
    vectorAlpha `dot`
    cmap log vectorPi

mStep :: Vector Double -> Vector Double -> ExpectationQZ -> Vector Double
mStep vectorPi vectorAlpha eQz = let vectorNewPi = fromList
                                             (map
                                                  (\k ->
                                                        -(vectorPi ! k) *
                                                         sumElements
                                                              (eqZt ! k) /
                                                         (vectorAlpha ! k - 1))
                                                  [0 .. nMixtures - 1])
    in norm vectorNewPi
    where eqZt = tr eQz
          norm v = v /
              scalar (sumElements v)

nMixtures :: Int
nMixtures = 10


{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE UndecidableInstances #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE BangPatterns #-}
module HLearn.Models.Distributions.Univariate.Binomial
    where

import HLearn.Algebra
import HLearn.Models.Distributions.Common
import HLearn.Models.Distributions.Univariate.Internal.Moments
import HLearn.Models.Distributions.Visualization.Gnuplot

import qualified Statistics.Distribution as S
import Statistics.Distribution.Binomial

-------------------------------------------------------------------------------
-- data types

newtype Binomial prob dp = Binomial {  bmoments :: (Moments3 dp) }
    deriving (Read,Show,Eq,Ord,Monoid,Group)

-------------------------------------------------------------------------------
-- Training

instance (Num dp) => HomTrainer (Binomial prob dp) where
    type Datapoint (Binomial prob dp) = dp
    train1dp dp = Binomial $ train1dp dp

-------------------------------------------------------------------------------
-- distribution

instance (Num dp) => Probabilistic (Binomial prob dp) where
    type Probability (Binomial prob dp) = prob

instance (Floating prob) => PDF (Binomial Double Int) where
    pdf (Binomial dist) dp = S.probability (binomial n p) dp
        where
            n = bin_n $ Binomial dist
            p = bin_p $ Binomial dist

bin_n :: Binomial Double Int -> Int
bin_n (Binomial dist) = round $ ((fromIntegral $ m1 dist :: Double) / (fromIntegral $ m0 dist)) / (bin_p $ Binomial dist)

bin_p :: Binomial Double Int -> Double
bin_p (Binomial dist) = ((fromIntegral $ m1 dist) / (fromIntegral $ m0 dist)) + 1 - (fromIntegral $ m2 dist)/(fromIntegral $ m1 dist)

instance
    ( PDF (Binomial prob dp)
--     , PlottableDataPoint dp
    , Show prob
    , Show dp
    , Ord dp
    , Ord prob
    , Num prob
    , Integral dp
    ) => PlottableDistribution (Binomial prob dp)
-- instance PlottableDistribution (Poisson Int Double)
        where

    plotType _ = Points

    samplePoints dist = [min..max]
        where
            min = 0
            max = maximum [20,floor $ 3*mu]
            mu = 5

-- instance (Fractional prob) => Mean (Binomial prob) where
--     mean (Binomial dist) =
--
-- instance (Fractional prob) => Variance (Binomiral prob) where
--     variance dist = mean dist


module ML.Unsupervised.Clustering.KMeans (
  fit,
  predict,
  KMeans(..)
) where

import Numeric.LinearAlgebra
import Numeric.LinearAlgebra.Data

import Data.Ratio ((%))
import Data.List (minimumBy)
import qualified Data.Vector.Storable as V

import Debug.Trace

data KMeans a l = KMeans {
  centroids :: [(l, V.Vector a)]
} deriving (Show, Eq)

fit :: (Fractional a, Ord a, Eq l, Show l, Show a, V.Storable a) => [l] -> [V.Vector a] -> KMeans a l
fit labels points = fit' (KMeans initialCentroids)
  where
    fit' kmeans = let
                    ass = assignment kmeans
                    updatedKmeans = update kmeans ass in
                    if (trace ("updated assignment: " ++ show (assignment updatedKmeans)) assignment updatedKmeans) == ass
                      then updatedKmeans
                      else fit' updatedKmeans
    update kmeans ass = KMeans $ zip labels $ fmap (\l -> centerOfMass $ snd <$> filter (\x -> fst x == l) ass) labels
    assignment kmeans = fmap (\x -> (assign kmeans x, x)) points
    assign kmeans point = fst $ minimumBy (\x1 x2 -> snd x1 `compare` snd x2) $ (fmap (\(l, cent) -> (l, distance cent point))) $ centroids kmeans
    distance v1 v2 = V.sum $ V.zipWith (\x1 x2 -> (x2 - x1) * (x2 - x1)) v1 v2
    initialCentroids = zip labels points

    centerOfMass :: (Fractional a, V.Storable a) => [V.Vector a] -> V.Vector a
    centerOfMass (p:pts) = V.map (/ (fromRational $ (toInteger $ length pts + 1) % 1)) $ foldr sumVectors p pts

    dim = V.length $ head points
    sumVectors :: (Num a, V.Storable a) => V.Vector a -> V.Vector a -> V.Vector a
    sumVectors = V.zipWith (+)

predict :: (Ord a, Num a, V.Storable a) => KMeans a l -> V.Vector a -> l
predict kmeans point = fst $ minimumBy (\x1 x2 -> snd x1 `compare` snd x2) $ (fmap (\(l, cent) -> (l, distance cent point))) $ centroids kmeans
  where
    distance v1 v2 = V.sum $ V.zipWith (\x1 x2 -> (x2 - x1) * (x2 - x1)) v1 v2

fitMatrix :: (Element a) => [l] -> Matrix a -> KMeans a l
fitMatrix = undefined


{-# LANGUAGE DefaultSignatures          #-}
{-# LANGUAGE DeriveGeneric              #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE StandaloneDeriving         #-}

{-|
Module: Data.Semiring
Description: Haskell semirings
License: MIT
Maintainer: mail@doisinkidney.com
Stability: experimental
-}

module Data.Semiring
  ( Semiring(..)
  , Add(..)
  , Mul(..)
  , mul
  , add
  ) where

import           Data.Coerce           (coerce)

import           Data.Functor.Identity (Identity (..))

import           Data.Monoid
import           Data.Semigroup        (Max (..), Min (..))

import           Data.Complex          (Complex)
import           Data.Fixed            (Fixed, HasResolution)
import           Data.Ratio            (Ratio)
import           Numeric.Natural       (Natural)

import           Data.Set              (Set)
import qualified Data.Set              as Set

import           Data.Int              (Int16, Int32, Int64, Int8)
import           Data.Word             (Word16, Word32, Word64, Word8)
import           Foreign.C.Types       (CChar, CClock, CDouble, CFloat, CInt,
                                        CIntMax, CIntPtr, CLLong, CLong,
                                        CPtrdiff, CSChar, CSUSeconds, CShort,
                                        CSigAtomic, CSize, CTime, CUChar, CUInt,
                                        CUIntMax, CUIntPtr, CULLong, CULong,
                                        CUSeconds, CUShort, CWchar)

import           Foreign.Ptr           (IntPtr, WordPtr)
import           System.Posix.Types    (CCc, CDev, CGid, CIno, CMode, CNlink,
                                        COff, CPid, CRLim, CSpeed, CSsize,
                                        CTcflag, CUid, Fd)

import           GHC.Generics          (Generic, Generic1)

import           Test.QuickCheck       (Arbitrary)

-- | A <https://en.wikipedia.org/wiki/Semiring Semiring> is like the
-- the combination of two 'Data.Monoid.Monoid's. The first
-- is called '<+>'; it has the identity element 'zero', and it is
-- commutative. The second is called '<.>'; it has identity element 'one',
-- and it must distribute over '<+>'.
--
-- = Laws
-- == Normal 'Precursor.Algebra.Monoid.Monoid' laws
-- * @(a '<+>' b) '<+>' c = a '<+>' (b '<+>' c)@
-- * @'zero' '<+>' a = a '<+>' 'zero' = a@
-- * @(a '<.>' b) '<.>' c = a '<.>' (b '<.>' c)@
-- * @'one' '<.>' a = a '<.>' 'one' = a@
--
-- == Commutativity of '<+>'
-- * @a '<+>' b = b '<+>' a@
--
-- == Distribution of '<.>' over '<+>'
-- * @a '<.>' (b '<+>' c) = (a '<.>' b) '<+>' (a '<.>' c)@
-- * @(a '<+>' b) '<.>' c = (a '<.>' c) '<+>' (b '<.>' c)@
--
-- Another useful law, annihilation, may be deduced from the axioms
-- above:
--
-- * @'zero' '<.>' a = a '<.>' 'zero' = 'zero'@
class Semiring a where
  -- | The identity of '<+>'.
  zero :: a
  -- | The identity of '<.>'.
  one :: a
  -- | An associative binary operation, which distributes over '<+>'.
  infixl 7 <.>
  (<.>) :: a -> a -> a
  -- | An associative, commutative binary operation.
  infixl 6 <+>
  (<+>) :: a -> a -> a

  default zero :: Num a => a
  default one :: Num a => a
  default (<+>) :: Num a => a -> a -> a
  default (<.>) :: Num a => a -> a -> a

  zero = 0
  one = 1
  (<+>) = (+)
  (<.>) = (*)

------------------------------------------------------------------------
-- Instances
------------------------------------------------------------------------
instance Semiring Bool where
  one = True
  zero = False
  (<+>) = (||)
  (<.>) = (&&)

instance Semiring () where
  one = ()
  zero = ()
  _ <+> _ = ()
  _ <.> _ = ()

cartProd :: (Ord a, Monoid a) => Set a -> Set a -> Set a
cartProd xs ys =
  Set.foldl' (\a x ->
                Set.foldl' (flip (Set.insert . mappend x)) a ys)
  Set.empty xs

-- | The 'Set' 'Semiring' is 'Data.Set.union' for '<+>', and a Cartesian
-- product for '<.>'.
instance (Ord a, Monoid a) => Semiring (Set a) where
  (<.>) = cartProd
  (<+>) = Set.union
  zero = Set.empty
  one = Set.singleton mempty

------------------------------------------------------------------------
-- Addition and multiplication newtypes
------------------------------------------------------------------------

type WrapBinary f a = (a -> a -> a) -> f a -> f a -> f a

-- | Monoid under '<+>'. Analogous to 'Data.Monoid.Sum', but uses the
-- 'Semiring' constraint, rather than 'Num'.
newtype Add a = Add
  { getAdd :: a
  } deriving (Eq, Ord, Read, Show, Enum, Bounded, Generic, Generic1, Num
             ,Arbitrary)

-- | Monoid under '<.>'. Analogous to 'Data.Monoid.Product', but uses the
-- 'Semiring' constraint, rather than 'Num'.
newtype Mul a = Mul
  { getMul :: a
  } deriving (Eq, Ord, Read, Show, Enum, Bounded, Generic, Generic1, Num
             ,Arbitrary)

add :: (Semiring a, Foldable f) => f a -> a
add = getAdd . foldMap Add

mul :: (Semiring a, Foldable f) => f a -> a
mul = getMul . foldMap Mul

instance Functor Add where fmap = coerce

instance Functor Mul where fmap = coerce

instance Foldable Add where
  foldr   =
    (coerce :: ((a -> b -> c) -> (b -> a -> c))
            -> (a -> b -> c)
            -> (b -> Add a -> c)) flip
  foldl   = coerce
  foldMap = coerce
  length  = const 1

instance Foldable Mul where
  foldr   =
    (coerce :: ((a -> b -> c) -> (b -> a -> c))
            -> (a -> b -> c)
            -> (b -> Mul a -> c)) flip
  foldl   = coerce
  foldMap = coerce
  length  = const 1

instance Applicative Add where
  pure = coerce
  (<*>) =
    (coerce :: ((a -> b) -> a -> b)
            -> (Add (a -> b) -> Add a -> Add b)) ($)

instance Applicative Mul where
  pure = coerce
  (<*>) =
    (coerce :: ((a -> b) -> a -> b)
            -> (Mul (a -> b) -> Mul a -> Mul b)) ($)

instance Monad Add where
  (>>=) = flip coerce

instance Monad Mul where
  (>>=) = flip coerce

instance Semiring a => Monoid (Add a) where
  mempty = Add zero
  mappend = (coerce :: WrapBinary Add a) (<+>)

instance Semiring a => Monoid (Mul a) where
  mempty = Mul one
  mappend = (coerce :: WrapBinary Mul a) (<.>)

instance Semiring a => Semiring (Add a) where
  zero = Add zero
  one = Add one
  (<+>) = (coerce :: WrapBinary Add a) (<+>)
  (<.>) = (coerce :: WrapBinary Add a) (<.>)

instance Semiring a => Semiring (Mul a) where
  zero = Mul zero
  one = Mul one
  (<+>) = (coerce :: WrapBinary Mul a) (<+>)
  (<.>) = (coerce :: WrapBinary Mul a) (<.>)

------------------------------------------------------------------------
-- Ord wrappers
------------------------------------------------------------------------

-- | The 'Semiring' for 'Max' uses the 'max' operation for '<+>', and
-- normal '+' for '<.>'.
instance (Ord a, Bounded a, Semiring a) => Semiring (Max a) where
  (<+>) = mappend
  zero = mempty
  (<.>) = (coerce :: WrapBinary Max a) (<+>)
  one = Max zero

-- | The 'Semiring' for 'Min' uses the 'min' operation for '<+>', and
-- normal '+' for '<.>'.
instance (Ord a, Bounded a, Semiring a) => Semiring (Min a) where
  (<+>) = mappend
  zero = mempty
  (<.>) = (coerce :: WrapBinary Min a) (<+>)
  one = Min zero

------------------------------------------------------------------------
-- (->) instance
------------------------------------------------------------------------

-- | The @(->)@ instance is analogous to the one for 'Monoid'.
instance Semiring b => Semiring (a -> b) where
  zero = const zero
  one = const one
  (f <+> g) x = f x <+> g x
  (f <.> g) x = f x <.> g x

------------------------------------------------------------------------
-- Endo instance
------------------------------------------------------------------------

-- | The 'Endo' semiring uses function composition for '<.>', and
-- pointwise 'mappend' for '<+>'. The underlying 'Monoid' needs to be
-- commutative.
instance Monoid a => Semiring (Endo a) where
  (<.>) = mappend
  one = mempty
  Endo f <+> Endo g = Endo (\x -> f x `mappend` g x)
  zero = Endo (const mempty)

------------------------------------------------------------------------
-- Boring instances
------------------------------------------------------------------------
instance Semiring Float
instance Semiring Double
instance Semiring Int
instance Semiring Int8
instance Semiring Int16
instance Semiring Int32
instance Semiring Int64
instance Semiring Integer
instance Semiring Word
instance Semiring Word8
instance Semiring Word16
instance Semiring Word32
instance Semiring Word64
instance Semiring CUIntMax
instance Semiring CIntMax
instance Semiring CUIntPtr
instance Semiring CIntPtr
instance Semiring CSUSeconds
instance Semiring CUSeconds
instance Semiring CTime
instance Semiring CClock
instance Semiring CSigAtomic
instance Semiring CWchar
instance Semiring CSize
instance Semiring CPtrdiff
instance Semiring CDouble
instance Semiring CFloat
instance Semiring CULLong
instance Semiring CLLong
instance Semiring CULong
instance Semiring CLong
instance Semiring CUInt
instance Semiring CInt
instance Semiring CUShort
instance Semiring CShort
instance Semiring CUChar
instance Semiring CSChar
instance Semiring CChar
instance Semiring IntPtr
instance Semiring WordPtr
instance Semiring Fd
instance Semiring CRLim
instance Semiring CTcflag
instance Semiring CSpeed
instance Semiring CCc
instance Semiring CUid
instance Semiring CNlink
instance Semiring CGid
instance Semiring CSsize
instance Semiring CPid
instance Semiring COff
instance Semiring CMode
instance Semiring CIno
instance Semiring CDev
instance Semiring Natural
instance Integral a => Semiring (Ratio a)
deriving instance Semiring a => Semiring (Product a)
deriving instance Semiring a => Semiring (Sum a)
instance RealFloat a => Semiring (Complex a)
instance HasResolution a => Semiring (Fixed a)
deriving instance Semiring a => Semiring (Identity a)

------------------------------------------------------------------------
-- Very boring instances
------------------------------------------------------------------------
instance (Semiring a, Semiring b) => Semiring (a,b) where
        zero = (zero, zero)
        (a1,b1) <+> (a2,b2) =
                (a1 <+> a2, b1 <+> b2)
        one = (one, one)
        (a1,b1) <.> (a2,b2) =
                (a1 <.> a2, b1 <.> b2)

instance (Semiring a, Semiring b, Semiring c) => Semiring (a,b,c) where
        zero = (zero, zero, zero)
        (a1,b1,c1) <+> (a2,b2,c2) =
                (a1 <+> a2, b1 <+> b2, c1 <+> c2)
        one = (one, one, one)
        (a1,b1,c1) <.> (a2,b2,c2) =
                (a1 <.> a2, b1 <.> b2, c1 <.> c2)

instance (Semiring a, Semiring b, Semiring c, Semiring d) => Semiring (a,b,c,d) where
        zero = (zero, zero, zero, zero)
        (a1,b1,c1,d1) <+> (a2,b2,c2,d2) =
                (a1 <+> a2, b1 <+> b2,
                 c1 <+> c2, d1 <+> d2)
        one = (one, one, one, one)
        (a1,b1,c1,d1) <.> (a2,b2,c2,d2) =
                (a1 <.> a2, b1 <.> b2, c1 <.> c2, d1 <.> d2)

instance (Semiring a, Semiring b, Semiring c, Semiring d, Semiring e) =>
                Semiring (a,b,c,d,e) where
        zero = (zero, zero, zero, zero, zero)
        (a1,b1,c1,d1,e1) <+> (a2,b2,c2,d2,e2) =
                (a1 <+> a2, b1 <+> b2, c1 <+> c2,
                 d1 <+> d2, e1 <+> e2)
        one = (one, one, one, one, one)
        (a1,b1,c1,d1,e1) <.> (a2,b2,c2,d2,e2) =
                (a1 <.> a2, b1 <.> b2, c1 <.> c2,
                 d1 <.> d2, e1 <.> e2)


module MachineLearning.Optimization.MinibatchGradientDescentTest
(
  tests
)

where

import Test.Framework (testGroup)
import Test.Framework.Providers.HUnit
import Test.HUnit
import Test.HUnit.Approx
import Test.HUnit.Plus

import qualified Data.Vector.Storable as V
import qualified Numeric.LinearAlgebra as LA

import MachineLearning.Types (Vector)
import qualified MachineLearning as ML
import MachineLearning.Regularization (Regularization(..))
import MachineLearning.LeastSquaresModel (LeastSquaresModel(..))
import MachineLearning.Optimization.MinibatchGradientDescent

import MachineLearning.DataSets (dataset1)

(x, y) = ML.splitToXY dataset1

muSigma = ML.meanStddev x
xNorm = ML.featureNormalization muSigma x
x1 = ML.addBiasDimension xNorm
initialTheta :: Vector
initialTheta = LA.konst 0 (LA.cols x1)
lsExpectedTheta = LA.vector [325010.120,113889.649,5234.404]
eps = 1e-3

isInDescendingOrder :: [Double] -> Bool
isInDescendingOrder lst = and . snd . unzip $ scanl (\(prev, _) current -> (current, prev >= current)) (1/0, True) lst

testMinibatchGradientDescent model expectedTheta = do
  let (theta, optPath) = minibatchGradientDescent 0 16 0.01 model eps 5000 RegNone x1 y initialTheta
      js = V.toList $ (LA.toColumns optPath) !! 1
  assertVector "theta" 0.01 expectedTheta theta
  assertBool "non-increasing errors" $ isInDescendingOrder js

tests = [testGroup "minibatchGradientDescent" [
            testCase "leastSquares" $ testMinibatchGradientDescent LeastSquares lsExpectedTheta
            ]
        ]


module Tool where
  import Numeric.LinearAlgebra
  import Data.List
  import System.Random

  {- Convenience functions -}

  -- zero vector of length n
  zeros n = n |> [0,0..]

  -- one vector of length n
  ones n = n |> [1,1..]

  -- random vector of length n
  rands range n = do
    randarr <- (getStdGen >>= (return.randomRs range))
    return $ n |> randarr

  -- zero matrix of (n, m) shape
  zeromat n m = (n><m) [0,0..]

  -- one matrix of (n, m) shape
  onemat n m = (n><m) [1,1..]

  -- random matrix of (n, m) shape
  randmat range n m = do
    randarr <- (getStdGen >>= (return.randomRs range))
    return $ (n><m) randarr

  -- append 1 at the end of a vector
  homo :: Vector Double -> Vector Double
  homo v = vjoin [v, vector [1]]

  -- remove 1 at the end of a vector
  invhomo :: Vector Double -> Vector Double
  invhomo v = subVector 0 (size v - 1) v

  -- Make a matrix by repeating a column vector
  repeatCols :: Int -> Vector Double -> Matrix Double
  repeatCols n v = repmat (asColumn v) 1 n

  -- Given one-hot label, convert it to a index
  stepLabel :: Vector Double -> Int
  stepLabel v = argmax (toList v)

  -- Check if two one-hot labels are equal
  labelEquals :: Vector Double -> Vector Double -> Bool
  labelEquals a b = (argmax $ toList a) == (argmax $ toList b)

  -- Find the index of maximum element in a list
  argmax l = case elemIndex (maximum l) l of Just result -> result

  -- Split a list into half
  split l = [take halfLength l, drop halfLength l]
    where halfLength = length l `quot` 2


module Main where
import Plot
import Graphics.GD
import Data.Complex

imageDim :: Size
imageDim = (600,600)

-- How many iterations before  we declare the point doesn't diverge
maxIter = 127

-- Coordinates in the complex plane to graph
-- First coordinate must be top left.
mbWindow :: Window
mbWindow = ( (-1.9, 1.3), (0.7, -1.3) ) -- Full view

main :: IO ()
main = drawPlot drawMandelbrot imageDim mbWindow "mandelbrot.png"

drawMandelbrot :: Coordinate -> Color
drawMandelbrot (x,y) = colorIterations $ mandelbrot (x :+ y) (0 :+ 0) 0

-- Count number of iterations for the given point to diverge
-- Non-diverging points return 0
mandelbrot :: Complex Double -- Coordinate to test
           -> Complex Double -- Iterating Z value
           -> Int -- Current iteration
           -> Int -- Iterations before diverging
mandelbrot c z iter
    | iter > maxIter = 0
    | otherwise = let z' = z^2 + c in
                  if magnitude z' > 2
                  then iter
                  else mandelbrot c z' (iter+1)

colorIterations :: Int -> Color
colorIterations x
    | x > maxIter = rgb 255 255 255
    | otherwise = let c = x + 30
                  in rgb c x (x^2)


module Main where

import           Numeric.LinearAlgebra
import           System.Environment (getArgs)
import           System.Random
import           System.ProgressBar
import           Control.Monad (replicateM)
import           Control.Monad.IO.Class
import           Data.Maybe
import           Data.List
import qualified Data.Vector as V
import           Polynomials
import           Genetic
import           Formulas

main :: IO ()
main = do
    (nGenerations:nPop:goal:nElit:_) <- getArgs

    ------------ Information ------------
    putStrLn ""
    putStrLn $ "No. of generatoins: "  ++ nGenerations
    putStrLn $ "Population Size: "     ++ nPop
    putStrLn $ "Score of goal: "       ++ goal
    putStrLn $ "Elite size (N-elit): " ++ nElit
    putStrLn ""

    ------------ Progress Bar -----------
    putStrLn "Progress:"
    (pb, asy) <- startProgress (msg "") (msg "") 10 (Progress 0 (read nGenerations))

    ------------ Algorithm --------------
    powers <- randomPowerGen 6 2
    let
        ------ Inps and outps of function ------
        inps  = inpsRosen
        outps = fromList $ V.toList yRosen
        ----------------------------------------
        individual = Article powers inps outps
        init = replicate (read nPop) individual

    algo <- runGeneticAlgorithm (read goal) (read nGenerations) (read nElit) init pb
    let
        finalPowers = _articlePowers algo
        x = calcX inps finalPowers
        y = outps
        coeffs = normEq x y
        ----------------------------------------
    grade <- show . fromJust <$> fitness algo

    putStrLn ""
    putStrLn $ showPolynomial coeffs finalPowers
    putStrLn $ "Fitness of final Polynomial is: " ++ grade
    putStrLn ""
    putStrLn "Attention: component with coefficient smaller tan 0.001 are rounded to nothing"

openVLists = V.map V.toList

-- Check x + 1:
simple x = x + 1.0

inpsSimple :: V.Vector Double
inpsSimple = V.fromList
       [ 1
       , 2
       , 3
       , 4]

powsSimple :: V.Vector Double
powsSimple = V.fromList
            [1.0,
             0]

ySimple = V.map simple inpsSimple

-- Check x**2 + y**2 -1:
circle [x, y] = x**2 + y**2 --  -1

inpsCircle :: Inplist
inpsCircle = V.fromList
     . map V.fromList
     $ [ [0.0001, 1.1]
       , [2, 0.001]
       , [2, 1.05]
       , [1, 2.11]
       , [0.1, 2]]

powsCircle :: Powers
powsCircle = V.fromList
          . map V.fromList
          $ [[2, 0],
             [0, 2],
             [0, 0]]

yCircle = V.map circle $ openVLists inpsCircle

-- Check Rosenbrock:
rosen [x, y] =
    1 - 2*x + x**2 + 100*y**2 + 200*x**2*y + 100*x**4

inpsRosen :: Inplist
inpsRosen = V.fromList
     . map V.fromList
     $ [[i, j] | i <- [-0.9, -0.8 .. -0.2], j <- [0.2, 0.3 .. 0.9]]

powsRosen :: Powers
powsRosen = V.fromList
        . map V.fromList
        $ [[0, 0],
           [1, 0],
           [2, 0],
           [0, 2],
           [2, 1],
           [4, 0]]

yRosen = V.map rosen $ openVLists inpsRosen