Require Import String.
Require Import RelationClasses.
Require Import EquivDec.
Require Import Streams.
Require Import List.
Require Import ListAdd.
Require Import Rbase Rtrigo Rpower Rbasic_fun.
Require Import DefinedFunctions.
Require Import Vector.
Require Import Lra Lia.

Require Import Floatish Utils.

Section GenNN.

  Context {floatish_impl:floatish}.

  Local Open Scope float.

  Definition UnitDefinedFunction := @DefinedFunction floatish_impl UnitAnn.

  Record FullNN : Type := mkNN { ldims : list nat; param_var : SubVar;
                                 f_activ : UnitDefinedFunction DTfloat; f_loss : UnitDefinedFunction DTfloat }.

  Definition mkSubVarVector (v : SubVar) (n : nat) : UnitDefinedFunction (DTVector n) :=
    DVector tt (fun i => Var (Sub v (proj1_sig i), DTfloat) tt).

  Definition mkVarVector (v : SubVar) (n : nat) : UnitDefinedFunction (DTVector n) :=
    Var (v, DTVector n) tt.

  Definition mkSubVarMatrix (v : SubVar) (n m : nat) : UnitDefinedFunction (DTMatrix n m) :=
    DMatrix tt (fun i j => Var (Sub (Sub v (proj1_sig i)) (proj1_sig j), DTfloat) tt).

  Definition mkVarMatrix (v : SubVar) (n m : nat) : UnitDefinedFunction (DTMatrix n m) :=
    Var (v, DTMatrix n m) tt.

  Definition unique_var {Ann} (df : DefinedFunction Ann DTfloat) : option var_type :=
    let fv := nodup vart_dec (df_free_variables df) in
    match fv with
    | nil => None
    | v :: nil => Some v
    | _ => None
    end.

  Definition activation (df : UnitDefinedFunction DTfloat) (vec : list (UnitDefinedFunction DTfloat)) : option (list (UnitDefinedFunction DTfloat)) :=
    match unique_var df with
    | Some (v, DTfloat) => Some (map (fun dfj => df_subst df (v, DTfloat) dfj) vec)
    | _ => None
    end.

  Definition create_activation_fun (df : UnitDefinedFunction DTfloat) : option (UnitDefinedFunction DTfloat -> UnitDefinedFunction DTfloat) :=
    match unique_var df with
    | Some (v, DTfloat) => Some (fun val => df_subst df (v, DTfloat) val)
    | _ => None
    end.

  Definition mkNN2 (n1 n2 n3 : nat) (ivar wvar : SubVar) (f_activ : UnitDefinedFunction DTfloat) (f_activ_var : SubVar) : (UnitDefinedFunction (DTVector n3)) :=
    let mat1 := mkSubVarMatrix (Sub wvar 1) n2 n1 in
    let mat2 := mkSubVarMatrix (Sub wvar 2) n3 n2 in
    let ivec := mkVarVector ivar n1 in
    let N1 := VectorApply tt f_activ_var f_activ (MatrixVectorMult  tt mat1 ivec) in
    VectorApply tt f_activ_var f_activ (MatrixVectorMult tt mat2 N1).

  Definition mkVarNN2 (n1 n2 n3 : nat) (ivar wvar : SubVar) (f_activ : UnitDefinedFunction DTfloat) (f_activ_var : SubVar) : (UnitDefinedFunction (DTVector n3) * list var_type) :=
    let mat1 := (Sub wvar 1, DTMatrix n2 n1) in
    let mat2 := (Sub wvar 2, DTMatrix n3 n2) in
    let ivec := Var (ivar, DTVector n1) tt in
    let N1 := VectorApply tt f_activ_var f_activ (MatrixVectorMult tt (Var mat1 tt) ivec) in
    (VectorApply tt f_activ_var f_activ (MatrixVectorMult tt (Var mat2 tt) N1),
        mat1 :: mat2 :: nil).

  Definition mkVarMatNN2 (n1 n2 n3 nsamp: nat) (ivar wvar : SubVar) (f_activ : UnitDefinedFunction DTfloat) (f_activ_var : SubVar) : (UnitDefinedFunction (DTMatrix n3 nsamp) * list var_type) :=
    let mat1 := (Sub wvar 1, DTMatrix n2 n1) in
    let mat2 := (Sub wvar 2, DTMatrix n3 n2) in
    let imat := Var (ivar, DTMatrix n1 nsamp) tt in
    let N1 := MatrixApply tt f_activ_var f_activ (MatrixMult tt (Var mat1 tt) imat) in
    (MatrixApply tt f_activ_var f_activ (MatrixMult tt (Var mat2 tt) N1),
     mat1 :: mat2 :: nil).

  Definition mkNN_bias_step (n1 n2 : nat) (ivec : UnitDefinedFunction (DTVector n1))
             (mat : UnitDefinedFunction (DTMatrix n2 n1))
             (bias : UnitDefinedFunction (DTVector n2))
             (f_activ_var : SubVar) (f_activ : UnitDefinedFunction DTfloat)
              : UnitDefinedFunction (DTVector n2) :=
    VectorApply tt f_activ_var f_activ (VectorPlus tt (MatrixVectorMult tt mat ivec) bias).

 Definition mkNN2_bias (n1 n2 n3 : nat) (ivar wvar : SubVar) (f_activ : UnitDefinedFunction DTfloat) (f_activ_var : SubVar) : UnitDefinedFunction (DTVector n3) :=
    let mat1 := mkSubVarMatrix (Sub wvar 1) n2 n1 in
    let b1 := mkSubVarVector (Sub wvar 1) n2 in
    let mat2 := mkSubVarMatrix (Sub wvar 2) n3 n2 in
    let b2 := mkSubVarVector (Sub wvar 2) n3 in
    let ivec := mkVarVector ivar n1 in
    let N1 := mkNN_bias_step n1 n2 ivec mat1 b1 f_activ_var f_activ in
    mkNN_bias_step n2 n3 N1 mat2 b2 f_activ_var f_activ.

 Definition mkNN2_Var_bias (n1 n2 n3 : nat) (ivar wvar bvar : SubVar) (f_activ : UnitDefinedFunction DTfloat) (f_activ_var : SubVar) : (UnitDefinedFunction (DTVector n3) * list var_type) :=
    let mat1 := (Sub wvar 1, DTMatrix n2 n1) in
    let b1 :=  (Sub bvar 1, DTVector n2) in
    let mat2 := (Sub wvar 2, DTMatrix n3 n2) in
    let b2 :=  (Sub bvar 2, DTVector n3) in
    let ivec := Var (ivar, DTVector n1) tt in
    let N1 := mkNN_bias_step n1 n2 ivec (Var mat1 tt) (Var b1 tt) f_activ_var f_activ in
    (mkNN_bias_step n2 n3 N1 (Var mat2 tt) (Var b2 tt) f_activ_var f_activ,
     mat1 :: b1 :: mat2 :: b2 :: nil).

  Definition mkNN_Mat_bias_step (n1 n2 nsamp : nat) (imat : UnitDefinedFunction (DTMatrix n1 nsamp))
             (mat : UnitDefinedFunction (DTMatrix n2 n1))
             (bias : UnitDefinedFunction (DTVector n2))
             (f_activ_var : SubVar) (f_activ : UnitDefinedFunction DTfloat)
              : UnitDefinedFunction (DTMatrix n2 nsamp) :=
    MatrixApply tt f_activ_var f_activ (MatrixVectorAdd tt (MatrixMult tt mat imat) bias).

 Definition mkNN2_Var_Mat_bias (n1 n2 n3 nsamp: nat) (ivar wvar bvar: SubVar) (f1_activ f2_activ : UnitDefinedFunction DTfloat) (f1_activ_var f2_activ_var : SubVar) : (UnitDefinedFunction (DTMatrix n3 nsamp) * list var_type)  :=
    let mat1 := (Sub wvar 1, DTMatrix n2 n1) in
    let b1 :=  (Sub bvar 1, DTVector n2) in
    let mat2 := (Sub wvar 2, DTMatrix n3 n2) in
    let b2 :=  (Sub bvar 2, DTVector n3) in
    let imat := Var (ivar, DTMatrix n1 nsamp) tt in
    let N1 := mkNN_Mat_bias_step n1 n2 nsamp imat (Var mat1 tt) (Var b1 tt) f1_activ_var f1_activ in
    (mkNN_Mat_bias_step n2 n3 nsamp N1 (Var mat2 tt) (Var b2 tt) f2_activ_var f2_activ,
     mat1 :: b1 :: mat2 :: b2 :: nil).

 Lemma vector_float_map_last_rewrite {B nvlist1 n2 v n1} :
   (DTVector (last ((@domain _ B) nvlist1) n2)) =
   (DTVector (last (domain((n2, v) :: nvlist1)) n1)).
 Proof.
   rewrite domain_cons.
   rewrite last_cons.
   reflexivity.
 Qed.

 Lemma matrix_float_map_last_rewrite {B nvlist1 n2 v n1 nsamp} :
   (DTMatrix (last ((@domain _ B) nvlist1) n2) nsamp) =
   (DTMatrix (last (domain((n2, v) :: nvlist1)) n1) nsamp).
 Proof.
   rewrite domain_cons.
   rewrite last_cons.
   reflexivity.
 Qed.

  Fixpoint mkNN_gen_0 (n1:nat) (nvlist : list (nat * SubVar))
           (ivec : (UnitDefinedFunction (DTVector n1)))
           (f_activ_var : SubVar ) (f_activ : UnitDefinedFunction DTfloat) :
    UnitDefinedFunction (DTVector (last (domain nvlist) n1))
:=
    match nvlist with
    | nil => ivec
    | cons (n2,v) nvlist1 =>
      let mat := mkSubVarMatrix v n2 n1 in
      let b := mkSubVarVector v n2 in
      let N := mkNN_bias_step n1 n2 ivec mat b f_activ_var f_activ in
      eq_rect _ UnitDefinedFunction (mkNN_gen_0 n2 nvlist1 N f_activ_var f_activ) _ vector_float_map_last_rewrite
    end.

  Fixpoint mkNN_Var_gen_0 (n1:nat) (nvlist : list (nat * SubVar))
           (ivec : (UnitDefinedFunction (DTVector n1)))
           (f_activ_var : SubVar ) (f_activ : UnitDefinedFunction DTfloat) :
    UnitDefinedFunction (DTVector (last (domain nvlist) n1))
:=
    match nvlist with
    | nil => ivec
    | cons (n2,v) nvlist1 =>
      let mat := mkVarMatrix v n2 n1 in
      let b := mkVarVector v n2 in
      let N := mkNN_bias_step n1 n2 ivec mat b f_activ_var f_activ in
      eq_rect _ UnitDefinedFunction (mkNN_Var_gen_0 n2 nvlist1 N f_activ_var f_activ) _ vector_float_map_last_rewrite
    end.

  Fixpoint mkNN_Mat_Var_gen_0 (nsamp n1:nat) (nvlist : list (nat * SubVar))
           (imat : (UnitDefinedFunction (DTMatrix n1 nsamp)))
           (f_activ_var : SubVar ) (f_activ : UnitDefinedFunction DTfloat) :
    UnitDefinedFunction (DTMatrix (last (domain nvlist) n1) nsamp)
