{-# 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?


module Main where

import           Data.Either
import           DistanceGeometry
import           Generators
import           Numeric.LinearAlgebra
import           Numeric.LinearAlgebra.Data
import           ReadWrite
import           System.Directory

idir = "test/molecules"
odir = "test/results"

main :: IO ()
main = do
  putStrLn ""
  putStrLn "generateOneMolecule: Begin"
  generateOneMolecule idir odir "example" 16
  putStrLn "generateOneMolecule: End"
  putStrLn ""
  putStrLn "generateManyMolecules: Begin"
  generateManyMolecules idir odir "methan" 16 100
  putStrLn "generateManyMolecules: End"


{-# OPTIONS_GHC -fno-warn-type-defaults #-}
{-# LANGUAGE FlexibleInstances    #-}
{-# LANGUAGE TypeSynonymInstances #-}
-- ProcessingTime.hs ---
--
-- Filename: ProcessingTime.hs
-- Description:
-- Author: Manuel Schneckenreither
-- Maintainer:
-- Created: Tue Nov 21 10:16:08 2017 (+0100)
-- Version:
-- Package-Requires: ()
-- Last-Updated:
--           By:
--     Update #: 69
-- URL:
-- Doc URL:
-- Keywords:
-- Compatibility:
--
--

-- Commentary:
--
--
--
--

-- Change Log:
--
--
--
--
--
--
--

-- Code:

module SimSim.ProcessingTime.Type where

import           ClassyPrelude
import qualified Data.Map.Strict                     as M
import           Statistics.Distribution
import           Statistics.Distribution.Exponential
import           System.Random.MWC

import           SimSim.Block
import           SimSim.ProductType
import           SimSim.Time

-- ^ Use for instance Statistics.Distribution.Exponential and Statistics.Distribution to generate random numbers.
type ProcessingTime = GenIO -> IO Time

type ProcessingTimes = M.Map Block (M.Map ProductType ProcessingTime)

type ProcTimes = [(Block, [(ProductType, ProcessingTime)])]

fromProcTimes :: ProcTimes -> ProcessingTimes
fromProcTimes xs = M.fromList $ fmap (second M.fromList) xs

-- d :: ExponentialDistribution
-- d = exponential (1/0.8)

-- test :: IO [Double]
-- test = do
--   g <- createSystemRandom
--   mapM (\_ -> genContVar d g) [1..7000]

--
-- ProcessingTime.hs ends here


{-# LANGUAGE RecordWildCards #-}
module DFA where

import qualified Data.Set as Set
import qualified Data.Map.Strict as Map
import Control.Monad.State
import Data.Maybe
import Data.Semigroup (stimes)
import Numeric.LinearAlgebra hiding (remap, (<>))
import NFA
import Utils
import qualified Partition as P
import LruCache

-- | The DFA type is parameterized on the state type, so that we can use
-- Sets during the subset construction and then map them to integers for
-- better efficiency.
data DFA s c = DFA
  { dfaStart :: s
  , dfaFinal :: Set.Set s
  , dfaStates :: Set.Set s
  , dfaTransitions :: Set.Set (c, s, s)
  }
  deriving Show

-- | At first, the DFA states are sets of NFA states
type DfaState1 = Set.Set NfaState

-- | Then, the sets are mapped into integers
type DfaState2 = Int

----------------------------------------------------------------------
--                      The subset construction
----------------------------------------------------------------------

-- | Return all epsilon-edges from an NFA state
epsilonEdges :: Ord c => NFA (Maybe c) -> NfaState -> Set.Set NfaState
epsilonEdges NFA{nfaTransitions} s =
  fromMaybe mempty . Map.lookup Nothing . fromMaybe mempty . Map.lookup s $
    transitionMap nfaTransitions

-- | Compute an epsilon-closure
epsilonClosure :: Ord c => NFA (Maybe c) -> Set.Set NfaState -> Set.Set NfaState
epsilonClosure nfa = go mempty
  where
    go !seen0 !new0
      | Set.null new0 = seen0
      | otherwise =
      let
        seen1 = Set.union seen0 new0
        new1 = Set.difference (sconcatMap (epsilonEdges nfa) new0) seen0
      in go seen1 new1

-- | @move nfa s a@ is a set of NFA states where we can get from one of the
-- states in @s@ after receiving @a@ as an input
move :: Ord c => NFA (Maybe c) -> DfaState1 -> c -> DfaState1
move NFA{nfaTransitions} ss c = sconcatMap
  (\s -> fromMaybe mempty . Map.lookup (Just c) .
    fromMaybe mempty . Map.lookup s $
    transitionMap nfaTransitions)
  ss

-- | Build a DFA from an NFA
nfaToDfa :: forall c . Ord c => [c] -> NFA (Maybe c) -> DFA DfaState1 c
nfaToDfa alphabet nfa =
  let
    start :: DfaState1
    start = epsilonClosure nfa (nfaStart nfa)
    transitions = go Set.empty (Set.singleton $ start) Set.empty
    all_states = concat [ [s1,s2] | (_, s1, s2) <- Set.toList transitions ]
    contains_final :: DfaState1 -> Bool
    contains_final dfa_st = not . Set.null $ Set.intersection (nfaFinal nfa) dfa_st
  in
    DFA
      { dfaStart = start
      , dfaFinal = Set.fromList $ filter contains_final all_states
      , dfaStates = Set.fromList all_states
      , dfaTransitions = transitions
      }
  where
    go
      :: Set.Set DfaState1
      -> Set.Set DfaState1
      -> Set.Set (c, DfaState1, DfaState1)
      -> Set.Set (c, DfaState1, DfaState1)
    go !seen0 !new0 !transitions0
      | Set.null new0 = transitions0
      | otherwise =
      let
        seen1 = Set.union seen0 new0
        -- where can we get from the DFA states in new1?
        (transitions1, new1) =
          flip foldMap new0 $ \s ->
          flip foldMap alphabet $ \c ->
          let
            next :: DfaState1
            next = epsilonClosure nfa (move nfa s c)
          in (Set.singleton (c, s, next), Set.singleton next)

        transitions2 = transitions0 `Set.union` transitions1
        new2 = new1 `Set.difference` seen1
      in go seen1 new2 transitions2

----------------------------------------------------------------------
--                         DFA optimization
----------------------------------------------------------------------

-- | Map all DFA states
mapStates
  :: (Ord s1, Ord s2, Ord c)
  => (s1 -> s2)
  -> DFA s1 c
  -> DFA s2 c
mapStates remap DFA{..} =
  DFA
  { dfaStart = remap dfaStart
  , dfaFinal = Set.map remap dfaFinal
  , dfaStates = Set.map remap dfaStates
  , dfaTransitions = Set.map (\(c,s0,s1) -> (c, remap s0, remap s1)) dfaTransitions
  }

-- | Convert all states of the DFA to consecutive integers starting from 0
mapStatesToInt
  :: (Ord s, Ord c)
  => DFA s c
  -> DFA DfaState2 c
mapStatesToInt dfa =
  let
    stateMap = Map.fromList $ zip (Set.toList . dfaStates $ dfa) [0..]
    remap = (stateMap Map.!)
  in
    mapStates remap dfa

minimizeDfa
  :: forall s c . (Ord c, Ord s)
  => DFA s c
  -> DFA Int c
minimizeDfa dfa@DFA{..} = flip evalState 0 $ do
  initialPartition <- liftM2 (<>)
    (P.singleton dfaFinal)
    (P.singleton $ dfaStates `Set.difference` dfaFinal)

  let
    transitionMap :: Map.Map s (Map.Map c s)
    transitionMap = Map.fromListWith Map.union
      [ (s0, Map.singleton c s1)
      | (c, s0, s1) <- Set.toList dfaTransitions
      ]

    classify :: P.Partition s -> s -> Map.Map c Int
    classify pt s = fmap (P.lookupState pt) $ transitionMap Map.! s

  finalPartition <- P.subpartitionFixpoint classify initialPartition

  return $ mapStates (P.lookupState finalPartition) dfa

----------------------------------------------------------------------
--                Transfer matrices and probabilities
----------------------------------------------------------------------

data TransferMatrix = TransferMatrix
  { tmMatrix :: Matrix Double
    -- ^ transfer matrix itself
  , tmStart :: Vector Double
    -- ^ the indicator vector of the start state
  , tmFinal :: Vector Double
    -- ^ the indicator vector of the final states
  }
  deriving (Show)

-- | Construct the DFA transfer matrix assuming the uniform distribution
-- over nucleotides.
--
-- Return the matrix and the 0-based indices of the start and end states.
dfaTransferMatrix
  :: (Ord c, Ord s)
  => Map.Map c Double -- ^ character probabilities
  -> DFA s c
  -> TransferMatrix
dfaTransferMatrix freqs (mapStatesToInt -> DFA{..}) =
  -- we remap the DFA once again to make sure the numbers are consecutive
  let
    n = Set.size dfaStates
    elts =
      [ ((s0, s1), freqs Map.! c)
      | (c, s0, s1) <- Set.toList dfaTransitions
      ]
  in
    TransferMatrix
    { tmMatrix = (assoc (n,n) 0 . assocSum) elts
    , tmStart = assoc n 0 [ (dfaStart, 1) ]
    , tmFinal = assoc n 0 [ (st, 1) | st <- Set.toList dfaFinal ]
    }
  where
    assocSum = Map.toList . Map.fromListWith (+)

dfaProbability
  :: TransferMatrix
  -> Int -- ^ length of the random string where they motif may occur
  -> Double
dfaProbability TransferMatrix{..} len =
  (tmStart <# stimes len tmMatrix) <.> tmFinal

dfaProbabilityCached
  :: MonadLru Int (Matrix Double) m
  => TransferMatrix
  -> Int -- ^ length of the random string where they motif may occur
  -> m Double
dfaProbabilityCached TransferMatrix{..} len = do
  pw <- lruPow len tmMatrix
  return $ (tmStart <# pw) <.> tmFinal


{-# LANGUAGE BangPatterns #-}

module Data.Benchmarks.UnorderedCollections.Distributions
  ( makeRandomData
  , makeRandomVariateData
    -- * Workloads
  , insertWorkload
  , deleteWorkload
  , uniformLookupWorkload
  , exponentialLookupWorkload

  , loadOnly
  , loadAndUniformLookup
  , loadAndSkewedLookup
  , loadAndDeleteAll
  , loadAndDeleteSome
  , uniformlyMixed
  ) where

import qualified Control.Concurrent.Thread as Th
import           Control.DeepSeq
import           Control.Monad
import           Control.Monad.Reader
import           Control.Monad.Trans (liftIO)
import           Data.Benchmarks.UnorderedCollections.Types
import qualified Data.Vector as V
import qualified Data.Vector.Mutable as MV
import qualified Data.Vector.Unboxed as VU
import           Data.Vector (Vector)
import qualified Data.Vector.Algorithms.Shuffle as V
import           GHC.Conc (numCapabilities)
import           Statistics.Distribution
import           Statistics.Distribution.Exponential
import           System.Random.MWC

import           Criterion.Collection.Types

------------------------------------------------------------------------------
debug :: (MonadIO m) => String -> m ()
debug s = liftIO $ putStrLn s

------------------------------------------------------------------------------
makeRandomData :: (NFData k) =>
                  (GenIO -> IO k)
               -> Int
               -> WorkloadMonad (Vector (k,Int))
makeRandomData !genFunc !n = do
    rng <- getRNG
    debug $ "making " ++ show n ++ " data items"
    keys <- liftIO $ vreplicateM n rng genFunc
    let !v = keys `V.zip` vals
    let !_ = forceVector v
    debug $ "made " ++ show n ++ " data items"
    return $! v

  where
    vals      = V.enumFromN 0 n

------------------------------------------------------------------------------
makeRandomVariateData :: (Ord k, NFData k, Variate k) =>
                         Int
                      -> WorkloadMonad (Vector (k,Int))
makeRandomVariateData = makeRandomData uniform

------------------------------------------------------------------------------
insertWorkload :: (NFData k) => Vector (k,Int) -> Vector (Operation k)
insertWorkload = mapForce $ \(k,v) -> Insert k v

------------------------------------------------------------------------------
deleteWorkload :: (NFData k) => Vector (k,Int) -> Vector (Operation k)
deleteWorkload = mapForce $ \(k,_) -> Delete k

------------------------------------------------------------------------------
uniformLookupWorkload :: (NFData k) =>
                         Vector (k,Int)
                      -> Int
                      -> WorkloadMonad (Vector (Operation k))
uniformLookupWorkload !vec !ntimes = do
    rng <- getRNG
    debug $ "uniformLookupWorkload: generating " ++ show ntimes ++ " lookups"
    v <- liftIO $ vreplicateM ntimes rng f
    debug $ "uniformLookupWorkload: done"
    return v

  where
    !n = V.length vec
    f r = do
        idx <- pick
        let (k,_) = V.unsafeIndex vec idx
        return $ Lookup k
      where
        pick = uniformR (0,n-1) r

------------------------------------------------------------------------------
exponentialLookupWorkload :: (NFData k) =>
                             Double
                          -> Vector (k,Int)
                          -> Int
                          -> WorkloadMonad (Vector (Operation k))
exponentialLookupWorkload !lambda !vec !ntimes = do
    rng <- getRNG
    liftIO $ vreplicateM ntimes rng f
  where
    !dist = exponential lambda
    !n    = V.length vec
    !n1   = n-1
    !nd   = fromIntegral n

    f r = do
        x <- uniformR (0.1, 7.0) r
        let idx = max 0 . min n1 . round $ nd * density dist x
        let (k,_) = V.unsafeIndex vec idx
        return $! Lookup k

------------------------------------------------------------------------------
loadOnly :: (NFData k) =>
            (GenIO -> IO k)     -- ^ rng for keys
         -> WorkloadGenerator (Operation k)
loadOnly !genFunc !n = return $ Workload V.empty f
  where
    f _ = liftM insertWorkload $ makeRandomData genFunc n

------------------------------------------------------------------------------
loadAndUniformLookup :: (NFData k) =>
                        (GenIO -> IO k)  -- ^ rng for keys
                     -> WorkloadGenerator (Operation k)
loadAndUniformLookup !genFunc !n = do
    !vals <- makeRandomData genFunc n
    let !inserts = insertWorkload vals

    return $! Workload inserts $ uniformLookupWorkload vals

------------------------------------------------------------------------------
loadAndSkewedLookup :: (NFData k) =>
                       (GenIO -> IO k)  -- ^ rng for keys
                    -> WorkloadGenerator (Operation k)
loadAndSkewedLookup !genFunc !n = do
    !vals <- makeRandomData genFunc n
    let !inserts = insertWorkload vals
    return $! Workload inserts $ exponentialLookupWorkload 1.5 vals

------------------------------------------------------------------------------
loadAndDeleteAll :: (NFData k) =>
                    (GenIO -> IO k)     -- ^ key generator
                 -> WorkloadGenerator (Operation k)
loadAndDeleteAll !genFunc !n = do
    rng <- getRNG
    !vals <- makeRandomData genFunc n
    let !inserts = insertWorkload vals
    let !deletes = deleteWorkload $ V.shuffle rng vals

    return $ Workload inserts (const $ return deletes)

------------------------------------------------------------------------------
loadAndDeleteSome :: (NFData k) =>
                     (GenIO -> IO k)
                  -> WorkloadGenerator (Operation k)
loadAndDeleteSome !genFunc !n = do
    !vals <- makeRandomData genFunc n
    let !inserts = insertWorkload vals

    return $ Workload inserts $ f vals

  where
    f vals k = do
        rng <- getRNG
        return $ deleteWorkload $ V.take k $ V.shuffle rng vals

------------------------------------------------------------------------------
uniformlyMixed :: (NFData k) =>
                  (GenIO -> IO k)
               -> Double
               -> Double
               -> WorkloadGenerator (Operation k)
uniformlyMixed !genFunc !lookupPercentage !deletePercentage !n = do
    let !numLookups = ceiling (fromIntegral n * lookupPercentage)
    let !numDeletes = ceiling (fromIntegral n * deletePercentage)

    !vals <- makeRandomData genFunc n
    let !inserts = insertWorkload vals
    !lookups <- uniformLookupWorkload vals numLookups

    rng <- getRNG
    let !deletes = deleteWorkload $ V.take numDeletes $ V.shuffle rng vals
    let !out = V.shuffle rng $ V.concat [inserts, lookups, deletes]

    return $! Workload V.empty $ const $ return $ forceVector out

------------------------------------------------------------------------------
-- utilities
------------------------------------------------------------------------------
forceVector :: (NFData k) => Vector k -> Vector k
forceVector !vec = V.foldl' force () vec `seq` vec
  where
    force x v = x `deepseq` v `deepseq` ()

mapForce :: (NFData b) => (a -> b) -> Vector a -> Vector b
mapForce !f !vIn = let !vOut = V.map f vIn
                   in forceVector vOut

-- split a GenIO
splitGenIO :: GenIO -> IO GenIO
splitGenIO rng = VU.replicateM 256 (uniform rng) >>= initialize

-- vector replicateM is slow as dogshit.
vreplicateM :: Int -> GenIO -> (GenIO -> IO a) -> IO (Vector a)
vreplicateM n origRng act = do
    rngs <- replicateM numCapabilities (splitGenIO origRng)
    mv <- MV.new n
    let actions = map (f mv) (parts `zip` rngs)
    results <- liftM (map snd) $ mapM Th.forkIO actions
    _ <- sequence results
    V.unsafeFreeze mv

  where
    parts = partition (n-1) numCapabilities

    f mv ((low,high),rng) = do
        f' low
      where
        f' !idx | idx > high = return ()
                | otherwise = do
                                x <- act rng
                                MV.unsafeWrite mv idx x
                                f' (idx+1)

partition :: Int -> Int -> [(Int,Int)]
partition n k = ys `zip` xs
  where
    xs = map f [1..k]
    ys = 0:(map (+1) xs)
    f i = (i * n) `div` k


{-# LANGUAGE CPP #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE UndecidableInstances #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE DeriveFunctor #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE DefaultSignatures #-}
{-# LANGUAGE Trustworthy #-}
{-# LANGUAGE TypeOperators #-}
{-# OPTIONS_GHC -fenable-rewrite-rules #-}
----------------------------------------------------------------------
-- |
-- Copyright   :  (c) Edward Kmett 2011-2014
-- License     :  BSD3
--
-- Maintainer  :  ekmett@gmail.com
-- Stability   :  experimental
--
-- Representable endofunctors over the category of Haskell types are
-- isomorphic to the reader monad and so inherit a very large number
-- of properties for free.
----------------------------------------------------------------------

module Data.Functor.Rep
  (
  -- * Representable Functors
    Representable(..)
  , tabulated
  -- * Wrapped representable functors
  , Co(..)
  -- * Default definitions
  -- ** Functor
  , fmapRep
  -- ** Distributive
  , distributeRep
  , collectRep
  -- ** Apply/Applicative
  , apRep
  , pureRep
  , liftR2
  , liftR3
  -- ** Bind/Monad
  , bindRep
  -- ** MonadFix
  , mfixRep
  -- ** MonadZip
  , mzipRep
  , mzipWithRep
  -- ** MonadReader
  , askRep
  , localRep
  -- ** Extend
  , duplicatedRep
  , extendedRep
  -- ** Comonad
  , duplicateRep
  , extendRep
  , extractRep
  -- ** Comonad, with user-specified monoid
  , duplicateRepBy
  , extendRepBy
  , extractRepBy
  -- ** WithIndex
  , imapRep
  , ifoldMapRep
  , itraverseRep

  -- ** Generics
  , GRep
  , gindex
  , gtabulate
  , WrappedRep(..)
  ) where

import Control.Applicative
import Control.Applicative.Backwards
import Control.Arrow ((&&&))
#if __GLASGOW_HASKELL__ >= 708
import Data.Coerce
#endif
import Control.Comonad
import Control.Comonad.Trans.Class
import Control.Comonad.Trans.Traced
import Control.Comonad.Cofree
import Control.Monad.Trans.Identity
import Control.Monad.Reader
#if MIN_VERSION_base(4,4,0)
import Data.Complex
#endif
import Data.Distributive
import Data.Foldable (Foldable(fold))
import Data.Functor.Bind
import Data.Functor.Identity
import Data.Functor.Compose
import Data.Functor.Extend
import Data.Functor.Product
import Data.Functor.Reverse
import qualified Data.Monoid as Monoid
import Data.Profunctor.Unsafe
import Data.Proxy
import Data.Sequence (Seq)
import qualified Data.Sequence as Seq
import Data.Semigroup hiding (Product)
import Data.Tagged
import Data.Traversable (Traversable(sequenceA))
import Data.Void
import GHC.Generics hiding (Rep)
import Prelude hiding (lookup)

-- | A 'Functor' @f@ is 'Representable' if 'tabulate' and 'index' witness an isomorphism to @(->) x@.
--
-- Every 'Distributive' 'Functor' is actually 'Representable'.
--
-- Every 'Representable' 'Functor' from Hask to Hask is a right adjoint.
--
-- @
-- 'tabulate' . 'index'  ≡ id
-- 'index' . 'tabulate'  ≡ id
-- 'tabulate' . 'return' ≡ 'return'
-- @

class Distributive f => Representable f where
  -- | If no definition is provided, this will default to 'GRep'.
  type Rep f :: *
  type Rep f = GRep f

  -- |
  -- @
  -- 'fmap' f . 'tabulate' ≡ 'tabulate' . 'fmap' f
  -- @
  --
  -- If no definition is provided, this will default to 'gtabulate'.
  tabulate :: (Rep f -> a) -> f a
  default tabulate :: (Generic1 f, GRep f ~ Rep f, GTabulate (Rep1 f))
                   => (Rep f -> a) -> f a
  tabulate = gtabulate

  -- | If no definition is provided, this will default to 'gindex'.
  index    :: f a -> Rep f -> a
  default index :: (Generic1 f, GRep f ~ Rep f, GIndex (Rep1 f))
                => f a -> Rep f -> a
  index = gindex

-- | A default implementation of 'Rep' for a datatype that is an instance of
-- 'Generic1'. This is usually composed of 'Either', tuples, unit tuples, and
-- underlying 'Rep' values. For instance, if you have:
--
-- @
-- data Foo a = MkFoo a (Bar a) (Baz (Quux a)) deriving ('Functor', 'Generic1')
-- instance 'Representable' Foo
-- @
--
-- Then you'll get:
--
-- @
-- 'GRep' Foo = Either () (Either ('WrappedRep' Bar) ('WrappedRep' Baz, 'WrappedRep' Quux))
-- @
--
-- (See the Haddocks for 'WrappedRep' for an explanation of its purpose.)
type GRep f = GRep' (Rep1 f)

-- | A default implementation of 'tabulate' in terms of 'GRep'.
gtabulate :: (Generic1 f, GRep f ~ Rep f, GTabulate (Rep1 f))
          => (Rep f -> a) -> f a
gtabulate = to1 . gtabulate'

-- | A default implementation of 'index' in terms of 'GRep'.
gindex :: (Generic1 f, GRep f ~ Rep f, GIndex (Rep1 f))
       => f a -> Rep f -> a
gindex = gindex' . from1

type family GRep' (f :: * -> *) :: *
class GTabulate f where
  gtabulate' :: (GRep' f -> a) -> f a
class GIndex f where
  gindex' :: f a -> GRep' f -> a

type instance GRep' (f :*: g) = Either (GRep' f) (GRep' g)
instance (GTabulate f, GTabulate g) => GTabulate (f :*: g) where
  gtabulate' f = gtabulate' (f . Left) :*: gtabulate' (f . Right)
instance (GIndex f, GIndex g) => GIndex (f :*: g) where
  gindex' (a :*: _) (Left  i) = gindex' a i
  gindex' (_ :*: b) (Right j) = gindex' b j

type instance GRep' (f :.: g) = (WrappedRep f, GRep' g)
instance (Representable f, GTabulate g) => GTabulate (f :.: g) where
  gtabulate' f = Comp1 $ tabulate $ fmap gtabulate' $ fmap (curry f) WrapRep
instance (Representable f, GIndex g) => GIndex (f :.: g) where
  gindex' (Comp1 fg) (i, j) = gindex' (index fg (unwrapRep i)) j

type instance GRep' Par1 = ()
instance GTabulate Par1 where
  gtabulate' f = Par1 (f ())
instance GIndex Par1 where
  gindex' (Par1 a) () = a

type instance GRep' (Rec1 f) = WrappedRep f
#if __GLASGOW_HASKELL__ >= 708
-- Using coerce explicitly here seems a bit more readable, and
-- likely a drop easier on the simplifier.
instance Representable f => GTabulate (Rec1 f) where
  gtabulate' = coerce (tabulate :: (Rep f -> a) -> f a)
                 :: forall a . (WrappedRep f -> a) -> Rec1 f a
instance Representable f => GIndex (Rec1 f) where
  gindex' = coerce (index :: f a -> Rep f -> a)
                 :: forall a . Rec1 f a -> WrappedRep f -> a
#else
instance Representable f => GTabulate (Rec1 f) where
  gtabulate' = Rec1 #. tabulate .# (. WrapRep)
instance Representable f => GIndex (Rec1 f) where
  gindex' = (. unwrapRep) #. index .# unRec1
#endif

type instance GRep' (M1 i c f) = GRep' f
instance GTabulate f => GTabulate (M1 i c f) where
  gtabulate' = M1 #. gtabulate'
instance GIndex f => GIndex (M1 i c f) where
  gindex' = gindex' .# unM1

-- | On the surface, 'WrappedRec' is a simple wrapper around 'Rep'. But it plays
-- a very important role: it prevents generic 'Representable' instances for
-- recursive types from sending the typechecker into an infinite loop. Consider
-- the following datatype:
--
-- @
-- data Stream a = a :< Stream a deriving ('Functor', 'Generic1')
-- instance 'Representable' Stream
-- @
--
-- With 'WrappedRep', we have its 'Rep' being:
--
-- @
-- 'Rep' Stream = 'Either' () ('WrappedRep' Stream)
-- @
--
-- If 'WrappedRep' didn't exist, it would be:
--
-- @
-- 'Rep' Stream = Either () (Either () (Either () ...))
-- @
--
-- An infinite type! 'WrappedRep' breaks the potentially infinite loop.
newtype WrappedRep f = WrapRep { unwrapRep :: Rep f }

{-# RULES
"tabulate/index" forall t. tabulate (index t) = t #-}

-- | 'tabulate' and 'index' form two halves of an isomorphism.
--
-- This can be used with the combinators from the @lens@ package.
--
-- @'tabulated' :: 'Representable' f => 'Iso'' ('Rep' f -> a) (f a)@
tabulated :: (Representable f, Representable g, Profunctor p, Functor h)
          => p (f a) (h (g b)) -> p (Rep f -> a) (h (Rep g -> b))
tabulated = dimap tabulate (fmap index)
{-# INLINE tabulated #-}

-- * Default definitions

fmapRep :: Representable f => (a -> b) -> f a -> f b
fmapRep f = tabulate . fmap f . index

pureRep :: Representable f => a -> f a
pureRep = tabulate . const

bindRep :: Representable f => f a -> (a -> f b) -> f b
bindRep m f = tabulate $ \a -> index (f (index m a)) a

mfixRep :: Representable f => (a -> f a) -> f a
mfixRep = tabulate . mfix . fmap index

mzipWithRep :: Representable f => (a -> b -> c) -> f a -> f b -> f c
mzipWithRep f as bs = tabulate $ \k -> f (index as k) (index bs k)

mzipRep :: Representable f => f a -> f b -> f (a, b)
mzipRep as bs = tabulate (index as &&& index bs)

askRep :: Representable f => f (Rep f)
askRep = tabulate id

localRep :: Representable f => (Rep f -> Rep f) -> f a -> f a
localRep f m = tabulate (index m . f)

apRep :: Representable f => f (a -> b) -> f a -> f b
apRep f g = tabulate (index f <*> index g)

distributeRep :: (Representable f, Functor w) => w (f a) -> f (w a)
distributeRep wf = tabulate (\k -> fmap (`index` k) wf)

collectRep :: (Representable f, Functor w) => (a -> f b) -> w a -> f (w b)
collectRep f w = tabulate (\k -> (`index` k) . f <$> w)

duplicateRepBy :: Representable f => (Rep f -> Rep f -> Rep f) -> f a -> f (f a)
duplicateRepBy plus w = tabulate (\m -> tabulate (index w . plus m))

extendRepBy :: Representable f => (Rep f -> Rep f -> Rep f) -> (f a -> b) -> f a -> f b
extendRepBy plus f w = tabulate (\m -> f (tabulate (index w . plus m)))

extractRepBy :: Representable f => (Rep f) -> f a -> a
extractRepBy = flip index

duplicatedRep :: (Representable f, Semigroup (Rep f)) => f a -> f (f a)
duplicatedRep = duplicateRepBy (<>)

extendedRep :: (Representable f, Semigroup (Rep f)) => (f a -> b) -> f a -> f b
extendedRep = extendRepBy (<>)

duplicateRep :: (Representable f, Monoid (Rep f)) => f a -> f (f a)
duplicateRep = duplicateRepBy mappend

extendRep :: (Representable f, Monoid (Rep f)) => (f a -> b) -> f a -> f b
extendRep = extendRepBy mappend

extractRep :: (Representable f, Monoid (Rep f)) => f a -> a
extractRep = extractRepBy mempty

imapRep :: Representable r => (Rep r -> a -> a') -> (r a -> r a')
imapRep f xs = tabulate (f <*> index xs)

ifoldMapRep :: forall r m a. (Representable r, Foldable r, Monoid m)
            => (Rep r -> a -> m) -> (r a -> m)
ifoldMapRep ix xs = fold (tabulate (\(i :: Rep r) -> ix i $ index xs i) :: r m)

itraverseRep :: forall r f a a'. (Representable r, Traversable r, Applicative f)
             => (Rep r -> a -> f a') -> (r a -> f (r a'))
itraverseRep ix xs = sequenceA $ tabulate (ix <*> index xs)

-- * Instances

instance Representable Proxy where
  type Rep Proxy = Void
  index Proxy = absurd
  tabulate _ = Proxy

instance Representable Identity where
  type Rep Identity = ()
  index (Identity a) () = a
  tabulate f = Identity (f ())

instance Representable (Tagged t) where
  type Rep (Tagged t) = ()
  index (Tagged a) () = a
  tabulate f = Tagged (f ())

instance Representable m => Representable (IdentityT m) where
  type Rep (IdentityT m) = Rep m
  index = index .# runIdentityT
  tabulate = IdentityT #. tabulate

instance Representable ((->) e) where
  type Rep ((->) e) = e
  index = id
  tabulate = id

instance Representable m => Representable (ReaderT e m) where
  type Rep (ReaderT e m) = (e, Rep m)
  index (ReaderT f) (e,k) = index (f e) k
  tabulate = ReaderT . fmap tabulate . curry

instance (Representable f, Representable g) => Representable (Compose f g) where
  type Rep (Compose f g) = (Rep f, Rep g)
  index (Compose fg) (i,j) = index (index fg i) j
  tabulate = Compose . tabulate . fmap tabulate . curry

instance Representable w => Representable (TracedT s w) where
  type Rep (TracedT s w) = (s, Rep w)
  index (TracedT w) (e,k) = index w k e
  tabulate = TracedT . unCo . collect (Co #. tabulate) . curry

instance (Representable f, Representable g) => Representable (Product f g) where
  type Rep (Product f g) = Either (Rep f) (Rep g)
  index (Pair a _) (Left i)  = index a i
  index (Pair _ b) (Right j) = index b j
  tabulate f = Pair (tabulate (f . Left)) (tabulate (f . Right))

instance Representable f => Representable (Cofree f) where
  type Rep (Cofree f) = Seq (Rep f)
  index (a :< as) key = case Seq.viewl key of
      Seq.EmptyL -> a
      k Seq.:< ks -> index (index as k) ks
  tabulate f = f Seq.empty :< tabulate (\k -> tabulate (f . (k Seq.<|)))

instance Representable f => Representable (Backwards f) where
  type Rep (Backwards f) = Rep f
  index = index .# forwards
  tabulate = Backwards #. tabulate

instance Representable f => Representable (Reverse f) where
  type Rep (Reverse f) = Rep f
  index = index .# getReverse
  tabulate = Reverse #. tabulate

instance Representable Monoid.Dual where
  type Rep Monoid.Dual = ()
  index (Monoid.Dual d) () = d
  tabulate f = Monoid.Dual (f ())

instance Representable Monoid.Product where
  type Rep Monoid.Product = ()
  index (Monoid.Product p) () = p
  tabulate f = Monoid.Product (f ())

instance Representable Monoid.Sum where
  type Rep Monoid.Sum = ()
  index (Monoid.Sum s) () = s
  tabulate f = Monoid.Sum (f ())

#if MIN_VERSION_base(4,4,0)
instance Representable Complex where
  type Rep Complex = Bool
  index (r :+ i) key = if key then i else r
  tabulate f = f False :+ f True
#endif

instance Representable U1 where
  type Rep U1 = Void
  index U1 = absurd
  tabulate _ = U1

instance (Representable f, Representable g) => Representable (f :*: g) where
  type Rep (f :*: g) = Either (Rep f) (Rep g)
  index (a :*: _) (Left  i) = index a i
  index (_ :*: b) (Right j) = index b j
  tabulate f = tabulate (f . Left) :*: tabulate (f . Right)

instance (Representable f, Representable g) => Representable (f :.: g) where
  type Rep (f :.: g) = (Rep f, Rep g)
  index (Comp1 fg) (i, j) = index (index fg i) j
  tabulate = Comp1 . tabulate . fmap tabulate . curry

instance Representable Par1 where
  type Rep Par1 = ()
  index (Par1 a) () = a
  tabulate f = Par1 (f ())

instance Representable f => Representable (Rec1 f) where
  type Rep (Rec1 f) = Rep f
  index = index .# unRec1
  tabulate = Rec1 #. tabulate

instance Representable f => Representable (M1 i c f) where
  type Rep (M1 i c f) = Rep f
  index = index .# unM1
  tabulate = M1 #. tabulate

newtype Co f a = Co { unCo :: f a } deriving Functor

instance Representable f => Representable (Co f) where
  type Rep (Co f) = Rep f
  tabulate = Co #. tabulate
  index = index .# unCo

instance Representable f => Apply (Co f) where
  (<.>) = apRep

instance Representable f => Applicative (Co f) where
  pure = pureRep
  (<*>) = apRep

instance Representable f => Distributive (Co f) where
  distribute = distributeRep
  collect = collectRep

instance Representable f => Bind (Co f) where
  (>>-) = bindRep

instance Representable f => Monad (Co f) where
  return = pure
  (>>=) = bindRep

#if defined(__GLASGOW_HASKELL__) && __GLASGOW_HASKELL__ >= 704
instance (Representable f, Rep f ~ a) => MonadReader a (Co f) where
  ask = askRep
  local = localRep
#endif

instance (Representable f, Semigroup (Rep f)) => Extend (Co f) where
  extended = extendedRep

instance (Representable f, Monoid (Rep f)) => Comonad (Co f) where
  extend = extendRep
  extract = extractRep

instance ComonadTrans Co where
  lower (Co f) = f

liftR2 :: Representable f => (a -> b -> c) -> f a -> f b -> f c
liftR2 f fa fb = tabulate $ \i -> f (index fa i) (index fb i)

liftR3 :: Representable f => (a -> b -> c -> d) -> f a -> f b -> f c -> f d
liftR3 f fa fb fc = tabulate $ \i -> f (index fa i) (index fb i) (index fc i)