cmdstanex

defmodule CmdStan do
  @moduledoc """
  Elixir interface to CmdStan.

  This library provides functionality to download, install, and manage
  CmdStan releases from GitHub, similar to the Python CmdStanPy library.
  """

  alias CmdStan.{Installer, Model, Summary}

  @doc """
  Install CmdStan from GitHub releases.

  Downloads and installs a CmdStan release, builds it, and tests the compilation.

  ## Options

    * `:version` - CmdStan version string (e.g., "2.32.2"). Defaults to latest.
    * `:dir` - Installation directory. Defaults to `~/.cmdstan`.
    * `:overwrite` - Boolean, whether to overwrite existing installation. Defaults to false.
    * `:progress` - Boolean, whether to show progress bars. Defaults to false.
    * `:verbose` - Boolean, whether to show verbose output. Defaults to false.
    * `:cores` - Number of CPU cores to use for building. Defaults to number of schedulers.
    * `:interactive` - Boolean, whether to run in interactive mode. Defaults to false.

  ## Examples

      iex> CmdStan.install_cmdstan(version: "2.32.2")
      :ok

      iex> CmdStan.install_cmdstan(dir: "/opt/cmdstan", progress: true)
      :ok

  """
  @spec install_cmdstan(keyword()) :: :ok | {:error, term()}
  def install_cmdstan(opts \\ []) do
    Installer.install_cmdstan(opts)
  end

  @doc """
  Compile a Stan model file to an executable.

  ## Parameters
  - `stan_file`: Path to the Stan model file (.stan)

  ## Returns
  A map containing model information or an error tuple.

  ## Examples

      iex> CmdStan.compile_model("model.stan")
      {:ok, %{name: "model", stan_file: "model.stan", exe_file: "model"}}

  """
  @spec compile_model(String.t()) :: {:ok, map()} | {:error, term()}
  def compile_model(stan_file) do
    Model.compile(stan_file)
  end

  @doc """
  Run gradient diagnostics on a compiled model.

  Compares automatic differentiation gradients with finite difference gradients
  to validate model implementation.

  ## Parameters
  - `model`: Model map returned by `compile_model/1`
  - `opts`: Diagnostic options
    - `:data` - Data map in Stan format (required)
    - `:inits` - Initial parameter values
    - `:epsilon` - Step size for finite difference gradients (default: 1e-6)
    - `:error` - Absolute error threshold (default: 1e-6)
    - `:sig_figs` - Numerical precision for output
    - `:require_gradients_ok` - Whether to raise error if gradients exceed threshold

  ## Returns
  A result map containing diagnostics and metadata.

  ## Examples

      iex> model = %{exe_file: "bernoulli"}
      iex> data = %{"N" => 10, "y" => [0,1,0,0,0,0,0,0,0,1]}
      iex> CmdStan.diagnose(model, data: data)
      {:ok, %{diagnostics: [%{param_idx: 0, value: 0.5, model: -0.123, finite_diff: -0.124, error: 0.001}], metadata: %{...}}}

  """
  @spec diagnose(map(), keyword()) :: {:ok, map()} | {:error, term()}
  def diagnose(model, opts \\ []) do
    Model.diagnose(model, opts)
  end

  @doc """
  Run MCMC sampling on a compiled model.

  ## Parameters
  - `model`: Model map returned by `compile_model/1`
  - `data`: Data map in Stan format
  - `opts`: Sampling options (chains, iterations, etc.)

  ## Returns
  A result map containing draws, metadata, and diagnostics.

  ## Examples

      iex> model = %{exe_file: "bernoulli"}
      iex> data = %{"N" => 10, "y" => [0,1,0,0,0,0,0,0,0,1]}
      iex> CmdStan.sample(model, data, chains: 1, iter: 100)
      {:ok, %{draws: %{"theta" => [0.2, 0.3, ...]}, metadata: %{...}, diagnostics: %{...}}}

  """
  @spec sample(map(), map(), keyword()) :: {:ok, map()} | {:error, term()}
  def sample(model, data, opts \\ []) do
    Model.sample(model, data, opts)
  end

  @doc """
  Run MCMC diagnostics on sampling results.

  Checks for sampling issues like divergent transitions, low E-BFMI values,
  low effective sample sizes, and high R-hat values.

  ## Parameters
  - `fit`: Fit result from `sample/3`

  ## Returns
  A string containing the diagnostic output.

  ## Examples

      iex> fit = CmdStan.sample(model, data, chains: 4, iter: 1000)
      iex> CmdStan.diagnose_fit(fit)
      "Checking sampler transitions treedepth.\\nTreedepth satisfactory for all transitions.\\n...\\n"

  """
  @spec diagnose_fit(map()) :: {:ok, String.t()} | {:error, term()}
  def diagnose_fit(fit) do
    Model.diagnose_fit(fit)
  end

  @doc """
  Compute summary statistics for MCMC fit results.

  ## Parameters
  - `fit`: Fit result from `sample/3`
  - `opts`: Summary options (percentiles, significant figures)

  ## Returns
  A list of maps containing summary statistics for each variable.

  ## Examples

      iex> fit = CmdStan.sample(model, data, chains: 4, iter: 1000)
      iex> {:ok, summary} = CmdStan.summary(fit)
      iex> Enum.find(summary, &(&1.variable == "theta"))

  """
  @spec summary(map(), keyword()) :: {:ok, [map()]} | {:error, term()}
  def summary(fit, opts \\ []) do
    Summary.summary(fit, opts)
  end
end