import os

# List of directories to search
#directories = ['/path/to/_xx/folder1', '/path/to/_xx/folder2', '/path/to/_xx/folder3']
directories = ["dense_phi090_theta000",
"dense_phi090_theta001",
"dense_phi090_theta002",
"dense_phi090_theta003",
"dense_phi090_theta004",
"dense_phi090_theta005",
"dense_phi090_theta006",
"dense_phi090_theta007",
"dense_phi090_theta008",
"dense_phi090_theta009",
"dense_phi090_theta010",
"dense_phi090_theta011",
"dense_phi090_theta012",
"dense_phi090_theta013",
"dense_phi090_theta014",
"dense_phi090_theta015",
"dense_phi090_theta016",
"dense_phi090_theta017",
"dense_phi090_theta018",
"dense_phi090_theta019",
"dense_phi090_theta020",
"dense_phi090_theta021",
"dense_phi090_theta022",
"dense_phi090_theta023",
"dense_phi090_theta024",
"dense_phi090_theta025",
"dense_phi090_theta026",
"dense_phi090_theta027",
"dense_phi090_theta028",
"dense_phi090_theta029",
"dense_phi090_theta030",
"dense_phi090_theta031",
"dense_phi090_theta032",
"dense_phi090_theta033",
"dense_phi090_theta034",
"dense_phi090_theta035",
"dense_phi090_theta036",
"dense_phi090_theta037",
"dense_phi090_theta038",
"dense_phi090_theta039",
"dense_phi090_theta040",
"dense_phi090_theta041",
"dense_phi090_theta042",
"dense_phi090_theta043",
"dense_phi090_theta044",
"dense_phi090_theta045",
"dense_phi090_theta046",
"dense_phi090_theta047",
"dense_phi090_theta048",
"dense_phi090_theta049",
"dense_phi090_theta050",
"dense_phi090_theta051",
"dense_phi090_theta052",
"dense_phi090_theta053",
"dense_phi090_theta054",
"dense_phi090_theta055",
"dense_phi090_theta056",
"dense_phi090_theta057",
"dense_phi090_theta058",
"dense_phi090_theta059",
"dense_phi090_theta060",
"dense_phi090_theta061",
"dense_phi090_theta062",
"dense_phi090_theta063",
"dense_phi090_theta064",
"dense_phi090_theta065",
"dense_phi090_theta066",
"dense_phi090_theta067",
"dense_phi090_theta068",
"dense_phi090_theta069",
"dense_phi090_theta070",
"dense_phi090_theta071",
"dense_phi090_theta072",
"dense_phi090_theta073",
"dense_phi090_theta074",
"dense_phi090_theta075",
"dense_phi090_theta076",
"dense_phi090_theta077",
"dense_phi090_theta078",
"dense_phi090_theta079",
"dense_phi090_theta080",
"dense_phi090_theta081",
"dense_phi090_theta082",
"dense_phi090_theta083",
"dense_phi090_theta084",
"dense_phi090_theta085",
"dense_phi090_theta086",
"dense_phi090_theta087",
"dense_phi090_theta088",
"dense_phi090_theta089",
"dense_phi090_theta090"]

def update_wt_in_file(file_path):
    # Read the contents of the file
    with open(file_path, 'r') as file:
        lines = file.readlines()
    
    # Modify the line containing 'Nslice_BTau_Max'
    with open(file_path, 'w') as file:
        for line in lines:
            if 'Nslice_BTau_Max' in line:
                file.write('Nslice_BTau_Max = 40000\n')
            elif 'BTauNum' in line: 
                file.write('BTauNum = 100 \n')
            else:
                file.write(line)

def update_job_archer2_file(file_path):
    # Read the contents of the file
    with open(file_path, 'r') as file:
        lines = file.readlines()

    with open(file_path, 'w') as file:
        for line in lines:
            if '#SBATCH --nodes' in line:
                file.write('#SBATCH --nodes=16 \n')
            else:
                file.write(line)

for directory in directories[:]:
    # Walk through each directory
    for dirpath, dirnames, filenames in os.walk(directory):
        # Check if 'wt.in' is in the list of filenames
        if 'wt.in' in filenames:
            # Construct full file path
            full_path = os.path.join(dirpath, 'wt.in')
            # Update the file
            update_wt_in_file(full_path)
            print(f"Updated file: {full_path}")
            #
            job_path = os.path.join(dirpath, 'job_archer2_std')
            update_job_archer2_file(job_path)
            print(f'Updated file: {job_path}')
