From f29f93f29f6910eca233ed3466051afbb5953192 Mon Sep 17 00:00:00 2001
From: Ivan <iv.borissov@gmail.com>
Date: Mon, 4 Mar 2024 23:09:01 +0300
Subject: [PATCH] piecewise added

---
 src/HetaSimulator.jl    |  2 ++
 src/heta_funcs.jl       | 15 +++++++++++++++
 test/heta_funcs_test.jl |  7 +++++++
 test/runtests.jl        |  1 +
 4 files changed, 25 insertions(+)
 create mode 100644 src/heta_funcs.jl
 create mode 100644 test/heta_funcs_test.jl

diff --git a/src/HetaSimulator.jl b/src/HetaSimulator.jl
index 76abc8b..30b1cb3 100644
--- a/src/HetaSimulator.jl
+++ b/src/HetaSimulator.jl
@@ -56,6 +56,7 @@ module HetaSimulator
   include("import_platform.jl")
   include("gsa.jl")
   include("save_as_heta.jl")
+  include("heta_funcs.jl")
   
   heta_update()
 
@@ -78,4 +79,5 @@ module HetaSimulator
   export gsa, pearson, partial, standard
   export save_as_heta
   export scale_params, unscale_params
+  export piecewise
 end
diff --git a/src/heta_funcs.jl b/src/heta_funcs.jl
new file mode 100644
index 0000000..26dc29a
--- /dev/null
+++ b/src/heta_funcs.jl
@@ -0,0 +1,15 @@
+piecewise(args...) = begin
+  len = length(args)
+  res = for i in 1:div(len, 2)
+      if args[2*i]
+          return args[2*i-1]
+      end
+  end
+  if !isnothing(res)
+      return res
+  elseif mod(len, 2) === 1
+      return args[len]
+  else
+      return NaN
+  end
+end
\ No newline at end of file
diff --git a/test/heta_funcs_test.jl b/test/heta_funcs_test.jl
new file mode 100644
index 0000000..db2de8a
--- /dev/null
+++ b/test/heta_funcs_test.jl
@@ -0,0 +1,7 @@
+@test piecewise(1, true, 2, false, 3, false, 100) === 1
+@test piecewise(1, false, 2, true, 3, false, 100) === 2
+@test piecewise(1, true, 100) === 1
+@test piecewise(1, false, 100) === 100
+@test piecewise(1, false) === NaN
+@test piecewise(100) == 100
+@test piecewise() === NaN
\ No newline at end of file
diff --git a/test/runtests.jl b/test/runtests.jl
index 5bdbc6b..4e83b96 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -12,5 +12,6 @@ end
 
 @testset "HetaSimulator" begin
   @testset "Single-compartment model without events" begin include("single_comp_test.jl") end
+  @testset "Functions used in heta models" begin include("heta_funcs_test.jl") end
 end