open Printf
open Bigarray
type vec = Odepack.vec
let vec a = Array1.of_array float64 fortran_layout a
let grav = 9.81

let matA ms mt th i l=
  let a= Lacaml_D.Mat.make 2 2 0. in
  a.{1,1} <- ms+.mt;
  a.{1,2} <- mt*.l*.(cos th);
  a.{2,1} <- mt *.l*.(cos th);
  a.{2,2} <- i+.mt*.l**2.;
  a

let matG f ms mt l v i =
  let g t x dx th dth =
    let tmp = matA ms mt th i l in
    Lacaml_D.getri tmp;(*inversion de A*)
    let tmp1 = Lacaml_D.Mat.make 2 1 0. in
    tmp1.{1,1} <- (f t x dx th dth) -. v*.dx +. ms*.l*.(sin th) *. (dth)**2.;
    tmp1.{2,1} <- -.ms*.grav*.l*.(sin th);
    let tmp2 = Lacaml_D.gemm tmp tmp1 in
    tmp2 in
  (fun t x dx th dth -> g t x dx th dth)

let vecU x x0 th0 =
  let u = vec [|x0;0.;th0;0.|] in
  u

(*ms = masse segway
  mt = masse tige
  v = coef de frottement du segway
  l = distance du point de rotation au centre de masse de la tige
  i = moment inertie de la tige
  f est la fonction*)
let solve ms mt v l i f =
  let g = matG f ms mt l v i in
  let newf t a b=
    let newg = g t a.{1} a.{2} a.{3} a.{4} in
    b.{1} <- a.{2};
    b.{2} <- newg.{1,1};
    b.{3} <- a.{4};
    b.{4} <- newg.{2,1};
    in
  let var x0 th0 t =
    let u = vecU x0 th0 in
    let tmp = vec (Odepack.lsoda newf u 0.) in
    let tmp1 = vec [|tmp.{2};tmp.{4}|] in
    tmp1 in
  (fun x0 th0 t -> var x0 th0 t)