from abaqus import mdb
from abaqusConstants import OFF
import os as _sw_os
_sw_before_jobs = {str(n): mdb.jobs[n] for n in mdb.jobs.keys()}
_sw_written_jobs = []
def _sw_record_write(_job, *args, **kwargs):
    _answer = _job.writeInput(*args, **kwargs)
    _name = str(getattr(_job, 'name', ''))
    if not _name or _name not in mdb.jobs.keys():
        _names = [str(n) for n in mdb.jobs.keys() if mdb.jobs[n] is _job]
        if len(_names) != 1:
            raise RuntimeError('Cannot identify the job which wrote the INP')
        _name = _names[0]
    _sw_written_jobs.append((_name, _sw_os.path.abspath(_name + '.inp')))
    return _answer

from abaqus import *
from abaqusConstants import *
import mesh
import math
import os
PREFIX = 'CraneJointNative'
GLOBAL_SIZE = 22.0
LOCAL_SIZE = 9.0
DURATION = 0.06
TARGET_DT = 2e-06
DENSITY = 7.85e-09
STAGE = 'Initialize'

def stage(value):
    global STAGE
    STAGE = value
    print('STAGE: ' + value)
index = 1
while True:
    base = PREFIX + '_%03d' % index
    available = True
    for case in ('Compression', 'Tension'):
        name = base + '_' + case
        if name in mdb.models.keys() or name in mdb.jobs.keys():
            available = False
        if os.path.exists(name + '.inp'):
            available = False
    if available:
        break
    index += 1

def revolve(model, name, start, segments, angle):
    stage(model.name + ': geometry ' + name)
    sk = model.ConstrainedSketch(name=name + '_Sketch', sheetSize=1000.0)
    sk.ConstructionLine(point1=(0.0, -400.0), point2=(0.0, 400.0))
    previous = start
    for item in segments:
        if item[0] == 'L':
            sk.Line(point1=previous, point2=item[1])
            previous = item[1]
        else:
            sk.ArcByCenterEnds(center=item[1], point1=previous, point2=item[2], direction=item[3])
            previous = item[2]
    sk.Line(point1=previous, point2=start)
    part = model.Part(name=name, dimensionality=THREE_D, type=DEFORMABLE_BODY)
    part.BaseSolidRevolve(sketch=sk, angle=angle, flipRevolveDirection=OFF)
    del model.sketches[sk.name]
    return part

def material(model, name, modulus, strength):
    mat = model.Material(name=name)
    mat.Density(table=((DENSITY,),))
    mat.Elastic(table=((modulus, 0.3),))
    mat.Plastic(table=((strength, 0.0),))
    model.HomogeneousSolidSection(name=name + '_Section', material=name, thickness=None)

def mesh_solid(part, section):
    stage('Mesh: ' + part.name)
    region = part.Set(name='SOLID', cells=part.cells[:])
    part.SectionAssignment(region=region, sectionName=section)
    part.setMeshControls(regions=part.cells[:], elemShape=TET, technique=FREE)
    part.setElementType(regions=(part.cells[:],), elemTypes=(mesh.ElemType(elemCode=C3D10M, elemLibrary=EXPLICIT),))
    part.seedPart(size=GLOBAL_SIZE, deviationFactor=0.12, minSizeFactor=0.15)
    selected = []
    for edge in part.edges:
        if 45.0 <= abs(edge.pointOn[0][1]) <= 135.0:
            selected.append(edge)
    if selected:
        part.seedEdgeBySize(edges=tuple(selected), size=LOCAL_SIZE, deviationFactor=0.12, minSizeFactor=0.15, constraint=FINER)
    part.generateMesh()
    if len(part.elements) == 0:
        raise RuntimeError('Empty mesh: ' + part.name)

def exterior_records(instance):
    records = []
    for face in instance.elements.getExteriorFaces():
        nodes = face.getNodes()
        center = [0.0, 0.0, 0.0]
        labels = []
        for node in nodes:
            labels.append(node.label)
            for k in range(3):
                center[k] += float(node.coordinates[k])
        for k in range(3):
            center[k] /= float(len(nodes))
        records.append((str(face.face), face.label, tuple(center), tuple(labels)))
    if not records:
        raise RuntimeError('No exterior faces: ' + instance.name)
    return records