:=
    match nvlist with
    | nil => imat
    | cons (n2,v) nvlist1 =>
      let mat := mkVarMatrix v n2 n1 in
      let b := mkVarVector v n2 in
      let N := mkNN_Mat_bias_step n1 n2 nsamp imat mat b f_activ_var f_activ in
      eq_rect _ UnitDefinedFunction (mkNN_Mat_Var_gen_0 nsamp n2 nvlist1 N f_activ_var f_activ) _ matrix_float_map_last_rewrite
    end.

  Program Definition mkNN_gen (n1:nat) (nlist : list nat) (ivar wvar f_activ_var : SubVar)
             (f_activ : UnitDefinedFunction DTfloat) :
    UnitDefinedFunction (DTVector (last nlist n1)) :=
    let vlist := map (fun i => Sub wvar i) (seq 1 (length nlist)) in
    let ivec := mkVarVector ivar n1 in
    eq_rect _ UnitDefinedFunction
            (mkNN_gen_0 n1 (combine nlist vlist) ivec f_activ_var f_activ) _ _.
  Next Obligation.
    f_equal.
    f_equal.
    rewrite combine_domain_eq; trivial.
    now rewrite map_length, seq_length.
  Qed.

  Program Definition mkNN_Var_gen (n1:nat) (nlist : list nat) (ivar wvar f_activ_var : SubVar)
             (f_activ : UnitDefinedFunction DTfloat) :
    UnitDefinedFunction (DTVector (last nlist n1)) :=
    let vlist := map (fun i => Sub wvar i) (seq 1 (length nlist)) in
    let ivec := mkVarVector ivar n1 in
    eq_rect _ UnitDefinedFunction
            (mkNN_Var_gen_0 n1 (combine nlist vlist) ivec f_activ_var f_activ) _ _.
  Next Obligation.
    f_equal.
    f_equal.
    rewrite combine_domain_eq; trivial.
    now rewrite map_length, seq_length.
  Qed.

  Program Definition mkNN_Mat_Var_gen (nsamp n1:nat) (nlist : list nat) (ivar wvar f_activ_var : SubVar)
             (f_activ : UnitDefinedFunction DTfloat) :
    UnitDefinedFunction (DTMatrix (last nlist n1) nsamp) :=
    let vlist := map (fun i => Sub wvar i) (seq 1 (length nlist)) in
    let imat := mkVarMatrix ivar n1 nsamp in
    eq_rect _ UnitDefinedFunction
            (mkNN_Mat_Var_gen_0 nsamp n1 (combine nlist vlist) imat f_activ_var f_activ) _ _.
  Next Obligation.
    f_equal.
    f_equal.
    rewrite combine_domain_eq; trivial.
    now rewrite map_length, seq_length.
  Qed.

  Definition softmax {n:nat} (NN : UnitDefinedFunction (DTVector n)) : UnitDefinedFunction (DTVector n) :=
    let expvar := Name "expvar" in
    let NNexp := VectorApply tt expvar (Exp tt (Var (expvar, DTfloat) tt)) NN in
    let NNexpscale := Divide tt (Number tt 1) (VectorSum tt NNexp) in
    VectorScalMult tt NNexpscale NNexp.

  Definition L2loss (nnvar ovar : SubVar) : UnitDefinedFunction DTfloat :=
    Square tt ( Minus tt (Var (nnvar, DTfloat) tt) (Var (ovar, DTfloat) tt) ).

  Definition L1loss (nnvar ovar : SubVar) : UnitDefinedFunction DTfloat :=
    Abs tt (Minus tt (Var (nnvar, DTfloat) tt) (Var (ovar, DTfloat) tt)).

  Definition Sigmoid (var : SubVar) : UnitDefinedFunction DTfloat :=
    Divide tt (Number tt 1)
           (Plus tt (Number tt 1) (Exp tt (Minus tt (Number tt 0) (Var (var, DTfloat) tt)))).

  Definition CrossEntropy (nnvar ovar : SubVar) : UnitDefinedFunction DTfloat :=
    let nnvar' := Var (nnvar, DTfloat) tt in
    let ovar' := Var (ovar, DTfloat) tt in
    Minus tt (Times tt (Minus tt ovar' (Number tt 1)) (Log tt (Minus tt (Number tt 1) nnvar')))
             (Times tt ovar' (Log tt nnvar')).

  Record testcases : Type := mkTest {ninput: nat; noutput: nat; ntest: nat;
                                     datavec : Vector ((Vector float ninput) * (Vector float noutput)) ntest}.

  Definition NNinstance1samp {ninput noutput : nat} (ivar : SubVar)
             (f_loss : UnitDefinedFunction DTfloat)
             (f_loss_NNvar f_loss_outvar : SubVar)
             (NN : UnitDefinedFunction (DTVector noutput)) (σ:df_env)
             (data: (Vector float ninput) * (Vector float noutput))
             : df_env * (UnitDefinedFunction DTfloat) :=
    let ipair := mk_env_entry (ivar, DTVector ninput) (fst data) in
    (cons ipair σ, VLossfun tt f_loss_NNvar f_loss_outvar f_loss NN (snd data)).

  Definition NNinstancebatch {ninput nsamp noutput : nat} (ivar : SubVar)
             (f_loss : UnitDefinedFunction DTfloat)
             (f_loss_NNvar f_loss_outvar : SubVar)
             (NN : UnitDefinedFunction (DTMatrix noutput nsamp)) (σ:df_env)
             (data: (Matrix float ninput nsamp) * (Matrix float noutput nsamp))
             : df_env * (UnitDefinedFunction DTfloat) :=
    let ipair := mk_env_entry (ivar, DTMatrix ninput nsamp) (fst data) in
    (cons ipair σ, MLossfun tt f_loss_NNvar f_loss_outvar f_loss NN (snd data)).

  Definition EvalNNinstance1samp {ninput noutput : nat} (ivar : SubVar)
             (f_loss : UnitDefinedFunction DTfloat)
             (f_loss_NNvar f_loss_outvar : SubVar)
             (NN : UnitDefinedFunction (DTVector noutput)) (σ:df_env)
             (data: (Vector float ninput) * (Vector float noutput))
             : option float :=
    let ipair := mk_env_entry (ivar, DTVector ninput) (fst data) in
    df_eval (cons ipair σ) (VLossfun tt f_loss_NNvar f_loss_outvar f_loss NN (snd data)).

  Definition EvalNNinstancebatch {ninput nsamp noutput : nat} (ivar : SubVar)
             (f_loss : UnitDefinedFunction DTfloat)
             (f_loss_NNvar f_loss_outvar : SubVar)
             (NN : UnitDefinedFunction (DTMatrix noutput nsamp)) (σ:df_env)
             (data: (Matrix float ninput nsamp) * (Matrix float noutput nsamp))
             : option float :=
    let ipair := mk_env_entry (ivar, DTMatrix ninput nsamp) (fst data) in
    df_eval (cons ipair σ) (MLossfun tt f_loss_NNvar f_loss_outvar f_loss NN (snd data)).

  (*
  Lemma NNinstance_unique_var (n1 n2 n3 : nat) (ivar : SubVar) (f_loss : DefinedFunction DTfloat)
        (NN2 : DefinedFunction (DTVector n3)) (inputs : (list float))
        (outputs : Vector float n3) (v:SubVar) :
    unique_var f_loss = Some v ->
    NNinstance n1 n2 n3 ivar f_loss NN2 inputs outputs =
    Some (
        let ipairs := (list_prod (map (fun n => (Sub ivar n)) (seq 1 n1))
                                 (map Number inputs)) in
        let losses := VectorMinus (df_subst_list NN2 ipairs)
                                  (DVector (vmap Number outputs)) in
        (VectorSum (VectorApply v f_loss losses))
      ).
  Proof.
    unfold NNinstance.
    intros.
    rewrite H.
    reflexivity.
  Qed.

  Lemma NNinstance_None (n1 n2 n3 : nat) (ivar : SubVar) (f_loss : DefinedFunction DTfloat)
        (NN2 : DefinedFunction (DTVector n3)) (inputs : (list float))
        (outputs : Vector float n3) :
    unique_var f_loss = None ->
    NNinstance n1 n2 n3 ivar f_loss NN2 inputs outputs = None.
  Proof.
    unfold NNinstance.
    intros.
    now rewrite H.
  Qed.
   *)

