import numpy as np
import matplotlib.pyplot as plt


def catenary(x, a):
    return a * np.cosh((x - 0.5) / a) - a * np.cosh((-0.5) / a)


x = np.linspace(0, 1, 100)
for a in [0.35, 0.5, 0.8]:
    plt.plot(x, catenary(x, a), label='a={:g}'.format(a))
plt.legend()
plt.show()

# 悬链线的力学模拟
import numpy as np


class Catenary:
    def __init__(self):
        self.length = 2
        self.n = 31
        self.dump = 0.2  # 阻尼系数
        self.k = 1  # 弹簧系数
        self.l = self.length / (self.n - 1)  # 每段长度
        self.g = 0.01  # 重力加速度

        self.x = np.concatenate([np.linspace(0, 1, self.n).reshape(-1, 1), np.zeros((self.n, 1))],
                                axis=1)
        self.v = np.zeros_like(self.x)

        # 重力加速度
        self.acc_gravity = np.zeros_like(self.x)
        self.acc_gravity[:, 1] = -self.g

    def cal_accelerate(self, a, b):
        # 输入x，和x左/右相邻点坐标，记为a，b
        # 输出弹力对应的加速度
        k, l = self.k, self.l
        ab = b - a
        ab_m = np.sqrt(np.square(ab[:, [0]]) + np.square(ab[:, [1]]))
        acc_scale_tmp = k * (ab_m - l)
        acc_scale = np.where(acc_scale_tmp > 0, acc_scale_tmp, 0)
        acc = ab / ab_m * k * acc_scale
        return acc

    def update_v(self, x, v):
        n, dump = self.n, self.dump
        acc_left = self.cal_accelerate(x[1:n - 1, :], x[:n - 2, :])  # 受到左边的拉力
        acc_right = self.cal_accelerate(x[1:n - 1, :], x[2:, :])  # 受到右边的拉力
        acc1 = np.concatenate([[[0, 0]], acc_left + acc_right, [[0, 0]]], axis=0)
        v = (1 - dump) * v + acc1 + self.acc_gravity
        v[0, :] = 0
        v[-1, :] = 0
        return v

    def update_location(self, x, v):
        x = x + v
        return x

    def update_all(self):
        self.v = self.update_v(self.x, self.v)
        self.x = self.update_location(self.x, self.v)
        return self.x


catenary = Catenary()

frame = [catenary.update_all() for i in range(100)]

# %%
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation

fig, ax = plt.subplots()

line, = ax.plot([0], [0], '.-', animated=True)
ax.set_xlim(0, 1)
ax.set_ylim(-5, 1)


def update_line(i):
    plt.setp(line, 'xdata', i[:, 0], 'ydata', i[:, 1])
    return [line]


ani = FuncAnimation(fig, update_line, blit=True, interval=25, frames=frame)
# ani.save('catenary.gif', writer='pillow')