(* port of taylor functions, taylor interval *)

(*
The first part of the file implements basic operations on type taylor_interval.

Then a type tfunction is defined that represents a twice continuously
differentiable function of six variables.  It can be evaluated, which
is the taylor_interval data associated with it.

Sometimes a tfunction f is used to represent an inequality f < 0.
(See recurse.hl.
*)


module Taylor = struct

open Interval;;
open Univariate;;
open Line_interval;;


(* general utilities *)


let center_form(x,z) =
  let ( + ) , ( - ), ( / ) = up(); upadd,upsub,updiv in
  let y = if x = z then x else (x + z) / 2.0 in
  let w = max (z - y) (y - x) in
  let _ = (w >= 0.0) or failwith "centerform" in
     (y,w);;

(* start with taylor interval operations *)

let make_taylor_interval (l1,w1,dd1) = {l = l1; w = w1; dd=dd1;};;

let ti_add (ti1,ti2) =
  let _ = (ti1.w = ti2.w) or failwith ("width mismatch in ti") in
   make_taylor_interval( ladd ti1.l ti2.l,ti1.w, iadd ti1.dd ti2.dd);;
   
let ti_sub (ti1,ti2) =
	let _ = (ti1.w = ti2.w) or failwith ("width mistmatch in ti") in
		make_taylor_interval( lsub ti1.l ti2.l, ti1.w, isub ti1.dd ti2.dd );;

let ti_scale (ti,t) =
   make_taylor_interval( smul ti.l t, ti.w, imul ti.dd t );;
   


let taylor_error ti =
  let ( + ), ( * ) , ( / )= up(); upadd, upmul, updiv in
	(ti.w * ti.w * iabs ti.dd) / 2.0;;

let upper_bound ti = 
  let e = taylor_error ti in
  let ( + ), ( * ) = up(); upadd, upmul in
	ti.l.f.hi + ti.w * iabs ti.l.df + e;;

let lower_bound ti = 
  let e = taylor_error ti in
  let ( + ), ( * ),( - ) = down(); downadd,downmul,downsub in
    ti.l.f.lo - e  + ( ~-. (ti.w)) * iabs ti.l.df;;
	
(*    t + List.fold_left2 (fun a b c -> a + ( ~-. b) * iabs c) 0.0 ti.w ti.l.df;; *)

let upper_partial ti = 
  let ( + ), ( * ) =   up(); upadd,upmul in
	let err = ti.w * (max ti.dd.hi (~-. (ti.dd.lo))) in
      err + Interval.sup ( ti.l.df );;

let lower_partial ti = 
  let ( + ), ( * ), ( - ) = down();downadd,downmul,downsub in
    let err = ti.w * min ti.dd.lo (~-. (ti.dd.hi)) in
      Interval.inf ( ti.l.df ) + err;;
	  
	  


let ti_mul (ti1, ti2) =
  let _ = (ti1.w = ti2.w) or failwith ("width mistmatch in ti") in
  let f = mk_interval (lower_bound ti1, upper_bound ti1) and
      g = mk_interval (lower_bound ti2, upper_bound ti2) and
      df = mk_interval (lower_partial ti1, upper_partial ti1) and
      dg = mk_interval (lower_partial ti2, upper_partial ti2) and
      ( + ), ( * ) = iadd, imul in
    make_taylor_interval( lmul ti1.l ti2.l, ti1.w, 
			  ti1.dd * g + two * df * dg + f * ti2.dd);;

	  
(* primitive A *)

type primitiveA = {
  hfn : float -> line;
  second : float -> float -> interval;
};;

let make_primitiveA (h1,s1) = {hfn = h1; second = s1; };;

let unitA = 
	make_primitiveA (
      (fun y -> line_unit),
      (fun x z -> zero)
);;

let evalf4A pA w x y z =
  make_taylor_interval(
    pA.hfn y,
    w,
    pA.second x z
  );;

let line_estimateA pA y = pA.hfn y;;

(* primitive U *)

type primitiveU = {
  uv: univariate;
};;

let mk_primitiveU uv1 = { uv = uv1; };;

let line_estimateU p y = 
  let y0 = y in
  let t = mk_interval(y0,y0) in
  let d = eval p.uv t 1 in
	mk_line (eval p.uv t 0, d);;

let evalf4U p w x y z =
	let t = mk_interval(x,z) in
	make_taylor_interval(
		line_estimateU p y,
		w,
		eval p.uv t 2
	);;

type tfunction = 
  |  Prim_a of primitiveA
  |  Uni of primitiveU
  |  Plus of tfunction * tfunction
  |  Minus of tfunction * tfunction
  |  Product of tfunction * tfunction
  |  Scale of tfunction * interval
  |  Uni_compose of univariate * tfunction
  |  Composite of tfunction * tfunction;;

let unit = Prim_a unitA;;

let x1 = Uni (mk_primitiveU ux);;


(* This is one of the most difficult functions in the interval code.
   It uses the chain rule to compute the second partial derivatives with
   respect to x(i) x(j), of a function composition

   F(x1,...,x6) = f(g1(x1,...x6),g2(x1,...x6),...,g6(x1,...x6)).

   (F i j) = sum {k m} (f k m) (gk i) (gm j)     + sum {r} (f r) (gr i j).

   Fast performance of this function is very important, especially
   when many of the functions g* are constant.
   There is a bit of imperative programming here, in computing the sums.

   Note that ( + ) and ( * ) have different types in various subsections.
*)

let eval_composite =
  fun hdr p w ->
	(* wide and narrow ranges of p *)
    let (aw,bw) = lower_bound p, upper_bound p in
    let (a,b) = p.l.f.lo, p.l.f.hi in
	(* wide and narrow widths from a to b *)
    let (u,wu,wf) = 
		let ( + ),( - ),( / ) = up();upadd,upsub,updiv in
		let u = (a + b) / 2.0 in
		let wu = max (bw - u) (u - aw) in
		let wf = max (b - u) (u - a) in
			(u, wu, wf) in
	let (fu:taylor_interval) = hdr wu aw u bw in
    let fpy = 
		let t = make_taylor_interval(fu.l,wf,fu.dd) in
			mk_line (
				mk_interval(lower_bound t, upper_bound t),
				mk_interval(lower_partial t, upper_partial t)) in
	(* use chain rule imperatively to compute narrow first derivative *)
	let lin = 
		let ( * ) = imul in
			mk_line (fpy.f, fpy.df * p.l.df) in
	(* second derivative init *)
	let fW_partial = mk_interval(lower_partial fu, upper_partial fu) in
	let pW_partial = mk_interval(lower_partial p, upper_partial p) in
	let dcw_list =
		let ( + ) = iadd in
		let ( * ) = imul in
			fW_partial * p.dd + fu.dd * pW_partial * pW_partial in
	make_taylor_interval(lin,w,dcw_list);;

let rec evalf4 tf w x y z = match tf with
  | Prim_a p -> evalf4A p w x y z
  | Uni p -> evalf4U p w x y z
  | Plus (tf1,tf2) -> ti_add(evalf4 tf1 w x y z, evalf4 tf2 w x y z)
  | Minus (tf1, tf2) -> ti_sub(evalf4 tf1 w x y z, evalf4 tf2 w x y z)
  | Product (tf1, tf2) -> ti_mul(evalf4 tf1 w x y z, evalf4 tf2 w x y z)
  | Composite(h,g) -> eval_composite (evalf4 h) (evalf4 g w x y z) w
  | Scale (tf,t) -> ti_scale ((evalf4 tf w x y z),t)
  | Uni_compose (uf,tf) -> 
      evalf4 (Composite(Uni (mk_primitiveU uf),tf)) w x y z;;

let evalf tf x z = 
  let (y,w) = center_form (x,z) in
    evalf4 tf w x y z;;

let line_estimate_composite =
  let ( * ) = imul in
    fun h p ->
      let (a,b) = (p.f.lo, p.f.hi) in
      let fN = evalf h a b in
      let fN_partial = mk_interval(lower_partial fN, upper_partial fN) in
		mk_line (fN.l.f, fN_partial * p.df);;

let rec line_estimate tf y = match tf with
  | Prim_a p -> line_estimateA p y
  | Uni p -> line_estimateU p y
  | Plus (p,q) -> ladd (line_estimate p y) (line_estimate q y)
  | Minus (p,q) -> lsub (line_estimate p y) (line_estimate q y)
  | Product (p,q) -> lmul (line_estimate p y) (line_estimate q y)
  | Scale (p,t) -> smul (line_estimate p y) t
  | Uni_compose (uf,tf) -> 
      line_estimate (Composite(Uni (mk_primitiveU uf),tf)) y
  | Composite(h,g) -> line_estimate_composite h (line_estimate g y);;

end;;