# -*- coding: ascii -*-
# BridgeWithTrain
# Suspension bridge + visible train + synchronized moving axle loads.
# Abaqus/Standard implicit dynamics. Units: m, kg, s, N, Pa.
# Creates model/mesh/job and writes INP only.

from abaqus import mdb
from abaqusConstants import *
import part, material, section, assembly, step, load, mesh, job
import math
import os

MODEL = 'BridgeWithTrain'
JOB = 'BridgeWithTrain'

# ------------------------------------------------------------------
# Parameters
# ------------------------------------------------------------------
L = 200.0
DX = 10.0
YS = (-5.0, 5.0)
TX = (40.0, 160.0)
TOP = 24.0
BASE_Z = -15.0
SAG = 12.0
XS = [i * DX for i in range(int(round(L / DX)) + 1)]

SPEED = 20.0
AXLE_FORCE = 180000.0
APPROACH = 10.0
CAR_LENGTH = 24.0
CAR_WIDTH = 3.2
CAR_HEIGHT = 3.5
CAR_PITCH = 26.0
CAR_COUNT = 2
LOCAL_AXLES = (22.0, 19.5, 4.5, 2.0)
RAMP_TIME = 0.20
RINGDOWN = 2.0
AMP_DT = 0.02

OFFSETS = tuple(
    c * CAR_PITCH + LOCAL_AXLES[0] - x
    for c in range(CAR_COUNT)
    for x in LOCAL_AXLES
)

END_TIME = (
    (L + APPROACH + max(OFFSETS)) / SPEED
    + 0.5 * RAMP_TIME + RINGDOWN
)

# ------------------------------------------------------------------
# New model
# ------------------------------------------------------------------
if MODEL in mdb.models.keys():
    del mdb.models[MODEL]
if JOB in mdb.jobs.keys():
    del mdb.jobs[JOB]

m = mdb.Model(name=MODEL, modelType=STANDARD_EXPLICIT)

for repo in (m.fieldOutputRequests, m.historyOutputRequests):
    for key in list(repo.keys()):
        del repo[key]

# ------------------------------------------------------------------
# Suspension bridge geometry
# ------------------------------------------------------------------
def cable_z(x):
    if x <= TX[0]:
        return TOP * x / TX[0]
    if x >= TX[1]:
        return TOP * (L - x) / (L - TX[1])
    r = (x - TX[0]) / (TX[1] - TX[0])
    return TOP - 4.0 * SAG * r * (1.0 - r)


groups = dict(
    (name, []) for name in
    ('LONG', 'CROSS', 'TOWER', 'TOPBEAM', 'CABLE', 'HANGER')
)

sk = m.ConstrainedSketch(name='BridgeSketch', sheetSize=500.0)

for y in YS:
    for xa, xb in zip(XS[:-1], XS[1:]):
        sk.Line(point1=(xa, y), point2=(xb, y))
        groups['LONG'].append(((xa, y, 0.0), (xb, y, 0.0)))

for x in XS:
    sk.Line(point1=(x, YS[0]), point2=(x, YS[1]))
    groups['CROSS'].append(((x, YS[0], 0.0), (x, YS[1], 0.0)))

p = m.Part(name='Bridge', dimensionality=THREE_D, type=DEFORMABLE_BODY)
p.BaseWire(sketch=sk)
del m.sketches['BridgeSketch']

for x in TX:
    for y in YS:
        groups['TOWER'].extend((
            ((x, y, BASE_Z), (x, y, 0.0)),
            ((x, y, 0.0), (x, y, TOP))
        ))
    groups['TOPBEAM'].append(
        ((x, YS[0], TOP), (x, YS[1], TOP))
    )

for y in YS:
    for xa, xb in zip(XS[:-1], XS[1:]):
        groups['CABLE'].append(
            ((xa, y, cable_z(xa)), (xb, y, cable_z(xb)))
        )
    for x in XS[1:-1]:
        if x not in TX:
            groups['HANGER'].append(
                ((x, y, 0.0), (x, y, cable_z(x)))
            )

for name in ('TOWER', 'TOPBEAM', 'CABLE', 'HANGER'):
    p.WirePolyLine(
        points=tuple(groups[name]),
        mergeType=IMPRINT,
        meshable=ON
    )

