from basic_function import format_parser
from basic_function import packaged_function
from basic_function import conformer_search
import time
import os
import sys

if __name__ == '__main__':

    time_start = time.time()

    # --- Load configuration from generation_config.sh ---
    script_dir = os.path.dirname(os.path.abspath(__file__))
    config_path = sys.argv[1] if len(sys.argv) > 1 else os.path.join(script_dir, "generation_config.sh")
    cfg = {}
    exec(open(config_path).read(), cfg)

    output_dir = cfg["OUTPUT_DIR"]
    work_dir = os.path.join(script_dir, output_dir)
    space_group_list = [int(x) for x in cfg["SPACE_GROUP_LIST"].split(",")]
    add_name = cfg["ADD_NAME"]
    num_generation = int(cfg["NUM_GENERATION"])
    max_workers = int(cfg["MAX_WORKERS"])

    # ##############################################################################################
    # # conformer search (uncomment to use)
    conformer_search.conformer_search(cfg["SMILES"], os.path.join(work_dir, "molecule_1"),
    num_conformers=int(cfg["GENERATE_CONFORMERS"]), max_attempts=10000, rms_thresh=0.1)
    # ##############################################################################################

    # ##############################################################################################
    # single crystal structure generate with Z'=1

    molecule1 = format_parser.read_xyz_file(os.path.join(work_dir, "molecule_1/conformers/conformer_0.xyz"))
    packaged_function.CSP_generater_parallel([molecule1], work_dir,
        need_structure=num_generation, space_group_list=space_group_list,
        add_name=f"{add_name}_C1", max_workers=max_workers, start_seed=1)
    # ##############################################################################################

    # ##############################################################################################
    # single crystal structure generate with Z'=2 (uncomment to use)

    # molecule1 = format_parser.read_xyz_file(os.path.join(work_dir, "molecule_1/conformers/conformer_0.xyz"))
    # packaged_function.CSP_generater_parallel([molecule1, molecule1], work_dir,
    #     need_structure=num_generation, space_group_list=space_group_list,
    #     add_name=f"{add_name}_C1", max_workers=max_workers, start_seed=1)
    # ##############################################################################################

    # ##############################################################################################
    # co-crystal structure generate (uncomment to use)

    # molecule1 = format_parser.read_xyz_file(os.path.join(work_dir, "molecule_1/conformers/conformer_0.xyz"))
    # molecule2 = format_parser.read_xyz_file(os.path.join(work_dir, "molecule_2/conformers/conformer_0.xyz"))
    # packaged_function.CSP_generater_parallel([molecule1, molecule2], work_dir,
    #     need_structure=num_generation, space_group_list=space_group_list,
    #     add_name=f"{add_name}_C1", max_workers=max_workers, start_seed=1)
    # ##############################################################################################

    time_end = time.time()
    print('time cost', time_end - time_start, 's')