# Copyright (c) 2026 JSOL CORPORATION
#
# This script is released under the MIT License.
# See the full license text at:
# https://www.jmag-international.com/scriptlibrary/jmag_script_library_mit/

app = designer.GetApplication()
geomApp = app.CreateGeometryEditor()
geomDoc = geomApp.GetDocument()

# Dimensional parameters
Slots = 24
StatorOuterDiameter = 204
StatorInnerDiameter = 102
ToothTipThickness = 1.6
SlotWidth1 = 6.0
SlotWidth2 = 3.6
SlotDepthFromOuter = 25.5

# Equation names
EqSlots = u"SLOTS"
EqStatorOuterDiameter = u"SD1"
EqStatorInnerDiameter = u"SD2"
EqToothTipThickness = u"ST"
EqSlotWidth1 = u"SW1"
EqSlotWidth2 = u"SW2"
EqSlotDepthFromOuter = u"SW4"


def setupEquation(equation, expression, equationType):
    """Configure the type and expression of an equation in DesignTable."""
    # Argument of SetType
    # 0: Value
    # 1: Expression
    equation.SetType(equationType)
    equation.SetExpression(str(expression))


def createStatorEquations():
    """Register equations used for the stator geometry in DesignTable."""
    designTable = geomDoc.GetDesignTable()
    equationSettings = [
        (EqSlots, Slots, 0),
        (EqStatorOuterDiameter, StatorOuterDiameter, 0),
        (EqStatorInnerDiameter, StatorInnerDiameter, 0),
        (EqToothTipThickness, ToothTipThickness, 0),
        (EqSlotWidth1, SlotWidth1, 0),
        (EqSlotWidth2, SlotWidth2, 0),
        (EqSlotDepthFromOuter, SlotDepthFromOuter, 0),
    ]
    designTable.EditStart()
    for equationName, expression, equationType in equationSettings:
        designTable.AddEquation(equationName)
        equation = designTable.GetEquation(equationName)
        setupEquation(equation, expression, equationType)
    designTable.EditEnd()


def createReference(item):
    """Create a ReferenceObject from a geometry item."""
    return geomDoc.CreateReferenceFromItem(item)


def addItemsToReferenceList(feature, propertyName, itemList):
    """Add geometry items to a stringlist property."""
    for item in itemList:
        feature.AddPropertyByReference(propertyName, createReference(item))


def createRegionFromItems(sketch, itemList):
    """Create a region on the sketch from all items in itemList."""
    selection = geomDoc.GetSelection()
    selection.Clear()
    for item in itemList:
        selection.Add(item)
    sketch.CreateRegions()
    selection.Clear()


def getRegionItemByOrder(sketch, createdOrder):
    """Get the region object at the specified creation order."""
    currentCount = 0
    for i in range(sketch.NumItems()):
        item = sketch.GetItem(i)
        if item.GetScriptTypeName() == u"RegionItem":
            if currentCount == createdOrder - 1:
                return item
            currentCount += 1
    return sketch


def calculateDirectionByAngle(sketch, angle):
    """Get a unit vector in the specified direction from the X-axis as a Point object."""
    dummyLine = sketch.CreateLineByLengthAngle(0, 0, 1, angle)
    endVertex = dummyLine.GetEndVertex()
    direction = geomApp.CreatePoint(endVertex.GetX(), endVertex.GetY(), 0)
    selection = geomDoc.GetSelection()
    selection.Clear()
    selection.Add(dummyLine)
    selection.Delete()
    selection.Clear()
    return direction


def calculatePointOnCircle(radius, y):
    """Get a point on a circle centered at the origin for the specified Y coordinate as a Point object."""
    x = (radius * radius - y * y) ** 0.5
    return geomApp.CreatePoint(x, y, 0)


def calculatePointOnRadialLine(radius, direction):
    """Get the intersection of a circle centered at the origin and a line through the origin in the specified direction."""
    x = radius * direction.GetX()
    y = radius * direction.GetY()
    return geomApp.CreatePoint(x, y, 0)