for name, members in groups.items():
    mids = tuple(
        ((0.5 * (a[0] + b[0]),
          0.5 * (a[1] + b[1]),
          0.5 * (a[2] + b[2])),)
        for a, b in members
    )
    p.Set(name=name, edges=p.edges.findAt(*mids))

p.Set(name='ALL_BRIDGE', edges=p.edges[:])

# ------------------------------------------------------------------
# Materials and sections
# ------------------------------------------------------------------
for name, young, rho in (
    ('DeckSteel', 210.0e9, 12000.0),
    ('FrameSteel', 210.0e9, 7850.0),
    ('CableSteel', 195.0e9, 7850.0),
    ('TrainSteel', 210.0e9, 7850.0)
):
    mat = m.Material(name=name)
    mat.Elastic(table=((young, 0.30),))
    mat.Density(table=((rho,),))
    if name != 'TrainSteel':
        mat.Damping(alpha=0.05, beta=0.0002)

for name, aa, bb, tt in (
    ('DeckBox', 2.0, 1.4, 0.04),
    ('CrossBox', 0.8, 0.6, 0.025),
    ('TowerBox', 2.0, 2.0, 0.06)
):
    m.BoxProfile(
        name=name, a=aa, b=bb,
        uniformThickness=ON, t1=tt
    )

for name, profile, mat in (
    ('DeckSec', 'DeckBox', 'DeckSteel'),
    ('CrossSec', 'CrossBox', 'FrameSteel'),
    ('TowerSec', 'TowerBox', 'FrameSteel')
):
    m.BeamSection(
        name=name,
        integration=DURING_ANALYSIS,
        poissonRatio=0.30,
        profile=profile,
        material=mat,
        temperatureVar=LINEAR,
        consistentMassMatrix=False
    )

m.TrussSection(name='CableSec', material='CableSteel', area=0.05)
m.TrussSection(name='HangerSec', material='CableSteel', area=0.004)

sec_map = {
    'LONG': 'DeckSec',
    'CROSS': 'CrossSec',
    'TOWER': 'TowerSec',
    'TOPBEAM': 'TowerSec',
    'CABLE': 'CableSec',
    'HANGER': 'HangerSec'
}

beam_type = mesh.ElemType(elemCode=B31, elemLibrary=STANDARD)
truss_type = mesh.ElemType(elemCode=T3D2, elemLibrary=STANDARD)

for name, sec_name in sec_map.items():
    region = p.sets[name]
    p.SectionAssignment(region=region, sectionName=sec_name)

    if name in ('CABLE', 'HANGER'):
        p.setElementType(
            regions=(region.edges,),
            elemTypes=(truss_type,)
        )
    else:
        n1 = (0.0, 1.0, 0.0) if name == 'LONG' else (1.0, 0.0, 0.0)
        p.assignBeamSectionOrientation(
            region=region, method=N1_COSINES, n1=n1
        )
        p.setElementType(
            regions=(region.edges,),
            elemTypes=(beam_type,)
        )

    p.seedEdgeByNumber(
        edges=region.edges,
        number=1,
        constraint=FIXED
    )

p.generateMesh()

if len(p.elements) == 0:
    raise RuntimeError('Bridge mesh is empty.')

# ------------------------------------------------------------------
# Visible train car
# ------------------------------------------------------------------
sk = m.ConstrainedSketch(name='CarSketch', sheetSize=60.0)
sk.rectangle(
    point1=(0.0, -0.5 * CAR_WIDTH),
    point2=(CAR_LENGTH, 0.5 * CAR_WIDTH)
)

car = m.Part(
    name='TrainCar',
    dimensionality=THREE_D,
    type=DEFORMABLE_BODY
)
car.BaseSolidExtrude(sketch=sk, depth=CAR_HEIGHT)
del m.sketches['CarSketch']

m.HomogeneousSolidSection(
    name='CarSection',
    material='TrainSteel',
    thickness=None
)
car.Set(name='BODY', cells=car.cells[:])
car.SectionAssignment(
    region=car.sets['BODY'],
    sectionName='CarSection'
)

rp = car.ReferencePoint(
    point=(0.5 * CAR_LENGTH, 0.0, 0.5 * CAR_HEIGHT)
)
car.Set(
    name='CONTROL',
    referencePoints=(car.referencePoints[rp.id],)
)

