{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE RecordWildCards, PatternGuards, BangPatterns #-}
-----------------------------------------------------------------------------
--
-- Module      :  AI.Network.RNN.RNN
-- Copyright   :  (c) JP Moresmau
-- License     :  BSD3
--
-- Maintainer  :  JP Moresmau <jp@moresmau.fr>
-- Stability   :  experimental
-- Portability :
--
-- | Recurrent neural network
--
-----------------------------------------------------------------------------

module AI.Network.RNN.RNN  where

import Numeric.LinearAlgebra.HMatrix hiding ((|>))
import Control.DeepSeq
import Control.Monad.Random hiding (fromList)

import AI.Network.RNN.Types
import AI.Network.RNN.Util

data RNNDimensions = RNNDimensions
  { rnndInput    :: Int
  , rnndInternal :: Int
  , rnndOutput   :: Int
  , rnndWithBack :: Bool
  }  deriving (Show,Read,Eq,Ord)

instance NFData RNNDimensions where
    rnf RNNDimensions{..} = rnf (rnndInput,rnndInternal,rnndOutput,rnndWithBack)

data RNNetwork = RNNetwork
  { rnnDimensions :: RNNDimensions
  , rnnMIn        :: Matrix Double
  , rnnM          :: Matrix Double
  , rnnMOut       :: Matrix Double
  , rnnMBack      :: Maybe (Matrix Double)
  , rnnState      :: Vector Double
  , rnnOutput     :: Vector Double
  } deriving (Show,Read,Eq)

networkDimensions :: RNNetwork -> RNNDimensions
networkDimensions = rnnDimensions

instance NFData RNNetwork where
    rnf RNNetwork{..} = rnf (rnnDimensions,rnnMIn,rnnM,rnnMOut,rnnMBack,rnnState,rnnOutput)

createNetwork
    :: RNNDimensions
    -> Matrix Double
    -> Matrix Double
    -> Matrix Double
    -> Maybe (Matrix Double)
    -> Vector Double
    -> Vector Double
    -> Either [String] RNNetwork
createNetwork dim mIn m mOut mmback st out =
    case checkNetworkDimensions dim mIn m mOut mmback st of
        []  -> Right $ RNNetwork dim mIn m mOut mmback st out
        err -> Left err

createNetworkFromArray
    :: RNNDimensions
    -> [Double]
    -> Either String RNNetwork
createNetworkFromArray dim@RNNDimensions{..} v =
    let
        enough = length v >= totalDataLength dim
    in if enough
        then Right $ fromVector dim $ fromList v
        else Left $ "Not enough data in array (needs at least"++show (totalDataLength dim)++")"

networkToArray :: RNNetwork -> [Double]
networkToArray RNNetwork{..} =
       concat (toLists rnnMIn)
    ++ concat (toLists rnnM)
    ++ concat (toLists rnnMOut)
    ++ (case rnnMBack of
            Just b -> concat $ toLists b
            _      -> [])
    ++ toList rnnState
    ++ toList rnnOutput

createNetworkFromVector
    :: RNNDimensions
    -> Vector Double
    -> Either String RNNetwork
createNetworkFromVector dim@RNNDimensions{..} v =
    let
        enough = size v >= totalDataLength dim
    in if enough
        then {-# SCC "createNetworkFromVector" #-} Right $ fromVector dim v
        else Left $ "Not enough data in vector (needs at least"++ show (totalDataLength dim) ++")"

inputMatrixLength :: RNNDimensions -> Int
inputMatrixLength RNNDimensions{..} = rnndInput * rnndInternal

internalMatrixLength :: RNNDimensions -> Int
internalMatrixLength RNNDimensions{..} = rnndInternal * rnndInternal

outputMatrixLength :: RNNDimensions -> Int
outputMatrixLength RNNDimensions{..} = rnndOutput * (rnndInput + rnndInternal)

backMatrixLength :: RNNDimensions -> Int
backMatrixLength RNNDimensions{..} = if rnndWithBack then rnndInternal * rnndOutput else 0

totalDataLength :: FullSize RNNetwork
totalDataLength dim@RNNDimensions{..} =
    inputMatrixLength dim + internalMatrixLength dim + outputMatrixLength dim
    + backMatrixLength dim
    + rnndInternal + rnndOutput

randNetwork :: (Monad m,RandomGen g) =>  RNNDimensions -> RandT g m RNNetwork
randNetwork dim@RNNDimensions{..} = do
    s <- getRandom
    -- vs <- sequence (replicate (totalDataLength dim) getRandom)
    return $ fromVector dim $ randomVector s Gaussian (totalDataLength dim)

collectErrors :: [(Bool,a)] ->  [a]
collectErrors = map snd . filter fst

checkDimensions
    :: RNNDimensions
    -> [String]
checkDimensions RNNDimensions{..} = collectErrors
    [(rnndInput    < 0, "input must be >=0")
    ,(rnndInternal < 1, "internal must be >0")
    ,(rnndOutput   < 1, "output must be >0")]

checkNetworkDimensions
    :: RNNDimensions
    -> Matrix Double
    -> Matrix Double
    -> Matrix Double
    -> Maybe (Matrix Double)
    -> Vector Double
    -> [String]
checkNetworkDimensions dim@RNNDimensions{..} mIn m mOut mmback st = checkDimensions dim ++ collectErrors (
    [(cols mIn /= rnndInput, "input matrix column count doesn't match input neurons count")
    ,(rows mIn /= rnndInternal, "input matrix row count doesn't match internal neurons count")
    ,(cols m /= rnndInternal, "internal matrix column count doesn't match internal neurons count")
    ,(rows m /= rnndInternal, "internal matrix row count doesn't match internal neurons count")
    ,(cols mOut /= (rnndInput + rnndInternal), "internal matrix column count doesn't match input + internal neurons count")
    ,(rows mOut /= rnndOutput, "output matrix row count doesn't match output neurons count")
    ,(size st /= rnndInternal,"internal state length doesn't match internal neurons count")]
    ++(case mmback of
        Nothing -> []
        Just mback ->
            [(rows mback /= rnndInternal,"back matrix row count doesn't match internal neurons count")
            ,(cols m /= rnndOutput,"back matrix column count doesn't match output neurons count")
            ]))

instance RNNEval RNNetwork where
    type Size RNNetwork = RNNDimensions
    evalStep rnn@RNNetwork{..} iv =
        let
            sum1 = (rnnMIn #> iv) + (rnnM #> rnnState)
            sum2 = case rnnMBack of
                Just mback -> sum1 + (mback #> rnnOutput)
                _          -> sum1
            s2 = cmap tanh sum2
            out = cmap sigmoid (rnnMOut #> vjoin [iv,s2])
        in (rnn{rnnState=s2,rnnOutput=out},out)
    toVector RNNetwork{..} = vjoin
        [ flatten rnnMIn
        , flatten rnnM
        , flatten rnnMOut
        , maybe (fromList []) flatten rnnMBack
        , rnnState
        , rnnOutput
        ]
    fromVector dim@RNNDimensions{..} v =
        let
            lenIn   = inputMatrixLength dim
            len     = internalMatrixLength dim
            lenOut  = outputMatrixLength dim
            lenBack = backMatrixLength dim
            [v1,v2,v3,v4,v5,v6] = takesV [lenIn,len,lenOut,lenBack,rnndInternal,rnndOutput] v
            mIn     = reshape rnndInput v1
            m       = reshape rnndInternal v2
            mOut    = reshape (rnndInput + rnndInternal) v3
            mback = if rnndWithBack
                            then Just $ reshape rnndOutput v4
                            else Nothing
        in RNNetwork dim mIn m mOut mback v5 v6
    rnnsize = rnnDimensions
    fullSize = totalDataLength . rnnDimensions


{-# LANGUAGE TemplateHaskell, GeneralizedNewtypeDeriving, ScopedTypeVariables #-}
module CNC.Figure(Pos, Figure, Path,
                  unPath, mapPath, mapPathM_, deltaPath,
                  solveFigure, declarePoints, figure,
                  xSize, ySize, len, xAngle, yAngle) where

import CNC.Geometry(RealT, eps)

import qualified Data.Map as M
import Data.Map((!))
import Data.Maybe(fromJust, isNothing)
import Data.Label
import Data.List(deleteBy,sort)
import qualified Data.Label.Monadic as LM
import qualified Control.Monad.State as S
import Data.Complex
import Text.Printf

newtype PointInd = PointInd Int deriving (Eq, Ord, Enum, Show)

type Pos = Complex RealT

type FigMap = M.Map PointInd (Maybe Pos)
data Figure = Figure { _fPoints ::  FigMap, _fNextPoint :: PointInd, _fConstraints :: [Constraint] } deriving Show
newtype Path = Path {unPath :: [Pos]} deriving Show

mkPos x y = x :+ y

-- applies a given transform to all points of path
mapPath f (Path p) = Path (map f p)
mapPathM_ f (Path p) = mapM_ f p
deltaPath (Path p) = last p - head p

zero_ind = PointInd 0
zero_point = mkPos 0 0

empty_figure = Figure M.empty zero_ind []
figNumPoints fig = case _fNextPoint fig of
  PointInd k -> k

instance (Ord a, RealFloat a) => Ord (Complex a) where
  compare x y = let c1 = compare (realPart x) (realPart y)
                in if c1 == EQ then compare (imagPart x) (imagPart y)
                   else c1

data UniConstraint pt = Disp {cPoint1 :: pt, cPoint2 :: pt, cPos :: Pos} -- A vector from first point to second
                      | Length {cPoint1 :: pt, cPoint2 :: pt, cLen :: RealT} -- distance between points
                      | Angle {cPoint1 :: pt, cPoint2 :: pt, cAngle :: RealT} deriving (Eq, Ord, Show)

instance Functor UniConstraint where
  fmap f c = case c of
    Disp p1 p2 d -> Disp (f p1) (f p2) d
    Length p1 p2 l -> Length (f p1) (f p2) l
    Angle p1 p2 a -> Angle (f p1) (f p2) a

type Constraint = UniConstraint PointInd

mkLabels [''Figure]

figure fig_m = let fig = S.execState fig_m empty_figure
               in solveFigure fig

-- Creates k new points
declarePoints k = do
  next <- LM.gets fNextPoint
  let pts = take k $ iterate succ next
      pairs = map (\p -> (p, Nothing)) pts
  LM.modify fPoints (\ps -> ps `M.union` M.fromList pairs)
  LM.puts fNextPoint (succ $ last pts)
  return pts

xSize p1 p2 d = LM.modify fConstraints (Disp p1 p2 (mkPos d 0) :) -- distance between horizontally aligned points
ySize p1 p2 d = LM.modify fConstraints (Disp p1 p2 (mkPos 0 d) :) -- distance between vertically aligned points
len p1 p2 d = LM.modify fConstraints (Length p1 p2 d :) -- distance between points
xAngle p1 p2 a = LM.modify fConstraints (Angle p1 p2 a :) -- angle between line and axe X
yAngle p1 p2 a = LM.modify fConstraints (Angle p1 p2 (pi/2 - a) :) -- angle between line and axe Y

-- solves system consisting of N declared points and user specified constraints
solveFigure :: Figure -> Path
solveFigure fig | figNumPoints fig == 0 = error "can't solve figure"
                | otherwise = let points = M.insert zero_ind (Just zero_point) $ get fPoints fig
                                  constraints = sort $ get fConstraints fig
                                  solution = solveConstraints points constraints [] False
                              in Path $ map fromJust $ M.elems $ checkSolution solution constraints

solveConstraints :: FigMap -> [Constraint] -> [Constraint] -> Bool -> FigMap
solveConstraints points [] [] _ = points -- no more constraints, finish here
solveConstraints points [] prev_cs True = solveConstraints points (reverse prev_cs) [] False -- iterating over constraints again
solveConstraints points [] prev_cs False = points -- no constraints was applied, finishing
solveConstraints points (c:cs) prev_cs shrinked = case c of -- process a single constraint
  Disp pi1 pi2 disp -> case (points ! pi1, points ! pi2) of
    (Nothing, Nothing) -> skipConstr
    (Just p1, Nothing) -> solveDisp p1 pi2 disp
    (Nothing, Just p2) -> solveDisp p2 pi1 (-disp)
    (Just p1, Just p2) -> checkConstr
  Length pi1 pi2 len -> case (points ! pi1, points ! pi2) of
    (Nothing, Nothing) -> skipConstr
    (Just p1, Nothing) -> solveLength p1 pi1 pi2 len
    (Nothing, Just p2) -> solveLength p2 pi2 pi1 len
    (Just p1, Just p2) -> checkConstr
  Angle pi1 pi2 a -> case (M.lookup pi1 points, M.lookup pi2 points) of -- only check these, they're used as part of Length constraints processing
    (Just p1, Just p2) -> checkConstr
    _ -> skipConstr
  where
    pointConstr = fmap (fromJust . (points !)) c
    checkConstr = checkEps (calcDiscrepancy pointConstr) (solveConstraints points cs prev_cs shrinked) ("disagreement while processing " ++ show c ++ " evaluated to " ++ show pointConstr)
    skipConstr = solveConstraints points cs (c:prev_cs) shrinked
    solveDisp p ind d = solveConstraints (M.insert ind (Just $ p + d) points) cs prev_cs True
    solveLength (p :: Pos) ind0 ind len = case findAngleConstr ind0 ind cs of
      Nothing -> skipConstr
      Just angle -> solveConstraints (M.insert ind (Just $ p + mkPolar len angle) points)
                                     (deleteAngleConstr ind0 ind cs) prev_cs True

checkEps :: RealT -> a -> String -> a
checkEps val cont msg = if abs val < eps then cont else error $ msg

calcDiscrepancy :: (UniConstraint Pos) -> RealT
calcDiscrepancy c =
  case c of
    Disp p1 p2 d -> magnitude $ p2 - p1 - d
    Length p1 p2 l -> magnitude (p2 - p1) - l
    Angle p1 p2 a -> phase (p2 - p1) - a

checkSolution points constrs | any isNothing $ M.elems points = error $ "couldn't find a complete solution, got " ++ show (M.elems points)
checkSolution points cs = checkConstrs points cs

checkConstrs :: FigMap -> [UniConstraint PointInd] -> FigMap
checkConstrs points [] = points
checkConstrs points (c:cs) = check $ calcDiscrepancy $ fmap (fromJust . (points !)) c
 where p1 = cPoint1 c
       p2 = cPoint2 c
       check x = checkEps x (checkConstrs points cs) $
                   printf "error checking constraint %s on points %s %s placed at %s %s"
                   (show c) (show p1) (show p2) (show $ points ! p1) (show $ points ! p2)

-- finds first angle constraint on points with specified indices
findAngleConstr :: PointInd -> PointInd -> [Constraint] -> Maybe RealT
findAngleConstr _ _ [] = Nothing
findAngleConstr ind1 ind2 (c:cs) = case c of
  Angle i1 i2 a | i1 == ind1 && i2 == ind2 -> Just a
                | i1 == ind2 && i2 == ind1 -> Just (-a)
  _ -> findAngleConstr ind1 ind2 cs

-- deletes first angle constraint
deleteAngleConstr ind1 ind2 = deleteBy findAngle
                              (Angle ind1 ind2 undefined)
  where findAngle (Angle d1 d2 _) c2 =
          case c2 of
            (Angle i1 i2 _) -> (i1, i2) == (d2,d1) || (i2, i1) == (d1,d2)
            _ -> False


{-# LANGUAGE CPP                   #-}
{-# LANGUAGE DataKinds             #-}
{-# LANGUAGE FlexibleContexts      #-}
{-# LANGUAGE FlexibleInstances     #-}
{-# LANGUAGE GADTs                 #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE RankNTypes            #-}
{-# LANGUAGE ScopedTypeVariables   #-}
{-# LANGUAGE TypeFamilies          #-}
{-# LANGUAGE TypeOperators         #-}
{-# LANGUAGE UndecidableInstances  #-}
{-|
Module      : Grenade.Layers.Deconvolution
Description : Deconvolution layer
Copyright   : (c) Huw Campbell, 2016-2017
License     : BSD2
Stability   : experimental

A deconvolution layer is in many ways a convolution layer in reverse.
It learns a kernel to apply to each pixel location, spreading it out
into a larger layer.

This layer is important for image generation tasks, such as GANs on
images.
-}
module Grenade.Layers.Deconvolution (
    Deconvolution (..)
  , Deconvolution' (..)
  ) where

import           Control.DeepSeq                     (NFData (..))
import           Data.Kind                           (Type)
import           Data.List                           (foldl1')
import           Data.Maybe
import           Data.Proxy
import           Data.Serialize
import           Data.Singletons.TypeLits            hiding (natVal)
import           GHC.TypeLits

import           Numeric.LinearAlgebra               hiding (konst,
                                                      uniformSample)
import qualified Numeric.LinearAlgebra               as LA
import           Numeric.LinearAlgebra.Static        hiding (build, toRows,
                                                      (|||))

import           Grenade.Core
import           Grenade.Layers.Internal.Convolution
import           Grenade.Layers.Internal.Update
import           Grenade.Utils.LinearAlgebra
import           Grenade.Utils.ListStore

-- | A Deconvolution layer for a neural network.
--   This uses the im2col Convolution trick popularised by Caffe.
--
--   The Deconvolution layer is a way of spreading out a single response
--   into a larger image, and is useful in generating images.
--
data Deconvolution :: Nat -- Number of channels, for the first layer this could be RGB for instance.
                   -> Nat -- Number of filters, this is the number of channels output by the layer.
                   -> Nat -- The number of rows in the kernel filter
                   -> Nat -- The number of column in the kernel filter
                   -> Nat -- The row stride of the Deconvolution filter
                   -> Nat -- The columns stride of the Deconvolution filter
                   -> Type where
  Deconvolution :: ( KnownNat channels
                   , KnownNat filters
                   , KnownNat kernelRows
                   , KnownNat kernelColumns
                   , KnownNat strideRows
                   , KnownNat strideColumns
                   , KnownNat kernelFlattened
                   , kernelFlattened ~ (kernelRows * kernelColumns * filters))
                 => !(L kernelFlattened channels) -- The kernel filter weights
                 -> !(ListStore (L kernelFlattened channels)) -- The last kernel update (or momentum)
                 -> Deconvolution channels filters kernelRows kernelColumns strideRows strideColumns

instance NFData (Deconvolution channels filters kernelRows kernelColumns strideRows strideColumns) where
  rnf (Deconvolution a b) = rnf a `seq` rnf b

instance ( KnownNat channels
         , KnownNat filters
         , KnownNat kernelRows
         , KnownNat kernelColumns
         , KnownNat strideRows
         , KnownNat strideColumns
         , KnownNat (kernelRows * kernelColumns * filters)
         ) => Serialize (Deconvolution channels filters kernelRows kernelColumns strideRows strideColumns) where
  put (Deconvolution w store) = do
    putListOf put . toList . flatten . extract $ w
    put (fmap (toList . flatten . extract) store)
  get = do
      let f  = fromIntegral $ natVal (Proxy :: Proxy channels)
      wN    <- maybe (fail "Vector of incorrect size") return . create . reshape f . LA.fromList =<< getListOf get
      store <- fmap (fromMaybe (error "Vector of incorrect size") . create . reshape f . LA.fromList)  <$> get
      return $ Deconvolution wN store

data Deconvolution' :: Nat -- Number of channels, for the first layer this could be RGB for instance.
                    -> Nat -- Number of filters, this is the number of channels output by the layer.
                    -> Nat -- The number of rows in the kernel filter
                    -> Nat -- The number of column in the kernel filter
                    -> Nat -- The row stride of the Deconvolution filter
                    -> Nat -- The columns stride of the Deconvolution filter
                    -> Type where
  Deconvolution' :: ( KnownNat channels
                    , KnownNat filters
                    , KnownNat kernelRows
                    , KnownNat kernelColumns
                    , KnownNat strideRows
                    , KnownNat strideColumns
                    , KnownNat kernelFlattened
                    , kernelFlattened ~ (kernelRows * kernelColumns * filters))
                 => !(L kernelFlattened channels) -- The kernel filter gradient
                 -> Deconvolution' channels filters kernelRows kernelColumns strideRows strideColumns

instance NFData (Deconvolution' channels filters kernelRows kernelColumns strideRows strideColumns) where
  rnf (Deconvolution' a) = rnf a `seq` ()

instance ( KnownNat channels
         , KnownNat filters
         , KnownNat kernelRows
         , KnownNat kernelColumns
         , KnownNat strideRows
         , KnownNat strideColumns
         , KnownNat (kernelRows * kernelColumns * filters)
         ) =>
         Serialize (Deconvolution' channels filters kernelRows kernelColumns strideRows strideColumns) where
  put (Deconvolution' w) = putListOf put . toList . flatten . extract $ w
  get = do
    let f = fromIntegral $ natVal (Proxy :: Proxy channels)
    wN <- maybe (fail "Vector of incorrect size") return . create . reshape f . LA.fromList =<< getListOf get
    return $ Deconvolution' wN

instance (KnownNat channels
         , KnownNat filters
         , KnownNat kernelRows
         , KnownNat kernelColumns
         , KnownNat strideRows
         , KnownNat strideColumns) => FoldableGradient (Deconvolution' channels filters kernelRows kernelColumns strideRows strideColumns) where
  mapGradient f (Deconvolution' kernelGradient) = Deconvolution' (dmmap f kernelGradient)
  squaredSums (Deconvolution' kernelGradient) = [sumM . squareM $ kernelGradient]

instance Show (Deconvolution c f k k' s s') where
  show (Deconvolution a _) = renderConv a
    where
      renderConv mm =
        let m = extract mm
            ky = fromIntegral $ natVal (Proxy :: Proxy k)
            rs = LA.toColumns m
            ms = map (take ky) $ toLists . reshape ky <$> rs
            render n'
              | n' <= 0.2 = ' '
              | n' <= 0.4 = '.'
              | n' <= 0.6 = '-'
              | n' <= 0.8 = '='
              | otherwise = '#'
            px = (fmap . fmap . fmap) render ms
         in unlines $ foldl1 (zipWith (\a' b' -> a' ++ "   |   " ++ b')) px

instance ( KnownNat channels
         , KnownNat filters
         , KnownNat kernelRows
         , KnownNat kernelColumns
         , KnownNat strideRows
         , KnownNat strideColumns
         , KnownNat ((kernelRows * kernelColumns) * filters)
         , KnownNat ((kernelRows * kernelColumns) * channels)
         , KnownNat (channels * ((kernelRows * kernelColumns) * filters))
         ) =>
         RandomLayer (Deconvolution channels filters kernelRows kernelColumns strideRows strideColumns) where
  createRandomWith m gen = do
    wN <- getRandomMatrix i i m gen
    return $ Deconvolution wN mkListStore
    where
      i = natVal (Proxy :: Proxy ((kernelRows * kernelColumns) * channels))

instance ( KnownNat channels
         , KnownNat filters
         , KnownNat kernelRows
         , KnownNat kernelColumns
         , KnownNat strideRows
         , KnownNat strideColumns
         , KnownNat (kernelRows * kernelColumns * filters)
         ) => UpdateLayer (Deconvolution channels filters kernelRows kernelColumns strideRows strideColumns) where
  type Gradient (Deconvolution channels filters kernelRows kernelColumns strideRows strideColumns) = (Deconvolution' channels filters kernelRows kernelColumns strideRows strideColumns)

  type MomentumStore (Deconvolution channels filters kernelRows kernelColumns strideRows strideColumns)  = ListStore (L (kernelRows * kernelColumns * filters) channels)

  runUpdate opt@OptSGD{} x@(Deconvolution oldKernel store) (Deconvolution' kernelGradient) =
    let oldMomentum = getData opt x store
        result = descendMatrix opt (MatrixValuesSGD oldKernel kernelGradient oldMomentum)
        newStore = setData opt x store (matrixMomentum result)
    in Deconvolution (matrixActivations result) newStore
  runUpdate opt@OptAdam{} x@(Deconvolution oldKernel store) (Deconvolution' kernelGradient) =
    let (m, v) = toTuple $ getData opt x store
        result = descendMatrix opt (MatrixValuesAdam (getStep store) oldKernel kernelGradient m v)
        newStore = setData opt x store [matrixM result, matrixV result]
    in Deconvolution (matrixActivations result) newStore
    where toTuple [m ,v] = (m, v)
          toTuple xs = error $ "unexpected input of length " ++ show (length xs) ++ "in toTuple in Convolution.hs"

  reduceGradient grads = Deconvolution' $ dmmap (/ (fromIntegral $ length grads)) (foldl1' add (map (\(Deconvolution' x) -> x) grads))

instance ( KnownNat channels
         , KnownNat filters
         , KnownNat kernelRows
         , KnownNat kernelColumns
         , KnownNat strideRows
         , KnownNat strideColumns
         , KnownNat (kernelRows * kernelColumns * filters)
         ) =>
         LayerOptimizerData (Deconvolution channels filters kernelRows kernelColumns strideRows strideColumns) (Optimizer 'SGD) where
  type MomentumExpOptResult (Deconvolution channels filters kernelRows kernelColumns strideRows strideColumns) (Optimizer 'SGD) = L (kernelRows * kernelColumns * filters) channels
  type MomentumDataType (Deconvolution channels filters kernelRows kernelColumns strideRows strideColumns) (Optimizer 'SGD) = L (kernelRows * kernelColumns * filters) channels
  getData opt x store = head $ getListStore opt x store
  setData opt x store = setListStore opt x store . return
  newData _ _ = konst 0

instance ( KnownNat channels
         , KnownNat filters
         , KnownNat kernelRows
         , KnownNat kernelColumns
         , KnownNat strideRows
         , KnownNat strideColumns
         , KnownNat (kernelRows * kernelColumns * filters)
         ) => LayerOptimizerData (Deconvolution channels filters kernelRows kernelColumns strideRows strideColumns) (Optimizer 'Adam) where
  type MomentumExpOptResult (Deconvolution channels filters kernelRows kernelColumns strideRows strideColumns) (Optimizer 'Adam)  = [L (kernelRows * kernelColumns * filters) channels]
  type MomentumDataType (Deconvolution channels filters kernelRows kernelColumns strideRows strideColumns) (Optimizer 'Adam) = L (kernelRows * kernelColumns * filters) channels
  getData = getListStore
  setData = setListStore
  newData _ _ = konst 0

-- | A two dimentional image may have a Deconvolution filter applied to it
instance ( KnownNat kernelRows
         , KnownNat kernelCols
         , KnownNat filters
         , KnownNat strideRows
         , KnownNat strideCols
         , KnownNat inputRows
         , KnownNat inputCols
         , KnownNat outputRows
         , KnownNat outputCols
         , ((inputRows - 1) * strideRows) ~ (outputRows - kernelRows)
         , ((inputCols - 1) * strideCols) ~ (outputCols - kernelCols)
         , KnownNat (kernelRows * kernelCols * filters)
         , KnownNat (outputRows * filters)
         ) => Layer (Deconvolution 1 filters kernelRows kernelCols strideRows strideCols) ('D2 inputRows inputCols) ('D3 outputRows outputCols filters) where
  type Tape (Deconvolution 1 filters kernelRows kernelCols strideRows strideCols) ('D2 inputRows inputCols) ('D3 outputRows outputCols filters) = S ('D3 inputRows inputCols 1)
  runForwards c (S2D input) =
    runForwards c (S3D input :: S ('D3 inputRows inputCols 1))

  runBackwards c tape grads =
    case runBackwards c tape grads of
      (c', S3D back :: S ('D3 inputRows inputCols 1)) ->  (c', S2D back)

-- | A two dimentional image may have a Deconvolution filter applied to it
instance ( KnownNat kernelRows
         , KnownNat kernelCols
         , KnownNat strideRows
         , KnownNat strideCols
         , KnownNat inputRows
         , KnownNat inputCols
         , KnownNat outputCols
         , ((inputRows - 1) * strideRows) ~ (outputRows - kernelRows)
         , ((inputCols - 1) * strideCols) ~ (outputCols - kernelCols)
         , KnownNat (kernelRows * kernelCols * 1)
         , KnownNat (outputRows * 1)
         ) => Layer (Deconvolution 1 1 kernelRows kernelCols strideRows strideCols) ('D2 inputRows inputCols) ('D2 outputRows outputCols) where
  type Tape (Deconvolution 1 1 kernelRows kernelCols strideRows strideCols) ('D2 inputRows inputCols) ('D2 outputRows outputCols) = S ('D3 inputRows inputCols 1)
  runForwards c (S2D input) =
    case runForwards c (S3D input :: S ('D3 inputRows inputCols 1)) of
      (tps, S3D fore :: S ('D3 outputRows outputCols 1)) ->  (tps, S2D fore)

  runBackwards c tape (S2D grads) =
    case runBackwards c tape (S3D grads :: S ('D3 outputRows outputCols 1)) of
      (c', S3D back :: S ('D3 inputRows inputCols 1)) ->  (c', S2D back)

-- | A two dimentional image may have a Deconvolution filter applied to it
instance ( KnownNat kernelRows
         , KnownNat kernelCols
         , KnownNat strideRows
         , KnownNat strideCols
         , KnownNat inputRows
         , KnownNat inputCols
         , KnownNat outputCols
         , ((inputRows - 1) * strideRows) ~ (outputRows - kernelRows)
         , ((inputCols - 1) * strideCols) ~ (outputCols - kernelCols)
         , KnownNat (kernelRows * kernelCols * 1)
         , KnownNat (outputRows * 1)
         , KnownNat channels
         ) => Layer (Deconvolution channels 1 kernelRows kernelCols strideRows strideCols) ('D3 inputRows inputCols channels) ('D2 outputRows outputCols) where
  type Tape (Deconvolution channels 1 kernelRows kernelCols strideRows strideCols) ('D3 inputRows inputCols channels) ('D2 outputRows outputCols) = S ('D3 inputRows inputCols channels)
  runForwards c input =
    case runForwards c input of
      (tps, S3D fore :: S ('D3 outputRows outputCols 1)) ->  (tps, S2D fore)

  runBackwards c tape (S2D grads) =
    runBackwards c tape (S3D grads :: S ('D3 outputRows outputCols 1))

-- | A three dimensional image (or 2d with many channels) can have
--   an appropriately sized Deconvolution filter run across it.
instance ( KnownNat kernelRows
         , KnownNat kernelCols
         , KnownNat filters
         , KnownNat strideRows
         , KnownNat strideCols
         , KnownNat inputRows
         , KnownNat inputCols
         , KnownNat outputRows
         , KnownNat outputCols
         , KnownNat channels
         , ((inputRows - 1) * strideRows) ~ (outputRows - kernelRows)
         , ((inputCols - 1) * strideCols) ~ (outputCols - kernelCols)
         , KnownNat (kernelRows * kernelCols * filters)
         , KnownNat (outputRows * filters)
         ) => Layer (Deconvolution channels filters kernelRows kernelCols strideRows strideCols) ('D3 inputRows inputCols channels) ('D3 outputRows outputCols filters) where

  type Tape (Deconvolution channels filters kernelRows kernelCols strideRows strideCols) ('D3 inputRows inputCols channels) ('D3 outputRows outputCols filters) = S ('D3 inputRows inputCols channels)

  runForwards (Deconvolution kernel _) (S3D input) =
    let ex = extract input
        ek = extract kernel
        ix = fromIntegral $ natVal (Proxy :: Proxy inputRows)
        iy = fromIntegral $ natVal (Proxy :: Proxy inputCols)
        kx = fromIntegral $ natVal (Proxy :: Proxy kernelRows)
        ky = fromIntegral $ natVal (Proxy :: Proxy kernelCols)
        sx = fromIntegral $ natVal (Proxy :: Proxy strideRows)
        sy = fromIntegral $ natVal (Proxy :: Proxy strideCols)
        ox = fromIntegral $ natVal (Proxy :: Proxy outputRows)
        oy = fromIntegral $ natVal (Proxy :: Proxy outputCols)

        c  = vid2col 1 1 1 1 ix iy ex

        mt = c LA.<> tr ek

        r  = col2vid kx ky sx sy ox oy mt
        rs = fromJust . create $ r
    in  (S3D input, S3D rs)
  runBackwards (Deconvolution kernel _) (S3D input) (S3D dEdy) =
    let ex = extract input
        ix = fromIntegral $ natVal (Proxy :: Proxy inputRows)
        iy = fromIntegral $ natVal (Proxy :: Proxy inputCols)
        kx = fromIntegral $ natVal (Proxy :: Proxy kernelRows)
        ky = fromIntegral $ natVal (Proxy :: Proxy kernelCols)
        sx = fromIntegral $ natVal (Proxy :: Proxy strideRows)
        sy = fromIntegral $ natVal (Proxy :: Proxy strideCols)
        ox = fromIntegral $ natVal (Proxy :: Proxy outputRows)
        oy = fromIntegral $ natVal (Proxy :: Proxy outputCols)

        c  = vid2col 1 1 1 1 ix iy ex

        eo = extract dEdy
        ek = extract kernel

        vs = vid2col kx ky sx sy ox oy eo

        kN = fromJust . create . tr $ tr c LA.<> vs

        dW = vs LA.<> ek

        xW = col2vid 1 1 1 1 ix iy dW
    in  (Deconvolution' kN, S3D . fromJust . create $ xW)


{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE ExistentialQuantification #-}
{-# LANGUAGE FlexibleContexts #-}
-----------------------------------------------------------------------------
--
-- Module      :  AI.Network.RNN.Genetic
-- Copyright   :  (c) JP Moresmau
-- License     :  BSD3
--
-- Maintainer  :  JP Moresmau <jp@moresmau.fr>
-- Stability   :  experimental
-- Portability :
--
-- | Code for genetic evolution of a network
--
-----------------------------------------------------------------------------

module AI.Network.RNN.Genetic where

import Control.Monad
import Control.Monad.Random hiding (fromList)
import qualified Control.Monad.Random as MR
import Control.DeepSeq
import Data.IORef
import AI.GeneticAlgorithm.Simple
import AI.Network.RNN.Types
import AI.Network.RNN.Util
import Numeric.LinearAlgebra.HMatrix

import Debug.Trace

-- | Keep all relevant data together
data RNNData a b = (RNNEval a) => RNNData a !(TrainData b) !Double
   -- deriving (Show,Read,Eq)

-- | evaluate network
instance NFData (RNNData a b) where
    rnf (RNNData rnn _ _) = rnf rnn

-- | Error threshold (not currently used)
errorThreshold :: Double
errorThreshold = 0.1

-- | Chromosome instance
instance Chromosome (RNNData a b) where
    -- | cross network by taking half the values of one and half of the other
    crossover (RNNData rnn1 td l) (RNNData rnn2 _ _) = do
        rnns <- crossNetworkFull rnn1 rnn2
        --idx <- getRandomR (0, size rnn1 - 1)
        --let (v3,v4) = crossVector rnn1 rnn2 idx
        return $ map (\r->RNNData r td l) rnns -- [v3,v4]

    mutation (RNNData rnn1 td l) = do
        rnn2 <- mutateNetwork rnn1
        --rnn2 <- stdNormalMutate rnn1
        return $ RNNData rnn2 td l
    -- | fitness is "higher is better", inverse of cost
    fitness (RNNData rnn1 td _) = 1 / cost rnn1 td -- (fromVector sz rnn1)

-- | Cross 2 networks by taking the first half of one, the second half of the other (not currently used)
crossNetworkFull :: (RandomGen g,RNNEval a) =>  a -> a -> Rand g [a]
crossNetworkFull rnn1 rnn2 = do
    let v1 = toVector rnn1
        v2 = toVector rnn2
    idx <- getRandomR (0, size v1 - 1)
    let (v3,v4) = crossVector v1 v2 idx
        rnn3 = fromVector (rnnsize rnn1) v3
        rnn4 = fromVector (rnnsize rnn1) v4
    return $[rnn3,rnn4]

-- | cross 2 networks by taking half the values of one and half of the other
crossNetworkHalf :: (RandomGen g,RNNEval a) =>  a -> a -> Rand g [a]
crossNetworkHalf rnn1 rnn2 = do
    let v1 = toVector rnn1
        v2 = toVector rnn2
    (v3,v4) <- mixVector v1 v2 0.5
    let
        rnn3 = fromVector (rnnsize rnn1) v3
        rnn4 = fromVector (rnnsize rnn1) v4
    return $ force [rnn3,rnn4]

crossMutate :: (RandomGen g,RNNEval a) =>  a -> a -> Rand g [a]
crossMutate rnn1 rnn2 = do
    let v1 = toVector rnn1
        v2 = toVector rnn2
    v3 <- stdNormalMutate v1
    v4 <- stdNormalMutate v2
    let
        rnn3 = fromVector (rnnsize rnn1) v3
        rnn4 = fromVector (rnnsize rnn1) v4
    return [rnn3,rnn4]

-- | Cross networks by taking the average of their values (not currently used)
avgNetwork :: (RandomGen g,RNNEval a) =>  a -> a -> Rand g [a]
avgNetwork rnn1 rnn2 = do
    let v1 = toVector rnn1
        v2 = toVector rnn2
        v3 = fromList $ zipWith (\x y -> (x+y)/2) (toList v1) (toList v2)
        rnn3 = fromVector (rnnsize rnn1) v3
      --  rnn4 = fromVector (rnnsize rnn1) v4 of
      --      Right n -> n
      --      Left e  -> error e
    return [rnn3]

crossMatrixEq :: Matrix Double -> Matrix Double -> (Matrix Double,Matrix Double)
crossMatrixEq m1 m2 =
    let row1      = upperHalf $ rows m1
        col1      = upperHalf $ cols m1
        [[a1,a2],[a3,a4]] = toBlocksEvery row1 col1 m1
        [[b1,b2],[b3,b4]] = toBlocksEvery row1 col1 m2
    in (fromBlocks [[a1,b2],[a3,b4]],fromBlocks [[b1,a2],[b3,a4]])
    where
        upperHalf a = (if odd a then a+1 else a) `div` 2

crossVectorEq :: Vector Double -> Vector Double -> (Vector Double,Vector Double)
crossVectorEq v1 v2 = crossVector v1 v2 (size v1 `div` 2)

-- | Cross vector at the given point
crossVector :: Vector Double -> Vector Double -> Int -> (Vector Double,Vector Double)
crossVector v1 v2 idx =
--    let m1 = asRow v1
--        m2 = asRow v2
--        [[a1,a2]] = toBlocks [1] [idx,size v1-idx]  m1
--        [[b1,b2]] = toBlocks [1] [idx,size v1-idx] m2
--    in (head $ toRows $ fromBlocks [[a1,b2]],head $ toRows $ fromBlocks [[b1,a2]])
--    let sz2 = size v2 - idx - 1
--        [a1,a2] = takesV [idx,sz2] v1
--        [b1,b2] = takesV [idx,sz2] v2
--    in (vjoin [a1,b2],vjoin [a2,b1])
 let sz = size v1-idx
 in force $ (vjoin [subVector 0 idx v1,subVector idx sz v2],
                    vjoin [subVector 0 idx v2,subVector idx sz v1])

-- | Mix vector values given a probability to take a value from the first one and not the second
mixVector :: (Monad m,RandomGen g) =>  Vector Double -> Vector Double -> Double -> RandT g m (Vector Double,Vector Double)
mixVector v1 v2 prob = do
    -- rs <- replicateM (size v1) (getRandomR (0, 1))
    let (l1,l2) = swapL (toList v1) (toList v2) True
        -- unzip $ map swapR $ zip3 (toList v1) (toList v2) [0..]
    return (fromList l1,fromList l2)
    where
        swapR (a,b,idx) = if (idx `mod` 2) == 0 then (a,b) else (b,a)
        --if idx <= prob then (a,b) else (b,a):
        swapL :: [a] -> [a] -> Bool -> ([a],[a])
        swapL [] _ _ = ([],[])
        swapL (x:xs) (y:ys) True = let
            (rs1,rs2)=swapL xs ys False in (x:rs1,y:rs2)
        swapL (x:xs) (y:ys) False = let
            (rs1,rs2)=swapL xs ys True in (y:rs1,x:rs2)

-- | Mix vector values given a probability to take a value from the first one and not the second
mixList :: (Monad m,RandomGen g) =>  [Double] -> [Double] -> Double -> RandT g m ([Double],[Double])
mixList v1 v2 prob = do
    --rs <- getRandomRs (0,1)
    return $ unzip $ map swapR $ zip3 v1 v2 [0..]
    where
        swapR (a,b,idx) = if (idx `mod` 2) == 0 then (a,b) else (b,a)

stdNormalMutate :: (Monad m,RandomGen g) =>  Vector Double -> RandT g m (Vector Double)
stdNormalMutate v1 = do
    --rs <- replicateM (size v1) stdNormal
    -- let mx = maxElement v1
    -- let mn = minElement v1
    --let sz= trace ("v1:"++ show (mn,mx)) $ size v1
    seed <- getRandom
    let rs = randomVector seed Gaussian $ size v1
    let v2 = (rs) + v1
--    let nv2 = cmap (\x->
--                            if x < 0
--                                then 0
--                                else if x > 1
--                                    then 1
--                                    else x)
--         v2
    -- let mx2 = maxElement v2
    -- let mn2 = minElement v2
    return v2
    -- return $ trace ("v2:"++show (mn2,mx2)) v2

-- | network mutation: select randomly from several mutation algorithms
mutateNetwork :: (RandomGen g,RNNEval a) =>  a -> Rand g a
mutateNetwork rnn = do
    f <- MR.fromList [(stdNormalMutation,1),(pointMutation,1),(swapMutation,1),(insertMutation,1),(flipMutation,0.5)]
    f rnn

-- | Mutate values around a standard deviation
stdNormalMutation :: (RandomGen g,RNNEval a) =>  a -> Rand g a
stdNormalMutation rnn = do
    let v1 = toVector rnn
    v2 <- stdNormalMutate v1
    return $ force $ fromVector (rnnsize rnn) v2

-- | Mutate a single value of the vector, taking a random value
pointMutation :: (RandomGen g,RNNEval a) =>  a -> Rand g a
pointMutation rnn = do
    let v1 = toVector rnn
    idx <- getRandomR (0, size v1 - 1)
    dbl <- getRandomR (0,1)
    let v2 = accum v1 const [(idx,dbl)]
    return $ fromVector (rnnsize rnn) v2

-- | Swap two values in the vector
swapMutation :: (RandomGen g,RNNEval a) =>  a -> Rand g a
swapMutation rnn = do
    let v1 = toVector rnn
    idx1 <- getRandomR (0, size v1 - 1)
    idx2 <- getRandomR (0, size v1 - 1)
    let v2 = accum v1 const [(idx1,atIndex v1 idx2),(idx2,atIndex v1 idx1)]
    return $ fromVector (rnnsize rnn) v2

-- | Insert a random value at a random point in the vector, discarding the last value
insertMutation :: (RandomGen g,RNNEval a) =>  a -> Rand g a
insertMutation rnn = do
    let v1 = toVector rnn
    idx <- getRandomR (0, size v1 - 1)
    dbl <- getRandomR (0,1)
    let v2 = vjoin [subVector 0 idx v1,fromList [dbl],subVector idx (size v1 - idx -1) v1]
    return $ fromVector (rnnsize rnn) v2

flipMutation :: (RandomGen g,RNNEval a) =>  a -> Rand g a
flipMutation rnn = do
    let v1 = toVector rnn
    let v2 = cmap (\a->(-a)) v1
    return $ fromVector (rnnsize rnn) v2

-- | Stop function
-- we take a list of past fitnesses and stop when the fitness 5 generations ago was not worse that the current one
stopf :: IORef [Double] -> Int -> (RNNData a b,Double) -> Int -> IO Bool
stopf fitnessList maxGen (nd,fit) gen = do
    mfit <- atomicModifyIORef' fitnessList (\l->
        let l2 = take 5 (fit:l)
        in if length l2 == 5
            then (l2, Just $ last l2)
            else (l2, Nothing)
        )

    print $ "Fitness (" ++ show gen ++"): " ++ show fit
    let converged = case mfit of
                        Nothing -> False
                        Just f  -> f >= fit
    when converged $ print "Converged!"
    return ( gen >= maxGen)

-- | Build a random network data
buildNetworkData :: (RNNEval re,Monad m,RandomGen g) => TrainData b -> Size re -> FullSize re -> RandT g m (RNNData re b)
buildNetworkData td@(TrainData is _ _ _) dim fullSz = do
  --n <- randomNetwork dim fullSz
  let n = startNetwork dim fullSz
  return $ RNNData n td (fromIntegral $ length is)


{-# LANGUAGE Rank2Types #-}
module CreateGalaxy
where

import Data.List (sort)
import Control.Monad.State
import qualified Data.Map as M
import qualified Data.Edison.Assoc.StandardMap as E

import Galaxy
import DataFunction
import Statistics
import Math
import Utils

gasGiantByMass :: Flt -> PlanetType
gasGiantByMass mass | mass < 20.0  = SmallGasGiant
                    | mass < 100.0 = MediumGasGiant
                    | mass < 400.0 = LargeGasGiant
                    | otherwise    = VeryLargeGasGiant

createPlanetType :: Flt -> Temperature -> Flt -> Rnd PlanetType
createPlanetType mass startemp orbitradius = do
  if mass < 0.001 then return Planetoid
   else if mass < 0.01 then return NoAtmosphere
   else if mass > 15.0 then return (gasGiantByMass mass)
   else createRockyPlanetAtmosphere mass startemp orbitradius

createRockyPlanetAtmosphere :: Flt -> Temperature -> Flt -> Rnd PlanetType
createRockyPlanetAtmosphere mass startemp orbitradius = do
  weather <- randomRM (1, 100 :: Int)
  let atm1 = if weather < 10
               then RockyPlanet WaterWeatherSystem
               else if weather < 30 then RockyPlanet MethaneWeatherSystem
               else if weather < 50 then RockyPlanet SulphurDioxide
               else if weather < 90 then RockyPlanet CarbonDioxide
               else RockyPlanet Nitrogen
  let ptemp = planetTemperature'' startemp orbitradius atm1
  return $! case atm1 of
    RockyPlanet WaterWeatherSystem   -> if ptemp < waterMaxTemperature && ptemp > waterMinTemperature then atm1 else RockyPlanet Nitrogen
    RockyPlanet MethaneWeatherSystem -> if ptemp < 120 && ptemp > 70  then atm1 else RockyPlanet Nitrogen
    _                                -> atm1

createSatellite :: Flt -> Flt -> (Planet () -> Rnd a) -> Temperature -> String -> Flt -> Rnd (Planet a)
createSatellite minmass maxmass genfunc startemp name orbitradius = do
  orbit <- createOrbit orbitradius
  mass <- randomRM (minmass, maxmass)
  atmosphere <- createPlanetType mass startemp orbitradius
  let ptemp = planetTemperature'' startemp orbitradius atmosphere
  let emptyplanet = Planet name orbit (BodyPhysics mass) atmosphere ptemp M.empty ()
  cont <- genfunc emptyplanet
  return $! Planet name orbit (BodyPhysics mass) atmosphere ptemp M.empty cont

createPlanet :: (Planet () -> Rnd a) -> Temperature -> String -> Flt -> Rnd (Planet a)
createPlanet genfunc startemp name orbitradius = do
  orbit <- createOrbit orbitradius
  gentype <- randomRM (1, 4 :: Int)
  mass <- case gentype of
            1 -> randomRM (50, 400)
            2 -> randomRM (20, 50)
            3 -> randomRM (0.0001, 0.1)
            _ -> randomRM (0.1, 20)
  atmosphere <- createPlanetType mass startemp orbitradius
  numsatellites <- if mass < 1.0 then return 0 else randomRM (0 :: Int, min 16 (floor (sqrt mass)))
  satellites <- zipWithM (createSatellite (0.00001 * mass) (0.01 * mass) genfunc startemp) (bodyNames name) (replicate numsatellites orbitradius)
  let temp = planetTemperature'' startemp orbitradius atmosphere
  cont <- genfunc (Planet name orbit (BodyPhysics mass) atmosphere temp M.empty ())
  return $! Planet name orbit (BodyPhysics mass) atmosphere temp (namedsToMap satellites) cont

bodyNames :: String -> [String]
bodyNames = namesFromBasenameNum

starprobs, starbinaryprobs :: [(Flt, SpectralType)]
starprobs       = [(0.001, SpectralTypeB), (0.004, SpectralTypeA), (0.02, SpectralTypeF), (0.06, SpectralTypeG), (0.12, SpectralTypeK), (1, SpectralTypeM)]
starbinaryprobs = [(0.001, SpectralTypeB), (0.010, SpectralTypeA), (0.10, SpectralTypeF), (0.25, SpectralTypeG), (0.40, SpectralTypeK), (1, SpectralTypeM)]

getProbTableValue :: (Ord a) => a -> [(a, b)] -> b
getProbTableValue n [(_, v)] = v
getProbTableValue n ((k, v):xs) = if n <= k then v else getProbTableValue n xs
getProbTableValue _ []       = error "Given probability >1?"

randomStarTemperature :: SpectralType -> Rnd Temperature
randomStarTemperature s = randomRM (specTempRange s)

createStar :: (Planet () -> Rnd a) -> String -> Flt -> Orbit -> Rnd (Star a)
createStar genfunc name maxplanetorbitradius orbit = do
  r <- randomRM (0, 1)
  let s = getProbTableValue r starprobs
  createStar' genfunc name maxplanetorbitradius orbit s

addPrefixToPlanetName :: String -> Planet a -> Planet a
addPrefixToPlanetName n p =
  let oname = planetname p
      nname = n ++ oname
  in p{planetname = nname, satellites = E.map (addPrefixToPlanetName n) (satellites p)}

createStar' :: (Planet () -> Rnd a) -> String -> Flt -> Orbit -> SpectralType -> Rnd (Star a)
createStar' genfunc name maxplanetorbitradius orbit spectraltype = do
  t <- randomStarTemperature spectraltype
  numplanets <- randomRM (64, 128)
  let planetnames = bodyNames name
  planetorbitradiuses <- separate `fmap` sort `fmap` replicateM numplanets (randomRM (0.0001, min (fromIntegral t * 2 / 120) maxplanetorbitradius))
  planets <- zipWithM (createPlanet genfunc t) (repeat "") planetorbitradiuses
  let planets' = filter (\p -> planetTemperature' t p > 30 && planetTemperature' t p < t `div` 4) planets
  let planets'' = zipWith addPrefixToPlanetName planetnames planets'
  return $! Star name t orbit (namedsToMap planets'')

namesFromBasenameCap :: String -> [String]
namesFromBasenameCap n = zipWith (++) (repeat (n ++ " ")) (map (:[]) ['A' .. 'Z'])

namesFromBasenameMin :: String -> [String]
namesFromBasenameMin n = zipWith (++) (repeat (n ++ " ")) (map (:[]) ['a' .. 'z'])

namesFromBasenameNum :: String -> [String]
namesFromBasenameNum n = zipWith (++) (repeat (n ++ " ")) (map show [(1 :: Int) ..])

createOrbit :: Flt -> Rnd Orbit
createOrbit oradius = return $! Orbit oradius 0 0 1 0

noOrbit :: Orbit
noOrbit = Orbit 0 0 0 1 0

createStars :: (Planet () -> Rnd a) -> String -> Int -> Rnd [Star a]
createStars genfunc basename 1 = do
  s <- createStar genfunc basename 3000 noOrbit
  return [s]
createStars genfunc basename numstars = do
  let names = namesFromBasenameCap basename
  orbitradiuses <- sort `fmap` replicateM numstars (randomRM (0.1, 12000))
  orbits <- mapM createOrbit orbitradiuses
  let dists = map (/4) (distances orbitradiuses)
  spectraltypes <- map (flip getProbTableValue starbinaryprobs) `fmap` replicateM numstars (randomRM (0, 1))
  stars <- zipWith4M (createStar' genfunc) names dists orbits (sort spectraltypes)
  return $! stars

createStarSystem :: (Planet () -> Rnd a) -> String -> Vector3 -> Rnd (StarSystem a)
createStarSystem genfunc ssname sspos = do
  -- http://www.cfa.harvard.edu/news/2006/pr200611.html
  -- "Most milky way stars are single"
  singular <- chance 2 3
  n <- if singular then return 1 else randomRM (2, 6 :: Int)
  stars <- (createStars genfunc) ssname n
  return $! StarSystem ssname sspos (namedsToMap stars)

create2DPoint :: (Flt, Flt) -> Rnd Vector3
create2DPoint (minc, maxc) = do
  x <- randomRM (minc, maxc)
  y <- randomRM (minc, maxc)
  return (x, y, 0)

create3DPoint :: (Flt, Flt) -> Rnd Vector3
create3DPoint (minc, maxc) = do
  x <- randomRM (minc, maxc)
  y <- randomRM (minc, maxc)
  z <- randomRM (minc, maxc)
  return (x, y, z)

ssSpacingCoefficient :: Float
ssSpacingCoefficient = 5

createGalaxy :: (Planet () -> Rnd a) -> String -> [String] -> Rnd (Galaxy a)
createGalaxy genfunc galname ssnames = do
  let numss = length ssnames
  let dim = sqrt (fromIntegral numss) * ssSpacingCoefficient -- TODO: when 3d galaxy, use cbrt
  points <- replicateM numss (create3DPoint (-dim, dim))
  sss <- zipWithM (createStarSystem genfunc) ssnames points
  return $! Galaxy galname (namedsToMap sss)


module BondCalculus.ODEExtraction
  (IVP(..), PrintStyle(..), ODE(..), AsSage(..),
  matlabExpr, sympyExpr, sageExpr, matlabODE, vectorFieldToODEs, sageODE, sageExprAbst,
  sageExprAbst', extractIVP, sympyODE, solveODEPython, printODEPython,
  sympySimplify, runPython, generateSage) where

import BondCalculus.Base
import BondCalculus.Symbolic
import BondCalculus.Processes
import BondCalculus.Vector
import BondCalculus.AST (pretty)
import qualified BondCalculus.AST as AST
import qualified Data.HashMap.Strict as H
import BondCalculus.Simulation (Trace)
import Data.String.Utils
import Text.Printf

-- import qualified Control.Exception as X
import qualified Data.Map as M
import qualified Data.List as L
import Data.Maybe

--import qualified Numeric.LinearAlgebra as LA

import System.IO.Unsafe
-- import qualified System.Process as OS
import System.Process (readProcess)
-- import Data.Maybe

newtype IVP a = IVP ([String], [Expr (Atom a)], [a]) deriving (Show, Eq)
newtype ODE a = ODE ([String], [Expr (Atom a)])

-------------------------------
-- MATLAB script output:
-------------------------------

extractODE :: (ExprConstant a) => AST.Env -> ConcreteAffinityNetwork a -> P' a -> ODE a
extractODE env network p = vectorFieldToODEs p'
  where -- p' must be expressed in the same basis of p
        -- otherwise the variable order becomes
        -- inconsistient, and we get extra variables
        -- p' :: Vect AST.Species (SymbolicExpr a)
        p' = fromList [(p'coeff i, i) | (_,i) <- toList p]
        -- p'coeff :: AST.Species -> SymbolicExpr a
        p'coeff i = fromMaybe 0.0 (H.lookup i v)
        Vect v = dPdt' tr network p
        tr = tracesGivenNetwork network env

extractIVP :: ExprConstant a => AST.Env -> ConcreteAffinityNetwork a -> P' a -> [a] -> IVP a
extractIVP env network p = fromODEToIVP (extractODE env network p)

vectorFieldToODEs :: ExprConstant a => P' a -> ODE a
vectorFieldToODEs v = ODE (vars, rhss)
  where
    vars = map (pretty.snd) kis
    rhss = map (simplify.fst) kis
    kis = toList v

fromODEToIVP :: ODE a -> [a] -> IVP a
fromODEToIVP (ODE (x,y)) z = IVP (x, y, take (length x) z)

-- modelToIVP :: BondCalculusModel SymbolicExpr Double -> IVP

matlabODE :: IVP Double -> (Int,(Double,Double)) -> Either String String
matlabODE ivp (n,(t0,tn)) =
  if any isNothing eqns
  then Left "An ODE equation has an unbound variable"
  else Right $
  "init = [" ++ L.intercalate ";" (map show inits) ++ "];\n" ++
    "t = linspace(" ++ show t0 ++ "," ++ show tn ++ "," ++ show n ++ ");\n" ++
    "function xdot = f(x,t) \n\n" ++
    L.intercalate ";\n" (catMaybes eqns) ++
    ";\nendfunction;\n\n" ++
    -- Need new support for calculating jacobian
    -- "function jac = jj(x,t)\n\n" ++
    -- matlabJac env p' ++ ";\nendfunction;\n\n" ++
    "x = lsode(f, init, t);\n" ++
    "save (\"-ascii\", \"-\", \"x\");"
      where
        xdots = ["xdot(" ++ show i ++ ")" | i <- [(1::Integer)..]]
        xvars = ["x(" ++ show i ++ ")" | i <- [(1::Integer)..]]
        varmap = M.fromList $ zip vars xvars
        IVP (vars,rhss,inits) = ivp
        eqns = [fmap (\y -> xdot ++ " = " ++ y) (matlabExpr varmap rhs)
               | (xdot,rhs) <- zip xdots rhss]

class Show a => AsSage a where
    asSage :: a -> String

instance AsSage Double where
    asSage = show

instance AsSage Interval where
    asSage x = printf "RIF('[%.20f .. %.20f]')" l u
        where (l, u) = endpoints x
              -- Do some fiddling with the endpoints to ensure sound enclosure
            --   l' = l - 6e-21
            --   u' = u + 6e-21

sageExpr :: AsSage a => M.Map String String -> Expr (Atom a) -> Maybe String
sageExpr mp (Atom (Var x)) = M.lookup x mp
sageExpr mp (Atom (Const x)) = return $ asSage x
sageExpr mp (x `Sum` y) = do
    x' <- sageExpr mp x
    y' <- sageExpr mp y
    return $ x' ++ " + " ++ y'
sageExpr mp (x `Prod` y) = do
    x' <- sageExpr mp x
    y' <- sageExpr mp y
    return $ "(" ++ x' ++ ") * (" ++ y' ++ ")"
sageExpr mp (x `Frac` y) = do
    x' <- sageExpr mp x
    y' <- sageExpr mp y
    return $ "(" ++ x' ++ ") / (" ++ y' ++ ")"
sageExpr mp (x `Pow` y) = do
    x' <- sageExpr mp x
    y' <- sageExpr mp y
    return $ "(" ++ x' ++ ") ** (" ++ y' ++ ")"
sageExpr mp (Abs x) = do
    x' <- sageExpr mp x
    return $ "abs(" ++ x' ++ ")"
sageExpr mp (Sin x) = do
    x' <- sageExpr mp x
    return $ "sin(" ++ x' ++ ")"
sageExpr mp (Cos x) = do
    x' <- sageExpr mp x
    return $ "sin(" ++ x' ++ ")"
sageExpr mp (Tan x) = do
    x' <- sageExpr mp x
    return $ "tan(" ++ x' ++ ")"
sageExpr mp (Exp x) = do
    x' <- sageExpr mp x
    return $ "exp(" ++ x' ++ ")"
sageExpr mp (Log x) = do
    x' <- sageExpr mp x
    return $ "log(" ++ x' ++ ")"

-- Convert a sage expression into a string with all non-simple double constants (e.g. intervals)
-- abstracted by symbolic variables
sageExprAbst' :: (AsSage a, Boundable a) => M.Map String String -> Expr (Atom a) -> Maybe (M.Map String String, String)
sageExprAbst' = sageExprAbst M.empty

sageExprAbst :: (AsSage a, Boundable a) => M.Map String String -> M.Map String String -> Expr (Atom a) -> Maybe (M.Map String String, String)
sageExprAbst cmp mp (Atom (Var x)) = do
    x' <- M.lookup x mp
    return (cmp, x')
sageExprAbst cmp mp (Atom (Const x)) = return $ case singleValue x of
    Just 1.0    -> (cmp, "1")
    Just (-1.0) -> (cmp, "-1")
    Just 0.0    -> (cmp, "0")
    Just v      -> (cmp, asSage v)
    Nothing     -> (cmp', s)
        where n = M.size cmp
              s = "a" ++ show n
              cmp' = M.insert s (asSage x) cmp
sageExprAbst cmp mp (x `Sum` y) = f2 cmp mp x y $ \x' y' -> x' ++ " + " ++ y'
sageExprAbst cmp mp (x `Prod` y) = f2 cmp mp x y $ \x' y' -> "(" ++ x' ++ ") * (" ++ y' ++ ")"
sageExprAbst cmp mp (x `Frac` y) = f2 cmp mp x y $ \x' y' -> "(" ++ x' ++ ") / (" ++ y' ++ ")"
sageExprAbst cmp mp (x `Pow` y) = f2 cmp mp x y $ \x' y' -> "(" ++ x' ++ ") ** (" ++ y' ++ ")"
sageExprAbst cmp mp (Abs x) = f1 cmp mp x $ \x' -> "abs(" ++ x' ++ ")"
sageExprAbst cmp mp (Sin x) = f1 cmp mp x $ \x' -> "sin(" ++ x' ++ ")"
sageExprAbst cmp mp (Cos x) = f1 cmp mp x $ \x' -> "cos(" ++ x' ++ ")"
sageExprAbst cmp mp (Tan x) = f1 cmp mp x $ \x' -> "tan(" ++ x' ++ ")"
sageExprAbst cmp mp (Exp x) = f1 cmp mp x $ \x' -> "exp(" ++ x' ++ ")"
sageExprAbst cmp mp (Log x) = f1 cmp mp x $ \x' -> "log(" ++ x' ++ ")"
sageExprAbst cmp mp x = error $ "Support for converting expression " ++ show x ++ " to sage/sympy not implemented!"

f2 cmp mp x y g = do
    (cmp', x') <- sageExprAbst cmp mp x
    (cmp'', y') <- sageExprAbst cmp' mp y
    return (cmp'', g x' y')

f1 cmp mp x g = do
    (cmp', x') <- sageExprAbst cmp mp x
    return (cmp', g x')

-- This assumes we can reduce the ODEs to sage
sageODE :: (AsSage a, Boundable a) => IVP a -> String -> Either String String
sageODE (IVP (vars, rhss, inits)) network = case odeExprs of
        Nothing -> Left "An ODE equation has an unbound variable"
        Just (cmp, odes) -> Right $
            "from sage.all import *\n" ++
            "import sympy as sym\n\n" ++
            "R, x = PolynomialRing(RIF, " ++ show nvars ++ ", ', '.join(map('x{}'.format, range(0, " ++ show nvars ++ ")))).objgens()\n" ++
            "xsr = [SR.var(str(x1)) for x1 in x]\n" ++
            "for v in xsr:\n"++
            "   assume(v, 'real')\n"++
            "xsym = sym.var(','.join(map('x{}'.format, range(0," ++ show nvars ++ "))))\n" ++
            (if (M.size cmp) > 0
             then "asym = sym.var(','.join(map('a{}'.format, range(0," ++ show (M.size cmp) ++ "))))\n"
             else "asym = []\n") ++
            "y0 = [" ++ L.intercalate ", " (map asSage inits) ++ "]\n" ++
            "ysymraw = [" ++ L.intercalate ", " odes ++ "]\n" ++
            "ysym = [sym.simplify(y1) for y1 in ysymraw]\n" ++
            "ysr = [y1._sage_().substitute(" ++ subsExpr ++ ") for y1 in ysym]\n" ++
            "varmap = {" ++ L.intercalate ", " varmappings ++ "}\n" ++
            "varmapsr = {" ++ L.intercalate ", " varmappingssr ++ "}\n" ++
            "affinity_network = \"" ++ network ++ "\"\n" ++
            "try:\n" ++
            "    y = vector([R(y1) for y1 in ysr])\n" ++
            "except TypeError:\n" ++
            "    y = None\n"
            where subsExpr = L.intercalate ","
                             [x ++ "=" ++ y | (x, y) <- M.toList cmp]
                  varmappings = [ "'" ++ p ++ "': x[" ++ show i ++ "]"
                                | p <- vars | i <- [0..] ]
                  varmappingssr = [ "'" ++ p ++ "': xsr[" ++ show i ++ "]"
                                  | p <- vars | i <- [0..] ]
    where
    --   odeExprs' :: AsSage a => M.Map String String -> [Expr (Atom a)] -> Maybe (M.Map String String, [String])
      odeExprs' cmp (x:xs) = do
        (cmp', expr) <- sageExprAbst cmp varmap x
        (cmp'', exprs) <- odeExprs' cmp' xs
        return (cmp'', expr:exprs)
      odeExprs' cmp [] = Just (cmp, [])
      odeExprs :: Maybe (M.Map String String, [String])
      odeExprs = odeExprs' M.empty rhss
      nvars = length vars
      xvars = ["x" ++ show n | (n, _) <- zip [0..] vars]
      varmap = M.fromList $ zip vars xvars
    --   odes = map odeExpr rhss
    --   odeExpr rhs = do
    --     expr <- sageExpr varmap rhs
    --     return $ "sym.simplify(" ++ expr ++ ")"

sympyODE :: IVP Double -> (Int,(Double,Double)) -> Either String String
sympyODE ivp (n,(t0,tn)) =
  if any isNothing eqns
  then Left "An ODE equation has an unbound variable"
  else Right $
  "import sympy as sym\n" ++
  "import numpy as np\n" ++
  "from sympy.abc import t\n" ++
  "from scipy.integrate import odeint\n" ++
  "import sys\n" ++
  "sys.setrecursionlimit(100000)\n\n" ++
  "xs = [" ++ L.intercalate ","
    ["sym.Function('" ++ show x ++ "')" | x <- vars ] ++ "]\n" ++
  "xts = [x(t) for x in xs]\n" ++
  "odes = [" ++ L.intercalate ", " (catMaybes eqns) ++ "]\n" ++
  "y0 = [" ++ L.intercalate ", " (map show inits) ++ "]\n" ++
  "ts = np.linspace(" ++ show t0 ++ "," ++ show tn ++ "," ++ show n ++ ")\n" ++
  "rhss = [eqn.rhs for eqn in odes]\n" ++
  "Jac = sym.Matrix(rhss).jacobian(xts)\n" ++
  "f = sym.lambdify((xts, t), rhss, modules='numpy')\n" ++
  "J = sym.lambdify((xts, t), Jac, modules='numpy')\n" ++
  -- Try to solve ODEs using symbolic Jacobian
  "try: ys = odeint(f, y0, ts, (), J)\n" ++
  -- Fallback to using method without Jacobian
  "except NameError: ys = odeint(f, y0, ts, ())\n" ++
  -- "print(ys)\n" ++
  "print('\\n'.join([' '.join([('%.18e' % y).replace('nan', '0') for y in ysa]) for ysa in ys]))"
    where
      xvars = ["xts[" ++ show i ++ "]" | i <- [(0::Integer)..]]
      xdots = [xvar ++ ".diff()" | xvar <- xvars]
      varmap = M.fromList $ zip vars xvars
      IVP (vars,rhss,inits) = ivp
      mkeqn xdot y = "sym.Eq(" ++ xdot ++ ", sym.simplify(" ++ y ++ "))"
      eqns = [fmap (mkeqn xdot) (sympyExpr varmap rhs)
             | (xdot,rhs) <- zip xdots rhss]

data PrintStyle = Plain | Pretty | LaTeX | MathML deriving (Eq)

sympyODEPrint :: ODE Double -> PrintStyle -> Either String String
sympyODEPrint ode style =
  if any isNothing eqns
  then Left "An ODE equation has an unbound variable"
  else Right $
  "import sympy as sym\n" ++
  "import numpy as np\n" ++
  "from sympy.abc import t\n" ++
  "import sys\n" ++
  "sys.setrecursionlimit(100000)\n\n" ++
  -- "sym.init_printing(" ++ printingoptions ++ ")\n\n" ++
  "names = [" ++ L.intercalate "," ["'[" ++ (if style == LaTeX then replace "->" "\\\\rightarrow " (replace " " "\\\\ " x) else x) ++ "]'" | x <- vars] ++ "]\n" ++
  "xs = list(map(sym.Function, names))\n" ++
  "xts = [x(t) for x in xs]\n" ++
  "odes = [" ++ L.intercalate ", " (catMaybes eqns) ++ "]\n" ++
  "rhss = [eqn.rhs for eqn in odes]\n" ++
  case style of
    Pretty -> "for ode in odes: print(sym.pretty(ode))\n"
    Plain -> "for ode in odes: print(ode)\n"
    -- print system of odes as latex (formatted for use in an align environment)
    LaTeX -> "print(sym.latex(odes)[7:-7].replace('=', '& =').replace(', \\\\  ', ' \\\\\\\\\\n').replace('{\\\\left(t \\\\right)}', '').replace('\\\\frac{d}{d t}', '\\\\frac{\\\\mathrm d}{\\\\mathrm d t}').replace('operatorname', 'mathrm'))\n"
    MathML -> "from sympy.printing import print_mathml\nprint_mathml(odes)\n"
    where
      xvars = ["xts[" ++ show i ++ "]" | i <- [(0::Integer)..]]
      xdots = [xvar ++ ".diff()" | xvar <- xvars]
      varmap = M.fromList $ zip vars xvars
      ODE (vars,rhss) = ode
      mkeqn xdot y = "sym.Eq(" ++ xdot ++ ", sym.simplify(" ++ y ++ "))"
      eqns = [fmap (mkeqn xdot) (sympyExpr varmap rhs)
             | (xdot,rhs) <- zip xdots rhss]

matlabExpr :: M.Map String String -> SymbolicExpr Double -> Maybe String
matlabExpr mp (Atom (Var x)) = M.lookup x mp
matlabExpr _ (Atom (Const x)) = return $ show x
matlabExpr mp (x `Sum` y) = do
  x' <- matlabExpr mp x
  y' <- matlabExpr mp y
  return $ x' ++ " .+ " ++ y'
matlabExpr mp (x `Prod` y) = do
  x' <- matlabExpr mp x
  y' <- matlabExpr mp y
  return $ "(" ++ x' ++ ") .* (" ++ y' ++ ")"
matlabExpr mp (x `Pow` y) = do
  x' <- matlabExpr mp x
  y' <- matlabExpr mp y
  return $ "(" ++ x' ++ ") .** (" ++ y' ++ ")"
matlabExpr mp (Abs x) = do
  x' <- matlabExpr mp x
  return $ "abs(" ++ x' ++ ")"

sympyExpr :: M.Map String String -> SymbolicExpr Double -> Maybe String
sympyExpr mp (Atom (Var x)) = M.lookup x mp
sympyExpr _ (Atom (Const x)) = return $ show x
sympyExpr mp (x `Sum` y) = do
  x' <- sympyExpr mp x
  y' <- sympyExpr mp y
  return $ x' ++ " + " ++ y'
sympyExpr mp (x `Prod` y) = do
  x' <- sympyExpr mp x
  y' <- sympyExpr mp y
  return $ "(" ++ x' ++ ") * (" ++ y' ++ ")"
sympyExpr mp (x `Frac` y) = do
  x' <- sympyExpr mp x
  y' <- sympyExpr mp y
  return $ "(" ++ x' ++ ") / (" ++ y' ++ ")"
sympyExpr mp (x `Pow` y) = do
  x' <- sympyExpr mp x
  y' <- sympyExpr mp y
  return $ "(" ++ x' ++ ") ** (" ++ y' ++ ")"
sympyExpr mp (Abs x) = do
  x' <- sympyExpr mp x
  -- return $ x'
  return $ "abs(" ++ x' ++ ")"
sympyExpr mp (Sin x) = do
  x' <- sympyExpr mp x
  return $ "sin(" ++ x' ++ ")"
sympyExpr mp (Cos x) = do
  x' <- sympyExpr mp x
  return $ "sin(" ++ x' ++ ")"
sympyExpr mp (Tan x) = do
  x' <- sympyExpr mp x
  return $ "tan(" ++ x' ++ ")"
sympyExpr mp (Exp x) = do
  x' <- sympyExpr mp x
  return $ "exp(" ++ x' ++ ")"
sympyExpr mp (Log x) = do
  x' <- sympyExpr mp x
  return $ "log(" ++ x' ++ ")"

---------------------------------------
-- Using Octave to execute the scripts
---------------------------------------

runPython :: String -> IO String
runPython = readProcess "python3" ["-q"]

callSolveODEPython :: AST.Env -> ConcreteAffinityNetwork Double -> P' Double -> [Double] -> (Int, (Double, Double)) -> IO String
callSolveODEPython env network p inits tr = case sympyODE (extractIVP env network p inits) tr of
  Right script -> do
    putStrLn $ "Python script:\n\n" ++ script
    writeFile "script.py" script
    runPython script
  Left _ -> undefined

generateSage :: (ExprConstant a, AsSage a)
             => String
             -> AST.Env
             -> ConcreteAffinityNetwork a
             -> P' a
             -> [a]
             -> IO String
generateSage filename env network p inits = case sageODE (extractIVP env network p inits) (pretty network) of
  Right script -> do
    putStrLn $ "Sage script:\n\n" ++ script
    writeFile filename script
    return script
  Left _ -> undefined

callPrintODEPython :: AST.Env -> ConcreteAffinityNetwork Double -> P' Double -> PrintStyle -> IO String
callPrintODEPython env network p style = case sympyODEPrint (extractODE env network p) style of
  Right script -> do
    -- putStrLn $ "Python script:\n\n" ++ script
    writeFile "print_script.py" script
    runPython script
  Left _ -> undefined

solveODEPython :: AST.Env -> ConcreteAffinityNetwork Double -> P' Double -> [Double] -> (Int, (Double, Double)) -> Trace
solveODEPython env network p inits tr@(n,(t0,tn))
  = let raw = unsafePerformIO (callSolveODEPython env network p inits tr)
        ts = [t0 + fromIntegral i*(tn-t0)/fromIntegral n | i <- [0..n]]
        yss = (map (map read.words) $ lines raw) :: [[Double]]
        ys = [fromList (xs `zip` pbasis) | xs <- yss]
        pbasis = map snd $ toList p
    in ts `zip` ys

sympySimplify :: String -> String
sympySimplify s = unsafePerformIO (runPython script)
  where script = "from sympy import simplify, sympify\nimport re\n" ++
                 "print(simplify(sympify(re.sub(r'([0-9]+)\\.0(?![0-9]*[1-9])', r'\\1', " ++ show s ++ "))))"

printODEPython :: AST.Env -> ConcreteAffinityNetwork Double -> P' Double -> PrintStyle -> String
printODEPython env network p style
  = let raw = unsafePerformIO (callPrintODEPython env network p style)
    in raw
-- callOctave env p mts p' ts = let
--     script = matlabODE env (wholeProc env p mts) p' ts
--   in do
--     putStrLn $ "Octave Script: \n" ++ script
--     OS.readProcess
--       "octave" ["-q", "--eval", script] []

-- | Solver which calculates the symbolic Jacobian, writes MATLAB code, and executes it with GNU Octave. (General purpose, deals with stiff systems, uses LSODE.)
-- solveODEoctave :: Solver
-- solveODEoctave env p mts p' ts@(n,(t0,tn))
--     = let raw = unsafePerformIO (callOctave env p mts p' ts)
--       in (n>< Map.size p') $ map s2d $ words raw

-- Return the MATLAB script for ODEs
-- matlabScript :: Env
--              -> Process
--              -> MTS
--              -> P'
--              -> (Int, (Double, Double))
--              -> String
-- matlabScript env p mts = matlabODE env (wholeProc env p mts)


{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE ViewPatterns     #-}
{-|
Module      : Grenade.Layers.Internal.Hmatrix
Description : Faster implementations of HMatrix functions
Maintainer  : Theo Charalambous
License     : BSD2
Stability   : experimental
-}
module Grenade.Layers.Internal.Hmatrix (
    reshapeMatrix
  ) where

import           Numeric.LinearAlgebra
import           Numeric.LinearAlgebra.Devel

-- | Reshape a row major matrix
reshapeMatrix :: (Element t, Num t, Container Vector t) => Int -> Int -> Matrix t -> Matrix t
reshapeMatrix r c m@(size->(r', c'))
    | r * c == r' * c' = matrixFromVector RowMajor r c $ flatten m
    | otherwise        = error $ "can't reshape matrix of shape dim = " ++ show (r', c') ++ " to matrix of shape " ++ show (r, c)
{-# INLINE reshapeMatrix #-}


module Tests.SpecFunctions where

import Test.QuickCheck  hiding (choose,within)
--import Test.Framework
--import Test.Framework.Providers.QuickCheck2
--import Numeric.SpecFunctions


{-# LANGUAGE FlexibleContexts    #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE Strict              #-}
module Pinwheel.FourierSeries2D where

import           Control.Monad                  as M
import           Control.Monad.IO.Class
import           Control.Monad.Trans.Resource
import qualified Data.Array.Accelerate          as A
import           Data.Array.Accelerate.LLVM.PTX as A
import           Data.Array.IArray              as IA
import           Data.Array.Repa                as R
import           Data.Complex
import           Data.Conduit                   as C
import           Data.Conduit.List              as CL
import           Data.List                      as L
import           Data.Vector.Generic            as VG
import           Data.Vector.Storable           as VS
import           Data.Vector.Unboxed            as VU
import           Foreign.CUDA.Driver            as CUDA
import           FourierMethod.FourierSeries2D
import           Math.Gamma
import           Pinwheel.Base
import           Pinwheel.List
import           Utils.BLAS
import           Utils.Distribution
import           Utils.List
import           Utils.Parallel                 hiding ((.|))
import           Utils.SimpsonRule
import           Utils.Time
import DFT.Plan
import Filter.Utils
import Control.Concurrent.Async
import Numeric.GSL.Special.Bessel

pinwheelFourierSeries ::
     ( Storable e
     , CUBLAS (Complex e)
     , Unbox e
     , RealFloat e
     , A.Elt e
     , A.Elt (Complex e)
     , Floating (A.Exp e)
     , A.FromIntegral Int e
     , VG.Vector vector (Complex e)
     , NFData (vector (Complex e))
     )
  => [Int]
  -> [PTX]
  -> Int
  -> Int
  -> e
  -> e
  -> Int
  -> Int
  -> Int
  -> Int
  -> e
  -> Int
  -> Int
  -> IO (IA.Array (Int, Int) (vector (Complex e)))
pinwheelFourierSeries deviceIDs ptxs numR2Freqs numPoints delta periodR2 phiFreq rhoFreq thetaFreq rFreq sigma numBatchR2 numBatchPinwheelFreqs = do
  let idxs =
        [ (radialFreq, angularFreq)
        | radialFreq <- pinwheelFreqs rhoFreq rFreq
        , angularFreq <- pinwheelFreqs phiFreq thetaFreq
        ]
      centerR2Freq = div numR2Freqs 2
  pinwheels <-
    computeUnboxedP .
    R.traverse
      (fromListUnboxed (Z :. (L.length idxs)) idxs)
      (\(Z :. freqs) -> (Z :. freqs :. numR2Freqs :. numR2Freqs)) $ \f (Z :. pFreq :. xFreq :. yFreq) ->
      fourierMellin
        sigma
        (snd $ f (Z :. pFreq))
        (fst $ f (Z :. pFreq))
        ( delta * fromIntegral (xFreq - centerR2Freq)
        , delta * fromIntegral (yFreq - centerR2Freq))
  -- pinwheels <-
  --   M.mapM
  --     (\idx ->
  --        fmap
  --          (CuMat (numR2Freqs ^ 2) (L.length idx) .
  --           CuVecHost . VU.convert . toUnboxed) .
  --        computeP .
  --        R.traverse
  --          (fromListUnboxed (Z :. (L.length idx)) idx)
  --          (\(Z :. freqs) -> (Z :. numR2Freqs :. numR2Freqs :. freqs)) $ \f (Z :. xFreq :. yFreq :. pFreq) ->
  --          fourierMellin
  --            sigma
  --            (snd $ f (Z :. pFreq))
  --            (fst $ f (Z :. pFreq))
  --            ( fromIntegral (xFreq - centerR2Freq)
  --            , fromIntegral (yFreq - centerR2Freq))) .
  --   divideListN numBatchPinwheelFreqs $
  --   idxs
  -- printCurrentTime "Compute Fourier Series... "
  -- arr <-
  --   computeFourierSeriesR2Stream
  --     deviceIDs
  --     ptxs
  --     numR2Freqs
  --     numPoints
  --     periodR2
  --     delta
  --     numBatchR2
  --     pinwheels
  -- printCurrentTime "Compute Fourier Series Done"
  let (radialLB, radialUB) = pinwheelFreqsBound rhoFreq rFreq
      (angularLB, angularUB) = pinwheelFreqsBound phiFreq thetaFreq
  return .
    listArray ((radialLB, angularLB), (radialUB, angularUB)) .
    parMap
      rdeepseq
      (\i ->
         VG.convert . toUnboxed . computeS . R.slice pinwheels $  -- arr $
         (Z :. i :. All :. All)) $
    [0 .. (L.length idxs - 1)]

-- Stream
{-# INLINE source #-}
source :: Int -> Int -> Int -> Int -> ConduitT () (Int, Int) (ResourceT IO) ()
source phiFreq rhoFreq thetaFreq rFreq =
  CL.sourceList
    [ (radialFreq, angularFreq)
    | radialFreq <- pinwheelFreqs rhoFreq rFreq
    , angularFreq <- pinwheelFreqs phiFreq thetaFreq
    ]

conduit ::
     ( Storable e
     , CUBLAS (Complex e)
     , Unbox e
     , RealFloat e
     , A.Elt e
     , A.Elt (Complex e)
     , Floating (A.Exp e)
     , A.FromIntegral Int e
     , VG.Vector vector (Complex e)
     )
  => [Int]
  -> Int
  -> Int
  -> e
  -> Int
  -> e
  -> [[CuMat (Complex e)]]
  -> ConduitT (Int, Int) [vector (Complex e)] (ResourceT IO) ()
conduit deviceIDs numR2Freqs numPoints periodR2 batchSize sigma inverseR2Harmonics = do
  liftIO $ printCurrentTime "conduit"
  idx <- CL.take batchSize
  unless
    (L.null idx)
    (do let centerR2Freq = div numR2Freqs 2
        pinwheels <-
          fmap
            (CuMat (L.length idx) (numR2Freqs ^ 2) .
             CuVecHost . VU.convert . toUnboxed) .
          liftIO .
          computeP .
          R.traverse
            (fromListUnboxed (Z :. (L.length idx)) idx)
            (\(Z :. freqs) -> (Z :. freqs :. numR2Freqs :. numR2Freqs)) $ \f (Z :. pFreq :. xFreq :. yFreq) ->
            fourierMellin
              sigma
              (snd $ f (Z :. pFreq))
              (fst $ f (Z :. pFreq))
              ( fromIntegral (xFreq - centerR2Freq)
              , fromIntegral (yFreq - centerR2Freq))
        arr <-
          liftIO $
          computeFourierSeriesR2
            deviceIDs
            numR2Freqs
            numPoints
            periodR2
            inverseR2Harmonics
            [pinwheels]
        yield .
          L.map
            (\i ->
               VG.convert . toUnboxed . computeS . R.slice arr $
               (Z :. i :. All :. All)) $
          [0 .. (L.length idx - 1)])

pinwheelFourierSeriesStream ::
     ( Storable e
     , CUBLAS (Complex e)
     , Unbox e
     , RealFloat e
     , A.Elt e
     , A.Elt (Complex e)
     , Floating (A.Exp e)
     , A.FromIntegral Int e
     , VG.Vector vector (Complex e)
     )
  => [Int]
  -> [PTX]
  -> Int
  -> Int
  -> e
  -> e
  -> Int
  -> Int
  -> Int
  -> Int
  -> e
  -> Int
  -> Int
  -> Int
  -> IO (IA.Array (Int, Int) (vector (Complex e)))
pinwheelFourierSeriesStream deviceIDs ptxs numR2Freqs numPoints delta periodR2 phiFreq rhoFreq thetaFreq rFreq sigma numBatchR2 batchSize numBatchPinwheelFreqs = do
  inverseR2Harmonics <-
    createInverseHarmonicMatriesGPU
      ptxs
      numBatchR2
      numPoints
      numR2Freqs
      periodR2
      delta
  let (radialLB, radialUB) = pinwheelFreqsBound rhoFreq rFreq
      (angularLB, angularUB) = pinwheelFreqsBound rhoFreq rFreq
  fmap (listArray ((radialLB, angularLB), (radialUB, radialUB)) . L.concat) .
    runConduitRes $
    source phiFreq rhoFreq thetaFreq rFreq .|
    conduit
      deviceIDs
      numR2Freqs
      numPoints
      periodR2
      batchSize
      sigma
      inverseR2Harmonics .|
    CL.consume

-- 2D spatial domain to frequency domain:
-- transform: cis (-freq * x)
-- inverse transform: cis (freq * x)

-- Analytical solution type 1
-- The pinwheel harmonics are cis (-freq * x)
{-# INLINE analyticalFourierSeriesFunc1 #-}
analyticalFourierSeriesFunc1 ::
     (Eq e, Fractional e, RealFloat e, Gamma (Complex e))
  => Int
  -> Int
  -> e
  -> e
  -> e
  -> e
  -> e
  -> Complex e
analyticalFourierSeriesFunc1 angularFreq radialFreq sigma periodR2 periodEnv phi rho =
  let radialConst = 2 * pi / log periodEnv
   in pi * ((0 :+ 1) ^ abs angularFreq) *
      cis (fromIntegral (-angularFreq) * phi) *
      ((periodR2 / (pi * rho) :+ 0) **
       ((2 + sigma) :+ (radialConst * fromIntegral (-radialFreq)))) *
      gamma
        (((2 + fromIntegral (abs angularFreq) + sigma) :+
          (radialConst * fromIntegral (-radialFreq))) /
         2) /
      gamma
        (((fromIntegral (abs angularFreq) - sigma) :+
          (radialConst * fromIntegral radialFreq)) /
         2)

-- The pinwheel is in the frequency domain, the function computes the Fourier series
{-# INLINE analyticalFourierSeries1 #-}
analyticalFourierSeries1 ::
     (Eq e, Fractional e, RealFloat e, Gamma (Complex e), Unbox e)
  => Int
  -> e
  -> Int
  -> Int
  -> e
  -> e
  -> e
  -> R.Array D DIM2 (Complex e)
analyticalFourierSeries1 numPoints delta angularFreq radialFreq sigma periodR2 periodEnv =
  let center = div numPoints 2
  in fromFunction (Z :. numPoints :. numPoints) $ \(Z :. i :. j) ->
       let x = fromIntegral $ i - center
           y = fromIntegral $ j - center
           rho = sqrt $ x ^ 2 + y ^ 2
           phi = atan2 y x
       in if rho == 0
            then 0
            else analyticalFourierSeriesFunc1
                   angularFreq
                   radialFreq
                   sigma
                   periodR2
                   periodEnv
                   phi
                   (rho * delta)

-- The pinwheel is in the spatial domain, the function computes the Fourier coefficients
{-# INLINE analyticalFourierCoefficients1 #-}
analyticalFourierCoefficients1 ::
     (Eq e, Fractional e, RealFloat e, Gamma (Complex e), Unbox e)
  => Int
  -> e
  -> Int
  -> Int
  -> e
  -> e
  -> e
  -> R.Array D DIM2 (Complex e)
analyticalFourierCoefficients1 numFreqs delta angularFreq radialFreq sigma periodR2 periodEnv =
  let c = ((-1) ^ (abs angularFreq)) :+ 0
  in R.map (* c) $
     analyticalFourierSeries1 numFreqs delta angularFreq radialFreq sigma periodR2 periodEnv

-- Analytical solution type 2
-- The pinwheel harmonics are cis (freq * x)
{-# INLINE analyticalFourierSeriesFunc2 #-}
analyticalFourierSeriesFunc2 ::
     (Eq e, Fractional e, RealFloat e, Gamma (Complex e))
  => Int
  -> Int
  -> e
  -> e
  -> e
  -> e
  -> e
  -> Complex e
analyticalFourierSeriesFunc2 angularFreq radialFreq sigma periodR2 periodEnv phi rho =
  let radialConst = 2 * pi / (log periodEnv)
  in pi * ((0 :+ 1) ^ (abs angularFreq)) *
     ((cis (fromIntegral angularFreq * phi))) *
     ((periodR2 / (pi * rho) :+ 0) **
      ((2 + sigma) :+ (radialConst * fromIntegral radialFreq))) *
     (gamma $
      ((2 + fromIntegral (abs angularFreq) + sigma) :+
       (radialConst * fromIntegral radialFreq)) /
      2) /
     (gamma $
      ((fromIntegral (abs angularFreq) - sigma) :+
       (radialConst * fromIntegral (-radialFreq))) /
      2)

-- The pinwheel is in the frequency domain, the function computes the Fourier series
{-# INLINE analyticalFourierSeries2 #-}
analyticalFourierSeries2 ::
     (Eq e, Fractional e, RealFloat e, Gamma (Complex e), Unbox e)
  => Int
  -> e
  -> Int
  -> Int
  -> e
  -> e
  -> e
  -> R.Array D DIM2 (Complex e)
analyticalFourierSeries2 numPoints delta angularFreq radialFreq sigma periodR2 periodEnv =
  let center = div numPoints 2
      arr =
        fromFunction (Z :. numPoints :. numPoints) $ \(Z :. i :. j) ->
          let x = delta * (fromIntegral $ i - center)
              y = delta * (fromIntegral $ j - center)
              rho = sqrt $ x ^ 2 + y ^ 2
              phi = atan2 y x
          in if rho == 0 -- || (rho > period / 2) -- || phi /= 0
               then 0
               else analyticalFourierSeriesFunc2
                      angularFreq
                      radialFreq
                      sigma
                      periodR2
                      periodEnv
                      phi
                      rho
     -- if angularFreq == 0
     --   then fromFunction (Z :. numPoints :. numPoints) $ \_ -> 0
     --   else
  in arr

-- The pinwheel is in the spatial domain, the function computes the Fourier coefficients
{-# INLINE analyticalFourierCoefficients2 #-}
analyticalFourierCoefficients2 ::
     (Eq e, Fractional e, RealFloat e, Gamma (Complex e), Unbox e)
  => Int
  -> e
  -> Int
  -> Int
  -> e
  -> e
  -> e
  -> R.Array D DIM2 (Complex e)
analyticalFourierCoefficients2 numFreqs delta angularFreq radialFreq sigma periodR2 periodEnv =
  let c = ((-1) ^ (abs angularFreq)) / periodR2 :+ 0
  in R.map (* c) $
     analyticalFourierSeries2
       numFreqs
       delta
       angularFreq
       radialFreq
       sigma
       periodR2
       periodEnv

-- Analytical solution type 3
-- The pinwheel harmonics are cis (-freq * x)
-- The envelope is r^2 * r^\alpha, where \alpha \in (-2,-0.5)
{-# INLINE analyticalFourierSeriesFunc3 #-}
analyticalFourierSeriesFunc3 ::
     (Eq e, Fractional e, RealFloat e, Gamma (Complex e))
  => Int
  -> Int
  -> e
  -> e
  -> e
  -> e
  -> e
  -> Complex e
analyticalFourierSeriesFunc3 angularFreq radialFreq sigma periodR2 periodEnv phi rho =
  let radialConst = 2 * pi / log periodEnv
   in ((fromIntegral angularFreq ^ 2 :+ 0) -
       ((sigma + 2) :+ (radialConst * fromIntegral (-radialFreq))) ^ 2) /
      (8 * pi) *
      ((0 :+ 1) ^ abs angularFreq) *
      cis (fromIntegral (-angularFreq) * phi) /
      (rho ^ 2 :+ 0) *
      ((periodR2 / (pi * rho) :+ 0) **
       ((2 + sigma) :+ (radialConst * fromIntegral (-radialFreq)))) *
      gamma
        (((2 + fromIntegral (abs angularFreq) + sigma) :+
          (radialConst * fromIntegral (-radialFreq))) /
         2) /
      gamma
        (((fromIntegral (abs angularFreq) - sigma) :+
          (radialConst * fromIntegral radialFreq)) /
         2)

{-# INLINE analyticalFourierSeries3 #-}
analyticalFourierSeries3 ::
     (Eq e, Fractional e, RealFloat e, Gamma (Complex e), Unbox e)
  => Int
  -> e
  -> Int
  -> Int
  -> e
  -> e
  -> e
  -> R.Array D DIM2 (Complex e)
analyticalFourierSeries3 numPoints delta angularFreq radialFreq sigma periodR2 periodEnv =
  let center = div numPoints 2
  in fromFunction (Z :. numPoints :. numPoints) $ \(Z :. i :. j) ->
       let x = fromIntegral $ i - center
           y = fromIntegral $ j - center
           rho = sqrt $ x ^ 2 + y ^ 2
           phi = atan2 y x
       in if rho == 0
            then 0
            else analyticalFourierSeriesFunc3
                   angularFreq
                   radialFreq
                   sigma
                   periodR2
                   periodEnv
                   phi
                   (rho * delta)

{-# INLINE analyticalFourierCoefficients3 #-}
analyticalFourierCoefficients3 ::
     (Eq e, Fractional e, RealFloat e, Gamma (Complex e), Unbox e)
  => Int
  -> e
  -> Int
  -> Int
  -> e
  -> e
  -> e
  -> R.Array D DIM2 (Complex e)
analyticalFourierCoefficients3 numFreqs delta angularFreq radialFreq sigma periodR2 periodEnv =
  let c = ((-1) ^ (abs angularFreq)) :+ 0
  in R.map (* c) $
     analyticalFourierSeries3 numFreqs delta angularFreq radialFreq sigma periodR2 periodEnv

pinwheelFourierCoefficientsAnatical ::
     ( Unbox e
     , RealFloat e
     , Gamma (Complex e)
     , VG.Vector vector (Complex e)
     , NFData (vector (Complex e))
     )
  => Int
  -> Int
  -> Int
  -> Int
  -> Int
  -> e
  -> e
  -> e
  -> IA.Array (Int, Int) (vector (Complex e))
pinwheelFourierCoefficientsAnatical numR2Freqs phiFreq rhoFreq thetaFreq rFreq sigma periodR2 periodEnv =
  let idxs =
        [ (radialFreq, angularFreq)
        | radialFreq <- pinwheelFreqs rhoFreq rFreq
        , angularFreq <- pinwheelFreqs phiFreq thetaFreq
        ]
      centerR2Freq = div numR2Freqs 2
      (radialLB, radialUB) = pinwheelFreqsBound rhoFreq rFreq
      (angularLB, angularUB) = pinwheelFreqsBound phiFreq thetaFreq
      pinwheels =
        parMap
          rdeepseq
          (\(radialFreq, angularFreq) ->
             VG.convert .
             toUnboxed .
             computeS -- .
             -- R.zipWith
             --   (*)
             --   (fromFunction (Z :. numR2Freqs :. numR2Freqs) $ \(Z :. i :. j) ->
             --      gaussian2D
             --        (fromIntegral $ i - centerR2Freq)
             --        (fromIntegral $ j - centerR2Freq)
             --        (fromIntegral $ div centerR2Freq 1))
            $
             analyticalFourierCoefficients2
               numR2Freqs
               1
               (angularFreq)
               (radialFreq)
               sigma
               periodR2
               periodEnv)
          idxs
  in listArray ((radialLB, angularLB), (radialUB, angularUB)) pinwheels

pinwheelFourierCoefficientsAnaticalList ::
     ( Unbox e
     , RealFloat e
     , Gamma (Complex e)
     , VG.Vector vector (Complex e)
     , NFData (vector (Complex e))
     )
  => Int
  -> Int
  -> Int
  -> Int
  -> Int
  -> e
  -> e
  -> e
  -> [vector (Complex e)]
pinwheelFourierCoefficientsAnaticalList numR2Freqs phiFreq rhoFreq thetaFreq rFreq sigma periodR2 periodEnv =
  let idxs =
        [ (radialFreq, angularFreq)
        | radialFreq <- [-rhoFreq .. rhoFreq]
        , angularFreq <- [-phiFreq .. phiFreq]
        ]
      centerR2Freq = div numR2Freqs 2
      pinwheels =
        parMap
          rdeepseq
          (\(radialFreq, angularFreq) ->
             VG.convert . toUnboxed . computeS $
             analyticalFourierCoefficients2
               numR2Freqs
               1
               (angularFreq)
               (radialFreq)
               sigma
               periodR2
               periodEnv)
          idxs
  in pinwheels

pinwheelFourierCoefficientsAnaticalList1 ::
     ( Unbox e
     , RealFloat e
     , Gamma (Complex e)
     , VG.Vector vector (Complex e)
     , NFData (vector (Complex e))
     )
  => Int
  -> Int
  -> Int
  -> Int
  -> Int
  -> e
  -> e
  -> e
  -> [vector (Complex e)]
pinwheelFourierCoefficientsAnaticalList1 numR2Freqs phiFreq rhoFreq thetaFreq rFreq sigma periodR2 periodEnv =
  let idxs =
        [ (radialFreq, angularFreq)
        | radialFreq <- [-rhoFreq .. rhoFreq]
        , angularFreq <- [-phiFreq .. phiFreq]
        ]
      pinwheels =
        parMap
          rdeepseq
          (\(radialFreq, angularFreq) ->
             VG.convert . toUnboxed . computeS $
             analyticalFourierCoefficients2
               numR2Freqs
               1
               (angularFreq)
               (radialFreq)
               sigma
               periodR2
               periodEnv)
          idxs
  in pinwheels

pinwheelFourierCoefficientsAnaticalList2 ::
     ( Unbox e
     , RealFloat e
     , Gamma (Complex e)
     , VG.Vector vector (Complex e)
     , NFData (vector (Complex e))
     )
  => Int
  -> Int
  -> Int
  -> Int
  -> Int
  -> e
  -> e
  -> e
  -> [vector (Complex e)]
pinwheelFourierCoefficientsAnaticalList2 numR2Freqs phiFreq rhoFreq thetaFreq rFreq sigma periodR2 periodEnv =
  let idxs =
        [ (radialFreq, angularFreq)
        | radialFreq <- [rFreq,(rFreq - 1) .. -rFreq]
        , angularFreq <- [thetaFreq,(thetaFreq - 1) .. -thetaFreq]
        ]
      centerR2Freq = div numR2Freqs 2
      radialConstant = 2 * pi / (log periodEnv)
      pinwheels =
        parMap
          rdeepseq
          (\(radialFreq, angularFreq) ->
             VG.convert . toUnboxed . computeS $
             fromFunction (Z :. numR2Freqs :. numR2Freqs) $ \(Z :. x :. y) ->
               let xFreq = fromIntegral $ x - centerR2Freq
                   yFreq = fromIntegral $ y - centerR2Freq
                   rho = sqrt $ xFreq ^ 2 + yFreq ^ 2
                   phi = atan2 yFreq xFreq
               in if rho == 0
                    then 0
                    else ((0 :+ (-1)) ^ 1) *
                         ((periodR2 / (pi * rho) :+ 0) **
                          (0 :+ fromIntegral radialFreq * radialConstant)) *
                         (cis $ phi * fromIntegral angularFreq))
          idxs
  in pinwheels

{-# INLINE analyticalFourierSeries2' #-}
analyticalFourierSeries2' ::
     (Eq e, Fractional e, RealFloat e, Gamma (Complex e), Unbox e)
  => Int
  -> e
  -> Int
  -> Int
  -> e
  -> e
  -> e
  -> R.Array D DIM2 (Complex e)
analyticalFourierSeries2' numPoints _ angularFreq radialFreq sigma periodR2 periodEnv =
  let center = div numPoints 2
      arr =
        fromFunction (Z :. numPoints :. numPoints) $ \(Z :. i :. j) ->
          let x' = i - center
              y' = j - center
          in if x' == 0 && y' == 0
               then 0
               else let x =
                          if x' == 0
                            then 0
                            else if x' > 0
                                   then 1 / (fromIntegral $ center + 1 - x')
                                   else (-1) / (fromIntegral $ center + 1 + x')
                        y =
                          if y' == 0
                            then 0
                            else if y' > 0
                                   then 1 / (fromIntegral $ center + 1 - y')
                                   else (-1) / (fromIntegral $ center + 1 + y')
                        rho = sqrt $ x ^ 2 + y ^ 2
                        phi = atan2 y x
                    in analyticalFourierSeriesFunc2
                         angularFreq
                         radialFreq
                         sigma
                         periodR2
                         periodEnv
                         phi
                         rho
  in arr

{-# INLINE analyticalFourierCoefficients2' #-}
analyticalFourierCoefficients2' ::
     (Eq e, Fractional e, RealFloat e, Gamma (Complex e), Unbox e)
  => Int
  -> e
  -> Int
  -> Int
  -> e
  -> e
  -> e
  -> R.Array D DIM2 (Complex e)
analyticalFourierCoefficients2' numFreqs delta angularFreq radialFreq sigma periodR2 periodEnv =
  let c = (-1) ^ (abs angularFreq) :+ 0
  in R.map (* c) $
     analyticalFourierSeries2'
       numFreqs
       delta
       angularFreq
       radialFreq
       sigma
       periodR2
       periodEnv

pinwheelFourierCoefficientsAnatical' ::
     ( Unbox e
     , RealFloat e
     , Gamma (Complex e)
     , VG.Vector vector (Complex e)
     , NFData (vector (Complex e))
     )
  => Int
  -> Int
  -> Int
  -> Int
  -> Int
  -> e
  -> e
  -> e
  -> IA.Array (Int, Int) (vector (Complex e))
pinwheelFourierCoefficientsAnatical' numR2Freqs phiFreq rhoFreq thetaFreq rFreq sigma periodR2 periodEnv =
  let idxs =
        [ (radialFreq, angularFreq)
        | radialFreq <- pinwheelFreqs rhoFreq rFreq
        , angularFreq <- pinwheelFreqs phiFreq thetaFreq
        ]
      centerR2Freq = div numR2Freqs 2
      (radialLB, radialUB) = pinwheelFreqsBound rhoFreq rFreq
      (angularLB, angularUB) = pinwheelFreqsBound phiFreq thetaFreq
      pinwheels =
        parMap
          rdeepseq
          (\(radialFreq, angularFreq) ->
             VG.convert .
             toUnboxed .
             computeS -- .
             -- R.zipWith
             --   (*)
             --   (fromFunction (Z :. numR2Freqs :. numR2Freqs) $ \(Z :. i :. j) ->
             --      gaussian2D
             --        (fromIntegral $ i - centerR2Freq)
             --        (fromIntegral $ j - centerR2Freq)
             --        (fromIntegral $ div centerR2Freq 2))
            $
             analyticalFourierCoefficients2'
               numR2Freqs
               1
               (angularFreq)
               (radialFreq)
               sigma
               periodR2
               periodEnv)
          idxs
  in listArray ((radialLB, angularLB), (radialUB, angularUB)) pinwheels

{-# INLINE idealHighPassFilter #-}
idealHighPassFilter ::
     (VG.Vector vector (Complex Double))
  => Double
  -> Int
  -> vector (Complex Double)
idealHighPassFilter radius numR2Freq =
  let a = 1 / (radius * 2)
      r2Freqs = L.map fromIntegral . getListFromNumber $ numR2Freq
      sinc x =
        if x == 0
          then 1
          else (sin (pi * x)) / (pi * x)
  in VG.fromList
       [ if xFreq == 0 && yFreq == 0
         then (1 - 1 / a ^ 2) :+ 0
         else ((-1) * (sinc (xFreq / a)) * (sinc (yFreq / a)) / (a ^ 2)) :+ 0
       | xFreq <- r2Freqs
       , yFreq <- r2Freqs
       ]

{-# INLINE idealHighPassFilter1 #-}
idealHighPassFilter1 :: Int -> VS.Vector (Complex Double)
idealHighPassFilter1 numR2Freq =
  let r2Freqs = L.map fromIntegral . getListFromNumber $ numR2Freq
  in VS.fromList
       [ if xFreq == 0 && yFreq == 0
         then 0
         else (-1) :+ 0
       | xFreq <- r2Freqs
       , yFreq <- r2Freqs
       ]

{-# INLINE gaussianHighPassFilter #-}
gaussianHighPassFilter :: Double -> Int -> VS.Vector (Complex Double)
gaussianHighPassFilter alpha numR2Freq =
  let r2Freqs = L.map fromIntegral . getListFromNumber $ numR2Freq
  in VS.fromList
       [ if xFreq == 0 && yFreq == 0
         then (1 - (pi / alpha)) :+ 0
         else ((-1) * (pi / alpha) *
               exp ((pi ^ 2 * (xFreq ^ 2 + yFreq ^ 2)) / (-alpha))) :+
              0
       | xFreq <- r2Freqs
       , yFreq <- r2Freqs
       ]

{-# INLINE gaussianHighPassFilter1 #-}
gaussianHighPassFilter1 :: Double -> Double -> Int -> VS.Vector (Complex Double)
gaussianHighPassFilter1 radius alpha numR2Freq =
  let a = 1 / (radius * 2)
      r2Freqs = L.map fromIntegral . getListFromNumber $ numR2Freq
      sinc x =
        if x == 0
          then 1
          else (sin (pi * x)) / (pi * x)
  in VS.fromList
       [ (sinc (xFreq / a) * sinc (yFreq / a) / (a ^ 2) -
          (pi / alpha) * exp ((pi ^ 2 * (xFreq ^ 2 + yFreq ^ 2)) / (-alpha))) :+
       0
       | xFreq <- r2Freqs
       , yFreq <- r2Freqs
       ]

{-# INLINE idealLowPassFilter #-}
idealLowPassFilter ::
     (VG.Vector vector (Complex Double))
  => Double
  -> Double
  -> Int
  -> vector (Complex Double)
idealLowPassFilter radius periodR2 numR2Freq =
  let r2Freqs = L.map fromIntegral . getListFromNumber $ numR2Freq
   in VG.fromList
        [ let rho = 2 * pi * sqrt (xFreq ^ 2 + yFreq ^ 2) / periodR2
           in if rho == 0
                then 0
                else radius / rho  * bessel_J1 (radius * rho) / periodR2 / (2 * pi)^2 :+ 0
        | xFreq <- r2Freqs
        , yFreq <- r2Freqs
        ]

{-# INLINE gaussianLowPassFilter #-}
gaussianLowPassFilter :: Double -> Int -> VS.Vector (Complex Double)
gaussianLowPassFilter std numR2Freq =
  let r2Freqs = L.map fromIntegral . getListFromNumber $ numR2Freq
  in VS.fromList
       [ ((std ^ 2) * exp ((-pi) * (xFreq ^ 2 + yFreq ^ 2) * (std ^ 2))) :+
       0
       | xFreq <- r2Freqs
       , yFreq <- r2Freqs
       ]

{-# INLINE laplacianLowPassFilter #-}
laplacianLowPassFilter :: Double -> Int -> VS.Vector (Complex Double)
laplacianLowPassFilter a numR2Freq =
  let r2Freqs = L.map fromIntegral . getListFromNumber $ numR2Freq
  in VS.fromList
       [ let s2 = xFreq ^ 2 + yFreq ^ 2
         in if xFreq == 0 && yFreq == 0
              then (1 - a / ((4 * pi * pi * s2 + a ^ 2) ** 1.5)) :+ 0
              else (-a) / ((4 * pi * pi * s2 + a ^ 2) ** 1.5) :+ 0
       | xFreq <- r2Freqs
       , yFreq <- r2Freqs
       ]

{-# INLINE laplacianHighPassFilter #-}
laplacianHighPassFilter :: Int -> VS.Vector (Complex Double)
laplacianHighPassFilter numR2Freq =
  let r2Freqs = L.map fromIntegral . getListFromNumber $ numR2Freq
  in VS.fromList
       [ if xFreq == 0 && yFreq == 0
         then 8 -- -3.33
         else if xFreq == 0 && abs yFreq == 1
                then -1 -- 0.67
                else if abs xFreq == 1 && yFreq == 0
                       then -1 -- 0.67
                       else if abs xFreq == 1 && abs yFreq == 1
                              then -1 -- 0.17
                              else 0
       | xFreq <- r2Freqs
       , yFreq <- r2Freqs
       ]

{-# INLINE convolveFrequency #-}
convolveFrequency ::
     (VG.Vector vector (Complex Double))
  => DFTPlan
  -> Int
  -> vector (Complex Double)
  -> IA.Array (Int, Int) (vector (Complex Double))
  -> IO (IA.Array (Int, Int) (vector (Complex Double)))
convolveFrequency plan numR2Freq filter' arr = do
  let dftPlanID = DFTPlanID DFT1DG [numR2Freq, numR2Freq] [0, 1]
      idftPlanID = DFTPlanID IDFT1DG [numR2Freq, numR2Freq] [0, 1]
      vecs = L.map VG.convert . IA.elems $ arr
      filter =
        VU.convert .
        toUnboxed .
        computeUnboxedS .
        makeFilter2D . fromUnboxed (Z :. numR2Freq :. numR2Freq) . VG.convert $
        filter'
  dftF <- dftExecute plan dftPlanID filter
  vecsF <- mapConcurrently (dftExecute plan dftPlanID) vecs
  outputs <-
    mapConcurrently (dftExecute plan idftPlanID . VS.zipWith (*) dftF) vecsF
  return . listArray (bounds arr) . L.map VG.convert $ outputs

{-# INLINE convolveFrequency1 #-}
convolveFrequency1 ::
     (VG.Vector vector (Complex Double))
  => DFTPlan
  -> Int
  -> vector (Complex Double)
  -> IA.Array (Int, Int) (vector (Complex Double))
  -> IO (IA.Array (Int, Int) (vector (Complex Double)))
convolveFrequency1 plan numR2Freq filter' arr = do
  let dftPlanID = DFTPlanID DFT1DG [numR2Freq, numR2Freq] [0, 1]
      idftPlanID = DFTPlanID IDFT1DG [numR2Freq, numR2Freq] [0, 1]
      vecs = L.map VG.convert . IA.elems $ arr
      filter =
        VU.convert .
        toUnboxed .
        computeUnboxedS .
        makeFilter2D . fromUnboxed (Z :. numR2Freq :. numR2Freq) . VG.convert $
        filter'
  dftF <- dftExecute plan dftPlanID filter
  vecsF <- mapConcurrently (dftExecute plan dftPlanID) vecs
  outputs <-
    mapConcurrently (dftExecute plan idftPlanID . VS.zipWith (*) dftF) vecsF
  return .
    listArray (bounds arr) . L.map VG.convert . L.zipWith (VG.zipWith (-)) vecs $
    outputs

{-# INLINE centerHollow #-}
centerHollow ::
     (VG.Vector vector (Complex Double))
  => Int
  -> IA.Array (Int, Int) (vector (Complex Double))
  -> IA.Array (Int, Int) (vector (Complex Double))
centerHollow numR2Freq arr = centerHollowVector numR2Freq <$> arr

{-# INLINE centerHollowVector #-}
centerHollowVector ::
     (VG.Vector vector (Complex Double))
  => Int
  -> vector (Complex Double)
  -> vector (Complex Double)
centerHollowVector numR2Freq vec =
  let s = VG.sum vec / (fromIntegral (numR2Freq ^ 2) :+ 0)
   in VG.map (\x -> x - s) vec

{-# INLINE centerHollowArray #-}
centerHollowArray ::
     (R.Source s (Complex e), Unbox e, RealFloat e)
  => Int
  -> R.Array s DIM2 (Complex e)
  -> R.Array D DIM2 (Complex e)
centerHollowArray numR2Freq arr =
  let s = sumAllS arr / (fromIntegral (numR2Freq ^ 2) :+ 0)
   in R.map (\x -> x - s) $ arr

{-# INLINE centerHollowArray' #-}
centerHollowArray' ::
     (R.Source s (Complex e), Unbox e, RealFloat e)
  => Int
  -> R.Array s DIM2 (Complex e)
  -> R.Array D DIM2 (Complex e)
centerHollowArray' numR2Freq arr =
  let s = sumAllS arr
      c = div numR2Freq 2
   in R.traverse arr id $ \f idx@(Z :. i :. j) ->
        if i == c && j == c
          then (-s)
          else f idx

envelopIntegral ::
     Double -> Double -> Double -> Double -> Double -> Int -> Complex Double
envelopIntegral a b delta s period radialFreq =
  let m = round $ (b - a) / delta
      n =
        if odd m
          then m
          else m - 1
      -- weights = VU.fromList $ weightsSimpsonRule n
      vec =
        VU.generate n $ \i ->
          let x = (a + fromIntegral i * delta)
          -- in   ((x :+ 0) ** (s :+ freq)) * ((pi :+ 0) ** (0 :+ freq)) *
          --      (gamma ((0.5 :+ 0) * (1 :+ (-freq)))) /
          --      (gamma ((0.5 :+ 0) * (1 :+ freq)))
          in analyticalFourierSeriesFunc2 0 radialFreq s period period 0 x
  -- in (delta / 3 :+ 0) * (VU.sum . VU.zipWith (*) vec $ weights)
  in VU.sum vec

envelopIntegral2D ::
     Int -> Double -> Double -> Double -> Int -> Complex Double
envelopIntegral2D numPoints delta s period radialFreq =
  let center = div numPoints 2
      -- weights = VU.fromList $ weightsSimpsonRule n
      arr = analyticalFourierSeries2 numPoints delta 0 radialFreq s period period
  in sumAllS arr

printEnvelopIntegral ::
     Double -> Double -> Double -> Double -> Double -> Int -> IO ()
printEnvelopIntegral a b delta s period radialFreq =
  let m = round $ (b - a) / delta
      n =
        if odd m
          then m
          else m - 1
      weights = VU.fromList $ weightsSimpsonRule n
      vec =
        VU.generate n $ \i ->
          let x = a + fromIntegral i * delta
          in analyticalFourierSeriesFunc2 0 radialFreq s period period 0 x
  in print $ VU.zipWith (*) weights vec

{-# INLINE createFrequencyArray #-}
createFrequencyArray ::
     (RealFloat e) => Int -> (e -> e -> Complex e) -> R.Array D DIM2 (Complex e)
createFrequencyArray numR2Freqs f =
  fromFunction (Z :. numR2Freqs :. numR2Freqs) $ \(Z :. i :. j) ->
    let x = fromIntegral $ i - div numR2Freqs 2
        y = fromIntegral $ j - div numR2Freqs 2
        rho = sqrt $ x ^ 2 + y ^ 2
        phi = atan2 y x
     in f phi rho

{-# INLINE createFrequencyArrayZeroCenter #-}
createFrequencyArrayZeroCenter ::
     (RealFloat e) => Int -> (e -> e -> Complex e) -> R.Array D DIM2 (Complex e)
createFrequencyArrayZeroCenter numR2Freqs f =
  fromFunction (Z :. numR2Freqs :. numR2Freqs) $ \(Z :. i :. j) ->
    let x = fromIntegral $ i - div numR2Freqs 2
        y = fromIntegral $ j - div numR2Freqs 2
        rho = sqrt $ x ^ 2 + y ^ 2
        phi = atan2 y x
     in if rho == 0
          then 0
          else f phi rho