import Choreography.Projection.Trans.Core

/-! # Trans Participation

First-branch and all-branch participation predicates and guardedness of projection
under participation.
-/

/-
The Problem. The `trans` projection function follows specific paths through branches,
but the standard `participates` predicate is existential over all branches. This
mismatch breaks contractiveness proofs.

Solution Structure. We define path-aware participation predicates:
1. `participatesFirstBranch`: follows the same path as `trans` (first branch for non-direct)
2. `participatesAllBranches`: requires participation in ALL branches (stronger)
3. Implication lemmas connecting these to standard participation
The key insight is matching the predicate's branch traversal to `trans`'s behavior.
-/

namespace Choreography.Projection.Trans
open SessionTypes.GlobalType
open SessionTypes.LocalTypeR
/-! ## First-Branch Participation

The theorem `trans_isContractive_of_participant` with existential `participates` is FALSE.
Counterexample: role participates in branch 2, but `trans` follows branch 1 where they
don't participate, producing `.mu "x" (.var "x")` which is not contractive.

The correct approach is to use a predicate that tracks the same path as `trans`:
`participatesFirstBranch` follows the first branch for non-direct participants,
matching exactly what `trans` does. -/

/-! ## Participation Predicates -/

mutual
  /-- Participation check following the same path as `trans`.
      For non-direct participants, follows the first branch like `trans` does. -/
  def participatesFirstBranch (role : String) : GlobalType → Bool
    | .end => false
    | .var _ => false
    | .mu _ body => participatesFirstBranch role body
    | .comm sender receiver branches =>
        SessionTypes.Participation.is_participant role sender receiver ||
        match branches with
        | [] => false
        | (_, cont) :: _ => participatesFirstBranch role cont
    | .delegate p q _ _ cont =>
        SessionTypes.Participation.is_participant role p q ||
        participatesFirstBranch role cont

  /-- Branches version for mutual recursion (unused but needed for termination). -/
  def participatesFirstBranchBranches (role : String) : List (Label × GlobalType) → Bool
    | [] => false
    | (_, cont) :: _ => participatesFirstBranch role cont
end

/-! ## First-Branch Implies Standard Participation -/

/-! ## First-Branch Comm Helper -/

private theorem participates_first_branch_imp_participates_comm
    (role sender receiver : String) (branches : List (Label × GlobalType))
    (h : participatesFirstBranch role (.comm sender receiver branches) = true)
    (ih : ∀ g : GlobalType, participatesFirstBranch role g = true →
      SessionTypes.Participation.participates role g = true) :
    SessionTypes.Participation.participates role (.comm sender receiver branches) = true := by
  unfold participatesFirstBranch at h
  unfold SessionTypes.Participation.participates
  cases hpart : SessionTypes.Participation.is_participant role sender receiver with
  | true =>
      simp
  | false =>
      match hbranches : branches with
      | [] =>
          simp [hpart] at h
      | (label, cont) :: rest =>
          have hcont : participatesFirstBranch role cont = true := by
            simp [hpart] at h
            exact h
          have hcont_part : SessionTypes.Participation.participates role cont = true :=
            ih cont hcont
          simp [SessionTypes.Participation.participatesBranches, hcont_part]

/-! ## First-Branch Delegate Helper -/

private theorem participates_first_branch_imp_participates_delegate
    (role p q r : String) (sid : Nat) (cont : GlobalType)
    (h : participatesFirstBranch role (.delegate p q sid r cont) = true)
    (ih : ∀ g : GlobalType, participatesFirstBranch role g = true →
      SessionTypes.Participation.participates role g = true) :
    SessionTypes.Participation.participates role (.delegate p q sid r cont) = true := by
  unfold participatesFirstBranch at h
  unfold SessionTypes.Participation.participates
  cases hpart : SessionTypes.Participation.is_participant role p q with
  | true =>
      simp
  | false =>
      simp [hpart] at h
      have hcont_part : SessionTypes.Participation.participates role cont = true :=
        ih cont h
      simp [hcont_part]

/-! ## First-Branch Main Theorem -/

