from abaqus import *
from abaqusConstants import *
import mesh
import math

modelName = 'ThreeBall_Vibration_Robust'
jobName = 'ThreeBall_Vibration_Robust_Job'

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

model = mdb.Model(name=modelName)

plateL = 300.0
plateW = 180.0
plateT = 12.0

columnR = 8.0
columnHeights = [120.0, 180.0, 240.0]
columnPos = [(70.0, 90.0), (150.0, 90.0), (230.0, 90.0)]

ballR = 20.0
plateOverlap = 2.0
ballOverlap = 3.0

frequencies = [80.0, 120.0, 180.0, 240.0, 320.0]
baseDisp = 0.15
cyclesPerStep = 5
pointsPerCycle = 25

steel = model.Material(name='Steel')
steel.Density(table=((7.85e-9,),))
steel.Elastic(table=((210000.0, 0.30),))
steel.Damping(alpha=0.20, beta=2.0e-6)

model.HomogeneousSolidSection(
    name='SteelSection',
    material='Steel',
    thickness=None
)

s = model.ConstrainedSketch(
    name='PlateSketch',
    sheetSize=500.0
)
s.rectangle(
    point1=(0.0, 0.0),
    point2=(plateL, plateW)
)

platePart = model.Part(
    name='PlatePart',
    dimensionality=THREE_D,
    type=DEFORMABLE_BODY
)
platePart.BaseSolidExtrude(
    sketch=s,
    depth=plateT
)
del model.sketches['PlateSketch']

columnParts = []

for i in range(3):
    sketchName = 'ColumnSketch_%d' % (i + 1)
    partName = 'ColumnPart_%d' % (i + 1)

    s = model.ConstrainedSketch(
        name=sketchName,
        sheetSize=100.0
    )
    s.CircleByCenterPerimeter(
        center=(0.0, 0.0),
        point1=(columnR, 0.0)
    )

    p = model.Part(
        name=partName,
        dimensionality=THREE_D,
        type=DEFORMABLE_BODY
    )

    totalColumnHeight = columnHeights[i] + plateOverlap

    p.BaseSolidExtrude(
        sketch=s,
        depth=totalColumnHeight
    )

    del model.sketches[sketchName]
    columnParts.append(p)

s = model.ConstrainedSketch(
    name='BallSketch',
    sheetSize=100.0
)

s.ConstructionLine(
    point1=(0.0, -ballR * 2.0),
    point2=(0.0, ballR * 2.0)
)

s.ArcByCenterEnds(
    center=(0.0, 0.0),
    point1=(0.0, -ballR),
    point2=(0.0, ballR),
    direction=COUNTERCLOCKWISE
)

s.Line(
    point1=(0.0, ballR),
    point2=(0.0, -ballR)
)

ballPart = model.Part(
    name='BallPart',
    dimensionality=THREE_D,
    type=DEFORMABLE_BODY
)

ballPart.BaseSolidRevolve(
    sketch=s,
    angle=360.0,
    flipRevolveDirection=OFF
)

del model.sketches['BallSketch']

a = model.rootAssembly
a.DatumCsysByDefault(CARTESIAN)

mergeInstances = []

plateInst = a.Instance(
    name='PlateRaw-1',
    part=platePart,
    dependent=ON
)
mergeInstances.append(plateInst)

for i in range(3):
    x = columnPos[i][0]
    y = columnPos[i][1]
    h = columnHeights[i]

    cInst = a.Instance(
        name='ColumnRaw-%d' % (i + 1),
        part=columnParts[i],
        dependent=ON
    )

    a.translate(
        instanceList=(cInst.name,),
        vector=(x, y, plateT - plateOverlap)
    )

    mergeInstances.append(cInst)

    bInst = a.Instance(
        name='BallRaw-%d' % (i + 1),
        part=ballPart,
        dependent=ON
    )

    ballCenterZ = plateT + h + ballR - ballOverlap

    a.translate(
        instanceList=(bInst.name,),
        vector=(x, y, ballCenterZ)
    )

    mergeInstances.append(bInst)

a.regenerate()

mergedInst = a.InstanceFromBooleanMerge(
    name='ThreeBallSystem',
    instances=tuple(mergeInstances),
    keepIntersections=ON,
    originalInstances=DELETE,
    domain=GEOMETRY
)

mergedPart = model.parts['ThreeBallSystem']

mergedPart.Set(
    name='ALL_CELLS',
    cells=mergedPart.cells[:]
)

