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

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

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

  * ``unconditional_build`` - Boolean flag to force building of conditionally ignored targets
  * ``abaqus`` - String path for the Abaqus executable
"""

import copy
import pathlib

import waves

from modsim_package.python.rectangle_compression_sobol_sequence 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
parameter_study_file = build_directory / "parameter_study.h5"
simulation_constants = {
    "global_seed": 1,
    "displacement": -0.01,
}
kwargs = {"scramble": False}

# 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 Sobol Sequence
parameter_generator = waves.parameter_generators.SobolSequence(
    parameter_schema(),
    output_file=parameter_study_file,
    previous_parameter_study=parameter_study_file,
    **kwargs,
)

# 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, **simulation_constants}

    # 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
    copy_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",
    ]
    workflow.extend(
        env.CopySubstfile(
            copy_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 / "rectangle_compression.inp",
        set_name / "assembly.inp",
        set_name / "boundary.inp",
        set_name / "field_output.inp",
        set_name / "materials.inp",
        set_name / "parts.inp",
        set_name / "history_output.inp",
        set_name / "rectangle_mesh.inp",
    ]

    datacheck.extend(
        env.AbaqusSolver(
            target=[
                set_name / "rectangle_compression_DATACHECK.odb",
                set_name / "rectangle_compression_DATACHECK.dat",
                set_name / "rectangle_compression_DATACHECK.msg",
                set_name / "rectangle_compression_DATACHECK.com",
                set_name / "rectangle_compression_DATACHECK.prt",
                set_name / "rectangle_compression_DATACHECK.023",
                set_name / "rectangle_compression_DATACHECK.mdl",
                set_name / "rectangle_compression_DATACHECK.sim",
                set_name / "rectangle_compression_DATACHECK.stt",
            ],
            source=solve_source_list,
            job="rectangle_compression_DATACHECK",
            program_options="-double both -datacheck",
        )
    )

    workflow.extend(
        env.AbaqusSolver(
            target=[
                set_name / "rectangle_compression.odb",
                set_name / "rectangle_compression.dat",
                set_name / "rectangle_compression.msg",
                set_name / "rectangle_compression.com",
                set_name / "rectangle_compression.prt",
                set_name / "rectangle_compression.sta",
            ],
            source=solve_source_list,
            job="rectangle_compression",
            program_options="-double both",
        )
    )

# Comment used in tutorial code snippets: marker-6


def write_parameter_study(target, source, env):
    """`SCons Python build function`_ wrapper for the parameter generator's write() function.

    Reference: https://scons.org/doc/production/HTML/scons-user/ch17s04.html
    """
    set_count = len(source)
    if set_count < 8:
        new_schema = copy.deepcopy(parameter_schema())
        new_schema["num_simulations"] = set_count + 2
        new_generator = waves.parameter_generators.SobolSequence(
            new_schema,
            output_file=parameter_study_file,
            **kwargs,
        )
        new_generator.write()
    else:
        parameter_generator.write()
    return None


workflow.extend(
    env.Command(
        target=[parameter_study_file.name],
        source=[
            pathlib.Path(set_name) / "rectangle_compression.odb"
            for set_name in parameter_generator.parameter_study_to_dict().keys()
        ],
        action=[write_parameter_study],
        parameter_generator=parameter_generator,
    )
)

# Collector alias based on parent directory name
env.Alias(workflow_name, workflow)
env.Alias(f"{workflow_name}_datacheck", datacheck)

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)
