import PackageCalculus.Extensions.PackageFormula.Definition
import PackageCalculus.Extensions.Conflict.Reduction.Definition

/-! # Package-formula extension: reduction

Translates a package-level Boolean formula into core dependencies by walking
the formula tree and introducing fresh names for each conjunct, disjunct, and
negation. -/

namespace PackageCalculus.PkgFormula

variable {N : Type*} {V : Type*}

inductive PFName (N V : Type*) where
  | orig : N → PFName N V
  | disjunct : Formula N V → Formula N V → PFName N V
  | negDep : N → Finset V → PFName N V
  deriving DecidableEq

variable {N' : Type*} [DecidableEq N'] {V' : Type*} [DecidableEq V']
variable [hpn : HasPFNames N V N'] [hpv : Conflict.HasConflictVersions V V']

def embedPkg (p : Package N V) : Package N' V' :=
  (hpn.origN p.1, hpv.origV p.2)

def embedVS (vs : Finset V) : Finset V' :=
  vs.map hpv.origV

/-- Negation triples so De Morgan transformations strictly decrease weight. -/
def Formula.weight : Formula N V → Nat
  | .dep _ _ => 0
  | .conj ψ₁ ψ₂ => ψ₁.weight + ψ₂.weight + 2
  | .disj ψ₁ ψ₂ => ψ₁.weight + ψ₂.weight + 2
  | .neg ψ => 3 * (ψ.weight + 1)

/-! ### NNF atoms

