#! /usr/bin/env python
"""Rectangle compression workflow

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

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

  * ``datacheck_alias`` - String for the alias collecting the datacheck workflow targets
  * ``regression_alias`` - String for the alias collecting the regression test suite targets
  * ``archive_prefix`` - String prefix for archive target(s) containing identifying project and version information
  * ``project_configuration`` - String absolute path to the project SCons configuration file
  * ``unconditional_build`` - Boolean flag to force building of conditionally ignored targets
  * ``abaqus`` - String path for the Abaqus executable
"""

import pathlib

import waves

from modsim_package.python.rectangle_compression_cartesian_product import parameter_schema

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

# Comment used in tutorial code snippets: marker-1

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

# Collect the target nodes to build a concise alias for all targets
workflow = []
datacheck = []

# Comment used in tutorial code snippets: marker-2

# 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()

# Comment used in tutorial code snippets: marker-3

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

    # Comment used in tutorial code snippets: marker-4

    # Geometry
    workflow.extend(
        env.AbaqusJournal(
            target=[set_name / "rectangle_geometry.cae", set_name / "rectangle_geometry.jnl"],
            source=["#/modsim_package/abaqus/rectangle_geometry.py"],
            subcommand_options="--width ${width} --height ${height}",
            **simulation_variables,
        )
    )

    # Partition
    workflow.extend(
        env.AbaqusJournal(
            target=[set_name / "rectangle_partition.cae", set_name / "rectangle_partition.jnl"],
            source=["#/modsim_package/abaqus/rectangle_partition.py", set_name / "rectangle_geometry.cae"],
            subcommand_options="--width ${width} --height ${height}",
            **simulation_variables,
        )
    )

    # Mesh
    workflow.extend(
        env.AbaqusJournal(
            target=[
                set_name / "rectangle_mesh.inp",
                set_name / "rectangle_mesh.cae",
                set_name / "rectangle_mesh.jnl",
            ],
            source=["#/modsim_package/abaqus/rectangle_mesh.py", set_name / "rectangle_partition.cae"],
            subcommand_options="--global-seed ${global_seed}",
            **simulation_variables,
        )
    )

    # SolverPrep
    abaqus_source_list = [
        "#/modsim_package/abaqus/rectangle_compression.inp.in",
        "#/modsim_package/abaqus/assembly.inp",
        "#/modsim_package/abaqus/boundary.inp",
        "#/modsim_package/abaqus/field_output.inp",
        "#/modsim_package/abaqus/materials.inp",
        "#/modsim_package/abaqus/parts.inp",
        "#/modsim_package/abaqus/history_output.inp",
    ]
    abaqus_source_list = [pathlib.Path(source_file) for source_file in abaqus_source_list]
    workflow.extend(
        env.CopySubstfile(
            abaqus_source_list,
            substitution_dictionary=env.SubstitutionSyntax(simulation_variables),
            build_subdirectory=set_name,
        )
    )

    # Comment used in tutorial code snippets: marker-5

    # Abaqus Solve
    solve_source_list = [set_name / source_file.name.rstrip(".in") for source_file in abaqus_source_list]
    solve_source_list.append(set_name / "rectangle_mesh.inp")
    datacheck_suffixes = ("023", "mdl", "sim", "stt")
    datacheck.extend(
        env.AbaqusPseudoBuilder(
            job=set_name / "rectangle_compression_DATACHECK",
            inp=set_name / "rectangle_compression.inp",
            extra_targets=[set_name / f"rectangle_compression_DATACHECK.{suffix}" for suffix in datacheck_suffixes],
            extra_sources=solve_source_list,
            extra_options="-datacheck",
        )
    )

    workflow.extend(
        env.AbaqusPseudoBuilder(
            job=set_name / "rectangle_compression",
            extra_sources=solve_source_list,
        )
    )

    # Extract Abaqus
    extract_source_list = [set_name / "rectangle_compression.odb"]
    workflow.extend(
        env.AbaqusExtract(
            target=[set_name / "rectangle_compression.h5"],
            source=extract_source_list,
        )
    )

# Comment used in tutorial code snippets: marker-6

# Post-processing
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 ${SOURCES[2:].abspath}"
script_options += " --output-file ${TARGET.file} --x-units mm/mm --y-units MPa"
script_options += " --parameter-study-file ${SOURCES[1].abspath}"
workflow.extend(
    env.PythonScript(
        target=["stress_strain_comparison.pdf", "stress_strain_comparison.csv"],
        source=["#/modsim_package/python/post_processing.py", parameter_study_file.name] + post_processing_source,
        subcommand_options=script_options,
    )
)

# Regression test
workflow.extend(
    env.PythonScript(
        target=["regression.yaml"],
        source=[
            "#/modsim_package/python/regression.py",
            "stress_strain_comparison.csv",
            "#/modsim_package/python/rectangle_compression_cartesian_product.csv",
        ],
        subcommand_options="${SOURCES[1:].abspath} --output-file ${TARGET.abspath}",
    )
)

# Data archival
archive_name = f"{env['archive_prefix']}-{workflow_name}.tar.bz2"
archive_target = env.Tar(
    target=archive_name,
    source=workflow + workflow_configuration,
)

# Collector alias based on parent directory name
env.Alias(workflow_name, workflow)
env.Alias(f"{workflow_name}_datacheck", datacheck)
env.Alias(env["datacheck_alias"], datacheck)
env.Alias(env["regression_alias"], datacheck)
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}_datacheck"], datacheck)
    Ignore([".", env["datacheck_alias"], env["regression_alias"]], datacheck)
