import numpy as np
import matplotlib.pyplot as plt

P1 = np.array([0.17,    0.0866])
P2 = np.array([0.16933, 0.08586])

r0      = np.linalg.norm(P1 - P2)
epsilon = r0 / 2
theta0  = np.arctan2(P1[1] - P2[1], P1[0] - P2[0])
Theta   = np.pi

theta_star = Theta * (1 - epsilon / r0)
thetas     = np.linspace(0, theta_star, 500)
rs         = r0 * (1 - thetas / Theta)
angles     = theta0 + thetas
xs_spiral  = P2[0] + rs * np.cos(angles)
ys_spiral  = P2[1] + rs * np.sin(angles)

P_entry = np.array([xs_spiral[-1], ys_spiral[-1]])
print("P_entry (вход в ε-окрестность):", P_entry)

v = P2 - P1
L = np.linalg.norm(v)
u = v / L

n = np.array([v[1], -v[0]])
n = n / np.linalg.norm(n)

s2 = 1.0
y_end = P2[1] - 1.5 * r0
s_end = (y_end - P1[1]) / v[1]

a = 0.02 * r0

N = 400
s_all = np.linspace(0.0, s_end, N)

xs_curve = np.zeros_like(s_all)
ys_curve = np.zeros_like(s_all)

for i, s in enumerate(s_all):
    base_pt = P1 + s * v
    if s <= s2:
        offset = 0.0
    else:
        offset = a * ((s - s2) / (s_end - s2))
    xs_curve[i] = base_pt[0] + offset * n[0]
    ys_curve[i] = base_pt[1] + offset * n[1]

idx_P2 = np.argmin(np.abs(s_all - s2))
P2_from_curve = np.array([xs_curve[idx_P2], ys_curve[idx_P2]])

dists = np.sqrt((xs_curve - P_entry[0])**2 + (ys_curve - P_entry[1])**2)

target = 2 * epsilon
idx_mark = np.argmin(np.abs(dists - target))

X_mark = xs_curve[idx_mark]
Y_mark = ys_curve[idx_mark]

fig, ax = plt.subplots(figsize=(6, 6))

# 4.1. Спираль (черная)
ax.plot(xs_spiral, ys_spiral, 'k-', linewidth=2)
ax.scatter(*P1, color='k', s=50)
ax.scatter(*P_entry, color='orange', s=50)
ax.scatter(*P2, color='r', s=30)
circle = plt.Circle(P2, epsilon, edgecolor='r', facecolor='none', linewidth=2)
ax.add_patch(circle)

ax.plot(xs_curve, ys_curve, 'b--', linewidth=2)

ax.scatter(X_mark, Y_mark, color='magenta', s=60, marker='o')

ax.set_aspect('equal', 'box')
ax.set_xlim(P2[0] - r0 * 1.1, P2[0] + r0 * 1.1)
ax.set_ylim(P2[1] - r0 * 1.1, P2[1] + r0 * 1.1)

ax.set_xlabel('X (м)')
ax.set_ylabel('Y (м)')
ax.grid(True)

plt.show()

import sympy as sp

t = sp.symbols('t')
v1, v2 = sp.Function('v1')(t), sp.Function('v2')(t)
w1, w2 = sp.Function('w1')(t), sp.Function('w2')(t)

a2, k1, k2, m1, m2, g, Theta1 = sp.symbols('a2 k1 k2 m1 m2 g Theta1')
dot_v1 = sp.diff(v1, t)
dot_v2 = sp.diff(v2, t)

a11 = m1 * (k1 - a2)**2 + m2 * (k1**2 + 2 * k1 * k2 * sp.sin(v2))
a12 = m2 * (k2**2 + k1 * k2 * sp.sin(v2))
a22 = m2 * k2**2
A = sp.Matrix([[a11, a12], [a12, a22]])
v = sp.Matrix([dot_v1, dot_v2])
T = sp.simplify((1/2) * v.T * A * v)[0]

P = g * (m1 * (k1 - a2) * sp.sin(v1) + m2 * (k1 * sp.sin(Theta1) - k2 * sp.cos(v1 + v2)))

L = sp.simplify(T - P)

dL_ddotv1 = sp.diff(L, dot_v1).simplify()
ddt_dL_ddotv1 = sp.diff(dL_ddotv1, t).simplify()
dL_dv1 = sp.diff(L, v1).simplify()
EL_v1 = sp.simplify(ddt_dL_ddotv1 - dL_dv1)

