#! /usr/bin/env python

import os
import pathlib

import waves

# Accept command line options with fall back default values
AddOption(
    "--build-dir",
    dest="build_dir",
    default="build",
    nargs=1,
    type="string",
    action="store",
    metavar="DIR",
    help="SCons build (variant) root directory. Relative or absolute path. (default: '%default')",
)
# Python optparse appends to the default list instead of overriding. Must implement default/override ourselves.
default_truchas_commands = [
    "truchas",
]
AddOption(
    "--truchas-command",
    dest="truchas_command",
    nargs=1,
    type="string",
    action="append",
    metavar="COMMAND",
    help=f"Override for the Truchas command. Repeat to specify more than one (default: {default_truchas_commands})",
)
default_cubit_commands = [
    "/apps/Cubit-16.12/cubit",
    "/usr/projects/ea/Cubit/Cubit-16.12/cubit",
    "cubit",
]
AddOption(
    "--cubit-command",
    dest="cubit_command",
    nargs=1,
    type="string",
    action="append",
    metavar="COMMAND",
    help=f"Override for the Cubit command. Repeat to specify more than one (default: {default_cubit_commands})",
)

# Inherit user's full environment and set project options
env = waves.scons_extensions.WAVESEnvironment(
    ENV=os.environ.copy(),
    build_dir=pathlib.Path(GetOption("build_dir")),
    truchas_commands=GetOption("truchas_command"),
    cubit_commands=GetOption("cubit_command"),
)
env["parameter_study_directory"] = env["build_dir"] / "parameter_studies"

# Always copy (no sym-links) when duplicating
env.SetOption("duplicate", "copy")

# Empty defaults list to avoid building all simulation targets by default
env.Default()

# Find required programs for conditional target ignoring and absolute path for use in target actions
env["TRUCHAS_PROGRAM"] = env.AddProgram(
    env["truchas_commands"] if env["truchas_commands"] is not None else default_truchas_commands
)
env["MPIRUN_PROGRAM"] = env.AddProgram(["mpirun"])

# Add project builders
env.Append(
    BUILDERS={
        "Truchas": waves.scons_extensions.truchas_builder_factory(
            program="mpirun" if env["MPIRUN_PROGRAM"] is not None else "",
            program_options="-np 1" if env["MPIRUN_PROGRAM"] is not None else "",
            subcommand="${TRUCHAS_PROGRAM}",
        )
    }
)

# Cubit Python requires a separate construction environment to avoid putting Cubit distributed Python packages in the
# parent Python 3 environment's PYTHONPATH
envCubit = Environment(ENV=os.environ.copy())
envCubit.AddMethod(waves.scons_extensions.add_cubit_python, "AddCubitPython")
envCubit["python"] = envCubit.AddCubitPython(
    env["cubit_commands"] if env["cubit_commands"] is not None else default_cubit_commands,
)
envCubit.Append(
    BUILDERS={
        "PythonScript": waves.scons_extensions.python_builder_factory(program=envCubit["python"]),
    }
)

# Print captured STDOUT/STDERR for failed task(s)
env.PrintBuildFailures()

# Define parameter studies
nominal_parameters = {
    "width": 1.0,
    "height": 1.0,
    "global_seed": 1.0,
    "displacement": -0.01,
}
mesh_convergence_parameter_study_file = env["parameter_study_directory"] / "mesh_convergence.h5"
mesh_convergence_parameter_generator = waves.parameter_generators.CartesianProduct(
    {
        "width": [1.0],
        "height": [1.0],
        "global_seed": [1.0, 0.5, 0.25, 0.125],
        "displacement": [-0.01],
    },
    output_file=mesh_convergence_parameter_study_file,
    previous_parameter_study=mesh_convergence_parameter_study_file,
)
mesh_convergence_parameter_generator.write()

# Add workflow(s)
workflow_configurations = [
    ("nominal", nominal_parameters),
    ("mesh_convergence", mesh_convergence_parameter_generator),
]
for study_name, study_definition in workflow_configurations:
    env.ParameterStudySConscript(
        "SConscript",
        variant_dir=env["build_dir"] / study_name,
        exports={"env": env, "envCubit": envCubit, "alias": study_name},
        study=study_definition,
        subdirectories=True,
        duplicate=True,
    )

# List all aliases in help message.
# This must come *after* all expected Alias definitions and SConscript files.
env.ProjectHelp()
