From 41620db1f938fc862331a659872b7ca3b25de943 Mon Sep 17 00:00:00 2001
From: Xianda Sun <5433119+sunxd3@users.noreply.github.com>
Date: Thu, 14 Nov 2024 07:37:06 +0000
Subject: [PATCH] Improve `show` function of BUGSModel (#236)

---
 Project.toml     |  2 +-
 src/JuliaBUGS.jl |  1 -
 src/model.jl     | 54 ++++++++++++++++++++++++++++++++++--------------
 3 files changed, 40 insertions(+), 17 deletions(-)

diff --git a/Project.toml b/Project.toml
index 01f959bea..982689e6c 100644
--- a/Project.toml
+++ b/Project.toml
@@ -1,6 +1,6 @@
 name = "JuliaBUGS"
 uuid = "ba9fb4c0-828e-4473-b6a1-cd2560fee5bf"
-version = "0.6.4"
+version = "0.6.5"
 
 [deps]
 AbstractMCMC = "80f14c24-f653-4e6a-9b94-39d6b0f70001"
diff --git a/src/JuliaBUGS.jl b/src/JuliaBUGS.jl
index db888c352..d2add3de3 100644
--- a/src/JuliaBUGS.jl
+++ b/src/JuliaBUGS.jl
@@ -16,7 +16,6 @@ using StaticArrays
 
 import Base: ==, hash, Symbol, size
 import Distributions: truncated
-import AbstractPPL: AbstractContext, evaluate!!
 
 export @bugs
 export compile, initialize!
diff --git a/src/model.jl b/src/model.jl
index 4c178d785..06f39eded 100644
--- a/src/model.jl
+++ b/src/model.jl
@@ -77,23 +77,47 @@ struct BUGSModel{base_model_T<:Union{<:AbstractBUGSModel,Nothing},T<:NamedTuple,
 end
 
 function Base.show(io::IO, model::BUGSModel)
-    if model.transformed
-        println(
-            io,
-            "BUGSModel (transformed, with dimension $(model.transformed_param_length)):",
-            "\n",
-        )
+    # Print model type and dimension
+    space_type =
+        model.transformed ? "transformed (unconstrained)" : "original (constrained)"
+    dim = if model.transformed
+        model.transformed_param_length
     else
-        println(
-            io,
-            "BUGSModel (untransformed, with dimension $(model.untransformed_param_length)):",
-            "\n",
-        )
+        model.untransformed_param_length
+    end
+    printstyled(io, "BUGSModel"; bold=true, color=:blue)
+    println(io, " (parameters are in ", space_type, " space, with dimension ", dim, "):\n")
+
+    # Group and print parameters
+    printstyled(io, "  Model parameters:\n"; bold=true, color=:yellow)
+    grouped_params = Dict{Symbol,Vector{VarName}}()
+    for param in model.parameters
+        sym = AbstractPPL.getsym(param)
+        push!(get!(grouped_params, sym, VarName[]), param)
+    end
+    for (sym, params) in grouped_params
+        param_str = length(params) == 1 ? string(params[1]) : "$(join(params, ", "))"
+        print(io, "    ")
+        printstyled(io, param_str; color=:cyan)
+        println(io)
+    end
+    println(io)
+
+    # Print variable info
+    printstyled(io, "  Variable sizes and types:\n"; bold=true, color=:yellow)
+    for (name, value) in pairs(model.evaluation_env)
+        type_str = if isa(value, Number)
+            "type = $(typeof(value))"
+        else
+            "size = $(size(value)), type = $(typeof(value))"
+        end
+        print(io, "    ")
+        printstyled(io, name; color=:cyan)
+        print(io, ": ")
+        printstyled(io, type_str; color=:green)
+        println(io)
     end
-    println(io, "  Model parameters:")
-    println(io, "    ", join(model.parameters, ", "), "\n")
-    println(io, "  Variable values:")
-    return println(io, "$(model.evaluation_env)")
+    return nothing
 end
 
 """