# -*- coding: utf-8 -*-
from abaqus import mdb
from abaqusConstants import *
import mesh
import math
import time

try:
    import builtins as py
except ImportError:
    import __builtin__ as py

# SI units: m, kg, s, Pa.
# Reference: Hsu et al., Ocean Engineering 72 (2013), 98-106.
# Paper: PMMA 50x50x5.9 mm; Eulerian domain 50x50x20 mm;
# spherical-head water jet: diameter 3 mm, length 12 mm, speed 570 m/s.
# Supplementary assumptions not explicitly specified in the PDF:
# tip clearance 0.5 mm; EOS s=1.9, Gamma0=0; viscosity 0.0031 Pa.s;
# elastic-perfectly plastic PMMA; total duration 32 microseconds.
# Coarser mesh targets overall flow morphology, not converged peak pressure.
# Explicit loops avoid generator arguments and Abaqus built-in name clashes.

W = 0.050
T = 0.0059
H = 0.020
R = 0.0015
L = 0.012
GAP = 0.0005
SPEED = 570.0
DURATION = 32.0e-6
NS = 4

prefix = 'JetCEL_Fix_' + time.strftime('%Y%m%d_%H%M%S')
name = prefix
suffix = 1
while name in mdb.models.keys() or name in mdb.jobs.keys():
    name = prefix + '_%d' % suffix
    suffix += 1

model = mdb.Model(name=name, modelType=STANDARD_EXPLICIT)

water = model.Material(name='Water')
water.Density(table=((1000.0,),))
water.Eos(type=USUP, table=((1500.0, 1.9, 0.0),))
water.Viscosity(table=((0.0031,),))

solid = model.Material(name='PMMA')
solid.Density(table=((1170.0,),))
solid.Elastic(table=((3.3e9, 0.36),))
solid.Plastic(table=((220.0e6, 0.0),))

model.EulerianSection(
    name='EulerianSection', data={'Water': 'Water'})
model.HomogeneousSolidSection(
    name='PMMASection', material='PMMA', thickness=None)

def cut(part, plane, offset):
    feature = part.DatumPlaneByPrincipalPlane(
        principalPlane=plane, offset=offset)
    part.PartitionCellByDatumPlane(
        cells=part.cells[:], datumPlane=part.datums[feature.id])

def block(part_name, depth, kind):
    sketch_name = part_name + '_Profile'
    sketch = model.ConstrainedSketch(
        name=sketch_name, sheetSize=0.12)
    sketch.rectangle(
        point1=(-W/2.0, -W/2.0), point2=(W/2.0, W/2.0))
    part = model.Part(
        name=part_name, dimensionality=THREE_D, type=kind)
    part.BaseSolidExtrude(sketch=sketch, depth=depth)
    del model.sketches[sketch_name]
    for plane in (YZPLANE, XZPLANE):
        for offset in (-0.008, -0.003, 0.003, 0.008):
            cut(part, plane, offset)
    return part

plate = block('PMMATarget', T, DEFORMABLE_BODY)
fluid = block('EulerianGrid', H, EULERIAN)
cut(plate, XYPLANE, T-0.001)
for z in (T, T+0.002):
    cut(fluid, XYPLANE, z)

for part, section in (
        (plate, 'PMMASection'), (fluid, 'EulerianSection')):
    part.Set(name='Cells', cells=part.cells[:])
    part.SectionAssignment(
        region=part.sets['Cells'], sectionName=section)

def mesh_block(part, is_fluid):
    part.setMeshControls(
        regions=part.cells[:], elemShape=HEX, technique=STRUCTURED)
    code = EC3D8R if is_fluid else C3D8R
    part.setElementType(
        regions=(part.cells[:],),
        elemTypes=(mesh.ElemType(
            elemCode=code, elemLibrary=EXPLICIT),))
    for edge in part.edges:
        vertex_ids = edge.getVertices()
        p0 = part.vertices[vertex_ids[0]].pointOn[0]
        p1 = part.vertices[vertex_ids[1]].pointOn[0]
        spans = []
        for axis in py.range(3):
            spans.append(py.abs(p1[axis]-p0[axis]))
        length = py.max(spans)
        axis = spans.index(length)
        middle = 0.5*(p0[axis]+p1[axis])
        if axis < 2:
            distance = py.abs(middle)
            if distance < 0.003:
                size = 0.00025
            elif distance < 0.008:
                size = 0.00065
            else:
                size = 0.002
        elif is_fluid:
            if middle < T:
                size = 0.001
            elif middle < T+0.002:
                size = 0.00025
            else:
                size = 0.00040
        else:
            size = 0.00025 if middle > T-0.001 else 0.0008
        count = py.max(1, py.int(math.ceil(length/size-1.0e-9)))
        part.seedEdgeByNumber(
            edges=(edge,), number=count, constraint=FIXED)
    part.generateMesh()
    if py.len(part.elements) == 0:
        raise RuntimeError('Mesh generation failed: ' + part.name)
    part.Set(name='Elements', elements=part.elements[:])
    part.Set(name='Nodes', nodes=part.nodes[:])

mesh_block(plate, False)
mesh_block(fluid, True)

assembly = model.rootAssembly
assembly.DatumCsysByDefault(CARTESIAN)
assembly.Instance(name='Target-1', part=plate, dependent=ON)
assembly.Instance(name='Fluid-1', part=fluid, dependent=ON)
assembly.translate(
    instanceList=('Target-1', 'Fluid-1'), vector=(0.0, 0.0, -T))
assembly.regenerate()
target = assembly.instances['Target-1']
domain = assembly.instances['Fluid-1']

