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

unless GSL::MultiFit.const_defined?("Ndlinear")
  puts("The extension library NDLINEAR is not installed.")
  exit()
end

N_DIM = 3
N_SUM_R = 10
N_SUM_THETA = 11
N_SUM_PHI = 9
R_MAX = 3.0

def psi_real_exact(k, l, m, r, theta, phi)
  rr = GSL::pow(r, l)*Math::exp(-r*r)*GSL::Sf::laguerre_n(k, l + 0.5, 2 * r * r)

  tt = GSL::Sf::legendre_sphPlm(l, m, Math::cos(theta))

  pp = Math::cos(m*phi)

  rr*tt*pp
end

basis_r = Proc.new { |r, y, params|
  params.eval(r, y)
}

basis_theta = Proc.new { |theta, y, params|
  for i in 0...N_SUM_THETA do
    y[i] = GSL::Sf::legendre_Pl(i, Math::cos(theta));
  end
}

basis_phi = Proc.new { |phi, y, params|
  for i in 0...N_SUM_PHI do
    if i%2 == 0
      y[i] = Math::cos(i*0.5*phi)
    else
      y[i] = Math::sin((i+1.0)*0.5*phi)
    end
  end
}


GSL::Rng::env_setup()

k = 5
l = 4
m = 2

NDATA = 3000

N = [N_SUM_R, N_SUM_THETA, N_SUM_PHI]
u = [basis_r, basis_theta, basis_phi]

rng = GSL::Rng.alloc()

bspline = GSL::BSpline.alloc(4, N_SUM_R - 2)
bspline.knots_uniform(0.0, R_MAX)

ndlinear = GSL::MultiFit::Ndlinear.alloc(N_DIM, N, u, bspline)
multifit = GSL::MultiFit.alloc(NDATA, ndlinear.n_coeffs)
vars = GSL::Matrix.alloc(NDATA, N_DIM)
data = GSL::Vector.alloc(NDATA)


for i in 0...NDATA do

  r = rng.uniform()*R_MAX
  theta = rng.uniform()*Math::PI
  phi = rng.uniform()*2*Math::PI

  psi = psi_real_exact(k, l, m, r, theta, phi)

  dpsi = rng.gaussian(0.05*psi)

  vars[i,0] = r
  vars[i,1] = theta
  vars[i,2] = phi

  data[i] = psi + dpsi
end

#GSL::MultiFit::Ndlinear::design(vars, X, ndlinear)
X = GSL::MultiFit::Ndlinear::design(vars, ndlinear)

coeffs, cov, chisq, = GSL::MultiFit::linear(X, data, multifit)

rsq = GSL::MultiFit::linear_Rsq(data, chisq)
STDERR.printf("chisq = %e, Rsq = %f\n", chisq, rsq)

eps_rms = 0.0
volume = 0.0
dr = 0.05;
dtheta = 5.0 * Math::PI / 180.0
dphi = 5.0 * Math::PI / 180.0
x = GSL::Vector.alloc(N_DIM)

r = 0.01
while r < R_MAX do
  theta = 0.0
  while theta < Math::PI do
    phi = 0.0
    while phi < 2*Math::PI do
      dV = r*r*Math::sin(theta)*dr*dtheta*dphi
      x[0] = r
      x[1] = theta
      x[2] = phi

      psi_model, err = GSL::MultiFit::Ndlinear.calc(x, coeffs, cov, ndlinear)
      psi = psi_real_exact(k, l, m, r, theta, phi)
      err = psi_model - psi
      eps_rms += err * err * dV;
      volume += dV;

      if phi == 0.0
         printf("%e %e %e %e\n", r, theta, psi, psi_model)
      end

      phi += dphi
    end
    theta += dtheta
  end
  printf("\n");

  r += dr
end

eps_rms /= volume
eps_rms = Math::sqrt(eps_rms)
STDERR.printf("rms error over all parameter space = %e\n", eps_rms)