car.seedPart(size=5.0, deviationFactor=0.1, minSizeFactor=0.1)
car.setMeshControls(
    regions=car.cells[:],
    elemShape=TET,
    technique=FREE
)
car.setElementType(
    regions=(car.cells[:],),
    elemTypes=(
        mesh.ElemType(elemCode=C3D10, elemLibrary=STANDARD),
    )
)
car.generateMesh()

if len(car.elements) == 0:
    raise RuntimeError('Train car mesh is empty.')

# ------------------------------------------------------------------
# Assembly
# ------------------------------------------------------------------
a = m.rootAssembly
a.DatumCsysByDefault(CARTESIAN)

bridge = a.Instance(name='BRIDGE', part=p, dependent=ON)
cars = []

for c in range(CAR_COUNT):
    iname = 'TRAIN_%02d' % (c + 1)
    ci = a.Instance(name=iname, part=car, dependent=ON)

    x0 = -APPROACH - LOCAL_AXLES[0] - c * CAR_PITCH
    a.translate(
        instanceList=(iname,),
        vector=(x0, 0.0, 0.70)
    )

    m.RigidBody(
        name='RigidTrain_%02d' % (c + 1),
        refPointRegion=ci.sets['CONTROL'],
        bodyRegion=ci.sets['BODY']
    )

    m.DisplacementBC(
        name='TrainMotion_%02d' % (c + 1),
        createStepName='Initial',
        region=ci.sets['CONTROL'],
        u1=0.0, u2=0.0, u3=0.0,
        ur1=0.0, ur2=0.0, ur3=0.0
    )
    cars.append(ci)

a.regenerate()

# ------------------------------------------------------------------
# Bridge node lookup / supports
# ------------------------------------------------------------------
lookup = {}
for nd in bridge.nodes:
    lookup[tuple(round(float(v), 7) for v in nd.coordinates)] = nd.label


def make_node_set(name, points):
    labels = []
    for xyz in points:
        key = tuple(round(float(v), 7) for v in xyz)
        if key not in lookup:
            raise RuntimeError('Missing bridge node: %s' % (xyz,))
        labels.append(lookup[key])
    return a.Set(
        name=name,
        nodes=bridge.nodes.sequenceFromLabels(labels=tuple(labels))
    )


foundations = make_node_set(
    'FOUNDATIONS',
    [(x, y, BASE_Z) for x in TX for y in YS]
)
anchors = make_node_set(
    'ANCHORS',
    [(x, y, 0.0) for x in (0.0, L) for y in YS]
)
midspan = make_node_set(
    'MIDSPAN',
    [(0.5 * L, y, 0.0) for y in YS]
)

m.EncastreBC(
    name='FixedBases',
    createStepName='Initial',
    region=foundations
)
m.DisplacementBC(
    name='EndSupports',
    createStepName='Initial',
    region=anchors,
    u1=0.0, u2=0.0, u3=0.0
)

# ------------------------------------------------------------------
# Standard analysis steps
# ------------------------------------------------------------------
m.StaticStep(
    name='DeadLoad',
    previous='Initial',
    nlgeom=ON,
    timePeriod=1.0,
    initialInc=0.02,
    minInc=1.0e-7,
    maxInc=0.10,
    maxNumInc=1000
)

m.Gravity(
    name='BridgeGravity',
    createStepName='DeadLoad',
    region=bridge.sets['ALL_BRIDGE'],
    comp3=-9.81
)

m.ImplicitDynamicsStep(
    name='TrainPass',
    previous='DeadLoad',
    nlgeom=ON,
    timePeriod=END_TIME,
    application=TRANSIENT_FIDELITY,
    initialInc=0.005,
    minInc=1.0e-7,
    maxInc=0.02,
    maxNumInc=30000
)

# ------------------------------------------------------------------
# Strictly increasing amplitude time grid
# No mixed event-time list is used.
# ------------------------------------------------------------------
def travel(t):
    if t <= RAMP_TIME:
        return 0.5 * SPEED * t * t / RAMP_TIME
    return SPEED * (t - 0.5 * RAMP_TIME)


n_regular = int(math.floor(END_TIME / AMP_DT))
times = [round(i * AMP_DT, 8) for i in range(n_regular + 1)]

end_value = round(END_TIME, 8)
if end_value > times[-1] + 1.0e-8:
    times.append(end_value)
else:
    times[-1] = end_value

