##################################################################
##
## AMSS-NCKU 数值相对论启动程序
## 小曲
## 2024/03/19
## 2025/01/21 修改
##
##################################################################


##################################################################

## 输出本程序的说明

import print_information

print_information.print_program_introduction()

##################################################################

## 程序开始运行前的提示

print(                                                                                )
print( " 计算即将开始，请确认在 AMSS_NCKU_Input.py 中设置了正确的参数，按回车继续！！！  " )
print( " 如果输入参数没有设置好，Ctrl+C 退出，调整 AMSS_NCKU_Input.py 中的输入参数！！！ " )

## 设定一个输入（回车），以便程序下一步运行
inputvalue = input()           
print()


##################################################################

## 导入参数文件

import AMSS_NCKU_Input as input_data


##################################################################

## 生成文件目录，用来程序运行的数据

import os
import shutil

# 根据输入文件来设置文件目录
File_directionary = os.path.join(input_data.File_directionary)   

# 如果文件目录已经存在，则去除它
shutil.rmtree(File_directionary, ignore_errors=True)

## 创建文件夹
os.mkdir(File_directionary)

## 复制 python 输入文件到程序运行目录
shutil.copy("AMSS_NCKU_Input.py", File_directionary)

# 生成文件目录，用来存放各类输出文件

output_directionary = os.path.join(File_directionary, "AMSS_NCKU_output")
os.mkdir(output_directionary)

binary_results_directionary = os.path.join(output_directionary, input_data.Output_directionary)
os.mkdir(binary_results_directionary)

figure_directionary = os.path.join(File_directionary, "figure")
os.mkdir(figure_directionary)

print(                     )
print( " 生成文件目录完成 " )
print(                     )


##################################################################

## 输出相关的参数信息

##################################################################

import setup

## 输出相关的参数信息

setup.print_input_data( File_directionary )
setup.generate_AMSSNCKU_input()

print(                                                                             )
print( " 检查网格大小和分辨率是否满足要求 "                                           )
print( " 如果网格大小和分辨率不合适，Ctrl+C 退出，调整网格层数和每层网格格点数目！！！ " )
print( " 如果网格大小和分辨率设定无误，按回车继续！！！ "                              )
inputvalue = input()  ## 设定一个输入（回车），以便程序下一步运行
print()

setup.print_puncture_information()


##################################################################

## 根据设定的参数生成 AMSS-NCKU 程序的输入文件

##################################################################

print(                                     )
print( " 正在生成 AMSS-NCKU 程序的输入文件 " )  
print(                                     ) 

## 根据格点信息生成 cgh 的相关输入

import numerical_grid        

numerical_grid.append_AMSSNCKU_cgh_input()  

print(                                                                 )
print( " 完成 AMSS-NCKU 程序的输入文件 "                                 )    
print( " 但是 AMSS-NCKU ABE 中 TwoPuncture 相关的输入要后几步运行后追加 " )
print(                                                                 )


##################################################################

## 画出初始的格点设置

##################################################################

print(                          )
print( " 正在画出初始格点的图像 " )    
print(                          )

numerical_grid.plot_initial_grid()


##################################################################

##  根据算法和参数生成 AMSS-NCKU 的宏文件

##################################################################

print(                                         ) 
print( " 根据算法和参数生成 AMSS-NCKU 的宏文件 " )        
print(                                         )

import generate_macrodef

generate_macrodef.generate_macrodef_h()
print(" AMSS-NCKU 程序的宏文件 macrodef.h 已生成 ")   
     
generate_macrodef.generate_macrodef_fh()
print(" AMSS-NCKU 程序的宏文件 macrodef.fh 已生成 ")  


##################################################################

##  根据用户的要求编译 AMSS-NCKU 程序

##################################################################

print(                                                         )
print( " 准备根据要求编译并运行 AMSS-NCKU 程序，按回车继续！！！ " )
inputvalue = input()           
print()

AMSS_NCKU_source_path = "AMSS_NCKU_source"
AMSS_NCKU_source_copy = os.path.join(File_directionary, "AMSS_NCKU_source_copy")

###############################

## 如果 AMSS-NCKU 的源文件夹不存在，则创建它

# if not os.path.exists(destination_folder):
#     os.makedirs(destination_folder)

if not os.path.exists(AMSS_NCKU_source_path):
    os.makedirs(AMSS_NCKU_source_path)
    print( " 缺少 AMSS-NCKU 源文件，请将代码文件复制到 AMSS_NCKU_source 文件夹，按回车继续！！！ " )
    ## 设定一个输入（回车），以便程序下一步运行
    inputvalue = input()
    
###############################

# 拷贝 AMSS-NCKU 的源文件，准备编译
shutil.copytree(AMSS_NCKU_source_path, AMSS_NCKU_source_copy)    

# 将整个src文件夹复制到dst位置
# shutil.copytree(src, dst)

# 将生成的宏文件拷贝进 AMSS-NCKU 的代码文件夹

