(*****************************************************************************
 * Calcul de l'expansion de Taylor d'un lambda-terme par la machine de Krivine
 *****************************************************************************)

open List;;
open Lambda;;
open Diff;;

type lambda_closure = Clos  of lambda_term * lambda_env
                    | Free  of int
and    diff_closure = Dclos of       bunch * diff_env

and      lambda_env = lambda_closure list
and        diff_env =   diff_closure list

and    lambda_stack = lambda_closure list
and      diff_stack =   diff_closure list

and    lambda_state = (lambda_term * lambda_env * lambda_stack)
and      diff_state = (  diff_term *   diff_env *   diff_stack)
;;


let rec diff_env_make : diff_closure -> int -> diff_env =
  fun c ->
    function
        0 -> [c]
      | n -> Dclos([], [])::(diff_env_make c (n-1))

and diff_env_contract : diff_env -> diff_env -> diff_env =
  fun e1 e2 ->
    match (e1, e2) with
        (e, []) | ([], e) -> e
      | (Dclos(b1, e'1)::e''1, Dclos(b2, e'2)::e''2) ->
          (Dclos(bunch_mult b1 b2, diff_env_contract e'1 e'2))::
            (diff_env_contract e''1 e''2)

and diff_env_power : diff_env -> int -> diff_env =
  fun e n ->
    match e with
        [] -> []
      | Dclos(b, e')::e'' ->
          (Dclos(bunch_power b n, diff_env_power e' n))::(diff_env_power e'' n)
;;

(* Prend en argument un terme clos t et un terme avec ressources simple et
   normal t0 approximant la forme normale de t. Retourne le terme avec resource
   simple du développement de Taylor de t qui se réduit en t0. *)

let diff_of_term : lambda_term -> diff_term -> diff_term = fun t t0 ->
  let rec kam : int -> int -> diff_term -> lambda_state -> diff_state =
    (* Les deux premiers arguments depth et lambdas sont respectivement
       le nombre d'appels récursifs, et le nombre de lambdas sous-lesquels la
       machine est entrée. *)
    fun depth lambdas t0 ->
      if depth > max_depth_limit then failwith "Max recursion depth exceeded"
      else
        function
          (App(u, v), e, s) -> (
            let (u_d, e_d, s'_d) =
              kam (depth + 1) lambdas t0 (u, e, Clos(v, e)::s)
            in
              match s'_d with 
                  Dclos(b, e'_d)::s_d ->
                    (DiffApp(u_d, b), diff_env_contract e_d e'_d, s_d)
                | _ -> failwith "kam failure (App case)"
          )

        | (Abs(x, u), e, c::s) -> (
            let (u_d, e'_d, s_d) = kam (depth + 1) lambdas t0 (u, c::e, s)
            in 
              match e'_d
              with
                  [] -> (DiffAbs(x, 0, u_d), [], Dclos([], [])::s_d)
                | (Dclos(b, _) as c_d)::e_d ->
                    (DiffAbs(x, bunch_length b, u_d), e_d, c_d::s_d)
          )

        | (Abs(x, u), e, []) -> (
            match t0 with
                DiffAbs(_, d, u0) ->
                  let (u_d, e_d, s_d) =
                    kam (depth + 1) (lambdas + 1) u0
                      (u, Free(lambdas + 1)::e, [])
                  in (match (e_d, s_d)
                      with
                          ([], []) -> (DiffAbs(x, d, u_d), [], [])
                        | (Dclos([], [])::e'_d, []) ->
                            (DiffAbs(x, d, u_d), e'_d, [])
                        | _ -> failwith "kam failure (Abs case)"
                     )
              | _ -> failwith "lambda mismatch"
          )

        | (Var(n, x), e, s) -> (
            try
              match nth e n
              with
                  Clos(vn, en) ->
                    let (vn_d, en_d, s_d) =
                      kam (depth + 1) lambdas t0 (vn, en, s)
                    in (DiffVar(n, x),
                        diff_env_make (Dclos([(vn_d, 1)], en_d)) n,
                        s_d)

                | Free(lambda_index) ->
                    let s_d = kam_args depth lambdas
                      (lambdas - lambda_index) (rev s, t0)
                    in (DiffVar(n, x), [], rev s_d)

            with Failure(_) -> failwith "kam failure (Var case)"
          )

        |  (Star, _, []) -> (
             match t0
             with
                 DiffStar -> (DiffStar, [], [])
               | _ -> failwith "star mismatch"
           )

        | (Star, _, _::_) -> failwith "star constant violation"

  and kam_args : int -> int -> int -> (lambda_stack * diff_term) -> diff_env =
    fun depth lambdas var_index ->
      function
          ([], DiffVar(n, _)) ->
            if n = var_index then []
            else failwith "variable mismatch"

        | (c::s, DiffApp(u0, b0)) ->
            let c_d = kam_bunch depth lambdas c b0
            and s_d = kam_args depth lambdas var_index (s, u0)
            in c_d::s_d

        | _ -> failwith "application mismatch"

  and kam_bunch : int -> int -> lambda_closure -> bunch -> diff_closure =
    fun depth lambdas c b0 ->
      match c
      with
          Clos(t, e) -> (
            match b0
            with
                [] -> Dclos([], [])
              | (t0, n)::b0' ->
                  let (t_d, e_d, _)   = kam (depth + 1) lambdas t0 (t, e, [])
                  and Dclos(b', e'_d) = kam_bunch depth lambdas c b0'
                  in let e_d         = diff_env_power e_d n
                  in Dclos(monomial_bunch_mult (t_d, n) b',
                           diff_env_contract e_d e'_d)
          )
        | _ -> failwith "kam failure (bunch case)"

  in
    match kam 0 0 t0 (t, [], [])
    with
        (t_d, [], []) -> t_d
      |        _      -> failwith "kam failure (toplevel)"
;;

(* synthèse *)
let linearize (t: string) (t0: string) : string =
  string_of_diff (diff_of_term (lambda_of_string t) (diff_of_string t0));;