def calculatePointOnOffsetLineAndCircle(radius, offset, direction):
    """Get the intersection of a circle and a line offset clockwise from the specified direction."""
    lengthOnDirection = (radius * radius - offset * offset) ** 0.5
    x = direction.GetX() * lengthOnDirection + direction.GetY() * offset
    y = direction.GetY() * lengthOnDirection - direction.GetX() * offset
    return geomApp.CreatePoint(x, y, 0)


def createArc(sketch, startPoint, endPoint):
    """Create an arc centered at the origin between the specified start and end points."""
    return sketch.CreateArc(0, 0, startPoint.GetX(), startPoint.GetY(), endPoint.GetX(), endPoint.GetY())


def createLine(sketch, startPoint, endPoint):
    """Create a line between the specified start and end points."""
    return sketch.CreateLine(startPoint.GetX(), startPoint.GetY(), endPoint.GetX(), endPoint.GetY())


def createDiameterConstraint(sketch, arc, expression):
    """Create a diameter constraint for an arc and set its equation."""
    constraint = sketch.CreateMonoConstraint(u"radius", createReference(arc))
    # Argument of SetType
    # 0: Radius
    # 1: Diameter
    constraint.SetType(1)
    constraint.SetProperty(u"Diameter", expression)


def createDistanceConstraint(sketch, firstItem, secondItem, expression, useDoubleWidth=False):
    """Create a distance constraint between two geometry items and set its equation."""
    constraint = sketch.CreateBiConstraint(u"distance", createReference(firstItem), createReference(secondItem))
    constraint.SetProperty(u"Distance", expression)
    if useDoubleWidth:
        constraint.SetProperty(u"UseDoubleWidth", True)


def createAngleConstraint(sketch, firstItem, secondItem, expression):
    """Create an angle constraint between two geometry items and set its equation."""
    constraint = sketch.CreateBiConstraint(u"angle", createReference(firstItem), createReference(secondItem))
    constraint.SetProperty(u"Angle", expression)


def setRegionAppearance(region, name, color):
    """Set the name and display color of a region."""
    region.SetName(name)
    region.SetProperty(u"Color", color)
    region.SetProperty(u"UseParentColor", 0)


