(******************************************************************************
 * Définitions du type des termes différentiels et utilitaires de bases
 *****************************************************************************)

open List;;
open Term_lexer;;

type diff_term   = DiffVar of int * string
                 | DiffApp of diff_term * bunch
                 | DiffAbs of string * int * diff_term
                 | DiffStar
and bunch        = monomial list
and monomial     = (diff_term * int)
;;

type var_degrees = int list;;

(* affichage *)
let rec string_of_diff : diff_term -> string =
  fun t ->
    match t with
        DiffVar(_, x)        -> x
      | DiffStar             -> "*"
      | DiffAbs(x, _, u)     -> "[" ^ x ^ "]" ^ (string_of_diff u)
      | DiffApp(f, b)        ->
          "<" ^ (string_of_diff f) ^ ">" ^
            match b with
                []         -> "{}"
              | [(a, n)]  -> string_of_monome a n
              | (a, n)::b'  -> "{" ^
                  (fold_left (fun s (t, p) ->  s ^ " . " ^
                                (string_of_monome t p))
                     (string_of_monome a n) b') ^ "}"

and string_of_monome : diff_term -> int -> string =
  fun t p ->
    if p = 1 then string_of_diff t
    else match t with
        DiffVar(_, x) -> x ^ "^" ^ (string_of_int p)
      | _ ->
          "(" ^ (string_of_diff t) ^ ")" ^ "^" ^ (string_of_int p)
;;

(* Opérations sur les degrés : une liste ne contenant que des 0 est égale à la
   liste vide ; toutes les opérations supposent que le dernier élément de
   chaque liste passée en argument est non nul et rendent un résultat de
   même. *)
let rec make_degrees : int -> int -> var_degrees =
  fun n d ->
    if d = 0 then []
    else let rec make n =
      if n = 0 then [d]
      else 0::(make (n-1))
    in make n
;;

let rec degrees_mult : int -> var_degrees -> var_degrees =
  fun n ->
    if n = 0 then fun _ -> []
    else let rec mult =
      function
          [] -> []
        | p::deg -> (n*p)::(mult deg)
    in mult
;;

let rec degrees_substract : var_degrees -> var_degrees -> var_degrees =
  fun deg1 deg2 ->
    match (deg1, deg2) with
        ([], []) -> []
      | (deg, []) -> deg
      | ([], _) -> failwith "degrees_substract : deg2 > deg1"
      | (n::deg1', p::deg2') ->
          if n < p then failwith "degrees_substract : deg2 > deg1"
          else if n = p
          then
            let deg = degrees_substract deg1' deg2'
            in if deg = [] then [] else 0::deg
          else (n-p)::(degrees_substract deg1' deg2')
;;

let rec degrees_add : var_degrees -> var_degrees -> var_degrees =
  fun deg1 deg2 ->
    match (deg1, deg2) with
        ([], deg) | (deg, []) -> deg
      | (n::deg1', p::deg2') ->
          let s = n+p
          and deg = degrees_add deg1' deg2'
          in if s = 0 && deg = [] then [] else s::deg
;;

let rec degrees_eq : var_degrees -> var_degrees -> bool =
  fun deg1 deg2 ->
    try degrees_substract deg1 deg2 = []
    with Failure(_) -> false
;;

(* Opérations de base sur les termes différentiels. *)
let rec diff_eq : diff_term -> diff_term -> bool =
  fun t1 t2 ->
    match (t1, t2) with
        (DiffStar, DiffStar) -> true
      | (DiffVar(n1, s1), DiffVar(n2, s2)) ->
          (* On ne compare les noms que si la variable est libre *)
          n1 = n2 && (n1 >= 0 || s1 = s2)
      | (DiffApp(f1, b1), DiffApp(f2, b2)) ->
          diff_eq f1 f2 && bunch_eq b1 b2
      | (DiffAbs(_, d1, u1),  DiffAbs(_, d2, u2)) ->
          d1 = d2 && diff_eq u1 u2
      | _ -> false

(* Égalité de bunch *)
and bunch_eq : bunch -> bunch -> bool =
  fun b1 b2 ->
    match b1 with
        [] -> b2 = []
      | m1::b1' ->
          try
            let b2' = bunch_divide b2 m1
            in bunch_eq b1' b2'
          with Not_found -> false

and bunch_divide : bunch -> monomial -> bunch =
  fun b ((t, p) as m) ->
    match b with
        [] -> if p = 0 then [] else raise Not_found
      | ((t', p') as m')::b' ->
        if diff_eq t t'
        then
          if p' > p then (t', p' - p)::b'
          else if p' = p  then b'
          else bunch_divide b' (t, (p'-p))
        else m'::(bunch_divide b' m)
;;

