"""
Verification of substrate optimization first-order conditions (Ch. II.8).
Renewal Libertarianism, Part 1, Chapter II, Module 2.

Each FOC the manuscript states is re-derived symbolically by differentiating
the Bellman objective F = Pi_S * (1-r) + delta_S * (1-r) * V_S with respect to
the relevant control variable and applying the envelope theorem (treating V_S
as the optimized continuation value, so its derivative with respect to the
control vanishes at the optimum).

The script then checks that SymPy's derivation produces the stated FOC form:
    (d Pi_S / d x) / (Pi_S + delta_S * V_S)  =  (d r / d x) / (1 - r)
for each control x in {tau, h, K_S, Q}, and verifies the manuscript's structural
claims about scope (discrete choice entry condition) and the asymptotic limits.

Run: uv run python substrate_optimization_focs.py
"""

import sympy as sp

from _helpers import banner


def main():
    banner("Chapter II.8 substrate optimization FOCs - symbolic verification")

    # ----- Symbols -----
    tau   = sp.symbols("tau",   positive=True)   # extraction rate
    h     = sp.symbols("h",     positive=True)   # opacity investment
    K_S   = sp.symbols("K_S",   positive=True)   # coordination-disruption investment
    Q     = sp.symbols("Q",     positive=True)   # quality investment allocation
    V_S_const = sp.symbols("V_S", positive=True)  # continuation value (envelope-theorem constant)
    delta_S = sp.symbols("delta_S", positive=True)

    # ----- Abstract per-period payoff and replacement probability -----
    # The book writes the per-period payoff Pi_S as a function of (tau, h, K_S, S, Q)
    # and the replacement probability r as a function of the same controls. We
    # carry both as abstract SymPy Functions so chain-rule differentiation does
    # the work without us assuming specific functional forms.
    Pi_S = sp.Function("Pi_S")(tau, h, K_S, Q)
    r    = sp.Function("r")(tau, h, K_S, Q)

    # ----- Bellman objective per period (envelope theorem: V_S held constant) -----
    F = Pi_S * (1 - r) + delta_S * (1 - r) * V_S_const

    print("\n  Bellman per-period objective being maximized:")
    print(f"      F = Pi_S * (1 - r) + delta_S * (1 - r) * V_S")
    print(f"      F = {F}")
    print()
    print("  Envelope theorem: at the optimum, the implicit dependence of V_S")
    print("  on each control vanishes, so V_S is held constant during partial")
    print("  differentiation. SymPy implements this by carrying V_S as a symbol.")

    # ============================================================
    # Generic FOC structure: same shape for any control variable x
    #
    #   d F / d x = (d Pi_S / d x) * (1 - r)
    #             + Pi_S * (- d r / d x)
    #             + delta_S * (- d r / d x) * V_S
    #
    #            = (d Pi_S / d x) * (1 - r)  -  (Pi_S + delta_S * V_S) * (d r / d x)
    #
    # Setting d F / d x = 0 and rearranging yields the manuscript's form:
    #
    #   (d Pi_S / d x) / (Pi_S + delta_S * V_S)  =  (d r / d x) / (1 - r)
    # ============================================================

    def verify_generic_foc(x, x_name):
        """For control variable x, derive the FOC by differentiating F and
        confirm it matches the manuscript's structural form."""
        banner(f"FOC for {x_name}: derive d F / d {x_name} = 0 and check the manuscript form")

        # Step 1: differentiate F symbolically with V_S held constant
        dF_dx = sp.diff(F, x)
        print(f"\n  d F / d {x_name} (raw SymPy output):")
        print(f"      {dF_dx}")

        # Step 2: collect terms into the structure
        #   (d Pi_S / d x) * (1 - r) - (Pi_S + delta_S * V_S) * (d r / d x)
        dPi_dx = sp.diff(Pi_S, x)
        dr_dx  = sp.diff(r, x)

        manuscript_form = dPi_dx * (1 - r) - (Pi_S + delta_S * V_S_const) * dr_dx
        print(f"\n  Manuscript form (unrearranged):")
        print(f"      (d Pi_S / d {x_name}) * (1 - r)  -  (Pi_S + delta_S * V_S) * (d r / d {x_name})")
        print(f"      = {manuscript_form}")

        # Step 3: verify the two expressions match
        residual = sp.simplify(sp.expand(dF_dx - manuscript_form))
        print(f"\n  Difference (SymPy derivation - manuscript form, expanded and simplified):")
        print(f"      = {residual}")
        assert residual == 0, (
            f"SymPy's d F / d {x_name} did not match the manuscript's FOC form")
        print(f"  [pass] d F / d {x_name} = manuscript form")

        # Step 4: verify the rearranged form
        #   (d Pi_S / d x) / (Pi_S + delta_S * V_S)  =  (d r / d x) / (1 - r)
        # Multiply both sides through and check the cross-multiplication identity
        lhs = dPi_dx * (1 - r)
        rhs = (Pi_S + delta_S * V_S_const) * dr_dx
        manuscript_section = {"tau": "1", "Q": "3"}.get(x_name, "*")
        print(f"\n  Rearranged form (manuscript II.8.{manuscript_section}):")
        print(f"      (d Pi_S / d {x_name}) / (Pi_S + delta_S * V_S)  =  (d r / d {x_name}) / (1 - r)")
        print(f"      equivalently: {lhs}  =  {rhs}")
        # The FOC sets d F / d x = 0, which (factoring out the negative on the second term)
        # is exactly lhs = rhs.
        # Verify: dF_dx = lhs - rhs (the unrearranged form has +(d Pi/d x)(1-r) - (Pi+delta V)(dr/dx))
        crossmul_residual = sp.simplify(sp.expand(dF_dx - (lhs - rhs)))
        assert crossmul_residual == 0, (
            "FOC rearrangement identity failed")
        print(f"  [pass] rearrangement to (manuscript II.8.x) form holds")

        return dF_dx, dPi_dx, dr_dx

    # Verify the FOC for each scalar control: tau, h, K_S, Q
    dF_dtau, dPi_dtau, dr_dtau = verify_generic_foc(tau, "tau")
    dF_dh,   _,        _      = verify_generic_foc(h,   "h")
    dF_dKS,  _,        _      = verify_generic_foc(K_S, "K_S")
    dF_dQ,   dPi_dQ,   dr_dQ  = verify_generic_foc(Q,   "Q")

    # ============================================================
    # Scope FOC: discrete entry condition (II.8.2)
    # The substrate compares expected rent from entering domain d against
    # expected cost. The book's condition (entering iff > 0):
    #
    #   tau_d_star * V_d  -  c_S(d)  -  r_S(d) * (Pi_S + delta_S * V_S)  >  0
    # ============================================================
    banner("Scope FOC (II.8.2): discrete entry condition")

    tau_d_star, V_d, c_S_d, r_S_d = sp.symbols(
        "tau_d_star V_d c_S_d r_S_d", positive=True)

    entry_value = tau_d_star * V_d - c_S_d - r_S_d * (Pi_S + delta_S * V_S_const)
    print(f"\n  Entry value for domain d:")
    print(f"      tau_d* * V_d  -  c_S(d)  -  r_S(d) * (Pi_S + delta_S * V_S)")
    print(f"      = {entry_value}")
    print(f"\n  Substrate enters domain d iff this expression is positive.")
    print(f"  The condition decomposes into:")
    print(f"      rent available    : tau_d* * V_d")
    print(f"      cost of authority : c_S(d)")
    print(f"      replacement-risk  : r_S(d) * (Pi_S + delta_S * V_S)")

    # Verify each component contributes with the manuscript's claimed sign
    components = {
        "tau_d_star": (sp.diff(entry_value, tau_d_star),  +1, "rent increases entry value"),
        "V_d":        (sp.diff(entry_value, V_d),         +1, "rent base increases entry value"),
        "c_S_d":      (sp.diff(entry_value, c_S_d),       -1, "cost decreases entry value"),
        "r_S_d":      (sp.diff(entry_value, r_S_d),       -1, "scope-detection risk decreases entry value"),
    }
    print()
    for name, (deriv, expected_sign, note) in components.items():
        # The partial is structural; check the sign of the leading expression.
        # At a typical interior point (all variables positive, Pi_S+delta_S*V_S > 0):
        test_pt = {
            tau_d_star: sp.Rational(1), V_d: sp.Rational(1),
            c_S_d: sp.Rational(1), r_S_d: sp.Rational(1, 2),
            Pi_S: sp.Rational(1), V_S_const: sp.Rational(1),
            delta_S: sp.Rational(1, 2),
        }
        val = float(sp.N(deriv.subs(test_pt)))
        sign = +1 if val > 0 else (-1 if val < 0 else 0)
        print(f"  d entry_value / d {name:11s} = {deriv}   value at test pt = {val:+.4f}   ({note})")
        assert sign == expected_sign, (
            f"Sign of d entry_value / d {name} does not match manuscript claim")
    print("  [pass] all four components contribute with the manuscript's claimed signs")

    # ============================================================
    # Quality-extraction complementarity (mentioned in II.8.1, II.9.4)
    #
    # Manuscript claim: "Higher Q raises tau_c, which means lower r_tau at any
    # given extraction level, which means higher optimal tau*."
    #
    # Symbolically: by the implicit function theorem applied to the extraction
    # FOC (treat it as G(tau, Q) = 0), we have
    #
    #   d tau* / d Q  =  - (d G / d Q) / (d G / d tau)
    #
    # At a regular maximum, d G / d tau = d^2 F / d tau^2  <  0. So the sign of
    # d tau* / d Q is the sign of  d^2 F / (d tau)(d Q).
    #
    # If higher Q reduces d r / d tau (the marginal replacement risk from raising
    # extraction), then d^2 F / (d tau d Q) > 0 and d tau* / d Q > 0.
    # ============================================================
    banner("Quality-extraction complementarity (II.8.1, II.9.4)")

    # Build the cross partial d^2 F / (d tau d Q)
    d2F_dtau_dQ = sp.diff(F, tau, Q)
    print(f"\n  Cross partial d^2 F / (d tau d Q) (raw):")
    print(f"      {d2F_dtau_dQ}")

    # The manuscript's sign claim depends on:
    #   (i)   d^2 r / (d tau d Q)  <  0   (higher Q reduces marginal extraction-detection)
    #   (ii)  d^2 Pi_S / (d tau d Q)  >=  0   (higher Q does not reduce extraction's marginal rent)
    #
    # With (i) and (ii), the cross partial d^2 F / (d tau d Q) > 0 at the optimum,
    # and by implicit function theorem d tau* / d Q > 0.
    print(f"\n  Sign of d tau* / d Q determined by sign of d^2 F / (d tau d Q).")
    print(f"  Substantive sufficient conditions (from manuscript):")
    print(f"      (i)  d^2 r / (d tau d Q) <= 0   (higher Q reduces marginal extraction-detection)")
    print(f"      (ii) d^2 Pi_S / (d tau d Q) >= 0  (higher Q does not reduce extraction's marginal rent)")

    # Verify the manuscript's mechanism directly. The book's reduced form for r
    # is the probit r = Phi((tau - tau_c)/sigma), with tau_c the natural revolt
    # threshold (Ch. I.3.4). Quality enters through tau_c(Q): higher Q raises
    # the threshold. Use this concrete form with Pi_S = tau (no quality cost,
    # to isolate the r-channel mechanism the manuscript focuses on).
    tau_sym, Q_sym, sigma_sym = sp.symbols("tau Q sigma", positive=True)
    # Standard normal CDF expressed via erf
    Phi = lambda x: (1 + sp.erf(x / sp.sqrt(2))) / 2
    # tau_c(Q) = Q (linear, increasing) — captures the manuscript's
    # "higher Q raises tau_c" assumption.
    tau_c = Q_sym
    r_concrete = Phi((tau_sym - tau_c) / sigma_sym)
    Pi_concrete = tau_sym

    print(f"\n  Concrete test family (matches the manuscript's reduced form):")
    print(f"      r        = Phi((tau - tau_c(Q))/sigma),  tau_c(Q) = Q")
    print(f"      Pi_S     = tau   (no quality cost; isolates the r channel)")

    # Verify the three building-block sign claims first
    dr_dQ_concrete = sp.diff(r_concrete, Q_sym)
    dr_dtau_concrete = sp.diff(r_concrete, tau_sym)
    d2r_dtau_dQ_concrete = sp.diff(r_concrete, tau_sym, Q_sym)

    # In the operating regime tau < tau_c = Q (substrate hedging below threshold,
    # which is where the manuscript's complementarity claim applies)
    operating_pts = [
        {tau_sym: sp.Rational(1, 4), Q_sym: sp.Rational(1, 2), sigma_sym: sp.Rational(1)},
        {tau_sym: sp.Rational(2, 5), Q_sym: sp.Rational(3, 4), sigma_sym: sp.Rational(1)},
        {tau_sym: sp.Rational(1, 10), Q_sym: sp.Rational(1, 2), sigma_sym: sp.Rational(1, 2)},
    ]

    print(f"\n  Three building-block signs in the operating regime tau < tau_c:")
    print(f"  (a) d r / d Q < 0  (higher Q lowers replacement risk at any tau)")
    for pt in operating_pts:
        val = float(sp.N(dr_dQ_concrete.subs(pt)))
        print(f"      d r / d Q at tau={float(pt[tau_sym]):.4f}, Q={float(pt[Q_sym]):.4f}"
              f"   =  {val:+.6f}")
        assert val < 0, "d r / d Q not negative in operating regime"
    print(f"      [pass] d r / d Q < 0 confirmed")

    print(f"\n  (b) d^2 r / (d tau d Q) < 0  (higher Q lowers marginal extraction-detection)")
    for pt in operating_pts:
        val = float(sp.N(d2r_dtau_dQ_concrete.subs(pt)))
        print(f"      d^2 r / (d tau d Q) at tau={float(pt[tau_sym]):.4f}, Q={float(pt[Q_sym]):.4f}"
              f"   =  {val:+.6f}")
        assert val < 0, "d^2 r / (d tau d Q) not negative in operating regime"
    print(f"      [pass] d^2 r / (d tau d Q) < 0 confirmed in operating regime")

    # Now build d^2 F / (d tau d Q) on this concrete family and check it's positive
    F_concrete = Pi_concrete * (1 - r_concrete) + delta_S * (1 - r_concrete) * V_S_const
    cross_partial = sp.diff(F_concrete, tau_sym, Q_sym)
    print(f"\n  (c) Cross partial d^2 F / (d tau d Q) > 0 in operating regime:")
    for pt in operating_pts:
        full_pt = {**pt, delta_S: sp.Rational(1, 2), V_S_const: sp.Rational(1)}
        val = float(sp.N(cross_partial.subs(full_pt)))
        print(f"      d^2 F / (d tau d Q) at tau={float(pt[tau_sym]):.4f}, Q={float(pt[Q_sym]):.4f},"
              f" delta_S=1/2, V_S=1  =  {val:+.6f}")
        assert val > 0, "Cross partial not positive in operating regime"
    print(f"      [pass] d^2 F / (d tau d Q) > 0 in operating regime")

    print(f"\n  By implicit function theorem applied to extraction FOC,")
    print(f"  d tau* / d Q = -[d^2 F / (d tau d Q)] / [d^2 F / d tau^2].")
    print(f"  At an interior maximum d^2 F / d tau^2 < 0 (second-order condition),")
    print(f"  so the sign of d tau* / d Q matches the sign of d^2 F / (d tau d Q).")
    print(f"  [pass] therefore d tau* / d Q > 0 in the operating regime")
    print(f"         (quality-extraction complementarity confirmed)")
    print(f"\n  Note: the complementarity is regime-dependent. The probit form is")
    print(f"  symmetric about tau = tau_c; the sign of d^2 r / (d tau d Q) flips")
    print(f"  for tau > tau_c. The manuscript's claim applies to the operating")
    print(f"  regime where substrate hedges below the natural revolt threshold,")
    print(f"  which is where any extraction-from-the-substrate analysis matters.")

    # ============================================================
    # Asymptotic limits (II.10)
    #
    # II.10.1: as opacity v -> infinity, optimal tau -> infinity
    # II.10.2: as scope-detection capacity r_S -> 0, optimal scope S -> D
    # II.10.3: as quality-detection capacity r_Q -> 0, optimal Q -> 0
    #
    # These follow directly from inspecting the FOCs as the relevant detection
    # terms vanish. Symbolically, when d r / d tau -> 0, the RHS of the
    # extraction FOC vanishes, forcing the LHS (d Pi_S / d tau / (Pi_S + delta V))
    # to vanish too, which under standard concavity requires tau to take its
    # boundary value (infinity).
    # ============================================================
    banner("Asymptotic limits (II.10): structural consequences of vanishing detection")

    # The two limits below are structural corollaries of the FOC structure
    # verified above; they are stated here as analytical consequences rather
    # than as separately asserted verification steps.
    print("\n  (II.10.1) Extraction limit as d r / d tau -> 0:")
    print(f"      Extraction FOC:  (d Pi_S / d tau) * (1 - r) = (Pi_S + delta_S * V_S) * (d r / d tau)")
    print(f"      As d r / d tau -> 0, RHS -> 0, so (d Pi_S / d tau) * (1 - r) -> 0.")
    print(f"      With (1 - r) > 0 in any viable equilibrium, this forces d Pi_S / d tau -> 0,")
    print(f"      which under increasing-in-tau Pi_S structure pushes tau to its unbounded boundary.")

    print("\n  (II.10.3) Quality limit as d r / d Q -> 0:")
    print(f"      Quality FOC:    (d Pi_S / d Q) * (1 - r) = (Pi_S + delta_S * V_S) * (d r / d Q)")
    print(f"      As d r / d Q -> 0, the marginal benefit of quality through risk reduction vanishes.")
    print(f"      With d Pi_S / d Q < 0 (quality costs rent directly), the LHS is negative while")
    print(f"      the RHS vanishes; the FOC cannot hold for any Q > 0, so the optimum is Q* -> 0.")

    # The asymptotic limits follow analytically from the FOC structure verified
    # above; a numerical demonstration would require a concrete parameterization
    # that actually makes extraction-detection or quality-detection capacity
    # vanish, which is not what the probit-with-threshold family above captures.
    # The structural arguments above complete the verification.

    banner("All verifications passed")
    print("""
  Result: Chapter II.8 FOCs are consistent with envelope-theorem differentiation
  of the Bellman objective.

    - Extraction FOC, opacity FOC, coordination-disruption FOC, and quality FOC
      all derive from d F / d x = 0 via SymPy, matching the manuscript's
      generic form  (d Pi_S / d x) / (Pi_S + delta_S * V_S) = (d r / d x) / (1 - r).
    - Scope discrete-choice entry condition decomposes into the four components
      the manuscript claims, each with the predicted sign.
    - Quality-extraction complementarity (d tau* / d Q > 0) confirmed on a
      concrete functional family that respects the manuscript's structural
      assumptions.
    - Asymptotic limits (extraction-detection vanishes => tau* unbounded,
      quality-detection vanishes => Q* -> 0) follow from inspection of the FOCs.
""")


if __name__ == "__main__":
    main()