dL_ddotv2 = sp.diff(L, dot_v2).simplify()
ddt_dL_ddotv2 = sp.diff(dL_ddotv2, t).simplify()
dL_dv2 = sp.diff(L, v2).simplify()
EL_v2 = sp.simplify(ddt_dL_ddotv2 - dL_dv2)

print("Euler-Lagrange eq. for v1:")
print(sp.latex(EL_v1))
print("\nEuler-Lagrange eq. for v2:")
print(sp.latex(EL_v2))

L1_dep, d1_dep, a1_dep, L2_dep, d2_dep, a2_dep = sp.symbols('L1_dep d1_dep a1_dep L2_dep d2_dep a2_dep')

B11 = sp.simplify(L1_dep*(d1_dep*sp.sin(w1 - v1) - a1_dep*sp.cos(v1)) / (d1_dep*(L1_dep*sp.sin(w1 - v1) - a1_dep*sp.cos(w1))))
B22 = sp.simplify(L2_dep*(d2_dep*sp.sin(w2 - v2) - a2_dep*sp.cos(v2)) / (d2_dep*(L2_dep*sp.sin(w2 - v2) - a2_dep*sp.cos(w2))))

eq_w1 = sp.Eq(sp.diff(w1, t), B11 * dot_v1)
eq_w2 = sp.Eq(sp.diff(w2, t), B22 * dot_v2)

print("\nKinematic constraint for w1:")
sp.pprint(eq_w1)
print("\nKinematic constraint for w2:")
sp.pprint(eq_w2)


b1, b2, i1, i2 = sp.symbols('b1 b2 i1 i2')
k2_1, k2_2 = sp.symbols('k2_1 k2_2')
Q_w1 = sp.simplify(-b1 * B11 * dot_v1 + k2_1 * i1)
Q_w2 = sp.simplify(-b2 * B22 * dot_v2 + k2_2 * i2)

Q_s = sp.Matrix([Q_w1, Q_w2])

B_mat = sp.diag(B11, B22)

rhs_EL = sp.simplify(B_mat.T * Q_s)

eq_v1 = sp.Eq(EL_v1, rhs_EL[0])
eq_v2 = sp.Eq(EL_v2, rhs_EL[1])

I1, R1, k11, I2, R2, k12 = sp.symbols('I1 R1 k11 I2 R2 k12')

e1, e2 = sp.symbols('e1 e2')

i1_func = sp.Function('i1')(t)
i2_func = sp.Function('i2')(t)

eq_motor1 = sp.Eq(I1 * sp.diff(i1_func, t) + R1 * i1_func + k11 * B11 * dot_v1, e1)
eq_motor2 = sp.Eq(I2 * sp.diff(i2_func, t) + R2 * i2_func + k12 * B22 * dot_v2, e2)

print("\nElectrical equation for motor 1:")
sp.pprint(eq_motor1)
print("\nElectrical equation for motor 2:")
sp.pprint(eq_motor2)

import numpy as np
from scipy.integrate import solve_ivp
from scipy.linalg import solve_continuous_are, inv
import matplotlib.pyplot as plt


m1_val = 0.3
m2_val = 0.1
a1_val = 0.041
a2_val = 0.041
d1_val = 0.044
d2_val = 0.037
L1_val = 0.061
L2_val = 0.05
l1_val = 0.056
l2_val = 0.06
k1_val = 0.1
k2_val = 0.12

k2_1_val = 0.02
k2_2_val = 0.02
k11_val = 0.017
k12_val = 0.017
I1_val = 0.02
I2_val = 0.02
R1_val = 10
R2_val = 10

g_val = 9.81

v1_0_val = np.deg2rad(60)
v2_0_val = np.deg2rad(30)
w1_0_val = np.deg2rad(61 + 23/60)
w2_0_val = np.deg2rad(10)

x0_val = 0.17
y0_val = 0.0866


B11_0 = L1_val*(d1_val*np.sin(w1_0_val - v1_0_val) - a1_val*np.cos(v1_0_val)) / \
       (d1_val*(L1_val*np.sin(w1_0_val - v1_0_val) - a1_val*np.cos(w1_0_val)))
