import sys
sys.path.insert(0, '/p/user_pub/xclim/persist/software/xsearch/')
import xsearch as xs
import argparse
import subprocess
import os

def get_args():
    parser = argparse.ArgumentParser()
    parser.add_argument('--model', type=str, default='CanESM5')
    parser.add_argument('--scenario', type=str, default='historical')
    parser.add_argument('--variability', type=str, choices=['PNA', 'NAM', 'SAM', 'NAO'])
    args = vars(parser.parse_args())
    return args

def modify_template_file(template_path, new_file_path, new_lines):
    # Read the content of the template file
    with open(template_path, 'r') as file:
        content = file.readlines()

    # Add new lines to the beginning
    new_content = new_lines + content

    # Save the modified content to a new file
    if os.path.exists(new_file_path):
        os.remove(new_file_path)
    with open(new_file_path, 'w') as file:
        file.writelines(new_content)


args = get_args()
model = args['model']
scenario = args['scenario']
variability = args['variability']
var = 'psl'
print('model: ', model, ' \nscenario: ', scenario)
files = list(xs.findPaths(scenario, var, 'mon', model = model).keys())
print(sorted(files))
print(len(files))
template = 'parameter_template_psl.py'
for file in files:
    member = file.split('/')[-6]
    # print(member)
    if 'i1p1f' in member:
        print(member)
        new_lines = [
            "scenario = "+"'"+scenario+"'\n",
            "gcm_name = "+"'"+model+"'\n",
            "realization = "+"'"+member+"'\n",
            "modpath = "+"'"+file+"'\n",
            "variability_mode = "+"'"+variability+"'\n"
        ]
        modify_template_file(template, variability+'-'+model+'-'+member+'.py', new_lines)

        command = f"variability_modes_driver.py -p "+variability+"-"+model+"-"+member+".py > "+variability+"-"+model+"-"+member+".out"
        subprocess.run(command, shell=True)