#! /usr/bin/env python
"""Rectangle mesh convergence workflow

Requires the following ``SConscript(..., exports={})``

* ``env`` - The SCons construction environment with the following required keys

  * ``modsim_package_abspath`` - Absolute path to the project package
  * ``solve_cpus`` - Integer number of cpus to use in the Abaqus solve task
  * ``datacheck_alias`` - String for the alias collecting the datacheck workflow targets
  * ``regression_alias`` - String for the alias collecting the regression workflow targets
  * ``abaqus`` - String path for the Abaqus executable
"""

import pathlib

import waves

from modsim_package import rectangle_compression

# Inherit the parent construction environment
Import("env")

# Set project-wide paths with os-agnostic path separators
modsim_package_abspath = env["modsim_package_abspath"]

# Simulation variables
build_directory = pathlib.Path(Dir(".").abspath)
workflow_name = build_directory.name
workflow_configuration = [env["project_configuration"], workflow_name]
parameter_schema = rectangle_compression.mesh_convergence()
parameter_study_file = build_directory / "mesh_convergence.h5"

# Collect the target nodes to build concise target alias(es)
artifacts = []
workflow = []

# Parameter Study with Cartesian Product
parameter_generator = waves.parameter_generators.CartesianProduct(
    parameter_schema,
    output_file=parameter_study_file,
    previous_parameter_study=parameter_study_file,
)
parameter_generator.write()

# Parameterized targets must live inside current parameters for loop
for set_name, parameters in parameter_generator.parameter_study_to_dict().items():
    set_name = pathlib.Path(set_name)

    # Rectangle part
    artifacts.extend(
        SConscript(
            "rectangle",
            variant_dir=set_name,
            exports={"env": env, "parameters": parameters},
        )
    )

    # Rectangle compression simulation
    artifacts.extend(
        SConscript(
            "rectangle_compression",
            variant_dir=set_name,
            exports={"env": env, "parameters": parameters, "workflow_name": workflow_name},
        )
    )

# Post-processing
plot_name = "stress_strain_comparison"
post_processing_source = [
    pathlib.Path(set_name) / "rectangle_compression_datasets.h5"
    for set_name in parameter_generator.parameter_study_to_dict().keys()
]
script_options = "--input-file ${post_processing_source}"
script_options += " --output-file ${TARGET.file} --x-units 'mm/mm' --y-units 'MPa'"
script_options += f" --parameter-study-file {parameter_study_file.name}"
workflow.extend(
    env.PythonScript(
        target=[f"{plot_name}.pdf", f"{plot_name}.csv"],
        source=[modsim_package_abspath / "post_processing.py", parameter_study_file.name] + post_processing_source,
        subcommand_options=script_options,
        post_processing_source=post_processing_source,
    )
)

plot_name = "mesh_convergence_stress"
selection_dict = modsim_package_abspath / f"{plot_name}.yaml"
script_options = "--input-file ${post_processing_source}"
script_options += " --output-file ${TARGET.file} --x-units 'mm' --y-units 'MPa' --x-var 'global_seed' --y-var 'S'"
script_options += f" --parameter-study-file {parameter_study_file.name}"
script_options += f" --selection-dict {selection_dict}"
workflow.extend(
    env.PythonScript(
        target=[f"{plot_name}.pdf", f"{plot_name}.csv"],
        source=[modsim_package_abspath / "post_processing.py", parameter_study_file.name] + post_processing_source,
        subcommand_options=script_options,
        post_processing_source=post_processing_source,
    )
)

# Data archival
artifacts.extend(workflow)
archive_name = f"{env['project_name']}-{workflow_name}-{env['version']}"
archive_target = env.Tar(
    target=archive_name,
    source=artifacts + workflow_configuration,
)

# Collector alias based on build directory name
env.Alias(workflow_name, workflow)
env.Alias(f"{workflow_name}-archive", archive_target)

if not env["unconditional_build"] and not env["ABAQUS_PROGRAM"]:
    print(f"Program 'abaqus' was not found in construction environment. Ignoring '{workflow_name}' target(s)")
    Ignore([".", workflow_name], workflow)
    Ignore([".", f"{workflow_name}-archive"], archive_target)
