Module FormalML.converge.finite_time

Require Import Reals Coq.Lists.List Coquelicot.Series Coquelicot.Hierarchy Coquelicot.SF_seq.
Require Import pmf_monad mdp Permutation fixed_point Finite LibUtils.
Require Import Sums Coq.Reals.ROrderedType.
Require Import micromega.Lra.
Require Import micromega.Lia.
Require Import Coq.Logic.FunctionalExtensionality.
Require Import Equivalence RelationClasses EquivDec Morphisms.
Require Import ExtLib.Structures.Monad.

Require Import ListAdd.

Set Bullet Behavior "Strict Subproofs".

Import ListNotations.

Section all_lists.

 Fixpoint all_lists_upto {A} (l:list A) (n:nat) : list (A * list A)
  := match n with
     | 0 => nil
     | 1 => map (fun x => (x,nil)) l
     | S m => (list_prod l (map (fun '(a,b) => a::b) (all_lists_upto l m)))
     end.

Lemma all_lists_upto_non_nil {A} (l:list A) n : nil <> l -> (n > 0)%nat -> nil <> all_lists_upto l n.

Lemma all_lists_upto_total {A} (l:list A) n :
  forall a' l', incl (a'::l') l -> length (a'::l') = n ->
           In (a',l') (all_lists_upto l n).

Lemma all_lists_upto_in {A} (l:list A) n x :
  In x (all_lists_upto l n) -> In (fst x) l /\ incl (snd x) l.

Lemma all_lists_length {A} (l:list A) n x:
  In x (all_lists_upto l n) -> length (fst x :: snd x) = n.

Lemma all_lists_upto_prod {A} (l:list A) n :
  let l' := all_lists_upto l n in
  equivlist (list_prod (map fst l') (map snd l')) l'.
Lemma all_lists_upto_fsts {A} (l:list A) n :
  equivlist (map fst (all_lists_upto l (S n))) l.

Lemma all_lists_upto_snds {A} (l:list A) n :
  equivlist (map snd (all_lists_upto l (S (S n)))) (map (fun '(h,t) => (h::t)) (all_lists_upto l (S n))).



End all_lists.

Section ltv_fin.
Open Scope R_scope.
Context (M : MDP) (γ : R).
Context (hγ : (0 <= γ < 1)%R).

Arguments t {_}.
Arguments reward {_}.


Definition kliesli_iter_left (Π : list(dec_rule M)) (init : M.(state)) : Pmf M.(state):=
  fold_left (fun p π => Pmf_bind p (fun s => t s (π s))) Π (ret init).

Definition kliesli_iter_left_vector (init : Pmf M.(state)) (Π : list(dec_rule M)) : Pmf M.(state):=
  fold_left (fun p π => Pmf_bind p (fun s => t s (π s))) Π init.


Lemma ff_left {A B: Type} (f: A -> B -> A) (π0: B) (πs: list B) (s: A): fold_left f πs (f s π0) = fold_left f (π0::πs) s.


Lemma kliesli_iter_left_cons0 π πs: forall p,
    kliesli_iter_left_vector (Pmf_bind p (fun s0 => t s0 (π s0))) πs
    = kliesli_iter_left_vector p (π :: πs).

Definition expt_reward_gen (p : Pmf M.(state)) (Π : list(dec_rule M)) (σ : dec_rule M)
  : R :=
  expt_value (kliesli_iter_left_vector p Π ) (step_expt_reward σ).


Definition ltv_fin (p: Pmf M.(state)) (Π : list(dec_rule M)) :=
    match Π with
    | nil => 0
    | π :: Π' => sum_f_R0' (fun n =>
                            match nth_error Π n with
                            | Some σ => γ^n*expt_reward_gen p (firstn n Π) σ
                            | None => 0
                            end) (length Π)
    end.

Lemma ltv_fin_cons_term1 p π:
  expt_reward_gen p [] π = expt_value p (fun s : state M => expt_value (t s (π s)) (reward s (π s))).


Lemma expt_reward_cons4 π πs σ vs:
  expt_reward_gen vs (π :: πs) σ =
  expt_reward_gen (Pmf_bind vs (fun s0 => t s0 (π s0))) πs σ.

Theorem ltv_fin_cons (π: dec_rule M) (πs: list(dec_rule M)) (p: Pmf M.(state)):
  ltv_fin p (π :: πs) = expt_value p (step_expt_reward π) +
                         γ*(ltv_fin (Pmf_bind (p) (fun s0 => (t s0 (π s0)))) πs).

Corollary ltv_fin_pure (π: dec_rule M) (πs: list(dec_rule M)) s:
  ltv_fin (Pmf_pure s) (π :: πs) = expt_value (Pmf_pure s) (step_expt_reward π) +
                                  γ*(ltv_fin (t s (π s)) πs).

Lemma combine_map_fst_snd {A B} (l:list (A*B)) : combine (map fst l) (map snd l) = l.

Theorem Rmax_list_prod_equivlist_split {A} (f:list A -> R) (l:list (A * list A)) :
  nil <> l ->
  equivlist (list_prod (map fst l) (map snd l)) l ->
  Max_{ l } (fun '(h,t) => f (h::t)) =
  Max_{ map fst l } (fun h => Max_{map snd l} (fun t => f (h::t)))
.
  
Theorem ltv_fin_28's (l:list (dec_rule M * list (dec_rule M))) s :
  nil <> l ->
  equivlist (list_prod (map fst l) (map snd l)) l ->
  Max_{ l } (fun '(h,t) => ltv_fin (Pmf_pure s) (h::t)) =
  Max_{ map fst l } (fun h => Max_{map snd l} (fun t => ltv_fin (Pmf_pure s) (h::t))).

Theorem Rmax_list_dec_rule_split n p :
  (n > 0)%nat ->
  let l' := all_lists_upto (@elms _ (dec_rule_finite M)) n in
  Max_{ l' } (fun '(h,t) => ltv_fin p (h::t)) =
  Max_{ map fst l' } (fun h => Max_{map snd l'} (fun t => ltv_fin p (h::t)))
.

Definition max_ltv_fin (n:nat) p :=
  let l' := all_lists_upto (@elms _ (dec_rule_finite M)) n in
  Max_{ l' } (fun '(h,t) => ltv_fin p (h::t)).

Theorem max_ltv_corec p n:
  max_ltv_fin (S (S n)) p =
  Max_{(@elms _ (dec_rule_finite M))}
      (fun h => expt_value p (step_expt_reward h) +
             γ * max_ltv_fin (S n) (Pmf_bind p (fun s0 => (t s0 (h s0))))).

End ltv_fin.