from abaqus import *
from abaqusConstants import *
import os
import re
import uuid
import zipfile
import shutil

source_dir = r'C:\SWTemp\jobs'
workdir = source_dir

suffix = uuid.uuid4().hex[:8]
model_name = 'Paper13_Normal_' + suffix
job_name = 'Paper13_Normal_' + suffix

input_member = 'plate_tension_compression_V1.inp'
umat_member = 'umat_Prandtl1D_creep.f'
archive_name = 'materials-18-02512-s001.zip'

inp_path = os.path.join(source_dir, input_member)
umat_path = os.path.join(source_dir, umat_member)
archive_path = os.path.join(source_dir, archive_name)

if not os.path.isdir(source_dir):
    raise RuntimeError('Source directory does not exist: ' + source_dir)

def find_file(root, wanted):
    wanted_lower = wanted.lower()
    for current, folders, files in os.walk(root):
        for filename in files:
            if filename.lower() == wanted_lower:
                return os.path.join(current, filename)
    return None

def extract_from_zip(zip_path, wanted, target_path):
    with zipfile.ZipFile(zip_path, 'r') as archive:
        matches = []
        for member in archive.namelist():
            name = member.replace('\\', '/').split('/')[-1]
            if name.lower() == wanted.lower():
                matches.append(member)

        if len(matches) != 1:
            raise RuntimeError(
                'Cannot uniquely find %s in archive.' % wanted
            )

        with open(target_path, 'wb') as stream:
            stream.write(archive.read(matches[0]))

if not os.path.isfile(inp_path):
    found = find_file(source_dir, input_member)
    if found is not None:
        inp_path = found

if not os.path.isfile(umat_path):
    found = find_file(source_dir, umat_member)
    if found is not None:
        umat_path = found

if (
    not os.path.isfile(inp_path) or
    not os.path.isfile(umat_path)
):
    if not os.path.isfile(archive_path):
        found_zip = find_file(source_dir, archive_name)
        if found_zip is not None:
            archive_path = found_zip

    if not os.path.isfile(archive_path):
        raise RuntimeError(
            'Original INP/UMAT and ZIP were not found in C:\\SWTemp\\jobs.'
        )

    if not os.path.isfile(inp_path):
        inp_path = os.path.join(
            source_dir,
            job_name + '_source.inp'
        )
        extract_from_zip(
            archive_path,
            input_member,
            inp_path
        )

    if not os.path.isfile(umat_path):
        umat_path = os.path.join(
            source_dir,
            job_name + '_UMAT.for'
        )
        extract_from_zip(
            archive_path,
            umat_member,
            umat_path
        )

if not os.path.isfile(inp_path):
    raise RuntimeError('INP file is missing.')

if not os.path.isfile(umat_path):
    raise RuntimeError('UMAT file is missing.')

with open(inp_path, 'rb') as stream:
    inp_bytes = stream.read()

with open(umat_path, 'rb') as stream:
    umat_bytes = stream.read()

if len(inp_bytes) < 100:
    raise RuntimeError('Source INP file is empty or invalid.')

if len(umat_bytes) < 100:
    raise RuntimeError('Source UMAT file is empty or invalid.')

inp_upper = inp_bytes.upper()
umat_upper = umat_bytes.upper()

if b'*NODE' not in inp_upper:
    raise RuntimeError('Source INP does not contain nodes.')

if b'*ELEMENT' not in inp_upper:
    raise RuntimeError('Source INP does not contain elements.')

if b'*STEP' not in inp_upper:
    raise RuntimeError('Source INP does not contain analysis steps.')

if b'SUBROUTINE UMAT' not in umat_upper:
    raise RuntimeError('Source Fortran file does not contain UMAT.')

if model_name in mdb.models.keys():
    del mdb.models[model_name]

model = mdb.ModelFromInputFile(
    name=model_name,
    inputFileName=inp_path
)

assembly = model.rootAssembly
assembly.regenerate()

node_count = 0
element_count = 0

for part_name in model.parts.keys():
    part = model.parts[part_name]
    node_count += len(part.nodes)
    element_count += len(part.elements)

if node_count == 0:
    raise RuntimeError('Imported model contains no mesh nodes.')

if element_count == 0:
    raise RuntimeError('Imported model contains no mesh elements.')

keyword = model.keywordBlock
keyword.synchVersions(
    storeNodesAndElements=False
)