(* Ajout d'un élément *)
let rec monomial_bunch_mult : monomial -> bunch -> bunch =
  fun ((t, n) as m) ->
    function
        [] -> if n = 0 then [] else [m]
      | (t', p)::b' ->
          if diff_eq t t'
          then (t', n + p)::b'
          else (t', p)::(monomial_bunch_mult m b')
;;

let rec bunch_mult : bunch -> bunch -> bunch =
  fun b1 b2 ->
    match b1 with
        [] -> b2
      | m::b1' -> bunch_mult b1' (monomial_bunch_mult m b2)
;;

let rec bunch_power : bunch -> int -> bunch =
  fun b n ->
    let rec pow =
      function
          [] -> []
        | (t, p)::b -> (t, n*p)::(pow b)
    in if n = 0 then [] else pow b
;;

let rec bunch_length : bunch -> int =
  function
      [] -> 0
    | (_, n):: b' -> n + (bunch_length b')
;;

(*****************************************************************************
 * Parsing
 *****************************************************************************
 *
 * diff   = '$' ident | '#' ident
          | var | '[' var ']' diff  | '<' diff  '>' bunch
 * bunch  = monome | '{' monome {'.' monome} '}'
 * monome = terme ['^' int] | '(' terme ')' ['^' int]
 *
 * Exemples : [x]<x>{x . x} = [x]<x>{x^2} = [x]<x>x^2
 *            [f][x]<f>{<f>{x . x} . <f>{x . x}} = [f][x]<f>(<f>x^2)^2
 *            [f][x]<f>{(<f>x)^2 . <f>x^2}
 *            [f]<f>([z]z)^2
 *
 * define_dmacro "f11" "<<f>{}>{}"
 * define_diff   "s1" "[f][x]<<f>#f11><<f>#f11>#f11"
 * define_diff   "s2" "[f][x]<<f><<f>#f11>#f11>#f11"
 *
 *****************************************************************************)


(* Tables des termes et macros différentiels prédéfinis *)
let diff_table : (string,  diff_term) Hashtbl.t = Hashtbl.create(17)
and diff_macros_table : (string, string) Hashtbl.t = Hashtbl.create(17)
;;

let rec diff_of_stream : token_stream -> diff_term = fun tokens ->
  let rec parse_term : string list -> token_stream -> (diff_term * var_degrees)
    = fun binders  ->
      parser
          [< 'Kwd '$'; 'Ident name >] -> (Hashtbl.find diff_table name, [])
        | [< 'Kwd '#'; 'Ident name >] ->
            parse_term binders
              (term_lexer (Hashtbl.find diff_macros_table name))
        | [< 'Kwd '[';
             'Ident x ?? "no variable in lambda";
             'Kwd ']' ?? "lambda not terminated" ;
             (t, deg_t) = parse_term (x::binders) >] -> (
            match deg_t with
                [] -> (DiffAbs(x, 0, t), [])
              | d::deg' -> (DiffAbs(x, d, t), deg')
          )
        | [< 'Kwd '<';
             (t, deg_t) = parse_term binders;
             'Kwd '>' ?? "application not terminated";
             (b, deg_b) = parse_bunch binders; >] ->
            (DiffApp(t, b), degrees_add deg_t deg_b)
        | [< 'Ident x >] ->
            let index = index_of x binders
            in (DiffVar(index, x), make_degrees index 1)
        | [< 'Kwd '*' >] -> (DiffStar, [])

  and parse_bunch : string list -> token_stream -> (bunch * var_degrees) =
    fun binders ->
      parser
          [< 'Kwd '{';
             b = parse_monome_list binders >] -> b
        | [< ((t, n), deg_t) = parse_monome binders >] ->
            ([(t, n)], degrees_mult n deg_t)

  and parse_monome_list : string list -> token_stream -> (bunch * var_degrees)
    = fun binders ->
      parser
          [< ((t, n)  as m, deg_t) = parse_monome binders; rest >] -> (
            match Stream.peek rest with
                Some(Kwd('.')) ->
                  Stream.junk rest;
                  let (b, deg_b) = parse_monome_list binders rest
                  in (monomial_bunch_mult m b,
                      degrees_add (degrees_mult n deg_t) deg_b)
              | Some(Kwd('}'))  ->
                  Stream.junk rest;
                  ([(t, n)], degrees_mult n deg_t)
              | _ -> raise (Stream.Error("invalid char in bunch"))
          )
        | [< 'Kwd '}' >] -> ([], [])

  and parse_monome : string list -> token_stream -> (monomial * var_degrees) =
    fun binders ->
      parser
        | [< 'Kwd '(';
             (t, deg_t) = parse_term binders;
             'Kwd ')' ?? "missing closing paren";
             n = parse_exponant >] -> ((t, n), deg_t)
        | [< (t, deg_t) = parse_term binders;
             n = parse_exponant >] -> ((t, n), deg_t)

  and parse_exponant : token_stream -> int =
    parser
        [< 'Kwd '^'; 'Int n >] -> n
      | [< >] -> 1

  in let (t, _) = parse_term [] tokens
  in match Stream.peek tokens with
      Some(_) -> raise (Stream.Error("garbage at end of stream"))
    | None    -> t

and diff_of_string : string -> diff_term =
  fun s -> diff_of_stream (term_lexer s)
;;

let define_diff : string -> string -> unit = fun name term ->
  Hashtbl.replace diff_table name (diff_of_string term)
and define_dmacro : string -> string -> unit = fun name term ->
  Hashtbl.replace diff_macros_table name term
;;