mergedPart.SectionAssignment(
    region=mergedPart.sets['ALL_CELLS'],
    sectionName='SteelSection'
)

baseFaces = mergedPart.faces.getByBoundingBox(
    xMin=-0.1,
    xMax=plateL + 0.1,
    yMin=-0.1,
    yMax=plateW + 0.1,
    zMin=-0.1,
    zMax=0.1
)

mergedPart.Set(
    name='BASE_FACE',
    faces=baseFaces
)

mergedPart.seedPart(
    size=12.0,
    deviationFactor=0.1,
    minSizeFactor=0.1
)

mergedPart.setMeshControls(
    regions=mergedPart.cells[:],
    elemShape=TET,
    technique=FREE
)

tetType = mesh.ElemType(
    elemCode=C3D10,
    elemLibrary=STANDARD
)

mergedPart.setElementType(
    regions=(mergedPart.cells[:],),
    elemTypes=(tetType,)
)

mergedPart.generateMesh()

a.regenerate()

systemInst = a.instances['ThreeBallSystem-1']

for i in range(3):
    x = columnPos[i][0]
    y = columnPos[i][1]
    h = columnHeights[i]

    targetZ = plateT + h + 2.0 * ballR - ballOverlap

    nearestNode = systemInst.nodes[0]
    bestDist = 1.0e30

    for nd in systemInst.nodes:
        dx = nd.coordinates[0] - x
        dy = nd.coordinates[1] - y
        dz = nd.coordinates[2] - targetZ
        dist2 = dx * dx + dy * dy + dz * dz

        if dist2 < bestDist:
            bestDist = dist2
            nearestNode = nd

    nodeSeq = systemInst.nodes.sequenceFromLabels(
        labels=(nearestNode.label,)
    )

    a.Set(
        name='BALL_%d_RESPONSE' % (i + 1),
        nodes=nodeSeq
    )

model.DisplacementBC(
    name='BaseVibrationBC',
    createStepName='Initial',
    region=systemInst.sets['BASE_FACE'],
    u1=0.0,
    u2=0.0,
    u3=0.0
)

previousStep = 'Initial'
stepNames = []

for freq in frequencies:
    stepName = 'Freq_%dHz' % int(freq)
    ampName = 'Amp_%dHz' % int(freq)

    timePeriod = float(cyclesPerStep) / freq
    dt = 1.0 / (freq * pointsPerCycle)

    ampData = []

    for j in range(cyclesPerStep * pointsPerCycle + 1):
        t = j * dt
        phase = 2.0 * math.pi * freq * t
        ramp = math.sin(math.pi * t / timePeriod) ** 2
        value = ramp * math.sin(phase)
        ampData.append((t, value))

    model.TabularAmplitude(
        name=ampName,
        timeSpan=STEP,
        data=tuple(ampData)
    )

    model.ImplicitDynamicsStep(
        name=stepName,
        previous=previousStep,
        timePeriod=timePeriod,
        nlgeom=ON,
        maxNumInc=10000,
        initialInc=dt,
        minInc=dt * 1.0e-5,
        maxInc=dt
    )

    model.boundaryConditions['BaseVibrationBC'].setValuesInStep(
        stepName=stepName,
        u1=baseDisp,
        amplitude=ampName
    )

    stepNames.append(stepName)
    previousStep = stepName

if 'F-Output-1' in model.fieldOutputRequests.keys():
    del model.fieldOutputRequests['F-Output-1']

if 'H-Output-1' in model.historyOutputRequests.keys():
    del model.historyOutputRequests['H-Output-1']

model.FieldOutputRequest(
    name='Cloud_Output',
    createStepName=stepNames[0],
    variables=('S', 'U', 'V', 'A'),
    frequency=2
)

for i in range(3):
    model.HistoryOutputRequest(
        name='Ball_%d_Response' % (i + 1),
        createStepName=stepNames[0],
        variables=(
            'U1', 'U2', 'U3',
            'V1', 'V2', 'V3',
            'A1', 'A2', 'A3'
        ),
        region=a.sets['BALL_%d_RESPONSE' % (i + 1)],
        frequency=1
    )

a.regenerate()

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

job = mdb.Job(
    name=jobName,
    model=modelName,
    description='Merged three-column three-ball vibration frequency study',
    type=ANALYSIS,
    memory=90,
    memoryUnits=PERCENTAGE,
    explicitPrecision=SINGLE,
    nodalOutputPrecision=SINGLE,
    numCpus=4,
    numDomains=4
)

job.writeInput()
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')