(*  Local Existing Instance floatish_interval. *)

  Definition lookup_list (σ:df_env) (lvar : list SubVar) : option (list float) :=
    listo_to_olist (map (fun v => (vartlookup σ (v, DTfloat)):option float) lvar).

  Definition combine_with {A:Type} {B:Type} {C:Type} (f: A -> B -> C )
             (lA : list A) (lB : list B) : list C :=
    map (fun '(a, b) => f a b) (combine lA lB).

  Definition combine3_with {A:Type} {B:Type} {C:Type} {D:Type} (f: A -> B -> C -> D)
             (lA : list A) (lB : list B) (lC : list C) : list D :=
    map (fun '(a, bc) => f a (fst bc) (snd bc)) (combine lA (combine lB lC)).

  Fixpoint streamtake (n : nat) {A : Type} (st : Stream A) : (list A) * (Stream A) :=
    match n with
    | 0 => (nil, st)
    | S n' => let rst := streamtake n' (Streams.tl st) in
              ((Streams.hd st)::(fst rst), snd rst)
    end.

  Lemma streamtake_n (n : nat) (A : Type) (st : Stream A) :
    length (fst (streamtake n st)) = n.
  Proof.
    generalize st.
    induction n.
    reflexivity.
    intros.
    simpl.
    f_equal.
    specialize IHn with (st := Streams.tl st0).
    apply IHn.
  Qed.

  Fixpoint env_update_first (l:df_env) (an:env_entry_type) : df_env
    := match l with
       | nil => nil
       | fv::os => if (projT1 an) == (projT1 fv) then an::os
                   else fv::(env_update_first os an)
       end.

  Definition env_update_list (l up:df_env) : df_env
    := fold_left (env_update_first) up l.

  Definition optimize_step
             (step : nat) (df : UnitDefinedFunction DTfloat) (σ:df_env) (lvar : list SubVar)
             (noise_st : Stream float) : (option df_env)*(Stream float) :=
    let lvart:list var_type := (map (fun v => (v, DTfloat)) lvar) in
    let ogradvec := df_eval_gradient σ df lvart in
    let alpha := 1 / (FfromZ (Z.of_nat (S step))) in
    let '(lnoise, nst) := streamtake (length lvar) noise_st in
    let olvals := lookup_list σ lvar in
    (match (ogradvec, olvals) with
     | (Some gradvec, Some lvals) =>
       Some (env_update_list
               σ
               (map (fun '(v,e) => mk_env_entry (v, DTfloat) (e:float))
                    (combine lvar (combine3_with
                                     (fun val grad noise => val - alpha*(grad + noise))
                                     lvals gradvec lnoise))))
     | (_, _) => None
     end, nst).

  Fixpoint get_noise_vector (n: nat) (noise_st: Stream float) :
    (Vector float n) * (Stream float) :=
    match n with
    | 0 => (vnil, noise_st)
    | S n' =>
      let noise := Streams.hd noise_st in
      let nst := Streams.tl noise_st in
      let '(vec, nst') := get_noise_vector n' nst in
      (vcons noise vec, nst')
    end.

  Fixpoint get_noise_matrix (n m: nat) (noise_st: Stream float) :
    (Matrix float n m) * (Stream float) :=
    let '(vec, nst) := get_noise_vector m noise_st in
    match n with
    | 0 => (fun i => vec, nst)
    | S n' =>
      let '(mat, nst') := get_noise_matrix n' m nst in
      (vcons vec mat, nst')
    end.

  Definition get_noise (t:definition_function_types) (noise_st:Stream float) :
    (definition_function_types_interp t) * (Stream float) :=
    match t with
    | DTfloat => (Streams.hd noise_st, Streams.tl noise_st)
    | DTVector n => get_noise_vector n noise_st
    | DTMatrix m n => get_noise_matrix m n noise_st
    end.

  Program Definition update_val_gradenv (grad_env:df_env) (x:var_type) (alpha:float) :
    definition_function_types_interp (snd x) ->
    definition_function_types_interp (snd x) ->
    definition_function_types_interp (snd x) :=
    (match snd x as y return snd x = y ->
                          definition_function_types_interp y ->
                          definition_function_types_interp y ->
                          definition_function_types_interp y with
       | DTfloat =>  fun pf val noise =>
                       match vartlookup grad_env x with
                       | Some grad => val - alpha * (coerce _ grad + noise)
                       | _ => val
                       end
       | DTVector n => fun pf val noise =>
                         match vartlookup grad_env x with
                         | Some grad =>
                           fun i =>
                             (val i) - alpha * (((coerce _ grad) i) + (noise i))
                         | _ => val
                         end
       | DTMatrix m n => fun pf val noise =>
                           match vartlookup grad_env x with
                           | Some grad =>
                             fun i j =>
                               (val i j) - alpha * (((coerce _ grad) i j) + (noise i j))
                         | _ => val
                         end
     end) (eq_refl _)
  .
    Next Obligation.
      rewrite pf; reflexivity.
    Qed.
    Next Obligation.
      rewrite pf; reflexivity.
    Qed.
    Next Obligation.
      rewrite pf; reflexivity.
    Qed.

  Definition update_entry (entry: env_entry_type) (grad_env:df_env) (alpha:float)
             (noise_st : Stream float) : (env_entry_type*Stream float) :=
    let x := projT1 entry in
    let val := projT2 entry in
    let '(noise, nst) := get_noise (snd x) noise_st in
    (mk_env_entry x (update_val_gradenv grad_env x alpha val noise), nst).

  Fixpoint list_arg_iter {A B} (f: A -> B -> A * B)
             (l:list A) (b: B) : (list A)*B :=
      match l with
      | nil => (l, b)
      | a :: l' =>
        let '(na, b') := f a b in
        let '(nl, nb) := list_arg_iter f l' b' in
        (na::nl, nb)
      end.

  Definition harmonic_lr (step : nat) : float := 1 / (FfromZ (Z.of_nat (S step))).

  Definition optimize_step_backprop
             (step : nat) (df : UnitDefinedFunction DTfloat) (σ:df_env) (lr : nat -> float)
             (noise_st : Stream float) (dvars : list var_type)
    : (option df_env)*(Stream float) :=
    match df_eval_backprop_deriv σ df (gradenv_init dvars) 1 with
    | Some gradenv =>
      let alpha := lr step in
      let '(env, nst) := list_arg_iter (fun a b => update_entry a gradenv alpha b) σ noise_st in
      (Some env, nst)
    | _ => (None, noise_st)
    end.

  Definition optimize_step_tree_backprop
             (step : nat) (df : UnitDefinedFunction DTfloat) (σ:df_env) (lr : nat -> float)
             (noise_st : Stream float) (dvars : list var_type)
    : (option df_env)*(Stream float) :=
    match df_eval_tree σ df with
    | Some df_tree =>
      match df_eval_tree_backprop_deriv (* σ not-needed *) nil df_tree (gradenv_init dvars) 1 with
      | Some gradenv =>
        let alpha := lr step in
        let '(env, nst) := list_arg_iter (fun a b => update_entry a gradenv alpha b) σ noise_st in
        (Some env, nst)
      | _ => (None, noise_st)
      end
    | _ => (None, noise_st)
    end.

  Fixpoint optimize_steps
           (start count:nat) (df : UnitDefinedFunction DTfloat) (σ:df_env) (lvar : list SubVar)
           (noise_st : Stream float) : (option df_env)*(Stream float) :=
    match count with
    | 0 => (Some σ, noise_st)
    | S n =>
      match optimize_step start df σ lvar noise_st with
      | (Some σ', noise_st') => optimize_steps (S start) n df σ' lvar noise_st'
      | (None, noise_st') => (None, noise_st')
      end
    end.

  Fixpoint optimize_steps_backprop
           (start count:nat) (df : UnitDefinedFunction DTfloat) (σ:df_env) (lr : nat -> float)
           (noise_st : Stream float) (dvars : list var_type)
    : (option df_env)*(Stream float) :=
    match count with
    | 0 => (Some σ, noise_st)
    | S n =>
      match optimize_step_backprop start df σ lr noise_st dvars with
      | (Some σ', noise_st') => optimize_steps_backprop (S start) n df σ' lr noise_st' dvars
      | (None, noise_st') => (None, noise_st')
      end
    end.

  Fixpoint optimize_steps_tree_backprop
           (start count:nat) (df : UnitDefinedFunction DTfloat) (σ:df_env) (lr : nat -> float)
           (noise_st : Stream float) (dvars : list var_type)
    : (option df_env)*(Stream float) :=
    match count with
    | 0 => (Some σ, noise_st)
    | S n =>
      match optimize_step_tree_backprop start df σ lr noise_st dvars with
      | (Some σ', noise_st') => optimize_steps_tree_backprop (S start) n df σ' lr noise_st' dvars
      | (None, noise_st') => (None, noise_st')
      end
    end.

Definition Fmax (a b:float) : float :=
  if Fgt a b then a else b.

Definition Fmin (a b:float) : float :=
  if Flt a b then a else b.

Definition vmax {n : nat} (vec : Vector float n)
          := vector_fold_right1 Fmax (FfromZ 0) id vec.

Definition vmin {n : nat} (vec : Vector float n)
          := vector_fold_right1 Fmin (FfromZ 0) id vec.

Definition nrows {A} (l : list (list A)) := List.length l.

Definition ncols {A} (l : list (list A)) :=
      match l with
      | nil => 0%nat
      | r :: _ => List.length r
      end.

Definition normalizeIntData (l:list (list Z)) : Matrix float (nrows l) (ncols l) :=
  let mat : Matrix float (nrows l) (ncols l) :=
      fun i j => FfromZ (List.nth (proj1_sig j) (List.nth (proj1_sig i) l nil) (0)%Z) in
  let tmat := transpose mat in
  let maxes := vmap vmax tmat in
  let mins := vmap vmin tmat in
  fun i j => ((mat i j) - (mins j))/((maxes j)- (mins j)).

Program Definition splitLastCol {nsamp ncols : nat}  (data : Matrix float nsamp ncols)
        (pf : (ncols > 0)%nat)
        : (Matrix float (ncols-1) nsamp) * (Matrix float 1 nsamp) :=
  (fun i j => data j i, fun i j => data j (i + ncols-1)%nat).
Next Obligation.
  lia.
Defined.
Next Obligation.
  lia.
Defined.

Definition init_env2 (dim1 dim2 dim3 : nat) (w b : string)
        (ranm1 : Matrix float dim2 dim1)
        (ranm2 : Matrix float dim3 dim2) : df_env :=
  let wvar := Name w in
  let bvar := Name b in
  let wvar1 := (Sub wvar 1, DTMatrix dim2 dim1) in
  let wvar2 := (Sub wvar 2, DTMatrix dim3 dim2) in
  let bvar1 := (Sub bvar 1, DTVector dim2) in
  let bvar2 := (Sub bvar 2, DTVector dim3) in
  (*  (mk_env_entry wvar1 ranm1) :: (mk_env_entry wvar2 ranm2) :: *)
  (mk_env_entry wvar1 (fun i j => (ranm1 i j) / (Fsqrt (FfromZ (Z.of_nat dim1))))) ::
  (mk_env_entry wvar2 (fun i j => (ranm2 i j) / (Fsqrt (FfromZ (Z.of_nat dim2))))) ::
  (mk_env_entry bvar1 (ConstVector dim2 0)) ::
  (mk_env_entry bvar2 (ConstVector dim3 0)) :: nil.

 Definition mkNN_wisconsin (nsamp:nat) (ivar wvar bvar f1v f2v : SubVar) : (UnitDefinedFunction (DTMatrix 1 nsamp) * list var_type) :=
   let f1_activ := Max tt (Var (f1v,DTfloat) tt) (Number tt 0) in
   let f2_activ := Sigmoid f2v in
   mkNN2_Var_Mat_bias 9 15 1 nsamp ivar wvar bvar f1_activ f2_activ f1v f2v.

 Program Definition eval_wisconsin_batch (nsamp : nat)
            (σ:df_env)
            (normaldata: Matrix float nsamp 10)   : option float :=
   let ivar := (Name "i") in
   let flnnv := (Name "NNv") in
   let outnnv := (Name "outnnv") in
   let '(NN, dvars) := mkNN_wisconsin nsamp ivar (Name "w") (Name "b") (Name "f1v") (Name "f2v") in
   @EvalNNinstancebatch 9 nsamp 1 ivar (CrossEntropy flnnv outnnv) flnnv outnnv
             NN σ (splitLastCol normaldata _).
 Next Obligation.
   lia.
 Qed.

Program Definition wisconsin_instance_batch (nsamp : nat)
            (σ:df_env)
            (normaldata: Matrix float nsamp 10)
             : (df_env * (UnitDefinedFunction DTfloat)) * (list var_type) :=
   let ivar := (Name "i") in
   let flnnv := (Name "NNv") in
   let outnnv := (Name "outnnv") in
   let '(NN, dvars) := mkNN_wisconsin nsamp ivar (Name "w") (Name "b") (Name "f1v") (Name "f2v") in
   (@NNinstancebatch 9 nsamp 1 ivar (CrossEntropy flnnv outnnv) flnnv outnnv
             NN
             σ
             (splitLastCol normaldata _), dvars).
Next Obligation.
  lia.
Qed.

CoFixpoint zeronoise : Stream float := Cons 0 zeronoise.

(*
https://towardsdatascience.com/predict-malignancy-in-breast-cancer-tumors-with-your-own-neural-network-and-the-wisconsin-dataset-76271a05e941
*)

Definition wisconsin_test (nsamp count : nat)
            (σ:df_env)
            (normaldata: Matrix float nsamp 10): list float :=
  let '(nninst, dvars) := wisconsin_instance_batch nsamp σ normaldata in
  let lr := fun _ => 1 / (FfromZ 100) in
  let onenv := fst (optimize_steps_tree_backprop 0 count (snd nninst) (fst nninst) lr
                                                zeronoise dvars) in
  match onenv with
  | Some nenv => match df_eval nenv (snd nninst) with
                 | Some val => val :: nil
                 | _ => nil
                 end
  | _ => nil
  end.

Definition wisconsin_test_env (nsamp count : nat)
            (σ:df_env)
            (normaldata: Matrix float nsamp 10): df_env :=
  let '(nninst, dvars) := wisconsin_instance_batch nsamp σ normaldata in
  let lr := fun _ => 1 / (FfromZ 100) in
  let onenv := fst (optimize_steps_tree_backprop 0 count (snd nninst) (fst nninst) lr
                                                zeronoise dvars) in
  match onenv with
  | Some nenv => nenv
  | _ => nil
  end.

Example xvar:var_type := (Name "x", DTfloat).
Example xfun:UnitDefinedFunction DTfloat := Var xvar tt.
Example tquad:UnitDefinedFunction DTfloat := Times tt xfun xfun.
Example quad:UnitDefinedFunction DTfloat := Minus tt (Times tt xfun xfun) (Number tt 1).
Example squad:UnitDefinedFunction DTfloat := Minus tt (Square tt xfun) (Number tt 1).
Example env : df_env := (mk_env_entry xvar (FfromZ 5))::nil.

Example gradenv := match df_eval_backprop_deriv env quad (gradenv_init (xvar :: nil)) 1 with
                   | Some gradenv => gradenv
                   | _ => nil
                   end.

Example gradenv_tree :=
  match df_eval_tree env quad with
  | Some df_tree =>
    match df_eval_tree_backprop_deriv nil df_tree (gradenv_init (xvar :: nil)) 1 with
    | Some gradenv => gradenv
    | _ => nil
    end
  | _ => nil
  end.

Example wisconsin_gradenv (nsamp : nat)
            (σ:df_env)
            (normaldata: Matrix float nsamp 10) : df_env :=
  let '((env,nn),dvars) := wisconsin_instance_batch nsamp σ normaldata in
  match df_eval_backprop_deriv env nn (gradenv_init dvars) 1 with
  | Some gradenv => gradenv
  | _ => nil
  end.

Example wisconsin_gradenv_tree (nsamp : nat)
            (σ:df_env)
            (normaldata: Matrix float nsamp 10) : df_env :=
  let '((env,nn),dvars) := wisconsin_instance_batch nsamp σ normaldata in
  match df_eval_tree env nn with
  | Some df_tree =>
    match df_eval_tree_backprop_deriv nil df_tree (gradenv_init dvars) 1 with
    | Some gradenv => gradenv
    | _ => nil
    end
  | _ => nil
  end.

 Definition mkperceptron (n1 n2 : nat) (ivar wvar bvar : SubVar) (f_activ : UnitDefinedFunction DTfloat) (f_activ_var : SubVar) : (UnitDefinedFunction (DTVector n2) * list var_type) :=
    let mat1 := (wvar, DTMatrix n2 n1) in
    let b1 :=  (bvar, DTVector n2) in
    let ivec := Var (ivar, DTVector n1) tt in
    let N1 := mkNN_bias_step n1 n2 ivec (Var mat1 tt) (Var b1 tt) f_activ_var f_activ in
    (N1,  mat1 :: b1 :: nil).

 Definition mkNN_test1 :=
   let ivar := Name "i" in
   let wvar := Name "w" in
   let bvar := Name "b" in
   let f1v := Name "f1" in
   let f1_activ := Max tt (Var (f1v,DTfloat) tt) (Number tt 0) in
   let '(nn, dvars) := mkperceptron 2 2 ivar wvar bvar f1_activ f1v in
   let winit := mk_env_entry (wvar, DTMatrix 2 2) (ConstMatrix 2 2 1) in
   let binit := mk_env_entry (bvar, DTVector 2) (ConstVector 2 1) in
   let datain := ConstVector 2 1 in
   let env := (mk_env_entry (ivar, DTVector 2) (ConstVector 2 1)):: winit :: binit :: nil in
   (env, nn).

 Definition mkNN_test :=
   let ivar := Name "i" in
   let wvar := Name "w" in
   let bvar := Name "b" in
   let f1v := Name "f1" in
(*   let f1_activ := Max tt (Var (f1v,DTfloat) tt) (Number tt 0) in *)
   let f1_activ := Sigmoid f1v in
   let '(nn, dvars) := mkperceptron 2 2 ivar wvar bvar f1_activ f1v in
   let datain := ConstVector 2 1 in
   let dataout := ConstVector 2 0 in
   let loss_nnvar := Name "lv" in
   let loss_ovar := Name "ov" in
   let floss := CrossEntropy loss_nnvar loss_ovar in
   let winit := mk_env_entry (wvar, DTMatrix 2 2) (ConstMatrix 2 2 (FfromZ 2)) in
   let binit := mk_env_entry (bvar, DTVector 2) (ConstVector 2 (FfromZ 3)) in
   let env := winit :: binit :: nil in
   (NNinstance1samp ivar floss loss_nnvar loss_ovar nn env (datain, dataout), dvars).

Definition NN_test (count : nat) : list float :=
  let '(nninst, dvars) := mkNN_test in
  let onenv := fst (optimize_steps_tree_backprop 0 count (snd nninst) (fst nninst) harmonic_lr
                                                zeronoise dvars) in
  match onenv with
  | Some nenv => match df_eval nenv (snd nninst) with
                 | Some val => val :: nil
                 | _ => nil
                 end
  | _ => nil
  end.

Example NN_test_gradenv : df_env :=
  let '((env, nn), dvars) := mkNN_test in
  match df_eval_backprop_deriv env nn (gradenv_init dvars) 1 with
  | Some gradenv => gradenv
  | _ => nil
  end.

Example NN_test_gradenv_tree :df_env :=
  let '((env, nn), dvars) := mkNN_test in
  match df_eval_tree env nn with
  | Some df_tree =>
    match df_eval_tree_backprop_deriv nil df_tree (gradenv_init dvars) 1 with
    | Some gradenv => gradenv
    | _ => nil
    end
  | _ => nil
  end.

Definition NN_test_env (count : nat) : df_env :=
  let '(nninst, dvars) := mkNN_test in
  let onenv := fst (optimize_steps_tree_backprop 0 count (snd nninst) (fst nninst)
                                  harmonic_lr zeronoise dvars) in
  match onenv with
  | Some nenv => nenv
  | _ => nil
  end.

Definition NN_test_NN  :=
  let '((env,nn), dvars) := mkNN_test in
  nn.

Definition NN_test_val : list float :=
  let '((env,nn), dvars) := mkNN_test in
  match df_eval env nn with
  | Some val => val :: nil
  | _ => nil
  end.

Definition test_update_val_gradenv
  := update_val_gradenv gradenv xvar 1 (FfromZ 5) 0.

Definition test_update : df_env :=
  (fst (update_entry (mk_env_entry xvar (FfromZ 5)) gradenv 1 zeronoise))::nil.

Definition test_optimize_step_backprop
             (step : nat) (df : UnitDefinedFunction DTfloat) (σ:df_env)
             (noise_st : Stream float) (dvars : list var_type) : df_env :=
  match fst (optimize_step_backprop step df σ harmonic_lr noise_st dvars) with
  | Some env => env
  | _ => nil
  end.

Definition test_optimize_step_tree_backprop
             (step : nat) (df : UnitDefinedFunction DTfloat) (σ:df_env)
             (noise_st : Stream float) (dvars : list var_type) : df_env :=
  match fst (optimize_step_tree_backprop step df σ harmonic_lr noise_st dvars) with
  | Some env => env
  | _ => nil
  end.

Example testopt := test_optimize_step_backprop 0 quad env zeronoise (xvar :: nil).
Example testreeopt := test_optimize_step_tree_backprop 0 quad env zeronoise (xvar :: nil).

Example opt := fst (optimize_steps 0 2 quad env ((fst xvar) :: nil) zeronoise).
Example opt2 := fst (optimize_steps_tree_backprop 0 2 quad env harmonic_lr zeronoise (xvar :: nil)).

Example val := 1+1.

End GenNN.

(*
Local Instance floatish_interval : floatish := floatish_interval_gen 53.
Compute df_eval NN_test_env NN_test_NN.
*)


Require Import String.
Require Import RelationClasses.
Require Import EquivDec.
Require Import Streams.
Require Import List.
Require Import ListAdd.
Require Import Rbase Rtrigo Rpower Rbasic_fun.
Require Import DefinedFunctions.
Require Import Lra Omega.

Require Import Floatish Utils.

Section GenNN.

  Context {floatish_impl:floatish}.

  Local Open Scope float.

  Definition UnitDefinedFunction := @DefinedFunction floatish_impl UnitAnn.

  Record FullNN : Type := mkNN { ldims : list nat; param_var : SubVar;
                                 f_activ : UnitDefinedFunction DTfloat; f_loss : UnitDefinedFunction DTfloat }.

  Definition mkSubVarVector (v : SubVar) (n : nat) : UnitDefinedFunction (DTVector n) :=
    DVector tt (fun i => Var (Sub v (proj1_sig i), DTfloat) tt).

  Definition mkVarVector (v : SubVar) (n : nat) : UnitDefinedFunction (DTVector n) :=
    Var (v, DTVector n) tt.

  Definition mkSubVarMatrix (v : SubVar) (n m : nat) : UnitDefinedFunction (DTMatrix n m) :=
    DMatrix tt (fun i j => Var (Sub (Sub v (proj1_sig i)) (proj1_sig j), DTfloat) tt).

  Definition mkVarMatrix (v : SubVar) (n m : nat) : UnitDefinedFunction (DTMatrix n m) :=
    Var (v, DTMatrix n m) tt.

  Definition unique_var {Ann} (df : DefinedFunction Ann DTfloat) : option var_type :=
    let fv := nodup vart_dec (df_free_variables df) in
    match fv with
    | nil => None
    | v :: nil => Some v
    | _ => None
    end.

  Definition activation (df : UnitDefinedFunction DTfloat) (vec : list (UnitDefinedFunction DTfloat)) : option (list (UnitDefinedFunction DTfloat)) :=
    match unique_var df with
    | Some (v, DTfloat) => Some (map (fun dfj => df_subst df (v, DTfloat) dfj) vec)
    | _ => None
    end.

  Definition create_activation_fun (df : UnitDefinedFunction DTfloat) : option (UnitDefinedFunction DTfloat -> UnitDefinedFunction DTfloat) :=
    match unique_var df with
    | Some (v, DTfloat) => Some (fun val => df_subst df (v, DTfloat) val)
    | _ => None
    end.

  Definition mkNN2 (n1 n2 n3 : nat) (ivar wvar : SubVar) (f_activ : UnitDefinedFunction DTfloat) (f_activ_var : SubVar) : (UnitDefinedFunction (DTVector n3)) :=
    let mat1 := mkSubVarMatrix (Sub wvar 1) n2 n1 in
    let mat2 := mkSubVarMatrix (Sub wvar 2) n3 n2 in
    let ivec := mkVarVector ivar n1 in
    let N1 := VectorApply tt f_activ_var f_activ (MatrixVectorMult  tt mat1 ivec) in
    VectorApply tt f_activ_var f_activ (MatrixVectorMult tt mat2 N1).

  Definition mkVarNN2 (n1 n2 n3 : nat) (ivar wvar : SubVar) (f_activ : UnitDefinedFunction DTfloat) (f_activ_var : SubVar) : (UnitDefinedFunction (DTVector n3) * list var_type) :=
    let mat1 := (Sub wvar 1, DTMatrix n2 n1) in
    let mat2 := (Sub wvar 2, DTMatrix n3 n2) in
    let ivec := Var (ivar, DTVector n1) tt in
    let N1 := VectorApply tt f_activ_var f_activ (MatrixVectorMult tt (Var mat1 tt) ivec) in
    (VectorApply tt f_activ_var f_activ (MatrixVectorMult tt (Var mat2 tt) N1),
        mat1 :: mat2 :: nil).

  Definition mkVarMatNN2 (n1 n2 n3 nsamp: nat) (ivar wvar : SubVar) (f_activ : UnitDefinedFunction DTfloat) (f_activ_var : SubVar) : (UnitDefinedFunction (DTMatrix n3 nsamp) * list var_type) :=
    let mat1 := (Sub wvar 1, DTMatrix n2 n1) in
    let mat2 := (Sub wvar 2, DTMatrix n3 n2) in
    let imat := Var (ivar, DTMatrix n1 nsamp) tt in
    let N1 := MatrixApply tt f_activ_var f_activ (MatrixMult tt (Var mat1 tt) imat) in
    (MatrixApply tt f_activ_var f_activ (MatrixMult tt (Var mat2 tt) N1),
     mat1 :: mat2 :: nil).

  Definition mkNN_bias_step (n1 n2 : nat) (ivec : UnitDefinedFunction (DTVector n1))
             (mat : UnitDefinedFunction (DTMatrix n2 n1))
             (bias : UnitDefinedFunction (DTVector n2))
             (f_activ_var : SubVar) (f_activ : UnitDefinedFunction DTfloat)
              : UnitDefinedFunction (DTVector n2) :=
    VectorApply tt f_activ_var f_activ (VectorPlus tt (MatrixVectorMult tt mat ivec) bias).

 Definition mkNN2_bias (n1 n2 n3 : nat) (ivar wvar : SubVar) (f_activ : UnitDefinedFunction DTfloat) (f_activ_var : SubVar) : UnitDefinedFunction (DTVector n3) :=
    let mat1 := mkSubVarMatrix (Sub wvar 1) n2 n1 in
    let b1 := mkSubVarVector (Sub wvar 1) n2 in
    let mat2 := mkSubVarMatrix (Sub wvar 2) n3 n2 in
    let b2 := mkSubVarVector (Sub wvar 2) n3 in
    let ivec := mkVarVector ivar n1 in
    let N1 := mkNN_bias_step n1 n2 ivec mat1 b1 f_activ_var f_activ in
    mkNN_bias_step n2 n3 N1 mat2 b2 f_activ_var f_activ.

 Definition mkNN2_Var_bias (n1 n2 n3 : nat) (ivar wvar bvar : SubVar) (f_activ : UnitDefinedFunction DTfloat) (f_activ_var : SubVar) : (UnitDefinedFunction (DTVector n3) * list var_type) :=
    let mat1 := (Sub wvar 1, DTMatrix n2 n1) in
    let b1 :=  (Sub bvar 1, DTVector n2) in
    let mat2 := (Sub wvar 2, DTMatrix n3 n2) in
    let b2 :=  (Sub bvar 2, DTVector n3) in
    let ivec := Var (ivar, DTVector n1) tt in
    let N1 := mkNN_bias_step n1 n2 ivec (Var mat1 tt) (Var b1 tt) f_activ_var f_activ in
    (mkNN_bias_step n2 n3 N1 (Var mat2 tt) (Var b2 tt) f_activ_var f_activ,
     mat1 :: b1 :: mat2 :: b2 :: nil).

  Definition mkNN_Mat_bias_step (n1 n2 nsamp : nat) (imat : UnitDefinedFunction (DTMatrix n1 nsamp))
             (mat : UnitDefinedFunction (DTMatrix n2 n1))
             (bias : UnitDefinedFunction (DTVector n2))
             (f_activ_var : SubVar) (f_activ : UnitDefinedFunction DTfloat)
              : UnitDefinedFunction (DTMatrix n2 nsamp) :=
    MatrixApply tt f_activ_var f_activ (MatrixVectorAdd tt (MatrixMult tt mat imat) bias).

 Definition mkNN2_Var_Mat_bias (n1 n2 n3 nsamp: nat) (ivar wvar bvar: SubVar) (f1_activ f2_activ : UnitDefinedFunction DTfloat) (f1_activ_var f2_activ_var : SubVar) : (UnitDefinedFunction (DTMatrix n3 nsamp) * list var_type)  :=
    let mat1 := (Sub wvar 1, DTMatrix n2 n1) in
    let b1 :=  (Sub bvar 1, DTVector n2) in
    let mat2 := (Sub wvar 2, DTMatrix n3 n2) in
    let b2 :=  (Sub bvar 2, DTVector n3) in
    let imat := Var (ivar, DTMatrix n1 nsamp) tt in
    let N1 := mkNN_Mat_bias_step n1 n2 nsamp imat (Var mat1 tt) (Var b1 tt) f1_activ_var f1_activ in
    (mkNN_Mat_bias_step n2 n3 nsamp N1 (Var mat2 tt) (Var b2 tt) f2_activ_var f2_activ,
     mat1 :: b1 :: mat2 :: b2 :: nil).

 Lemma vector_float_map_last_rewrite {B nvlist1 n2 v n1} :
   (DTVector (last ((@domain _ B) nvlist1) n2)) =
   (DTVector (last (domain((n2, v) :: nvlist1)) n1)).
 Proof.
   rewrite domain_cons.
   rewrite last_cons.
   reflexivity.
 Qed.

 Lemma matrix_float_map_last_rewrite {B nvlist1 n2 v n1 nsamp} :
   (DTMatrix (last ((@domain _ B) nvlist1) n2) nsamp) =
   (DTMatrix (last (domain((n2, v) :: nvlist1)) n1) nsamp).
 Proof.
   rewrite domain_cons.
   rewrite last_cons.
   reflexivity.
 Qed.

  Fixpoint mkNN_gen_0 (n1:nat) (nvlist : list (nat * SubVar))
           (ivec : (UnitDefinedFunction (DTVector n1)))
           (f_activ_var : SubVar ) (f_activ : UnitDefinedFunction DTfloat) :
    UnitDefinedFunction (DTVector (last (domain nvlist) n1))
:=
    match nvlist with
    | nil => ivec
    | cons (n2,v) nvlist1 =>
      let mat := mkSubVarMatrix v n2 n1 in
      let b := mkSubVarVector v n2 in
      let N := mkNN_bias_step n1 n2 ivec mat b f_activ_var f_activ in
      eq_rect _ UnitDefinedFunction (mkNN_gen_0 n2 nvlist1 N f_activ_var f_activ) _ vector_float_map_last_rewrite
    end.

  Fixpoint mkNN_Var_gen_0 (n1:nat) (nvlist : list (nat * SubVar))
           (ivec : (UnitDefinedFunction (DTVector n1)))
           (f_activ_var : SubVar ) (f_activ : UnitDefinedFunction DTfloat) :
    UnitDefinedFunction (DTVector (last (domain nvlist) n1))
:=
    match nvlist with
    | nil => ivec
    | cons (n2,v) nvlist1 =>
      let mat := mkVarMatrix v n2 n1 in
      let b := mkVarVector v n2 in
      let N := mkNN_bias_step n1 n2 ivec mat b f_activ_var f_activ in
      eq_rect _ UnitDefinedFunction (mkNN_Var_gen_0 n2 nvlist1 N f_activ_var f_activ) _ vector_float_map_last_rewrite
    end.

  Fixpoint mkNN_Mat_Var_gen_0 (nsamp n1:nat) (nvlist : list (nat * SubVar))
           (imat : (UnitDefinedFunction (DTMatrix n1 nsamp)))
           (f_activ_var : SubVar ) (f_activ : UnitDefinedFunction DTfloat) :
    UnitDefinedFunction (DTMatrix (last (domain nvlist) n1) nsamp)
:=
    match nvlist with
    | nil => imat
    | cons (n2,v) nvlist1 =>
      let mat := mkVarMatrix v n2 n1 in
      let b := mkVarVector v n2 in
      let N := mkNN_Mat_bias_step n1 n2 nsamp imat mat b f_activ_var f_activ in
      eq_rect _ UnitDefinedFunction (mkNN_Mat_Var_gen_0 nsamp n2 nvlist1 N f_activ_var f_activ) _ matrix_float_map_last_rewrite
    end.

  Program Definition mkNN_gen (n1:nat) (nlist : list nat) (ivar wvar f_activ_var : SubVar)
             (f_activ : UnitDefinedFunction DTfloat) :
    UnitDefinedFunction (DTVector (last nlist n1)) :=
    let vlist := map (fun i => Sub wvar i) (seq 1 (length nlist)) in
    let ivec := mkVarVector ivar n1 in
    eq_rect _ UnitDefinedFunction
            (mkNN_gen_0 n1 (combine nlist vlist) ivec f_activ_var f_activ) _ _.
  Next Obligation.
    f_equal.
    f_equal.
    rewrite combine_domain_eq; trivial.
    now rewrite map_length, seq_length.
  Qed.

  Program Definition mkNN_Var_gen (n1:nat) (nlist : list nat) (ivar wvar f_activ_var : SubVar)
             (f_activ : UnitDefinedFunction DTfloat) :
    UnitDefinedFunction (DTVector (last nlist n1)) :=
    let vlist := map (fun i => Sub wvar i) (seq 1 (length nlist)) in
    let ivec := mkVarVector ivar n1 in
    eq_rect _ UnitDefinedFunction
            (mkNN_Var_gen_0 n1 (combine nlist vlist) ivec f_activ_var f_activ) _ _.
  Next Obligation.
    f_equal.
    f_equal.
    rewrite combine_domain_eq; trivial.
    now rewrite map_length, seq_length.
  Qed.

  Program Definition mkNN_Mat_Var_gen (nsamp n1:nat) (nlist : list nat) (ivar wvar f_activ_var : SubVar)
             (f_activ : UnitDefinedFunction DTfloat) :
    UnitDefinedFunction (DTMatrix (last nlist n1) nsamp) :=
    let vlist := map (fun i => Sub wvar i) (seq 1 (length nlist)) in
    let imat := mkVarMatrix ivar n1 nsamp in
    eq_rect _ UnitDefinedFunction
            (mkNN_Mat_Var_gen_0 nsamp n1 (combine nlist vlist) imat f_activ_var f_activ) _ _.
  Next Obligation.
    f_equal.
    f_equal.
    rewrite combine_domain_eq; trivial.
    now rewrite map_length, seq_length.
  Qed.

  Definition softmax {n:nat} (NN : UnitDefinedFunction (DTVector n)) : UnitDefinedFunction (DTVector n) :=
    let expvar := Name "expvar" in
    let NNexp := VectorApply tt expvar (Exp tt (Var (expvar, DTfloat) tt)) NN in
    let NNexpscale := Divide tt (Number tt 1) (VectorSum tt NNexp) in
    VectorScalMult tt NNexpscale NNexp.

  Definition L2loss (nnvar ovar : SubVar) : UnitDefinedFunction DTfloat :=
    Square tt ( Minus tt (Var (nnvar, DTfloat) tt) (Var (ovar, DTfloat) tt) ).

  Definition L1loss (nnvar ovar : SubVar) : UnitDefinedFunction DTfloat :=
    Abs tt (Minus tt (Var (nnvar, DTfloat) tt) (Var (ovar, DTfloat) tt)).

  Definition Sigmoid (var : SubVar) : UnitDefinedFunction DTfloat :=
    Divide tt (Number tt 1)
           (Plus tt (Number tt 1) (Exp tt (Minus tt (Number tt 0) (Var (var, DTfloat) tt)))).

  Definition CrossEntropy (nnvar ovar : SubVar) : UnitDefinedFunction DTfloat :=
    let nnvar' := Var (nnvar, DTfloat) tt in
    let ovar' := Var (ovar, DTfloat) tt in
    Minus tt (Times tt (Minus tt ovar' (Number tt 1)) (Log tt (Minus tt (Number tt 1) nnvar')))
             (Times tt ovar' (Log tt nnvar')).

  Record testcases : Type := mkTest {ninput: nat; noutput: nat; ntest: nat;
                                     datavec : Vector ((Vector float ninput) * (Vector float noutput)) ntest}.

  Definition NNinstance1samp {ninput noutput : nat} (ivar : SubVar)
             (f_loss : UnitDefinedFunction DTfloat)
             (f_loss_NNvar f_loss_outvar : SubVar)
             (NN : UnitDefinedFunction (DTVector noutput)) (σ:df_env)
             (data: (Vector float ninput) * (Vector float noutput))
             : df_env * (UnitDefinedFunction DTfloat) :=
    let ipair := mk_env_entry (ivar, DTVector ninput) (fst data) in
    (cons ipair σ, VLossfun tt f_loss_NNvar f_loss_outvar f_loss NN (snd data)).

  Definition NNinstancebatch {ninput nsamp noutput : nat} (ivar : SubVar)
             (f_loss : UnitDefinedFunction DTfloat)
             (f_loss_NNvar f_loss_outvar : SubVar)
             (NN : UnitDefinedFunction (DTMatrix noutput nsamp)) (σ:df_env)
             (data: (Matrix float ninput nsamp) * (Matrix float noutput nsamp))
             : df_env * (UnitDefinedFunction DTfloat) :=
    let ipair := mk_env_entry (ivar, DTMatrix ninput nsamp) (fst data) in
    (cons ipair σ, MLossfun tt f_loss_NNvar f_loss_outvar f_loss NN (snd data)).

  Definition EvalNNinstance1samp {ninput noutput : nat} (ivar : SubVar)
             (f_loss : UnitDefinedFunction DTfloat)
             (f_loss_NNvar f_loss_outvar : SubVar)
             (NN : UnitDefinedFunction (DTVector noutput)) (σ:df_env)
             (data: (Vector float ninput) * (Vector float noutput))
             : option float :=
    let ipair := mk_env_entry (ivar, DTVector ninput) (fst data) in
    df_eval (cons ipair σ) (VLossfun tt f_loss_NNvar f_loss_outvar f_loss NN (snd data)).

  Definition EvalNNinstancebatch {ninput nsamp noutput : nat} (ivar : SubVar)
             (f_loss : UnitDefinedFunction DTfloat)
             (f_loss_NNvar f_loss_outvar : SubVar)
             (NN : UnitDefinedFunction (DTMatrix noutput nsamp)) (σ:df_env)
             (data: (Matrix float ninput nsamp) * (Matrix float noutput nsamp))
             : option float :=
    let ipair := mk_env_entry (ivar, DTMatrix ninput nsamp) (fst data) in
    df_eval (cons ipair σ) (MLossfun tt f_loss_NNvar f_loss_outvar f_loss NN (snd data)).

  (*
  Lemma NNinstance_unique_var (n1 n2 n3 : nat) (ivar : SubVar) (f_loss : DefinedFunction DTfloat)
        (NN2 : DefinedFunction (DTVector n3)) (inputs : (list float))
        (outputs : Vector float n3) (v:SubVar) :
    unique_var f_loss = Some v ->
    NNinstance n1 n2 n3 ivar f_loss NN2 inputs outputs =
    Some (
        let ipairs := (list_prod (map (fun n => (Sub ivar n)) (seq 1 n1))
                                 (map Number inputs)) in
        let losses := VectorMinus (df_subst_list NN2 ipairs)
                                  (DVector (vmap Number outputs)) in
        (VectorSum (VectorApply v f_loss losses))
      ).
  Proof.
    unfold NNinstance.
    intros.
    rewrite H.
    reflexivity.
  Qed.

  Lemma NNinstance_None (n1 n2 n3 : nat) (ivar : SubVar) (f_loss : DefinedFunction DTfloat)
        (NN2 : DefinedFunction (DTVector n3)) (inputs : (list float))
        (outputs : Vector float n3) :
    unique_var f_loss = None ->
    NNinstance n1 n2 n3 ivar f_loss NN2 inputs outputs = None.
  Proof.
    unfold NNinstance.
    intros.
    now rewrite H.
  Qed.
   *)