def mesh_surface(assembly, instance, name, records):
    groups = {}
    for (side, label, center, nodes) in records:
        if side not in groups:
            groups[side] = {}
        groups[side][label] = True
    arguments = {}
    for (side, labels) in groups.items():
        if side not in ('FACE1', 'FACE2', 'FACE3', 'FACE4'):
            raise RuntimeError('Unexpected tetrahedral face: ' + side)
        sequence = instance.elements.sequenceFromLabels(labels=tuple(sorted(labels.keys())))
        arguments['face' + side[-1] + 'Elements'] = sequence
    if not arguments:
        raise RuntimeError('Empty mesh surface: ' + name)
    return assembly.Surface(name=name, **arguments)

def select_patch(records, point):
    for radius in (22.0, 30.0, 40.0, 55.0):
        (selected, labels) = ([], {})
        for record in records:
            center = record[2]
            distance2 = (center[0] - point[0]) ** 2 + (center[1] - point[1]) ** 2 + (center[2] - point[2]) ** 2
            if distance2 <= radius * radius:
                selected.append(record)
                for label in record[3]:
                    labels[label] = True
        if len(selected) >= 3 and len(labels) >= 6:
            return (selected, radius)
    raise RuntimeError('No adequate local coupling patch')

def couple(model, name, control, surface):
    model.Coupling(name=name, controlPoint=control, surface=surface, influenceRadius=WHOLE_SURFACE, couplingType=DISTRIBUTING, weightingMethod=UNIFORM, u1=ON, u2=ON, u3=ON, ur1=ON, ur2=ON, ur3=ON)

