-- file: Astro/Benchmarks.hs
-- Benchmarks for some of the functions

module Astro.Benchmarks
	where

import Astro.Stumpff
import Astro.Kepler
import Astro.Elements

import Numeric.GSL
import Numeric.LinearAlgebra
import Control.Monad

keplerPropagateBM :: Int -> IO ()
keplerPropagateBM nsteps = step 0 (rv,vv)
	where
	(rv, vv) 	= (3|>[1,0,0], 3|>[0,0.1,0])
	oe 		= posvelToOrbel rv vv
	h  		= 10*2*pi/(fromIntegral nsteps::Double)
	step n (r,v) = do
		putStrLn $ (show r)
		(when (n < nsteps) $ step (n+1) newS)
		where
		newS = keplerPropagate r v h (1e-15,100)


-- This file is part of Quipper. Copyright (C) 2011-2014. Please see the
-- file COPYRIGHT for a list of authors, copyright holders, licensing,
-- and other details. All rights reserved.
--
-- ======================================================================

{-# LANGUAGE NoMonomorphismRestriction #-}

-- | This module contains an implementation of the oracle and its
-- automatic lifting to quantum circuits using Template Haskell.
module Algorithms.QLS.TemplateOracle where

import Data.Complex
import Libraries.Auxiliary

import Control.Monad

import QuipperLib.Arith hiding (template_symb_plus_)
import Algorithms.QLS.Utils
import Algorithms.QLS.QDouble
import Algorithms.QLS.RealFunc
import Algorithms.QLS.QSignedInt
import Algorithms.QLS.CircLiftingImport

import Quipper
import Quipper.CircLifting

-- | Lifted version of @'any'@ (local version).
build_circuit
local_any :: (a -> Bool) -> [a] -> Bool
local_any f l = case l of
            []    -> False
            (h:t) -> if (f h) then True else local_any f t

-- | Auxiliary function.
build_circuit
itoxy :: Int -> Int -> Int -> (Int, Int)
itoxy i nx ny =
   if (i <= 0) then (0,0)
   else if (i > (nx-1)*ny + nx*(ny-1)) then (0,0)
   else ((mod (i - 1) (2*nx - 1)) + 1, ceiling ((fromIntegral i)/(2.0*(fromIntegral nx)-1.0)))

-- | The function sin /x/ \/ /x/.
build_circuit
sinc :: Double -> Double
sinc x = if (x /= 0.0) then (sin x) / x else 1.0

-- | Auxiliary function.
build_circuit
edgetoxy :: Int -> Int -> Int -> (Double, Double)
edgetoxy e nx ny =
  let (ex,ey) = itoxy e nx ny in
     if (ex < nx) then ((fromIntegral ex) + 0.5, fromIntegral ey)
     else (fromIntegral (ex - nx + 1), (fromIntegral ey) + 0.5)

-- | Auxiliary function. The inputs are:
--
-- * /y1/ :: 'Int' - Global edge index of row index of desired matrix element;
--
-- * /y2/ :: 'Int' - Global edge index of column index of desired matrix element;
--
-- * /nx/ :: 'Int' - Number of vertices left to right;
--
-- * /ny/ :: 'Int' - Number of vertices top to bottom;
--
-- * /lx/ :: 'Double' - Length of horizontal edges (distance between vertices in /x/ direction);
--
-- * /ly/ :: 'Double' - Length of vertical edges (distance between vertices in /y/ direction);
--
-- * /k/ :: 'Double' - Plane wave wavenumber.
--
-- The output is the matrix element /A/(/y1/, /y2/).
build_circuit
calcmatrixelement :: --
	-- Inputs
	Int 	-- int y1 - Global edge index of row index of desired matrix element
	-> Int 	-- int y2 - Global edge index of column index of desired matrix element
	-> Int 	-- int nx - Number of vertices left to right
	-> Int	-- int ny - Number of vertices top to bottom
	-> Double -- float lx - Length of horizontal edges (distance between vertices in x direction)
	-> Double -- float ly - Length of vertical edges (distance between vertices in y direction)
	-> Double -- float k - Plane wave wavenumber
	-- Outputs
	-> Complex Double -- float A - Matrix element A(y1,y2)
calcmatrixelement y1 y2 nx ny lx ly k =
	let (xg1, yg1) = itoxy y1 nx ny in
        let (xg2, yg2) = itoxy y2 nx ny in
        let b = if ( (y1==y2) && (xg1 >= nx) ) then ly/lx - k*k*lx*ly/3.0
                -- B11 and B22
                else if ( (y1==y2) && (xg1<nx) ) then lx/ly - k*k*lx*ly/3.0
                -- B12 and B21
                else if ( (abs(yg1-yg2) == 1) && (abs(xg1-xg2) == 0) && (xg1<nx) ) then -lx/ly + k*k*lx*ly/12.0
                -- B34 and B43
                else if ( (abs(yg1-yg2)==0) && (abs(xg1-xg2) == 1) && (xg1>=nx) ) then -ly/lx + k*k*lx*ly/12.0
                -- B13
                else if ( (yg1==(yg2+1)) && (xg1==(xg2-nx+1)) && (xg2>=nx) ) then -1.0
                -- B31
                else if ( (yg2==(yg1+1)) && (xg2==(xg1-nx+1)) && (xg1>=nx) ) then -1.0
                -- B14
                else if ( (yg1==(yg2+1)) && (xg1==(xg2-nx)) && (xg1<nx) ) then 1.0
                -- B41
                else if ( (yg2==(yg1+1)) && (xg2==(xg1-nx)) && (xg2<nx) ) then 1.0
                -- B42
                else if ( (yg1==yg2) && (xg1==(xg2+nx)) && (xg1>=nx) ) then -1.0
                -- B24
                else if ( (yg2==yg1) && (xg2==(xg1+nx)) && (xg2>=nx) ) then -1.0
                -- B32
                else if ( (yg1==yg2) && (xg1==(xg2+nx-1)) && (xg2<nx) ) then 1.0
                -- B23
                else if ( (yg2==yg1) && (xg2==(xg1+nx-1)) && (xg1<nx) ) then 1.0
                else -1.0
        in
        let c = if ( (y1==y2) && ( (xg1==nx) || (xg1==(2*nx-1)) ) ) then 0.0 :+ (k*ly)
                else if  ( (y1==y2) && ( ((yg1==1) && (xg1<nx)) || ((yg1==ny) && (xg1<nx)) ) ) then 0.0 :+ (k*lx)
                else 0.0 :+ 0.0
	in (b :+ 0.0) + c

-- | Auxiliary function.
build_circuit
get_edges l = case l of
                [] -> []
                (x:tt) -> case tt of
                   [] -> []
                   (y:t) -> (x,y):(get_edges (y:t))

-- | Auxiliary function.
build_circuit
checkedge :: Int -> [(Double,Double)] -> Int -> Int -> Bool
checkedge e scatteringnodes nx ny =
     let (xi,yi) = edgetoxy e nx ny in
     let test_elt ((x1,y1),(x2, y2)) = xi >= x1 && yi >= y1 && xi <= x2 && yi <= y2 in
     let half = take_half scatteringnodes in
     let test_list = local_any test_elt (get_edges half) in (not test_list)

-- | Oracle /r/.
build_circuit
calcRweights :: Int -> Int -> Int -> Double -> Double -> Double -> Double -> Double -> Complex Double
calcRweights y nx ny lx ly k theta phi =
     let (xc',yc') = edgetoxy y nx ny in
     let xc = (xc'-1.0)*lx - ((fromIntegral nx)-1.0)*lx/2.0 in
     let yc = (yc'-1.0)*ly - ((fromIntegral ny)-1.0)*ly/2.0 in
     let (xg,yg) = itoxy y nx ny in

     if (xg == nx) then

         let i = (mkPolar ly (k*xc*(cos phi)))*
                 (mkPolar 1.0 (k*yc*(sin phi)))*
                 ((sinc (k*ly*(sin phi)/2.0)) :+ 0.0) in

         let r = ( cos(phi) :+ k*lx )*((cos (theta - phi))/lx :+ 0.0) in i * r

     else if (xg==2*nx-1) then

         let i = (mkPolar ly (k*xc*cos(phi)))*
                 (mkPolar 1.0 (k*yc*sin(phi)))*
                 ((sinc (k*ly*sin(phi)/2.0)) :+ 0.0) in

         let r = ( cos(phi) :+ (- k*lx))*((cos (theta - phi))/lx :+ 0.0) in i * r

     else if ( (yg==1) && (xg<nx) ) then

         let i = (mkPolar lx (k*yc*sin(phi)))*
                 (mkPolar 1.0 (k*xc*cos(phi)))*
                 ((sinc (k*lx*(cos phi)/2.0)) :+ 0.0) in

         let r = ( (- sin phi) :+ k*ly )*((cos(theta - phi))/ly :+ 0.0) in i * r

     else if ( (yg==ny) && (xg<nx) ) then

         let i = (mkPolar lx (k*yc*sin(phi)))*
                 (mkPolar 1.0 (k*xc*cos(phi)))*
                 ((sinc (k*lx*(cos phi)/2.0)) :+ 0.0) in

         let r = ( (- sin phi) :+ (- k*ly) )*((cos(theta - phi)/ly) :+ 0.0) in i * r

     else 0.0 :+ 0.0

-- | Auxiliary function for oracle /A/.
build_circuit
convertband :: Int -> Int -> Int -> Int -> Int
convertband y b nx ny =
 let nedges = (nx - 1)*ny + nx*(ny - 1) in
 let (ex,ey) = itoxy y nx ny in
 let x = if ( (ex < nx) && (ey /= 1) ) then
           case b of
                1 -> y-2*nx+1
                2 -> y-nx
                3 -> y-nx+1
                5 -> y
                7 -> y+nx-1
                8 -> y+nx
                9 -> y+2*nx-1
                _ -> -1
         else if ( (ex < nx) && (ey == 1) ) then
           case b of
                5 -> y
                7 -> y+nx-1
                8 -> y+nx
                9 -> y+2*nx-1
                _ -> -1
         else if ( (ex >= nx) && (ex /= nx) && (ex /= 2*nx-1) ) then
           case b of
                    2 -> y-nx
                    3 -> y-nx+1
                    4 -> y-1
                    5 -> y
                    6 -> y+1
                    7 -> y+nx-1
                    8 -> y+nx
                    _ -> -1
         else if ( (ex >= nx) && (ex == nx) ) then
                 case b of
                    3 -> y-nx+1
                    5 -> y
                    6 -> y+1
                    8 -> y+nx
                    _ -> -1
         else if ( (ex >= nx) && (ex == 2*nx-1) ) then
                case b of
                    2 -> y-nx
                    4 -> y-1
                    5 -> y
                    7 -> y+nx-1
                    _ -> -1
         else -1
 in if ( (x < 1) || (x > nedges) ) then -1 else x

-- | Oracle /A/. It is equivalent to the Matlab function
-- 'getBandNodeValues'.
--
-- 'getNodeValuesMoreOutputs v b ...'  outputs the node of the edge
-- connected to vertex v in band b, and a real number parameterized by
-- the 'BoolParam' parameter: the magnitude (PFalse) or the phase
-- (PTrue) of the complex value at the corresponding place in the matrix A.
build_circuit
getNodeValuesMoreOutputs ::
    Int -> Int -> Int -> Int -> [(Double,Double)] -> Double -> Double -> Double -> BoolParam
    -> Int -> (Int, Double)
getNodeValuesMoreOutputs v' b' nx ny scatteringnodes lx ly k argflag maxConnectivity =
   let maxC = getIntFromParam maxConnectivity in
   let b = getIntFromParam b' in
   let nedges = (nx - 1)*ny + nx*(ny - 1) in
   let flag = v' <= nedges in
   let v = (if flag then v' else (v' - nedges)) in
   let nodeDefault = (if flag then v + nedges else v) in
   let valueDefault = 0.0 in
   let indicesDefault = (-1, -1) in
   let isvalid = checkedge v scatteringnodes nx ny
   in
   if ( (not isvalid) && b == 5 ) then

     let indices = (if flag then (v, v + nedges) else (v + nedges, v)) in
     case argflag of
        PTrue  -> (nodeDefault, valueDefault)
        PFalse -> (nodeDefault, 1.0)

   else if ( (not isvalid) && b /= 5) then (nodeDefault, valueDefault)

   else if ((b > maxC + 2) || (b <= 0)) then (nodeDefault, valueDefault)

   else

   let x = convertband v b' nx ny in

   if (x == -1) then (nodeDefault, valueDefault)
   else

   let isvalid = checkedge x scatteringnodes nx ny in

   if isvalid then

     let ax = calcmatrixelement v x nx ny lx ly k in
     let (node, indices) = if flag then (x + nedges, (v, x + nedges))
                           else (x, (v+nedges,x)) in
     let value = case argflag of
                   PTrue -> atan2  (imagPart ax) (realPart ax)
                   PFalse -> magnitude ax
     in (node, value)

   else
     (nodeDefault, valueDefault)

-- | Auxiliary function for oracle /b/. The inputs are:
--
-- * /y/ :: 'Int' - Global edge index.  Note this is the unmarked /y/
-- coordinate, i.e. the coordinate without scattering regions removed;
--
-- * /nx/ :: 'Int' - Number of vertices left to right;
--
-- * /ny/ :: 'Int' - Number of vertices top to bottom;
--
-- * /lx/ :: 'Double' - Length of horizontal edges (distance between vertices in /x/ direction);
--
-- * /ly/ :: 'Double' - Length of vertical edges (distance between vertices in /y/ direction);
--
-- * /k/ :: 'Double' - Plane wave wavenumber;
--
-- * θ :: 'Double' - Direction of wave propagation;
--
-- * /E0/ :: 'Double' - Magnitude of incident plane wave.
--
-- The output is the magnitude of the electric field on edge /y/.
build_circuit
calcincidentfield :: --
	-- Inputs
		Int -- int y - Global edge index.  Note this is the unmarked y coordinate,
			-- i.e. the coordinate without scattering regions removed.
	-> Int -- int nx - Number of vertices left to right
	-> Int --  int ny - Number of vertices top to bottom
	-> Double -- float lx - Length of horizontal edges (distance between vertices in x direction)
	-> Double -- float ly - Length of vertical edges (distance between vertices in y direction)
	-> Double -- float k - Plane wave wavenumber
	-> Double -- float theta - Direction of wave propagation
	-> Double -- float E0 - Magnitude of incident plane wave
	-- Outputs
	-> Complex Double -- complex float e - Magnitude of electric field on edge y
calcincidentfield y nx ny lx ly k theta e0 =
        let (xg, yg) = itoxy y nx ny in
        --Determine whether edge is horizontal or vertical
        let isvertical = xg >= nx in
        let (xvalueTmp, yvalueTmp) = edgetoxy y nx ny in
        let xvalue = xvalueTmp * lx in
        let yvalue = yvalueTmp * ly in
        -- Convert x and y edge coordinates to x and y values and caluculate field
        if isvertical then
          mkPolar (-cos(theta)*e0) ( -k*(xvalue*cos(theta)+yvalue*sin(theta)))
        else
          mkPolar (sin(theta)*e0) ( -k*(xvalue*cos(theta)+yvalue*sin(theta)))

-- | Auxiliary function for oracle /b/.
build_circuit
getconnection :: Int -> Int -> Int -> Int -> Int -> Int
getconnection y i' nx ny maxConnectivity =
  let i = getIntFromParam i' in
  let maxC = getIntFromParam maxConnectivity in
  let (ex,ey) = itoxy y nx ny in
  let x = if ( (ex < nx) && (ey /= 1) ) then
            case i' of
                1 -> y-2*nx+1
                2 -> y-nx
                3 -> y-nx+1
                4 -> y
                5 -> y+nx-1
                6 -> y+nx
                7 -> y+2*nx-1
                _ -> -1
          else if ( (ex < nx) && (ey == 1) ) then
            case i' of
                1 -> y
                2 -> y+nx-1
                3 -> y+nx
                4 -> y+2*nx-1
                _ -> -1
          else if ( (ex >= nx) && (ex /= nx) && (ex /= 2*nx-1) ) then
             case i' of
                1 -> y-nx
                2 -> y-nx+1
                3 -> y-1
                4 -> y
                5 -> y+1
                6 -> y+nx-1
                7 -> y+nx
                _ -> -1
          else if ( (ex >= nx) && (ex == nx) ) then
             case i' of
                1 -> y-nx+1
                2 -> y
                3 -> y+1
                4 -> y+nx
                _ -> -1
          else if ( (ex >= nx) && (ex == 2*nx-1) ) then
             case i' of
                1 -> y-nx
                2 -> y-1
                3 -> y
                4 -> y+nx-1
                _ -> -1
          else -1
  in
  if (i > maxC) then -1
  else if (x > nx*(ny-1)+ny*(nx-1)) then -1
  else x

-- | Auxiliary function to @'template_paramZero'@.
build_circuit
local_loop_with_index_aux :: Int -> Int -> t -> (Int -> t -> t) -> t
local_loop_with_index_aux i n x f =
   case paramMinus n i of
     0 -> x
     _ -> local_loop_with_index_aux (paramSucc i) n (f i x) f

-- | Local version of @'loop_with_index'@, for lifting.
build_circuit
local_loop_with_index :: Int -> t -> (Int -> t -> t) -> t
local_loop_with_index n x f = local_loop_with_index_aux paramZero n x f

-- | Oracle /b/.
build_circuit
getKnownWeights :: Int -> Int -> Int -> [(Double,Double)] -> Double -> Double -> Double -> Double -> Double -> Int -> Complex Double
getKnownWeights y nx ny scatteringnodes lx ly k theta e0 maxConnectivity =
   let makeConnections i connections = let x = getconnection y (paramSucc i) nx ny maxConnectivity in
                                       let t = not $ checkedge x scatteringnodes nx ny in
                                       (x,t):connections
   in
   let calcTang b (c,t) = if t
                          then let matElt = calcmatrixelement y c nx ny lx ly k in
                               let incField = calcincidentfield c nx ny lx ly k theta e0
                               in b - matElt * incField
                          else b
   in
   let connections = local_loop_with_index maxConnectivity [] makeConnections in
   if (not $ checkedge y scatteringnodes nx ny) then 0.0 :+ 0.0
   else foldl calcTang (0.0 :+ 0.0) connections

----------------------------------------------------------------------
----------------------------------------------------------------------
-- Testing functions

test_template_sinc = do
      f <- template_sinc
      r <- qinit (0 :: FDouble)
      f r
      return ()

test_template_itoxy = do
      f <- template_itoxy
      x <- qinit (0 :: FSignedInt)
      y <- qinit (0 :: FSignedInt)
      z <- qinit (0 :: FSignedInt)
      g <- f x
      h <- g y
      k <- h z
      return ()

test_template_edgetoxy = do
      f <- template_edgetoxy
      x <- qinit (0 :: FSignedInt)
      y <- qinit (0 :: FSignedInt)
      z <- qinit (0 :: FSignedInt)
      g <- f x
      h <- g y
      k <- h z
      return ()

test_template_calcRweights = do
      f <- template_calcRweights
      n1 <- qinit (0 :: FSignedInt)
      n2 <- qinit (0 :: FSignedInt)
      n3 <- qinit (0 :: FSignedInt)
      x1 <- qinit (0 :: FDouble)
      x2 <- qinit (0 :: FDouble)
      x3 <- qinit (0 :: FDouble)
      x4 <- qinit (0 :: FDouble)
      x5 <- qinit (0 :: FDouble)
      f1 <- f n1
      f2 <- f1 n2
      f3 <- f2 n3
      g1 <- f3 x1
      g2 <- g1 x2
      g3 <- g2 x3
      g4 <- g3 x4
      g5 <- g4 x5
      return ()

test_template_calcincidentfield = do
   y' <- qinit (0 :: FSignedInt)
   nx' <- qinit (0 :: FSignedInt)
   ny' <- qinit (0 :: FSignedInt)
   lx' <- qinit (0 :: FDouble)
   ly' <- qinit (0 :: FDouble)
   k' <- qinit (0 :: FDouble)
   theta' <- qinit (0 :: FDouble)
   e0' <- qinit (0 :: FDouble)
   f <- template_calcincidentfield
   f1 <- f y'
   f2 <- f1 nx'
   f3 <- f2 ny'
   f4 <- f3 lx'
   f5 <- f4 ly'
   f6 <- f5 k'
   f7 <- f6 theta'
   f8 <- f7 e0'
   return ()

test_template_calcmatrixelement = do
   y1 <- qinit (0 :: FSignedInt)
   y2 <- qinit (0 :: FSignedInt)
   nx <- qinit (0 :: FSignedInt)
   ny <- qinit (0 :: FSignedInt)
   lx <- qinit (0 :: FDouble)
   ly <- qinit (0 :: FDouble)
   k  <- qinit (0 :: FDouble)
   f <- template_calcmatrixelement
   f1 <- f y1
   f2 <- f1 y2
   f3 <- f2 nx
   f4 <- f3 ny
   f5 <- f4 lx
   f6 <- f5 ly
   f7 <- f6 k
   return ()

test_template_getconnection = do
   y <- qinit (0 :: FSignedInt)
   nx <- qinit (0 :: FSignedInt)
   ny <- qinit (0 :: FSignedInt)
   f <- template_getconnection
   f1 <- f y
   f2 <- f1 6
   f3 <- f2 nx
   f4 <- f3 ny
   f5 <- f4 7
   return ()

test_template_checkedge = do
   y <- qinit (0 :: FSignedInt)
   nx <- qinit (0 :: FSignedInt)
   ny <- qinit (0 :: FSignedInt)
   s <- qinit [(0 :: FDouble,0 :: FDouble),(0 :: FDouble,0 :: FDouble)]
   f <- template_checkedge
   f1 <- f y
   f2 <- f1 s
   f3 <- f2 nx
   f4 <- f3 ny
   return ()


-- An abstract data type for performing multiple linear regression,
-- using HMatrix as the backend for matrix computations.
--
-- types:
--   Matrix a (2D-array)
--   Vector a (vector from the standard vector package)
--
-- Both types are dense, immutable, and strict in all elements and are
-- manipulated as whole blocks.
--
-- matrix product is a <> b, where a and b are matrices
-- matrix-vector product is a #> x, where a is a matrix and x a vector
-- vector-vector dot product is x <.> y
-- the (conjugate) transpose is trans a, where a is a matrix or vector
-- the general linear system solver is <\>
--
-- to construct a 4 by 3 matrix:
-- let a = (4><3) [1,2,3,4,5,6,7,8,9,10,11,12] :: Matrix R
--
-- to construct a vector of double-precision floating points:
-- let x = vector [1,2,3]
--
-- to multiply vectors:
-- x * y
--
-- to add two vectors elementwise:
-- x + y
--
-- to compute vector dot product:
-- x <.> y
--
-- to perform matrix-vector product:
-- m  #> x
--
-- to perform matrix-matrix product
-- m <> n
--
-- to sum elements of a vector or matrix:
-- sumelements m
--

module Regression.HMatrix (
    model_features,
    model_outputs,
    model_weights,
    model_predictions,
    model_rss,
    create_model,
    create_features,
    create_weights,
    predict,
    powers,
    normalize,
    newtons_method,
    newtons_method_with_norm,
    cross_validate,
    gradient_descent
) where

import Numeric.LinearAlgebra hiding (magnitude)
import Numeric.LinearAlgebra.HMatrix hiding (magnitude)
import Data.List hiding (transpose)
import Data.Maybe (fromJust)
import Debug.Trace (trace)
import Stat (range, mean, stdev)
import qualified Data.Vector as V

-- A model is:
-- * a feature matrix with n rows and d columns
-- * a vector of length n representing the observed output
-- * a weight vector containing the calculated vector of optimized weights
-- * the predicted outputs using the optimized weights
-- * the
data Model = MO {
    model_features    :: FeatureMatrix,
    model_outputs     :: FeatureVector,
    model_weights     :: WeightVector,
    model_predictions :: FeatureVector,
    model_rss         :: Double
} deriving (Show)

-- A two-dimensional matrix whose columns contain feature values.
data FeatureMatrix = FM {
    fm_name_indexes:: [(String, Int)],
    fm_values :: Matrix Double
} deriving (Show)

-- A vector of values for a named feature.
data FeatureVector = FV {
    fv_name :: String,
    fv_values :: Vector Double
}
instance Show FeatureVector where
    show (FV name values) =
        name ++ " = " ++ show (toList values)

type Feature a = (String, a -> Double)

type Output a = (String, a -> Double)

type Optimizer = FeatureMatrix -> FeatureVector -> WeightVector

-- A vector of weights.  Each weight corresponds to a feature in the
-- feature matrix used to calculate the weights.
data WeightVector = WV {
    wv_name_indexes :: [(String, Int)],
    wv_values :: Vector Double
}

instance Show WeightVector where
    show (WV name_indexes values) =
        concatMap showWeight (zip name_indexes (toList values))
        where showWeight ((name, i), value) = name ++ " = " ++ show value ++ "\n"

-- Creates a model from a list of records, a list of features, and the output
-- accessor function.  This computes the weight vector as well.  This function
-- does not scale the features nor does it add an intercept feature (a feature
-- whose values are all 1).
create_model :: [a] -> [Feature a] -> Output a -> Optimizer -> Model
create_model rows features (output_name, output) optimizer =
    let fmat         = create_features features rows
        nn           = length rows
        observations = FV output_name (fromList (Data.List.map output rows))
        weights      = optimizer fmat observations
        predictions  = predict fmat weights
        residuals    = rss observations predictions
    in MO fmat observations weights predictions residuals

-- Multiplies a feature matrix by a weights vector to obtain a prediction of
-- output.
predict :: FeatureMatrix -> WeightVector -> FeatureVector
predict (FM _ h) (WV _ w) = FV "" (h #> w)

-- Generates the feature matrix, usually denoted as H, with N rows and
-- D features where D is the length of the feature list.
create_features :: [Feature a] -> [a] -> FeatureMatrix
create_features hs inputs =
    let n = length inputs
        d = length hs
        names = (Data.List.map fst hs)
        name_indexes = Prelude.zip names [0..]
        dat = [h(row) | row <- inputs, (_,h) <- hs]
        h = (n><d) dat
    in FM name_indexes h

-- Computes the residual sum of squares from the observed output and
-- the predicted output.
rss :: (Num a) => FeatureVector -> FeatureVector -> Double
rss (FV _ v1) (FV _ v2) = rss' v1 v2

rss' :: Vector Double -> Vector Double -> Double
rss' v1 v2 = diff <.> diff
    where diff = v1 - v2

-- Takes a list of numbers and turns them into a vector of weights.
create_weights :: FeatureMatrix -> [Double] -> WeightVector
create_weights (FM name_indexes h) weights = WV name_indexes (vector weights)

newtons_method :: Optimizer
newtons_method h y = newtons_method_with_norm 0 h y

newtons_method_with_norm :: Matrix Double -> Optimizer
newtons_method_with_norm l (FM n h) (FV o y) = WV n w
    where w = newtons_method_with_norm' l h y

-- Performs newtons method, finding an estimate of the minimum.  This method
-- uses the formula:
-- w_hat = (H^t * H - lambda*I)^-1 * H^T * y
newtons_method_with_norm' :: Matrix Double -> Matrix Double -> Vector Double -> Vector Double
newtons_method_with_norm' l h y = ((inv (hth - mod)) <> th) #> y
    where th = tr h
          hth = th <> h
          (d,_) = size hth
          mod = l * (ident d) :: Matrix Double

-- Performs gradient descent, updating the weights until the residual
-- sum of squares is less than the epsilon value e, at which point the
-- weight matrix is returned.  n is the step size.
gradient_descent :: Double -> Double -> WeightVector -> Optimizer
gradient_descent e n (WV name w) (FM _ f) (FV _ o) =
    let ft      = tr f
        weights = gradient_descent' e n w f ft o
    in WV name weights

gradient_descent' :: Double -> Double -> Vector Double -> Matrix Double ->
    Matrix Double -> Vector Double -> Vector Double
gradient_descent' e n w h ht y =
    let grad = gradient h ht y w -- (-2H^t(y-Hw))
        grad_len = magnitude grad  -- grad RSS(w) == ||2H^t(y-HW)||
    --in if grad_len < e
    in if (trace ("gradient = " ++ show grad_len) grad_len) < e
        then w
        else let delta = cmap (*(-n)) grad -- (2nH^t(y-Hw))
                 w' = w + delta
             in gradient_descent' e n w' h ht y

-- Compute the magnitude of the given vector.
magnitude :: Vector Double -> Double
magnitude vec = sumElements $ cmap (\y -> y^2) vec

-- Calculates the gradient of the residual sum of squares (-2H^t(y-Hw)).
-- This is used to compute the magnitude of the gradient, to see if the
-- function is minimized.  It is also used to update the weights of the
-- features.
gradient :: Matrix Double -> Matrix Double -> Vector Double ->
    Vector Double -> Vector Double
gradient h ht y w =
    let yhat = h #> w
        err = y - yhat
        prod = ht #> err
        grad = cmap (*(-2)) prod
    in grad

-- Given a feature, computes polynomial powers of that feature from
-- 1 (the original feature) up to and including n.
powers :: Feature a -> Int -> [Feature a]
powers (name, f) n = Data.List.map  (\i -> (name ++ (show i), (\a -> (f a)^i))) [1..n]

-- given a list of values and a feature, return a normalized version of that
-- feature where the mean of the feature is subtracted from each feature and
-- the result is divided by the standard deviation of the feature.
normalize :: [a] -> (a -> Double) -> a -> Double
normalize xs f = (/sdxs) . (+(-meanxs)) . f
    where meanxs = mean (Data.List.map f xs)
          sdxs = stdev (Data.List.map f xs)

-- Compute the average k-fold cross-validation error for a given L2 penalty and
-- a given split size k.
-- Note that the rows should be shuffled before they are provided.
cross_validate :: Int -> [a] -> [Feature a] -> Output a  -> Optimizer -> Double
cross_validate k rows features o@(output_name, output) optimizer = kmean
    where ksplits         = ksplit k rows
          training_data   = map fst ksplits
          valid_data      = map snd ksplits
          trained_models  = map (\rows -> create_model rows features o optimizer) training_data
          trained_weights = map model_weights trained_models
          valid_features  = map (create_features features) valid_data
          predictions     = zipWith (\h w -> predict h w) valid_features trained_weights
          observations    = map (\dat -> FV output_name (fromList (Data.List.map output dat))) valid_data
          valid_rss       = zipWith (\o p -> rss o p) observations predictions
          kmean           = mean valid_rss

-- prepare splits of the data for k-means where k is the number of
-- clusters and the number of values in the test set (the second list
-- in the pair) is n/k where n is the size of the input list.
ksplit :: Int -> [a] -> [([a], [a])]
ksplit k xs = map (ksplit' k xs) [0..k-1]

-- prepare a k split ([a], [a]) for the chunk at index i with list xs
-- and with number of chunks k
ksplit' :: Int -> [a] -> Int -> ([a], [a])
ksplit' k xs i = ((take (sz*i) xs) ++ (drop (sz*(i+1)) xs), take sz (drop (sz*i) xs))
    where sz = length xs `div` k


{-# OPTIONS_GHC -fplugin GHC.TypeLits.KnownNat.Solver #-}
{-# OPTIONS_GHC -fno-warn-missing-signatures #-}
{-# OPTIONS_GHC -fno-warn-incomplete-patterns #-}

{-# LANGUAGE CPP                 #-}
{-# LANGUAGE ConstraintKinds     #-}
{-# LANGUAGE DataKinds           #-}
{-# LANGUAGE GADTs               #-}
{-# LANGUAGE KindSignatures      #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TemplateHaskell     #-}
{-# LANGUAGE TypeApplications    #-}
{-# LANGUAGE TypeOperators       #-}

module Test.Grenade.Layers.Mul where

import           System.Random.MWC

import           Data.Proxy

import           GHC.TypeLits

import           Grenade.Core
import           Grenade.Layers.Mul
import qualified Numeric.LinearAlgebra        as LA
import           Numeric.LinearAlgebra.Static (L, R)
import qualified Numeric.LinearAlgebra.Static as H

import           Hedgehog
import qualified Hedgehog.Range               as Range

import           Test.Hedgehog.Compat
import           Test.Hedgehog.Hmatrix

import           Data.Either
import           Data.Serialize

prop_mul_scalar_one_does_nothing = property $ do
  height   :: Int <- forAll $ choose 2 100
  width    :: Int <- forAll $ choose 2 100
  channels :: Int <- forAll $ choose 2 100

  case (someNatVal (fromIntegral height), someNatVal (fromIntegral width), someNatVal (fromIntegral channels)) of
    (Just (SomeNat (Proxy :: Proxy h)), Just (SomeNat (Proxy :: Proxy w)), Just (SomeNat (Proxy :: Proxy c))) -> do
      input :: S ('D3 h w c) <- forAll genOfShape

      let layer    = initMul :: Mul 1 1 1
          S3D out  = snd $ runForwards layer input :: S ('D3 h w c)
          S3D inp' = input :: S ('D3 h w c)

      H.extract inp' === H.extract out

prop_mul_random_scalar_as_expected = property $ do
  height   :: Int <- forAll $ choose 2 100
  width    :: Int <- forAll $ choose 2 100
  channels :: Int <- forAll $ choose 2 100

  scalar :: R 1 <- forAll randomVector
  let scale = H.extract scalar LA.! 0

  case (someNatVal (fromIntegral height), someNatVal (fromIntegral width), someNatVal (fromIntegral channels)) of
    (Just (SomeNat (Proxy :: Proxy h)), Just (SomeNat (Proxy :: Proxy w)), Just (SomeNat (Proxy :: Proxy c))) -> do
      input :: S ('D3 h w c) <- forAll genOfShape

      let layer    = Mul scalar :: Mul 1 1 1
          S3D out  = snd $ runForwards layer input :: S ('D3 h w c)
          inp' = (\(S3D x) -> H.dmmap (* scale) x) input :: L (h * c) w

      H.extract inp' === H.extract out

prop_mul_has_show = withTests 1 $ property $ do
  gen <- evalIO create
  mul :: Mul 1 1 1 <- evalIO $ createRandomWith UniformInit gen
  show mul `seq` success

prop_mul_can_update = withTests 1 $ property $ do
  gen <- evalIO create
  mul :: Mul 1 1 1 <- evalIO $ createRandomWith UniformInit gen
  runUpdate defSGD mul () `seq` success
  runUpdate defAdam mul () `seq` success

prop_mul_can_be_used_with_batch = withTests 1 $ property $ do
  reduceGradient @(Mul 1 1 1) [()] `seq` success

prop_mul_can_be_serialized = withTests 1 $ property $ do
  s <- forAll $ genRealNum $ Range.constant 1 10
  let mul :: Mul 1 1 1 = Mul (H.fromList [s])
      bs = encode mul
      dec = decode bs :: Either String (Mul 1 1 1)
  assert $ isRight dec
  let Right mul' = dec
      (Mul scalarMat)  = mul
      (Mul scalarMat') = mul'
  (H.extract scalarMat) === (H.extract scalarMat')

prop_mul_fails_serialize_with_incorrect_size = withTests 1 $ property $ do
  s <- forAll $ genRealNum $ Range.constant 1 10
  let mul :: Mul 1 1 1 = Mul (H.fromList [s])
  let bs = encode mul
  let dec = decode bs :: Either String (Mul 1 2 1)
  assert $ isLeft dec

tests :: IO Bool
tests = checkParallel $$(discover)


{-# LANGUAGE ForeignFunctionInterface, GeneralizedNewtypeDeriving #-}

module Numerical.HBLAS.BLAS.FFI.Level3  where

import Foreign.Ptr
import Foreign()
import Foreign.C.Types
import Data.Complex
import Numerical.HBLAS.BLAS.FFI

--------------------------------------------------------------------------------
------------------------------ | BLAS LEVEL 3 ROUTINES
--------------------------------------------------------------------------------
-----------------------  |  Level 3 ops are faster than Levels 1 or 2
--------------------------------------------------------------------------------

--void cblas_sgemm(const enum CBLAS_ORDER Order, const enum CBLAS_TRANSPOSE TransA, const enum CBLAS_TRANSPOSE TransB, const blasint M, const blasint N, const blasint K,
--         const float alpha, const float *A, const blasint lda, const float *B, const blasint ldb, const float beta, float *C, const blasint ldc);

--void cblas_dgemm(const enum CBLAS_ORDER Order, const enum CBLAS_TRANSPOSE TransA, const enum CBLAS_TRANSPOSE TransB, const blasint M, const blasint N, const blasint K,
--         const double alpha, const double *A, const blasint lda, const double *B, const blasint ldb, const double beta, double *C, const blasint ldc);
--void cblas_cgemm(const enum CBLAS_ORDER Order, const enum CBLAS_TRANSPOSE TransA, const enum CBLAS_TRANSPOSE TransB, const blasint M, const blasint N, const blasint K,
--         const float *alpha, const float *A, const blasint lda, const float *B, const blasint ldb, const float *beta, float *C, const blasint ldc);
--void cblas_zgemm(const enum CBLAS_ORDER Order, const enum CBLAS_TRANSPOSE TransA, const enum CBLAS_TRANSPOSE TransB, const blasint M, const blasint N, const blasint K,
--         const double *alpha, const double *A, const blasint lda, const double *B, const blasint ldb, const double *beta, double *C, const blasint ldc);

-- |  Matrix mult for general dense matrices
type GemmFunFFI scale el = CBLAS_ORDERT ->   CBLAS_TRANSPOSET -> CBLAS_TRANSPOSET->
        CInt -> CInt -> CInt -> {- scal A * B -} scale  -> {- Matrix A-} Ptr el  -> CInt -> {- B -}  Ptr el -> CInt->
            scale -> {- C -}  Ptr el -> CInt -> IO ()

{- C := alpha*op( A )*op( B ) + beta*C ,  -}

-- matrix mult!
foreign import ccall unsafe "cblas_sgemm"
    cblas_sgemm_unsafe :: GemmFunFFI Float Float

foreign import ccall unsafe "cblas_dgemm"
    cblas_dgemm_unsafe :: GemmFunFFI Double Double

foreign import ccall unsafe "cblas_cgemm"
    cblas_cgemm_unsafe :: GemmFunFFI (Ptr(Complex Float)) (Complex Float)

foreign import ccall unsafe "cblas_zgemm"
    cblas_zgemm_unsafe :: GemmFunFFI (Ptr (Complex Double)) (Complex Double)

-- safe ffi variant for large inputs
foreign import ccall "cblas_sgemm"
    cblas_sgemm_safe :: GemmFunFFI Float Float

foreign import ccall "cblas_dgemm"
    cblas_dgemm_safe :: GemmFunFFI Double Double

foreign import ccall "cblas_cgemm"
    cblas_cgemm_safe :: GemmFunFFI (Ptr(Complex Float)) (Complex Float)

foreign import ccall "cblas_zgemm"
    cblas_zgemm_safe :: GemmFunFFI (Ptr (Complex Double)) (Complex Double)

-----------------------------------------
----- |  Matrix mult for Symmetric Matrices
-----------------------------------------

--void cblas_ssymm(const enum CBLAS_ORDER Order, const enum CBLAS_SIDE Side, const enum CBLAS_UPLO Uplo, const blasint M, const blasint N,
--                 const float alpha, const float *A, const blasint lda, const float *B, const blasint ldb, const float beta, float *C, const blasint ldc);
--void cblas_dsymm(const enum CBLAS_ORDER Order, const enum CBLAS_SIDE Side, const enum CBLAS_UPLO Uplo, const blasint M, const blasint N,
--                 const double alpha, const double *A, const blasint lda, const double *B, const blasint ldb, const double beta, double *C, const blasint ldc);
--void cblas_csymm(const enum CBLAS_ORDER Order, const enum CBLAS_SIDE Side, const enum CBLAS_UPLO Uplo, const blasint M, const blasint N,
--                 const float *alpha, const float *A, const blasint lda, const float *B, const blasint ldb, const float *beta, float *C, const blasint ldc);
--void cblas_zsymm(const enum CBLAS_ORDER Order, const enum CBLAS_SIDE Side, const enum CBLAS_UPLO Uplo, const blasint M, const blasint N,
--                 const double *alpha, const double *A, const blasint lda, const double *B, const blasint ldb, const double *beta, double *C, const blasint ldc);

type SymmFunFFI scale el = CBLAS_ORDERT -> CBLAS_SIDET -> CBLAS_UPLOT ->
     CInt->CInt -> scale -> Ptr el -> CInt -> Ptr el -> CInt -> scale ->Ptr el -> CInt -> IO ()

foreign import ccall unsafe "cblas_ssymm"
    cblas_ssymm_unsafe :: SymmFunFFI Float Float

foreign import ccall unsafe "cblas_dsymm"
    cblas_dsymm_unsafe :: SymmFunFFI Double Double

foreign import ccall unsafe "cblas_csymm"
    cblas_csymm_unsafe :: SymmFunFFI (Ptr (Complex Float )) (Complex Float)

foreign import ccall unsafe "cblas_zsymm"
    cblas_zsymm_unsafe :: SymmFunFFI (Ptr (Complex Double)) (Complex Double)

-- safe ffi variant,
foreign import ccall  "cblas_ssymm"
    cblas_ssymm_safe :: SymmFunFFI Float Float

foreign import ccall  "cblas_dsymm"
    cblas_dsymm_safe :: SymmFunFFI Double Double

foreign import ccall  "cblas_csymm"
    cblas_csymm_safe :: SymmFunFFI (Ptr (Complex Float )) (Complex Float)

foreign import ccall  "cblas_zsymm"
    cblas_zsymm_safe :: SymmFunFFI (Ptr (Complex Double)) (Complex Double)

-----------------------------------
--- |  symmetric rank k  matrix update, C := alpha*A*A' + beta*C
--- or C = alpha*A'*A + beta*C
------------------------------------

--void cblas_ssyrk(const enum CBLAS_ORDER Order, const enum CBLAS_UPLO Uplo, const enum CBLAS_TRANSPOSE Trans,
--         const blasint N, const blasint K, const float alpha, const float *A, const blasint lda, const float beta, float *C, const blasint ldc);
--void cblas_dsyrk(const enum CBLAS_ORDER Order, const enum CBLAS_UPLO Uplo, const enum CBLAS_TRANSPOSE Trans,
--         const blasint N, const blasint K, const double alpha, const double *A, const blasint lda, const double beta, double *C, const blasint ldc);
--void cblas_csyrk(const enum CBLAS_ORDER Order, const enum CBLAS_UPLO Uplo, const enum CBLAS_TRANSPOSE Trans,
--         const blasint N, const blasint K, const float *alpha, const float *A, const blasint lda, const float *beta, float *C, const blasint ldc);
--void cblas_zsyrk(const enum CBLAS_ORDER Order, const enum CBLAS_UPLO Uplo, const enum CBLAS_TRANSPOSE Trans,
--         const blasint N, const blasint K, const double *alpha, const double *A, const blasint lda, const double *beta, double *C, const blasint ldc);

type SyrkFunFFI scale el = CBLAS_ORDERT -> CBLAS_UPLOT -> CBLAS_TRANSPOSET ->
     CInt -> CInt  -> scale -> Ptr el -> CInt -> scale -> Ptr el -> CInt -> IO ()
foreign import ccall unsafe "cblas_ssyrk"
    cblas_ssyrk_unsafe :: SyrkFunFFI Float Float
foreign import ccall unsafe "cblas_dsyrk"
    cblas_dsyrk_unsafe :: SyrkFunFFI Double Double
foreign import ccall unsafe "cblas_csyrk"
    cblas_csyrk_unsafe :: SyrkFunFFI (Ptr(Complex Float)) (Complex Float)
foreign import ccall unsafe "cblas_zsyrk"
    cblas_zsyrk_unsafe :: SyrkFunFFI (Ptr(Complex Double)) (Complex Double)

foreign import ccall safe "cblas_ssyrk"
    cblas_ssyrk_safe :: SyrkFunFFI Float Float
foreign import ccall safe "cblas_dsyrk"
    cblas_dsyrk_safe :: SyrkFunFFI Double Double
foreign import ccall safe "cblas_csyrk"
    cblas_csyrk_safe :: SyrkFunFFI (Ptr(Complex Float)) (Complex Float)
foreign import ccall safe "cblas_zsyrk"
    cblas_zsyrk_safe :: SyrkFunFFI (Ptr(Complex Double)) (Complex Double)
----------------------
----- | Symmetric Rank 2k matrix update, C= alpha* A*B' + alpha* B*A' + beta * C
----- or C= alpha* A'*B + alpha* B'*A + beta * C
-------------------

--void cblas_ssyr2k(const enum CBLAS_ORDER Order, const enum CBLAS_UPLO Uplo, const enum CBLAS_TRANSPOSE Trans,
--          const blasint N, const blasint K, const float alpha, const float *A, const blasint lda, const float *B, const blasint ldb, const float beta, float *C, const blasint ldc);
--void cblas_dsyr2k(const enum CBLAS_ORDER Order, const enum CBLAS_UPLO Uplo, const enum CBLAS_TRANSPOSE Trans,
--          const blasint N, const blasint K, const double alpha, const double *A, const blasint lda, const double *B, const blasint ldb, const double beta, double *C, const blasint ldc);
--void cblas_csyr2k(const enum CBLAS_ORDER Order, const enum CBLAS_UPLO Uplo, const enum CBLAS_TRANSPOSE Trans,
--          const blasint N, const blasint K, const float *alpha, const float *A, const blasint lda, const float *B, const blasint ldb, const float *beta, float *C, const blasint ldc);
--void cblas_zsyr2k(const enum CBLAS_ORDER Order, const enum CBLAS_UPLO Uplo, const enum CBLAS_TRANSPOSE Trans,
         --const blasint N, const blasint K, const double *alpha, const double *A, const blasint lda, const double *B, const blasint ldb, const double *beta, double *C, const blasint ldc);

type Syr2kFunFFI scale el = CBLAS_ORDERT -> CBLAS_UPLOT -> CBLAS_TRANSPOSET  ->
     CInt->CInt -> scale -> Ptr el -> CInt -> Ptr el -> CInt ->
     scale ->Ptr el -> CInt -> IO ()

foreign  import ccall unsafe "cblas_ssyr2k"
    cblas_ssyr2k_unsafe :: Syr2kFunFFI Float Float
foreign import ccall unsafe "cblas_dsyr2k"
    cblas_dsyr2k_unsafe :: Syr2kFunFFI Double Double
foreign  import ccall unsafe "cblas_csyr2k"
    cblas_csyr2k_unsafe :: Syr2kFunFFI (Ptr (Complex Float)) (Complex Float)
foreign  import ccall unsafe "cblas_zsyr2k"
    cblas_zsyr2k_unsafe :: Syr2kFunFFI (Ptr (Complex Double)) (Complex Double)

foreign  import ccall safe "cblas_ssyr2k"
    cblas_ssyr2k_safe :: Syr2kFunFFI Float Float
foreign import ccall safe "cblas_dsyr2k"
    cblas_dsyr2k_safe :: Syr2kFunFFI Double Double
foreign  import ccall safe "cblas_csyr2k"
    cblas_csyr2k_safe :: Syr2kFunFFI (Ptr (Complex Float)) (Complex Float)
foreign  import ccall safe "cblas_zsyr2k"
    cblas_zsyr2k_safe :: Syr2kFunFFI (Ptr (Complex Double)) (Complex Double)

-------------------------------
--------  |  matrix matrix product for triangular matrices
------------------------------

type TrmmFunFFI scale el = CBLAS_ORDERT -> CBLAS_SIDET -> CBLAS_UPLOT -> CBLAS_TRANSPOSET -> CBLAS_DIAGT ->
     CInt -> CInt -> scale -> Ptr el -> CInt -> Ptr el -> CInt -> IO ()
foreign  import ccall unsafe "cblas_strmm"
    cblas_strmm_unsafe :: TrmmFunFFI Float Float
foreign  import ccall unsafe "cblas_dtrmm"
    cblas_dtrmm_unsafe :: TrmmFunFFI Double Double
foreign  import ccall unsafe "cblas_ctrmm"
    cblas_ctrmm_unsafe :: TrmmFunFFI (Ptr (Complex Float )) (Complex Float)
foreign  import ccall unsafe "cblas_ztrmm"
    cblas_ztrmm_unsafe :: TrmmFunFFI (Ptr (Complex Double )) (Complex Double)

foreign  import ccall safe "cblas_strmm"
    cblas_strmm_safe :: TrmmFunFFI Float Float
foreign  import ccall safe "cblas_dtrmm"
    cblas_dtrmm_safe :: TrmmFunFFI Double Double
foreign  import ccall safe "cblas_ctrmm"
    cblas_ctrmm_safe :: TrmmFunFFI (Ptr (Complex Float )) (Complex Float)
foreign  import ccall safe "cblas_ztrmm"
    cblas_ztrmm_safe :: TrmmFunFFI (Ptr (Complex Double )) (Complex Double)
--void cblas_strmm(  enum CBLAS_ORDER Order,   enum CBLAS_SIDE Side,   enum CBLAS_UPLO Uplo,   enum CBLAS_TRANSPOSE TransA,
--                   enum CBLAS_DIAG Diag,   CInt M,   CInt N,   Float alpha,   Float *A,   CInt lda, Float *B,   CInt ldb);
--void cblas_dtrmm(  enum CBLAS_ORDER Order,   enum CBLAS_SIDE Side,   enum CBLAS_UPLO Uplo,   enum CBLAS_TRANSPOSE TransA,
--                   enum CBLAS_DIAG Diag,   CInt M,   CInt N,   Double alpha,   Double *A,   CInt lda, Double *B,   CInt ldb);
--void cblas_ctrmm(  enum CBLAS_ORDER Order,   enum CBLAS_SIDE Side,   enum CBLAS_UPLO Uplo,   enum CBLAS_TRANSPOSE TransA,
--                   enum CBLAS_DIAG Diag,   CInt M,   CInt N,   Float *alpha,   Float *A,   CInt lda, Float *B,   CInt ldb);
--void cblas_ztrmm(  enum CBLAS_ORDER Order,   enum CBLAS_SIDE Side,   enum CBLAS_UPLO Uplo,   enum CBLAS_TRANSPOSE TransA,
--                   enum CBLAS_DIAG Diag,   CInt M,   CInt N,   Double *alpha,   Double *A,   CInt lda, Double *B,   CInt ldb);

------------------------
--  |  triangular solvers
-----------------------

--
--TRSM solves  op(A)*X = alpha*B or  X*op(A) = alpha*B
--op(A) is one of op(A) = A, or op(A) = A', or op(A) = conjg(A').
-- A is a unit, or non-unit, upper or lower triangular matrix
----
type TrsmFunFFI scale el = CBLAS_ORDERT -> CBLAS_SIDET -> CBLAS_UPLOT -> CBLAS_TRANSPOSET -> CBLAS_DIAGT ->
     CInt -> CInt -> scale -> Ptr el -> CInt -> Ptr el -> CInt -> IO ()
foreign  import ccall unsafe "cblas_strsm"
    cblas_strsm_unsafe :: TrsmFunFFI Float Float
foreign  import ccall unsafe "cblas_dtrsm"
    cblas_dtrsm_unsafe :: TrsmFunFFI Double Double
foreign  import ccall unsafe "cblas_ctrsm"
    cblas_ctrsm_unsafe :: TrsmFunFFI (Ptr (Complex Float )) (Complex Float)
foreign  import ccall unsafe "cblas_ztrsm"
    cblas_ztrsm_unsafe :: TrsmFunFFI (Ptr (Complex Double )) (Complex Double)

foreign  import ccall safe "cblas_strsm"
    cblas_strsm_safe :: TrsmFunFFI Float Float
foreign  import ccall safe "cblas_dtrsm"
    cblas_dtrsm_safe :: TrsmFunFFI Double Double
foreign  import ccall safe "cblas_ctrsm"
    cblas_ctrsm_safe :: TrsmFunFFI (Ptr (Complex Float )) (Complex Float)
foreign  import ccall safe "cblas_ztrsm"
    cblas_ztrsm_safe :: TrsmFunFFI (Ptr (Complex Double )) (Complex Double)
--void cblas_strsm(  enum CBLAS_ORDER Order,   enum CBLAS_SIDE Side,   enum CBLAS_UPLO Uplo,   enum CBLAS_TRANSPOSE TransA,
--                   enum CBLAS_DIAG Diag,   CInt M,   CInt N,   Float alpha,   Float *A,   CInt lda, Float *B,   CInt ldb);
--void cblas_dtrsm(  enum CBLAS_ORDER Order,   enum CBLAS_SIDE Side,   enum CBLAS_UPLO Uplo,   enum CBLAS_TRANSPOSE TransA,
--                   enum CBLAS_DIAG Diag,   CInt M,   CInt N,   Double alpha,   Double *A,   CInt lda, Double *B,   CInt ldb);
--void cblas_ctrsm(  enum CBLAS_ORDER Order,   enum CBLAS_SIDE Side,   enum CBLAS_UPLO Uplo,   enum CBLAS_TRANSPOSE TransA,
--                   enum CBLAS_DIAG Diag,   CInt M,   CInt N,   Float *alpha,   Float *A,   CInt lda, Float *B,   CInt ldb);
--void cblas_ztrsm(  enum CBLAS_ORDER Order,   enum CBLAS_SIDE Side,   enum CBLAS_UPLO Uplo,   enum CBLAS_TRANSPOSE TransA,
--                   enum CBLAS_DIAG Diag,   CInt M,   CInt N,   Double *alpha,   Double *A,   CInt lda, Double *B,   CInt ldb);

-------------------------
-- | hermitian matrix mult
------------------------

type HemmFunFFI  el = CBLAS_ORDERT -> CBLAS_SIDET -> CBLAS_UPLOT ->
     CInt -> CInt -> Ptr el -> Ptr el -> CInt -> Ptr el -> CInt -> Ptr el -> Ptr el -> CInt -> IO ()

foreign  import ccall unsafe "cblas_chemm"
    cblas_chemm_unsafe :: HemmFunFFI (Complex Float)
foreign  import ccall unsafe "cblas_zhemm"
    cblas_zhemm_unsafe :: HemmFunFFI  (Complex Double)

foreign  import ccall safe "cblas_chemm"
    cblas_chemm_safe :: HemmFunFFI (Complex Float)
foreign  import ccall safe "cblas_zhemm"
    cblas_zhemm_safe :: HemmFunFFI  (Complex Double)

--void cblas_chemm(  enum CBLAS_ORDER Order,   enum CBLAS_SIDE Side,   enum CBLAS_UPLO Uplo,   CInt M,   CInt N,
--                   Float *alpha,   Float *A,   CInt lda,   Float *B,   CInt ldb,   Float *beta, Float *C,   CInt ldc);
--void cblas_zhemm(  enum CBLAS_ORDER Order,   enum CBLAS_SIDE Side,   enum CBLAS_UPLO Uplo,   CInt M,   CInt N,
--                   Double *alpha,   Double *A,   CInt lda,   Double *B,   CInt ldb,   Double *beta, Double *C,   CInt ldc);

type HerkFunFFI scale el = CBLAS_ORDERT -> CBLAS_UPLOT -> CBLAS_TRANSPOSET ->
     CInt -> CInt -> scale -> Ptr el -> CInt -> scale -> Ptr el -> CInt -> IO ()

foreign  import ccall unsafe "cblas_cherk"
    cblas_cherk_unsafe :: HerkFunFFI Float (Complex Float)
foreign  import ccall unsafe "cblas_zherk"
    cblas_zherk_unsafe :: HerkFunFFI Double (Complex Double)

foreign  import ccall safe "cblas_cherk"
    cblas_cherk_safe :: HerkFunFFI Float (Complex Float)
foreign  import ccall safe "cblas_zherk"
    cblas_zherk_safe :: HerkFunFFI Double (Complex Double)
--void cblas_cherk(  enum CBLAS_ORDER Order,   enum CBLAS_UPLO Uplo,   enum CBLAS_TRANSPOSE Trans,   CInt N,   CInt K,
--                   Float alpha,   Float *A,   CInt lda,   Float beta, Float *C,   CInt ldc);
--void cblas_zherk(  enum CBLAS_ORDER Order,   enum CBLAS_UPLO Uplo,   enum CBLAS_TRANSPOSE Trans,   CInt N,   CInt K,
--                   Double alpha,   Double *A,   CInt lda,   Double beta, Double *C,   CInt ldc);

type Her2kFunFFI scale el = CBLAS_ORDERT -> CBLAS_UPLOT -> CBLAS_TRANSPOSET ->
     CInt -> CInt -> Ptr el  -> Ptr el -> CInt -> Ptr el -> CInt -> scale ->Ptr el -> CInt -> IO ()

foreign  import ccall unsafe "cblas_cher2k"
    cblas_cher2k_unsafe :: Her2kFunFFI  Float  (Complex Float)
foreign  import ccall unsafe "cblas_zher2k"
    cblas_zher2k_unsafe :: Her2kFunFFI  Double  (Complex Double)

foreign  import ccall safe "cblas_cher2k"
    cblas_cher2k_safe :: Her2kFunFFI Float (Complex Float)
foreign  import ccall safe "cblas_zher2k"
    cblas_zher2k_safe :: Her2kFunFFI Double (Complex Double)
--void cblas_cher2k(  enum CBLAS_ORDER Order,   enum CBLAS_UPLO Uplo,   enum CBLAS_TRANSPOSE Trans,   CInt N,   CInt K,
--                    Float *alpha,   Float *A,   CInt lda,   Float *B,   CInt ldb,   Float beta, Float *C,   CInt ldc);
--void cblas_zher2k(  enum CBLAS_ORDER Order,   enum CBLAS_UPLO Uplo,   enum CBLAS_TRANSPOSE Trans,   CInt N,   CInt K,
--                    Double *alpha,   Double *A,   CInt lda,   Double *B,   CInt ldb,   Double beta, Double *C,   CInt ldc);

----void cblas_xerbla(CInt p, char *rout, char *form, ...);


{-# LANGUAGE BangPatterns        #-}
{-# LANGUAGE DataKinds           #-}
{-# LANGUAGE FlexibleContexts    #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TupleSections       #-}
{-# LANGUAGE TypeFamilies        #-}
{-# LANGUAGE TypeOperators       #-}

module Main where
import           Debug.Trace
import           Control.Monad
import qualified Data.ByteString              as B
import qualified Data.ByteString.Lazy         as L
import           Data.List
import           Data.Serialize               (decodeLazy, encodeLazy)
import qualified Data.Vector.Storable         as V
import           Grenade
import           Numeric.LinearAlgebra        (maxIndex)
import qualified Numeric.LinearAlgebra.Static as S
import           System.Random
import           Graphics.Image              as I hiding (map)
import           Graphics.Image.Interface    as I hiding (map)
import           Graphics.Image.Processing   as I hiding (map)
import qualified Numeric.LinearAlgebra as NLA
import qualified Numeric.LinearAlgebra.Static as NLA
import qualified Numeric.LinearAlgebra.Data as NLA

readF::FilePath->IO (S ('D3 100 100 3))
readF path =do
     Right x<- (fmap (resize Bilinear Edge (100,100)))<$>readImage path::IO (Either String (Image VS RGB Double))
     let Just res=fromStorable $V.concatMap (\(PixelRGB r g b) -> V.fromList [r,g,b]) $ toVector x
     return res

scons a !b = V.cons a b

elemAt = (V.!)

toImg::S ('D3 100 100 3)->Image VS RGB Double
toImg x = fromVector (100,100) $groupRGB $ext x
          where groupRGB v= if V.null v then V.empty else (PixelRGB (v `elemAt` 0) (v `elemAt` 1) (v `elemAt` 2)) `scons` groupRGB (V.drop 3 v)
                ext::S ('D3 100 100 3)->V.Vector Double
                ext (S3D m) = NLA.flatten$NLA.extract m

 type FE
  = Network
    '[ Convolution 3 24 5 5 1 1,  Relu, Pooling 2 2 2 2
     , Reshape
     , FullyConnected 55296 1024, Logit
     , FullyConnected 1024 55296, Relu
     , Reshape
     , Deconvolution 24 3 6 6 2 2, Logit
     ]
    '[ 'D3 100 100 3, 'D3 96 96 24, 'D3 96 96 24, 'D3 48 48 24
     , 'D1 55296
     , 'D1 1024, 'D1 1024
     , 'D1 55296, 'D1 55296
     , 'D3 48 48 24
     , 'D3 100 100 3,'D3 100 100 3]

trainOne::LearningParameters->FE->S ('D3 100 100 3)->FE
trainOne rate !net i=
  train rate net i i

trainT::FE->Int->S ('D3 100 100 3)->FE
trainT !net x img=
    let lp = LearningParameters 0.0001 0.9 0.0005
    in  foldl' (\n i ->
      trace (show i) $trainOne (lp{ learningRate = learningRate lp *(0.999 ^ i)}) n img)
      net [0..x]

main' = do
  i<-readF "image/n11669921_50878.JPEG"
  writeImage "newimg.jpg" $ toImg i

main = do
  i<-readF "image/n11669921_50878.JPEG"
  x<-Prelude.read<$>getLine
  net<-randomNetwork::IO FE
  let net'=trainT net x i
  writeImage "newimg.jpg" $ toImg $runNet net' i


{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE FlexibleContexts #-}
{- |
This module provides a simple way to train
the transition matrix and initial probability vector
using simple patterns of state sequences.

You may create a trained model using semigroup combinators like this:

> let a = atom $ HMM.state 0
>     b = atom $ HMM.state 1
>     distr =
>        Distr.DiscreteTrained $ Map.fromList $
>        ('a', Vector.fromList [1,2]) :
>        ('b', Vector.fromList [4,3]) :
>        ('c', Vector.fromList [0,1]) :
>        []
> in  finish 2 distr $ replicate 5 $ replicate 10 a <> replicate 20 b
-}
module Math.HiddenMarkovModel.Pattern (
   T,
   atom,
   append,
   replicate,
   finish,
   ) where

import qualified Math.HiddenMarkovModel.Distribution as Distr
import qualified Math.HiddenMarkovModel as HMM
import Math.HiddenMarkovModel.Private (Trained(..))
import Math.HiddenMarkovModel.Distribution (State(State))

import qualified Numeric.LinearAlgebra.Algorithms as Algo
import qualified Numeric.Container as NC
import qualified Data.Packed.Vector as Vector
import Data.Packed.Matrix (Matrix)
import Data.Packed.Vector (Vector)

import qualified Data.Map as Map
import Data.Semigroup (Semigroup, (<>), stimes)

import Prelude hiding (replicate)

newtype T prob = Cons (Int -> (State, Matrix prob, State))

atom ::
   (NC.Container Vector prob) =>
   State -> T prob
atom s = Cons $ \n -> (s, NC.konst 0 (n,n), s)

instance (Algo.Field prob) => Semigroup (T prob) where
   (<>) = append
   stimes k = replicate $ fromIntegral k

infixl 5 `append`

append ::
   (NC.Container Vector prob) =>
   T prob -> T prob -> T prob
append (Cons f) (Cons g) =
   Cons $ \n ->
      case (f n, g n) of
         ((sai, ma, sao), (sbi, mb, sbo)) ->
            (sai, increment (sbi,sao) 1 $ NC.add ma mb, sbo)

replicate ::
   (NC.Container Vector prob) =>
   Int -> T prob -> T prob
replicate ki (Cons f) =
   Cons $ \n ->
      case f n of
         (si, m, so) ->
            let k = fromIntegral ki
            in  (si, increment (si,so) (k-1) $ NC.scale k m, so)

increment ::
   (NC.Container Vector a) =>
   (State, State) -> a -> Matrix a -> Matrix a
increment (State i, State j) x m  =  NC.accum m (+) [((i,j), x)]

finish ::
   (NC.Container Vector prob) =>
   Int -> tdistr -> T prob -> Trained tdistr prob
finish n tdistr (Cons f) =
   case f n of
      (State si, m, _so) ->
         Trained {
            trainedInitial = NC.assoc n 0 [(si,1)],
            trainedTransition = m,
            trainedDistribution = tdistr
         }

_example :: HMM.DiscreteTrained Double Char
_example =
   let a = atom $ HMM.state 0
       b = atom $ HMM.state 1
       distr =
          Distr.DiscreteTrained $ Map.fromList $
          ('a', Vector.fromList [1,2]) :
          ('b', Vector.fromList [4,3]) :
          ('c', Vector.fromList [0,1]) :
          []
   in  finish 2 distr $ replicate 5 $ replicate 10 a <> replicate 20 b


{-# LANGUAGE RankNTypes #-}

module Test.Grenade.Layers.Internal.Reference where

import           Grenade.Types
import           Numeric.LinearAlgebra
import qualified Numeric.LinearAlgebra.Static as H
import           GHC.TypeLits

im2col :: Int -> Int -> Int -> Int -> Matrix RealNum -> Matrix RealNum
im2col nrows ncols srows scols m =
  let starts = fittingStarts (rows m) nrows srows (cols m) ncols scols
  in  im2colFit starts nrows ncols m

vid2col :: Int -> Int -> Int -> Int -> Int -> Int -> [Matrix RealNum] -> Matrix RealNum
vid2col nrows ncols srows scols inputrows inputcols ms =
  let starts = fittingStarts inputrows nrows srows inputcols ncols scols
      subs   = fmap (im2colFit starts nrows ncols) ms
  in  foldl1 (|||) subs

im2colFit :: [(Int,Int)] -> Int -> Int -> Matrix RealNum -> Matrix RealNum
im2colFit starts nrows ncols m =
  let imRows = fmap (\start -> flatten $ subMatrix start (nrows, ncols) m) starts
  in  fromRows imRows

col2vid :: Int -> Int -> Int -> Int -> Int -> Int -> Matrix RealNum -> [Matrix RealNum]
col2vid nrows ncols srows scols drows dcols m =
  let starts = fittingStart (cols m) (nrows * ncols) (nrows * ncols)
      r      = rows m
      mats   = fmap (\s -> subMatrix (0,s) (r, nrows * ncols) m) starts
      colSts = fittingStarts drows nrows srows dcols ncols scols
  in  fmap (col2imfit colSts nrows ncols drows dcols) mats

col2im :: Int -> Int -> Int -> Int -> Int -> Int -> Matrix RealNum -> Matrix RealNum
col2im krows kcols srows scols drows dcols m =
  let rs       = map toList $ toColumns m
      rs'      = zip [0..] rs
      indicies = (\[a,b] -> (a,b)) <$> sequence [[0..(krows-1)], [0..(kcols-1)]]
      accums   = concatMap (\(offset, column) -> zipWith (comb offset) indicies column) rs'
  in accum (konst 0 (drows, dcols)) (+) accums
  where
    comb o (i, j) x = let w = (div (dcols - kcols) scols) + 1
                          (a, b) = divMod o w
                      in  ((i + srows * a, j + scols * b), x)

col2imfit :: [(Int,Int)] -> Int -> Int -> Int -> Int -> Matrix RealNum -> Matrix RealNum
col2imfit starts krows kcols drows dcols m =
  let indicies   = (\[a,b] -> (a,b)) <$> sequence [[0..(krows-1)], [0..(kcols-1)]]
      convs      = fmap (zip indicies . toList) . toRows $ m
      pairs      = zip convs starts
      accums     = concatMap (\(conv',(stx',sty')) -> fmap (\((ix,iy), val) -> ((ix + stx', iy + sty'), val)) conv') pairs
  in  accum (konst 0 (drows, dcols)) (+) accums

poolForward :: Int -> Int -> Int -> Int -> Int -> Int -> Matrix RealNum -> Matrix RealNum
poolForward nrows ncols srows scols outputRows outputCols m =
  let starts = fittingStarts (rows m) nrows srows (cols m) ncols scols
  in  poolForwardFit starts nrows ncols outputRows outputCols m

poolForwardList :: Functor f => Int -> Int -> Int -> Int -> Int -> Int -> Int -> Int -> f (Matrix RealNum) -> f (Matrix RealNum)
poolForwardList nrows ncols srows scols inRows inCols outputRows outputCols ms =
  let starts = fittingStarts inRows nrows srows inCols ncols scols
  in  poolForwardFit starts nrows ncols outputRows outputCols <$> ms

poolForwardFit :: [(Int,Int)] -> Int -> Int -> Int -> Int -> Matrix RealNum -> Matrix RealNum
poolForwardFit starts nrows ncols _ outputCols m =
  let els    = fmap (\start -> maxElement $ subMatrix start (nrows, ncols) m) starts
  in  matrix outputCols els

poolBackward :: Int -> Int -> Int -> Int -> Matrix RealNum -> Matrix RealNum -> Matrix RealNum
poolBackward krows kcols srows scols inputMatrix gradientMatrix =
  let inRows     = rows inputMatrix
      inCols     = cols inputMatrix
      starts     = fittingStarts inRows krows srows inCols kcols scols
  in  poolBackwardFit starts krows kcols inputMatrix gradientMatrix

poolBackwardList :: Functor f => Int -> Int -> Int -> Int -> Int -> Int -> f (Matrix RealNum, Matrix RealNum) -> f (Matrix RealNum)
poolBackwardList krows kcols srows scols inRows inCols inputMatrices =
  let starts     = fittingStarts inRows krows srows inCols kcols scols
  in  uncurry (poolBackwardFit starts krows kcols) <$> inputMatrices

poolBackwardFit :: [(Int,Int)] -> Int -> Int -> Matrix RealNum -> Matrix RealNum -> Matrix RealNum
poolBackwardFit starts krows kcols inputMatrix gradientMatrix =
  let inRows     = rows inputMatrix
      inCols     = cols inputMatrix
      inds       = fmap (\start -> maxIndex $ subMatrix start (krows, kcols) inputMatrix) starts
      grads      = toList $ flatten gradientMatrix
      grads'     = zip3 starts grads inds
      accums     = fmap (\((stx',sty'),grad,(inx, iny)) -> ((stx' + inx, sty' + iny), grad)) grads'
  in  accum (konst 0 (inRows, inCols)) (+) accums

-- | These functions are not even remotely safe, but it's only called from the statically typed
--   commands, so we should be good ?!?!?
--   Returns the starting sub matrix locations which fit inside the larger matrix for the
--   convolution. Takes into account the stride and kernel size.
fittingStarts :: Int -> Int -> Int -> Int -> Int -> Int -> [(Int,Int)]
fittingStarts nrows kernelrows steprows ncols kernelcols stepcolsh =
  let rs = fittingStart nrows kernelrows steprows
      cs = fittingStart ncols kernelcols stepcolsh
      ls = sequence [rs, cs]
  in  fmap (\[a,b] -> (a,b)) ls

-- | Returns the starting sub vector which fit inside the larger vector for the
--   convolution. Takes into account the stride and kernel size.
fittingStart :: Int -> Int -> Int -> [Int]
fittingStart width kernel steps =
  let go left | left + kernel < width
              = left : go (left + steps)
              | left + kernel == width
              = [left]
              | otherwise
              = []
  in go 0

convBackProp :: Matrix RealNum -> Int -> Int -> Int
             -> Matrix RealNum -> Int -> Int -> Int
             -> Matrix RealNum -> Int -> Int
             -> Int -> Int
             -> (Matrix RealNum, Matrix RealNum)
convBackProp input channels inRows inCols kernel filters kernelRows kernelCols dout outRows outCols strideRows strideCols =
  let fs       = [0..filters-1]
      hs       = [let h_start = h * strideRows in (h, h_start) | h <- [0..outRows - 1]]
      ws       = [let w_start = w * strideCols in (w, w_start) | w <- [0..outCols - 1]]
      fhws     = [(f, h, w) | f <- fs, h <- hs, w <- ws]
      dX_accum = concatMap dxLoop fhws
      dX       = accum (konst 0 (channels * inRows, inCols)) (+) dX_accum
      dW_accum = concatMap dwLoop fhws
      dW       = accum (konst 0 (kernelRows * kernelCols * channels, filters)) (+) dW_accum
  in (dX, dW)
  where
    dxLoop (f, (h, h_start), (w, w_start))
      = [ ((c * inRows + h_start + i, w_start + j), (indexAtdOut f h w) * (indexAtKernel f c i j)) | i <- [0..kernelRows-1], j <- [0..kernelCols-1], c <- [0..channels-1]]

    dwLoop (f, (h, h_start), (w, w_start))
      = [ ((c * kernelRows * kernelCols + i * kernelCols + j, f), (indexAtdOut f h w) * (indexAtIn c (h_start + i) (w_start + j))) | i <- [0..kernelRows-1], j <- [0..kernelCols-1], c <- [0..channels-1]]

    indexAtIn       c x y = input `atIndex` (c * inRows + x, y)
    indexAtdOut     c x y = dout `atIndex` (c * outRows + x, y)
    indexAtKernel f c x y = kernel `atIndex` (c * kernelRows * kernelCols + x * kernelCols + y, f)

convForwards :: Matrix RealNum -> Int -> Int -> Int
             -> Matrix RealNum -> Int -> Int -> Int
             -> Int -> Int
             -> Int -> Int
             -> Matrix RealNum
convForwards input channels inRows inCols kernel filters kernelRows kernelCols outRows outCols strideRows strideCols =
  let fs     = [0..filters-1]
      hs     = [(h, h * strideRows) | h <- [0..outRows - 1], h * strideRows + kernelRows - 1 < inRows]
      ws     = [(w, w * strideCols) | w <- [0..outCols - 1], w * strideCols + kernelCols - 1 < inCols]
      fhws   = [(f, h, w) | f <- fs, h <- hs, w <- ws]
      accums = map loopIter fhws
  in accum (konst 0 (filters * outRows, outCols)) (+) accums
  where
    loopIter (f, (h, h_start), (w, w_start))
      = ((f * outRows + h, w), sum [ (indexAtKernel f c i j) * (indexAtIn c (h_start + i) (w_start + j)) | i <- [0..kernelRows-1], j <- [0..kernelCols-1], c <- [0..channels-1]])

    indexAtIn c x y = input `atIndex` (c * inRows + x, y)
    indexAtKernel f c x y = kernel `atIndex` (c * kernelRows * kernelCols + x * kernelCols + y, f)

convForwardsWithPadding :: Matrix RealNum -> Int -> Int -> Int
                        -> Matrix RealNum -> Int -> Int -> Int
                        -> Int -> Int
                        -> Int -> Int
                        -> Int -> Int -> Int -> Int
                        -> Matrix RealNum
convForwardsWithPadding input channels inRows inCols kernel filters kernelRows kernelCols outRows outCols strideRows strideCols padl padt padr padb =
  let accums      = [((c * padded_r + x + padt, y + padl), input `atIndex` (c * inRows + x, y)) | x <- [0..inRows-1], y <- [0..inCols-1], c <- [0..channels-1]]
      padded_r    = inRows + padt + padb
      padded_c    = inCols + padl + padr
      paddedInput = accum (konst 0 (channels * padded_r, padded_c)) (+) accums
  in convForwards paddedInput channels padded_r padded_c kernel filters kernelRows kernelCols outRows outCols strideRows strideCols

convBackPropWithPadding :: Matrix RealNum -> Int -> Int -> Int
                        -> Matrix RealNum -> Int -> Int -> Int
                        -> Matrix RealNum -> Int -> Int
                        -> Int -> Int
                        -> Int -> Int -> Int -> Int
                        -> (Matrix RealNum, Matrix RealNum)
convBackPropWithPadding input channels inRows inCols kernel filters kernelRows kernelCols dout outRows outCols strideRows strideCols padl padt padr padb =
  let accums      = [((c * padded_r + x + padt, y + padl), input `atIndex` (c * inRows + x, y)) | x <- [0..inRows-1], y <- [0..inCols-1], c <- [0..channels-1]]
      padded_r    = inRows + padt + padb
      padded_c    = inCols + padl + padr
      paddedInput = accum (konst 0 (channels * padded_r, padded_c)) (+) accums
      (dX', dW)   = convBackProp paddedInput channels padded_r padded_c kernel filters kernelRows kernelCols dout outRows outCols strideRows strideCols
      accums'     = [((c * inRows + x, y), dX' `atIndex` (c * padded_r + x + padt, y + padl)) | x <- [0..inRows-1], y <- [0..inCols-1], c <- [0..channels-1]]
      dX          = accum (konst 0 (channels * inRows, inCols)) (+) accums'
  in  (dX, dW)

biasConvForwards :: Matrix RealNum -> Int -> Int -> Int
                 -> Matrix RealNum -> Int -> Int -> Int
                 -> Vector RealNum
                 -> Int -> Int
                 -> Int -> Int
                 -> Matrix RealNum
biasConvForwards input channels inRows inCols kernel filters kernelRows kernelCols bias outRows outCols strideRows strideCols =
  let fs     = [0..filters-1]
      hs     = [(h, h * strideRows) | h <- [0..outRows - 1], h * strideRows + kernelRows - 1 < inRows]
      ws     = [(w, w * strideCols) | w <- [0..outCols - 1], w * strideCols + kernelCols - 1 < inCols]
      fhws   = [(f, h, w) | f <- fs, h <- hs, w <- ws]
      accums = map loopIter fhws
  in accum (konst 0 (filters * outRows, outCols)) (+) accums
  where
    loopIter (f, (h, h_start), (w, w_start))
      = ((f * outRows + h, w), indexAtBias f + sum [ (indexAtKernel f c i j) * (indexAtIn c (h_start + i) (w_start + j)) | i <- [0..kernelRows-1], j <- [0..kernelCols-1], c <- [0..channels-1]])

    indexAtBias i = bias `atIndex` i
    indexAtIn c x y = input `atIndex` (c * inRows + x, y)
    indexAtKernel f c x y = kernel `atIndex` (c * kernelRows * kernelCols + x * kernelCols + y, f)

biasConvForwardsWithPadding :: Matrix RealNum -> Int -> Int -> Int
                        -> Matrix RealNum -> Int -> Int -> Int
                        -> Vector RealNum
                        -> Int -> Int
                        -> Int -> Int
                        -> Int -> Int -> Int -> Int
                        -> Matrix RealNum
biasConvForwardsWithPadding input channels inRows inCols kernel filters kernelRows kernelCols bias outRows outCols strideRows strideCols padl padt padr padb =
  let accums      = [((c * padded_r + x + padt, y + padl), input `atIndex` (c * inRows + x, y)) | x <- [0..inRows-1], y <- [0..inCols-1], c <- [0..channels-1]]
      padded_r    = inRows + padt + padb
      padded_c    = inCols + padl + padr
      paddedInput = accum (konst 0 (channels * padded_r, padded_c)) (+) accums
  in biasConvForwards paddedInput channels padded_r padded_c kernel filters kernelRows kernelCols bias outRows outCols strideRows strideCols

biasConvBackProp :: Matrix RealNum -> Int -> Int -> Int
                 -> Matrix RealNum -> Int -> Int -> Int
                 -> Matrix RealNum -> Int -> Int
                 -> Int -> Int
                 -> (Matrix RealNum, Matrix RealNum, Vector RealNum)
biasConvBackProp input channels inRows inCols kernel filters kernelRows kernelCols dout outRows outCols strideRows strideCols =
  let (dX, dW)  = convBackProp input channels inRows inCols kernel filters kernelRows kernelCols dout outRows outCols strideRows strideCols
      db_accums = [ (f, sum [ indexAtdOut f i j | i <- [0..outRows-1], j <- [0..outCols-1]]) | f <- [0..filters-1]]
      db        = accum (konst 0 filters) (+) db_accums
  in (dX, dW, db)
  where
    indexAtdOut     c x y = dout `atIndex` (c * outRows + x, y)

biasConvBackPropWithPadding :: Matrix RealNum -> Int -> Int -> Int
                            -> Matrix RealNum -> Int -> Int -> Int
                            -> Matrix RealNum -> Int -> Int
                            -> Int -> Int
                            -> Int -> Int -> Int -> Int
                            -> (Matrix RealNum, Matrix RealNum, Vector RealNum)
biasConvBackPropWithPadding input channels inRows inCols kernel filters kernelRows kernelCols dout outRows outCols strideRows strideCols padl padt padr padb =
  let accums        = [((c * padded_r + x + padt, y + padl), input `atIndex` (c * inRows + x, y)) | x <- [0..inRows-1], y <- [0..inCols-1], c <- [0..channels-1]]
      padded_r      = inRows + padt + padb
      padded_c      = inCols + padl + padr
      paddedInput   = accum (konst 0 (channels * padded_r, padded_c)) (+) accums
      (dX', dW, db) = biasConvBackProp paddedInput channels padded_r padded_c kernel filters kernelRows kernelCols dout outRows outCols strideRows strideCols
      accums'       = [((c * inRows + x, y), dX' `atIndex` (c * padded_r + x + padt, y + padl)) | x <- [0..inRows-1], y <- [0..inCols-1], c <- [0..channels-1]]
      dX            = accum (konst 0 (channels * inRows, inCols)) (+) accums'
  in  (dX, dW, db)

naiveFullyConnectedRunForwards :: forall i o. (KnownNat i, KnownNat o)
                               => H.L o i       -- Weights
                               -> H.R o         -- Biases
                               -> H.R i         -- Input
                               -> (H.R i, H.R o)  -- (Tape, Output)
naiveFullyConnectedRunForwards w b i = (i, b + (w H.#> i))

naiveFullyConnectedBackprop :: forall i o. (KnownNat i, KnownNat o)
                            => H.L o i             -- Weights
                            -> (H.R i, H.R o)        -- (Tape, Output)
                            -> (H.L o i, H.R o, H.R i) -- (NablaW, NablaB, derivatives)
naiveFullyConnectedBackprop w (tape, out) = (w', b', d)
  where
    b' = out
    w' = H.outer out tape
    d  = (H.tr w) H.#> out

-- Implementation reference: https://papers.nips.cc/paper/2012/file/c399862d3b9d6b76c8436e924a68c45b-Paper.pdf
naiveLRNForwards :: RealNum -> RealNum -> RealNum -> Int -> [[[RealNum]]] -> [[[RealNum]]]
naiveLRNForwards a b k n values = [
    [[ g ch ro co | co <- [0..length (values!!ch!!ro) - 1] ] | ro <- [0..length (values!!ch) - 1]] | ch <- [0..cs-1]
  ]
  where
    cs = length values
    f ch ro co = values!!ch!!ro!!co
    g ch ro co = (f ch ro co) / den
      where
        den  = den' ** b
        sub = floor ((fromIntegral n) / 2     :: RealNum)
        add = floor ((fromIntegral n - 1) / 2 :: RealNum)
        lower = maximum [0, ch - sub]
        upper = minimum [cs - 1, ch + add]
        summation = sum [ (f j ro co) ** 2 | j <- [lower..upper]]
        den' = k + a * summation

naiveLRNBackwards :: RealNum         -- a
                    -> RealNum       -- b
                    -> RealNum       -- k
                    -> Int          -- n
                    -> [[[RealNum]]] -- inputs
                    -> [[[RealNum]]] -- backpropagated error
                    -> [[[RealNum]]] -- error to propagate further
naiveLRNBackwards a b k n values errs = [
    [[ ng ch ro co | co <- [0..length (values!!ch!!ro) - 1] ] | ro <- [0..length (values!!ch) - 1]] | ch <- [0..cs-1]
  ]
  where
    cs = length values
    f  ch ro co = values!!ch!!ro!!co
    nf ch ro co = errs!!ch!!ro!!co
    c ch ro co = den
      where
        sub = floor ((fromIntegral n) / 2     :: RealNum)
        add = floor ((fromIntegral n - 1) / 2 :: RealNum)
        lower = maximum [0, ch - sub]
        upper = minimum [cs - 1, ch + add]
        summation = sum [ (f j ro co) ** 2 | j <- [lower..upper]]
        den = k + a * summation
    ng ch ro co = t1 - t2
      where
        t1 = (c ch ro co) ** (-b) * (nf ch ro co)
        t2 = 2 * b * a * (f ch ro co) * (c ch ro co) ** (-b - 1) * s
        s  = sum [(f q ro co) * (nf q ro co) | q <- [lower..upper]]

        sub = floor ((fromIntegral n) / 2     :: RealNum)
        add = floor ((fromIntegral n - 1) / 2 :: RealNum)
        lower = maximum [0, ch - sub]
        upper = minimum [cs - 1, ch + add]


{-# LANGUAGE BangPatterns          #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE QuantifiedConstraints #-}
{-# LANGUAGE ScopedTypeVariables   #-}
{-# LANGUAGE TupleSections         #-}
module Q.Stochastic.Process
        where
import           Control.Monad
import           Control.Monad.State
import           Data.List             (foldl')
import           Data.RVar
import           Data.Random
import           Numeric.LinearAlgebra

rwalkState :: RVarT (State Double) Double
rwalkState = do
    prev <- lift get
    change  <- rvarT StdNormal

    let new = prev + change
    lift (put new)
    return new

type Time = Double

-- Dont know why this wasn't done.
-- Is there an easier way to do this where we either lift or return?
instance (Num a) => Num (RVarT m a) where
  (+) = liftM2 (+)
  (-) = liftM2 (-)
  (*) = liftM2 (*)
  abs = liftM abs
  signum = liftM signum
  fromInteger x = return $ fromInteger x

-- |Discretization of stochastic process over given interval
class (Num b) => Discretize d b where
  -- |Discretization of the drift process.
  dDrift  :: (StochasticProcess a b) => a -> d -> (Time, b) -> RVar b
  -- |Discretization of the diffusion process.
  dDiff   :: (StochasticProcess a b) => a -> d -> (Time, b) -> RVar b
  -- |dt used.
  dDt     :: (StochasticProcess a b) => a -> d -> (Time, b) -> Time

-- |A stochastic process of the form \(dX_t = \mu(X_t, t)dt + \sigma(S_t, t)dB_t \)
class (Num b) => StochasticProcess a b where
  -- |The process drift.
  pDrift  :: a -> (Time, b) -> RVar b
  -- |The process diffusion.
  pDiff   :: a -> (Time, b) -> RVar b

  -- |Evolve a process from a given state to a given time.
  pEvolve :: (Discretize d b) => a         -- ^The process
                             -> d         -- ^Discretization scheme
                             -> (Time, b) -- ^Initial state
                             -> Time      -- ^Target time t.
                             -> RVar b    -- ^\(dB_i\).
                             -> RVar b    -- ^\(X(t)\).
  pEvolve p disc s0@(t0, x0) t dw = do
    if t0 >= t then return x0 else do
      s'@(t', b') <- pEvolve' p disc s0 dw
      if t' >= t then return b' else pEvolve p disc s' t dw

  -- |Similar to evolve, but evolves the process with the discretization scheme \(dt\).
  pEvolve' :: (Discretize d b, Num b) => a -> d -> (Time, b) -> RVar b -> RVar (Time, b)
  pEvolve' process discr s@(t, b) dw = do
    let !newT = t + dDt process discr s
        !newX = do
               drift <- dDrift process discr s
               diff  <- dDiff process discr s
               dw' <- dw
               return $ b + drift + diff * dw'
        newX :: RVar b

    (newT,) <$>  newX

-- |Geometric Brownian motion
data GeometricBrownian = GeometricBrownian {
    gbDrift :: Double -- ^Drift
  , gbDiff  :: Double -- ^Vol
} deriving (Show)

instance StochasticProcess GeometricBrownian Double where
--  pDrift :: GeometricBrownian -> (Time, Double) -> RVar Double
  pDrift p (_, x) = return $ gbDrift p * x -- drift is prpotional to the spot.
  pDiff  p (_, x) = return $ gbDiff p  * x -- diffisuion is also prportional to the spot.

-- | Ito process
data ItoProcess = ItoProcess {
        ipDrift :: (Time, Double) -> Double,
        ipDiff  :: (Time, Double) -> Double
}


-- |
-- Module      : Tiler.hs
-- Description :
-- Copyright   : (c) Jonatan H Sundqvist, 2015
-- License     : MIT
-- Maintainer  : Jonatan H Sundqvist
-- Stability   : experimental|stable
-- Portability : POSIX (not sure)
--

-- Created July 9 2015

-- TODO | -
--        -

-- SPEC | -
--        -

module Tiler (renderPathWithJoints) where

---------------------------------------------------------------------------------------------------
-- We'll need these
---------------------------------------------------------------------------------------------------
import Control.Monad (mapM_, forM_)
import Data.Complex

import qualified Graphics.Rendering.Cairo as Cairo

import qualified Southpaw.Picasso.Palette as Palette
import Southpaw.Utilities.Utilities (pairwise)

---------------------------------------------------------------------------------------------------
-- Data
---------------------------------------------------------------------------------------------------
π = pi
τ = 2*π

---------------------------------------------------------------------------------------------------
-- Types
---------------------------------------------------------------------------------------------------
type Path = [Complex Double]

---------------------------------------------------------------------------------------------------
-- Functions
---------------------------------------------------------------------------------------------------
renderPathWithJoints :: Palette.Colour -> Palette.Colour -> Double -> Double -> Path -> Cairo.Render ()
renderPathWithJoints line joint radius thickness path = do

	-- Render the lines
	Cairo.setLineWidth thickness
	forM_ (pairwise $ path ++ [head path]) $ \ (fr, to) -> do
		Palette.choose line
		Cairo.moveTo (realPart fr) (imagPart fr)
		Cairo.lineTo (realPart to) (imagPart to)
		Cairo.stroke

	-- Render the 'joints'
	forM_ path $ \ (x:+y) -> do
		Palette.choose joint
		Cairo.arc x y radius 0 τ
		Cairo.fill


{-# language TypeFamilies #-}
{-# language ScopedTypeVariables #-}
{-# language RankNTypes #-}

-- | A <https://en.wikipedia.org/wiki/Monus commutative monoid with monus>
-- is a 'Monoid' equipped with a subtraction operator.
module Data.Monoid.Monus
  ( Monus(..)
  , (-)
  ) where

import Prelude hiding ((-))
import Data.Set (Set)
import Data.Complex (Complex(..))
import Data.Monoid (Any(..),All(..),Sum(..), Endo(..))
import Control.Applicative (liftA2)
import Numeric.Natural (Natural)
import Data.Foldable
import Data.Coerce

import qualified Prelude as P
import qualified Data.Set as S

infixl 6 -

-- | A commutative monoid that supports subtraction. The following
-- laws must hold:
--
-- > x <> (y - x) = y <> (x - y)
-- > (x - y) - z = x - (y <> z)
-- > x - x = mempty
-- > mempty - x = mempty
class Monoid a => Monus a where
  monus :: a -> a -> a

-- | An infix synonym for 'subtraction'.
(-) :: Monus a => a -> a -> a
(-) = monus
{-# INLINE (-) #-}

instance Ord a => Monus (Set a) where
  monus = S.difference
  {-# INLINE monus #-}

-- | Unlike the subtraction provided by the 'Num' instance of
-- 'Natural', this subtraction is total.
instance (a ~ Natural) => Monus (Sum a) where
  monus (Sum x) (Sum y) = Sum (if x > y then x P.- y else 0)
  {-# INLINE monus #-}

-- | Defined as @P - Q = P ∧ ¬Q@
instance Monus Any where
  monus (Any x) (Any y) = case x of
    False -> Any False
    True -> Any (not y)
  {-# INLINE monus #-}

-- | Defined as @P - Q = P ∨ ¬Q@
instance Monus All where
  monus (All x) (All y) = case x of
    False -> All (not y)
    True -> All True
  {-# INLINE monus #-}

instance forall a. Monus a => Monus (Endo a) where
  monus = coerce (liftA2 monus :: (a -> a) -> (a -> a) -> (a -> a))
  {-# INLINE monus #-}

instance Monus () where
  monus _ _ = ()
  {-# INLINE monus #-}

instance (Monus a, Monus b) => Monus (a,b) where
  monus (a1,b1) (a2,b2) = (monus a1 a2,monus b1 b2)
  {-# INLINE monus #-}

instance (Monus a, Monus b, Monus c) => Monus (a,b,c) where
  monus (a1,b1,c1) (a2,b2,c2) = (monus a1 a2,monus b1 b2,monus c1 c2)
  {-# INLINE monus #-}

instance (Monus a, Monus b, Monus c,Monus d) => Monus (a,b,c,d) where
  monus (a1,b1,c1,d1) (a2,b2,c2,d2) =
    (monus a1 a2,monus b1 b2,monus c1 c2,monus d1 d2)
  {-# INLINE monus #-}

instance (Monus a, Monus b, Monus c,Monus d,Monus e) => Monus (a,b,c,d,e) where
  monus (a1,b1,c1,d1,e1) (a2,b2,c2,d2,e2) =
    (monus a1 a2,monus b1 b2,monus c1 c2,monus d1 d2,monus e1 e2)
  {-# INLINE monus #-}

instance Monus b => Monus (a -> b) where
  monus = liftA2 monus
  {-# INLINE monus #-}

instance Monus a => Monus (IO a) where
  monus = liftA2 monus
  {-# INLINE monus #-}

instance Monus a => Monus [a] where
  monus [] _ = []
  monus x [] = x
  monus (x:xs) (y:ys) = monus x y : monus xs ys

instance Monus a => Monus (Maybe a) where
  monus = liftA2 monus
  {-# INLINE monus #-}


{-# LANGUAGE CPP                   #-}
{-# LANGUAGE FlexibleInstances     #-}
{-# LANGUAGE FlexibleContexts      #-}
{-# LANGUAGE OverloadedStrings     #-}
{-# LANGUAGE ScopedTypeVariables   #-}
{-# LANGUAGE TypeFamilies          #-}
{-# LANGUAGE DataKinds             #-}
{-# LANGUAGE PolyKinds             #-}
{-# LANGUAGE GADTs                 #-}
{-# LANGUAGE TypeApplications      #-}
{-# LANGUAGE TypeOperators         #-}
{-# LANGUAGE AllowAmbiguousTypes   #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE UndecidableInstances  #-}
{-# LANGUAGE ConstraintKinds #-}
module Frames.VegaLite.Regression
  (
    Frame2DRegressionScatterFit (..)
  , frameScatterWithFit
  , keyedLayeredFrameScatterWithFit
  , scatterWithFit
  , FitToPlot (..)
  , regressionCoefficientPlot
  , regressionCoefficientPlotMany
  ) where

import qualified Frames.VegaLite.Utils as FV
import qualified Frames.Transform as FT
import qualified Frames.Regression as FR
import qualified Math.Regression.Regression as RE

import qualified Control.Foldl          as FL
import qualified Data.Text              as T
import qualified Data.Vinyl             as V
import qualified Data.Vinyl.Functor     as V
import qualified Data.Vinyl.TypeLevel   as V
import qualified Frames                 as F
import qualified Frames.Melt            as F
import qualified Graphics.Vega.VegaLite as GV
import qualified Data.List              as List
import Text.Printf (printf)

import qualified Statistics.Types               as S

#if MIN_VERSION_hvega(0,4,0)
gvTitle :: Text -> GV.PropertySpec
gvTitle x = GV.title x []
#else
gvTitle :: Text -> (GV.VLProperty, GV.VLSpec)
gvTitle = GV.title
#endif

-- | Plot regression coefficents with error bars
-- | Flex version handles a foldable of results so we can, e.g.,
-- | 1. Compare across time or other variable in the data
-- | 2. Compare different fitting methods for the same data

-- TODO: Fix replicated y-axis ticks.  There since we need the difference between them (trailing "'"s) to provide y-offset
-- to the different results.  Otherwise they would overlap in y.  Maybe fix by switching to layers and only having Y-axis labels
-- on 0th?

-- TODO: Make this typed?  Using a regression result with typed parameters?  Power-to-weight?

regressionCoefficientPlot :: T.Text -> [T.Text] -> RE.RegressionResult Double -> S.CL Double -> GV.VegaLite
regressionCoefficientPlot title names r cl = regressionCoefficientPlotFlex False id title names (V.Identity ("",r)) cl

regressionCoefficientPlotMany :: Foldable f
                              => (k -> T.Text) -> T.Text -> [T.Text] -> f (k, RE.RegressionResult Double) -> S.CL Double -> GV.VegaLite
regressionCoefficientPlotMany = regressionCoefficientPlotFlex True

regressionCoefficientPlotFlex :: Foldable f
                              => Bool -> (k -> Text) -> T.Text -> [T.Text] -> f (k, RE.RegressionResult Double) -> S.CL Double -> GV.VegaLite
regressionCoefficientPlotFlex haveLegend printKey title names results cl =
  let toRow m (RE.NamedEstimate n e eci _) = [("Parameter",GV.Str (n <> T.replicate m "'")), ("Estimate",GV.Number e), ("Confidence",GV.Number eci)]
      addKey k l = ("Key", GV.Str $ printKey k) : l
      dataRowFold = FL.Fold (\(l,n) (k,regRes) -> (l ++ fmap (flip GV.dataRow [] . addKey k . toRow n) (RE.namedEstimates names regRes cl), n+1)) ([],0) (GV.dataFromRows [] . List.concat . fst)
      dat = FL.fold dataRowFold results
      xLabel = "Estimate (with " <> T.pack (printf "%2.0f" (100 * S.confidenceLevel cl)) <> "% confidence error bars)"
      estimateXEnc = GV.position GV.X [GV.PName "Estimate", GV.PmType GV.Quantitative, GV.PAxis [GV.AxTitle xLabel]]
      estimateYEnc = GV.position GV.Y [GV.PName "Parameter", GV.PmType GV.Ordinal]
      handleLegend l = if haveLegend then l else GV.MLegend [] : l
      estimateColorEnc = GV.color $ handleLegend  [GV.MName "Key", GV.MmType GV.Nominal]
      estimateEnc = estimateXEnc . estimateYEnc . estimateColorEnc
      estLoCalc = "datum.Estimate - datum.Confidence/2"
      estHiCalc = "datum.Estimate + datum.Confidence/2"
      calcEstConf = GV.calculateAs estLoCalc "estLo" . GV.calculateAs estHiCalc "estHi"
      estConfLoEnc = GV.position GV.X [GV.PName "estLo", GV.PmType GV.Quantitative, GV.PAxis [GV.AxTitle ""]]
      estConfHiEnc = GV.position GV.X2 [GV.PName "estHi", GV.PmType GV.Quantitative, GV.PAxis [GV.AxTitle ""]]
      estConfEnc = estConfLoEnc . estConfHiEnc . estimateYEnc . estimateColorEnc
      estSpec = GV.asSpec [(GV.encoding . estimateEnc) [], GV.mark GV.Point []]
      confSpec = GV.asSpec [(GV.encoding . estConfEnc) [], GV.mark GV.Rule []]
      configuration = GV.configure
        . GV.configuration (GV.ViewStyle [GV.ViewContinuousWidth 800, GV.ViewContinuousHeight 400]) . GV.configuration (GV.PaddingStyle $ GV.PSize 50)
      vl = GV.toVegaLite
        [
          gvTitle title
        , (GV.transform . calcEstConf) []
        , GV.layer [estSpec, confSpec]
        , dat
        , configuration []
        ]
  in vl

--
{-
type family WeightElemOf (rs :: [(Symbol, Type)]) (w :: (Symbol, Type)) :: Constraint where
  WeightElemOf _ FR.Unweighted = ()
  WeightElemOf rs w  = F.ElemOf rs w
-}

frameRegressionError :: forall y wc as w rs. FR.FrameRegressionResult y wc as w rs -> Error rs
frameRegressionError (FR.FrameUnweightedRegressionResult _) = const 0
frameRegressionError (FR.FrameWeightedRegressionResult wf _) = (\r -> 1/(wf $ F.rgetField @w r))

type ScatterFitConstraints x y = ( F.ColumnHeaders '[x]
                                 , F.ColumnHeaders '[y]
                                 , FV.ToVLDataValue (F.ElField x)
                                 , FV.ToVLDataValue (F.ElField y)
                                 , V.KnownField y
                                 , Real (V.Snd y)
                                 )

type ScatterFitC1 rs x y = ( F.ElemOf (rs V.++ [YError,YFit,YFitError]) x
                           , F.ElemOf (rs V.++ [YError,YFit,YFitError]) y
                           , F.ElemOf (rs V.++ [YError,YFit,YFitError]) YError
                           , F.ElemOf (rs V.++ [YError,YFit,YFitError]) YFit
                           , F.ElemOf (rs V.++ [YError,YFit,YFitError]) YFitError)

class Frame2DRegressionScatterFit rs a where
  regressionResultToError :: a -> Error rs
  regressionResultToFit :: Maybe T.Text -> a -> S.CL Double -> FitToPlot rs
--  scatterPlot :: (Foldable f, Functor f) => T.Text -> Maybe T.Text -> a -> Double -> f (F.Record rs) -> GV.VegaLite
  scatterPlotSpec :: (Foldable f, Functor f) => Maybe T.Text -> a -> S.CL Double -> f (F.Record rs) -> GV.VLSpec

frameScatterWithFit :: (Frame2DRegressionScatterFit rs a, Foldable f, Functor f)
                    => T.Text -> Maybe T.Text -> a -> S.CL Double -> f (F.Record rs) -> GV.VegaLite
frameScatterWithFit title fitNameM frr cl frame =
  let configuration = GV.configure
        . GV.configuration (GV.ViewStyle [GV.ViewContinuousWidth 800, GV.ViewContinuousHeight 400]) . GV.configuration (GV.PaddingStyle $ GV.PSize 50)
      swfSpec = GV.layer [scatterPlotSpec fitNameM frr cl frame]
  in GV.toVegaLite [configuration [], swfSpec, gvTitle title]

keyedLayeredFrameScatterWithFit :: (Frame2DRegressionScatterFit rs a, Foldable f, Functor f, Foldable g, Functor g)
  => T.Text -> (k -> T.Text) -> g (k, a) -> S.CL Double -> f (F.Record rs) -> GV.VegaLite
keyedLayeredFrameScatterWithFit title keyText keyedFits cl dat =
  let toSpec (k, a) = scatterPlotSpec (Just $ keyText k) a cl dat
      specs = FL.fold FL.list (fmap toSpec keyedFits)
      configuration = GV.configure
        . GV.configuration (GV.ViewStyle [GV.ViewContinuousWidth 800, GV.ViewContinuousHeight 400]) . GV.configuration (GV.PaddingStyle $ GV.PSize 50)
  in GV.toVegaLite [configuration [], GV.layer specs, gvTitle title]

instance ( V.KnownField x
         , Real (V.Snd x)
         , ScatterFitConstraints x y
         , ScatterFitC1 rs x y
         , F.ElemOf (rs V.++ [YError,YFit,YFitError]) x
         , F.ElemOf (rs V.++ [YError,YFit,YFitError]) y
         , F.ElemOf rs x
         , F.ElemOf rs y
         , V.KnownField w
         , rs F.⊆ rs
         ) => Frame2DRegressionScatterFit rs (FR.FrameRegressionResult y 'True '[x] w rs) where
  regressionResultToError = frameRegressionError
  regressionResultToFit fitNameM frr cl =
    let label = fromMaybe "fit" fitNameM
        dof = RE.degreesOfFreedom (FR.regressionResult frr)
        predF = RE.predictFromEstimateAtConfidence dof frr cl
    in FitToPlot label predF

--  scatterPlot title fitNameM frr ci frame =
--    scatterWithFit @x @y @rs @rs (regressionResultToError frr) (regressionResultToFit fitNameM frr ci) Nothing title frame

  scatterPlotSpec fitNameM frr cl frame =
      scatterWithFitSpec @x @y @rs @rs (regressionResultToError frr) (regressionResultToFit fitNameM frr cl) Nothing frame

-- in this case we have y = a(x1) + b(x2) so we plot (y/x1) vs (a + b(x1/x2))
type YOverX1 = "y_over_x1" F.:-> Double
type X2OverX1 = "x2_over_x1" F.:-> Double

instance ( F.ElemOf rs x1
         , F.ElemOf rs x2
         , F.ElemOf rs y
         , V.KnownField y
         , V.KnownField x1
         , V.KnownField x2
         , V.KnownField w
         , Real (V.Snd x1)
         , Real (V.Snd x2)
         , Real (V.Snd y)
         , F.ElemOf [x1,x2] x1
         , F.ElemOf [x1,x2] x2
         , F.ElemOf (rs V.++ [YOverX1, X2OverX1]) x1
         , F.ElemOf (rs V.++ [YOverX1, X2OverX1]) x2
         , F.ElemOf (rs V.++ [YOverX1, X2OverX1]) y
         , F.ElemOf (rs V.++ [YOverX1, X2OverX1]) X2OverX1
         , F.ElemOf (rs V.++ [YOverX1, X2OverX1]) YOverX1
         , rs F.⊆ (rs V.++  [YOverX1, X2OverX1])
         , ScatterFitC1 (rs V.++ [YOverX1, X2OverX1]) YOverX1 X2OverX1
         ) => Frame2DRegressionScatterFit rs (FR.FrameRegressionResult y 'False '[x1,x2] w rs) where
  regressionResultToError frr =
    let x1 = realToFrac . F.rgetField @x1
    in (\r -> frameRegressionError frr r/x1 r)

  regressionResultToFit fitNameM frr cl =
    let label = fromMaybe "fit" fitNameM
        dof = RE.degreesOfFreedom (FR.regressionResult frr)
        predF r =
          let x1 = realToFrac $ F.rgetField @x1 r
              (y,dy) = RE.predictFromEstimateAtConfidence dof frr cl r
          in (y/x1, dy/x1)
    in FitToPlot label predF

  scatterPlotSpec fitNameM frr cl frame =
    let mut :: F.Record rs -> F.Record [YOverX1, X2OverX1]
        mut r =
          let y = F.rgetField @y r
              x1 = F.rgetField @x1 r
              x2 = F.rgetField @x2 r
          in (realToFrac y/realToFrac x1) F.&: (realToFrac x2/realToFrac x1) F.&: V.RNil
        mutData = fmap (FT.mutate mut) frame
        yName = FV.colName @y
        x1Name = FV.colName @x1
        x2Name = FV.colName @x2
        xLabel = x2Name <> "/" <> x1Name
        yLabel = yName <> "/" <> x1Name
    in scatterWithFitSpec @X2OverX1 @YOverX1 @(rs V.++ [YOverX1,X2OverX1]) @rs (regressionResultToError @rs frr) (regressionResultToFit @rs fitNameM frr cl) (Just (xLabel,yLabel)) mutData

{-
frame2DRegressionScatter :: forall x y rs a f. (Foldable f, Frame2DRegressionScatterFit rs x y a)
  => T.Text -> Maybe T.Text -> a -> Double -> f (F.Record rs) -> GV.VegaLite
frame2DRegressionScatter title fitNameM frr ci frame =
-}

--
type YError = "yError" F.:-> Double
type YFit = "yFit" F.:-> Double
type YFitError = "yFitError" F.:->Double

type Error rs = (F.Record rs -> Double)
data FitToPlot rs = FitToPlot { fitLabel :: T.Text, fitFunction :: F.Record rs -> (Double, Double) }

-- | 2D Scatter of Data with calculated error and fit function.
-- | Use TypeApplications to specify x and y columns for scatter and then provide calculated error and fit.
-- | Since both calculations use the record itself as the domain, you can put the error and fit into the record and just use
-- | field selection.  But this allows more flexibility and doesn't require adding things to the input frame record just for the plot.

scatterWithFit :: forall x y rs as f. ( ScatterFitConstraints x y
                                      , as F.⊆ rs
                                      , ScatterFitC1 rs x y
                                      , Foldable f)
               => Error as -> FitToPlot as -> Maybe (T.Text, T.Text) -> Text -> f (F.Record rs) -> GV.VegaLite
scatterWithFit err fit axisLabelsM title frame =
  let configuration = GV.configure
        . GV.configuration (GV.ViewStyle [GV.ViewContinuousWidth 800, GV.ViewContinuousHeight 400]) . GV.configuration (GV.PaddingStyle $ GV.PSize 50)
      swfSpec = GV.specification $ scatterWithFitSpec @x @y @rs @as @f err fit axisLabelsM frame
  in GV.toVegaLite [configuration [], swfSpec, gvTitle title]

scatterWithFitSpec :: forall x y rs as f. ( ScatterFitConstraints x y
                                          , as F.⊆ rs
                                          , ScatterFitC1 rs x y
                                          , Foldable f)
               => Error as -> FitToPlot as -> Maybe (T.Text, T.Text) -> f (F.Record rs) -> GV.VLSpec
scatterWithFitSpec err fit axisLabelsM frame =
  let mut :: F.Record rs -> F.Record '[YError, YFit, YFitError]
      mut r =
        let a = F.rcast r
            (f,fe) = fitFunction fit a
        in err a F.&: f F.&: fe F.&: V.RNil
      vegaDat = FV.recordsToVLData (F.rcast @[x,y,YError,YFit,YFitError] . FT.mutate mut) frame
  in scatterWithFitSpec' @x @y @YError @YFit @YFitError axisLabelsM (fitLabel fit) vegaDat

{-
scatterWithFit' :: forall x y ye fy fye. ( F.ColumnHeaders '[x]
                                         , F.ColumnHeaders '[y]
                                         , F.ColumnHeaders '[ye]
                                         , F.ColumnHeaders '[fy]
                                         , F.ColumnHeaders '[fye])
  => Maybe (T.Text, T.Text) -> Text -> Text -> GV.Data -> GV.VegaLite
scatterWithFit' axisLabelsM title fitLbl dat =
  let configuration = GV.configure
        . GV.configuration (GV.ViewStyle [GV.ViewContinuousWidth 800, GV.ViewContinuousHeight 400]) . GV.configuration (GV.PaddingStyle $ GV.PSize 50)
      swfSpec = GV.specification $ scatterWithFitSpec' @x @y @ye @fy @fye axisLabelsM fitLbl dat
  in GV.toVegaLite [configuration [], swfSpec, gvTitle title]
-}
-- TODO: Add xErrors as well, in scatter and in fit
scatterWithFitSpec' :: forall x y ye fy fye. ( F.ColumnHeaders '[x]
                                             , F.ColumnHeaders '[y]
                                             , F.ColumnHeaders '[ye]
                                             , F.ColumnHeaders '[fy]
                                             , F.ColumnHeaders '[fye])
  => Maybe (T.Text, T.Text) -> Text -> GV.Data -> GV.VLSpec
scatterWithFitSpec' axisLabelsM fitLbl dat =
-- create 4 new cols so we can use rules/areas for errors
  let yLoCalc yName yErrName = "datum." <> yName <> " - (datum." <> yErrName <> ")/2"
      yHiCalc yName yErrName = "datum." <> yName <> " + (datum." <> yErrName <> ")/2"
      calcs = GV.calculateAs (yLoCalc (FV.colName @y) (FV.colName @ye)) "yLo"
              . GV.calculateAs (yHiCalc (FV.colName @y) (FV.colName @ye)) "yHi"
              . GV.calculateAs (yLoCalc (FV.colName @fy) (FV.colName @fye)) "fyLo"
              . GV.calculateAs (yHiCalc (FV.colName @fy) (FV.colName @fye)) "fyHi"
              . GV.calculateAs ("\"" <> fitLbl <> "\"")  "fitLabel"
      xLabel = maybe (FV.colName @x) fst axisLabelsM
      yLabel = maybe (FV.colName @y) snd axisLabelsM
      xEnc = GV.position GV.X [FV.pName @x, GV.PmType GV.Quantitative, GV.PAxis [GV.AxTitle xLabel]]
      yEnc = GV.position GV.Y [FV.pName @y, GV.PmType GV.Quantitative, GV.PAxis [GV.AxTitle yLabel]]
      yLEnc = GV.position GV.Y [GV.PName "yLo", GV.PmType GV.Quantitative]
      yHEnc = GV.position GV.Y2 [GV.PName "yHi", GV.PmType GV.Quantitative]
--      yErrorEnc = GV.position GV.YError [FV.pName @ye, GV.PmType GV.Quantitative]
      yFitEnc t = GV.position GV.Y [FV.pName @fy, GV.PmType GV.Quantitative, GV.PAxis [GV.AxTitle t]]
      yFitLEnc = GV.position GV.Y [GV.PName "fyLo", GV.PmType GV.Quantitative]
      yFitHEnc = GV.position GV.Y2 [GV.PName "fyHi", GV.PmType GV.Quantitative]
      colorEnc = GV.color [GV.MName "fitLabel", GV.MmType GV.Nominal]
  --    yFitErrorEnc = GV.position GV.YError [FV.pName @fye, GV.PmType GV.Quantitative]
      scatterEnc = xEnc . yEnc
      scatterBarEnc = xEnc . yLEnc . yHEnc
      fitEnc t = xEnc . yFitEnc t . colorEnc
      fitBandEnc = xEnc . yFitLEnc . yFitHEnc . colorEnc

      selectScalesS = GV.select ("scalesS" <> fitLbl) GV.Interval [GV.BindScales]
      selectScalesSE = GV.select("scalesSE" <> fitLbl) GV.Interval [GV.BindScales]
      selectScalesF = GV.select ("scalesF" <> fitLbl) GV.Interval [GV.BindScales]
      selectScalesFE = GV.select ("scalesFE" <> fitLbl) GV.Interval [GV.BindScales]
      scatterSpec = GV.asSpec
        [
          (GV.encoding . scatterEnc) []
        , GV.mark GV.Point []
        , (GV.selection . selectScalesS) []
        ]
      scatterBarSpec = GV.asSpec
        [
          (GV.encoding . scatterBarEnc) []
        , GV.mark GV.Rule []
        , (GV.selection . selectScalesSE) []
        ]
      fitLineSpec = GV.asSpec
        [
          (GV.encoding . fitEnc fitLbl) []
        , GV.mark GV.Line []
        , (GV.selection . selectScalesF) []
        ]
      fitBandSpec = GV.asSpec
        [
          (GV.encoding . fitBandEnc) []
        , GV.mark GV.Area [GV.MOpacity 0.5, GV.MFillOpacity 0.3]
        , (GV.selection . selectScalesFE) []
        ]
      layers = GV.layer [ scatterSpec, scatterBarSpec, fitLineSpec, fitBandSpec]
      configuration = GV.configure
        . GV.configuration (GV.ViewStyle [GV.ViewContinuousWidth 800, GV.ViewContinuousHeight 400]) . GV.configuration (GV.PaddingStyle $ GV.PSize 50)
      spec =
        GV.asSpec
        [
          (GV.transform . calcs) []
        , layers
        , dat
        , configuration []
        ]
  in spec