
import Mathlib
import Mathlib.Data.List.Forall2
import SessionCoTypes.Coinductive.LocalTypeC
import SessionCoTypes.Coinductive.Observable
import SessionTypes.GlobalType
import SessionCoTypes.CoinductiveRel

/- ## Structured Block 1 -/
set_option linter.dupNamespace false

/-! # Coinductive Projection

Projection from global to local types as a greatest fixed point. -/

/-
The Problem. Projection extracts a local type for a role from a global protocol.
For coinductive types, projection must be defined as a greatest fixed point to
handle infinite/recursive protocols correctly.

Solution Structure. Defines UnfoldsG/UnfoldsToG for global type unfolding,
BranchesProjRelC for branch alignment with label equality, ProjectC_step as
one-step generator allowing finite unfolding, and ProjectC as greatest fixed
point via CoinductiveRel. Proves stability under mu-unfolding.
-/

namespace SessionCoTypes.Coinductive

open SessionTypes.GlobalType
open SessionCoTypes.CoinductiveRel

-- Global type unfolding

/-- One-step global unfolding: µt.G ↦ G[µt.G/t]. -/
def UnfoldsG (g g' : GlobalType) : Prop :=
  ∃ t body, g = .mu t body ∧ g' = body.substitute t g

/-- Finite global unfolding closure. -/
def UnfoldsToG (g g' : GlobalType) : Prop :=
  Relation.ReflTransGen UnfoldsG g g'

-- Branch relations

/-- Projection relation type (global × role × coinductive local). -/
abbrev ProjRelC := GlobalType → String → LocalTypeC → Prop

/-- Branch-wise projection alignment with label equality. -/
def BranchesProjRelC (R : ProjRelC) (gbs : List (Label × GlobalType)) (role : String)
    (lbs : List (Label × LocalTypeC)) : Prop :=
  List.Forall₂ (fun gb lb => gb.1 = lb.1 ∧ R gb.2 role lb.2) gbs lbs

/-- All branches project to the same candidate (non-participant case). -/
def AllBranchesProjC (R : ProjRelC) (gbs : List (Label × GlobalType)) (role : String)
    (cand : LocalTypeC) : Prop :=
  ∀ gb ∈ gbs, R gb.2 role cand

private lemma branches_proj_rel_c_mono {R S : ProjRelC}
    (h : ∀ g role cand, R g role cand → S g role cand) :
    ∀ {gbs lbs role}, BranchesProjRelC R gbs role lbs → BranchesProjRelC S gbs role lbs := by
  intro gbs lbs role hrel
  induction hrel with
  | nil => exact List.Forall₂.nil
  | cons hhead htail ih =>
      refine List.Forall₂.cons ?_ ih
      exact ⟨hhead.1, h _ _ _ hhead.2⟩

private lemma all_branches_proj_c_mono {R S : ProjRelC}
    (h : ∀ g role cand, R g role cand → S g role cand) :
    ∀ {gbs role cand}, AllBranchesProjC R gbs role cand → AllBranchesProjC S gbs role cand := by
  intro gbs role cand hall gb hmem
  exact h _ _ _ (hall gb hmem)

-- Projection step generator

-- # Communication-Case Helper

private def ProjectC_comm (R : ProjRelC) (sender receiver : String)
    (gbs : List (Label × GlobalType)) (role : String) (l' : LocalTypeC) : Prop :=
  (role = sender ∧
      ∃ labels, head l' = .send receiver labels ∧
        BranchesProjRelC R gbs role (branchesOf l')) ∨
  (role = receiver ∧
      ∃ labels, head l' = .recv sender labels ∧
        BranchesProjRelC R gbs role (branchesOf l')) ∨
  (role ≠ sender ∧ role ≠ receiver ∧ AllBranchesProjC R gbs role l')

private lemma project_c_comm_mono {R S : ProjRelC}
    (h : ∀ g role cand, R g role cand → S g role cand) :
    ∀ {sender receiver gbs role l'},
      ProjectC_comm R sender receiver gbs role l' →
        ProjectC_comm S sender receiver gbs role l' := by
  intro sender receiver gbs role l' hcomm
  cases hcomm with
  | inl hsender =>
      rcases hsender with ⟨hrole, labels, hhead, hbr⟩
      exact Or.inl ⟨hrole, labels, hhead, branches_proj_rel_c_mono h hbr⟩
  | inr hrest =>
      cases hrest with
/- ## Structured Block 2 -/
      | inl hreceiver =>
          rcases hreceiver with ⟨hrole, labels, hhead, hbr⟩
          exact Or.inr (Or.inl ⟨hrole, labels, hhead, branches_proj_rel_c_mono h hbr⟩)
      | inr hother =>
          rcases hother with ⟨hns, hnr, hall⟩
          exact Or.inr (Or.inr ⟨hns, hnr, all_branches_proj_c_mono h hall⟩)

-- # Step Operator Definition

/-- One-step generator for coinductive projection, unfolding both sides finitely. -/
def ProjectC_step (R : ProjRelC) : ProjRelC := fun g role cand =>
  ∃ g' l', UnfoldsToG g g' ∧ UnfoldsToC cand l' ∧
    match g' with
    | .end => head l' = .end
    | .var t => head l' = .var t
    | .mu _ _ => False
    | .comm sender receiver gbs => ProjectC_comm R sender receiver gbs role l'

-- # Step Operator Monotonicity

private theorem project_c_step_mono : Monotone ProjectC_step := by
  intro R S h g role cand hrel
  rcases hrel with ⟨g', l', hg, hl, hstep⟩
  refine ⟨g', l', hg, hl, ?_⟩
  cases g' with
  | «end» =>
      simpa using hstep
  | var t =>
      simpa using hstep
  | mu t body =>
      cases hstep
  | comm sender receiver gbs =>
      cases hstep with
      | inl hsender =>
          rcases hsender with ⟨hrole, labels, hhead, hbr⟩
          exact Or.inl ⟨hrole, labels, hhead, branches_proj_rel_c_mono h hbr⟩
      | inr hrest =>
          cases hrest with
          | inl hreceiver =>
              rcases hreceiver with ⟨hrole, labels, hhead, hbr⟩
              exact Or.inr (Or.inl ⟨hrole, labels, hhead, branches_proj_rel_c_mono h hbr⟩)
          | inr hother =>
              rcases hother with ⟨hns, hnr, hall⟩
              exact Or.inr (Or.inr ⟨hns, hnr, all_branches_proj_c_mono h hall⟩)

instance : CoinductiveRel ProjRelC ProjectC_step := ⟨project_c_step_mono⟩

-- # Greatest Fixed Point


/-- Coinductive projection as the greatest fixed point of `ProjectC_step`. -/
def ProjectC : ProjRelC :=
  SessionCoTypes.CoinductiveRel.gfp (F := ProjectC_step)

-- Fixed-point helpers

theorem project_c_destruct {g : GlobalType} {role : String} {cand : LocalTypeC}
    (h : ProjectC g role cand) : ProjectC_step ProjectC g role cand := by
  have hle : ProjectC ≤ ProjectC_step ProjectC :=
    (SessionCoTypes.CoinductiveRel.unfold (F := ProjectC_step))
  exact hle _ _ _ h

theorem project_c_fold {g : GlobalType} {role : String} {cand : LocalTypeC}
    (h : ProjectC_step ProjectC g role cand) : ProjectC g role cand := by
  have hle : ProjectC_step ProjectC ≤ ProjectC :=
    (SessionCoTypes.CoinductiveRel.fold (F := ProjectC_step))
  exact hle _ _ _ h

-- Unfolding utilities

/- ## Structured Block 3 -/
private lemma unfolds_g_right_unique : Relator.RightUnique UnfoldsG := by
  intro a b c hab hac
  rcases hab with ⟨t1, body1, rfl, rfl⟩
  rcases hac with ⟨t2, body2, hmu, rfl⟩
  cases hmu
  rfl

private lemma unfolds_to_g_eq_of_nonmu {g g' : GlobalType} (h : UnfoldsToG g g')
    (hn : ∀ t body, g ≠ .mu t body) : g' = g := by
  induction h with
  | refl => rfl
  | tail h1 hstep ih =>
      cases ih
      rcases hstep with ⟨t, body, hmu, _⟩
      exact (False.elim (hn t body hmu))

private lemma unfolds_to_g_compare_nonmu {g g1 g' : GlobalType} (hg1 : UnfoldsToG g g1)
    (hg' : UnfoldsToG g g') (hn : ∀ t body, g1 ≠ .mu t body) : UnfoldsToG g' g1 := by
  have htot := Relation.ReflTransGen.total_of_right_unique (U := unfolds_g_right_unique) hg1 hg'
  cases htot with
  | inl h1 =>
      have hEq : g' = g1 := unfolds_to_g_eq_of_nonmu h1 hn
      cases hEq
      exact Relation.ReflTransGen.refl
  | inr h2 => exact h2

private lemma unfolds_c_head_mu {t u : LocalTypeC} (h : UnfoldsC t u) : ∃ x, head t = .mu x := by
  rcases h with ⟨x, f, hdest, _⟩
  refine ⟨x, ?_⟩
  simp [head, hdest]

private lemma unfolds_c_right_unique : Relator.RightUnique UnfoldsC := by
  intro a b c hab hac
  rcases hab with ⟨x, f, hdest, rfl⟩
  rcases hac with ⟨y, g, hdest', rfl⟩
  have hpair :
      (⟨LocalTypeHead.mu x, f⟩ : LocalTypeF LocalTypeC) =
        ⟨LocalTypeHead.mu y, g⟩ := by
    exact hdest.symm.trans hdest'
  cases hpair
  rfl

private lemma unfolds_to_c_eq_of_head_ne_mu {l l' : LocalTypeC} (h : UnfoldsToC l l')
    (hn : ∀ x, head l ≠ .mu x) : l' = l := by
  rcases (Relation.ReflTransGen.cases_head h) with (hEq | hstep)
  · cases hEq
    rfl
  · rcases hstep with ⟨c, hstep, _⟩
    rcases unfolds_c_head_mu hstep with ⟨x, hmu⟩
    exact (False.elim (hn x hmu))

private lemma unfolds_to_c_total {l l1 l2 : LocalTypeC} (h1 : UnfoldsToC l l1) (h2 : UnfoldsToC l l2) :
    UnfoldsToC l1 l2 ∨ UnfoldsToC l2 l1 :=
  Relation.ReflTransGen.total_of_right_unique (U := unfolds_c_right_unique) h1 h2

-- Unfolding stability

-- ProjectC Stability under Global Unfolding

theorem project_c_unfold_g {g g' : GlobalType} {role : String} {l : LocalTypeC}
    (hproj : ProjectC g role l) (hg : UnfoldsToG g g') : ProjectC g' role l := by
/- ## Structured Block 4 -/
  rcases project_c_destruct hproj with ⟨g1, l1, hg1, hl1, hmatch⟩
  cases g1 with
  | mu t body =>
      cases hmatch
  | «end» =>
      have hcomp : UnfoldsToG g' .end :=
        unfolds_to_g_compare_nonmu hg1 hg (by intro t body hmu; cases hmu)
      exact project_c_fold ⟨.end, l1, hcomp, hl1, by simpa using hmatch⟩
  | var t =>
      have hcomp : UnfoldsToG g' (.var t) :=
        unfolds_to_g_compare_nonmu hg1 hg (by intro t' body hmu; cases hmu)
      exact project_c_fold ⟨.var t, l1, hcomp, hl1, by simpa using hmatch⟩
  | comm sender receiver gbs =>
      have hcomp : UnfoldsToG g' (.comm sender receiver gbs) :=
        unfolds_to_g_compare_nonmu hg1 hg (by intro t body hmu; cases hmu)
      exact project_c_fold ⟨.comm sender receiver gbs, l1, hcomp, hl1, by simpa using hmatch⟩

-- ProjectC Stability under Local Unfolding

theorem project_c_unfold_c {g : GlobalType} {role : String} {l l' : LocalTypeC}
    (hproj : ProjectC g role l) (hl : UnfoldsToC l l') : ProjectC g role l' := by
  let R : ProjRelC := fun g role cand => ∃ l, ProjectC g role l ∧ UnfoldsToC l cand
  have hinc : ∀ g role cand, ProjectC g role cand → R g role cand := by
    intro g role cand hproj
    exact ⟨cand, hproj, Relation.ReflTransGen.refl⟩
  have hpost : ∀ g role cand, R g role cand → ProjectC_step R g role cand := by
    intro g role cand hR
    rcases hR with ⟨l0, hproj0, hstep0⟩
    rcases project_c_destruct hproj0 with ⟨g1, l1, hg1, hl1, hmatch⟩
    have hcomp := unfolds_to_c_total hl1 hstep0
    cases hcomp with
    -- Local Unfolding: Candidate Already Unfolds to Witness
    | inr h_cand_l1 =>
        refine ⟨g1, l1, hg1, h_cand_l1, ?_⟩
        cases g1 with
        | mu t body =>
            cases hmatch
        | «end» =>
            simpa using hmatch
        | var t =>
            simpa using hmatch
        | comm sender receiver gbs =>
            exact project_c_comm_mono hinc hmatch
    -- Local Unfolding: Witness Unfolds to Candidate
    | inl h_l1_cand =>
        cases g1 with
        | mu t body =>
            cases hmatch
        | «end» =>
            have hne : ∀ x, head l1 ≠ .mu x := by
              intro x hx
              simp_all
            have hEq : cand = l1 := unfolds_to_c_eq_of_head_ne_mu h_l1_cand hne
            subst cand
            exact ⟨.end, l1, hg1, Relation.ReflTransGen.refl, by simpa using hmatch⟩
/- ## Structured Block 5 -/
        | var t =>
            have hne : ∀ x, head l1 ≠ .mu x := by
              intro x hx
              simp_all
            have hEq : cand = l1 := unfolds_to_c_eq_of_head_ne_mu h_l1_cand hne
            subst cand
            exact ⟨.var t, l1, hg1, Relation.ReflTransGen.refl, by simpa using hmatch⟩
        -- Local Unfolding: Communication Subcases
        | comm sender receiver gbs =>
            cases hmatch with
            | inl hsender =>
                rcases hsender with ⟨hrole, labels, hhead, hbr⟩
                have hne : ∀ x, head l1 ≠ .mu x := by
                  intro x hx
                  simp_all
                have hEq : cand = l1 := unfolds_to_c_eq_of_head_ne_mu h_l1_cand hne
                subst cand
                refine ⟨.comm sender receiver gbs, l1, hg1, Relation.ReflTransGen.refl, ?_⟩
                exact Or.inl ⟨hrole, labels, hhead, branches_proj_rel_c_mono hinc hbr⟩
            | inr hrest =>
                cases hrest with
                | inl hreceiver =>
                    rcases hreceiver with ⟨hrole, labels, hhead, hbr⟩
                    have hne : ∀ x, head l1 ≠ .mu x := by
                      intro x hx
                      simp_all
                    have hEq : cand = l1 := unfolds_to_c_eq_of_head_ne_mu h_l1_cand hne
                    subst cand
                    refine ⟨.comm sender receiver gbs, l1, hg1, Relation.ReflTransGen.refl, ?_⟩
                    exact Or.inr (Or.inl ⟨hrole, labels, hhead, branches_proj_rel_c_mono hinc hbr⟩)
                | inr hother =>
                    rcases hother with ⟨hns, hnr, hall⟩
                    have hall' : AllBranchesProjC R gbs role cand := by
                      intro gb hmem
                      exact ⟨l1, hall gb hmem, h_l1_cand⟩
                    refine ⟨.comm sender receiver gbs, cand, hg1, Relation.ReflTransGen.refl, ?_⟩
                    exact Or.inr (Or.inr ⟨hns, hnr, hall'⟩)
  have hle : R ≤ ProjectC_step R := by
    intro g role cand hR
    exact hpost g role cand hR
  have hco : R ≤ ProjectC := by
    simpa [ProjectC] using (SessionCoTypes.CoinductiveRel.coind (F := ProjectC_step) hle)
  exact hco _ _ _ ⟨l, hproj, hl⟩

theorem project_c_mu_lift {g : GlobalType} {role : String} {x : String} {body : LocalTypeC}
    (hproj : ProjectC g role (mkMu x body)) : ProjectC g role body := by
  have hstep : UnfoldsC (mkMu x body) body := by
    refine ⟨x, (fun _ => body), ?_, rfl⟩
    simp [mkMu, PFunctor.M.dest_mk]
  exact project_c_unfold_c hproj (Relation.ReflTransGen.single hstep)

end SessionCoTypes.Coinductive