def build(case, force):
    name = base + '_' + case
    stage(name + ': create clean model')
    model = mdb.Model(name=name, modelType=STANDARD_EXPLICIT)
    model.setValues(description='Approximate reconstruction of doi:10.1038/s41598-024-54351-y. Native mesh surfaces and coupling objects. Assumed bore, fit, friction, equivalent unpreloaded bolts and perfect plasticity.')
    material(model, 'Q550', 212000.0, 550.0)
    material(model, 'BushingSteel', 210000.0, 760.0)
    jp = [('L', (117.5, 0.0)), ('L', (117.5, 53.0)), ('A', (105.5, 53.0), (105.5, 65.0), COUNTERCLOCKWISE), ('L', (99.5, 65.0)), ('A', (99.5, 72.0), (92.5, 72.0), CLOCKWISE), ('L', (92.5, 110.0)), ('A', (112.5, 110.0), (112.5, 130.0), CLOCKWISE), ('L', (149.5, 130.0)), ('L', (149.5, 230.0)), ('L', (120.0, 240.0)), ('L', (120.0, 250.0)), ('L', (30.0, 250.0))]
    sp = [('L', (149.5, -129.5)), ('L', (149.5, 129.5)), ('L', (113.0, 129.5)), ('L', (92.8, 109.0)), ('L', (92.8, 65.3)), ('L', (105.7, 65.3)), ('A', (105.7, 53.3), (117.7, 53.3), CLOCKWISE), ('L', (117.7, -53.3)), ('A', (105.7, -53.3), (105.7, -65.3), CLOCKWISE), ('L', (92.8, -65.3)), ('L', (92.8, -109.0))]
    joint = revolve(model, 'Joint', (30.0, 0.0), jp, 360.0)
    bush = revolve(model, 'Bushing', (113.0, -129.5), sp, 179.0)
    mesh_solid(joint, 'Q550_Section')
    mesh_solid(bush, 'BushingSteel_Section')
    a = model.rootAssembly
    a.DatumCsysByDefault(CARTESIAN)
    upper = a.Instance(name='UPPER', part=joint, dependent=ON)
    lower = a.Instance(name='LOWER', part=joint, dependent=ON)
    a.rotate(instanceList=('LOWER',), axisPoint=(0.0, 0.0, 0.0), axisDirection=(1.0, 0.0, 0.0), angle=180.0)
    for k in range(2):
        nm = 'BUSH_%d' % (k + 1)
        a.Instance(name=nm, part=bush, dependent=ON)
        a.rotate(instanceList=(nm,), axisPoint=(0.0, 0.0, 0.0), axisDirection=(0.0, 1.0, 0.0), angle=0.5 + 180.0 * k)
    a.regenerate()
    (records, halves, surfaces) = ({}, {}, [])
    for nm in ('UPPER', 'LOWER', 'BUSH_1', 'BUSH_2'):
        inst = a.instances[nm]
        stage(name + ': exterior mesh surface ' + nm)
        records[nm] = exterior_records(inst)
        surfaces.append(mesh_surface(a, inst, 'CONTACT_' + nm, records[nm]))
        if nm.startswith('BUSH'):
            total = 0.0
            for node in inst.nodes:
                total += float(node.coordinates[2])
            halves[1 if total > 0.0 else -1] = inst
    if len(halves) != 2:
        raise RuntimeError('Bushing orientation check failed')
    gap = 130.0 * math.tan(math.radians(0.5))
    mat = model.Material(name='EquivalentBoltSteel')
    mat.Density(table=((DENSITY * 80.0 / (2.0 * gap),),))
    mat.Elastic(table=((210000.0 * 2.0 * gap / 80.0, 0.3),))
    model.CircularProfile(name='BoltProfile', r=12.0)
    model.BeamSection(name='BoltSection', integration=DURING_ANALYSIS, poissonRatio=0.3, profile='BoltProfile', material='EquivalentBoltSteel')
    bolt = model.Part(name='Bolt', dimensionality=THREE_D, type=DEFORMABLE_BODY)
    bolt.WirePolyLine(points=(((0.0, 0.0, -gap), (0.0, 0.0, gap)),), mergeType=IMPRINT, meshable=ON)
    br = bolt.Set(name='BEAM', edges=bolt.edges[:])
    bolt.SectionAssignment(region=br, sectionName='BoltSection')
    bolt.assignBeamSectionOrientation(region=br, method=N1_COSINES, n1=(1.0, 0.0, 0.0))
    bolt.setElementType(regions=(bolt.edges[:],), elemTypes=(mesh.ElemType(elemCode=B31, elemLibrary=EXPLICIT),))
    bolt.seedEdgeByNumber(edges=bolt.edges[:], number=1, constraint=FIXED)
    bolt.generateMesh()
    number = 0
    for xx in (-130.0, 130.0):
        for yy in (-100.0, 100.0):
            number += 1
            nm = 'BOLT_%d' % number
            bi = a.Instance(name=nm, part=bolt, dependent=ON)
            a.translate(instanceList=(nm,), vector=(xx, yy, 0.0))
            a.regenerate()
            for sign in (-1, 1):
                suffix = '%d_%d' % (number, sign + 1)
                stage(name + ': native bolt coupling ' + suffix)
                labels = []
                for node in bi.nodes:
                    if abs(node.coordinates[2] - sign * gap) < 0.0001:
                        labels.append(node.label)
                if len(labels) != 1:
                    raise RuntimeError('Invalid bolt endpoint: ' + suffix)
                end = a.Set(name='END_' + suffix, nodes=bi.nodes.sequenceFromLabels(labels=tuple(labels)))
                target = halves[sign]
                (patch, radius) = select_patch(records[target.name], (xx, yy, sign * gap))
                surface = mesh_surface(a, target, 'PATCH_' + suffix, patch)
                couple(model, 'BoltCoupling_' + suffix, end, surface)
                print('%s: facets=%d radius=%.1f' % (suffix, len(patch), radius))
    bottom_nodes = lower.nodes.getByBoundingBox(xMin=-160.0, xMax=160.0, yMin=-250.001, yMax=-249.999, zMin=-160.0, zMax=160.0)
    if len(bottom_nodes) == 0:
        raise RuntimeError('Empty fixed end')
    fixed = a.Set(name='FIXED_END', nodes=bottom_nodes)
    top = []
    for record in records['UPPER']:
        if abs(record[2][1] - 250.0) < 0.0001:
            top.append(record)
    loaded = mesh_surface(a, upper, 'LOADED_END', top)
    rp = a.ReferencePoint(point=(0.0, 250.0, 0.0))
    control = a.Set(name='LOAD_RP', referencePoints=(a.referencePoints[rp.id],))
    couple(model, 'LoadDistribution', control, loaded)
    model.EncastreBC(name='FixBase', createStepName='Initial', region=fixed)
    model.DisplacementBC(name='AxialGuide', createStepName='Initial', region=control, u1=0.0, u2=UNSET, u3=0.0, ur1=0.0, ur2=0.0, ur3=0.0)
    model.ExplicitDynamicsStep(name='Load', previous='Initial', timePeriod=DURATION)
    model.steps['Load'].setValues(massScaling=((SEMI_AUTOMATIC, MODEL, AT_BEGINNING, 0.0, TARGET_DT, BELOW_MIN, 0, 0, 0.0, 0.0, 0, None),))
    prop = model.ContactProperty('SteelContact')
    prop.NormalBehavior(pressureOverclosure=HARD, allowSeparation=ON)
    prop.TangentialBehavior(formulation=PENALTY, directionality=ISOTROPIC, table=((0.15,),), maximumElasticSlip=FRACTION, fraction=0.005)
    pairs = []
    for i in range(len(surfaces)):
        for j in range(i, len(surfaces)):
            pairs.append((surfaces[i], surfaces[j]))
    contact = model.ContactExp(name='GeneralContact', createStepName='Initial')
    contact.includedPairs.setValuesInStep(stepName='Initial', useAllstar=OFF, addPairs=tuple(pairs))
    contact.contactPropertyAssignments.appendInStep(stepName='Initial', assignments=((GLOBAL, SELF, 'SteelContact'),))
    model.SmoothStepAmplitude(name='Ramp', timeSpan=STEP, data=((0.0, 0.0), (0.8 * DURATION, 1.0), (DURATION, 1.0)))
    model.ConcentratedForce(name='AxialForce', createStepName='Load', region=control, cf2=force, amplitude='Ramp')
    model.fieldOutputRequests['F-Output-1'].setValues(variables=('S', 'PEEQ', 'U', 'RF', 'CSTRESS'), numIntervals=30)
    model.historyOutputRequests['H-Output-1'].setValues(variables=('ALLIE', 'ALLKE', 'ALLAE', 'ALLVD', 'ALLFD', 'ALLWK', 'ETOTAL'), numIntervals=150)
    stage(name + ': write input')
    a.regenerate()
    job = mdb.Job(name=name, model=name, type=ANALYSIS, numCpus=4, numDomains=4, memory=80, memoryUnits=PERCENTAGE, explicitPrecision=DOUBLE, nodalOutputPrecision=SINGLE)
    _sw_record_write(job, consistencyChecking=ON)
    with open(name + '.inp', 'r') as stream:
        for line in stream:
            if line.lstrip().upper().startswith('*CONFLICTS'):
                raise RuntimeError('Conflict keyword detected: ' + name)
    print('Input conflict check passed: ' + name)
