package scientific.ode

import std.util.*
import std.math.*

public func fun1(x: Float64, y: Array<Float64>, j: Int64): Float64 {
  if (j == 1) {
    return 2.0 * y[1] - y[0]
  } else {
    return y[1]
  }
}

public func fun2(x: Float64, y: Array<Float64>, j: Int64): Float64 {
  if (j == 2) {
    return 2.0 * y[2] - y[1] + 2.0 * y[0]
  }
  if (j == 1) {
    return y[2]
  } else {
    return y[1]
  }
}

public func rungekutta(fun: (Float64, Array<Float64>, Int64) -> Float64, x: Float64, y: Array<Float64>, step: Int32, n: Int64): Float64 {
  var ywork = Array<Float64>(n, {i => Float64(0)})
  var k0 = Array<Float64>(n, {i => Float64(0)})
  var k1 = Array<Float64>(n, {i => Float64(0)})
  var k2 = Array<Float64>(n, {i => Float64(0)})
  var k3 = Array<Float64>(n, {i => Float64(0)})
  var h: Float64 = (x - Float64(0)) / Float64(step)
  var X: Float64 = 0.0
  while (abs(X - x) >= 1e-6) {
    for (j in 0..n) {
	  k0[j] = h * fun(X, y, j)
	}
	for (j in 0..n) {
	  ywork[j] = y[j] + 0.5 * k0[j]
	}
	for (j in 0..n) {
	  k1[j] = h * fun(X + 0.5 * h, ywork, j)
    }
	for (j in 0..n) {
	  ywork[j] = y[j] + 0.5 * k1[j]
    }
	for (j in 0..n) {
	  k2[j] = h * fun(X + 0.5 * h, ywork, j)
	}
	for (j in 0..n) {
	  ywork[j] = y[j] + k2[j]
	}
	for (j in 0..n) {
	  k3[j] = h * fun(X + h, ywork, j)
	}
	for (j in 0..n) {
	  y[j] = y[j] + (k0[j] + 2.0 * k1[j] + 2.0 * k2[j] + k3[j]) / 6.0
	}
	X = X + h
  }
  return y[0]
}

public func testOde() {
  print("微分方程为(y'' - 2y' + y = 0), y0 = 1, y1 = 1 \n")
  var n: Int64 = 2
  var step: Int32 = 20
  print("函数在X = 0.2的值为 ${ rungekutta(fun1, 0.2, Array<Float64>([1.0, 1.0]), step, n) }\n")//1.2214027
  print("函数在X = 0.5的值为 ${ rungekutta(fun1, 0.5, Array<Float64>([1.0, 1.0]), step, n) }\n\n")//1.6487212

  print("微分方程为(y''' - 2y'' + y' - 2y = 0), y0 = -1, y1 = 3, y2 = 6 \n")
  n = 3
  step = 50
  print("函数在X = 0.7的值为 ${ rungekutta(fun2, 0.7, Array<Float64>([-1.0, 3.0, 6.0]), step, n) }\n")//3.169733
  print("函数在X = 1.5的值为 ${ rungekutta(fun2, 1.5, Array<Float64>([-1.0, 3.0, 6.0]), step, n) }\n\n")//20.941558
}