The encoding erases conjunction structure (a conjunction's encoding is the
union of its conjuncts') and drives negations inward by De Morgan / double
negation. What it preserves of a formula is its set of *NNF atoms*: positive
literals, negative literals, and whole disjunctions (whose synthetic names
carry their subformulas verbatim). This is the normal form onto which the
transpiling retraction lands (`liftAtoms_pfDeps`). -/

/-- An NNF atom: a positive literal, a negative literal, or a disjunction
kept whole. -/
inductive Atom (N V : Type*) where
  | pos : N → Finset V → Atom N V
  | neg : N → Finset V → Atom N V
  | disj : Formula N V → Formula N V → Atom N V
  deriving DecidableEq

/-- The formula an atom denotes. -/
def Atom.toFormula : Atom N V → Formula N V
  | .pos n vs => .dep n vs
  | .neg n vs => .neg (.dep n vs)
  | .disj ψ₁ ψ₂ => .disj ψ₁ ψ₂

variable [DecidableEq N] [DecidableEq V] in
/-- The NNF atoms of a formula: flatten conjunctions and push negations
inward, mirroring the recursion of `encodeNNF`. -/
def atoms : Formula N V → Finset (Atom N V)
  | .dep n vs => {.pos n vs}
  | .conj ψ_L ψ_R => atoms ψ_L ∪ atoms ψ_R
  | .disj ψ_L ψ_R => {.disj ψ_L ψ_R}
  | .neg (.dep n vs) => {.neg n vs}
  | .neg (.conj ψ_L ψ_R) => {.disj (.neg ψ_L) (.neg ψ_R)}
  | .neg (.disj ψ_L ψ_R) => atoms (.neg ψ_L) ∪ atoms (.neg ψ_R)
  | .neg (.neg ψ) => atoms ψ
termination_by ψ => ψ.weight
decreasing_by all_goals simp only [Formula.weight]; omega

/-- Encoding function E: negation handled by inlining De Morgan / double-negation cases. -/
def encodeNNF (p : Package N' V') :
    (ψ : Formula N V) →
    Finset (Package N' V' × N' × Finset V')
  | .dep n vs =>
    { (p, hpn.origN n, embedVS vs) }
  | .conj ψ_L ψ_R =>
    encodeNNF p ψ_L ∪ encodeNNF p ψ_R
  | .disj ψ_L ψ_R =>
    { (p, hpn.disjunctN ψ_L ψ_R, {hpv.zeroV, hpv.oneV}) } ∪
    encodeNNF (hpn.disjunctN ψ_L ψ_R, hpv.zeroV) ψ_L ∪
    encodeNNF (hpn.disjunctN ψ_L ψ_R, hpv.oneV) ψ_R
  | .neg (.dep n vs) =>
    { (p, hpn.syntheticN n vs, ({hpv.oneV} : Finset V')) } ∪
    vs.image (fun u => ((hpn.origN n, hpv.origV u),
      hpn.syntheticN n vs, ({hpv.zeroV} : Finset V')))
  | .neg (.conj ψ_L ψ_R) =>
    encodeNNF p (.disj (.neg ψ_L) (.neg ψ_R))
  | .neg (.disj ψ_L ψ_R) =>
    encodeNNF p (.conj (.neg ψ_L) (.neg ψ_R))
  | .neg (.neg ψ) =>
    encodeNNF p ψ
termination_by ψ => ψ.weight
decreasing_by all_goals simp only [Formula.weight]; omega

def encode (p : Package N' V')
    (ψ : Formula N V) :
    Finset (Package N' V' × N' × Finset V') :=
  encodeNNF p ψ

/-- Collect all witness packages generated by encoding a formula. -/
def witnessPackages (p : Package N' V') :
    (ψ : Formula N V) → Finset (Package N' V')
  | .dep _ _ => ∅
  | .conj ψ_L ψ_R =>
    witnessPackages p ψ_L ∪ witnessPackages p ψ_R
  | .disj ψ_L ψ_R =>
    {(hpn.disjunctN ψ_L ψ_R, hpv.zeroV), (hpn.disjunctN ψ_L ψ_R, hpv.oneV)} ∪
    witnessPackages (hpn.disjunctN ψ_L ψ_R, hpv.zeroV) ψ_L ∪
    witnessPackages (hpn.disjunctN ψ_L ψ_R, hpv.oneV) ψ_R
  | .neg (.dep n vs) =>
    {(hpn.syntheticN n vs, hpv.zeroV), (hpn.syntheticN n vs, hpv.oneV)}
  | .neg (.conj ψ_L ψ_R) =>
    witnessPackages p (.disj (.neg ψ_L) (.neg ψ_R))
  | .neg (.disj ψ_L ψ_R) =>
    witnessPackages p (.conj (.neg ψ_L) (.neg ψ_R))
  | .neg (.neg ψ) =>
    witnessPackages p ψ
termination_by ψ => ψ.weight
decreasing_by all_goals simp only [Formula.weight]; omega

variable [DecidableEq N] [DecidableEq V]

def pfReal (R_Ψ : Real N V) (Δ_Ψ : PFDepRel N V) : Real N' V' :=
  R_Ψ.image embedPkg ∪
  Δ_Ψ.biUnion (fun ⟨p, ψ⟩ => witnessPackages (embedPkg p) ψ)

def pfDeps (Δ_Ψ : PFDepRel N V) : DepRel N' V' :=
  Δ_Ψ.biUnion (fun ⟨p, ψ⟩ => encode (embedPkg p) ψ)

end PackageCalculus.PkgFormula

namespace PackageCalculus

open Function

variable {N V : Type*}

instance : PkgFormula.HasPFNames N V (PkgFormula.PFName N V) where
  origN := ⟨PkgFormula.PFName.orig, fun _ _ h => PkgFormula.PFName.orig.inj h⟩
  syntheticN := PkgFormula.PFName.negDep
  syntheticN_injective := by
    intro a₁ a₂ b₁ b₂ h
    exact ⟨PkgFormula.PFName.negDep.inj h |>.1,
           PkgFormula.PFName.negDep.inj h |>.2⟩
  origN_ne_syntheticN := fun _ _ _ => nofun
  syntheticN_ne_origN := fun _ _ _ => nofun
  tryOrigN := fun
    | .orig n => some n
    | _ => none
  tryOrigN_origN := fun _ => rfl
  tryOrigN_some := fun n' n h => by
    cases n' with
    | orig m => simp at h; subst h; rfl
    | _ => simp at h
  trySyntheticN := fun
    | .negDep n vs => some (n, vs)
    | _ => none
  trySyntheticN_syntheticN := fun _ _ => rfl
  trySyntheticN_some := fun n' p h => by
    cases n' with
    | negDep n vs => simp at h; obtain ⟨rfl, rfl⟩ := h; rfl
    | _ => simp at h
  disjunctN := PkgFormula.PFName.disjunct
  disjunctN_injective := by
    intro a₁ a₂ b₁ b₂ h
    exact ⟨PkgFormula.PFName.disjunct.inj h |>.1,
           PkgFormula.PFName.disjunct.inj h |>.2⟩
  origN_ne_disjunctN := fun _ _ _ => nofun
  disjunctN_ne_origN := fun _ _ _ => nofun
  disjunctN_ne_syntheticN := fun _ _ _ _ => nofun
  syntheticN_ne_disjunctN := fun _ _ _ _ => nofun
  tryDisjunctN := fun
    | .disjunct ψ₁ ψ₂ => some (ψ₁, ψ₂)
    | _ => none
  tryDisjunctN_disjunctN := fun _ _ => rfl
  tryDisjunctN_some := fun n' q h => by
    cases n' with
    | disjunct a b => simp at h; obtain ⟨rfl, rfl⟩ := h; rfl
    | _ => simp at h

end PackageCalculus