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

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

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

  * ``abaqus_source_abspath`` - String absolute path to the EABM's Abaqus journal files
  * ``python_source_abspath`` - String absolute path to the EABM's Python 3 files
  * ``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 eabm_package.python.rectangle_compression_sobol_sequence import parameter_schema

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

# Comment used in tutorial code snippets: marker-1

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

# Simulation variables
build_directory = pathlib.Path(Dir('.').abspath)
workflow_name = build_directory.name
output_file_type = "h5"
parameter_study_file = build_directory / f"parameter_study.{output_file_type}"
previous_parameter_study = str(parameter_study_file) if parameter_study_file.exists() else None
model = "rectangle"
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,
    output_file_type=output_file_type,
    previous_parameter_study=previous_parameter_study,
    **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
    journal_file = f"{model}_geometry"
    journal_options = "--width ${width} --height ${height}"
    workflow.extend(env.AbaqusJournal(
        target=[f"{set_name / journal_file}.cae",
                f"{set_name / journal_file}.jnl"],
        source=[f"{abaqus_source_abspath / journal_file}.py"],
        journal_options=journal_options,
        width=simulation_variables['width'],
        height=simulation_variables['height']))

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

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

    # SolverPrep
    abaqus_source_list = [
        abaqus_source_abspath / f"{model}_compression.inp.in",
        abaqus_source_abspath / "assembly.inp",
        abaqus_source_abspath / "boundary.inp",
        abaqus_source_abspath / "field_output.inp",
        abaqus_source_abspath / "materials.inp",
        abaqus_source_abspath / "parts.inp",
        abaqus_source_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),
        build_subdirectory=set_name))

# Comment used in tutorial code snippets: marker-5

    # Abaqus Solve
    solve_source_list = [f"{set_name / source_file.name.rstrip('.in')}" for source_file in abaqus_source_list]
    solve_source_list.append([f"{set_name / 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"{set_name / 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"{set_name / job_name}.sta"],
        source=solve_source_list,
        job_name=job_name,
        abaqus_options=abaqus_options))

# 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,
            output_file_type=output_file_type,
            **kwargs)
        new_generator.write()
    else:
        parameter_generator.write()
    return None


workflow.extend(env.Command(
    target=[parameter_study_file.name],
    source=[f"{pathlib.Path(set_name) / job_name}.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']:
    print(f"Program 'abaqus' was not found in construction environment. Ignoring '{workflow_name}' target(s)")
    Ignore(['.', workflow_name], workflow)
    Ignore(['.', model], datacheck)