def createStatorBaseShape(sketch):
    """Create the stator base geometry and constraints for half a slot pitch."""
    statorOuterRadius = StatorOuterDiameter / 2.0
    statorInnerRadius = StatorInnerDiameter / 2.0
    toothTipOuterRadius = statorInnerRadius + ToothTipThickness
    slotBottomRadius = statorOuterRadius - SlotDepthFromOuter
    halfSlotAngle = 360.0 / Slots / 2.0
    halfSlotWidth1 = SlotWidth1 / 2.0
    halfSlotWidth2 = SlotWidth2 / 2.0
    slotCenterDirection = calculateDirectionByAngle(sketch, halfSlotAngle)

    # Create and fix the origin.
    originVertex = sketch.CreateVertex(0, 0)
    sketch.CreateMonoConstraint(u"fixture", createReference(originVertex))

    # Calculate the vertex positions of each boundary.
    pointInnerOnXAxis = geomApp.CreatePoint(statorInnerRadius, 0, 0)
    pointOuterOnXAxis = geomApp.CreatePoint(statorOuterRadius, 0, 0)
    pointInnerAtSlotOpening = calculatePointOnOffsetLineAndCircle(statorInnerRadius, halfSlotWidth2, slotCenterDirection)
    pointToothTipAtSlotOpening = calculatePointOnOffsetLineAndCircle(toothTipOuterRadius, halfSlotWidth2, slotCenterDirection)
    pointToothTipAtToothSide = calculatePointOnCircle(toothTipOuterRadius, halfSlotWidth1)
    pointSlotBottomAtToothSide = calculatePointOnCircle(slotBottomRadius, halfSlotWidth1)
    pointInnerAtSlotCenter = calculatePointOnRadialLine(statorInnerRadius, slotCenterDirection)
    pointToothTipAtSlotCenter = calculatePointOnRadialLine(toothTipOuterRadius, slotCenterDirection)
    pointSlotBottomAtSlotCenter = calculatePointOnRadialLine(slotBottomRadius, slotCenterDirection)
    pointOuterAtSlotCenter = calculatePointOnRadialLine(statorOuterRadius, slotCenterDirection)

    # Create arcs for the stator inner and outer diameters.
    innerArcForCore = createArc(sketch, pointInnerOnXAxis, pointInnerAtSlotOpening)
    createDiameterConstraint(sketch, innerArcForCore, EqStatorInnerDiameter)
    outerArc = createArc(sketch, pointOuterOnXAxis, pointOuterAtSlotCenter)
    createDiameterConstraint(sketch, outerArc, EqStatorOuterDiameter)

    # Create arcs for the tooth tip and slot bottom.
    toothTipArcForCore = createArc(sketch, pointToothTipAtToothSide, pointToothTipAtSlotOpening)
    toothTipArcForSlot = createArc(sketch, pointToothTipAtSlotOpening, pointToothTipAtSlotCenter)
    slotBottomArc = createArc(sketch, pointSlotBottomAtToothSide, pointSlotBottomAtSlotCenter)

    # Create the slot opening and tooth side.
    slotOpeningLine = createLine(sketch, pointInnerAtSlotOpening, pointToothTipAtSlotOpening)
    toothSideLine = createLine(sketch, pointToothTipAtToothSide, pointSlotBottomAtToothSide)
    sketch.CreateMonoConstraint(u"horizontality", createReference(toothSideLine))

    # Create the boundaries on the slot-center side.
    slotCenterOpeningLine = createLine(sketch, pointInnerAtSlotCenter, pointToothTipAtSlotCenter)
    sketch.CreateBiConstraint(u"coincident", createReference(slotCenterOpeningLine), createReference(originVertex))
    slotCenterSlotLine = createLine(sketch, pointToothTipAtSlotCenter, pointSlotBottomAtSlotCenter)
    sketch.CreateBiConstraint(u"coincident", createReference(slotCenterOpeningLine), createReference(slotCenterSlotLine))
    slotCenterOuterLine = createLine(sketch, pointSlotBottomAtSlotCenter, pointOuterAtSlotCenter)
    sketch.CreateBiConstraint(u"coincident", createReference(slotCenterSlotLine), createReference(slotCenterOuterLine))

    # Create the stator-core boundary on the X-axis.
    radialLineOnXAxis = createLine(sketch, pointInnerOnXAxis, pointOuterOnXAxis)
    sketch.CreateMonoConstraint(u"horizontality", createReference(radialLineOnXAxis))
    sketch.CreateBiConstraint(u"coincident", createReference(radialLineOnXAxis), createReference(originVertex))

    # Constrain the arcs to be concentric and set the stator dimensions.
    for arc in [toothTipArcForCore, toothTipArcForSlot, slotBottomArc, outerArc]:
        sketch.CreateBiConstraint(u"concentricity", createReference(innerArcForCore), createReference(arc))
    createDistanceConstraint(sketch, innerArcForCore, toothTipArcForCore, EqToothTipThickness)
    createDistanceConstraint(sketch, slotBottomArc, outerArc, EqSlotDepthFromOuter)
    createDistanceConstraint(sketch, toothSideLine, radialLineOnXAxis, EqSlotWidth1, useDoubleWidth=True)
    # Define the slot-opening width by the distance to the symmetry axis over the same radial range as the opening.
    createDistanceConstraint(sketch, slotOpeningLine, slotCenterOpeningLine, EqSlotWidth2, useDoubleWidth=True)
    createAngleConstraint(sketch, radialLineOnXAxis, slotCenterOuterLine, u"360/{0}/2".format(EqSlots))

    # Use the tooth-tip outer arc as the inner-side boundary of the slot region.
    # Include the inner-diameter arc and slot-opening line only in the stator-core region.
    slotEdgeList = [
        toothTipArcForCore,
        toothTipArcForSlot,
        toothSideLine,
        slotBottomArc,
        slotCenterSlotLine,
    ]
    statorCoreEdgeList = [
        innerArcForCore,
        slotOpeningLine,
        toothTipArcForCore,
        toothSideLine,
        slotBottomArc,
        slotCenterOuterLine,
        outerArc,
        radialLineOnXAxis,
    ]
    return slotEdgeList, statorCoreEdgeList, slotCenterSlotLine, slotCenterOuterLine