# Analytic jet shape: lower hemisphere plus cylindrical body.
# Total tip-to-tail length is L.
tip = GAP
cap_center = GAP+R
top = GAP+L

def occupied(x, y, z):
    rr = x*x+y*y
    if z < tip or z > top or rr > R*R:
        return False
    if z < cap_center:
        return rr+(z-cap_center)*(z-cap_center) <= R*R
    return True

pad = 0.0006
candidates = domain.elements.getByBoundingBox(
    xMin=-R-pad, xMax=R+pad,
    yMin=-R-pad, yMax=R+pad,
    zMin=tip-pad, zMax=top+pad)
buckets = {}
volume = 0.0
sample_count = py.float(NS*NS*NS)

for element in candidates:
    nodes = element.getNodes()
    lo = [1.0e30, 1.0e30, 1.0e30]
    hi = [-1.0e30, -1.0e30, -1.0e30]
    for node in nodes:
        xyz = node.coordinates
        for axis in py.range(3):
            if xyz[axis] < lo[axis]:
                lo[axis] = xyz[axis]
            if xyz[axis] > hi[axis]:
                hi[axis] = xyz[axis]
    increments = []
    for axis in py.range(3):
        increments.append((hi[axis]-lo[axis])/NS)
    hits = 0
    for i in py.range(NS):
        x = lo[0]+(i+0.5)*increments[0]
        for j in py.range(NS):
            y = lo[1]+(j+0.5)*increments[1]
            for k in py.range(NS):
                z = lo[2]+(k+0.5)*increments[2]
                if occupied(x, y, z):
                    hits += 1
    if hits > 0:
        if hits not in buckets:
            buckets[hits] = []
        buckets[hits].append(element.label)
        cell_volume = (hi[0]-lo[0])*(hi[1]-lo[1])*(hi[2]-lo[2])
        volume += cell_volume*hits/sample_count

if py.len(buckets) == 0:
    raise RuntimeError('No water-filled elements were found.')

assignment_rows = []
for hits in py.sorted(buckets.keys()):
    elements = domain.elements.sequenceFromLabels(
        labels=py.tuple(buckets[hits]))
    region = assembly.Set(name='VF_%02d' % hits, elements=elements)
    assignment_rows.append((region, (hits/sample_count,)))

# Unassigned Eulerian volume is void.
model.MaterialAssignment(
    name='WaterVolumeFractions', instanceList=(domain,),
    useFields=OFF, assignmentList=py.tuple(assignment_rows))

theoretical_volume = math.pi*R*R*(L-R/3.0)
error = py.abs(volume/theoretical_volume-1.0)
if error > 0.05:
    raise RuntimeError('Water initialization volume error exceeds 5 percent.')

model.Velocity(
    name='JetInitialVelocity', region=domain.sets['Nodes'],
    velocity1=0.0, velocity2=0.0, velocity3=-SPEED, omega=0.0)

bottom = target.nodes.getByBoundingBox(
    xMin=-0.026, xMax=0.026, yMin=-0.026, yMax=0.026,
    zMin=-T-1.0e-8, zMax=-T+1.0e-8)
if py.len(bottom) == 0:
    raise RuntimeError('Target bottom node set is empty.')
assembly.Set(name='TargetBottom', nodes=bottom)
model.EncastreBC(
    name='FixedBottom', createStepName='Initial',
    region=assembly.sets['TargetBottom'])

model.ExplicitDynamicsStep(
    name='Impact', previous='Initial',
    timePeriod=DURATION, nlgeom=ON)
model.ContactProperty('CELProperty')
model.interactionProperties['CELProperty'].NormalBehavior(
    pressureOverclosure=HARD, allowSeparation=ON)
contact = model.ContactExp(name='CELContact', createStepName='Initial')
contact.includedPairs.setValuesInStep(
    stepName='Initial', useAllstar=ON)
contact.contactPropertyAssignments.appendInStep(
    stepName='Initial', assignments=((GLOBAL, SELF, 'CELProperty'),))

for key in py.list(model.fieldOutputRequests.keys()):
    del model.fieldOutputRequests[key]
for key in py.list(model.historyOutputRequests.keys()):
    del model.historyOutputRequests[key]
model.FieldOutputRequest(
    name='WaterResults', createStepName='Impact',
    region=domain.sets['Elements'],
    variables=('EVF', 'V', 'S'), numIntervals=80)
model.FieldOutputRequest(
    name='TargetResults', createStepName='Impact',
    region=target.sets['Elements'],
    variables=('U', 'V', 'S', 'PEEQ'), numIntervals=80)
model.HistoryOutputRequest(
    name='EnergyBalance', createStepName='Impact',
    variables=('ALLKE', 'ALLIE', 'ALLAE', 'ALLVD', 'ETOTAL'),
    numIntervals=200)

assembly.regenerate()
job = mdb.Job(
    name=name, model=name, type=ANALYSIS,
    description='Reduced mesh CEL water jet impact on PMMA',
    numCpus=4, numDomains=4,
    parallelizationMethodExplicit=DOMAIN,
    explicitPrecision=DOUBLE, nodalOutputPrecision=SINGLE,
    memory=80, memoryUnits=PERCENTAGE)
job.writeInput(consistencyChecking=ON)
mdb.saveAs(pathName=name+'.cae')
print('Created job: '+name)
print('Elements: target=%d, Eulerian=%d' %
      (py.len(plate.elements), py.len(fluid.elements)))
print('Water volume error: %.3f percent' % (100.0*error))
print('INP written. Submit this existing job through CAE Job Manager.')
# After solving: display V magnitude on the EVF=0.5 material surface.
# Keep the PMMA target visible and hide the Eulerian void mesh.