from abaqus import *
from abaqusConstants import *
from LoadCase.InteractionHelpers import (
    create_RP, sum_regions, coupling_constraint, attach_spring, pick_region,
    couple_nearest_aMTs, find_nearest)
from SpindleAssembly import AddComponents
from SpindleAssembly.AddComponents import return_assembly
from SpindleAssembly.assembly_random import generate_assembly


def create_interactions(**kwargs):
    """
    Define all interactions and connections within model

    :param kwargs: object

    :return: kwargs

    :rtype: object

    :return: data

    :rtype: list
    """
    MTdata, aMTnames, kwargs, data = generate_assembly(**kwargs)
    kwargs.update(dict(ipMTnames=MTdata['MTnames'], aMTnames=aMTnames))

    CoupleIpMTsToSring(data, **kwargs)
    CoupleIpMTsToCentrosomes(**kwargs)
    CoupleAMTs(**kwargs)
    CoupleAMTsToCentrosomes(**kwargs)
    return kwargs, data


def CoupleIpMTsToSring(data, **kwargs):
    """
    Create distributed elastic spring and couple interpolar microtubules to it

    :param kwargs: object

    :return: None
    """
    spindleLength = kwargs['spindleLength']
    kwargs.update({'xpos': 0, 'ypos': 0.5, 'zpos': spindleLength / 2})
    a = AddComponents.return_assembly(**kwargs)

    # Create a reference point to represent ipMT sub-assembly
    RP = create_RP(**kwargs)
    # Catch the id of the defined reference point
    r1 = a.referencePoints
    # Create an associated region by id
    region1 = a.Set(referencePoints=(r1[RP.id], ), name='ipMTsRP')
    # Define a collection of ipMT edges to be assigned coupling connection
    combinedEdges = sum_regions(
        verts_index='None',
        regionType='edge',
        collectionName='ipMTnames',
        separate='False',
        **kwargs)
    region2 = a.Set(edges=combinedEdges, name='ipMT-set')
    # Create a coupling constraint between region 1 and region 2
    influenceRadius = WHOLE_SURFACE
    couplingType = DISTRIBUTING
    weightingMethod = UNIFORM
    coupling_constraint(region1, region2, influenceRadius, couplingType,
                        weightingMethod, 'ipMTtoRPcoupling', **kwargs)
    # Specify spring parameters
    attach_spring(region1, 1, 'x-spring-ground', 'Ground', **kwargs)
    attach_spring(region1, 2, 'y-spring-ground', 'Ground', **kwargs)


def CoupleIpMTsToCentrosomes(**kwargs):
    """
    Connect and couple interpolar microtubules to centrosomes

    :param kwargs: object

    :return: None
    """
    # Provide coupling constraint parameters
    influenceRadius = WHOLE_SURFACE
    couplingType = DISTRIBUTING
    weightingMethod = LINEAR
    a = return_assembly(**kwargs)
    for index, instance in enumerate(a.instances.keys()):
        if 'ipMT' in instance and 'connector' not in instance:
            vert = a.instances[instance].vertices[0].pointOn
            masterPoint = a.instances[instance].vertices.findAt(vert, )
            masterRegion = a.Set(vertices=masterPoint, name='master-' + instance)
            if index % 2 == 0:
                s1 = a.instances['centrosome-right'].faces
                side1Faces1 = s1.getSequenceFromMask(mask=('[#1 ]',), )
                slaveRegion = a.Surface(side1Faces=side1Faces1, name='RightCentrosomeTo' + instance + 'coupling')
            else:
                s1 = a.instances['centrosome-left'].faces
                side1Faces1 = s1.getSequenceFromMask(mask=('[#1 ]',), )
                slaveRegion = a.Surface(side1Faces=side1Faces1, name='LeftCentrosomeTo' + instance + 'coupling')
            coupling_constraint(region1=masterRegion,
                                region2=slaveRegion,
                                influenceRadius=influenceRadius,
                                couplingType=couplingType,
                                weightingMethod=weightingMethod,
                                name=instance + 'toCentrosomeCoupling',
                                **kwargs)


def CoupleAMTs(**kwargs):
    """
    Couple astral microtubules with springs

    :param kwargs: object

    :return: None
    """
    regionsRight = []
    regionsLeft = []
    pointsRight = []
    pointsLeft = []

    # Loop over the positions of each aMT and create coupling regions
    aMTnames = kwargs['aMTnames']
    for i in range(len(aMTnames)):
        if i % 2 == 0:  # right pole
            refPoint, region = couple_nearest_aMTs(i, **kwargs)
            # Store reference point objects and associated regions
            pointsRight.append(refPoint)
            regionsRight.append(region)
        else:  # left pole
            refPoint, region = couple_nearest_aMTs(i, **kwargs)
            # Store reference point objects and associated regions
            pointsLeft.append(refPoint)
            regionsLeft.append(region)
    # Define coupling between picked regions
    position = kwargs['aMTposition']
    positionRight = position[0::2]
    region = find_nearest(positionRight, regionsRight)
    attach_spring(region, 1, 'springs-right', 'Pair', **kwargs)
    positionLeft = position[1::2]
    region = find_nearest(positionLeft, regionsLeft)
    attach_spring(region, 1, 'springs-left', 'Pair', **kwargs)


def CoupleAMTsToCentrosomes(**kwargs):
    """
    Connect astral microtubules to centrosomes

    :param kwargs: object

    :return: None
    """
    # parametyers
    influenceRadius = WHOLE_SURFACE
    couplingType = DISTRIBUTING
    weightingMethod = LINEAR
    a = return_assembly(**kwargs)
    for index, instance in enumerate(a.instances.keys()):
        if 'aMT' in instance:
            vert = a.instances[instance].vertices[0].pointOn
            masterPoint = a.instances[instance].vertices.findAt(vert,)
            masterRegion = a.Set(vertices=masterPoint, name='master-'+instance)
            if index % 2 == 0:
                s1 = a.instances['centrosome-right'].faces
                side1Faces1 = s1.getSequenceFromMask(mask=('[#1 ]',), )
                slaveRegion = a.Surface(side1Faces=side1Faces1, name='RightCentrosomeTo'+instance+'coupling')
            else:
                s1 = a.instances['centrosome-left'].faces
                side1Faces1 = s1.getSequenceFromMask(mask=('[#1 ]',), )
                slaveRegion = a.Surface(side1Faces=side1Faces1, name='LeftCentrosomeTo' + instance + 'coupling')

            coupling_constraint(region1=masterRegion,
                                region2=slaveRegion,
                                influenceRadius=influenceRadius,
                                couplingType=couplingType,
                                weightingMethod=weightingMethod,
                                name=instance+'toCentrosomeCoupling',
                                **kwargs)
