module Network : sig
  open Matrix
  type t
  val make : int -> int -> float -> t
  val formalize : t -> int -> Matrix.matrix -> unit
  val feedforward : t -> float array -> float array
  val train : t -> Matrix.matrix -> Matrix.matrix -> float -> unit
  val ssq : t -> Matrix.matrix -> Matrix.matrix -> float
end = struct 
  open Matrix;;

  type t = {
    ni         : int;
    no         : int;
    weight     : Matrix.matrix; 
    output     : float array;
    delta      : float array;
  };;
  
  let sigmoid x    = 1. /. (1. +. exp (-.x));;
  let dsigmoid x   = x *. (1. -. x);;
  let rand _ _     = (-.1.) +. (2. *. (Random.float 1.0));;
  
  let seq start stop len =
    let by = (stop -. start) /. (float (len - 1)) in
      let step i = start +. (by *. (float i)) in Array.init len step; 
  ;;

  let make ni no k = 
    let ni         = ni + 1 (** for bias *) in
    let weight     = Matrix.init ni no rand
    and output     = Array.make no 0.0 
    and delta      = Array.make no 0.0  in

    (** renormalize weights to be in range of sigmoid *)
    let sum     = Matrix.l1norm weight in 
    let norm w  = w *. (k /. sum) in 
      Matrix.dmap norm weight;
      {
	ni      = ni; 
	no      = no; 
	weight  = weight; 
	output  = output;
	delta   = delta;
      };
  ;;

  let unitsum net j input = 
    (** start with bias *)
    let sum = ref net.weight.(net.ni-1).(j) in 
      (** skip bias in count *)
      for i = 0 to net.ni - 2 do 
	sum := !sum +. input.(i) *. net.weight.(i).(j);
      done;
      net.output.(j) <- !sum;
      !sum;
  ;;

  let feedforward net input = 
    for j = 0 to net.no - 1 do 
      (** copied from unit sum - put here to make fast *)
      let sum = ref net.weight.(net.ni-1).(j) in 
	(** skip bias in count *)
	for i = 0 to net.ni - 2 do 
	  sum := !sum +. input.(i) *. net.weight.(i).(j);
	done;
	net.output.(j) <- sigmoid !sum;
    done;
    Array.copy net.output;
  ;;
  
  let formalize net index inputs = 
    for j = 0 to net.no - 1 do 
      let output = Array.map (unitsum net j) inputs in
	Array.sort compare output;
	net.weight.(net.ni-1).(j) <- -.output.(index);
    done;
  ;;

  let backprop net output input target mu = 
    for k = 0 to net.no - 1 do 
      net.delta.(k) <- dsigmoid output.(k) *. (target.(k) -. output.(k));
    done;
    
    for i = 0 to net.ni - 1 do 
      let wi  = net.weight.(i) 
      and act = if i == net.ni - 1 then 1.0 else input.(i) in
	for k = 0 to net.no - 1 do 
	  let change = net.delta.(k) *. act in
	    wi.(k) <- wi.(k) +. mu *. change;
	done;
    done;
  ;;

  let train net inputs targets mu = 
    for i = 0 to Array.length inputs - 1 do 
      let output = feedforward net inputs.(i); in
	backprop net output inputs.(i) targets.(i) mu;
    done;
  ;;

  let ssq net inputs targets = 
    let error = ref 0. in
      for i = 0 to Array.length inputs - 1 do
	let output = feedforward net inputs.(i); in
	  for k = 0 to net.no - 1 do 
	    error := !error +. ((targets.(i).(k) -. output.(k)) ** 2.);
	  done;
      done;
      !error;
  ;;

end;;