macrodef_h_path  = os.path.join(File_directionary, "macrodef.h") 
macrodef_fh_path = os.path.join(File_directionary, "macrodef.fh") 

shutil.copy2(macrodef_h_path,  AMSS_NCKU_source_copy)
shutil.copy2(macrodef_fh_path, AMSS_NCKU_source_copy)

# 复制文件到目标位置
# shutil.copy2(source_file_path, target_location)
# shutil.copy2保留文件的元数据，如修改时间。如果你只想复制文件内容，不保留元数据，可以使用shutil.copy。

###############################

# 编译相关程序

import makefile_and_run

## 先切换到目标文件夹
os.chdir(AMSS_NCKU_source_copy)
 
## 编译 AMSS-NCKU 的主程序 ABE/ABEGPU
makefile_and_run.makefile_ABE()

## 如果初值方法选为 TwoPuncture，编译可执行文件 TwoPunctureABE 

if (input_data.Initial_Data_Method == "Ansorg-TwoPuncture" ): 
    makefile_and_run.makefile_TwoPunctureABE()
    
###########################

## 改变当前工作目录到上一级目录
os.chdir('..')
os.chdir('..')

print()

##################################################################

## 将 AMSS-NCKU 程序可执行文件 ABE 复制到程序运行的文件夹

##################################################################

if (input_data.GPU_Calculation == "no"):
    ABE_file = os.path.join(AMSS_NCKU_source_copy, "ABE")
elif (input_data.GPU_Calculation == "yes"):
    ABE_file = os.path.join(AMSS_NCKU_source_copy, "ABEGPU")

if not os.path.exists( ABE_file ):
    print(                                                                                             )
    print( " 缺少 AMSS-NCKU 可执行文件 ABE/ABEGPU，请将 AMSS_NCKU_source 编译，编译完成后按回车继续！！！ " )
    ## 设定一个输入（回车），以便程序下一步运行
    inputvalue = input() 

## 复制可执行文件 ABE 到程序运行目录
shutil.copy2(ABE_file, output_directionary)

###########################

## 如果初值方法选为 TwoPuncture，将 AMSS-NCKU 程序可执行文件 TwoPunctureABE 复制到程序运行的文件夹

TwoPuncture_file = os.path.join(AMSS_NCKU_source_copy, "TwoPunctureABE")

if (input_data.Initial_Data_Method == "Ansorg-TwoPuncture" ):

    if not os.path.exists( TwoPuncture_file ):
        print(                                                                                                                     )
        print( " 缺少 AMSS-NCKU 可执行文件 TwoPunctureABE，请将 AMSS_NCKU_source 中的 TwoPunctureABE 编译，编译完成后按回车继续！！！ " )
        inputvalue = input() 

    ## 复制可执行文件 TwoPunctureABE 到程序运行目录
    shutil.copy2(TwoPuncture_file, output_directionary)

###########################


##################################################################

## 如果初值方法选为 TwoPuncture，生成 TwoPuncture 的输入文件

##################################################################

if (input_data.Initial_Data_Method == "Ansorg-TwoPuncture" ):

    print(                                        )
    print( " 初始值类型选取为  Ansorg-Twopuncture" )
    print(                                        )
    
    print(                                                 )
    print( " 正在生成 AMSS-NCKU TwoPuncture 程序的输入文件 " )  
    print(                                                 ) 
    
    import generate_TwoPuncture_input
    
    generate_TwoPuncture_input.generate_AMSSNCKU_TwoPuncture_input()
    
    print(                                             )
    print( " 完成 AMSS-NCKU TwoPuncture 程序的输入文件 " )    
    print(                                             )
    
    ## 生成的 AMSS-NCKU 输入文件名
    AMSS_NCKU_TwoPuncture_inputfile      = 'AMSS-NCKU-TwoPuncture.input'
    AMSS_NCKU_TwoPuncture_inputfile_path = os.path.join(File_directionary, AMSS_NCKU_TwoPuncture_inputfile)
 
    ## 复制并重命名文件
    shutil.copy2( AMSS_NCKU_TwoPuncture_inputfile_path, os.path.join(output_directionary, 'TwoPunctureinput.par') )
    
    ###########################

    ## 运行 TwoPuncture 来生成初值文件
    
    print()
    ## print( " 准备启动 AMSS-NCKU TwoPuncture 程序 ，按回车继续！！！ " )
    ## inputvalue = input()                    
    print()
    
    ## 先切换到目标文件夹
    os.chdir(output_directionary)

    ## 运行 TwoPuncture 程序
    
    makefile_and_run.run_TwoPunctureABE()
    
    ###########################
    
    ## 改变当前工作目录到上两级目录
    os.chdir('..')
    os.chdir('..')
    
##################################################################
    
## 根据 TwoPuncture 的运行结果来更新 puncture data

##################################################################

import renew_puncture_parameter
    
renew_puncture_parameter.append_AMSSNCKU_BSSN_input(File_directionary, output_directionary)