for i in range(1, len(times)):
    if not times[i] > times[i - 1]:
        raise RuntimeError(
            'Amplitude time is not strictly increasing at index %d.' % i
        )

m.TabularAmplitude(
    name='TrainTravel',
    timeSpan=STEP,
    smooth=SOLVER_DEFAULT,
    data=tuple((t, travel(t)) for t in times)
)

for c in range(CAR_COUNT):
    m.boundaryConditions[
        'TrainMotion_%02d' % (c + 1)
    ].setValuesInStep(
        stepName='TrainPass',
        u1=1.0,
        amplitude='TrainTravel'
    )

# ------------------------------------------------------------------
# Moving axle forces
# Each axle is linearly distributed between adjacent deck stations.
# ------------------------------------------------------------------
def axle_weight(xnode, t):
    total = 0.0
    d = travel(t)

    for offset in OFFSETS:
        xpos = d - APPROACH - offset
        if -DX < xpos < L + DX:
            w = 1.0 - abs(xpos - xnode) / DX
            if w > 0.0:
                total += w

    return total


for i, x in enumerate(XS):
    amp_name = 'AxleAmp_%02d' % i
    amp_data = tuple(
        (times[j], axle_weight(x, times[j]))
        for j in range(len(times))
    )

    m.TabularAmplitude(
        name=amp_name,
        timeSpan=STEP,
        smooth=SOLVER_DEFAULT,
        data=amp_data
    )

    for side, y in enumerate(YS):
        region = make_node_set(
            'LOAD_%02d_%d' % (i, side),
            [(x, y, 0.0)]
        )
        m.ConcentratedForce(
            name='TrainForce_%02d_%d' % (i, side),
            createStepName='TrainPass',
            region=region,
            cf3=-0.5 * AXLE_FORCE,
            amplitude=amp_name
        )

# ------------------------------------------------------------------
# Output
# ------------------------------------------------------------------
dead_output = m.FieldOutputRequest(
    name='DeadLoadOutput',
    createStepName='DeadLoad',
    variables=('S', 'U', 'RF'),
    frequency=LAST_INCREMENT
)
dead_output.deactivate('TrainPass')

m.FieldOutputRequest(
    name='TrainPassOutput',
    createStepName='TrainPass',
    variables=('S', 'U', 'V', 'A', 'RF'),
    timeInterval=0.10
)

m.HistoryOutputRequest(
    name='MidspanResponse',
    createStepName='TrainPass',
    region=midspan,
    variables=('U3', 'V3', 'A3'),
    timeInterval=0.02
)

for c, ci in enumerate(cars):
    m.HistoryOutputRequest(
        name='TrainHistory_%02d' % (c + 1),
        createStepName='TrainPass',
        region=ci.sets['CONTROL'],
        variables=('U1', 'V1'),
        timeInterval=0.05
    )

# ------------------------------------------------------------------
# Initial cable/hanger tension
# ------------------------------------------------------------------
m.keywordBlock.synchVersions(storeNodesAndElements=False)

first_step = None
for i, block in enumerate(m.keywordBlock.sieBlocks):
    if block.lstrip().upper().startswith('*STEP'):
        first_step = i
        break

if first_step is None or first_step < 1:
    raise RuntimeError('Cannot locate first analysis step.')

m.keywordBlock.insert(
    position=first_step - 1,
    text=(
        '*INITIAL CONDITIONS, TYPE=STRESS\n'
        'BRIDGE.CABLE, 2.0E8\n'
        'BRIDGE.HANGER, 3.0E7'
    )
)

# ------------------------------------------------------------------
# Create the CAE Job and write input only
# ------------------------------------------------------------------
analysis_job = mdb.Job(
    name=JOB,
    model=MODEL,
    type=ANALYSIS,
    description='Visible train passing suspension bridge with moving axles',
    numCpus=4,
    numDomains=4,
    memory=80,
    memoryUnits=PERCENTAGE
)

analysis_job.writeInput(consistencyChecking=ON)

print('Input written: ' + os.path.abspath(JOB + '.inp'))
print('Amplitude points: %d' % len(times))
print('First amplitude times: %s' % str(times[:5]))
print('Last amplitude times: %s' % str(times[-5:]))
print('Visible train cars: %d' % CAR_COUNT)
print('Moving axles: %d' % len(OFFSETS))
print('BridgeWithTrain is ready in Abaqus/CAE Job Manager.')
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')