import pandas as pd
import numpy as np
import matplotlib.pyplot as plt

def read(s):
    tmp = pd.read_csv(s, delimiter=' ', names=['Index', 'Data', 'Time'], skiprows=3)
    tmp.set_index('Index', inplace=True)

    eps= 5 if '3ATM' in s else (0 if '0ATM' in s else 10)

    if 'dvv+pave' in s:
        tmp['Data'] /= 1.0e6

    if 'vzave' in s:
        negIdx = (tmp['Data'] < 0).idxmax() if (tmp['Data'] < 0).any() else len(tmp['Data'])
        tmp = tmp.loc[:negIdx]

    if 'dvv+pave' in s:
        tmp['Mean'] = (tmp['Data'] + tmp['Data'].shift(1)) / 2.0
        tmp['DiffTime'] = tmp['Time'].diff()
        tmp['DiffImpluse'] = tmp['Mean'] * tmp['DiffTime']
        tmp['AccImpluse'] = tmp['DiffImpluse'].cumsum()

    tmp.dropna(inplace=True)
    #tmp = tmp.loc[1:]

    return {'data': tmp, 'eps': eps, 'epsname': f'ε={eps}'}

# 在此处添加要读取的文件（左侧工具栏内上传后）

ATM3M3_Thick_dvv = read("./边界层厚/3M3ATM,边界层厚,reportdvv+pave-rfile.out")
ATM3M3_Thick_vzave = read("./边界层厚/3M3ATM,边界层厚,reportvzave-rfile.out")

ATM3M6_Thick_dvv = read("./边界层厚/6M3ATM,边界层厚,reportdvv+pave-rfile.out")
ATM3M6_Thick_vzave = read("./边界层厚/6M3ATM,边界层厚,reportvzave-rfile.out")

ATM3M0_dvv = read('./3ATM0M,reportdvv+pave-rfile.out')
ATM3M0_vzave = read('./3ATM0M,reportvzave-rfile.out')

ATM5M6_Thin_dvv = read('./边界层薄/5ATM6M,reportdvv+pave-rfile.out')
ATM5M6_Thin_vzave = read('./边界层薄/5ATM6M,reportvzave-rfile.out')

pass

ATM3M3_Thin_dvv = read('./边界层薄/3ATM3M,reportdvv+pave-rfile.out')
ATM3M3_Thin_vzave = read('./边界层薄/3ATM3M,reportvzave-rfile.out')

ATM3M3_Thin_Tess_dvv = read('./边界层薄/3ATM3M,加密,reportdvv+pave-rfile.out')
ATM3M3_Thin_Tess_vzave = read('./边界层薄/3ATM3M,加密,reportvzave-rfile.out')

ATM3M3_Thin_Conv_dvv = read('./边界层薄/3ATM3M,收敛到1e-3,reportdvv+pave-rfile.out')
ATM3M3_Thin_Conv_vzave = read('./边界层薄/3ATM3M,收敛到1e-3,reportvzave-rfile.out')

ATM3M6_Thin_dvv = read('./边界层薄/3ATM6M,reportdvv+pave-rfile.out')
ATM3M6_Thin_vzave = read('./边界层薄/3ATM6M,reportvzave-rfile.out')

ATM3M6_Thin_Tess_dvv = read('./边界层薄/3ATM6M,加密,reportdvv+pave-rfile.out')
ATM3M6_Thin_Tess_vzave = read('./边界层薄/3ATM6M,加密,reportvzave-rfile.out')

ATM5M3_Thin_dvv = read('./边界层薄/5ATM3M,reportdvv+pave-rfile.out')
ATM5M3_Thin_vzave = read('./边界层薄/5ATM3M,reportvzave-rfile.out')

ATM5M6_Thin_dvv = read('./边界层薄/5ATM6M,reportdvv+pave-rfile.out')
ATM5M6_Thin_vzave = read('./边界层薄/5ATM6M,reportvzave-rfile.out')