## 生成的 AMSS-NCKU 输入文件名
AMSS_NCKU_inputfile      = 'AMSS-NCKU.input'
AMSS_NCKU_inputfile_path = os.path.join(File_directionary, AMSS_NCKU_inputfile)
 
## 复制并重命名文件
shutil.copy2( AMSS_NCKU_inputfile_path, os.path.join(output_directionary, 'input.par') )


print(                                                    )
print( " 成功将生成的 AMSS-NCKU 程序输入文件复制到运行目录 " )    
print(                                                    )
    

##################################################################

##  启动 AMSS-NCKU 程序

##################################################################

print()
## print(" 准备启动 AMSS-NCKU 程序，按回车继续！！！ ")
## inputvalue = input()           
print()

## 先切换到目标文件夹
os.chdir(output_directionary)

makefile_and_run.run_ABE()

## 改变当前工作目录到上两级目录
os.chdir('..')
os.chdir('..')


##################################################################

## 将一些基本输入拷贝出来，方便进行 debug

##################################################################

## 用于输出计算所用参数的文件地址
AMSS_NCKU_error_file_path = os.path.join(binary_results_directionary, "setting.par")
## 复制并重命名文件
shutil.copy( AMSS_NCKU_error_file_path, os.path.join(output_directionary, "AMSSNCKU_setting_parameter") )

## 用于报错的文件地址
AMSS_NCKU_error_file_path = os.path.join(binary_results_directionary, "Error.log")
## 复制并重命名文件
shutil.copy( AMSS_NCKU_error_file_path, os.path.join(output_directionary, "Error.log") )

## 程序基本输出
AMSS_NCKU_BH_data         = os.path.join(binary_results_directionary, "bssn_BH.dat"        )
AMSS_NCKU_ADM_data        = os.path.join(binary_results_directionary, "bssn_ADMQs.dat"     )
AMSS_NCKU_psi4_data       = os.path.join(binary_results_directionary, "bssn_psi4.dat"      )
AMSS_NCKU_constraint_data = os.path.join(binary_results_directionary, "bssn_constraint.dat")
## 复制并重命名文件
shutil.copy( AMSS_NCKU_BH_data,         os.path.join(output_directionary, "bssn_BH.dat"        ) )
shutil.copy( AMSS_NCKU_ADM_data,        os.path.join(output_directionary, "bssn_ADMQs.dat"     ) )
shutil.copy( AMSS_NCKU_psi4_data,       os.path.join(output_directionary, "bssn_psi4.dat"      ) )
shutil.copy( AMSS_NCKU_constraint_data, os.path.join(output_directionary, "bssn_constraint.dat") )

## 程序其它输出
if (input_data.Equation_Class == "BSSN-EM"):
    AMSS_NCKU_phi1_data = os.path.join(binary_results_directionary, "bssn_phi1.dat" )
    AMSS_NCKU_phi2_data = os.path.join(binary_results_directionary, "bssn_phi2.dat" )
    shutil.copy( AMSS_NCKU_phi1_data, os.path.join(output_directionary, "bssn_phi1.dat" ) )
    shutil.copy( AMSS_NCKU_phi2_data, os.path.join(output_directionary, "bssn_phi2.dat" ) )
elif (input_data.Equation_Class == "BSSN-EScalar"):
    AMSS_NCKU_maxs_data = os.path.join(binary_results_directionary, "bssn_maxs.dat" )
    shutil.copy( AMSS_NCKU_maxs_data, os.path.join(output_directionary, "bssn_maxs.dat" ) )

##################################################################

##  对 AMSS-NCKU 程序运行结果进行画图

##################################################################

print(                                         )
print( " 准备对 AMSS-NCKU 程序运行结果进行画图 " )         
print(                                         )


import plot_xiaoqu

## 画出黑洞轨迹图
plot_xiaoqu.generate_puncture_orbit_plot(   binary_results_directionary, figure_directionary )
plot_xiaoqu.generate_puncture_orbit_plot3D( binary_results_directionary, figure_directionary )

## 画出黑洞间距随时间的变化图
plot_xiaoqu.generate_puncture_distence_plot( binary_results_directionary, figure_directionary )

## 画出引力波波形图
for i in range(input_data.Detector_Number):
    plot_xiaoqu.generate_gravitational_waveform_plot( binary_results_directionary, figure_directionary, i )

## 画出时空 ADM 质量和角动量随时间的变化图
for i in range(input_data.Detector_Number):
    plot_xiaoqu.generate_ADMmass_plot( binary_results_directionary, figure_directionary, i )

## 画出哈密顿约束违反性况的变化图
for i in range(input_data.grid_level):
    plot_xiaoqu.generate_constraint_check_plot( binary_results_directionary, figure_directionary, i )

## 对储存的二进制数据画出图像
plot_xiaoqu.generate_binary_data_plot( binary_results_directionary, figure_directionary )


##################################################################

print(                              )
print( " 程序顺利结束，谢谢您的使用 " )
print(                              )

##################################################################