(*  Local Existing Instance floatish_interval. *)

  Definition lookup_list (σ:df_env) (lvar : list SubVar) : option (list float) :=
    listo_to_olist (map (fun v => (vartlookup σ (v, DTfloat)):option float) lvar).

  Definition combine_with {A:Type} {B:Type} {C:Type} (f: A -> B -> C )
             (lA : list A) (lB : list B) : list C :=
    map (fun '(a, b) => f a b) (combine lA lB).

  Definition combine3_with {A:Type} {B:Type} {C:Type} {D:Type} (f: A -> B -> C -> D)
             (lA : list A) (lB : list B) (lC : list C) : list D :=
    map (fun '(a, bc) => f a (fst bc) (snd bc)) (combine lA (combine lB lC)).

  Fixpoint streamtake (n : nat) {A : Type} (st : Stream A) : (list A) * (Stream A) :=
    match n with
    | 0 => (nil, st)
    | S n' => let rst := streamtake n' (Streams.tl st) in
              ((Streams.hd st)::(fst rst), snd rst)
    end.

  Lemma streamtake_n (n : nat) (A : Type) (st : Stream A) :
    length (fst (streamtake n st)) = n.
  Proof.
    generalize st.
    induction n.
    reflexivity.
    intros.
    simpl.
    f_equal.
    specialize IHn with (st := Streams.tl st0).
    apply IHn.
  Qed.

  Fixpoint env_update_first (l:df_env) (an:env_entry_type) : df_env
    := match l with
       | nil => nil
       | fv::os => if (projT1 an) == (projT1 fv) then an::os
                   else fv::(env_update_first os an)
       end.

  Definition env_update_list (l up:df_env) : df_env
    := fold_left (env_update_first) up l.

  Definition optimize_step
             (step : nat) (df : UnitDefinedFunction DTfloat) (σ:df_env) (lvar : list SubVar)
             (noise_st : Stream float) : (option df_env)*(Stream float) :=
    let lvart:list var_type := (map (fun v => (v, DTfloat)) lvar) in
    let ogradvec := df_eval_gradient σ df lvart in
    let alpha := 1 / (FfromZ (Z.of_nat (S step))) in
    let '(lnoise, nst) := streamtake (length lvar) noise_st in
    let olvals := lookup_list σ lvar in
    (match (ogradvec, olvals) with
     | (Some gradvec, Some lvals) =>
       Some (env_update_list
               σ
               (map (fun '(v,e) => mk_env_entry (v, DTfloat) (e:float))
                    (combine lvar (combine3_with
                                     (fun val grad noise => val - alpha*(grad + noise))
                                     lvals gradvec lnoise))))
     | (_, _) => None
     end, nst).

  Fixpoint get_noise_vector (n: nat) (noise_st: Stream float) :
    (Vector float n) * (Stream float) :=
    match n with
    | 0 => (vnil, noise_st)
    | S n' =>
      let noise := Streams.hd noise_st in
      let nst := Streams.tl noise_st in
      let '(vec, nst') := get_noise_vector n' nst in
      (vcons noise vec, nst')
    end.

  Fixpoint get_noise_matrix (n m: nat) (noise_st: Stream float) :
    (Matrix float n m) * (Stream float) :=
    let '(vec, nst) := get_noise_vector m noise_st in
    match n with
    | 0 => (fun i => vec, nst)
    | S n' =>
      let '(mat, nst') := get_noise_matrix n' m nst in
      (vcons vec mat, nst')
    end.

  Definition get_noise (t:definition_function_types) (noise_st:Stream float) :
    (definition_function_types_interp t) * (Stream float) :=
    match t with
    | DTfloat => (Streams.hd noise_st, Streams.tl noise_st)
    | DTVector n => get_noise_vector n noise_st
    | DTMatrix m n => get_noise_matrix m n noise_st
    end.

  Program Definition update_val_gradenv (grad_env:df_env) (x:var_type) (alpha:float) :
    definition_function_types_interp (snd x) ->
    definition_function_types_interp (snd x) ->
    definition_function_types_interp (snd x) :=
    (match snd x as y return snd x = y ->
                          definition_function_types_interp y ->
                          definition_function_types_interp y ->
                          definition_function_types_interp y with
       | DTfloat =>  fun pf val noise =>
                       match vartlookup grad_env x with
                       | Some grad => val - alpha * (coerce _ grad + noise)
                       | _ => val
                       end
       | DTVector n => fun pf val noise =>
                         match vartlookup grad_env x with
                         | Some grad =>
                           fun i =>
                             (val i) - alpha * (((coerce _ grad) i) + (noise i))
                         | _ => val
                         end
       | DTMatrix m n => fun pf val noise =>
                           match vartlookup grad_env x with
                           | Some grad =>
                             fun i j =>
                               (val i j) - alpha * (((coerce _ grad) i j) + (noise i j))
                         | _ => val
                         end
     end) (eq_refl _)
  .
    Next Obligation.
      rewrite pf; reflexivity.
    Qed.
    Next Obligation.
      rewrite pf; reflexivity.
    Qed.
    Next Obligation.
      rewrite pf; reflexivity.
    Qed.

  Definition update_entry (entry: env_entry_type) (grad_env:df_env) (alpha:float)
             (noise_st : Stream float) : (env_entry_type*Stream float) :=
    let x := projT1 entry in
    let val := projT2 entry in
    let '(noise, nst) := get_noise (snd x) noise_st in
    (mk_env_entry x (update_val_gradenv grad_env x alpha val noise), nst).

  Fixpoint list_arg_iter {A B} (f: A -> B -> A * B)
             (l:list A) (b: B) : (list A)*B :=
      match l with
      | nil => (l, b)
      | a :: l' =>
        let '(na, b') := f a b in
        let '(nl, nb) := list_arg_iter f l' b' in
        (na::nl, nb)
      end.

  Definition harmonic_lr (step : nat) : float := 1 / (FfromZ (Z.of_nat (S step))).

  Definition optimize_step_backprop
             (step : nat) (df : UnitDefinedFunction DTfloat) (σ:df_env) (lr : nat -> float)
             (noise_st : Stream float) (dvars : list var_type)
    : (option df_env)*(Stream float) :=
    match df_eval_backprop_deriv σ df (gradenv_init dvars) 1 with
    | Some gradenv =>
      let alpha := lr step in
      let '(env, nst) := list_arg_iter (fun a b => update_entry a gradenv alpha b) σ noise_st in
      (Some env, nst)
    | _ => (None, noise_st)
    end.

  Definition optimize_step_tree_backprop
             (step : nat) (df : UnitDefinedFunction DTfloat) (σ:df_env) (lr : nat -> float)
             (noise_st : Stream float) (dvars : list var_type)
    : (option df_env)*(Stream float) :=
    match df_eval_tree σ df with
    | Some df_tree =>
      match df_eval_tree_backprop_deriv (* σ not-needed *) nil df_tree (gradenv_init dvars) 1 with
      | Some gradenv =>
        let alpha := lr step in
        let '(env, nst) := list_arg_iter (fun a b => update_entry a gradenv alpha b) σ noise_st in
        (Some env, nst)
      | _ => (None, noise_st)
      end
    | _ => (None, noise_st)
    end.

  Fixpoint optimize_steps
           (start count:nat) (df : UnitDefinedFunction DTfloat) (σ:df_env) (lvar : list SubVar)
           (noise_st : Stream float) : (option df_env)*(Stream float) :=
    match count with
    | 0 => (Some σ, noise_st)
    | S n =>
      match optimize_step start df σ lvar noise_st with
      | (Some σ', noise_st') => optimize_steps (S start) n df σ' lvar noise_st'
      | (None, noise_st') => (None, noise_st')
      end
    end.

  Fixpoint optimize_steps_backprop
           (start count:nat) (df : UnitDefinedFunction DTfloat) (σ:df_env) (lr : nat -> float)
           (noise_st : Stream float) (dvars : list var_type)
    : (option df_env)*(Stream float) :=
    match count with
    | 0 => (Some σ, noise_st)
    | S n =>
      match optimize_step_backprop start df σ lr noise_st dvars with
      | (Some σ', noise_st') => optimize_steps_backprop (S start) n df σ' lr noise_st' dvars
      | (None, noise_st') => (None, noise_st')
      end
    end.

  Fixpoint optimize_steps_tree_backprop
           (start count:nat) (df : UnitDefinedFunction DTfloat) (σ:df_env) (lr : nat -> float)
           (noise_st : Stream float) (dvars : list var_type)
    : (option df_env)*(Stream float) :=
    match count with
    | 0 => (Some σ, noise_st)
    | S n =>
      match optimize_step_tree_backprop start df σ lr noise_st dvars with
      | (Some σ', noise_st') => optimize_steps_tree_backprop (S start) n df σ' lr noise_st' dvars
      | (None, noise_st') => (None, noise_st')
      end
    end.

Definition Fmax (a b:float) : float :=
  if Fgt a b then a else b.

Definition Fmin (a b:float) : float :=
  if Flt a b then a else b.

Definition vmax {n : nat} (vec : Vector float n)
          := vector_fold_right1 Fmax (FfromZ 0) id vec.

Definition vmin {n : nat} (vec : Vector float n)
          := vector_fold_right1 Fmin (FfromZ 0) id vec.

Definition nrows {A} (l : list (list A)) := List.length l.

Definition ncols {A} (l : list (list A)) :=
      match l with
      | nil => 0%nat
      | r :: _ => List.length r
      end.

Definition normalizeIntData (l:list (list Z)) : Matrix float (nrows l) (ncols l) :=
  let mat : Matrix float (nrows l) (ncols l) :=
      fun i j => FfromZ (List.nth (proj1_sig j) (List.nth (proj1_sig i) l nil) (0)%Z) in
  let tmat := transpose mat in
  let maxes := vmap vmax tmat in
  let mins := vmap vmin tmat in
  fun i j => ((mat i j) - (mins j))/((maxes j)- (mins j)).

Program Definition splitLastCol {nsamp ncols : nat}  (data : Matrix float nsamp ncols)
        (pf : (ncols > 0)%nat)
        : (Matrix float (ncols-1) nsamp) * (Matrix float 1 nsamp) :=
  (fun i j => data j i, fun i j => data j (i + ncols-1)%nat).
Next Obligation.
  omega.
Defined.
Next Obligation.
  omega.
Defined.

Definition init_env2 (dim1 dim2 dim3 : nat) (w b : string)
        (ranm1 : Matrix float dim2 dim1)
        (ranm2 : Matrix float dim3 dim2) : df_env :=
  let wvar := Name w in
  let bvar := Name b in
  let wvar1 := (Sub wvar 1, DTMatrix dim2 dim1) in
  let wvar2 := (Sub wvar 2, DTMatrix dim3 dim2) in
  let bvar1 := (Sub bvar 1, DTVector dim2) in
  let bvar2 := (Sub bvar 2, DTVector dim3) in
  (*  (mk_env_entry wvar1 ranm1) :: (mk_env_entry wvar2 ranm2) :: *)
  (mk_env_entry wvar1 (fun i j => (ranm1 i j) / (Fsqrt (FfromZ (Z.of_nat dim1))))) ::
  (mk_env_entry wvar2 (fun i j => (ranm2 i j) / (Fsqrt (FfromZ (Z.of_nat dim2))))) ::
  (mk_env_entry bvar1 (ConstVector dim2 0)) ::
  (mk_env_entry bvar2 (ConstVector dim3 0)) :: nil.

 Definition mkNN_wisconsin (nsamp:nat) (ivar wvar bvar f1v f2v : SubVar) : (UnitDefinedFunction (DTMatrix 1 nsamp) * list var_type) :=
   let f1_activ := Max tt (Var (f1v,DTfloat) tt) (Number tt 0) in
   let f2_activ := Sigmoid f2v in
   mkNN2_Var_Mat_bias 9 15 1 nsamp ivar wvar bvar f1_activ f2_activ f1v f2v.

 Program Definition eval_wisconsin_batch (nsamp : nat)
            (σ:df_env)
            (normaldata: Matrix float nsamp 10)   : option float :=
   let ivar := (Name "i") in
   let flnnv := (Name "NNv") in
   let outnnv := (Name "outnnv") in
   let '(NN, dvars) := mkNN_wisconsin nsamp ivar (Name "w") (Name "b") (Name "f1v") (Name "f2v") in
   @EvalNNinstancebatch 9 nsamp 1 ivar (CrossEntropy flnnv outnnv) flnnv outnnv
             NN σ (splitLastCol normaldata _).
 Next Obligation.
   omega.
 Qed.

Program Definition wisconsin_instance_batch (nsamp : nat)
            (σ:df_env)
            (normaldata: Matrix float nsamp 10)
             : (df_env * (UnitDefinedFunction DTfloat)) * (list var_type) :=
   let ivar := (Name "i") in
   let flnnv := (Name "NNv") in
   let outnnv := (Name "outnnv") in
   let '(NN, dvars) := mkNN_wisconsin nsamp ivar (Name "w") (Name "b") (Name "f1v") (Name "f2v") in
   (@NNinstancebatch 9 nsamp 1 ivar (CrossEntropy flnnv outnnv) flnnv outnnv
             NN
             σ
             (splitLastCol normaldata _), dvars).
Next Obligation.
  omega.
Qed.

CoFixpoint zeronoise : Stream float := Cons 0 zeronoise.

(*
https://towardsdatascience.com/predict-malignancy-in-breast-cancer-tumors-with-your-own-neural-network-and-the-wisconsin-dataset-76271a05e941
*)

Definition wisconsin_test (nsamp count : nat)
            (σ:df_env)
            (normaldata: Matrix float nsamp 10): list float :=
  let '(nninst, dvars) := wisconsin_instance_batch nsamp σ normaldata in
  let lr := fun _ => 1 / (FfromZ 100) in
  let onenv := fst (optimize_steps_tree_backprop 0 count (snd nninst) (fst nninst) lr
                                                zeronoise dvars) in
  match onenv with
  | Some nenv => match df_eval nenv (snd nninst) with
                 | Some val => val :: nil
                 | _ => nil
                 end
  | _ => nil
  end.

Definition wisconsin_test_env (nsamp count : nat)
            (σ:df_env)
            (normaldata: Matrix float nsamp 10): df_env :=
  let '(nninst, dvars) := wisconsin_instance_batch nsamp σ normaldata in
  let lr := fun _ => 1 / (FfromZ 100) in
  let onenv := fst (optimize_steps_tree_backprop 0 count (snd nninst) (fst nninst) lr
                                                zeronoise dvars) in
  match onenv with
  | Some nenv => nenv
  | _ => nil
  end.

Example xvar:var_type := (Name "x", DTfloat).
Example xfun:UnitDefinedFunction DTfloat := Var xvar tt.
Example tquad:UnitDefinedFunction DTfloat := Times tt xfun xfun.
Example quad:UnitDefinedFunction DTfloat := Minus tt (Times tt xfun xfun) (Number tt 1).
Example squad:UnitDefinedFunction DTfloat := Minus tt (Square tt xfun) (Number tt 1).
Example env : df_env := (mk_env_entry xvar (FfromZ 5))::nil.

Example gradenv := match df_eval_backprop_deriv env quad (gradenv_init (xvar :: nil)) 1 with
                   | Some gradenv => gradenv
                   | _ => nil
                   end.

Example gradenv_tree :=
  match df_eval_tree env quad with
  | Some df_tree =>
    match df_eval_tree_backprop_deriv nil df_tree (gradenv_init (xvar :: nil)) 1 with
    | Some gradenv => gradenv
    | _ => nil
    end
  | _ => nil
  end.

Example wisconsin_gradenv (nsamp : nat)
            (σ:df_env)
            (normaldata: Matrix float nsamp 10) : df_env :=
  let '((env,nn),dvars) := wisconsin_instance_batch nsamp σ normaldata in
  match df_eval_backprop_deriv env nn (gradenv_init dvars) 1 with
  | Some gradenv => gradenv
  | _ => nil
  end.

Example wisconsin_gradenv_tree (nsamp : nat)
            (σ:df_env)
            (normaldata: Matrix float nsamp 10) : df_env :=
  let '((env,nn),dvars) := wisconsin_instance_batch nsamp σ normaldata in
  match df_eval_tree env nn with
  | Some df_tree =>
    match df_eval_tree_backprop_deriv nil df_tree (gradenv_init dvars) 1 with
    | Some gradenv => gradenv
    | _ => nil
    end
  | _ => nil
  end.

 Definition mkperceptron (n1 n2 : nat) (ivar wvar bvar : SubVar) (f_activ : UnitDefinedFunction DTfloat) (f_activ_var : SubVar) : (UnitDefinedFunction (DTVector n2) * list var_type) :=
    let mat1 := (wvar, DTMatrix n2 n1) in
    let b1 :=  (bvar, DTVector n2) in
    let ivec := Var (ivar, DTVector n1) tt in
    let N1 := mkNN_bias_step n1 n2 ivec (Var mat1 tt) (Var b1 tt) f_activ_var f_activ in
    (N1,  mat1 :: b1 :: nil).

 Definition mkNN_test1 :=
   let ivar := Name "i" in
   let wvar := Name "w" in
   let bvar := Name "b" in
   let f1v := Name "f1" in
   let f1_activ := Max tt (Var (f1v,DTfloat) tt) (Number tt 0) in
   let '(nn, dvars) := mkperceptron 2 2 ivar wvar bvar f1_activ f1v in
   let winit := mk_env_entry (wvar, DTMatrix 2 2) (ConstMatrix 2 2 1) in
   let binit := mk_env_entry (bvar, DTVector 2) (ConstVector 2 1) in
   let datain := ConstVector 2 1 in
   let env := (mk_env_entry (ivar, DTVector 2) (ConstVector 2 1)):: winit :: binit :: nil in
   (env, nn).

 Definition mkNN_test :=
   let ivar := Name "i" in
   let wvar := Name "w" in
   let bvar := Name "b" in
   let f1v := Name "f1" in
(*   let f1_activ := Max tt (Var (f1v,DTfloat) tt) (Number tt 0) in *)
   let f1_activ := Sigmoid f1v in
   let '(nn, dvars) := mkperceptron 2 2 ivar wvar bvar f1_activ f1v in
   let datain := ConstVector 2 1 in
   let dataout := ConstVector 2 0 in
   let loss_nnvar := Name "lv" in
   let loss_ovar := Name "ov" in
   let floss := CrossEntropy loss_nnvar loss_ovar in
   let winit := mk_env_entry (wvar, DTMatrix 2 2) (ConstMatrix 2 2 (FfromZ 2)) in
   let binit := mk_env_entry (bvar, DTVector 2) (ConstVector 2 (FfromZ 3)) in
   let env := winit :: binit :: nil in
   (NNinstance1samp ivar floss loss_nnvar loss_ovar nn env (datain, dataout), dvars).

Definition NN_test (count : nat) : list float :=
  let '(nninst, dvars) := mkNN_test in
  let onenv := fst (optimize_steps_tree_backprop 0 count (snd nninst) (fst nninst) harmonic_lr
                                                zeronoise dvars) in
  match onenv with
  | Some nenv => match df_eval nenv (snd nninst) with
                 | Some val => val :: nil
                 | _ => nil
                 end
  | _ => nil
  end.

Example NN_test_gradenv : df_env :=
  let '((env, nn), dvars) := mkNN_test in
  match df_eval_backprop_deriv env nn (gradenv_init dvars) 1 with
  | Some gradenv => gradenv
  | _ => nil
  end.

Example NN_test_gradenv_tree :df_env :=
  let '((env, nn), dvars) := mkNN_test in
  match df_eval_tree env nn with
  | Some df_tree =>
    match df_eval_tree_backprop_deriv nil df_tree (gradenv_init dvars) 1 with
    | Some gradenv => gradenv
    | _ => nil
    end
  | _ => nil
  end.

Definition NN_test_env (count : nat) : df_env :=
  let '(nninst, dvars) := mkNN_test in
  let onenv := fst (optimize_steps_tree_backprop 0 count (snd nninst) (fst nninst)
                                  harmonic_lr zeronoise dvars) in
  match onenv with
  | Some nenv => nenv
  | _ => nil
  end.

Definition NN_test_NN  :=
  let '((env,nn), dvars) := mkNN_test in
  nn.

Definition NN_test_val : list float :=
  let '((env,nn), dvars) := mkNN_test in
  match df_eval env nn with
  | Some val => val :: nil
  | _ => nil
  end.

Definition test_update_val_gradenv
  := update_val_gradenv gradenv xvar 1 (FfromZ 5) 0.

Definition test_update : df_env :=
  (fst (update_entry (mk_env_entry xvar (FfromZ 5)) gradenv 1 zeronoise))::nil.

Definition test_optimize_step_backprop
             (step : nat) (df : UnitDefinedFunction DTfloat) (σ:df_env)
             (noise_st : Stream float) (dvars : list var_type) : df_env :=
  match fst (optimize_step_backprop step df σ harmonic_lr noise_st dvars) with
  | Some env => env
  | _ => nil
  end.

Definition test_optimize_step_tree_backprop
             (step : nat) (df : UnitDefinedFunction DTfloat) (σ:df_env)
             (noise_st : Stream float) (dvars : list var_type) : df_env :=
  match fst (optimize_step_tree_backprop step df σ harmonic_lr noise_st dvars) with
  | Some env => env
  | _ => nil
  end.

Example testopt := test_optimize_step_backprop 0 quad env zeronoise (xvar :: nil).
Example testreeopt := test_optimize_step_tree_backprop 0 quad env zeronoise (xvar :: nil).

Example opt := fst (optimize_steps 0 2 quad env ((fst xvar) :: nil) zeronoise).
Example opt2 := fst (optimize_steps_tree_backprop 0 2 quad env harmonic_lr zeronoise (xvar :: nil)).

Example val := 1+1.

End GenNN.

(*
Local Instance floatish_interval : floatish := floatish_interval_gen 53.
Compute df_eval NN_test_env NN_test_NN.
*)


Require Import HoTT.
From GR.bicategories Require Import
     bicategory.bicategory
     bicategory.bicategory_laws
     lax_functor.lax_functor.

Section ConstantFunctor.
  Context `{Univalence}
          {C D : BiCategory}.
  Variable (Y : D).

  Definition constant_d
    : PseudoFunctor_d C D.
  Proof.
    make_pseudo_functor.
    - exact (fun _ => Y).
    - exact (fun _ _ _ => id₁ Y).
    - exact (fun _ _ _ _ _ => id₂ (id₁ Y)).
    - exact (fun _ _ _ _ _ => left_unit (id₁ Y)).
    - exact (fun _ => id₂ (id₁ Y)).
    - exact (fun _ _ _ _ _ => left_unit_inv (id₁ Y)).
    - exact (fun _ => id₂(id₁ Y)).
  Defined.

  Definition constant_is_pseudo
    : is_pseudo_functor_p constant_d.
  Proof.
    make_is_pseudo ; intros ; cbn.
    - reflexivity.
    - rewrite vcomp_left_identity.
      reflexivity.
    - rewrite hcomp_id₂, vcomp_right_identity, vcomp_left_identity.
      reflexivity.
    - rewrite hcomp_id₂, vcomp_left_identity, vcomp_right_identity.
      apply right_unit_id_is_left_unit_id.
    - rewrite hcomp_id₂, vcomp_right_identity, vcomp_left_identity.
      reflexivity.
    - rewrite vcomp_left_identity.
      rewrite vcomp_assoc.
      rewrite <- triangle_r.
      rewrite right_unit_id_is_left_unit_id.
      reflexivity.
    - apply left_unit_right.
    - apply left_unit_left.
    - apply vcomp_left_identity.
    - apply vcomp_right_identity.
  Qed.

  Definition lax_constant
    : LaxFunctor C D
    := Build_PseudoFunctor constant_d constant_is_pseudo.

  Global Instance lax_ap_functor_pseudo
    : is_pseudo lax_constant
    := _.
End ConstantFunctor.


From lrust.lang Require Export notation.
From lrust.lang Require Import heap proofmode.
From iris.prelude Require Import options.

Definition swap : val :=
  rec: "swap" ["p1";"p2";"len"] :=
    if: "len" ≤ #0 then #☠
    else let: "x" := !"p1" in
         "p1" <- !"p2";;
         "p2" <- "x";;
         "swap" ["p1" +ₗ #1 ; "p2" +ₗ #1 ; "len" - #1].

Lemma wp_swap `{!lrustGS Σ} E l1 l2 vl1 vl2 (n : Z):
  Z.of_nat (length vl1) = n → Z.of_nat (length vl2) = n →
  {{{ l1 ↦∗ vl1 ∗ l2 ↦∗ vl2 }}}
    swap [ #l1; #l2; #n] @ E
  {{{ RET #☠; l1 ↦∗ vl2 ∗ l2 ↦∗ vl1 }}}.
Proof.
  iIntros (Hvl1 Hvl2 Φ) "(Hl1 & Hl2) HΦ".
  iLöb as "IH" forall (n l1 l2 vl1 vl2 Hvl1 Hvl2). wp_rec.
  wp_op; case_bool_decide; wp_if.
  - iApply "HΦ". assert (n = O) by lia; subst.
    destruct vl1, vl2; try discriminate. by iFrame.
  - destruct vl1 as [|v1 vl1], vl2 as [|v2 vl2], n as [|n|]; try (discriminate || lia).
    revert Hvl1 Hvl2. intros [= Hvl1] [= Hvl2]; rewrite !heap_mapsto_vec_cons. subst n.
    iDestruct "Hl1" as "[Hv1 Hl1]". iDestruct "Hl2" as "[Hv2 Hl2]".
    Local Opaque Zminus.
    wp_read; wp_let; wp_read; do 2 wp_write. do 3 wp_op.
    iApply ("IH" with "[%] [%] Hl1 Hl2"); [lia..|].
    iIntros "!> [Hl1 Hl2]"; iApply "HΦ"; by iFrame.
Qed.


From Categories Require Import Essentials.Notations.
From Categories Require Import Essentials.Types.
From Categories Require Import Essentials.Facts_Tactics.
From Categories Require Import Category.Main.
From Categories Require Import Functor.Functor Functor.Functor_Ops
        Functor.Representable.Hom_Func_Prop.
From Categories Require Import Ext_Cons.Prod_Cat.Prod_Cat.
From Categories Require Import NatTrans.NatTrans NatTrans.Operations.
From Categories Require Import Adjunction.Adjunction Adjunction.Duality
        Adjunction.Adj_Facts.
From Categories Require Import KanExt.Global KanExt.Local
     KanExt.LocalFacts.ConesToHom KanExt.LocalFacts.HomToCones
     KanExt.GlobalDuality.

Local Open Scope functor_scope.

(** This module contains conversion from global to local kan extensions. *)
Section Global_to_Local_Right.
  Context {C C' : Category}
          (p : C --> C')
          (D : Category)
          (rke : Right_KanExt p D)
          (F : C --> D).

  (** The cone which (we will prove) is the local kan extension. *)
  Definition Cone_for_LoKan : LoKan_Cone p F :=
    {|
      cone_apex := (rke _o F)%object;
      cone_edge :=
        @adj_morph_ex _ _ _ _
                      (right_kan_ext_adj rke) (rke _o F) F (NatTrans_id _)
    |}.

  Section Cone_Morph_to_Cone_for_LoKan.
    Context (Cn : LoKan_Cone p F).

    (** We show that any natural transformation from any cone (the apex functor
        of a cone) to (apex functor of) the cone we constructed above remains
        the same under the following transformation:

        morph : Cn ——––>  (rke _o F)

        morph ∘_h (NatTrans_id p) : (Cn ∘ p) ————> ((rke _o F) ∘ p)

        η_{NatTrans_id (rke _o F)} ∘ (morph ∘_h (NatTrans_id p))
           : (Cn ∘ p) ————> F

        rke @_a (Cn ∘ p) F (η_{NatTrans_id (rke _o F)}
                             ∘ (morph ∘_h (NatTrans_id p))) :
          (rke _o (Cn ∘ p)) ————> (rke _o F)

        (rke @_a (Cn ∘ p) F (η_{NatTrans_id (rke _o F)}
                               ∘ (morph ∘_h (NatTrans_id p))))
          ∘ (Trans (adj_unit (right_kan_ext_adj rke)) Cn) :
          Cn ————> (rke _o F)

        This result is used to show existence and unique ness of cones from Cn
        to the cone constructed above.
     *)

    Lemma Cone_Morph_to_Cone_for_LoKan_adj_unit_rke_id
          (morph : (Cn --> ((rke _o)%object F))%nattrans) :
      morph =
      (
        (
          (rke @_a)%morphism (Cn ∘ p)%functor F
                 (
                   (adj_morph_ex
                      (right_kan_ext_adj rke) (NatTrans_id ((rke _o) F)%object))
                     ∘ (morph ∘_h (NatTrans_id p))
                 )
        ) ∘ (Trans (adj_unit (right_kan_ext_adj rke)) Cn)
      )%nattrans.
    Proof.
      rewrite (@F_compose); cbn.
      rewrite NatTrans_compose_assoc.
      cbn_rewrite <- (@Trans_com
                       _ _ _ _
                       (@adj_unit _ _ _ _ (right_kan_ext_adj rke)) _ _ morph).
      rewrite <- NatTrans_compose_assoc.
      cbn_rewrite <- (
                    @adj_morph_com _ _ _ _
                                   (right_kan_ext_adj rke)
                                   _
                                   _
                                   (NatTrans_id ((rke _o)%object F))
                  ).
      rewrite NatTrans_id_unit_left.
      trivial.
    Qed.

    (** Given a cone, we construct a cone morph to the cone morph that we
        constructed above. *)
    Program Definition Cone_Morph_to_Cone_for_LoKan :
      LoKan_Cone_Morph Cn Cone_for_LoKan :=
      {|
        cone_morph :=
          (((rke _a (cone_edge Cn))%morphism)
             ∘ (Trans (adj_unit (right_kan_ext_adj rke)) Cn))%nattrans
      |}.

    Next Obligation.
    Proof.
      match goal with
        [|- _ = NatTrans_compose (NatTrans_hor_comp _ ?X) _] =>
        apply (@adj_morph_unique _ _ _ _ (right_kan_ext_adj rke) _ _ X); trivial
      end.
      apply Cone_Morph_to_Cone_for_LoKan_adj_unit_rke_id.
    Qed.

  End Cone_Morph_to_Cone_for_LoKan.

  Section Cone_Morph_to_Cone_for_LoKan_Unique.
    Context {Cn : LoKan_Cone p F} (M M' : LoKan_Cone_Morph Cn Cone_for_LoKan).

    (** Cone morph to the cone constructed is unique. *)
    Theorem Cone_Morph_to_Cone_for_LoKan_Unique : (M = M' :> (_ --> _)%nattrans).
    Proof.
      rewrite (Cone_Morph_to_Cone_for_LoKan_adj_unit_rke_id Cn M).
      rewrite (Cone_Morph_to_Cone_for_LoKan_adj_unit_rke_id Cn M').
      do 2 apply f_equal.
      set (H := cone_morph_com M'); rewrite (cone_morph_com M) in H; exact H.
    Qed.

  End Cone_Morph_to_Cone_for_LoKan_Unique.

  (** The conversion from global kan extensions to local kan extensions *)
  Definition Global_to_Local_Right : Local_Right_KanExt p F :=
    {|
      LRKE := Cone_for_LoKan;
      LRKE_morph_ex := Cone_Morph_to_Cone_for_LoKan;
      LRKE_morph_unique := @Cone_Morph_to_Cone_for_LoKan_Unique
    |}.

End Global_to_Local_Right.

(** Teh conversion from global left kan extensions to local left kan extensions
    is jsut the dual what we just proved. *)
Section Global_to_Local_Left.
  Context {C C' : Category}
          (p : C --> C')
          (D : Category)
          (lke : Left_KanExt p D)
          (F : C --> D).

  Definition Global_to_Local_Left : Local_Left_KanExt p F :=
    Global_to_Local_Right _ _ (KanExt_Left_to_Right _ _ lke) (F^op).

End Global_to_Local_Left.


Module M.
  Inductive I : Set := C : nat -> I.
End M.

Module M1 := M.

Goal forall x, M.C x = M1.C 0 -> x = 0 .
  intros x H.
  (*
     injection sur deux constructeurs egaux mais appeles
     par des modules differents
     *)
  injection H.
  tauto.
Qed.

Goal  M.C 0 <> M1.C 1.
  (*
     Discriminate sur deux constructeurs egaux mais appeles
     par des modules differents
     *)
  intro H;discriminate H.
Qed.

Goal forall x,  M.C x = M1.C 0 -> x = 0.
  intros x H.
  (*
     inversion sur deux constructeurs egaux mais appeles
     par des modules differents
     *)
  inversion H. reflexivity.
Qed.


(*
  Author(s):
    Andrej Dudenhefner (1)
  Affiliation(s):
    (1) Saarland University, Saarbrücken, Germany
*)

Require Import Undecidability.Synthetic.EnumerabilityFacts.

Require Import List Lia.
Require Cantor.
Require Import ssreflect ssrbool ssrfun.

Set Default Goal Selector "!".

(* Enumerable X allows injection into nat *)
Class Enumerable (X: Type) :=
  {
    to_nat : X -> nat;
    of_nat : nat -> X;
    enumP {x: X} : of_nat (to_nat x) = x
  }.

#[export] Instance nat_Enumerable : Enumerable nat.
Proof. by exists id id. Qed.

#[export] Instance bool_Enumerable : Enumerable bool.
Proof.
  exists (fun b => if b then 1 else 0) (fun n => if n is 0 then false else true).
  by case.
Qed.

#[export] Instance prod_Enumerable {X Y: Type} {enumX: Enumerable X} {enumY: Enumerable Y} : Enumerable (X * Y).
Proof.
  exists
    (fun '(x, y) => Cantor.to_nat (to_nat x, to_nat y))
    (fun n => match Cantor.of_nat n with | (n1, n2) => (of_nat n1, of_nat n2) end).
  move=> [x y]. by rewrite Cantor.cancel_of_to !enumP.
Qed.

#[export] Instance sum_Enumerable {X Y: Type} {enumX: Enumerable X} {enumY: Enumerable Y} : Enumerable (X + Y).
Proof.
  exists
    (fun t => match t with | inl x => to_nat (0, to_nat x) | inr y => to_nat (1, to_nat y) end)
    (fun n => match of_nat n with | (0, n) => inl (of_nat n) | (1, n) => inr (of_nat n) | _ => inl (of_nat n) end).
  by case => ?; rewrite ?enumP.
Qed.

Opaque RoseTree.of_nat RoseTree.to_nat.

#[export] Instance list_Enumerable {X: Type} {enumX: Enumerable X} : Enumerable (list X).
Proof.
  exists
    (fun l => RoseTree.to_nat (RoseTree.mk (map (fun x => RoseTree.mk (repeat (RoseTree.mk nil) (to_nat x))) l)))
    (fun n => match (RoseTree.of_nat n) with RoseTree.mk rs => map (fun '(RoseTree.mk r's) => of_nat (length r's)) rs end).
  move=> l /=. rewrite RoseTree.cancel_of_to map_map -[RHS]map_id.
  apply: map_ext => ?. by rewrite repeat_length enumP.
Qed.


Set Warnings "-notation-overridden".

Require Import Category.Lib.
Require Export Category.Theory.Natural.Transformation.
Require Export Category.Functor.Product.
Require Import Category.Construction.Opposite.
Require Import Category.Instance.Fun.
Require Import Category.Instance.Sets.

Generalizable All Variables.
Set Primitive Projections.
Set Universe Polymorphism.
Unset Transparent Obligations.

(* Bifunctors can be curried:

  C × D ⟶ E --> C ⟶ [D, E] ~~~ (C, D) -> E --> C -> D -> E

  Where ~~~ should be read as "Morally equivalent to".

  Note: We do not need to define Bifunctors as a separate class, since they
  can be derived from functors mapping to a category of functors. So in the
  following two definitions, [P] is effectively our bifunctor.

  The trick to [bimap] is that both the [Functor] instances we need (for
  [fmap] and [fmap1]), and the [Natural] instance, can be found in the
  category of functors we're mapping to by applying [P]. *)

Program Definition Hom `(C : Category) : C^op ∏ C ⟶ Sets := {|
  fobj := fun p => {| carrier   := @hom C (fst p) (snd p)
                    ; is_setoid := @homset (C) (fst p) (snd p) |};
  fmap := fun x y (f : x ~{C^op ∏ C}~> y) =>
            {| morphism := fun g => snd f ∘ g ∘ fst f |}
|}.

Program Definition Curried_Hom `(C : Category) : C^op ⟶ [C, Sets] := {|
  fobj := fun x => {|
    fobj := fun y => {| carrier := @hom C x y
                      ; is_setoid := @homset C x y |};
    fmap := fun y z (f : y ~{C}~> z) =>
              {| morphism := fun (g : x ~{C}~> y) =>
                               (f ∘ g) : x ~{C}~> z |}
  |};
  fmap := fun x y (f : x ~{C^op}~> y) => {|
    transform := fun _ => {| morphism := fun g => g ∘ op f |}
  |}
|}.
Next Obligation.
  simpl; intros.
  unfold op.
  apply comp_assoc.
Qed.

Coercion Curried_Hom : Category >-> Functor.

Notation "[Hom A , ─]" := (@Curried_Hom _ A) : functor_scope.

Program Definition CoHom_Alt `(C : Category) : C ∏ C^op ⟶ Sets :=
  Hom C ○ Swap.

Program Definition CoHom `(C : Category) : C ∏ C^op ⟶ Sets := {|
  fobj := fun p => {| carrier   := @hom (C^op) (fst p) (snd p)
                    ; is_setoid := @homset (C^op) (fst p) (snd p) |};
  fmap := fun x y (f : x ~{C ∏ C^op}~> y) =>
    {| morphism := fun g => snd f ∘ g ∘ fst f |}
|}.

Program Definition Curried_CoHom `(C : Category) : C ⟶ [C^op, Sets] := {|
  fobj := fun x => {|
    fobj := fun y => {| carrier := @hom (C^op) x y
                      ; is_setoid := @homset (C^op) x y |};
    fmap := fun y z (f : y ~{C^op}~> z) =>
              {| morphism := fun (g : x ~{C^op}~> y) =>
                               (f ∘ g) : x ~{C^op}~> z |}
  |};
  fmap := fun x y (f : x ~{C}~> y) => {|
    transform := fun _ => {| morphism := fun g => g ∘ op f |}
  |}
|}.
Next Obligation.
  simpl; intros.
  symmetry.
  apply comp_assoc.
Qed.

Notation "[Hom ─ , A ]" := (@Curried_CoHom _ A) : functor_scope.


Require Export Coq.Strings.String.
From Coq Require Import Lists.List.
Require Import Ascii.
From Coq Require Import Strings.String.
Require Import Strings.String.
Open Scope string_scope.
Scheme Equality for string.

(*SYNTAX*)

(*Arithmetic expressions*)

Inductive Exp :=
| anum : nat -> Exp
| avar : string -> Exp
| aplus : Exp -> Exp -> Exp
| aminus : Exp -> Exp -> Exp
| amul : Exp -> Exp -> Exp
| adiv : Exp -> Exp -> Exp
| btrue : Exp
| bfalse : Exp
| bless : Exp -> Exp -> Exp
| bgreat : Exp -> Exp -> Exp
| bleq : Exp -> Exp -> Exp
| bgeq : Exp -> Exp -> Exp
| beq : Exp -> Exp -> Exp
| bnot : Exp -> Exp
| band : Exp -> Exp -> Exp
| bor : Exp -> Exp -> Exp.

Coercion anum : nat >-> Exp.
Coercion avar : string >-> Exp.

Notation "A +' B" := (aplus A B) (at level 50).
Notation "A -' B" := (aminus A B) (at level 50).
Notation "A *' B" := (amul A B) (at level 46).
Notation "A /' B" := (adiv A B) (at level 46).
Notation "A <' B" := (bless A B) (at level 53).
Notation "A >' B" := (bgreat A B) (at level 53).
Notation "A <=' B" := (bleq A B) (at level 53).
Notation "A >=' B" := (bgeq A B) (at level 53).
Notation "A ==' B" := (beq A B) (at level 53, no associativity).
Notation "'~' A" := (bnot A) (at level 75, right associativity).
Infix "and'" := band (at level 60).
Infix "or'" := bor (at level 60).

Compute 1 +' 2.
Compute "x" +' "y".
Compute 5 -' 3.
Compute "a" -' 7.
Compute 10 *' 100.
Compute 23 *' "b".
Compute 80 /' 5.
Compute "x" /' "y".

Compute 5 <' 7.
Compute "x" >' 10.
Compute 8 <=' "i".
Compute "j" >=' 25.
Compute 6 ==' 6.
Compute "a" ==' "b".
Compute btrue ==' 2.

Compute ~btrue.
Compute ~(btrue) and' bfalse.
Compute 5 or' 0.
Compute ~(btrue or' 0) and' "x".

(*functions*)

Inductive list_variable : Type :=
| nil
| cons (s : string) (l : list_variable).

Infix "::" := cons.

Notation "[]" := nil (format "[]") : list_scope.
Notation "[ A ]" := (cons A nil) : list_scope.
Notation "[ A , B , .. , C ]" := (cons A (cons B .. (cons C nil) ..)) : list_scope.

Inductive Function :=
| function : string -> list_variable -> Function.

Notation "A 'takes' B" := (function A B) (at level 50).

Compute "cmmdc" takes ["x", "y"].

Compute "printWelcome" takes [].

Compute "add1toit" takes ["number"].

(*objects*)

Inductive Object :=
| object : string -> list_variable -> Function -> Object.

Compute object ("student") (["age","name","grade"]) ("getGrade" takes []).

Compute object ("car") (["make","model","color"]) ("setColor" takes ["color"]).

(*Memory*)

Inductive Value :=
| undef : Value
| nat_val : nat -> Value
| bool_val : bool -> Value
| string_val : string -> Value.

Scheme Equality for Value.

Definition Variabila := string -> nat.
Definition Adresa := nat -> Value.

(*statements*)

Inductive Stmt :=
| var_declare : string -> Stmt
| object_declare : Object -> Stmt
| function_declare : Function -> Stmt -> Stmt
| assignment : string -> Exp -> Stmt
| atoi_assignment : string -> string -> Stmt
| itoa_assignment : string -> nat -> Stmt
| sequence : Stmt -> Stmt -> Stmt
| ifthenelse : Exp -> Stmt -> Stmt -> Stmt
| ifthen : Exp -> Stmt -> Stmt
| while : Exp -> Stmt -> Stmt.

Notation "'fie' A" := (var_declare A) (at level 90).
Notation "X ::= A" := (assignment X A) (at level 90).
Notation "S ;; S'" := (sequence S S') (at level 93, right associativity).
Notation "A <{ B }>" := (function_declare A B) (at level 90).
Notation "'class' A" := (object_declare A) (at level 90).
Notation "X ::= 'atoi'( Y )" := (atoi_assignment X Y) (at level 94).
Notation "X ::= 'itoa'( Y )" := (itoa_assignment X Y) (at level 94).

Example ex1 :=
fie "c" ;;
"someFunction" takes [ "c" ] <{ "c" ::= 10 }>;;
fie "b" ;;
"c" ::= 10 ;;
"b" ::= "ana"
.
Compute ex1.

Example ex2 :=
fie "obiect" ;;
class object ("student") (["age","name","grade"]) ("getGrade" takes []);;
ifthenelse ("a"=='"x") ("a"::=10) ("a"::=1);;
while ("a">='0) ("a"::="a"-'1);;
fie "b";;
ifthen (1) ("b"::=10+'5)
.
Compute ex2.

(*SEMANTICS*)

(*atoi - itoa*)

Open Scope char_scope.

Definition digitToNat (c : ascii) : nat :=
  match c with
    | "0" => 0
    | "1" => 1
    | "2" => 2
    | "3" => 3
    | "4" => 4
    | "5" => 5
    | "6" => 6
    | "7" => 7
    | "8" => 8
    | "9" => 9
    | _ => 0
  end.

Open Scope string_scope.

Fixpoint readNatAux (s : string) (acc : nat) : nat :=
  match s with
    | "" => acc
    | String c s' =>
      match digitToNat c with
        | n => readNatAux s' (10 * acc + n)
      end
  end.

Definition readNat (s : string) : nat :=
  readNatAux s 0.

Open Scope string_scope.

Definition natToDigit (n : nat) : ascii :=
  match n with
    | 0 => "0"
    | 1 => "1"
    | 2 => "2"
    | 3 => "3"
    | 4 => "4"
    | 5 => "5"
    | 6 => "6"
    | 7 => "7"
    | 8 => "8"
    | _ => "9"
  end.

Fixpoint writeNatAux (time n : nat) (acc : string) : string :=
  let acc' := String (natToDigit (Nat.modulo n 10)) acc in
  match time with
    | 0 => acc'
    | S time' =>
      match (Nat.div n 10) with
        | 0 => acc'
        | n' => writeNatAux time' n' acc'
      end
  end.

Definition writeNat (n : nat) : string :=
  writeNatAux n n "".

Notation "atoi( X )" := (readNat X) (at level 90).

Notation "itoa( X )" := (writeNat X) (at level 90).

Compute atoi("10").

Compute itoa(10).

(*Memory*)

Definition nval (n : Value) : nat :=
match n with
| nat_val n => n
| _ => 0
end.

Definition bval (b : Value) : bool :=
match b with
| bool_val b => b
| _ => false
end.

Definition sval (s : Value) : string :=
match s with
| string_val s => s
| _ => ""
end.

Definition State : Variabila := fun n => 0.
Definition Values : Adresa := fun n => undef.

Definition updateState (state : Variabila) (s : string) (value : nat) : Variabila :=
  fun s' =>
    if (string_eq_dec s' s)
    then value
    else (state s').

Definition updateValues (values : Adresa) (n : nat) (value : Value) : Adresa :=
  fun n' =>
    if (Nat.eqb n' n)
    then value
    else (values n').

Definition is_declared (s : string) (state : Variabila) (values : Adresa) :=
  if (Value_beq (values (state s)) undef)
  then false
  else true.

Compute (updateState State "valAdresa" 0).
Compute (updateState State "memPointer" 1)("memPointer").

(*Exp*)

Fixpoint expEval (e : Exp) (s : Variabila) (v : Adresa) : nat :=
match e with
| avar var => nval(v(s var))
| anum n' => n'
| btrue => 1
| bfalse => 0
| aplus e1 e2 => (expEval e1 s v) + (expEval e2 s v)
| aminus e1 e2 => (expEval e1 s v) - (expEval e2 s v)
| amul e1 e2 => (expEval e1 s v) * (expEval e2 s v)
| adiv e1 e2 => Nat.div (expEval e1 s v) (expEval e2 s v)
| bless b1 b2 => if (Nat.ltb (expEval b1 s v) (expEval b2 s v)) then 1 else 0
| bgreat b1 b2=> if (Nat.ltb (expEval b2 s v) (expEval b1 s v)) then 1 else 0
| bleq b1 b2=> if (Nat.leb (expEval b1 s v) (expEval b2 s v)) then 1 else 0
| bgeq b1 b2=> if (Nat.leb (expEval b2 s v) (expEval b1 s v)) then 1 else 0
| beq b1 b2=> if (Nat.eqb (expEval b1 s v) (expEval b1 s v)) then 1 else 0
| bnot b' => match (expEval b' s v) with | 0 => 0 | _ => 1 end
| band b1 b2 => 0
| bor b1 b2=> if (Nat.eqb (expEval b1 s v) (expEval b2 s v)) then 1 else 0
end.

Definition typeOf (e : Exp) : Value :=
match e with
| anum a => nat_val a
| btrue => bool_val true
| bfalse => bool_val false
| avar s => string_val s
| _ => undef
end.

(*functions*)

Fixpoint getElement (l : list_variable) (s : Variabila) (v : Adresa) {struct l} : Adresa :=
match l with
| nil => v
| str :: nil => (updateValues v ((updateState (updateState s "memPointer" (s("memPointer")+1)) str (s("memPointer"))) "memPointer") undef)
| str :: l => getElement l s v
end.

Definition functionEval (f : Function) (s : Variabila) (v : Adresa) : Adresa :=
match f with
|function str l => getElement l s v
end.

(*objects*)

Definition objEval (obj : Object) (s : Variabila) (v : Adresa) : Adresa :=
match obj with
|object name list func => getElement list s (functionEval func s v) (*defines both method variables and member variables*)
end.

(*statements*)

Fixpoint eval (pgm : Stmt) (s : Variabila) (v : Adresa) (gas : nat) : Adresa :=
match gas with
| 0 => v
| S gas' => match pgm with
            | var_declare var => if(is_declared var s v)
                                 then v
                                 else updateValues v
                                      ((updateState (updateState s "memPointer" (s("memPointer")+1)) var (s("memPointer"))) "memPointer")
                                      (nat_val 1) (*This should be undef - changed to nat_val 1 for the sake of explaining the issue*)
            | object_declare obj => objEval obj s v
            | function_declare funct stmt => eval stmt s (functionEval funct s v) gas'
            | assignment var exp => if(is_declared var s v)  (*Might need further testing on what "x"::="y" returns, should add nval,bval and sval*)
                                    then updateValues v (s var) (typeOf exp)
                                    else v
            | atoi_assignment var str => if(is_declared var s v)
                                         then updateValues v (s var) (nat_val (atoi(str)))
                                         else v
            | itoa_assignment var n => if(is_declared var s v)
                                       then updateValues v (s var) (string_val (itoa(n)))
                                       else v
            | sequence stmt1 stmt2 => eval stmt2 s (eval stmt1 s v gas') gas'
            | ifthen cond stmt =>  if (expEval cond s v)
                                   then eval stmt s v gas'
                                   else v
            | ifthenelse cond stmt stmt' => if (expEval cond s v)
                                            then eval stmt s v gas'
                                            else eval stmt s v gas'
            | while cond stmt' => if (expEval cond s v)
                               then eval (stmt' ;; while cond stmt') s v gas'
                               else v
            end
end.

(*Missing class method/member accesing such as object.member*)

Example ex3 :=
fie "c".

Definition Values' := (eval ex3 State Values 100).

Compute State "memPointer". (*Although memPointer is still 0-*)
Compute Values' 1. (*-the value after one definition works fine. The issue is that eval does not return the updated eval.
                    Either a different env that encapsulates both of them should exist, or a way to return or keep tracks of both State and Values
                    one way of fixing this would be through "Variable ValAdresa : nat." - global counter, unrelated to any state/value - this would fix
                    the value problem, but we would still need a way to keep track of state for variable names*)

Example ex4 :=
class object ("student") (["age","name","grade"]) ("getGrade" takes []);;
ifthenelse ("a"=='"x") ("a"::=10) ("a"::=1);;
while ("a">='0) ("a"::="a"-'1);;
fie "b";;
ifthen (1) (itoa_assignment "b" 20) ;;
ifthenelse("b"=='15) ("a"::=2) ("a"::=3)
.

Example ex5 :=
fie "c" ;;
"someFunction" takes [ "d" ] <{ "d" ::= 10 }>;; (*Scope is not simulated, thus function variables should have different names for desired result*)
fie "b" ;;
atoi_assignment "c" "10" ;;
"b" ::= "ana"
.

Definition Values'' := eval ex4 State Values 20.


Require Import CoRN.algebra.RSetoid.
Require Import CoRN.metric2.Metric.
Require Import CoRN.metric2.UniformContinuity.
Require Import
  Coq.QArith.Qabs
  CoRN.model.metric2.Qmetric
  CoRN.model.totalorder.QposMinMax
  CoRN.model.totalorder.QMinMax
  CoRN.reals.fast.CRabs
  CoRN.reals.faster.ARArith.

(** Absolute Value *)

Section ARabs.
Context `{AppRationals AQ}.

Definition AQabs : AQ -> AQ
  := fun q => proj1_sig (abs_sig q).

Lemma aq_abs : forall q:AQ, '(proj1_sig (abs_sig q)) == Qabs ('q).
Proof.
  intro q. destruct (abs_sig q). unfold proj1_sig.
  destruct a.
  destruct aq_ring_morphism.
  destruct aq_order_embed.
  destruct semiringmor_plus_mor.
  destruct monmor_sgmor.
  destruct sgmor_setmor.
  assert (AQtoQ 0 == 0) by (apply rings.preserves_0).
  destruct (Qlt_le_dec ('q) 0).
  - rewrite Qabs_neg. rewrite <- aq_opp.
    apply sm_proper, H6.
    apply order_embedding_reflecting.
    rewrite H7.
    apply Qlt_le_weak, q0.
    apply Qlt_le_weak, q0.
  - rewrite Qabs_pos. 2: exact q0.
    apply sm_proper, H5.
    apply order_embedding_reflecting.
    rewrite H7. exact q0.
Qed.

Lemma AQabs_triangle_reverse
  : forall a b : AQ, Qabs (' AQabs a - ' AQabs b) <= Qabs ('a - 'b).
Proof.
  intros a b. unfold AQabs.
  rewrite aq_abs, aq_abs.
  apply Qabs_case.
  - intros _. apply Qabs_triangle_reverse.
  - intros.
    apply (Qle_trans _ (Qabs ('b) - Qabs ('a))).
    ring_simplify; apply Qle_refl.
    rewrite <- (Qabs_opp (' a - ' b)).
    setoid_replace (- (' a - ' b))%Q with (' b - ' a)%Q.
    apply Qabs_triangle_reverse.
    unfold equiv, stdlib_rationals.Q_eq.
    ring.
Qed.

Lemma AQabs_uc_prf
  : is_UniformlyContinuousFunction
      (AQabs : AQ_as_MetricSpace -> AQ_as_MetricSpace) Qpos2QposInf.
Proof.
  intros d a b Hab.
  simpl in Hab.
  unfold Qball in Hab.
  rewrite <- AbsSmall_Qabs in Hab.
  simpl. unfold Qball.
  rewrite <- AbsSmall_Qabs.
  exact (Qle_trans _ _ _ (AQabs_triangle_reverse _ _) Hab).
Qed.

Local Open Scope uc_scope.

Definition AQabs_uc : AQ_as_MetricSpace --> AQ_as_MetricSpace
  := Build_UniformlyContinuousFunction AQabs_uc_prf.

Definition ARabs : AR --> AR := Cmap AQPrelengthSpace AQabs_uc.

(* The approximations of the absolute value are what we expect. *)
Lemma ARabs_approx : forall (e:Qpos) (x:AR),
    approximate (ARabs x) e = AQabs (approximate x e).
Proof.
  reflexivity.
Qed.

Lemma ARtoCR_preserves_abs
  : forall x : AR, cast AR CR (ARabs x) = CRabs (cast AR CR x).
Proof.
  intro x.
  unfold cast, ARtoCR, ARtoCR_uc.
  unfold MetricMorphisms.Eembed.
  unfold ARabs, CRabs.
  rewrite <- (fast_MonadLaw2
               AQPrelengthSpace
               (MetricMorphisms.EPrelengthSpace QPrelengthSpace (cast AQ Q_as_MetricSpace))
               (MetricMorphisms.metric_embed_uc (cast AQ Q_as_MetricSpace))
               AQabs_uc).
  rewrite <- (fast_MonadLaw2
               (MetricMorphisms.EPrelengthSpace QPrelengthSpace (cast AQ Q_as_MetricSpace))
               QPrelengthSpace Qabs_uc
               (MetricMorphisms.metric_embed_uc (cast AQ Q_as_MetricSpace)) x).
  apply Cmap_wd. 2: reflexivity.
  apply ucEq_equiv.
  intro d. simpl. apply Qball_0, aq_abs.
Qed.

Lemma ARabs_AbsSmall : forall a b : AR, ARle (ARabs b) a <-> (ARle (-a) b /\ ARle b a).
Proof.
  intros a b. pose proof CRabs_AbsSmall. split.
  - intro H6.
    assert ((CRabs (cast AR CR b) <= (cast AR CR a))%CR).
    { apply ARtoCR_preserves_le in H6.
      rewrite ARtoCR_preserves_abs in H6. exact H6. }
    specialize (H5 ('a) ('b)) as [H5 _].
    specialize (H5 H7).
    split; apply ARtoCR_preserves_le.
    rewrite ARtoCR_preserves_opp.
    apply H5. apply H5.
  - intros [H6 H7].
    apply (ARtoCR_preserves_le (-a) b) in H6.
    apply (ARtoCR_preserves_le b a) in H7.
    specialize (H5 ('a) ('b)) as [_ H5].
    apply ARtoCR_preserves_le.
    rewrite ARtoCR_preserves_abs. apply H5.
    split. rewrite <- ARtoCR_preserves_opp.
    exact H6. exact H7.
Qed.

Lemma ARle_abs : forall x:AR, ARle x (ARabs x).
Proof.
  intro x. apply ARtoCR_preserves_le.
  rewrite ARtoCR_preserves_abs.
  apply CRle_abs.
Qed.

Lemma ARabs_opp : forall x:AR, ARabs (-x) = ARabs x.
Proof.
  intro x.
  apply (injective ARtoCR).
  pose proof (ARtoCR_preserves_abs (-x)).
  unfold cast in H5. rewrite H5. clear H5.
  pose proof (ARtoCR_preserves_abs x).
  unfold cast in H5. rewrite H5.
  rewrite (ARtoCR_preserves_opp x).
  apply CRabs_opp.
Qed.

Lemma ARabs_pos : forall x:AR, ARle 0 x -> ARabs x = x.
Proof.
  intros x xpos.
  apply (injective ARtoCR).
  pose proof (ARtoCR_preserves_abs x).
  unfold cast in H5. rewrite H5. clear H5.
  apply CRabs_pos.
  rewrite <- ARtoCR_preserves_0.
  apply (ARtoCR_preserves_le AR0 x), xpos.
Qed.

Lemma ARabs_neg : forall x:AR, ARle x 0 -> ARabs x = -x.
Proof.
  intros x xneg.
  apply (injective ARtoCR).
  pose proof (ARtoCR_preserves_abs x).
  unfold cast in H5. rewrite H5. clear H5.
  rewrite CRabs_neg.
  symmetry. apply ARtoCR_preserves_opp.
  rewrite <- ARtoCR_preserves_0.
  apply (ARtoCR_preserves_le x AR0), xneg.
Qed.

End ARabs.