Module Qcert.Translation.Lang.NNRCtoDNNRC

Require Import String.
Require Import List.
Require Import Arith.
Require Import EquivDec.
Require Import Morphisms.
Require Import Utils.
Require Import DataRuntime.
Require Import NNRCSystem.
Require Import DNNRCSystem.
Require Import NRAEnvRuntime.
Require Import Dataframe.

Section NNRCtoDNNRC.
  Context {fruntime:foreign_runtime}.
  
  Fixpoint nnrc_to_dnnrc_base {A:Set} {plug_type:Set} (annot:A) (ctenv tenv:list (var*dlocalization)) (n:nnrc) : @dnnrc_base _ A plug_type :=
    match n with
    | NNRCGetConstant v =>
      match assoc_lookupr equiv_dec ctenv v with
      | None => DNNRCGetConstant annot v
      | Some Vlocal => DNNRCGetConstant annot v
      | Some Vdistr => DNNRCCollect annot (DNNRCGetConstant annot v)
      end
    | NNRCVar v =>
      match lookup equiv_dec tenv v with
      | None => DNNRCVar annot v
      | Some Vlocal => DNNRCVar annot v
      | Some Vdistr => DNNRCCollect annot (DNNRCVar annot v)
      end
    | NNRCConst d => DNNRCConst annot d
    | NNRCBinop b n1 n2 => DNNRCBinop annot b (nnrc_to_dnnrc_base annot ctenv tenv n1) (nnrc_to_dnnrc_base annot ctenv tenv n2)
    | NNRCUnop u n1 => DNNRCUnop annot u (nnrc_to_dnnrc_base annot ctenv tenv n1)
    | NNRCLet v n1 n2 => DNNRCLet annot v (nnrc_to_dnnrc_base annot ctenv tenv n1) (nnrc_to_dnnrc_base annot ctenv ((v,Vlocal)::tenv) n2)
    | NNRCFor v n1 n2 => DNNRCFor annot v (nnrc_to_dnnrc_base annot ctenv tenv n1) (nnrc_to_dnnrc_base annot ctenv ((v,Vlocal)::tenv) n2)
    | NNRCIf n1 n2 n3 => DNNRCIf annot (nnrc_to_dnnrc_base annot ctenv tenv n1) (nnrc_to_dnnrc_base annot ctenv tenv n2) (nnrc_to_dnnrc_base annot ctenv tenv n3)
    | NNRCEither n0 v1 n1 v2 n2 =>
      DNNRCEither annot (nnrc_to_dnnrc_base annot ctenv tenv n0) v1 (nnrc_to_dnnrc_base annot ctenv ((v1,Vlocal)::tenv) n1) v2 (nnrc_to_dnnrc_base annot ctenv ((v2,Vlocal)::tenv) n2)
    | NNRCGroupBy g sl n1 =>
      DNNRCGroupBy annot g sl (nnrc_to_dnnrc_base annot ctenv tenv n1)
    end.

  Definition wf_localization (tl:option dlocalization) (dl:option ddata) :=
    match tl, dl with
    | None, None => True
    | Some Vlocal, Some (Dlocal _) => True
    | Some Vlocal, Some (Ddistr _) => False
    | Some Vdistr, Some (Ddistr _) => True
    | Some Vdistr, Some (Dlocal _) => False
    | _, _ => False
    end.
    
  Definition wf_denv (tenv:list (var*dlocalization)) (denv:list (var*ddata)) :=
    (domain tenv = domain denv) /\
    forall v, wf_localization (lookup equiv_dec tenv v) (lookup equiv_dec denv v).

  Definition wf_cdenv (tenv:list (var*dlocalization)) (denv:list (var*ddata)) :=
    (domain tenv = domain denv) /\
    forall v, wf_localization (assoc_lookupr equiv_dec tenv v) (assoc_lookupr equiv_dec denv v).

  Lemma wf_denv_cons v d tenv denv :
    wf_denv tenv denv -> wf_denv ((v, Vlocal) :: tenv) ((v, Dlocal d) :: denv).

  Lemma domain_lookup_none {A B} (l1:list (string*A)) (l2:list (string*B)) (v:string) :
    domain l1 = domain l2 ->
    assoc_lookupr equiv_dec l1 v = None ->
    assoc_lookupr equiv_dec l2 v = None.

  Lemma domain_lookup_some {A B} (l1:list (string*A)) (l2:list (string*B)) (v:string) :
    domain l1 = domain l2 ->
    (exists x, assoc_lookupr equiv_dec l1 v = Some x) ->
    (exists x', assoc_lookupr equiv_dec l2 v = Some x').

  Lemma wf_cdenv_cons v d tenv denv :
    wf_cdenv tenv denv -> wf_cdenv ((v, Vlocal) :: tenv) ((v, Dlocal d) :: denv).

  Lemma lookup_denv_local v d denv :
    lookup equiv_dec denv v = Some (Dlocal d) ->
    lift Dlocal (lookup equiv_dec (unlocalize_constants denv) v) =
    lookup equiv_dec denv v.

  Lemma rev_localize_comm denv :
    rev (unlocalize_constants denv) = unlocalize_constants (rev denv).

  Lemma lookupr_denv_local v d denv :
    assoc_lookupr equiv_dec denv v = Some (Dlocal d) ->
    lift Dlocal (assoc_lookupr equiv_dec (unlocalize_constants denv) v) =
    assoc_lookupr equiv_dec denv v.

  Lemma lookup_denv_distr v l denv :
    lookup equiv_dec denv v = Some (Ddistr l) ->
    lift Dlocal (lookup equiv_dec (unlocalize_constants denv) v) =
    lift Dlocal (olift checkDistrAndCollect (lookup equiv_dec denv v)).

  Lemma lift_map_nnrc_to_dnnrc_base_correct {A:Set} {plug_type:Set} {plug:AlgPlug plug_type} (h:brand_relation_t) (annot:A) ctenv tenv cdenv denv v l n2 :
    wf_denv tenv denv ->
    (forall (tenv : list (var * dlocalization))
            (denv : list (var * ddata)),
        wf_denv tenv denv ->
        lift Dlocal (nnrc_core_eval h (unlocalize_constants cdenv) (unlocalize_constants denv) n2) =
        dnnrc_base_eval h cdenv denv (nnrc_to_dnnrc_base annot ctenv tenv n2)) ->
    lift_map
      (fun d1 : data =>
         olift checkLocal
               (dnnrc_base_eval h cdenv ((v, Dlocal d1) :: denv)
                          (nnrc_to_dnnrc_base annot ctenv ((v, Vlocal) :: tenv) n2))) l = lift_map (fun d1 : data => nnrc_core_eval h (unlocalize_constants cdenv) ((v, d1) :: unlocalize_constants denv) n2) l.

  Lemma assoc_lookupr_map_snd_comm {A B} (l:list (string*A)) (f:A->B) (v:string) :
    assoc_lookupr equiv_dec (map (fun x: string * A => (fst x, f (snd x))) l) v =
    lift f (assoc_lookupr equiv_dec l v).

  Lemma lookup_map_snd_comm {A B} (l:list (string*A)) (f:A->B) (v:string) :
    lookup equiv_dec (map (fun x: string * A => (fst x, f (snd x))) l) v =
    lift f (lookup equiv_dec l v).

  Lemma nnrc_to_dnnrc_base_get_constant_correct
        h
        (A plug_type:Set) (plug:AlgPlug plug_type) (annot:A) (v : var)
        (ctenv : list (var * dlocalization))
        (cdenv denv : list (var * ddata))
        (Hcdenv : wf_cdenv ctenv cdenv) :
    lift Dlocal (edot (unlocalize_constants cdenv) v) =
    dnnrc_base_eval h cdenv denv
               match assoc_lookupr equiv_dec ctenv v with
               | Some Vlocal => DNNRCGetConstant annot v
               | Some Vdistr => DNNRCCollect annot (DNNRCGetConstant annot v)
               | None => DNNRCGetConstant annot v
               end.
    
  Global Arguments dnnrc_base : clear implicits.
  Lemma nnrc_to_dnnrc_base_correct {A plug_type:Set} (annot:A) {plug:AlgPlug plug_type} h (ctenv tenv:list (var*dlocalization)) (n:nnrc) :
    forall cdenv denv:list (var*ddata),
      wf_cdenv ctenv cdenv ->
      wf_denv tenv denv ->
      lift Dlocal (nnrc_core_eval h (unlocalize_constants cdenv) (unlocalize_constants denv) n)
      = dnnrc_base_eval h cdenv denv (nnrc_to_dnnrc_base annot ctenv tenv n).

  Lemma nnrc_to_dnnrc_base_correct_lift {A plug_type:Set} (annot:A) {plug:AlgPlug plug_type} h (ctenv tenv:list (var*dlocalization)) (n:nnrc) :
    forall cdenv denv:list (var*ddata),
      wf_cdenv ctenv cdenv ->
      wf_denv tenv denv ->
      nnrc_core_eval h (unlocalize_constants cdenv) (unlocalize_constants denv) n
      = lift unlocalize_data (dnnrc_base_eval h cdenv denv (nnrc_to_dnnrc_base annot ctenv tenv n)).

  Section Top.
    Context {ftype: ForeignType.foreign_type}.

    Definition nnrc_to_dnnrc {A:Set} (annot:A) (tenv:vdbindings) (q:nnrc) :=
      @nnrc_to_dnnrc_base A dataframe annot (rec_sort tenv) nil
                     (nnrc_to_nnrc_base q).

    Definition nnrc_to_dnnrc_top (tenv:vdbindings) (q:nnrc) : dnnrc :=
      nnrc_to_dnnrc tt tenv q.

    Lemma rec_sort_unlocalize_constants_comm l :
      (rec_sort (unlocalize_constants l) = (unlocalize_constants (rec_sort l))).

    Lemma wf_cdenv_rec_sort tenv cdenv :
      wf_cdenv tenv cdenv -> wf_cdenv (rec_sort tenv) (rec_sort cdenv).

    Theorem nnrc_to_dnnrc_top_correct h (tenv:vdbindings) (q:nnrc) :
      forall cdenv:dbindings,
        wf_cdenv tenv cdenv ->
        nnrc_eval_top h q (unlocalize_constants cdenv) =
        dnnrc_eval_top h (nnrc_to_dnnrc_top tenv q) cdenv.
  End Top.
  
End NNRCtoDNNRC.