pass

ATM3M0_dvv = read('./3ATM0M,reportdvv+pave-rfile.out')
ATM3M0_vzave = read('./3ATM0M,reportvzave-rfile.out')

ATM5M0_dvv = read('./5ATM0M,reportdvv+pave-rfile.out')
ATM5M0_vzave = read('./5ATM0M,reportvzave-rfile.out')

pass

class Counter:
    i = 0
    def get(self):
        self.i += 1
        return self.i
    
nxter = Counter()

import matplotlib.font_manager
matplotlib.font_manager.findSystemFonts(fontpaths=None, fontext='ttf')

VZ, FORCE, IMPLUSE_ACCU = 0x01, 0x02, 0x04
names = {VZ: '出口速度', FORCE: '瞬时力', IMPLUSE_ACCU: '冲量'} # 此处更改不同列的名称

matplotlib.rcParams.update({
    "text.usetex": False,
    "pgf.texsystem": "xelatex",
    "font.size": 32
    })

plt.rcParams['text.latex.preamble']=r'\makeatletter \newcommand*{\rom}[1]{\expandafter\@slowromancap\romannumeral #1@} \makeatother'

matplotlib.use('pgf')

def plot(dvv, vz, item, legend=None):
    if item == VZ:
        plt.plot(vz['data']['Time'], vz['data']['Data'], label=vz['epsname'] if legend is None else legend)
    elif item == FORCE:
        plt.plot(dvv['data']['Time'], dvv['data']['Data'], label=dvv['epsname'] if legend is None else legend)
    elif item == IMPLUSE_ACCU:
        plt.plot(dvv['data']['Time'], dvv['data']['AccImpluse'], label=dvv['epsname'] if legend is None else legend)

def comparePlot(dvvs, vzs, title, xlabel, ylabel, item, legends=None):
    assert len(dvvs) == len(vzs)
    for i in range(len(dvvs)):
        plot(dvvs[i], vzs[i], item, None if legends is None else legends[i])
    plt.title(title)
    plt.xlabel(xlabel)
    plt.ylabel(ylabel)
    plt.legend()
    plt.ticklabel_format(style='sci', axis='x', scilimits=(0,0), useMathText=True)
    plt.ticklabel_format(style='sci', axis='y', scilimits=(0,0), useMathText=True)

for ele in [VZ, FORCE, IMPLUSE_ACCU]:
    plt.figure(nxter.get())
    comparePlot([ATM3M3_Thin_dvv, ATM3M3_Thick_dvv], 
                [ATM3M3_Thin_vzave, ATM3M3_Thick_vzave], 
                'ATM3M3 Thick - Thin', 'Time', names[ele], ele, legends=['Thin', 'Thick'])
    #plt.savefig(f'ATM3M3 Thick - Thin - {names[ele]}.png') # save file
    plt.show()

# ATM3M6_Thin_dvv 和 ATM3M6_Thick_dvv 关于 VZ, FORCE, IMPLUSE_ACCU 分别对比，产生3张图
# ATM3M6_Thin_dvv 和 ATM3M6_Thick_dvv 分别对应标注 'Thin', 'Thick'
# 横轴为 names[ele], 纵轴为 'Time', title为 'ATM3M6 Thick - Thin', 'Time'

for ele in [VZ, FORCE, IMPLUSE_ACCU]:
    plt.figure(nxter.get())
    comparePlot([ATM3M3_Thin_dvv, ATM3M3_Thin_Tess_dvv], 
                [ATM3M3_Thin_vzave, ATM3M3_Thin_Tess_vzave], # dvv 和 vzave 一一对应
                'ε=5,Ma=3,tess', 'Time', names[ele], ele, legends=['normal', 'tess'])
    plt.savefig(f'Ma=3,H=0km - {names[ele]}.png') # 保存文件
    plt.show()