from scipy.optimize import root,fsolve
import numpy as np
from matplotlib import pyplot as plt

flag = 1
result,axis = [],[]                         ###result是最后的y轴数据，axis是最后的x轴数据
while flag<100:
    mu = 0.7
    sigma = 10
    tau = 0.6
    f = flag
    L1 = f * mu / 100
    L2 = mu - L1                            ###设定参数
    def  f1(w):
        return np.array([w[0] * L1 - mu * (((L1 / L2) * (w[0] * tau / w[1]) ** (1 - sigma)) / (
            1 + (L1 / L2) * (w[0] * tau / w[1]) ** (1 - sigma)) * ((1 - mu) / 2 + w[0] * L1) +
                                       ((L1 / L2) * (w[0] / (w[1] * tau)) ** (1 - sigma)) / (
                                               1 + (L1 / L2) * (w[0] / (w[1] * tau)) ** (1 - sigma)) * (
                                               (1 - mu) / 2 + w[1] * L2)),
                     w[1] * L2 - mu * (1 / (1 + (L1 / L2) * (w[0] * tau / w[1]) ** (1 - sigma)) * (
                             (1 - mu) / 2 + w[0] * L1) +
                                       1 / (1 + (L1 / L2) * (w[0] / (w[1] * tau)) ** (1 - sigma)) * (
                                               (1 - mu) / 2 + w[1] * L2))])                     ###求解方程--1


    sol_fsolve = fsolve(f1, [1, 1])         ###求解方程--2
    sol_list = list(sol_fsolve)             ###将方程的解集的格式转化为list
    sol_w1 = sol_list[0]
    sol_w2 = sol_list[1]
    W1 = float(sol_w1)
    W2 = float(sol_w2)                      ###将方程的解的格式转化为float
    print('W1,W2:',W1,W2)
    print('type of W1 and W2:',type(W1),type(W2))
    P1 = ((f / 100) * (W1 ** (-(sigma - 1))) + (1 - f / 100) * (W2 / tau) ** (-(sigma - 1)))**(1/(1-sigma))
    print('P1:',P1)
    P2 = ((f / 100) * ((W1/tau) ** (-(sigma - 1))) + (1 - f / 100) * W2 ** (-(sigma - 1)))**(1/(1-sigma))
    print('P2:',P2)                         ###计算价格指数
    omega1 = W1 * P1**(-mu)
    omega2 = W2 * P2**(-mu)                 ###将名义价格转化为实际价格
    real_wage_ratio = omega1 / omega2
    result.append(real_wage_ratio)
    axis.append(f/100)                      ###将运算结果添加入x轴和y轴数据库
    flag += 1





print('result:',result)
Y = np.array(result)
print('Y:',Y)
X = np.array(axis)                          ###将list转化为np.array格式
print('X:',X)
print(result[49])

plt.plot(X,Y,'.',color = 'green')           ###散点图
plt.plot(X,Y,color = 'red',linewidth = 2)   ###拟合图

plt.xlim(0,1)
plt.ylim(0.9,1.1)                           ###限定坐标轴上下限
plt.xlabel('f')
plt.ylabel('ω1/ω2')
plt.title('labor distribution and real wage ratio')         ###坐标轴标签和表格标题
plt.text(0.5,1.08,r'$\mu=0.7,\ \sigma=10,\ \tau=0.6$')        ###表明参数值
plt.xticks([0.1,0.2,0.3,0.4,0.5,0.6,0.7,0.8,0.9,1])         ###限定横坐标标志点的值
t = 0.5
plt.plot([t,t],[0,result[49]],color = 'blue',linewidth = 2.5,linestyle = "--")
plt.plot([0,1],[result[49],result[49]],color = 'blue',linewidth = 2.5,linestyle = '--')
plt.annotate('stable equlibrium',xy = (0.5,result[49]),xytext = (0.6,1.02),arrowprops = dict(facecolor = 'black'))      ###标注均衡点
plt.annotate('unstable equlibrium',xy = (axis[21],result[21]),xytext = (0.2,0.96),arrowprops = dict(facecolor = 'black'))
plt.annotate('unstable equlibrium',xy = (axis[-28],result[-28]),xytext = (0.7,0.96),arrowprops = dict(facecolor = 'black'))
plt.show()                                  ###显示绘图结果