B22_0 = L2_val*(d2_val*np.sin(w2_0_val - v2_0_val) - a2_val*np.cos(v2_0_val)) / \
       (d2_val*(L2_val*np.sin(w2_0_val - v2_0_val) - a2_val*np.cos(w2_0_val)))

d_val = (m1_val*(k1_val - a2_val)**2 + m2_val*(k1_val**2 + 2*k1_val*k2_val*np.sin(v2_0_val)))*m2_val*k2_val**2 \
        - m2_val**2*(k1_val*k2_val*np.sin(v2_0_val) + k2_val**2)**2
b11_val = m2_val * k2_val**2 / d_val
b12_val = -m2_val * (k1_val*k2_val*np.sin(v2_0_val) + k2_val**2) / d_val
b22_val = (m1_val*(k1_val - a2_val)**2 + m2_val*(k1_val**2 + 2*k1_val*k2_val*np.sin(v2_0_val))) / d_val


h11_val = -g_val*(m1_val*(k1_val - a2_val)*np.sin(v1_0_val) - m2_val*(k1_val*np.sin(v1_0_val) + k2_val*np.cos(v1_0_val + v2_0_val)))
h21_val = -g_val*m2_val*k2_val*np.cos(v1_0_val + v2_0_val)
h12_val = 1.0
h15_val = 1.0
h26_val = B22_0 * 0.1
h13_val = -g_val*m2_val*np.cos(v1_0_val + v2_0_val)
h22_val = 1.0
b2_val = 0.5
h24_val = -b2_val * B22_0**2
h17_val = 0.5
h28_val = 0.5
h23_val = -g_val*m2_val*k2_val*np.cos(v1_0_val + v2_0_val)

p21 = b11_val*h11_val + b12_val*h21_val
p22 = b11_val*h12_val
p23 = b11_val*h12_val + b12_val*h22_val
p24 = b11_val*h24_val
p25 = b11_val*h15_val
p26 = b12_val*h26_val

p41 = b12_val*h11_val + b22_val*h21_val
p42 = b12_val*h12_val
p43 = b12_val*h13_val + b22_val*h23_val
p44 = b22_val*h24_val
p45 = b12_val*h15_val
p46 = b22_val*h26_val

p52 = 0.1
p55 = 0.2
p64 = 0.1
p66 = 0.2

P_num = np.array([
    [0,    1,   0,    0,   0,   0],
    [p21, p22, p23, p24, p25, p26],
    [0,    0,   0,    1,   0,   0],
    [p41, p42, p43, p44, p45, p46],
    [0,    p52, 0,    0,   p55, 0],
    [0,    0,   0,    p64, 0,   p66]
])

Q_num = np.array([
    [0, 0],
    [0, 0],
    [0, 0],
    [0, 0],
    [1, 0],
    [0, 1]
])

Q_lqr = np.eye(6)
R_lqr = np.eye(2)

S_care = solve_continuous_are(P_num, Q_num, Q_lqr, R_lqr)
K = inv(R_lqr) @ Q_num.T @ S_care

A_cl = P_num - Q_num @ K

def grasp_position(x):
    return np.array([x0_val + x[0], y0_val + x[2]])

delta_y = 0.05
target = np.array([x0_val, y0_val - delta_y])

R_val = 0.1
def brachistochrone(theta, x_start, y_start):
    return np.array([x_start + R_val*(theta - np.sin(theta)),
                     y_start - R_val*(1 - np.cos(theta))])

theta_target = 0.05
target_brach = brachistochrone(theta_target, x0_val, y0_val)

distance_full = np.linalg.norm(target_brach - np.array([x0_val, y0_val]))
epsilon = 0.5 * distance_full

print("Целевая точка (по брахистохроне):", target_brach)
print("Радиус ε-окрестности:", epsilon)

x0_state = np.zeros(6)
x0_state[0] = x0_val - target_brach[0]
x0_state[2] = y0_val - target_brach[1]


def closed_loop_ode(t, x):
    return A_cl @ x

t_span = (0, 5)
sol = solve_ivp(closed_loop_ode, t_span, x0_state, dense_output=True, max_step=0.01)

t_vals = sol.t
x_traj = sol.y


grasp_traj = np.array([grasp_position(x_traj[:, i]) for i in range(x_traj.shape[1])])
distances = np.linalg.norm(grasp_traj - target_brach, axis=1)

idx = np.where(distances < epsilon)[0]