From ce8b92543f0b685ef56f334edfcf9bb324c680e8 Mon Sep 17 00:00:00 2001
From: edwinnglabs <edwinnglabs@gmail.com>
Date: Fri, 8 Mar 2024 13:59:09 -0800
Subject: [PATCH] stan and exe file loading refinement

---
 orbit/__version__.py |  2 +-
 orbit/utils/stan.py  | 30 +++++++++++++++++++-----------
 setup.py             |  3 +--
 3 files changed, 21 insertions(+), 14 deletions(-)

diff --git a/orbit/__version__.py b/orbit/__version__.py
index 2994a6fc..e1666b22 100644
--- a/orbit/__version__.py
+++ b/orbit/__version__.py
@@ -1 +1 @@
-__version__ = "1.1.4.4.dev12"
+__version__ = "1.1.4.4.dev13"
diff --git a/orbit/utils/stan.py b/orbit/utils/stan.py
index ec3b5c0c..773e620e 100644
--- a/orbit/utils/stan.py
+++ b/orbit/utils/stan.py
@@ -1,12 +1,12 @@
 import json
 import os
 import platform
-from multiprocessing import cpu_count
 
-import cmdstanpy
 import importlib_resources
 from cmdstanpy import CmdStanModel
 
+from typing import Optional
+
 from ..utils.logger import get_logger
 
 logger = get_logger("orbit")
@@ -31,27 +31,33 @@
 
 
 def get_compiled_stan_model(
-    stan_model_name: str = "", stan_file_path: str = ""
+    stan_model_name: str = "", 
+    stan_file_path: Optional[str] = None,
+    exe_file_path: Optional[str] = None,
+    force_compile: bool = False,
 ) -> CmdStanModel:
     """Return a compiled Stan model using CmdStan.
     This includes both prepackaged models as well as user provided models through stan_file_path.
 
     Parameters
     ----------
-    stan_model_name : str
+    stan_model_name :
         The name of the Stan model to use. Use this for the built in models (dlt, ets, ktrlite, lgt)
-    stan_file_path : str, optional
+    stan_file_path : 
         The path to the Stan file to use. If not provided, the default is to search for the file in the 'orbit' package.
         If provided, function will ignore the stan_model_name parameter, and will compile the provide stan_file_path
         into executable in place (same folder as stan_file_path)
+    exe_file_path : 
+        The path to the Stan-exe file to use. If not provided, the default is to search for the file 
+        in the 'orbit' package. If provided, function will ignore the stan_model_name parameter,
+        and will compile the provide stan_file_path into executable in place (same folder as stan_file_path)
     Returns
     -------
     sm : CmdStanModel
         A compiled Stan model.
     """
-    if stan_file_path != "":
-        stan_file = stan_file_path
-        sm = CmdStanModel(stan_file=stan_file)
+    if (stan_file_path is not None) or (exe_file_path is not None):
+        sm = CmdStanModel(stan_file=stan_file, exe_file=exe_file_path, force_compile=force_compile)
     else:
         # Load orbit included cmdstan models
         # Some oddities here. if not providing exe_file, CmdStanModel would delete the actual executable file.
@@ -63,9 +69,11 @@ def get_compiled_stan_model(
         )
         # Check if exe is older than .stan file.
         # This behavior is default on CmdStanModel if we don't have to specify the exe_file.
-        if not os.path.isfile(exe_file) or (
-            os.path.getmtime(exe_file) <= os.path.getmtime(stan_file)
-        ):
+        # if not os.path.isfile(exe_file) or (
+        #     os.path.getmtime(exe_file) <= os.path.getmtime(stan_file)
+        # ):
+
+        if not os.path.isfile(exe_file) or force_compile:
             logger.info(f"Compiling stan model:{stan_file}. ETA 3 - 5 mins.")
             sm = CmdStanModel(stan_file=stan_file)
         else:
diff --git a/setup.py b/setup.py
index 6ba12af2..72d609b0 100644
--- a/setup.py
+++ b/setup.py
@@ -3,7 +3,6 @@
 import platform
 import shutil
 from pathlib import Path
-from shutil import copy
 
 from setuptools import find_packages, setup
 from setuptools.command.build_py import build_py
@@ -90,7 +89,7 @@ def build_model(model: str, model_dir: str):
     model_name = f"{model}.stan"
     model_path = os.path.join(model_dir, model_name)
     print(f"Compiling stan file in place: {model_path}")
-    sm = cmdstanpy.CmdStanModel(stan_file=model_path)
+    _ = cmdstanpy.CmdStanModel(stan_file=model_path)
 
 
 def build_stan_models():