import pandas as pd
from scipy.optimize import fsolve
import glob
import os
import math

# RCE方程函数
def RCE(rce):
    try:
        RC = float(rce[0])
        E = float(rce[1])
        return [
            2*K*E-RC*RC,
            2*K1*E-RC*RC
        ]
    except:
        return ['null', 'null']

# R123方程函数
def R123(r):
    try:
        R1 = float(r[0])
        R2 = float(r[1])
        R3 = float(r[2])
        Er = float(r[3])
        V = float(r[4])
        return [
            # 2*弹性模量（E）*机械比能-R1*R1-R2*R2-R3*R3+2*v*(R1*R2+R1*R3+R2*R3)
            2*Er*K-R1*R1-R2*R2-R3*R3+2*V*(R1*R2+R1*R3+R2*R3),
            2*Er*K1-R1*R1-R2*R2-R3*R3+2*V*(R1*R2+R1*R3+R2*R3),
            2*Er*K2-R1*R1-R2*R2-R3*R3+2*V*(R1*R2+R1*R3+R2*R3),
            2*Er*K3-R1*R1-R2*R2-R3*R3+2*V*(R1*R2+R1*R3+R2*R3),
            2*Er*K4-R1*R1-R2*R2-R3*R3+2*V*(R1*R2+R1*R3+R2*R3)
        ]
    except:
        return ['null', 'null', 'null', 'null', 'null']

# 对数据为null 的进行平均值填补
def nullData(arr):
    num = 0
    lens = len(arr)
    for j in arr:
        if j != 'null':
            num += j
    ag = num/lens
    for z in range(lens):
        el = arr[z]
        if el == 'null' or el == 0:
            arr[z] = ag


if __name__ == "__main__":
    path = 'C:\\Users\\caner\\Desktop\\countExcel'
    colTitle = '机械比能'
    try:
        files = glob.glob(os.path.join(path, "*.xls"))
        for file in files:
            array = pd.read_excel(file)  # 读取excel数据
            try:
                # 删除一次old
                array.drop('单轴抗压强度', axis=1, inplace=True)
                array.drop('弹性模量', axis=1, inplace=True)
                array.drop('第一主应力', axis=1, inplace=True)
                array.drop('第二主应力', axis=1, inplace=True)
                array.drop('第三主应力', axis=1, inplace=True)
                array.drop('岩石强度应力比', axis=1, inplace=True)
            except:
                print('无旧数据')
                
            # 获取数据
            col = array[colTitle]  # 某列
            colLen = len(col)  # 一列的长度
            rowLen = len(array.iloc[1])  # 一行的长度
            RC = [0]*(colLen)  # RC初始化
            E = [0]*(colLen)  # E初始化
            RK1 = [0]*(colLen)  # R1初始化
            RK2 = [0]*(colLen)  # R2初始化
            RK3 = [0]*(colLen)  # R3初始化
            GRA = [0]*(colLen)  # 强度初始化

            # 两行计算RCE
            for i in range(colLen):
                if (i+1) < len(col):
                    K = col[i]
                    K1 = col[i+1]
                    res = fsolve(RCE, [1, 1000])  # 其它参数直接用
                    rc = RCE(res)[0]
                    e = RCE(res)[1]
                    RC[i] = math.fabs(rc)  # 绝对值
                    RC[i+1] = math.fabs(rc)  # 列最后一个值使用上一个的值
                    E[i] = math.fabs(e)
                    E[i+1] = math.fabs(e)

                # 单独对数据为null的进行平均值计算
            nullData(RC)
            nullData(E)

            # 三行计算R123
            for j in range(colLen):
                if (j+4) < len(col):
                    K = col[j]
                    K1 = col[j+1]
                    K2 = col[j+2]
                    K3 = col[j+3]
                    K4 = col[j+4]
                    rea = fsolve(R123, [1, 1, 1, 100, 1])  # 其它参数直接用
                    rea = R123(rea)
                    arr = sorted(rea, reverse=True)  # 倒叙
                    print(arr)
                    RK1[j] = math.fabs(arr[0])
                    RK1[j+1] = math.fabs(arr[0])
                    RK1[j+2] = math.fabs(arr[0])
                    RK1[j+3] = math.fabs(arr[0])
                    RK1[j+4] = math.fabs(arr[0])
                    RK2[j] = math.fabs(arr[1])
                    RK2[j+1] = math.fabs(arr[1])
                    RK2[j+2] = math.fabs(arr[1])
                    RK2[j+3] = math.fabs(arr[1])
                    RK2[j+4] = math.fabs(arr[1])
                    RK3[j] = math.fabs(arr[2])
                    RK3[j+1] = math.fabs(arr[2])
                    RK3[j+2] = math.fabs(arr[2])
                    RK3[j+3] = math.fabs(arr[2])
                    RK3[j+4] = math.fabs(arr[2])

            # 单独对数据为null的进行平均值计算
            nullData(RK1)
            nullData(RK2)
            nullData(RK3)

            # 对整体数据进行倍数计算↓
            for G in range(colLen):
                RC[G] = RC[G]/100
                E[G] = E[G]
                RK1[G] = RK1[G]/100
                RK2[G] = RK2[G]/100
                RK3[G] = RK3[G]/100

            # 计算强度
            for z in range(colLen):
                try:
                    GRA[z] = RC[z] / RK1[z]
                except:
                    GRA[z] = 'null'

            # 单独对数据为null的进行平均值计算
            nullData(GRA)

            # 保存excel
            array.insert(rowLen, '单轴抗压强度', RC)
            array.insert(rowLen+1, '弹性模量', E)
            array.insert(rowLen+2, '第一主应力', RK1)
            array.insert(rowLen+3, '第二主应力', RK2)
            array.insert(rowLen+4, '第三主应力', RK3)
            array.insert(rowLen+5, '岩石强度应力比', GRA)
            pd.DataFrame(array).to_excel(
                file, sheet_name='Sheet1', index=False, header=True)
            print('保存完毕')

    except ZeroDivisionError as err:
        print('计算错误', err)