(*

  skeg - Sex, Kinematics, Elegance and Glory.

  Copyright (C) 2004 David Baelde and Samuel Mimram.



  This program is free software; you can redistribute it and/or modify

  it under the terms of the GNU General Public License as published by

  the Free Software Foundation; either version 2 of the License, or

  (at your option) any later version.



  This program is distributed in the hope that it will be useful,

  but WITHOUT ANY WARRANTY; without even the implied warranty of

  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the

  GNU General Public License for more details.



  You should have received a copy of the GNU General Public License

  along with this program; if not, write to the Free Software

  Foundation, Inc., 59 Temple Place - Suite 330,

  Boston, MA 02111-1307, USA.

*)

(** Inverse Kinematic solvers.

@author David Baelde, Samuel Mimram *)





open Skel

(**

Definitions





let mhalfpi = asin (-1.)

let halfpi = asin 1.

let pi = acos (-1.)

let twopi = pi+.pi

let rad2deg = 360./.twopi



let mod2pi t =

  mod_float t twopi



type angle = { theta: float ; phi: float }



module type Param =

sig

  val skel : skel

  val pos : position

  val root : int

end



module type Solver_param =

sig

  val epsilon : float

end



module type Solver =

sig

  val skel : skel

  val pos : position

(** Applies the constraints. *)



  val elementary_move : Constraints.t -> unit

  val elementary_move_precise : float option array -> unit

(** Split a movement into small steps of length epsilon. *)



  val begin_move : ?steps:int -> Constraints.t -> unit

  val begin_move_precise : ?steps:int -> float option array -> unit

  val move : unit->bool

end



module type Acyclic_t =

sig

  val skel : skel

  val pos : position

  val nb_bones : int

  val update_angles : unit -> unit

  val init : int -> unit

  val bones : (int * int * float) array

  val angles : angle array

  val root : int ref

  val equations : float option array -> float array array * float array

end

(**

Utilities for acyclic skeletons





module Acyclic (P : Param) : Acyclic_t =

struct



  let skel = P.skel

  let n = Array.length skel

  let pos = Array.copy P.pos



  let nb_bones =

    let n = ref 0 in

      Array.iter

        (fun line ->

           Array.iter (fun e -> if e then incr n) line)

        skel ;

      assert ((!n/2)*2 = !n) ;

      !n/2



  let root = ref (-1)

(** The directed bones (source, destination, length). *)



  let bones = Array.create nb_bones (0,0,0.)

(** Angles (orientation) of the bones. *)



  let angles = Array.create nb_bones { theta=0. ; phi=0. }



  let do_theta x y =

    let a = atan (y/.x) in

      if x > 0. then a else a+.pi

 

  let update_angles () =

    let do_angles index bone =

      let (i,j,d) = bone in

      let (xi,yi,zi) = pos.(i) in

      let (xj,yj,zj) = pos.(j) in

      let (x,y,z) = (xj-.xi, yj-.yi, zj-.zi) in

      let newangles =

        { theta =

            if x = 0. then

              if y>=0. then halfpi else mhalfpi

            else

              do_theta x y ;

          phi =

            acos (z/.d) } in

        Printf.printf "Angle %d (%f,%f,%f) : (%f,%f)\n%!"

          index x y z

          newangles.theta newangles.phi ;

        angles.(index) <- newangles

    in

      Array.iteri do_angles bones

(** Parse the skel, starting from the root, in order to get the new bones, topologicaly sorted. *)



  let init r =

    (* We'll width-first parse the skel,

     * starting from the root,

     * in order to get the new bones, topologicaly sorted. *)

    let index = ref 0 in

    let add_bones i father =

      let xi,yi,zi = pos.(i) in

        Array.iteri

          (fun j connected ->

             if connected && j <> father then

               let xj,yj,zj = pos.(j) in

               let x,y,z = xj-.xi,yj-.yi,zj-.zi in

               let d = sqrt (x*.x+.y*.y+.z*.z) in

                 bones.(!index) <- (i,j,d) ;

                 incr index ;

                 Printf.printf "Bone (%d,%d,%f)\n" i j d) skel.(i)

    in

    let rec round backup =

      let new_backup = !index in

        for i = backup to !index - 1 do

          let (a,b,_) = bones.(i) in

            try

              add_bones b a

            with

              | Invalid_argument ("index out of bounds") ->

                  failwith "Skel is cyclic!"

        done ;

        if new_backup <> !index then

          round new_backup

    in

      if !root <> r then

        begin

          root := r ;

          add_bones r (-1) ;

          round 0 ;

          if !index <> nb_bones then

            failwith "Skel is not connex!" ;

          update_angles () ;

          Printf.printf "Root changed to %d.\n" !root

        end

        

  let () = init P.root



  (*

    Spherical to cartesian

    ---------------------------------------------------------------------------

    x = r sin(phi) cos(theta)

    y = r sin(phi) sin(theta)

    z = r cos(phi)

    

    So the variation is (r remaining constant):

    dx = r (dphi cos(theta) cos(phi) - dtheta sin(phi) sin(theta))

    dy = r (dphi sin(theta) cos(phi) + dtheta sin(phi) cos(theta))

    dz = - r dphi sin(phi)                                                   *)



  (*

    Encoding the state of the skel,

    where b is the number of bones, p the number of points.

    --------------------------------------------------------------------------

    2*b floats : dtheta_1, dphi_1, dtheta_2, dphi_2 ... dtheta_b, dphi_b

    3*p floats : dx_1, dy_1, dz_1 ...

    So we have a l=(2*b+3*p) vector.

    The equations will be of two sorts.

    1) m moving constraints

    2) 3*b skeleton constraints : the bones's constraints (one per axis)     *)



  let equations move =

    move.(3* !root+0) <- Some 0. ;

    move.(3* !root+1) <- Some 0. ;

    move.(3* !root+2) <- Some 0. ;

    let m =

      let m = ref 0 in

        Array.iter (fun a -> if a <> None then incr m) move ;

        !m

    in

    let p = Array.length pos in

    let b = nb_bones in

    let l = 2*b+3*p in (* vector's length *)

    let eq = Array.make (m+3*b) [||] in

    let cst = Array.make (m+3*b) 0. in

    let eq_i = ref 0 in



      (* Moving constraints *)

      Array.iteri

        (fun i -> function

           | None -> ()

           | Some d ->

               let e = Array.make l 0. in

                 e.(2*b+i) <- 1. ;

                 eq.(!eq_i) <- e ; cst.(!eq_i) <- d ; incr eq_i )

        move ;



      (* Skeleton constraints *)

      for k = 0 to b-1 do

        let (i,j,d) = bones.(k) in

        let {theta=theta;phi=phi} = angles.(k) in

        let ex = Array.make l 0. in

        let ey = Array.make l 0. in

        let ez = Array.make l 0. in

        let o = 2*k in



          (* Here we encode

           * dx_j = dx_i - dtheta sin phi sin theta + dphi cos phi cos theta *)

          ex.(o+0) <- -. d *. (sin phi) *. (sin theta) ;

          ex.(o+1) <-    d *. (cos phi) *. (cos theta) ;

          ex.(2*b+3*j+0) <- -1. ;

          ex.(2*b+3*i+0) <- 1. ;

          eq.(!eq_i) <- ex ; incr eq_i ;



          (* Similar stuff for dy ... *)

          ey.(o+0) <- d *. (sin phi) *. (cos theta) ;

          ey.(o+1) <- d *. (cos phi) *. (sin theta) ;

          ey.(2*b+3*j+1) <- -1. ;

          ey.(2*b+3*i+1) <- 1. ;

          eq.(!eq_i) <- ey ; incr eq_i ;



          (* ... and dz. *)

          ez.(o+1) <- -. d *. (sin phi) ;

          ez.(2*b+3*j+2) <- -1. ;

          ez.(2*b+3*i+2) <- 1. ;

          eq.(!eq_i) <- ez ; incr eq_i ;



      done ;



      assert (!eq_i = m+3*b) ;

      (eq, cst)

end



module Gauss (P : Solver_param) (A : Acyclic_t) : Solver =

struct



  open P

  include A



  let print_rotations angles =

    Array.iteri

      (fun k (i,j,d) ->

         let (dt,dp) = angles.(2*k), angles.(2*k+1) in

           Printf.printf "(%d,%d) rot (%f,%f)\n" i j dt dp) bones



  let elementary_move_precise constraints =

    let (m,b) = equations constraints in

    let a = Matrix.solve m b in

    let oldpos = Array.copy pos in

      print_rotations a ;

      for k = 0 to nb_bones - 1 do

        let (i,j,d) = bones.(k) in

        let (xi,yi,zi) = pos.(i) in

        let (xj,yj,zj) = pos.(j) in

        let old = angles.(k) in

        let (theta,phi) = (mod2pi (old.theta +. a.(2*k)),

                           mod2pi (old.phi +. a.(2*k+1))) in

        let (nxj,nyj,nzj) =

          (xi +. d*.(sin phi)*.(cos theta),

           yi +. d*.(sin phi)*.(sin theta),

           zi +. d*.(cos phi)) in



          pos.(j) <- (nxj,nyj,nzj) ;

          angles.(k) <- {theta=theta;phi=phi} ; 



          Printf.printf

            "dM_%d :\n\tEffectif :\t%f %f %f\n\tAnnonc\233 :\t%f %f %f\n%!"

            j (nxj-.xj) (nyj-.yj) (nzj-.zj)

            a.(2*nb_bones+3*j+0)

            a.(2*nb_bones+3*j+1)

            a.(2*nb_bones+3*j+2) ;



      done



  let elementary_move constraints =

    elementary_move_precise (one_of_three constraints)



  let goal = ref [||]

  let rem_steps = ref 0



  let goal_length g =

    let n = ref 0 in

    let s = ref 0. in

      Array.iter (function

                    | None -> ()

                    | Some a -> incr n ; s := !s+.a*.a) g ;

      !s/.(float !n)



  let get_pos_of_precise i =

    let j = i/3 in

    let c = i-3*j in

    let (x,y,z) = pos.(j) in

      match c with

        | 0 -> x

        | 1 -> y

        | 2 -> z

        | _ -> failwith "get_pos_of_precise"



  let begin_move_in constraints steps =

    goal := constraints ;

    rem_steps := (match steps with None -> 10 | Some i -> i) ;

    for i = 0 to (Array.length !goal)-1 do

      match !goal.(i) with

        | None -> ()

        | Some d -> !goal.(i) <- Some (d+.(get_pos_of_precise i))

    done



  let begin_move_precise ?steps c = begin_move_in c steps

  let begin_move ?steps c = begin_move_in (one_of_three c) steps

(** move () returns true if the move is finished. *)



  let move () =

    let move = Array.copy !goal in

      for i = 0 to (Array.length move)-1 do

        match move.(i) with

          | None -> ()

          | Some d -> move.(i) <- Some (d-.(get_pos_of_precise i))

      done ;

      let length = goal_length move in

        if length < epsilon || !rem_steps = 0 then

          true

        else

          let movee = Array.make (Array.length !goal) None in

            for i = 0 to (Array.length move)-1 do

              match move.(i) with

                | None -> ()

                | Some x ->

                    let xx = x*.epsilon/.length in

                      movee.(i) <- Some xx

            done ;

            try

              Printf.printf "===== Step %d ... \n" !rem_steps ;

              decr rem_steps ;

              elementary_move_precise movee ;

              false

            with

              | e ->

                  Printf.printf "*** %s ***\n%!"

                  (Printexc.to_string e) ;

                  true



end

(** Values for the evolved Y test. Branches are (1,2,3) and (1,2,0,4)*)





module TestLongY =

struct

  let skel =

    [|

      [|false; false; true ; false; true |]; (* 0 is part of branch 1 *)

      [|false; false; true ; false; false|]; (* 1 is the root *)

      [|true ; true ; false; true ; false|]; (* 2 is the branching point *)

      [|false; false; true ; false; false|]; (* 3 is branch 2 *)

      [|true ; false; false; false; false|]; (* 4 is the end of branch 1 *)

    |]

  let root = 1

  let pos = [|

    (0.,1.,-1.) ; (0.,0.,0.) ; (0.,1.,0.) ; (0.,1.,1.) ; (0.,2.,-1.) |]

  let all = skel, pos, root

end

(** Values for the basic Y test. *)





module TestY =

struct

  let skel =

    [|

      [|false; false; true ; false|];

      [|false; false; true ; false|];

      [|true ; true ; false; true |];

      [|false; false; true ; false|]

    |]

  let root = 3

  let pos = [|(1., 1., 0.); (0., 1., 1.); (0., 1., 0.); (0., 0., 0.)|]

  let all = skel, pos, root

end