try:
    for (case, force) in (('Compression', -6864655.0), ('Tension', 3922660.0)):
        build(case, force)
    mdb.saveAs(pathName=base + '.cae')
    print('Ready for Abaqus/CAE Job Manager.')
except Exception:
    print('FAILED STAGE: ' + STAGE)
    raise
print('SOUTHWEITION_JOB_PREPARED')

if not _sw_written_jobs:
    _sw_candidates = [n for n in mdb.jobs.keys() if n not in _sw_before_jobs or mdb.jobs[n] is not _sw_before_jobs[n]]
    if len(_sw_candidates) != 1:
        raise RuntimeError('Cannot identify this script job: call writeInput() on exactly the intended job')
    _sw_record_write(mdb.jobs[_sw_candidates[0]], consistencyChecking=OFF)
_sw_name, _sw_inp = _sw_written_jobs[-1]
if not _sw_os.path.isfile(_sw_inp):
    raise RuntimeError('This script did not produce its INP: ' + _sw_inp)
print('SOUTHWEITION_CWD ' + _sw_os.path.dirname(_sw_inp))
print('SOUTHWEITION_JOB ' + _sw_name)
print('SOUTHWEITION_MODEL ' + str(getattr(mdb.jobs[_sw_name], 'model', '')))
print('SOUTHWEITION_INP ' + _sw_inp)
print('SOUTHWEITION_JOB_PREPARED')