from pydrake.solvers import MathematicalProgram, Solve
import numpy as np

prog = MathematicalProgram()
x = prog.NewContinuousVariables(2)
cost1 = prog.AddLinearCost(2*x[0] + 3 * x[1] + 2)
print(f"original cost: {cost1}")
prog.AddBoundingBoxConstraint(-1, 1, x)
result = Solve(prog)
print(f"optimal solution: {result.GetSolution(x)}")
print(f"original optimal cost: {result.get_optimal_cost()}")

# Now update the cost to 3x[0] - 4x[1] + 5
cost1.evaluator().UpdateCoefficients(new_a=[3, -4], new_b=5)
print(f"updated cost: {cost1}")
# Solve the optimization problem again with the updated cost.
result = Solve(prog)
print(f"updated optimal solution: {result.GetSolution(x)}")
print(f"updated optimal cost: {result.get_optimal_cost()}")prog = MathematicalProgram()
x = prog.NewContinuousVariables(2)
cost1 = prog.AddQuadraticCost(x[0]**2 + 2 * x[1]**2 + x[0]*x[1] + 3*x[0] + 5)
print(f" original cost: {cost1}")
cost1.evaluator().UpdateCoefficients(new_Q=[[1., 2], [2., 4]], new_b=[1, 2.], new_c= 4)
print(f" updated cost: {cost1}")prog = MathematicalProgram()
x = prog.NewContinuousVariables(2)
constraint1 = prog.AddLinearConstraint(x[0] + 3 * x[1] <= 2)
print(f"original constraint: {constraint1}")
constraint1.evaluator().UpdateLowerBound([-1])
print(f"updated constraint: {constraint1}")
constraint1.evaluator().UpdateUpperBound([3])
print(f"updated constraint: {constraint1}")
constraint1.evaluator().set_bounds(new_lb=[-5], new_ub=[10])
print(f"updated constraint: {constraint1}")prog = MathematicalProgram()
x = prog.NewContinuousVariables(2)
linear_constraint = prog.AddLinearConstraint(3 * x[0] + 4 * x[1] <= 5)
linear_eq_constraint = prog.AddLinearConstraint(5 * x[0] + 2 * x[1] == 3)
print(f"original linear constraint: {linear_constraint}")
linear_constraint.evaluator().UpdateCoefficients(new_A = [[1, 3]], new_lb=[-2], new_ub=[3])
print(f"updated linear constraint: {linear_constraint}")

print(f"original linear equality constraint: {linear_eq_constraint}")
linear_eq_constraint.evaluator().UpdateCoefficients(Aeq=[[3, 4]], beq=[2])
print(f"updated linear equality constraint: {linear_eq_constraint}")