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

Requires the following ``SConscript(..., exports=[])``

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

  * ``eabm_package_abspath`` - String absolute path to the EABM 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
  * ``abaqus`` - String path for the Abaqus executable
"""

import pathlib

import waves

from eabm_package import rectangle_compression

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

# Set project-wide paths with os-agnostic path separators
eabm_package_abspath = pathlib.Path(env["eabm_package_abspath"])

# Simulation variables
build_directory = pathlib.Path(Dir(".").abspath)
workflow_name = build_directory.name
workflow_configuration = [env["project_configuration"], workflow_name]
model = "rectangle"
simulation_variables = rectangle_compression.nominal

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

# Geometry
journal_file = f"{model}_geometry"
journal_options = f"--width {simulation_variables['width']} --height {simulation_variables['height']}"
workflow.extend(env.AbaqusJournal(
    target=[f"{journal_file}.cae", f"{journal_file}.jnl"],
    source=[f"{eabm_package_abspath / journal_file}.py"],
    journal_options=journal_options))

# Partition
journal_file = f"{model}_partition"
journal_options = f"--width {simulation_variables['width']} --height {simulation_variables['height']}"
workflow.extend(env.AbaqusJournal(
    target=[f"{journal_file}.cae", f"{journal_file}.jnl"],
    source=[f"{eabm_package_abspath / journal_file}.py", f"{model}_geometry.cae"],
    journal_options=journal_options))

# Mesh
journal_file = f"{model}_mesh"
journal_options = f"--global-seed {simulation_variables['global_seed']}"
workflow.extend(env.AbaqusJournal(
    target=[f"{journal_file}.inp", f"{journal_file}.cae", f"{journal_file}.jnl"],
    source=[f"{eabm_package_abspath / journal_file}.py", f"{model}_partition.cae"],
    journal_options=journal_options))

# SolverPrep
abaqus_source_list = [
    eabm_package_abspath / f"{model}_compression.inp.in",
    eabm_package_abspath / "assembly.inp",
    eabm_package_abspath / "boundary.inp",
    eabm_package_abspath / "field_output.inp",
    eabm_package_abspath / "materials.inp",
    eabm_package_abspath / "parts.inp",
    eabm_package_abspath / "history_output.inp"
]
abaqus_source_list = [pathlib.Path(source_file) for source_file in abaqus_source_list]
workflow.extend(waves.scons_extensions.copy_substitute(
    abaqus_source_list,
    substitution_dictionary=waves.scons_extensions.substitution_syntax(simulation_variables)))

# Abaqus Solve
solve_source_list = [source_file.name.rstrip(".in") for source_file in abaqus_source_list]
solve_source_list.append([f"{journal_file}.inp"])
job_name = pathlib.Path(solve_source_list[0]).stem
datacheck_name = f"{job_name}_DATACHECK"
datacheck_suffixes = ("023", "mdl", "sim", "stt")
abaqus_options="-double both"
datacheck.extend(env.AbaqusSolver(
    target=[f"{datacheck_name}.{suffix}" for suffix in datacheck_suffixes],
    source=solve_source_list,
    job_name=datacheck_name,
    abaqus_options=f"{abaqus_options} -datacheck"))

workflow.extend(env.AbaqusSolver(
    target=[f"{job_name}.sta"],
    source=solve_source_list,
    job_name=job_name,
    abaqus_options=f"{abaqus_options} -cpus $({env['solve_cpus']}$)"))

# Abaqus Extract
extract_source_list = [f"{job_name}.odb"]
workflow.extend(env.AbaqusExtract(
    target=[f"{job_name}.h5"],
    source=extract_source_list))

# Post-processing
plot_name = "stress_strain_comparison"
regression_file = eabm_package_abspath / f"{workflow_name}_{plot_name}.csv"
post_processing_source = [f"{job_name}_datasets.h5"]
script_options = "--input-file " + " ".join(str(path) for path in post_processing_source)
script_options += f" --output-file ${{TARGET.file}} --x-units 'mm/mm' --y-units 'MPa'"
script_options += f" --csv-regression-file {regression_file}"
workflow.extend(env.PythonScript(
    target=[f"{plot_name}.pdf", f"{plot_name}.csv"],
    source=[str(eabm_package_abspath / "post_processing.py")] + [str(regression_file)] + post_processing_source,
    script_options=script_options))

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

# Collector alias based on build directory name
env.Alias(workflow_name, workflow)
env.Alias(f"{workflow_name}_datacheck", datacheck)
env.Alias(env["datacheck_alias"], datacheck)
env.Alias(f"{workflow_name}_archive", archive_target)

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