/-- First-branch participation implies standard participation. -/
theorem participates_first_branch_imp_participates (g : GlobalType) (role : String) :
    participatesFirstBranch role g = true → SessionTypes.Participation.participates role g = true := by
  intro h
  match g with
  | .end =>
      simp [participatesFirstBranch] at h
  | .var _ =>
      simp [participatesFirstBranch] at h
  | .mu t body =>
      simpa [participatesFirstBranch, SessionTypes.Participation.participates] using
        (participates_first_branch_imp_participates body role h)
  | .comm sender receiver branches =>
      unfold participatesFirstBranch at h
      unfold SessionTypes.Participation.participates
      cases hpart : SessionTypes.Participation.is_participant role sender receiver with
      | true =>
          simp
      | false =>
          cases hbranches : branches with
          | nil =>
              simp [hpart, hbranches] at h
          | cons head tail =>
              cases head with
              | mk label cont =>
                  have hcont : participatesFirstBranch role cont = true := by
                    simpa [hpart, hbranches] using h
                  have hcont_part : SessionTypes.Participation.participates role cont = true :=
                    participates_first_branch_imp_participates cont role hcont
                  have hbranches_part :
                      SessionTypes.Participation.participatesBranches role ((label, cont) :: tail) = true := by
                    simp [SessionTypes.Participation.participatesBranches, hcont_part]
                  simpa [hpart] using hbranches_part
  | .delegate p q sid r cont =>
      unfold participatesFirstBranch at h
      unfold SessionTypes.Participation.participates
      cases hpart : SessionTypes.Participation.is_participant role p q with
      | true =>
          simp
      | false =>
          have hcont_part : SessionTypes.Participation.participates role cont = true :=
            participates_first_branch_imp_participates cont role (by simpa [hpart] using h)
          exact hcont_part
termination_by sizeOf g
decreasing_by
  all_goals
    subst_vars
    simp_wf
    omega

mutual
  /-- Participation that continues through ALL branch continuations, not just first branch.
      This is needed for contractiveness: when role is sender/receiver, we need
      participation to continue in ALL continuations, not just the outer level. -/
  def participatesAllBranches (role : String) : GlobalType → Bool
/- ## Structured Block 1 -/
    | .end => false
    | .var _ => false
    | .mu _ body => participatesAllBranches role body
    | .comm sender receiver branches =>
        SessionTypes.Participation.is_participant role sender receiver &&
        participatesAllBranchesList role branches ||
        -- OR: not direct participant but participates in first branch
        (!(SessionTypes.Participation.is_participant role sender receiver) &&
         match branches with
         | [] => false
         | (_, cont) :: _ => participatesAllBranches role cont)
    | .delegate p q _ _ cont =>
        SessionTypes.Participation.is_participant role p q &&
        participatesAllBranches role cont ||
        (!(SessionTypes.Participation.is_participant role p q) &&
         participatesAllBranches role cont)

  /-- Helper for branch list participation. -/
  def participatesAllBranchesList (role : String) : List (Label × GlobalType) → Bool
    | [] => true
    | (_, cont) :: rest =>
        participatesAllBranches role cont && participatesAllBranchesList role rest
end

/-! ## Guardedness Primitives -/

private theorem is_guarded_send (p : String) (bs : List BranchR) (v : String) :
    (LocalTypeR.send p bs).isGuarded v = true := by
  simp [LocalTypeR.isGuarded]

private theorem is_guarded_recv (p : String) (bs : List BranchR) (v : String) :
    (LocalTypeR.recv p bs).isGuarded v = true := by
  simp [LocalTypeR.isGuarded]

private theorem is_guarded_end (v : String) : (LocalTypeR.end).isGuarded v = true := by
  simp [LocalTypeR.isGuarded]

/-! ## Comm Participant Guardedness -/

/-- Helper: trans result is guarded when role participates via first-branch path.

    The key insight: participation means we eventually reach a .send/.recv which
    guards any variable. The guardedness propagates up through the structure.

    Note: This does NOT require closedness - participation alone is sufficient. -/