has_initial_solution = False
first_step_index = -1

for index in range(len(keyword.sieBlocks)):
    text = keyword.sieBlocks[index].upper()

    if (
        '*INITIAL CONDITIONS' in text and
        'TYPE=SOLUTION' in text and
        'USER' in text
    ):
        has_initial_solution = True

    if (
        first_step_index < 0 and
        '*STEP' in text
    ):
        first_step_index = index

if not has_initial_solution:
    if first_step_index > 0:
        keyword.insert(
            first_step_index - 1,
            '*INITIAL CONDITIONS, TYPE=SOLUTION, USER'
        )

if job_name in mdb.jobs.keys():
    del mdb.jobs[job_name]

job = mdb.Job(
    name=job_name,
    model=model_name,
    type=ANALYSIS,
    numCpus=4,
    numDomains=4,
    multiprocessingMode=MPI,
    memory=60,
    memoryUnits=PERCENTAGE,
    userSubroutine=umat_path,
    resultsFormat=ODB
)

job.writeInput(
    consistencyChecking=OFF
)

written_inp = os.path.join(
    workdir,
    job_name + '.inp'
)

if not os.path.isfile(written_inp):
    raise RuntimeError(
        'The new INP file was not generated.'
    )

with open(written_inp, 'rb') as stream:
    written_data = stream.read().upper()

if b'*NODE' not in written_data:
    raise RuntimeError(
        'Written INP does not contain nodes.'
    )

if b'*ELEMENT' not in written_data:
    raise RuntimeError(
        'Written INP does not contain elements.'
    )

if b'*STEP' not in written_data:
    raise RuntimeError(
        'Written INP does not contain analysis steps.'
    )

cae_path = os.path.join(
    workdir,
    job_name + '.cae'
)

mdb.saveAs(
    pathName=cae_path
)

if len(session.viewports) > 0:
    viewport = session.viewports[
        session.currentViewportName
    ]

    viewport.setValues(
        displayedObject=assembly
    )

    viewport.assemblyDisplay.setValues(
        mesh=ON
    )

    viewport.view.setValues(
        viewOffsetX=0.0,
        viewOffsetY=0.0
    )

    viewport.view.fitView()

try:
    session.defaultOdbDisplay.commonOptions.setValues(
        deformationScaling=UNIFORM,
        uniformScaleFactor=1.0
    )
except:
    pass

print('MODEL: ' + model_name)
print('JOB: ' + job_name)
print('SOURCE INP: ' + inp_path)
print('UMAT: ' + umat_path)
print('NEW INP: ' + written_inp)
print('CAE: ' + cae_path)
print('NODES: %d' % node_count)
print('ELEMENTS: %d' % element_count)
print('DEFAULT ODB DEFORMATION SCALE: 1.0')
print('INPUT FILE WRITTEN')
print('NO ANALYSIS SUBMITTED')
print('SOUTHWEITION_JOB_PREPARED')
import os as _sw_os
try:
    print('SOUTHWEITION_CWD ' + _sw_os.getcwd())
    def _sw_job_model(_job):
        _model = getattr(_job, 'model', '')
        try:
            return str(_model.name)
        except Exception:
            return str(_model or '')
    _sw_models = list(mdb.models.keys())
    _sw_latest_model = _sw_models[-1] if _sw_models else ''
    _sw_name = ''
    for _cand in reversed(list(mdb.jobs.keys())):
        if _sw_latest_model and _sw_job_model(mdb.jobs[_cand]) == _sw_latest_model:
            _sw_name = str(_cand)
            break
    if not _sw_name:
        _sw_names = list(mdb.jobs.keys())
        if _sw_names:
            _sw_name = str(_sw_names[-1])
    if _sw_name:
        try:
            mdb.jobs[_sw_name].setValues(numCpus=4, numDomains=4)
        except Exception:
            pass
        mdb.jobs[_sw_name].writeInput(consistencyChecking=OFF)
        print('SOUTHWEITION_JOB ' + str(_sw_name))
        print('SOUTHWEITION_MODEL ' + str(_sw_latest_model))
        print('SOUTHWEITION_INP ' + _sw_os.path.abspath(str(_sw_name) + '.inp'))
    print('SOUTHWEITION_JOB_PREPARED')
except Exception:
    print('SOUTHWEITION_JOB_PREPARED')