# =============================================================================
# patch_debug_identity_origin.py
# -----------------------------------------------------------------------------
# Fixes a DANGEROUS-BUT-BENIGN diagnostic mislabel in
# prometheus_vqe_engine_penalized_latest.py.
#
# BUG:
#   debug_energy_bookkeeping() (line ~4763) prints:
#       id_coeff = float(self.qubit_h.terms.get((), 0.0))
#   But self.qubit_h is mutated IN PLACE by the sector penalties
#   (self.qubit_h += penalty_op at ~4567/4639). So after penalization the
#   printed "qubit_h identity coeff" is the PENALIZED identity (e.g.
#   -450.59 Ha, +74004 kcal/mol vs frozen core) — not the physical origin.
#   The published physical energy is correctly measured from
#   self.qubit_h_physical (the pre-penalty snapshot, line ~4900), so the
#   science is unaffected. But the misleading print cost real debugging time
#   (a phantom "variational violation") and will make any reviewer running
#   the engine distrust every energy.
#
# FIX:
#   Print the identity coefficient of self.qubit_h_physical (pre-penalty,
#   the operator the publishable energy and the exact reference both use).
#   If qubit_h_physical is unavailable (no penalties / external H), fall
#   back to self.qubit_h. Label both cases explicitly so the log is
#   self-documenting.
#
#   This does NOT change any computed/reported energy. Diagnostic-only.
#
# ENABLE: apply LAST is fine; it only wraps debug_energy_bookkeeping, not
# optimize(), so it composes with patch_sector_aware_selection regardless
# of order. Add near the other patch imports:
#       import patch_debug_identity_origin
#       patch_debug_identity_origin.apply(globals())
# =============================================================================


def apply(g):
    VQESolverV2 = g["VQESolverV2"]

    def debug_energy_bookkeeping(self):
        """Print minimal diagnostics to validate reference-energy
        consistency. PATCHED: identity coeff is read from the PRE-PENALTY
        self.qubit_h_physical so it matches the publishable energy and the
        exact same-active-space reference (self.qubit_h is mutated in place
        by penalties and must NOT be used for this diagnostic)."""
        phys = getattr(self, "qubit_h_physical", None)
        id_source = None
        id_coeff = None
        try:
            if phys is not None:
                id_coeff = float(phys.terms.get((), 0.0))
                id_source = "qubit_h_physical (pre-penalty)"
            else:
                id_coeff = float(self.qubit_h.terms.get((), 0.0))
                id_source = "qubit_h (no pre-penalty snapshot available)"
        except Exception:
            id_coeff = None
            id_source = "unavailable"

        # Also surface the penalized identity for transparency, clearly
        # labelled, so the old confusing value is explained rather than
        # hidden.
        pen_id = None
        try:
            if phys is not None:
                pen_id = float(self.qubit_h.terms.get((), 0.0))
        except Exception:
            pen_id = None

        fc = self.ham_info.get("frozen_core_energy", None)
        nuc = self.ham_info.get("nuclear_repulsion", None)

        print("\n" + "=" * 70)
        print("ENERGY BOOKKEEPING DIAGNOSTICS  [patched: physical-origin]")
        print("=" * 70)
        if fc is not None:
            print(f"  frozen_core_energy:      {float(fc): .12f} Ha")
        else:
            print("  frozen_core_energy:      <missing>")
        if nuc is not None:
            print(f"  nuclear_repulsion:       {float(nuc): .12f} Ha")
        else:
            print("  nuclear_repulsion:       <missing>")

        if id_coeff is None:
            print("  qubit_h identity coeff:  <unavailable>")
        else:
            print(f"  identity coeff (PHYS):   {id_coeff: .12f} Ha"
                  f"   [{id_source}]")
            if fc is not None:
                delta = (id_coeff - float(fc)) * 627.509
                print(f"  (PHYS identity - frozen_core): {delta:+.6f} "
                      f"kcal/mol")
        if pen_id is not None:
            print(f"  identity coeff (PENALIZED, self.qubit_h, "
                  f"NOT for energy): {pen_id: .12f} Ha")
            print("    ^ expected to differ from PHYS by the penalty "
                  "constants; this is correct, not a bug.")
        print("=" * 70)

    VQESolverV2.debug_energy_bookkeeping = debug_energy_bookkeeping
    print("  [PATCH] patch_debug_identity_origin applied "
          "(diagnostic identity coeff now reads qubit_h_physical).")