private theorem trans_is_guarded_of_participates_first_branch_comm_participant
    (sender receiver : String) (branches : List (Label × GlobalType))
    (v : String) (role : String)
    (hpart : SessionTypes.Participation.is_participant role sender receiver = true) :
    (trans (.comm sender receiver branches) role).isGuarded v = true := by
  -- Direct participant: sender or receiver.
  unfold SessionTypes.Participation.is_participant at hpart
  cases hrole_s : role == sender with
  | true =>
      have heq : role = sender := beq_iff_eq.mp hrole_s
      rw [trans_comm_sender sender receiver role branches heq]
      exact is_guarded_send receiver (transBranches branches role) v
  | false =>
      simp only [hrole_s, Bool.false_or] at hpart
      have heq : role = receiver := beq_iff_eq.mp hpart
      have hne : role ≠ sender := by
        intro heq'
        rw [heq'] at hrole_s
        simp at hrole_s
      rw [trans_comm_receiver sender receiver role branches heq hne]
      exact is_guarded_recv sender (transBranches branches role) v

/-! ## Main Guardedness Theorem -/

/-- If a role participates along the first-branch path, then its projection is guarded. -/
theorem trans_is_guarded_of_participates_first_branch
    (g : GlobalType) (v : String) (role : String)
    (hpart : participatesFirstBranch role g = true) :
    (trans g role).isGuarded v = true := by
  match g with
  | .end =>
      simp [participatesFirstBranch] at hpart
  | .var _ =>
      simp [participatesFirstBranch] at hpart
  -- trans_is_guarded_of_participates_first_branch: Mu Case
  | .mu t body =>
      -- Mu case: follow the body and propagate guardedness.
      unfold participatesFirstBranch at hpart
      simp only [trans]
      by_cases hguard : (trans body role).isGuarded t
      · simp only [hguard, ↓reduceIte, LocalTypeR.isGuarded]
        by_cases hvt : v = t
        · simp [hvt]
        · have hvne : (v == t) = false := beq_eq_false_iff_ne.mpr hvt
          simp only [hvne, Bool.false_eq_true, ↓reduceIte]
          exact trans_is_guarded_of_participates_first_branch body v role hpart
      · simp [hguard, Bool.false_eq_true, ↓reduceIte, LocalTypeR.isGuarded]
  -- trans_is_guarded_of_participates_first_branch: Comm Case
  | .comm sender receiver branches =>
      -- Comm case: direct participant or follow the first branch.
      cases hpart_direct : SessionTypes.Participation.is_participant role sender receiver with
      | true =>
          exact trans_is_guarded_of_participates_first_branch_comm_participant sender receiver branches v role
            hpart_direct
      | false =>
          unfold SessionTypes.Participation.is_participant at hpart_direct
          have hne_s : role ≠ sender := by
            intro heq
            rw [heq] at hpart_direct
            simp at hpart_direct
          have hne_r : role ≠ receiver := by
            intro heq
            rw [heq] at hpart_direct
            simp at hpart_direct
          rw [trans_comm_other sender receiver role branches hne_s hne_r]
          cases branches with
          | nil =>
              exact is_guarded_end v
          | cons head tail =>
              cases head with
              | mk label cont =>
                  have hcont' :
                      SessionTypes.Participation.is_participant role sender receiver = true ∨
                        participatesFirstBranch role cont = true := by
                    have hpart' := hpart
                    simpa [participatesFirstBranch, Bool.or_eq_true] using hpart'
                  have hcont : participatesFirstBranch role cont = true := by
                    cases hcont' with
                    | inl hleft =>
/- ## Structured Block 2 -/
                        have : False := by
                          unfold SessionTypes.Participation.is_participant at hleft
                          simp [hpart_direct] at hleft
                        exact this.elim
                    | inr hright =>
                        exact hright
                  exact trans_is_guarded_of_participates_first_branch cont v role hcont
  -- trans_is_guarded_of_participates_first_branch: Delegate Case
  | .delegate p q sid r cont =>
      -- Delegate case: direct participant or follow the continuation.
      cases hpart_direct : SessionTypes.Participation.is_participant role p q with
      | true =>
          -- Direct participant: projection is send or recv, which are guarded
          simp only [trans]
          by_cases hp : role == p
          · simp only [hp, ↓reduceIte, LocalTypeR.isGuarded]
          · by_cases hq : role == q
            · simp only [hp, Bool.false_eq_true, ↓reduceIte, hq, LocalTypeR.isGuarded]
            · -- Contradiction: is_participant but neither p nor q
              unfold SessionTypes.Participation.is_participant at hpart_direct
              simp only [Bool.or_eq_true] at hpart_direct
              cases hpart_direct with
              | inl heqp => exact absurd heqp hp
              | inr heqq => exact absurd heqq hq
      | false =>
          -- Not direct participant: follow continuation
          simp only [trans]
          have hp : (role == p) = false := by
            by_contra h
            simp only [Bool.not_eq_false] at h
            unfold SessionTypes.Participation.is_participant at hpart_direct
            simp [h] at hpart_direct
          have hq : (role == q) = false := by
            by_contra h
            simp only [Bool.not_eq_false] at h
            unfold SessionTypes.Participation.is_participant at hpart_direct
            simp [h] at hpart_direct
          simp only [hp, Bool.false_eq_true, ↓reduceIte, hq]
          have hcont : participatesFirstBranch role cont = true := by
            simp [participatesFirstBranch, hpart_direct] at hpart
            exact hpart
          exact trans_is_guarded_of_participates_first_branch cont v role hcont

end Choreography.Projection.Trans