def createSlotPattern(sketch, slotRegion, slotSymmetryLine):
    """Mirror-copy the slot region and circular-pattern it over the 90 deg partial model."""
    # Mirror-copy the half-slot-pitch region to create one full slot.
    slotMirror = sketch.CreateRegionMirrorCopy()
    slotMirror.SetName(u"SlotMirror1Slot")
    addItemsToReferenceList(slotMirror, u"Region", [slotRegion])
    slotMirror.SetSymmetry(createReference(slotSymmetryLine))

    # Circular-pattern the one-slot geometry for the number of slots in the 90 deg partial model.
    slotPattern = sketch.CreateRegionCircularPattern()
    slotPattern.SetName(u"SlotCircularPattern")
    addItemsToReferenceList(slotPattern, u"Region", [slotRegion])
    addItemsToReferenceList(slotPattern, u"Feature", [slotMirror])
    slotPattern.SetProperty(u"CenterType", u"DefaultOrigin")
    slotPattern.SetProperty(u"Angle", u"360/{0}".format(EqSlots))
    slotPattern.SetProperty(u"Instance", u"{0}/4".format(EqSlots))


def createStatorCorePattern(sketch, statorCoreRegion, coreSymmetryLine):
    """Mirror-copy the stator-core region and circular-pattern it over the 90 deg partial model."""
    # Mirror-copy the half-slot-pitch region to create one full slot.
    statorCoreMirror = sketch.CreateRegionMirrorCopy()
    statorCoreMirror.SetName(u"StatorCoreMirror1Slot")
    statorCoreMirror.SetMerge(True)
    addItemsToReferenceList(statorCoreMirror, u"Region", [statorCoreRegion])
    statorCoreMirror.SetProperty(u"OverlappedAreaError", True)
    statorCoreMirror.SetSymmetry(createReference(coreSymmetryLine))

    # Circular-pattern the one-slot geometry for the number of slots in the 90 deg partial model.
    statorCorePattern = sketch.CreateRegionCircularPattern()
    statorCorePattern.SetName(u"StatorCoreCircularPattern")
    statorCorePattern.SetMerge(True)
    addItemsToReferenceList(statorCorePattern, u"Region", [statorCoreRegion])
    addItemsToReferenceList(statorCorePattern, u"Feature", [statorCoreMirror])
    statorCorePattern.SetProperty(u"CenterType", u"DefaultOrigin")
    statorCorePattern.SetProperty(u"Angle", u"360/{0}".format(EqSlots))
    statorCorePattern.SetProperty(u"Instance", u"{0}/4".format(EqSlots))


def drawStator(sketch):
    """Create a 2D stator geometry with tapered slots on the sketch."""
    slotEdgeList, statorCoreEdgeList, slotSymmetryLine, coreSymmetryLine = createStatorBaseShape(sketch)

    regionOrder = 0

    # Create the slot region and set its name and color.
    createRegionFromItems(sketch, slotEdgeList)
    regionOrder += 1
    slotRegion = getRegionItemByOrder(sketch, regionOrder)
    if slotRegion.GetScriptTypeName() == u"RegionItem":
        setRegionAppearance(slotRegion, u"Slot", u"#a0a0a4")

    # Create the stator-core region and set its name and color.
    createRegionFromItems(sketch, statorCoreEdgeList)
    regionOrder += 1
    statorCoreRegion = getRegionItemByOrder(sketch, regionOrder)
    if statorCoreRegion.GetScriptTypeName() == u"RegionItem":
        setRegionAppearance(statorCoreRegion, u"StatorCore", u"lime")

    createSlotPattern(sketch, slotRegion, slotSymmetryLine)
    createStatorCorePattern(sketch, statorCoreRegion, coreSymmetryLine)


def drawInSketch(sketch):
    """Open the sketch in editable status and create the geometry."""
    sketch.OpenSketch()
    drawStator(sketch)
    sketch.CloseSketch()


def main():
    """Create a 2D stator geometry with tapered slots."""
    createStatorEquations()
    assm = geomDoc.GetAssembly()
    refXYPlane = assm.GetPlaneXY()
    statorSketch = assm.CreateSketch(refXYPlane)
    drawInSketch(statorSketch)


main()