#!/usr/bin/env ruby
require("gsl")
include MultiFit

procf = Proc.new { |x, t, y, sigma, f|
  a = x[0]; lambda = x[1]; b = x[2]
  n = t.size
  for i in 0...n do
    yi = a*Math::exp(-lambda*t[i]) + b
    f[i] = (yi - y[i])/sigma[i]
  end
}

procdf = Proc.new { |x, t, y, sigma, jac|
  a = x[0]; lambda = x[1]
  n = t.size
  for i in 0...n do
    ti = t[i]
    si = sigma[i]
    ei = Math::exp(-lambda*ti)
    jac.set(i, 0, ei/si)
    jac.set(i, 1, -ti*a*ei/si)
    jac.set(i, 2, 1.0/si)
  end
}

n = 40
np = 3

f = MultiFit::Function_fdf.new(procf, procdf, np)

r = Rng.new()
t = Vector.new(n)
y = Vector.new(n)
sigma = Vector.new(n)
File.open("expdata.dat", "w") do |fp|
  for i in 0...n do
    t[i] = i
    y[i] = 1.0 + 5*Math::exp(-0.1*t[i]) + r.gaussian(0.1)
    sigma[i] = 0.1
    fp.printf("%d %g %g\n", t[i], y[i], sigma[i])
  end
end

f.set_data(t, y, sigma)
x = Vector.new([1.0, 0.0, 0.0])

solver = FdfSolver.new(FdfSolver::LMSDER, n, np)
solver.set(f, x)

iter = 0
solver.print_state(iter)
begin
  iter += 1
  status = solver.iterate
  solver.print_state(iter)
  status = solver.test_delta(1e-4, 1e-4)
end while status == GSL::CONTINUE and iter < 500

covar = solver.covar(0.0)
position = solver.position

printf("A      = %.5f +/- %.5f\n", position[0], Math::sqrt(covar.get(0, 0)))
printf("lambda = %.5f +/- %.5f\n", position[1], Math::sqrt(covar.get(1, 1)))
printf("b      = %.5f +/- %.5f\n", position[2], Math::sqrt(covar.get(2, 2)))

str = "#{position[0]}*exp(-#{position[1]}*x)+#{position[2]}"

File.open("nonlinearfit.gp", "w") do |fp|
  fp.puts("set grid; plot 'expdata.dat' pt 7 ps 2, #{str} lw 3")
end

#puts("\nTry \"gnuplot -persist nonlinearfit.gp\"")
system("gnuplot -persist nonlinearfit.gp")
File.delete("expdata.dat")
