##################################################################
##
## AMSS-NCKU Plot-Only Restart Script
## Author: Xiaoqu / Claude
## 2026/05/12
##
## This script checks for existing output data from AMSS_NCKU_Program.py.
## If data exists, it skips all computation and goes directly to plotting,
## saving time when plotting was interrupted.
## If no data is found, it exits with a message.
##
##################################################################

## Guard against re-execution by multiprocessing child processes.
if __name__ != '__main__':
    import sys as _sys
    _sys.exit(0)


import os
import sys

import AMSS_NCKU_Input as input_data

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

## Construct paths from input configuration
File_directory         = os.path.join(input_data.File_directory)
output_directory       = os.path.join(File_directory, "AMSS_NCKU_output")
binary_results_directory = os.path.join(output_directory, input_data.Output_directory)
figure_directory       = os.path.join(File_directory, "figure")

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

## Check whether the required output data files exist

required_files = [
    os.path.join(binary_results_directory, "bssn_BH.dat"),
    os.path.join(binary_results_directory, "bssn_ADMQs.dat"),
    os.path.join(binary_results_directory, "bssn_psi4.dat"),
    os.path.join(binary_results_directory, "bssn_constraint.dat"),
]

missing_files = [f for f in required_files if not os.path.exists(f)]

if missing_files:
    print(" No existing AMSS_NCKU_Program.py output data found. ")
    print(" The following required files are missing: ")
    for f in missing_files:
        print(f"   {f}")
    print()
    print(" Please run AMSS_NCKU_Program.py first to generate the simulation data. ")
    print(" Exiting. ")
    sys.exit(1)

print(" Found existing AMSS_NCKU_Program.py output data. "          )
print(" Skipping all computation and going directly to plotting. " )
print()

## Ensure the figure directory exists (it should, but be safe)
os.makedirs(figure_directory, exist_ok=True)

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

## Plot the AMSS-NCKU program results

import plot_xiaoqu
import plot_GW_strain_amplitude_xiaoqu
from parallel_plot_helper import run_plot_tasks_parallel

plot_tasks = []

## Plot black hole trajectory
plot_tasks.append((plot_xiaoqu.generate_puncture_orbit_plot,   (binary_results_directory, figure_directory)))
plot_tasks.append((plot_xiaoqu.generate_puncture_orbit_plot3D, (binary_results_directory, figure_directory)))

## Plot black hole separation vs. time
plot_tasks.append((plot_xiaoqu.generate_puncture_distence_plot, (binary_results_directory, figure_directory)))

## Plot gravitational waveforms (psi4 and strain amplitude)
for i in range(input_data.Detector_Number):
    plot_tasks.append((plot_xiaoqu.generate_gravitational_wave_psi4_plot, (binary_results_directory, figure_directory, i)))
    plot_tasks.append((plot_GW_strain_amplitude_xiaoqu.generate_gravitational_wave_amplitude_plot, (binary_results_directory, figure_directory, i)))

## Plot ADM mass evolution
for i in range(input_data.Detector_Number):
    plot_tasks.append((plot_xiaoqu.generate_ADMmass_plot, (binary_results_directory, figure_directory, i)))

## Plot Hamiltonian constraint violation over time
for i in range(input_data.grid_level):
    plot_tasks.append((plot_xiaoqu.generate_constraint_check_plot, (binary_results_directory, figure_directory, i)))

run_plot_tasks_parallel(plot_tasks)

## Plot stored binary data (runs serially, not in the parallel pool)
plot_xiaoqu.generate_binary_data_plot(binary_results_directory, figure_directory)

print()
print(" Plotting completed successfully. ")
print()