text:
{-# LANGUAGE BangPatterns #-}
-- | The GenOptProblem module performs several passes on the translation generated
-- by the Style compiler to generate the initial state (fields and GPIs) and optimization problem
-- (objectives, constraints, and computations) specified by the Substance/Style pair.

{-# OPTIONS_HADDOCK prune #-}
{-# LANGUAGE AllowAmbiguousTypes, RankNTypes, UnicodeSyntax, NoMonomorphismRestriction #-}
{-# LANGUAGE DeriveGeneric #-}
-- Mostly for autodiff

module GenOptProblem where

import Utils
import Shapes
import Transforms
import qualified SubstanceJSON as J
import qualified Substance as C
import Env
import Style
import Functions
import Text.Show.Pretty (ppShow, pPrint)
import System.Random
import Debug.Trace
import qualified Data.Map.Strict as M
import Control.Monad (foldM, forM_)
import Data.List (foldl', minimumBy, intercalate, partition)
import Data.Array (assocs)
import Data.Either (partitionEithers)
import           System.Console.Pretty (Color (..), Style (..), bgColor, color, style, supportsPretty)
import qualified Data.Set as Set
import qualified Data.Graph as Graph
import GHC.Float (float2Double, double2Float)
import qualified Data.Maybe as DM (fromJust)
import qualified Data.Aeson as A
import GHC.Generics
import qualified Numeric.LinearAlgebra as L

default (Int, Float)

-------------------- Type definitions

type StyleOptFn = (String, [Expr]) -- Objective or constraint

data OptType = Objfn | Constrfn
     deriving (Show, Eq)

data Fn = Fn { fname :: String,
               fargs :: [Expr],
               optType :: OptType }
     deriving (Show, Eq)

data FnDone a = FnDone { fname_d :: String,
                         fargs_d :: [ArgVal a],
                         optType_d :: OptType }
     deriving (Show, Eq)

-- A map from the varying path to its value, used to look up values in the translation
type VaryMap a = M.Map Path (TagExpr a)

------- State type definitions

-- Stores the last EP varying state (that is, the state when the unconstrained opt last converged)
type LastEPstate = [Float] -- Note: NOT polymorphic (due to system slowness with polymorphism)

data OptStatus = NewIter
               | UnconstrainedRunning LastEPstate
               | UnconstrainedConverged LastEPstate
               | EPConverged

instance Show OptStatus where
         show NewIter = "New iteration"
         show (UnconstrainedRunning lastEPstate) =
              "Unconstrained running" -- with last EP state:\n" ++ show lastEPstate
         show (UnconstrainedConverged lastEPstate) =
              "Unconstrained converged" -- with last EP state:\n" ++ show lastEPstate
         show EPConverged = "EP converged"

instance Eq OptStatus where
         x == y = case (x, y) of
                  (NewIter, NewIter) -> True
                  (EPConverged, EPConverged) -> True
                  (UnconstrainedRunning a, UnconstrainedRunning b) -> a == b
                  (UnconstrainedConverged a, UnconstrainedConverged b) -> a == b
                  (_, _) -> False

data Params = Params { weight :: Float,
                       optStatus :: OptStatus,
                       overallObjFn :: forall a . (Autofloat a) => StdGen -> a -> [a] -> a,
                       bfgsInfo :: BfgsParams
                     }

instance Show Params where
         show p = "Weight: " ++ show (weight p) ++ " | Opt status: " ++ show (optStatus p)
                             -- ++ "\nBFGS info:\n" ++ show (bfgsInfo p)

data BfgsParams = BfgsParams {
     lastState :: Maybe (L.Vector L.R), -- x_k
     lastGrad :: Maybe (L.Vector L.R),  -- gradient of f(x_k)
     invH :: Maybe (L.Matrix L.R),  -- (BFGS only) estimate of the inverse of the hessian, H_k (TODO: are these indices right?)
     s_list :: [L.Vector L.R], -- (L-BFGS only) s_i (state difference) from k-1 to k-m
     y_list :: [L.Vector L.R],  -- (L-BFGS only) y_i (grad difference) from k-1 to k-m
     numUnconstrSteps :: Int, -- (L-BFGS only) number of steps so far, starting at 0
     memSize :: Int -- (L-BFGS only) number of vectors to retain
}

instance Show BfgsParams where
         show s = "\nBFGS params:\n" ++
                  "\nlastState: \n" ++ ppShow (lastState s) ++
                  "\nlastGrad: \n" ++ ppShow (lastGrad s) ++
                  "\ninvH: \n" ++ ppShow (invH s) ++
                  -- This is a lot of output (can be 2 * defaultBfgsMemSize * state size)
                  -- "\ns_list:\n" ++ ppShow (s_list s) ++
                  -- "\ny_list:\n" ++ ppShow (y_list s) ++
                  "\nlength of s_list:\n" ++ (show $ length $ s_list s) ++
                  "\nlength of y_list:\n" ++ (show $ length $ y_list s) ++
                  "\nnumUnconstrSteps:\n" ++ ppShow (numUnconstrSteps s) ++
                  "\nmemSize:\n" ++ ppShow (memSize s) ++ "\n\n"

defaultBfgsMemSize :: Int
defaultBfgsMemSize = 17
-- Shorter memory seems to work better in practice; Nocedal says between 3 and 30 is a good `m` (see p227)
-- but the choice of `m` is also problem-dependent

defaultBfgsParams = BfgsParams { lastState = Nothing, lastGrad = Nothing, invH = Nothing,
                                 s_list = [], y_list = [], numUnconstrSteps = 0, memSize = defaultBfgsMemSize }

type PolicyState = String -- Should this include the functions that it returned last time?
type Policy = [Fn] -> [Fn] -> PolicyParams -> (Maybe [Fn], PolicyState)

data PolicyParams = PolicyParams { policyState :: String,
                                   policySteps :: Int,
                                   currFns :: [Fn]
                                 }

instance Show PolicyParams where
         show p = "Policy state: " ++ policyState p ++ " | Policy steps: " ++ show (policySteps p)
                          -- ++ "\nFunctions:\n" ++ ppShow (currFns p)

data OptMethod = Newton | BFGS | LBFGS | GradientDescent
     deriving (Eq, Show, Generic)

instance A.ToJSON OptMethod where
             toEncoding = A.genericToEncoding A.defaultOptions

instance A.FromJSON OptMethod

data OptConfig = OptConfig {
               optMethod :: OptMethod
     } deriving (Eq, Show, Generic)

defaultOptConfig = OptConfig { optMethod = BFGS }

instance A.ToJSON OptConfig where
             toEncoding = A.genericToEncoding A.defaultOptions

instance A.FromJSON OptConfig

data State = State { shapesr :: forall a . (Autofloat a) => [Shape a],
                     shapeNames :: [(String, Field)], -- TODO Sub name type
                     shapeOrdering :: [String],
                     shapeProperties :: [(String, Field, Property)],
                     transr :: forall a . (Autofloat a) => Translation a,
                     varyingPaths :: [Path],
                     uninitializedPaths :: [Path],
                     varyingState :: [Float], -- Note: NOT polymorphic
                     paramsr :: Params,
                     objFns :: [Fn],
                     constrFns :: [Fn],
                     rng :: StdGen,
                     autostep :: Bool,
                     policyFn :: Policy,
                     policyParams :: PolicyParams,
                     oConfig :: OptConfig }

instance Show State where
         show s = "Shapes: \n" ++ ppShow (shapesr s) ++
                  "\nShape names: \n" ++ ppShow (shapeNames s) ++
                  "\nTranslation: \n" ++ ppShow (transr s) ++
                  "\nVarying paths: \n" ++ ppShow (varyingPaths s) ++
                  "\nUninitialized paths: \n" ++ ppShow (uninitializedPaths s) ++
                  "\nVarying state: \n" ++ ppShow (varyingState s) ++
                  "\nParams: \n" ++ ppShow (paramsr s) ++
                  "\nObjective Functions: \n" ++ ppShowList (objFns s) ++
                  "\nConstraint Functions: \n" ++ ppShowList (constrFns s) ++
                  "\nAutostep: \n" ++ ppShow (autostep s)

-- Reimplementation of 'ppShowList' from pretty-show. Not sure why it cannot be imported at all
ppShowList = concatMap ((++) "\n" . ppShow)

--------------- Constants

-- For evaluating expressions
startingIteration, maxEvalIteration :: Int
startingIteration = 0
maxEvalIteration  = 500 -- Max iteration depth in case of cycles

evalIterRange :: (Int, Int)
evalIterRange = (startingIteration, maxEvalIteration)

initRng :: StdGen
initRng = mkStdGen seed
    where seed = 17 -- deterministic RNG with seed

--------------- Parameters used in optimization
-- Should really be in Optimizer, but need to fix module import structure

constrWeight :: Floating a => a
constrWeight = 10 ^ 4

-- for use in barrier/penalty method (interior/exterior point method)
-- seems if the point starts in interior + weight starts v small and increases, then it converges
-- not quite: if the weight is too small then the constraint will be violated
initWeight :: Autofloat a => a
-- initWeight = 10 ** (-5)

-- Converges very fast w/ constraints removed (function-composition.sub)
-- initWeight = 0

-- Steps very slowly with a higher weight; does not seem to converge but looks visually OK (function-composition.sub)
-- initWeight = 1
initWeight = 10 ** (-3)

policyToUse :: Policy
policyToUse = optimizeSumAll
-- policyToUse = optimizeConstraintsThenObjectives
-- policyToUse = optimizeConstraints
-- policyToUse = optimizeObjectives

--------------- Utility functions

declaredVarying :: (Autofloat a) => TagExpr a -> Bool
declaredVarying (OptEval (AFloat Vary)) = True
declaredVarying _                       = False

sumMap :: Floating b => (a -> b) -> [a] -> b -- common pattern in objective functions
sumMap f l = sum $ map f l

-- TODO: figure out what to do with sty vars
mkPath :: [String] -> Path
mkPath [name, field]           = FieldPath (BSubVar (VarConst name)) field
mkPath [name, field, property] = PropertyPath (BSubVar (VarConst name)) field property

pathToList :: Path -> [String]
pathToList (FieldPath (BSubVar (VarConst name)) field)             = [name, field]
pathToList (PropertyPath (BSubVar (VarConst name)) field property) = [name, field, property]
pathToList _ = error "pathToList should not handle Sty vars"

isFieldPath :: Path -> Bool
isFieldPath (FieldPath _ _)      = True
isFieldPath (PropertyPath _ _ _) = False

bvarToString :: BindingForm -> String
bvarToString (BSubVar (VarConst s)) = s
bvarToString (BStyVar (StyVar' s)) = s -- For namespaces
             -- error ("bvarToString: cannot handle Style variable: " ++ show v)

getShapeName :: String -> Field -> String
getShapeName subName field = subName ++ "." ++ field

-- For varying values to be inserted into varyMap
floatToTagExpr :: (Autofloat a) => a -> TagExpr a
floatToTagExpr n = Done (FloatV n)

-- | converting from Value to TagExpr
toTagExpr :: (Autofloat a) => Value a -> TagExpr a
toTagExpr v = Done v

-- | converting from TagExpr to Value
toVal :: (Autofloat a) => TagExpr a -> Value a
toVal (Done v)    = v
toVal (OptEval _) = error "Shape properties were not fully evaluated"

toFn :: OptType -> StyleOptFn -> Fn
toFn otype (name, args) = Fn { fname = name, fargs = args, optType = otype }

toFns :: ([StyleOptFn], [StyleOptFn]) -> ([Fn], [Fn])
toFns (objfns, constrfns) = (map (toFn Objfn) objfns, map (toFn Constrfn) constrfns)

list2 (a, b) = [a, b]

mkVaryMap :: (Autofloat a) => [Path] -> [a] -> VaryMap a
mkVaryMap varyPaths varyVals = M.fromList $ zip varyPaths (map floatToTagExpr varyVals)

------------------- Translation helper functions

------ Generic functions for folding over a translation

foldFields :: (Autofloat a) => (String -> Field -> FieldExpr a -> [b] -> [b]) ->
                                       Name -> FieldDict a -> [b] -> [b]
foldFields f name fieldDict acc =
    let name' = nameStr name in -- TODO do we need do anything with Sub vs Gen names?
    let res = M.foldrWithKey (f name') [] fieldDict in
    res ++ acc

foldSubObjs :: (Autofloat a) => (String -> Field -> FieldExpr a -> [b] -> [b]) -> Translation a -> [b]
foldSubObjs f trans = M.foldrWithKey (foldFields f) [] (trMap trans)

------- Inserting into a translation

insertGPI :: (Autofloat a) =>
    Translation a -> String -> Field -> ShapeTypeStr -> PropertyDict a
    -> Translation a
insertGPI trans n field t propDict = case M.lookup (Sub n) $ trMap trans of
    Nothing        -> error "Substance ID does not exist"
    Just fieldDict ->
        let fieldDict' = M.insert field (FGPI t propDict) fieldDict
            trMap'     = M.insert (Sub n) fieldDict' $ trMap trans
        in trans { trMap = trMap' }

insertPath :: (Autofloat a) => Translation a -> (Path, TagExpr a) -> Either [Error] (Translation a)
insertPath trans (path, expr) =
         let overrideFlag = False in -- These paths should not exist in trans
         addPath overrideFlag trans path expr

insertPaths :: (Autofloat a) => [Path] -> [TagExpr a] -> Translation a -> Translation a
insertPaths varyingPaths varying trans =
         if length varying /= length varyingPaths
         then error "not the same # varying paths as varying variables"
         else case foldM insertPath trans (zip varyingPaths varying) of
              Left errs -> error $ "Error while adding varying paths: " ++ intercalate "\n" errs
              Right tr -> tr

------- Looking up fields/properties in a translation

-- First check if the path is a varying path. If so then use the varying value
-- (The value in the translation is stale and should be ignored)
-- If not then use the expr in the translation
lookupFieldWithVarying :: (Autofloat a) => BindingForm -> Field -> Translation a -> VaryMap a -> FieldExpr a
lookupFieldWithVarying bvar field trans varyMap =
    case M.lookup (mkPath [bvarToString bvar, field]) varyMap of
    Just varyVal -> {-trace "field lookup was vary" $ -} FExpr varyVal
    Nothing -> {-trace "field lookup was not vary" $ -} lookupField bvar field trans

lookupPropertyWithVarying :: (Autofloat a) => BindingForm -> Field -> Property
                                              -> Translation a -> VaryMap a -> TagExpr a
lookupPropertyWithVarying bvar field property trans varyMap =
    case M.lookup (mkPath [bvarToString bvar, field, property]) varyMap of
    Just varyVal -> {-trace "property lookup was vary" $ -} varyVal
    Nothing -> {- trace "property lookup was not vary" $ -} lookupProperty bvar field property trans

lookupProperty :: (Autofloat a) => BindingForm -> Field -> Property -> Translation a -> TagExpr a
lookupProperty bvar field property trans =
    let name = trName bvar in
    case lookupField bvar field trans of
    FExpr e ->
        -- to deal with path synonyms, e.g. `y.f = some GPI with property p; z.f = y.f; z.f.p = some value`
        -- if we're looking for `z.f.p` and we find out that `z.f = y.f`, then look for `y.f.p` instead
        -- NOTE: this makes a recursive call!
        case e of
        OptEval (EPath (FieldPath bvarSynonym fieldSynonym)) ->
                if bvar == bvarSynonym && field == fieldSynonym
                then error ("nontermination in lookupProperty with path '" ++ pathStr3 name field property ++ "' set to itself")
                else lookupProperty bvarSynonym fieldSynonym property trans
        -- the only thing that might have properties is another field path
        _ -> error ("path '" ++ pathStr3 name field property ++ "' has no properties")
    FGPI ctor properties ->
        case M.lookup property properties of
        Nothing -> error ("path '" ++ pathStr3 name field property ++ "'s property does not exist")
        Just texpr -> texpr

lookupPaths :: (Autofloat a) => [Path] -> Translation a -> [a]
lookupPaths paths trans = map lookupPath paths
    where
        lookupPath p@(FieldPath v field) = case lookupField v field trans of
            FExpr (OptEval (AFloat (Fix n))) -> r2f n
            FExpr (Done (FloatV n))          -> r2f n
            xs -> error ("varying path \"" ++ pathStr p ++ "\" is invalid: is '" ++ show xs ++ "'")
        lookupPath p@(PropertyPath v field pty) = case lookupProperty v field pty trans of
            OptEval (AFloat (Fix n)) -> r2f n
            Done (FloatV n) -> n
            xs -> error ("varying path \"" ++ pathStr p ++ "\" is invalid: is '" ++ show xs ++ "'")

-- TODO: resolve label logic here?
shapeExprsToVals :: (Autofloat a) => (String, Field) -> PropertyDict a -> Properties a
shapeExprsToVals (subName, field) properties =
          let shapeName   = getShapeName subName field
              properties' = M.map toVal properties
          in M.insert "name" (StrV shapeName) properties'

getShapes :: (Autofloat a) => [(String, Field)] -> Translation a -> [Shape a]
getShapes shapenames trans = map (getShape trans) shapenames
          -- TODO: fix use of Sub/Sty name here
          where getShape trans (name, field) =
                    let fexpr = lookupField (BSubVar $ VarConst name) field trans in
                    case fexpr of
                    FExpr _ -> error "expected GPI, got field"
                    FGPI ctor properties -> (ctor, shapeExprsToVals (name, field) properties)

----- GPI helper functions

shapes2vals :: (Autofloat a) => [Shape a] -> [Path] -> [Value a]
shapes2vals shapes paths = reverse $ foldl' (lookupPath shapes) [] paths
    where
        lookupPath shapes acc (PropertyPath s field property) =
            let subID = bvarToString s
                shapeName = getShapeName subID field in
            get (findShape shapeName shapes) property : acc
        lookupPath _ acc (FieldPath _ _) = acc

-- Given a set of new shapes (from the frontend) and a varyMap (for varying field values):
-- look up property values in the shapes and field values in the varyMap
-- NOTE: varyState is constructed using a foldl, so to preserve its order, we must reverse the list of values!
shapes2floats :: (Autofloat a) => [Shape a] -> VaryMap a -> [Path] -> [a]
shapes2floats shapes varyMap varyingPaths = reverse $ foldl' (lookupPathFloat shapes varyMap) [] varyingPaths
    where
        lookupPathFloat :: (Autofloat a) => [Shape a] -> VaryMap a -> [a] -> Path -> [a]
        lookupPathFloat shapes _ acc (PropertyPath s field property) =
            let subID = bvarToString s
                shapeName = getShapeName subID field in
            getNum (findShape shapeName shapes) property : acc
        lookupPathFloat _ varyMap acc fp@(FieldPath _ _) =
             case M.lookup fp varyMap of
             Just (Done (FloatV num)) -> num : acc
             Just _ -> error ("wrong type for varying field path (expected float): " ++ show fp)
             Nothing -> error ("could not find varying field path '" ++ show fp ++ "' in varyMap")

--------------------------------- Analyzing the translation

--- Find varying (float) paths

-- For now, don't optimize these float-valued properties of a GPI
-- (use whatever they are initialized to in Shapes or set to in Style)
unoptimizedFloatProperties :: [String]
unoptimizedFloatProperties = ["rotation", "strokeWidth", "thickness",
                              "transform", "transformation"]

-- If any float property is not initialized in properties,
-- or it's in properties and declared varying, it's varying
findPropertyVarying :: (Autofloat a) => String -> Field -> M.Map String (TagExpr a) ->
                                                                 String -> [Path] -> [Path]
findPropertyVarying name field properties floatProperty acc =
    case M.lookup floatProperty properties of
    Nothing -> if floatProperty `elem` unoptimizedFloatProperties then acc
               else mkPath [name, field, floatProperty] : acc
    Just expr -> if declaredVarying expr then mkPath [name, field, floatProperty] : acc
                 else acc

findFieldVarying :: (Autofloat a) => String -> Field -> FieldExpr a -> [Path] -> [Path]
findFieldVarying name field (FExpr expr) acc =
    if declaredVarying expr
    then mkPath [name, field] : acc -- TODO: deal with StyVars
    else acc
findFieldVarying name field (FGPI typ properties) acc =
    let ctorFloats    = propertiesOf FloatT typ
        varyingFloats = filter (not . isPending typ) ctorFloats
        vs = foldr (findPropertyVarying name field properties) [] varyingFloats
    in vs ++ acc

findVarying :: (Autofloat a) => Translation a -> [Path]
findVarying = foldSubObjs findFieldVarying

--- Find uninitialized (non-float) paths

findPropertyUninitialized :: (Autofloat a) => String -> Field -> M.Map String (TagExpr a) ->
                                                                       String -> [Path] -> [Path]
findPropertyUninitialized name field properties nonfloatProperty acc =
    case M.lookup nonfloatProperty properties of
    -- nonfloatProperty is a non-float property that is NOT set by the user and thus we can sample it
    Nothing   -> mkPath [name, field, nonfloatProperty] : acc
    Just expr -> acc

findFieldUninitialized :: (Autofloat a) => String -> Field -> FieldExpr a -> [Path] -> [Path]
-- NOTE: we don't find uninitialized field because you can't leave them uninitialized. Plus, we don't know what types they are
findFieldUninitialized name field (FExpr expr) acc = acc
findFieldUninitialized name field (FGPI typ properties) acc =
    let ctorNonfloats  = filter (/= "name") $ propertiesNotOf FloatT typ in
    -- TODO: add a separate field (e.g. pendingPaths) in State to store these special properties that needs frontend updates
    let uninitializedProps = pendingProperties typ ++ ctorNonfloats in
    let vs = foldr (findPropertyUninitialized name field properties) [] uninitializedProps in
    vs ++ acc

-- | Find the paths to all uninitialized, non-float, non-name properties
findUninitialized :: (Autofloat a) => Translation a -> [Path]
findUninitialized = foldSubObjs findFieldUninitialized

--- Find various kinds of functions

findObjfnsConstrs :: (Autofloat a) => Translation a -> [Either StyleOptFn StyleOptFn]
findObjfnsConstrs = foldSubObjs findFieldFns
    where findFieldFns :: (Autofloat a) => String -> Field -> FieldExpr a -> [Either StyleOptFn StyleOptFn]
                                        -> [Either StyleOptFn StyleOptFn]
          findFieldFns name field (FExpr (OptEval expr)) acc =
              case expr of
              ObjFn fname args -> Left (fname, args) : acc
              ConstrFn fname args -> Right (fname, args) : acc
              _ -> acc -- Not an optfn
          -- COMBAK: what should we do if there's a constant field?
          findFieldFns name field (FExpr (Done _)) acc = acc
          findFieldFns name field (FGPI _ _) acc = acc

findDefaultFns :: (Autofloat a) => Translation a -> [Either StyleOptFn StyleOptFn]
findDefaultFns = foldSubObjs findFieldDefaultFns
    where findFieldDefaultFns :: (Autofloat a) => String -> Field -> FieldExpr a ->
                                            [Either StyleOptFn StyleOptFn] -> [Either StyleOptFn StyleOptFn]
          findFieldDefaultFns name field gpi@(FGPI typ props) acc =
              let args    = [EPath $ FieldPath (BSubVar (VarConst name)) field]
                  objs    = map (Left . addArgs args) $ defaultObjFnsOf typ
                  constrs = map (Right . addArgs args) $ defaultConstrsOf typ
              in constrs ++ objs ++ acc
              where addArgs arguments f = (f, arguments)
          findFieldDefaultFns _ _ _ acc = acc

--- Find shapes and their properties

findShapeNames :: (Autofloat a) => Translation a -> [(String, Field)]
findShapeNames = foldSubObjs findGPIName
    where findGPIName :: (Autofloat a) => String -> Field -> FieldExpr a ->
                                          [(String, Field)] -> [(String, Field)]
          findGPIName name field (FGPI _ _) acc = (name, field) : acc
          findGPIName _ _ (FExpr _) acc = acc

findShapesProperties :: (Autofloat a) => Translation a -> [(String, Field, Property)]
findShapesProperties = foldSubObjs findShapeProperties
    where findShapeProperties :: (Autofloat a) => String -> Field -> FieldExpr a -> [(String, Field, Property)]
                                                  -> [(String, Field, Property)]
          findShapeProperties name field (FGPI ctor properties) acc =
               let paths = map (\property -> (name, field, property)) (M.keys properties)
               in paths ++ acc
          findShapeProperties _ _ (FExpr _) acc = acc

------------------------------ Evaluating the translation and expressions/GPIs in it

-- TODO: write a more general typechecking mechanism
evalUop :: (Autofloat a) => UnaryOp -> ArgVal a -> Value a
evalUop UMinus v = case v of
                  Val (FloatV a) -> FloatV (-a)
                  Val (IntV i) -> IntV (-i)
                  GPI _ -> error "cannot negate a GPI"
                  Val _ -> error "wrong type to negate"
evalUop UPlus v = error "unary + doesn't make sense" -- TODO remove from parser

evalBinop :: (Autofloat a) => BinaryOp -> ArgVal a -> ArgVal a -> Value a
evalBinop op v1 v2 =
        case (v1, v2) of
        (Val (FloatV n1), Val (FloatV n2)) ->
                  case op of
                  BPlus -> FloatV $ n1 + n2
                  BMinus -> FloatV $ n1 - n2
                  Multiply -> FloatV $ n1 * n2
                  Divide -> if n2 == 0 then error "divide by 0!" else FloatV $ n1 / n2
                  Exp -> FloatV $ n1 ** n2
        (Val (IntV n1), Val (IntV n2)) ->
                  case op of
                  BPlus -> IntV $ n1 + n2
                  BMinus -> IntV $ n1 - n2
                  Multiply -> IntV $ n1 * n2
                  Divide -> if n2 == 0 then error "divide by 0!" else IntV $ n1 `quot` n2 -- NOTE: not float
                  Exp -> IntV $ n1 ^ n2
        -- Cannot mix int and float
        (Val _, Val _) -> error ("wrong field types for binary op: " ++ show v1 ++ show op ++ show v2)
        (GPI _, Val _) -> error "binop cannot operate on GPI"
        (Val _, GPI _) -> error "binop cannot operate on GPI"
        (GPI _, GPI _) -> error "binop cannot operate on GPIs"

-- | Given a path that is a computed property of a shape (e.g. A.shape.transformation), evaluate each of its arguments (e.g. A.shape.sizeX), pass the results to the property-computing function, and return the result (e.g. an HMatrix)
computeProperty :: (Autofloat a) => (Int, Int) -> BindingForm -> Field -> Property -> VaryMap a -> Translation a -> StdGen -> ComputedValue a -> (ArgVal a, Translation a, StdGen)
computeProperty limit bvar field property varyMap trans g (props, compFn) =
                let args = map (\p -> EPath $ PropertyPath bvar field p) props
                    (argVals, trans', g') = evalExprs limit args trans varyMap g
                    propertyValue = compFn $ map fromGPI argVals in
                (Val propertyValue, trans', g')
                where fromGPI (Val x) = x
                      fromGPI (GPI x) = error "expected value as prop fn arg, got GPI"

evalProperty :: (Autofloat a)
    => (Int, Int) -> BindingForm -> Field -> VaryMap a -> ([(Property, TagExpr a)], Translation a, StdGen) -> (Property, TagExpr a)
    -> ([(Property, TagExpr a)], Translation a, StdGen)
evalProperty (i, n) bvar field varyMap (propertiesList, trans, g) (property, expr) =
    let path = EPath $ PropertyPath bvar field property in -- factor out?
    let (res, trans', g') = evalExpr (i, n) path trans varyMap g in
    -- This check might be redundant with the later GPI conversion in evalExpr, TODO factor out
    case res of
        Val val -> {-trace ("Evaled property " ++ show path)-} ((property, Done val) : propertiesList, trans', g')
        GPI _ -> error "GPI property should not evaluate to GPI argument" -- TODO: true later? references?

evalGPI_withUpdate :: (Autofloat a)
    => (Int, Int) -> BindingForm -> Field -> (GPICtor, PropertyDict a) -> Translation a -> VaryMap a -> StdGen
    -> ((GPICtor, PropertyDict a), Translation a, StdGen)
evalGPI_withUpdate (i, n) bvar field (ctor, properties) trans varyMap g =
        -- Fold over the properties, evaluating each path, which will update the translation each time,
        -- and accumulate the new property-value list (WITH varying looked up)
        let (propertyList', trans', g') = foldl' (evalProperty (i, n) bvar field varyMap) ([], trans, g) (M.toList properties) in
        let properties' = M.fromList propertyList' in
        {-trace ("Start eval GPI: " ++ show properties ++ " " ++ "\n\tctor: " ++ "\n\tfield: " ++ show field)-}
        ((ctor, properties'), trans', g')

-- recursively evaluate, tracking iteration depth in case there are cycles in graph
evalExpr :: (Autofloat a) => (Int, Int) -> Expr -> Translation a -> VaryMap a -> StdGen -> (ArgVal a, Translation a, StdGen)
evalExpr (i, n) arg trans varyMap g =
    if i >= n then error ("evalExpr: iteration depth exceeded (" ++ show n ++ ")")
        else {-trace ("Evaluating expression: " ++ show arg ++ "\n(i, n): " ++ show i ++ ", " ++ show n)-} argResult
    where limit = (i + 1, n)
          argResult = case arg of
            -- Already done values; don't change trans
            IntLit i -> (Val $ IntV i, trans, g)
            StringLit s -> (Val $ StrV s, trans, g)
            BoolLit b -> (Val $ BoolV b, trans, g)
            AFloat (Fix f) -> (Val $ FloatV (r2f f), trans, g) -- TODO: note use of r2f here. is that ok?
            AFloat Vary -> error "evalExpr should not encounter an uninitialized varying float!"

            -- Inline computation, needs a recursive lookup that may change trans, but not a path
            -- TODO factor out eval / trans computation?
            UOp op e ->
                let (val, trans', g') = evalExpr limit e trans varyMap g in
                let compVal = evalUop op val in
                (Val compVal, trans', g')
            BinOp op e1 e2 ->
                let ([v1, v2], trans', g') = evalExprs limit [e1, e2] trans varyMap g in
                let compVal = evalBinop op v1 v2 in
                (Val compVal, trans', g')
            CompApp fname args ->
                -- NOTE: the goal of all the rng passing in this module is for invoking computations with randomization
                let (vs, trans', g') = evalExprs limit args trans varyMap g
                    (compRes, g'')   = invokeComp fname vs compSignatures g'
                in (compRes, trans', g'')
                -- -- TODO: invokeComp should be used here
                -- case M.lookup fname compDict of
                -- Nothing -> error ("computation '" ++ fname ++ "' doesn't exist")
                -- Just f -> let res = f vs in
                --           (res, trans')
            List es ->
                let (vs, trans', g') = evalExprs limit es trans varyMap g
                    floatvs = map checkFloatType vs
                in (Val $ ListV floatvs, trans', g')

            ListAccess p i -> error "TODO list accesses"

            Tuple e1 e2 ->
                let (vs, trans', g') = evalExprs limit [e1, e2] trans varyMap g
                    [v1, v2] = map checkFloatType vs
                in (Val $ TupV (v1, v2), trans', g')

            -- Needs a recursive lookup that may change trans. The path case is where trans is actually changed.
            EPath p ->
                  case p of
                  FieldPath bvar field ->
                     -- Lookup field expr, evaluate it if necessary, cache the evaluated value in the trans,
                     -- return the evaluated value and the updated trans
                     let fexpr = lookupFieldWithVarying bvar field trans varyMap in
                     case fexpr of
                     FExpr (Done v) -> (Val v, trans, g)
                     FExpr (OptEval e) ->
                         let (v, trans', g') = evalExpr limit e trans varyMap g in
                         case v of
                             Val fval ->
                                 case insertPath trans' (p, Done fval) of
                                 Right trans' -> (v, trans', g')
                                 Left err -> error $ concat err
                             gpiVal@(GPI _) -> (gpiVal, trans', g') -- to deal with path synonyms, e.g. "y.f = some GPI; z.f = y.f"
                     FGPI ctor properties ->
                     -- Eval each property in the GPI, storing each property result in a new dictionary
                     -- No need to update the translation because each path should update the translation
                         let (gpiVal@(ctor', propertiesVal), trans', g') =
                                 evalGPI_withUpdate limit bvar field (ctor, properties) trans varyMap g in
                         (GPI (ctor', shapeExprsToVals (bvarToString bvar, field) propertiesVal), trans', g')

                  PropertyPath bvar field property ->
                      let gpiType = shapeType bvar field trans in
                      case M.lookup (gpiType, property) computedProperties of
                      Just computeValueInfo -> computeProperty limit bvar field property varyMap trans g computeValueInfo
                      Nothing -> -- Compute the path as usual
                          let texpr = lookupPropertyWithVarying bvar field property trans varyMap in
                          case texpr of
                          Done v -> (Val v, trans, g)
                          OptEval e ->
                             let (v, trans', g') = evalExpr limit e trans varyMap g in
                             case v of
                             Val fval ->
                                 case insertPath trans' (p, Done fval) of
                                 Right trans' -> (v, trans', g')
                                 Left err -> error $ concat err
                             GPI _ -> error ("path to property expr '" ++ pathStr p ++ "' evaluated to a GPI")

            -- GPI argument
            Ctor ctor properties -> error "no anonymous/inline GPIs allowed as expressions!"

            -- Error
            Layering _ _ -> error "layering should not be an objfn arg (or in the children of one)"
            ObjFn _ _ -> error "objfn should not be an objfn arg (or in the children of one)"
            ConstrFn _ _ -> error "constrfn should not be an objfn arg (or in the children of one)"
            AvoidFn _ _ -> error "avoidfn should not be an objfn arg (or in the children of one)"
            PluginAccess _ _ _ -> error "plugin access should not be evaluated at runtime"
            -- xs -> error ("unmatched case in evalExpr with argument: " ++ show xs)

checkFloatType :: (Autofloat a) => ArgVal a -> a
checkFloatType (Val (FloatV x)) = x
checkFloatType _ = error "expected float type"

-- Any evaluated exprs are cached in the translation for future evaluation
-- The varyMap is not changed because its values are final (set by the optimization)
evalExprs :: (Autofloat a)
    => (Int, Int) -> [Expr] -> Translation a -> VaryMap a -> StdGen
    -> ([ArgVal a], Translation a, StdGen)
evalExprs limit args trans varyMap g =
    foldl' (evalExprF limit varyMap) ([], trans, g) args
    where evalExprF :: (Autofloat a) => (Int, Int) -> VaryMap a -> ([ArgVal a], Translation a, StdGen) -> Expr -> ([ArgVal a], Translation a, StdGen)
          evalExprF limit varyMap (argvals, trans, rng) arg =
                       let (argVal, trans', rng') = evalExpr limit arg trans varyMap rng in
                       (argvals ++ [argVal], trans', rng') -- So returned exprs are in same order

------------------- Generating and evaluating the objective function

evalFnArgs :: (Autofloat a) => (Int, Int) -> VaryMap a -> ([FnDone a], Translation a, StdGen) -> Fn -> ([FnDone a], Translation a, StdGen)
evalFnArgs limit varyMap (fnDones, trans, g) fn =
    let args = fargs fn in
    let (argsVal, trans', g') = evalExprs limit (fargs fn) trans varyMap g in
    let fn' = FnDone { fname_d = fname fn, fargs_d = argsVal, optType_d = optType fn } in
    (fnDones ++ [fn'], trans', g') -- TODO factor out this pattern

evalFns :: (Autofloat a)
    => (Int, Int) -> [Fn] -> Translation a -> VaryMap a -> StdGen
    -> ([FnDone a], Translation a, StdGen)
evalFns limit fns trans varyMap g = foldl' (evalFnArgs limit varyMap) ([], trans, g) fns

applyOptFn :: (Autofloat a) =>
    M.Map String (OptFn a) -> OptSignatures -> FnDone a -> a
applyOptFn dict sigs finfo =
    let (name, args) = (fname_d finfo, fargs_d finfo)
    in invokeOptFn dict name args sigs

applyCombined :: (Autofloat a) => a -> [FnDone a] -> a
applyCombined penaltyWeight fns =
        let (objfns, constrfns) = partition (\f -> optType_d f == Objfn) fns in
        sumMap (applyOptFn objFuncDict objSignatures) objfns
               + constrWeight * penaltyWeight * sumMap (applyOptFn constrFuncDict constrSignatures) constrfns

-- Main function: generates the objective function, partially applying it with some info

genObjfn :: (Autofloat a)
    => Translation a -> [Fn] -> [Fn] -> [Path]
    -> StdGen -> a -> [a]
    -> a
genObjfn trans objfns constrfns varyingPaths =
     \rng penaltyWeight varyingVals ->
         let varyMap = tr "varyingMap: " $ mkVaryMap varyingPaths varyingVals in
         let (fnsE, transE, rng') = evalFns evalIterRange (objfns ++ constrfns) trans varyMap rng in
         applyCombined penaltyWeight fnsE

--------------- Generating an initial state (concrete values for all fields/properties needed to draw the GPIs)
-- 1. Initialize all varying fields
-- 2. Initialize all properties of all GPIs

-- NOTE: since we store all varying paths separately, it is okay to mark the default values as Done -- they will still be optimized, if needed.
-- TODO: document the logic here (e.g. only sampling varying floats) and think about whether to use translation here or [Shape a] since we will expose the sampler to users later
initProperty :: (Autofloat a) => (PropertyDict a, StdGen) -> String ->
                                               (ValueType, SampledValue a) -> (PropertyDict a, StdGen)
initProperty (properties, g) pID (typ, sampleF) =
    let (v, g')    = sampleF g
        autoRndVal = Done v in
    case M.lookup pID properties of
      Just (OptEval (AFloat Vary)) -> (M.insert pID autoRndVal properties, g')
      Just (OptEval e) -> (properties, g)
      Just (Done v)    -> (properties, g)
      Nothing          -> (M.insert pID autoRndVal properties, g')

initShape :: (Autofloat a) => (Translation a, StdGen) -> (String, Field) -> (Translation a, StdGen)
initShape (trans, g) (n, field) =
    case lookupField (BSubVar (VarConst n)) field trans of
        FGPI shapeType propDict ->
            let def = findDef shapeType
                (propDict', g') = foldlPropertyMappings initProperty (propDict, g) def
                -- NOTE: getShapes resolves the names + we don't use the names of the shapes in the translation
                -- The name-adding logic can be removed but is left in for debugging
                shapeName = getShapeName n field
                propDict'' =  M.insert "name" (Done $ StrV shapeName) propDict'
            in (insertGPI trans n field shapeType propDict'', g')
        _   -> error "expected GPI but got field"

initShapes :: (Autofloat a) =>
    Translation a -> [(String, Field)] -> StdGen -> (Translation a, StdGen)
initShapes trans shapePaths gen = foldl' initShape (trans, gen) shapePaths

resampleFields :: (Autofloat a) => [Path] -> StdGen -> ([a], StdGen)
resampleFields varyingPaths g =
    let varyingFields = filter isFieldPath varyingPaths in
    Functions.randomsIn g (fromIntegral $ length varyingFields) Shapes.canvasDims

-- sample varying fields only (from the range defined by canvas dims) and store them in the translation
-- example: A.val = OPTIMIZED
initFields :: (Autofloat a) => [Path] -> Translation a -> StdGen -> (Translation a, StdGen)
initFields varyingPaths trans g =
    let varyingFields = filter isFieldPath varyingPaths
        (sampledVals, g') = Functions.randomsIn g (fromIntegral $ length varyingFields) Shapes.canvasDims
        trans' = insertPaths varyingFields (map (Done . FloatV) sampledVals) trans in
    (trans', g')

------------- Evaluating all shapes in a translation

evalShape :: (Autofloat a) =>
    (Int, Int) -> VaryMap a
    -> ([Shape a], Translation a, StdGen) -> Path
    -> ([Shape a], Translation a, StdGen)
evalShape limit varyMap (shapes, trans, g) shapePath =
    let (res, trans', g') = evalExpr limit (EPath shapePath) trans varyMap g in
    case res of
    GPI shape -> (shape : shapes, trans', g')
    _ -> error "evaluating a GPI path did not result in a GPI"

-- recursively evaluate every shape property in the translation
evalShapes :: (Autofloat a) => (Int, Int) -> [Path] -> Translation a -> VaryMap a -> StdGen -> ([Shape a], Translation a, StdGen)
evalShapes limit shapeNames trans varyMap rng =
           let (shapes, trans', rng') = foldl' (evalShape limit varyMap) ([], trans, rng) shapeNames in
           (reverse shapes, trans', rng')

-- Given the shape names, use the translation and the varying paths/values in order to evaluate each shape
-- with respect to the varying values
evalTranslation :: (Autofloat a) => State -> ([Shape a], Translation a, StdGen)
evalTranslation s =
    let varyMap = mkVaryMap (varyingPaths s) (map r2f $ varyingState s) in
    evalShapes evalIterRange (map (mkPath . list2) $ shapeNames s) (transr s) varyMap (rng s)

------------- Compute global layering of GPIs

lookupGPIName :: (Autofloat a) => Path -> Translation a -> String
lookupGPIName path@(FieldPath v field) trans =
    case lookupField v field trans of
        FExpr e  ->
           -- to deal with path synonyms in a layering statement (see `lookupProperty` for more explanation)
           case e of
               OptEval (EPath pathSynonym@(FieldPath vSynonym fieldSynonym)) ->
                   if v == vSynonym && field == fieldSynonym
                   then error ("nontermination in lookupGPIName w/ path '" ++ show path ++ "' set to itself")
                   else lookupGPIName pathSynonym trans
               _ -> notGPIError
        FGPI _ _ -> getShapeName (bvarToString v) field

lookupGPIName _ _ = notGPIError
notGPIError = error "Layering expressions can only operate on GPIs."

-- | Walk the translation to find all layering statements.
findLayeringExprs :: (Autofloat a) => Translation a -> [Expr]
findLayeringExprs t = foldSubObjs findLayeringExpr t
  where findLayeringExpr :: (Autofloat a) => String -> Field -> FieldExpr a -> [Expr] -> [Expr]
        findLayeringExpr name field fexpr acc =
          case fexpr of
          FExpr (OptEval x@(Layering _ _)) -> x : acc
          _ -> acc

-- | Calculates all the nodes that are part of cycles in a graph.
cyclicNodes :: Graph.Graph -> [Graph.Vertex]
cyclicNodes graph =
  map fst . filter isCyclicAssoc . assocs $ graph
  where
    isCyclicAssoc = uncurry $ reachableFromAny graph

-- | In the specified graph, can the specified node be reached, starting out
-- from any of the specified vertices?
reachableFromAny :: Graph.Graph -> Graph.Vertex -> [Graph.Vertex] -> Bool
reachableFromAny graph node =
  elem node . concatMap (Graph.reachable graph)

-- | 'computeLayering' takes in a list of all GPI names and a list of directed edges [(a -> b)] representing partial layering orders as input and outputs a linear layering order of GPIs
topSortLayering :: [String] -> [(String, String)] -> [String]
topSortLayering names partialOrderings =
    let orderedNodes = nodesFromEdges partialOrderings
        freeNodes = Set.difference (Set.fromList names) orderedNodes
        edges = map (\(x, y) -> (x, x, y)) $ adjList partialOrderings
                    ++ (map (\x -> (x, [])) $ Set.toList freeNodes)
        (graph, nodeFromVertex, vertexFromKey) = Graph.graphFromEdges edges
        cyclic = not . null $ cyclicNodes graph
    in if cyclic then error "The graph is cyclic!" else map (getNodePart . nodeFromVertex) $ Graph.topSort graph
    where
        getNodePart (n, _, _) = n

nodesFromEdges edges = Set.fromList $ concatMap (\(a, b) -> [a, b]) edges

adjList :: [(String, String)] -> [(String, [String])]
adjList edges =
    let nodes = Set.toList $ nodesFromEdges edges
    in map (\x -> (x, findNeighbors x)) nodes
    where findNeighbors node = map snd $ filter ((==) node . fst) edges

computeLayering :: (Autofloat a) => Translation a -> [String]
computeLayering trans =
    let layeringExprs = findLayeringExprs trans
        partialOrderings = map findNames layeringExprs
        gpiNames  = map (uncurry getShapeName) $ findShapeNames trans
    in topSortLayering gpiNames partialOrderings
    where
        unused = -1
        substitute res (block, substs) =
            let block'  = (block, unused)
                substs' = map (\s -> (s, unused)) substs
            in res ++ map (`substituteBlock` block') substs'
        findNames (Layering path1 path2) = (lookupGPIName path1 trans, lookupGPIName path2 trans)

------------- Main function: what the Style compiler generates

genOptProblemAndState :: (forall a. (Autofloat a) => Translation a) -> OptConfig -> State
genOptProblemAndState trans optConfig =
    -- Save information about the translation
    let !varyingPaths       = findVarying trans in
    -- NOTE: the properties in uninitializedPaths are NOT floats. Floats are included in varyingPaths already
    let uninitializedPaths = findUninitialized trans in
    let shapeNames         = findShapeNames trans in

    -- sample varying fields
    let (transInitFields, g') = initFields varyingPaths trans initRng in

    -- sample varying vals and instantiate all the non-float base properties of every GPI in the translation
    let (!transInit, g'') = initShapes transInitFields shapeNames g' in
    let shapeProperties  = transInit `seq` findShapesProperties transInit in

    let (objfns, constrfns) = (toFns . partitionEithers . findObjfnsConstrs) transInit in
    let (defaultObjFns, defaultConstrs) = (toFns . partitionEithers . findDefaultFns) transInit in
    let (!objFnsWithDefaults, !constrsWithDefaults) = (objfns ++ defaultObjFns, constrfns ++ defaultConstrs) in
    let overallFn = genObjfn transInit objFnsWithDefaults constrsWithDefaults varyingPaths in
    -- NOTE: this does NOT use transEvaled because it needs to be re-evaled at each opt step
    -- the varying values are re-inserted at each opt step

    -- Evaluate all expressions once to get the initial shapes
    let initVaryingMap = M.empty in -- No optimization has happened. Sampled varying vals are in transInit
    let (initialGPIs, transEvaled, _) = evalShapes evalIterRange (map (mkPath . list2) shapeNames) transInit initVaryingMap g'' in -- intentially discarding the new random feed, since we want the computation result to be consistent within one optimization session
    let initState = lookupPaths varyingPaths transEvaled in

    if null initState then error "empty state in genopt" else

    -- This is the final Style compiler output
    let s = trace "genOptProblem init state: " $
                        State { shapesr = initialGPIs,
                                 shapeNames = shapeNames,
                                 shapeProperties = shapeProperties,
                                 shapeOrdering = [], -- NOTE: to be populated later
                                 transr = transInit, -- note: NOT transEvaled
                                 varyingPaths = varyingPaths,
                                 uninitializedPaths = uninitializedPaths,
                                 varyingState = initState,
                                 objFns = objFnsWithDefaults,
                                 constrFns = constrsWithDefaults,
                                 paramsr = Params { weight = initWeight,
                                                    optStatus = NewIter,
                                                    overallObjFn = overallFn,
                                                    bfgsInfo = defaultBfgsParams },
                                 rng = g'',
                                 autostep = False, -- default
                                 policyParams = initPolicyParams,
                                 policyFn = policyToUse,
                                 oConfig = optConfig
                               } in

    initPolicy s
    -- NOTE: we do not resample the very first initial state. Not sure why the shapes / labels are rendered incorrectly.
    -- resampleBest numStateSamples initFullState

-- | 'compileStyle' runs the main Style compiler on the AST of Style and output from the Substance compiler and outputs the initial state for the optimization problem. This function is a top-level function used by "Server" and "ShadowMain"
-- NOTE: this function also print information out to stdout
-- TODO: enable logger
compileStyle :: StyProg -> C.SubOut -> [J.StyVal] -> OptConfig -> IO State
compileStyle styProg (C.SubOut subProg (subEnv, eqEnv) labelMap) styVals optConfig = do
   putStrLn "Running Style semantics\n"
   let selEnvs = checkSels subEnv styProg

   putStrLn "Selector static semantics and local envs:\n"
   forM_ selEnvs pPrint
   divLine

   let subss = find_substs_prog subEnv eqEnv subProg styProg selEnvs
   putStrLn "Selector matches:\n"
   forM_ subss pPrint
   divLine

   let !trans = translateStyProg subEnv eqEnv subProg styProg labelMap styVals
                       :: Either [Error] (Translation Float)
                       -- NOT :: forall a . (Autofloat a) => Either [Error] (Translation a)
                       -- We intentionally specialize/monomorphize the translation to Float so it can be fully evaluated
                       -- and is not trapped under the lambda of the typeclass (Autofloat a) => ...
                       -- This greatly improves the performance of the system. See #166 for more details.
   let transAuto = castTranslation $ fromRight trans
                       :: forall a . (Autofloat a) => Translation a
   putStrLn "Translated Style program:\n"
   pPrint trans
   divLine

   let initState = genOptProblemAndState transAuto optConfig
   putStrLn "Generated initial state:\n"
   print initState
   divLine

   -- global layering order computation
   let gpiOrdering = computeLayering transAuto
   putStrLn "Generated GPI global layering:\n"
   print gpiOrdering
   divLine

   let initState' = initState { shapeOrdering = gpiOrdering }

   putStrLn (bgColor Cyan $ style Italic "   Style program warnings   ")
   let warns = warnings transAuto
   putStrLn (color Red $ intercalate "\n" warns ++ "\n")
   return initState'

-- | After monomorphizing the translation's type (to make sure it's computed), we generalize the type again, which means
-- | it's again under a typeclass lambda. (#166)
castTranslation :: Translation Float -> (forall a . Autofloat a => Translation a)
castTranslation t =
      let res = M.map castFieldDict (trMap t) in
      t { trMap = res }
      where
        castFieldDict :: FieldDict Float -> (forall a . Autofloat a => FieldDict a)
        castFieldDict dict = M.map castFieldExpr dict

        castFieldExpr :: FieldExpr Float -> (forall a . (Autofloat a) => FieldExpr a)
        castFieldExpr e =
          case e of
             FExpr te -> FExpr $ castTagExpr te
             FGPI n props -> FGPI n $ M.map castTagExpr props

        castTagExpr :: TagExpr Float -> (forall a . Autofloat a => TagExpr a)
        castTagExpr e =
           case e of
             Done v ->
                let res = case v of
                          FloatV x -> FloatV (r2f x)
                          PtV (x, y) -> PtV (r2f x, r2f y)
                          PtListV pts -> PtListV $ map (app2 r2f) pts
                          PathDataV d -> PathDataV $ map castPath d
                          -- More boilerplate not involving floats
                          IntV x -> IntV x
                          BoolV x -> BoolV x
                          StrV x -> StrV x
                          FileV x -> FileV x
                          StyleV x -> StyleV x
                in Done res
             OptEval e -> OptEval e -- Expr only contains floats

        castPath :: Path' Float -> (forall a . Autofloat a => Path' a)
        castPath p = case p of
                     Closed elems -> Closed $ map castElem elems
                     Open elems -> Open $ map castElem elems

        castElem :: Elem Float -> (forall a . Autofloat a => Elem a)
        castElem e = case e of
                     Pt pt -> Pt $ app2 r2f pt
                     CubicBez pts -> CubicBez $ app3 (app2 r2f) pts
                     CubicBezJoin pts -> CubicBezJoin $ app2 (app2 r2f) pts
                     QuadBez pts -> QuadBez $ app2 (app2 r2f) pts
                     QuadBezJoin pt -> QuadBezJoin $ app2 r2f pt

-------------------------------
-- Sampling code
-- TODO: should this code go in the optimizer?

numStateSamples :: Int
numStateSamples = 500

-- | Resample the varying state.
-- | We are intentionally using a monomorphic type (float) and NOT using the translation, to avoid slowness.
resampleVState :: [Path] -> [Shape Double] -> StdGen -> (([Shape Double], [Double], [Double]), StdGen)
resampleVState varyPaths shapes g =
    let (resampledShapes, rng') = sampleShapes g shapes
        (resampledFields, rng'') = resampleFields varyPaths rng'
        -- make varying map using the newly sampled fields (we do not need to insert the shape paths)
        varyMapNew = mkVaryMap (filter isFieldPath $ varyPaths) resampledFields
        varyingState = shapes2floats resampledShapes varyMapNew $ varyPaths
    in ((resampledShapes, varyingState, resampledFields), rng'')

-- | Update the translation to get the full state.
updateVState :: State -> (([Shape Double], [Double], [Double]), StdGen) -> State
updateVState s ((resampledShapes, varyingState', fields'), g) =
    let polyShapes = toPolymorphics resampledShapes
        uninitVals = map toTagExpr $ shapes2vals polyShapes $ uninitializedPaths s
        trans' = insertPaths (uninitializedPaths s) uninitVals (transr s)
                    -- TODO: shapes', rng' = sampleConstrainedState (rng s) (shapesr s) (constrs s)
        varyMapNew = mkVaryMap (filter isFieldPath $ varyingPaths s) fields'
    in s { shapesr = polyShapes,
           rng = g,
           transr = trans' { warnings = [] }, -- Clear the warnings, since they aren't relevant anymore
           varyingState = map r2f varyingState',
           paramsr = (paramsr s) { weight = initWeight, optStatus = NewIter } }
    -- NOTE: for now we do not update the new state with the new rng from eval.
    -- The results still look different because resampling updated the rng.
    -- Therefore, we do not have to update rng here.

-- | Iterate a function that uses a generator, generating an infinite list of results with their corresponding updated generators.
iterateS :: (a -> (b, a)) -> a -> [(b, a)]
iterateS f g = let (res, g') = f g in
               (res, g') : iterateS f g'

-- | Compare two states and return the one with less energy.
lessEnergyOn :: ([Double] -> Double) -> (([Shape Double], [Double], [Double]), StdGen)
                             -> (([Shape Double], [Double], [Double]), StdGen) -> Ordering
lessEnergyOn f ((_, vs1, _), _) ((_, vs2, _), _) = compare (f vs1) (f vs2)

-- | Resample the varying state some number of times (sampling each new state from the original state, but with an updated rng).
-- | Pick the one with the lowest energy and update the original state with the lowest-energy-state's info.
resampleBest :: Int -> State -> State
resampleBest n s =
          if n < 2 then error "Need to sample at least two states" else
          let optInfo = paramsr s
              -- Take out the relevant information for resampling
              f       = (overallObjFn optInfo) (rng s) (float2Double $ weight optInfo)
              (varyPaths, shapes, g) = (varyingPaths s, shapesr s, rng s)
              -- Partially apply resampleVState with the params that don't change over a resampling
              resampleVStateConst = resampleVState varyPaths shapes
              sampledResults = take n $ iterateS resampleVStateConst g
              res = minimumBy (lessEnergyOn f) sampledResults
              {- (trace ("energies: " ++ (show $ map (\((_, x, _), _) -> f x) sampledResults)) -}
          in initPolicy $ updateVState s res

------- Other possibly-useful utility functions (not currently used)

-- | Evaluate the objective function on the varying state (with the penalty weight, which should be the same between state).
evalFnOn :: State -> Double
evalFnOn s = let optInfo = paramsr s
                 f       = (overallObjFn optInfo) (rng s) (float2Double $ weight optInfo)
                 args    = map float2Double $ varyingState s
             in f args

-- | Compare two states and return the one with less energy.
lessEnergy :: State -> State -> Ordering
lessEnergy s1 s2 = compare (evalFnOn s1) (evalFnOn s2)

---------- List of policies that can be used with the optimizer

-- Policy stops when value is None
-- Note: if there are no objectives/constraints, policy may return an empty list of functions
-- Policy step = one optimization through to convergence
-- TODO: factor out number of policy steps / other boilerplate? or let it remain dynamic?
-- TODO: factor out the weights on the objective functions / method of combination (in genObjFn)

initPolicyParams :: PolicyParams
initPolicyParams = PolicyParams { policyState = "", policySteps = 0, currFns = [] }

initPolicy :: State -> State
initPolicy s = -- TODO: make this less verbose
    let (policyRes, pstate) = (policyFn s) (objFns s) (constrFns s) initPolicyParams in
    let newFns = DM.fromJust policyRes in
    let stateWithPolicy = s { paramsr = (paramsr s) { overallObjFn = genObjfn (transr s) (filter isObjFn newFns)
                                                                              (filter isConstr newFns) (varyingPaths s) },
                              policyParams = initPolicyParams { policyState = pstate, currFns = newFns } } in
    stateWithPolicy

optimizeConstraints :: Policy
optimizeConstraints objfns constrfns params =
    let (pstate, psteps) = (policyState params, policySteps params) in
    if psteps == 0 then (Just constrfns, "")
    else (Nothing, "") -- Take 1 policy step

optimizeObjectives :: Policy
optimizeObjectives objfns constrfns params =
    let (pstate, psteps) = (policyState params, policySteps params) in
    if psteps == 0 then (Just objfns, "")
    else (Nothing, "") -- Take 1 policy step

-- This is the typical/old Penrose policy
optimizeSumAll :: Policy
optimizeSumAll objfns constrfns params =
    let (pstate, psteps) = (policyState params, policySteps params) in
    if psteps == 0 then (Just $ objfns ++ constrfns, "")
    else (Nothing, "") -- Take 1 policy step

optimizeConstraintsThenObjectives :: Policy
optimizeConstraintsThenObjectives objfns constrfns params =
     let (pstate, psteps) = (policyState params, policySteps params) in
     if psteps == 0 then (Just constrfns, "Constraints") -- Initial policy state
     else if psteps >= 2 then (Nothing, "Done") -- Just constraints then objectives for now, then done
     else if pstate == "Constraints" then (Just objfns, "Objectives")
     else if pstate == "Objectives" then (Just constrfns, "Constraints")
     else error "invalid policy state"

isObjFn f  = optType f == Objfn
isConstr f = optType f == Constrfn

-- TODO: does genObjFns work with an empty list?

max_stars_repo_name: daodaoliang/penrose

max_issues_repo_name: daodaoliang/penrose

max_forks_repo_name: daodaoliang/penrose


text:
{-# LANGUAGE TemplateHaskell #-}
{-# LANGUAGE UndecidableInstances #-}
{-# LANGUAGE DeriveGeneric #-}
{-# LANGUAGE InstanceSigs #-}
{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE StandaloneDeriving #-}
-- {-# LANGUAGE PolyKinds #-}
-- {-# LANGUAGE AllowAmbiguousTypes #-}

{-# OPTIONS_GHC -ddump-splices #-}
{-# OPTIONS_GHC -fprint-explicit-kinds #-}

module Deep.Example where

import           Data.List

import           Deep.Expr

import           Language.Haskell.TH

import           Data.Char (ord)

import           GHC.Generics
import           Data.Void

import           Data.Bifunctor

import           Data.Complex

import           Data.Ratio

import           GHC.Real

import           Data.Proxy

-- import           Debug.Trace

-- deriving instance Generic (Ratio a) -- XXX: This instance should probably be in 'base'

data Nat = Z | S Nat deriving (Generic, Show)

data Example = N Nat | B Bool deriving (Generic, Show)

thExample1 :: Q Exp
thExample1 = do
  exp <-
    [|
      case the @(Either Int Bool) (Right False) of
        Left x -> x * 2
        Right y -> fromEnum y
    |]
  transformExpr exp

thExample2 :: Q Exp
thExample2 = do
  exp <-
    [| case (True, the @Int 7) of
         (x, y) -> fromEnum x + y
    |]
  transformExpr exp

data Example' = N' Int | B' Bool deriving (Show, Generic)

thExample3 :: Q Exp
thExample3 = do
  exp <-
    [|
      case B' False of
        N' n -> N' (n+2)
        B' b -> B' (not b)
    |]
  transformExpr exp

instance GPURep Example' where
  type TypeHead Example' = Example'

data Example3 = X Int Float deriving (Show, Generic)

instance GPURep Example3 where
  type TypeHead Example3 = Example3

data Example4 = E1 Int | E2 Float | E3 Bool deriving (Show, Generic)

instance GPURep Example4 where
  type TypeHead Example4 = Example4

thExample4 :: Q Exp
thExample4 = do
  exp <-
    [| case E2 23.0 of
        E1 x -> 2
        E2 y -> 4
        E3 z -> 6
    |]
  transformExpr exp

data Example5 = A1 Float Float | A2 Int deriving (Show, Generic)

instance GPURep Example5 where
  type TypeHead Example5 = Example5

thExample5 :: Q Exp
thExample5 = do
  exp <-
    [| case A1 2.3 7.5 of
        A2 x -> fromIntegral x
        A1 x y -> x + y
    |]

  transformExpr exp

-- instance GPURep a => GPURep (Ratio a)

transformDecTailRec
  [d|
  thExample6 :: Int -> Bool
  thExample6 x =
    if x == 0
      then True
      else if x == 1
            then False
            else thExample6 (x - 2)
  |]

data IntPair = IntPair Int Int deriving (Show, Generic)

instance GPURep IntPair where
  type TypeHead IntPair = IntPair

thExample7 :: Q Exp
thExample7 = do
  exp <-
    [| case IntPair 1 2 of
        IntPair x y ->
          if x == 0
            then y
            else x
    |]
  r <- transformExpr exp
  return r

-- transformDecTailRec
--   [d|
thExample8 :: IntPair -> Int
thExample8 p =
  case p of
    IntPair x y -> y

-- -- thExample9_aj5t :: IntPair -> Int
-- -- thExample9_aj5t p_aj5x
-- --   = (gpuAbs
-- --        (TailRec
-- --           ((\ thExample9_aj5t
-- --               -> \ p_aj5x
-- --                    -> (CaseExp p_aj5x)
-- --                         ((SafeSumMatch (Data.Proxy.Proxy :: Proxy (TypeHead (Tagged2 (TypeHead IntPair :: *) IntPair (Int, Int)) :: *)))
-- --                            (OneSumMatch
-- --                               (ProdMatch
-- --                                  (\ x_aj5y
-- --                                     -> OneProdMatch
-- --                                          (\ y_aj5z
-- --                                             -> ((IfThenElse ((Equal x_aj5y) (Lit 0)))
-- --                                                   (DoneExp
-- --                                                      ((construct thExample8) p_aj5x)))
-- --                                                  (thExample9_aj5t
-- --                                                     (((construct IntPair)
-- --                                                         ((Sub x_aj5y) (Lit 1)))
-- --                                                        ((Mul x_aj5y) y_aj5z)))))))))
-- --              StepExp)))
-- --       p_aj5x

--   thExample9 :: IntPair -> Int
--   thExample9 p =
--     case p of
--       IntPair x y ->
--         if x == 0
--           then thExample8 p
--           else thExample9 (IntPair (x-1) (x*y))

--   |]

thExample9_aijJ :: IntPair -> Int
thExample9_aijJ p_aijN
  = (gpuAbs
       (TailRec
          ((\ thExample9_aijJ
              -> \ p_aijN
                   -> (CaseExp p_aijN)
                        ((SafeSumMatch (Proxy @(KindOf IntPair)) (Proxy @IntPair))
                           (OneSumMatch
                              (ProdMatch
                                 (\ x_aijO
                                    -> OneProdMatch
                                         (\ y_aijP -> undefined))))))
                                            -- -> ((IfThenElse ((Equal x_aijO) (Lit 0)))
                                            --       (DoneExp
                                            --          ((construct thExample8) p_aijN)))
                                            --      (thExample9_aijJ
                                            --         (((construct IntPair)
                                            --             ((Sub x_aijO) (Lit 1)))
                                            --            ((Mul x_aijO) y_aijP)))))))))
             StepExp)))
      p_aijN

-- transformDecTailRec
--   [d|
--   mandelbrot_nextZ :: (Complex (Ratio Integer), Complex (Ratio Integer)) -> Complex (Ratio Integer)
--   mandelbrot_nextZ t =
--     case t of
--       (c, z) -> (z*z) + c

--   -- shouldFail :: Complex (Ratio Integer) -> (Ratio Integer)
--   -- shouldFail t =
--   --   case t of
--   --     (x, y) -> x + y

--   mandelbrot_helper :: (Int, Complex (Ratio Integer), Complex (Ratio Integer)) -> Maybe Int
--   mandelbrot_helper t =
--     case t of
--       (iters, c, z) ->
--         if iters == 50
--           then Nothing
--           else
--             case z of
--               (:+) real imag ->
--                 if ((real*real) + (imag*imag)) > 4
--                   then Just iters
--                   else mandelbrot_helper (iters+1, c, mandelbrot_nextZ (c, z))

--   mandelbrot_point :: Complex (Ratio Integer) -> Maybe Int
--   mandelbrot_point c = mandelbrot_helper (0, c, 0)

--   |]

ext: hs


text:
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE InstanceSigs #-}
{-# LANGUAGE Rank2Types #-}
{-# LANGUAGE FlexibleInstances    #-}
{-# LANGUAGE UndecidableInstances #-}
{-# LANGUAGE IncoherentInstances  #-}
{-# LANGUAGE StandaloneDeriving  #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE RecordWildCards #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE ForeignFunctionInterface #-}
{-# OPTIONS_GHC -fno-warn-orphans #-}

module Data.Vector.Image.FFT (
  fft1d
, fft2d
, fft3d
, Direction(..)
) where

import Prelude hiding(map)
import qualified Prelude as P
import qualified Data.List as L
import Data.Vector.Storable hiding(forM_,map,replicate,fromList,toList,(!?),(!))
import Data.Maybe
import Numeric.LinearAlgebra.Data hiding (fromList,toList,(!))
import Numeric.LinearAlgebra.HMatrix hiding (fromList,toList,(!))
import qualified Data.Vector.Storable as V
import Foreign
import GHC.Real
import Foreign.C.String
import Foreign.C.Types
import qualified Control.Monad as C
import Data.Vector.Image.Color.RGB
import System.IO.Unsafe(unsafePerformIO)

foreign import ccall unsafe "fft1d" c_fft1d :: Ptr (Complex Double) -> CInt -> CInt -> IO (Ptr (Complex Double))
foreign import ccall unsafe "fft2d" c_fft2d :: Ptr (Complex Double) -> CInt -> CInt -> CInt -> IO (Ptr (Complex Double))
foreign import ccall unsafe "fft3d" c_fft3d :: Ptr (Complex Double) -> CInt -> CInt -> CInt -> CInt -> IO (Ptr (Complex Double))
foreign import ccall unsafe "&fftw_free" c_p_freeComplexDouble :: FunPtr(Ptr (Complex Double) -> IO ())

data Direction = Forward | Backword

fft1d :: Vector (Complex Double) -> Direction -> Vector (Complex Double)
fft1d vec dir = unsafePerformIO $ do
  let n = V.length vec
  p <- unsafeWith vec $ \v -> do
    case dir of
      Forward -> c_fft1d v (fromIntegral n) 1
      Backword -> c_fft1d v (fromIntegral n) 0
  fp <- newForeignPtr c_p_freeComplexDouble p
  return $ unsafeFromForeignPtr fp 0 n

fft2d :: Vector (Complex Double) -> Int -> Direction -> Vector (Complex Double)
fft2d vec m dir = unsafePerformIO $ do
  let nm = V.length vec
  let n = nm `div` m
  p <- unsafeWith vec $ \v -> do
    case dir of
      Forward -> c_fft2d v (fromIntegral m) (fromIntegral n) 1
      Backword -> c_fft2d v (fromIntegral m) (fromIntegral n) 0
  fp <- newForeignPtr c_p_freeComplexDouble p
  return $ unsafeFromForeignPtr fp 0 nm

fft3d :: Vector (Complex Double) -> Int -> Int -> Direction -> Vector (Complex Double)
fft3d vec l m dir = unsafePerformIO $ do
  let lmn = V.length vec
  let n = lmn `div` (l*m)
  p <- unsafeWith vec $ \v -> do
    case dir of
      Forward -> c_fft3d v (fromIntegral l) (fromIntegral m) (fromIntegral n) 1
      Backword -> c_fft3d v (fromIntegral l) (fromIntegral m) (fromIntegral n) 0
  fp <- newForeignPtr c_p_freeComplexDouble p
  return $ unsafeFromForeignPtr fp 0 lmn

max_stars_repo_name: junjihashimoto/img-vector

max_issues_repo_name: junjihashimoto/img-vector

max_forks_repo_name: junjihashimoto/img-vector


text:
{- MathHmatrix.hs
 - Mapping of Linear algebra routines via the hMatrix library.
 - hMatrix uses GSL and a BLAS implementation such as ATLAS.
 -
 - Timothy A. Chagnon
 - CS 636 - Spring 2009
 -}
module MathHmatrix where

import Numeric.LinearAlgebra

deg2rad :: RealT -> RealT
deg2rad d = d * pi / 180

type RealT = Double
type Vec3f = Vector Double

vec3f :: Double -> Double -> Double -> Vector Double
vec3f x y z = (3 |>) [x,y,z]

zeroVec3f :: Vec3f
zeroVec3f = vec3f 0 0 0
svMul   :: Double -> Vec3f -> Vec3f
svMul = scale
dot     :: Vec3f -> Vec3f -> Double
dot  = Numeric.LinearAlgebra.dot
cross   :: Vec3f -> Vec3f -> Vec3f
cross v1 v2 =
    let [a, b, c] = toList v1 in
    let [d, e, f] = toList v2 in
    vec3f (b*f-c*e) (c*d-a*f) (a*e-b*d)

mag     :: Vec3f -> Double
mag v   = sqrt (magSq v)
magSq   :: Vec3f -> Double
magSq v = v `Numeric.LinearAlgebra.dot` v
norm    :: Vec3f -> Vec3f
norm v  = (1/(mag v)) `svMul` v

type Mat4f = Matrix Double
mat4f :: Vec4f -> Vec4f -> Vec4f -> Vec4f -> Mat4f
mat4f a b c d = fromRows [a,b,c,d]

type Vec4f = Vector Double
vec4f :: Double -> Double -> Double -> Double -> Vector Double
vec4f x y z w = (4 |>) [x,y,z,w]

mvMul     :: Mat4f -> Vec4f -> Vec4f
mvMul  = (<>)
mmMul     :: Mat4f -> Mat4f -> Mat4f
mmMul  = (<>)

id4f :: Mat4f
id4f = ident 4

point4f :: Vec3f -> Vec4f
point4f v = 4 |> ((toList v) ++ [1.0])

point3f :: Vec4f -> Vec3f
point3f v =
    let [x, y, z, w] = toList v in
    if w == 0
        then error "point3f divide by zero"
        else vec3f (x/w) (y/w) (z/w)

direction4f :: Vec3f -> Vec4f
direction4f v = 4 |> ((toList v) ++ [0.0])

direction3f :: Vec4f -> Vec3f
direction3f = subVector 0 3

vec3fElts :: Vec3f -> (RealT, RealT, RealT)
vec3fElts v =
    let [x,y,z] = toList v in
    (x,y,z)

type ColMat3f = Matrix Double
colMat3f :: Vec3f -> Vec3f -> Vec3f -> ColMat3f
colMat3f a b c = fromColumns [a,b,c]

detMat3f :: ColMat3f -> RealT
detMat3f = det


text:
{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE DeriveAnyClass   #-}
{-# LANGUAGE DeriveGeneric    #-}
{-# LANGUAGE FlexibleContexts #-}
module FokkerPlanck.FourierSeries
  ( fourierMellin
  , sampleScale
  -- , computeFourierCoefficients
  , computeHarmonicsArray
  , computeHarmonicsArraySparse
  -- , computeHarmonicsArrayGPU
  , getHarmonics
  , computeThetaRHarmonics
  , computeFourierSeriesThetaR
  -- , computeFourierSeriesR2
  , normalizeFreqArr
  -- , normalizeFreqArr'
  , plotThetaDimension
  , computeFourierSeriesOfLogPolarHarmonicsArray
  , computeRectangularInverseHarmonics
  ) where

import           Array.UnboxedArray                     as UA
import           Control.DeepSeq
import           Data.Array.IArray                      as IA
import           Data.Array.Repa                        as R
import           Data.Binary
import           Data.Complex
import           Data.DList                             as DL
import           Data.List                              as L
import           Data.Vector.Generic                    as VG
import           Data.Vector.Unboxed                    as VU
import           FokkerPlanck.BrownianMotion
import           FokkerPlanck.Histogram
import           GHC.Generics                           (Generic)
import           Text.Printf
import           Utils.Array
import           Utils.Parallel
import           Utils.Time
import           Graphics.Rendering.Chart.Backend.Cairo
import           Graphics.Rendering.Chart.Easy
import           Image.IO
import           System.FilePath
import           Utils.SimpsonRule
import Numeric.LinearAlgebra as NL
import Numeric.LinearAlgebra.Data as NL
import            Utils.Distribution
import Pinwheel.Base

-- {-# INLINE fourierMellin' #-}
-- fourierMellin' :: Int -> Int -> (Double, Double) -> Complex Double
-- fourierMellin' !angularFreq !radialFreq (!x, !y) =
--   if x == 0 && y == 0
--     then 0
--     else exp $ (-1 * log rho) :+ (tf * atan2 y x + rf * (log rho))

{-# INLINE sampleScale #-}
sampleScale ::
     Double
  -> Double
  -> Double
  -> Double
  -> Double
  -> [Particle]
  -> Complex Double
sampleScale !phiFreq !rhoFreq !thetaFreq !rFreq !halfLogPeriod =
  L.sum .
  L.map
    (\(Particle newPhi newRho theta newR _) ->
       if newRho == 0 -- || newR == 0
         then 0
         else (cos (phiFreq * newPhi + thetaFreq * (theta - newPhi)) :+ 0) *
              (cis $
               (-pi) * (rhoFreq * (newRho) + rFreq * (newR - newRho)) /
               halfLogPeriod))

-- {-# INLINE computeFourierCoefficients #-}
-- computeFourierCoefficients ::
--      [Double]
--   -> [Double]
--   -> [Double]
--   -> [Double]
--   -> Double
--   -> Double
--   -> [DList Particle]
--   -> Histogram (Complex Double)
-- computeFourierCoefficients !phiFreqs !rhoFreqs !thetaFreqs !rFreqs !halfLogPeriod !deltaLogRho !xs =
--   Histogram
--     [L.length phiFreqs, L.length rhoFreqs, L.length thetaFreqs, L.length rFreqs]
--     (L.length ys) .
--   toUnboxedVector .
--   UA.accum
--     (+)
--     0
--     ( (1, 1, 1, 1)
--     , ( L.length rFreqs
--       , L.length thetaFreqs
--       , L.length rhoFreqs
--       , L.length phiFreqs)) .
--   L.concat .
--   parMap
--     rdeepseq
--     (\(particle@(Particle phi rho theta r)) ->
--        let !n = floor $ (log r + halfLogPeriod) / deltaLogRho
--            !samples =
--              -- L.map moveParticle  $
--              particle :
--              [ -- (Particle
--              --      phi
--              --      rho
--              --      theta
--              --      (exp $ (fromIntegral i * deltaLogRho - halfLogPeriod)))
--              -- | i <- [0 .. n]
--              ]
--        in L.map
--             (\((rFreq, i), (thetaFreq, j), (rhoFreq, k), (phiFreq, l)) ->
--                ( (i, j, k, l)
--                , -- (deltaLogRho :+ 0) *
--                  1/ (16 * pi^2 * halfLogPeriod :+ 0) *
--                  sampleScale
--                    phiFreq
--                    rhoFreq
--                    thetaFreq
--                    rFreq
--                    halfLogPeriod
--                    samples))
--             freqs) $
--   ys
--   where
--     !ys = DL.toList . DL.concat $ xs
--     !freqs =
--       [ (rFreq', thetaFreq', rhoFreq', phiFreq')
--       | rFreq' <- L.zip rFreqs [1 ..]
--       , thetaFreq' <- L.zip thetaFreqs [1 ..]
--       , rhoFreq' <- L.zip rhoFreqs [1 ..]
--       , phiFreq' <- L.zip phiFreqs [1 ..]
--       ]

-- {-# INLINE normalizeFreqArr #-}
-- normalizeFreqArr ::
--      Double
--   -> [Double]
--   -> [Double]
--   -> R.Array U DIM4 (Complex Double)
--   -> R.Array U DIM4 (Complex Double)
-- normalizeFreqArr !std !phiFreqs !rhoFreqs arr =
--   computeUnboxedS .
--   R.traverse3
--     arr
--     (fromListUnboxed (Z :. L.length phiFreqs) phiFreqs)
--     (fromListUnboxed (Z :. L.length rhoFreqs) rhoFreqs)
--     (\sh _ _ -> sh) $ \fArr fPhi fRho idx@(Z :. r :. theta :. rho :. phi) ->
--     fArr idx *
--     ((exp $
--       (-1) *
--       ((fPhi (Z :. phi)) ^ 2 + (fPhi (Z :. theta)) ^ 2 + (fRho (Z :. rho)) ^ 2 +
--        (fRho (Z :. r)) ^ 2) /
--       2 /
--       (std ^ 2)) :+
--      0)

{-# INLINE normalizeFreqArr #-}
normalizeFreqArr ::
     Double
  -> Double
  -> [Double]
  -> [Double]
  -> R.Array U DIM4 (Complex Double)
  -> R.Array U DIM4 (Complex Double)
normalizeFreqArr !stdA !stdR !phiFreqs !rhoFreqs arr
  | stdA == 0 && stdR == 0 = arr
  | stdA == 0 =
    computeUnboxedS .
    R.traverse3
      arr
      (fromListUnboxed (Z :. L.length phiFreqs) phiFreqs)
      (fromListUnboxed (Z :. L.length rhoFreqs) rhoFreqs)
      (\sh _ _ -> sh) $ \fArr fPhi fRho idx@(Z :. r :. theta :. rho :. phi) ->
      fArr idx *
      (exp $
       ((-1) * ((fRho (Z :. rho)) ^ 2 + (fRho (Z :. r)) ^ 2) / (2 * stdR ^ 2)) :+
       0)
  | stdR == 0 =
    computeUnboxedS .
    R.traverse3
      arr
      (fromListUnboxed (Z :. L.length phiFreqs) phiFreqs)
      (fromListUnboxed (Z :. L.length rhoFreqs) rhoFreqs)
      (\sh _ _ -> sh) $ \fArr fPhi fRho idx@(Z :. _ :. theta :. rho :. phi) ->
      fArr idx *
      ((exp $
        (-1) * ((fPhi (Z :. phi)) ^ 2 + (fPhi (Z :. theta)) ^ 2) /
        (2 * stdA ^ 2)) :+
       0)
  | otherwise =
    computeUnboxedS .
    R.traverse3
      arr
      (fromListUnboxed (Z :. L.length phiFreqs) phiFreqs)
      (fromListUnboxed (Z :. L.length rhoFreqs) rhoFreqs)
      (\sh _ _ -> sh) $ \fArr fPhi fRho idx@(Z :. r :. theta :. rho :. phi) ->
      fArr idx *
      ((exp $
        (-1) * ((fPhi (Z :. phi)) ^ 2 + (fPhi (Z :. theta)) ^ 2) /
        (2 * stdA ^ 2) -
        ((fRho (Z :. rho)) ^ 2 + (fRho (Z :. r)) ^ 2) / (2 * stdR ^ 2)) :+
       0)

-- {-# INLINE computeHarmonicsArray #-}
computeHarmonicsArray ::
     (VG.Vector vector (Complex Double), NFData (vector (Complex Double)))
  => Int
  -> Double
  -> Int
  -> Double
  -> [Double]
  -> [Double]
  -> [Double]
  -> [Double]
  -> Double
  -> Double
  -> IA.Array (Int, Int) (vector (Complex Double))
computeHarmonicsArray !numRows !deltaRow !numCols !deltaCol !phiFreqs !rhoFreqs !thetaFreqs !rFreqs !halfLogPeriod !cutoff =
  let !centerRow = div numRows 2
      !centerCol = div numCols 2
      rangeFunc1 xs ys = [(L.head xs - L.last ys) .. (L.last xs - L.head ys)]
      rangeFunc2 xs ys =
        (round (L.head xs - L.last ys), round (L.last xs - L.head ys))
      (!tfLB, !tfUB) = rangeFunc2 phiFreqs thetaFreqs
      (!rfLB, !rfUB) = rangeFunc2 rhoFreqs rFreqs
      !xs =
        parMap
          rdeepseq
          (\(!rf, !tf) ->
             let !vec =
                   VG.convert .
                   toUnboxed . computeS . fromFunction (Z :. numCols :. numRows) $ \(Z :. c :. r) ->
                     let !x = fromIntegral (c - centerCol) * deltaCol
                         !y = fromIntegral (r - centerRow) * deltaRow
                         !rho = 0 + (sqrt $ x ^ 2 + y ^ 2)
                         !rho2 =
                           fromIntegral $
                           (c - centerCol) ^ 2 + (r - centerRow) ^ 2
                     in if rho2 > cutoff ^ 2 || (rho <= 0) -- || pi * rho < (abs tf) -- || log ((rho + 1) / rho) > (1 / (2 * rf))
                          then 0
                               -- (x :+ y) ** (tf :+ 0) *
                               -- ((x ^ 2 + y ^ 2) :+ 0) **
                               -- (((-tf - 0.5) :+ rf) / 2)
                          else ((rho :+ 0) ** ((-1) :+ rf)) *
                               (cis (tf * atan2 y x))
             in ((round rf, round tf), vec))
          [ (rf, tf)
          | rf <- rangeFunc1 rhoFreqs rFreqs
          , tf <- rangeFunc1 phiFreqs thetaFreqs
          ]
  in IA.array ((rfLB, tfLB), (rfUB, tfUB)) xs

-- {-# INLINE computeHarmonicsArraySparse #-}
computeHarmonicsArraySparse ::
     Int
  -> Double
  -> Int
  -> Double
  -> [Double]
  -> [Double]
  -> [Double]
  -> [Double]
  -> Double
  -> Double
  -> Double
  -> IA.Array (Int, Int) (R.Array U DIM2 (Complex Double))
computeHarmonicsArraySparse !numRows !deltaRow !numCols !deltaCol !phiFreqs !rhoFreqs !thetaFreqs !rFreqs !halfLogPeriod !cutoff !envelopeSigma =
  let !centerRow = div numRows 2
      !centerCol = div numCols 2
      rangeFunc1 xs ys =
        [(round $ L.head xs - L.last ys) .. (round $ L.last xs - L.head ys)]
      rangeFunc2 xs ys =
        ((round $ (L.head xs - L.last ys)), (round $ (L.last xs - L.head ys)))
      (!tfLB, !tfUB) = rangeFunc2 phiFreqs thetaFreqs
      (!rfLB, !rfUB) = rangeFunc2 rhoFreqs rFreqs
      !xs =
        parMap
          rseq
          (\(!rf, !tf) ->
             let !arr =
                   computeS . fromFunction (Z :. numCols :. numRows) $ \(Z :. c :. r) ->
                     let !x = fromIntegral (c - centerCol) * deltaCol
                         !y = fromIntegral (r - centerRow) * deltaRow
                         !rho = ((sqrt $ x ^ 2 + y ^ 2))
                         !rho2 =
                           fromIntegral $
                           (c - centerCol) ^ 2 + (r - centerRow) ^ 2
                     in if (rho <= 0) || rho2 > cutoff ^ 2
                          then 0
                          else (x :+ y) ** (fromIntegral tf :+ 0) *
                               (((x ^ 2 + y ^ 2) :+ 0) **
                                (((-(fromIntegral tf) - envelopeSigma) :+ (2 * pi / halfLogPeriod * fromIntegral rf)) /
                                 2))
             in deepSeqArray arr ((rf, tf), arr))
          [ (rf, tf)
          | rf <- rangeFunc1 rhoFreqs rFreqs
          , tf <- rangeFunc1 phiFreqs thetaFreqs
          ]
  in IA.array ((rfLB, tfLB), (rfUB, tfUB)) xs

-- {-# INLINE computeHarmonicsArrayGPU #-}
-- computeHarmonicsArrayGPU ::
--      (VG.Vector vector (Complex Double), NFData (vector (Complex Double)))
--   => Int
--   -> Double
--   -> Int
--   -> Double
--   -> [Double]
--   -> [Double]
--   -> [Double]
--   -> [Double]
--   -> Double
--   -> Double
--   -> IA.Array (Int, Int) (vector (Complex Double))
-- computeHarmonicsArrayGPU !numRows !deltaRow !numCols !deltaCol !phiFreqs !rhoFreqs !thetaFreqs !rFreqs !halfLogPeriod !cutoff =
--   let !centerRow = div numRows 2
--       !centerCol = div numCols 2
--       rangeFunc1 xs ys =
--         [round (L.head xs - L.last ys) .. round (L.last xs - L.head ys)]
--       rangeFunc2 xs ys =
--         (round (L.head xs - L.last ys), round (L.last xs - L.head ys))
--       (!tfLB, !tfUB) = rangeFunc2 phiFreqs thetaFreqs
--       (!rfLB, !rfUB) = rangeFunc2 rhoFreqs rFreqs
--       !xs =
--         parMap
--           rdeepseq
--           (\(!rf, !tf) ->
--              let !vec =
--                    VG.convert .
--                    toUnboxed . computeS . fromFunction (Z :. numCols :. numRows) $ \(Z :. c :. r) ->
--                      let !x = fromIntegral (c - centerCol) * deltaCol
--                          !y = fromIntegral (r - centerRow) * deltaRow
--                          !r2 = x ^ 2 + y ^ 2
--                      in if (x == 0 && y == 0) || r2 > cutoff ^ 2
--                           then 0
--                           else cis $
--                                fromIntegral tf * atan2 y x +
--                                (pi * (fromIntegral rf * 0.5 * log r2) /
--                                halfLogPeriod - pi)
--              in ((rf, tf), vec))
--           [ (rf, tf)
--           | rf <-  (rangeFunc1 rhoFreqs rFreqs)
--           , tf <- L.reverse (rangeFunc1 phiFreqs thetaFreqs)
--           ]
--   in IA.array ((rfLB, tfLB), (rfUB, tfUB)) xs

{-# INLINE getHarmonics #-}
getHarmonics ::
     (VG.Vector vector e)
  => IA.Array (Int, Int) (vector e)
  -> Double
  -> Double
  -> Double
  -> Double
  -> vector e
getHarmonics harmonicsArray phiFreq rhoFreq thetaFreq rFreq =
  harmonicsArray IA.! (round (rhoFreq - rFreq), round $ (phiFreq - thetaFreq))

-- {-# INLINE computeFourierSeriesR2 #-}
-- computeFourierSeriesR2 ::
--      Int
--   -> Double
--   -> Int
--   -> Double
--   -> [Double]
--   -> [Double]
--   -> [Double]
--   -> [Double]
--   -> Double
--   -> Double
--   -> IA.Array (Int, Int) (VU.Vector (Complex Double))
--   -> R.Array U DIM4 (Complex Double)
--   -> [VU.Vector (Complex Double)]
-- computeFourierSeriesR2 numRows deltaRow numCols deltaCol phiFreqs rhoFreqs thetaFreqs rFreqs halfLogPeriod cutoff harmonicsArray arr =
--   let (Z :. (!numRFreq) :. (!numThetaFreq) :. _ :. _) = extent arr
--       !initVec = VU.replicate (VU.length (harmonicsArray IA.! (0, 0))) 0
--   in parMap
--        rdeepseq
--        (\((!r, !rFreq), (!theta, !thetaFreq)) ->
--           L.foldl'
--             (\(!vec) ((!rho, !rhoFreq), (!phi, !phiFreq)) ->
--                VU.zipWith
--                  (+)
--                  vec
--                  (VU.map
--                     (* (arr R.! (Z :. r :. theta :. rho :. phi)))
--                     (getHarmonics harmonicsArray phiFreq rhoFreq thetaFreq rFreq)))
--             initVec $
--           (,) <$> (L.zip [0 ..] rhoFreqs) <*> (L.zip [0 ..] phiFreqs)) $
--      (,) <$> (L.zip [0 ..] rFreqs) <*> (L.zip [0 ..] thetaFreqs)

-- {-# INLINE computeThetaRHarmonics #-}
computeThetaRHarmonics ::
     Int -> Int -> [Double] -> [Double] -> Double -> [[(Complex Double)]]
computeThetaRHarmonics !numOrientation !numScale !thetaFreqs !rFreqs !halfLogPeriod =
  let !deltaTheta = 2 * pi / fromIntegral numOrientation
      !deltaScale = 2 * halfLogPeriod / fromIntegral numScale
      !numRFreq = L.length rFreqs
      !numThetaFreq = L.length thetaFreqs
  in if numScale == 1
       then parMap
              rdeepseq
              (\o ->
                 L.map
                   (\freq -> cis $ freq * (fromIntegral o * deltaTheta))
                   thetaFreqs) $
            [0 .. numOrientation - 1]
       else parMap
              rdeepseq
              (\(!o, !s) ->
                 R.toList .
                 R.traverse2
                   (fromListUnboxed (Z :. numRFreq) rFreqs)
                   (fromListUnboxed (Z :. numThetaFreq) thetaFreqs)
                   (\_ _ -> (Z :. numRFreq :. numThetaFreq)) $ \fRFreq fThetaFreq idx@(Z :. rFreq :. thetaFreq) ->
                   -- exp $
                   -- (-0.5 * (fromIntegral s * deltaScale - halfLogPeriod)) :+
                   -- ((fRFreq (Z :. rFreq)) *
                   --  (fromIntegral s * deltaScale - halfLogPeriod) +
                   --  fThetaFreq (Z :. thetaFreq) * fromIntegral o * deltaTheta)
                   cis
                     ((fRFreq (Z :. rFreq)) *
                      (pi * (fromIntegral s * deltaScale - halfLogPeriod) /
                       halfLogPeriod) +
                      fThetaFreq (Z :. thetaFreq) * fromIntegral o * deltaTheta)
               ) $
            (,) <$> [0 .. numOrientation - 1] <*> [0 .. numScale - 1]

{-# INLINE computeFourierSeriesThetaR #-}
computeFourierSeriesThetaR ::
     (VG.Vector vector (Complex Double), NFData (vector (Complex Double)))
  => [[(Complex Double)]]
  -> [vector (Complex Double)]
  -> [vector (Complex Double)]
computeFourierSeriesThetaR !harmonics !vecs =
  let !initVec = VG.replicate (VG.length . L.head $ vecs) 0
  in parMap
       rdeepseq
       (L.foldl'
          (\vec0 (!vec1, !v) -> VG.zipWith (+) vec0 . VG.map (* v) $ vec1)
          initVec .
        L.zip vecs) $
     harmonics

plotThetaDimension ::
     (R.Source s Double)
  => FilePath
  -> String
  -> (Int, Int)
  -> R.Array s DIM3 Double
  -> IO ()
plotThetaDimension folderPath prefix (x', y') inputArr = do
  let centerCol = div cols 2
      centerRow = div rows 2
      x = x' + centerCol
      y = y' + centerRow
      (Z :. numOrientation :. cols :. rows) = extent inputArr
      centerArr =
        fromFunction (Z :. (1 :: Int) :. cols :. rows) $ \(Z :. _ :. i :. j) ->
          if i == x && j == y
            then 1 :: Double
            else 0
  plotImageRepa (folderPath </> printf "%s(%d,%d)_center.png" prefix x' y') .
    ImageRepa 8 . computeS $
    centerArr
  let ys' = R.toList . R.slice inputArr $ (Z :. R.All :. x :. y)
      !maxY = L.maximum ys'
      -- ys = L.map (/maxY) ys'
      ys = ys'
      deltaTheta = (360 ::Double) / fromIntegral numOrientation
      xs = [fromIntegral x * deltaTheta | x <- [0 .. numOrientation - 1]]
  toFile def (folderPath </> printf "%s(%d,%d)_theta.png" prefix x' y') $ do
    layout_title .= printf "%s(%d,%d)" prefix x' y'
    layout_x_axis . laxis_generate .= scaledAxis def (0, 359)
    layout_y_axis . laxis_generate .= scaledAxis def (0, maxY)
    plot (line "" [L.zip xs ys])

-- The codes below compute Fourier series in R2

computeFourierSeriesOfLogPolarHarmonicsArray ::
     (VG.Vector vector (Complex Double), NFData (vector (Complex Double)))
  => Double
  -> Double
  -> Int
  -> Int
  -> Int
  -> Int
  -> Int
  -> Double
  -> IA.Array (Int, Int) (vector (Complex Double))
computeFourierSeriesOfLogPolarHarmonicsArray !radius !delta !r2Freq !phiFreq !rhoFreq !thetaFreq !rFreq !halfLogPeriod =
  let ((!angularFreqLB, !angularFreqUB), angularFreqs) =
        freqRange phiFreq thetaFreq
      ((!radialFreqLB, !radialFreqUB), radialFreqs) =
        freqRange phiFreq thetaFreq
      !numPoints' = round $ (2 * radius + 1) / delta
      !numPoints =
        if odd numPoints'
          then numPoints'
          else numPoints' - 1
      !center = div numPoints 2
      !period = 2 * radius + 1
      !periodConstant = delta * 2 * pi / period
      !numR2Freq = 2 * r2Freq + 1
      !std = fromIntegral $ div r2Freq 2
      !simpsonWeights =
        toUnboxed . computeS $
        (computeWeightArrFromListOfShape [numPoints, numPoints] :: R.Array D DIM2 (Complex Double))
      !weightedRectangularHarmonics =
        parMap
          rdeepseq
          (\(freq1, freq2) ->
             VU.zipWith (*) simpsonWeights .
             toUnboxed . computeS . fromFunction (Z :. numPoints :. numPoints) $ \(Z :. i :. j) ->
               cis
                 (-periodConstant *
                   fromIntegral (freq1 * (i - center) + freq2 * (j - center))))
          [ (freq1, freq2)
          | freq1 <- [-r2Freq .. r2Freq]
          , freq2 <- [-r2Freq .. r2Freq]
          ]
      !cartesianGrid =
        toUnboxed . computeS . fromFunction (Z :. numPoints :. numPoints) $ \(Z :. c :. r) ->
          let !x = fromIntegral (c - center) * delta
              !y = fromIntegral (r - center) * delta
          in (x, y)
      !simpsonNorm = (delta / 3) ^ 2 :+ 0
      !xs =
        parMap
          rdeepseq
          (\(!radialFreq, !angularFreq) ->
             let !logPolarHarmonics =
                   VU.map (fourierMellin 0.5 angularFreq radialFreq) cartesianGrid
                 coefficients =
                   fromListUnboxed (Z :. numR2Freq :. numR2Freq) .
                   L.map (VU.sum . VU.zipWith (*) logPolarHarmonics) $
                   weightedRectangularHarmonics
                 !coefficientsGaussian =
                   VG.convert .
                   toUnboxed . computeS . R.traverse coefficients id $ \f idx@(Z :. xFreq :. yFreq) ->
                     simpsonNorm * (f idx) *
                     (gaussian2D
                        (fromIntegral $ xFreq - r2Freq)
                        (fromIntegral $ yFreq - r2Freq)
                        std :+
                      0)
             in ((radialFreq, angularFreq), coefficientsGaussian))
          [ (radialFreq, angularFreq)
          | radialFreq <- radialFreqs
          , angularFreq <- angularFreqs
          ]
  in IA.array ((radialFreqLB, angularFreqLB), (radialFreqUB, angularFreqUB)) xs

computeRectangularInverseHarmonics ::
     (VG.Vector vector (Complex Double), NFData (vector (Complex Double)))
  => Int
  -> Int
  -> Double
  -> Double
  -> Int
  -> [vector (Complex Double)]
computeRectangularInverseHarmonics !numRows !numCols !delta !radius !maxFreq =
  let !period = 2 * radius + 1
      !periodConstant = delta * 2 * pi / period
      !centerRow = div numRows 2
      !centerCol = div numCols 2
  in parMap
       rdeepseq
       (\(col, row) ->
          VG.fromList
            [ cis
              (periodConstant *
               fromIntegral
                 (freq1 * (row - centerRow) + freq2 * (col - centerCol)))
            | freq1 <- [-maxFreq .. maxFreq]
            , freq2 <- [-maxFreq .. maxFreq]
            ])
       [(col, row) | col <- [0 .. numCols - 1], row <- [0 .. numRows - 1]]

-- Utilities

{-# INLINE freqRange #-}
freqRange :: Int -> Int -> ((Int, Int), [Int])
freqRange n m =
  let !x = n + m
  in ((-x, x), [-x .. x])

max_stars_repo_name: XinhuaZhang/Stochastic-Completion-Field

max_issues_repo_name: XinhuaZhang/Stochastic-Completion-Field

max_forks_repo_name: XinhuaZhang/Stochastic-Completion-Field