import itertools
from abaqus import *
from abaqusConstants import *
from interaction import *
from SpindleAssembly import AddComponents as ah
from SpindleAssembly.AddComponents import return_assembly


def generate_connector_positions( pos1, pos2, MTtype, **kwargs ):
    """
    Generate random points along the connection region of two ipMTs

    :param pos1: (x, y, zBeginPos, zEndPos, l) -> coordinates of the free end of first ipMT

    :param pos2: (x, y, zBeginPos, zEndPos, l) -> coordinates of the free end of the second ipMT

    :param MTtype: 'parallel' or 'antiparallel' -> defines if the MT pair is parallel or antiparallel

    :param kwargs: dictionary -> keyword arguments

    :return: numpy array of connector positions along z coordinate
    """
    import random

    # Calculate the extent of connectors positioning
    interval = abs(abs(pos1[3]) - abs(pos2[3]))

    # Return an array with random z positions for the connectors
    z = [random.uniform(0, interval) for n in range(kwargs['Nconnectors'])]

    if MTtype == 'antiparallel':
        # Calculate positions of connectors on ipMT-1
        z1 = sorted([pos1[4] - zi for zi in z])
        # Calculate position of connectors on ipMT-2
        z2 = sorted([pos2[4] - interval + zi for zi in z])
        # Calculate connector positions in global coordinates
        z1_global = sorted([pos1[3] - zi for zi in z])
        z2_global = sorted([pos2[3] - interval + zi for zi in z])
    elif MTtype == 'parallel':
        # Calculate positions of connectors on ipMT-1
        z1 = sorted([pos1[3] - zi for zi in z])
        # Calculate position of connectors on ipMT-2
        z2 = sorted([pos2[3] - zi for zi in z])
    else:
        raise ValueError('Unsupported MT orientation type')

    return z1, z2, z1_global, z2_global


def generate_partition_points(z, MTname, **kwargs):
    """
    Partition individual ipMT with points corresponding to connector positions

    :param z: numpy array of connector positions along z coordinate

    :param MTname: name of the ipMT to be partitioned

    :param kwargs: dictionary -> keyword arguments

    :return: list -> partition names
    """
    modelname = kwargs['modelname']
    # Loop over connector position array and partition ipMT
    for index, zcord in enumerate(z):
        # Create datum points
        p = mdb.models[modelname].parts[MTname]
        p.DatumPointByCoordinate(coords=(zcord, 0, 0))
        # Partition by datum point
        p = mdb.models[modelname].parts[MTname]
        e, v, d = p.edges, p.vertices, p.datums
        p.PartitionEdgeByPoint(edge=e[index], point=d[4+2*index])


def position_connectors(pos1, pos2, z, connectornames, alpha, **kwargs):
    """
    Position connectors on the associated ipMT by z coordinate

    :param pos1: (x, y, zBeginPos, zEndPos, l) -> coordinates of the first mt of the pair

    :param pos2: (x, y, zBeginPos, zEndPos, l) -> coordinates of the second mt of the pair

    :param z: list connector positions

    :param connectornames: list connector names

    :param alpha: float orientation of connectors

    :param kwargs: dictionary -> keyword arguments

    :return: global coordinates of connectors between a pair of ipMTs
    """
    a = ah.return_assembly(**kwargs)
    # Position connectors
    x0 = pos1[0]
    x1 = pos2[0]
    y0 = pos1[1]
    z_global = [zi + pos1[3] - pos1[4] for zi in z]
    for index, connectorname in enumerate(connectornames):
        if x1 - x0 > 0:
            angle = alpha
        elif x1 - x0 < 0:
            angle = 180 + alpha
        else:
            alpha = 0
        a.rotate(instanceList=(connectorname, ),
                 axisPoint=(0.0, 0.0, 0.0),
                 axisDirection=(0.0, 0.0, 1.0),
                 angle=-90)
        a.rotate(instanceList=(connectorname, ),
                 axisPoint=(0.0, 0.0, 0.0),
                 axisDirection=(0.0, 0.0, 1.0),
                 angle=angle)
        a.translate(instanceList=(connectorname, ), vector=(x0, y0, z_global[index]))
    return z_global


def pick_vertices(mtname, data, **kwargs):
    """
    Pick the connecting ends of all connectors associated with the current ipMT

    :param mtname: string name of the ipMT

    :param data: list containing all ipMTs and connectors of the model

    :param kwargs: dictionary -> keyword arguments

    :return: list -> connecting vertices that belong to connectors associated with the current ipMT
    """
    a = return_assembly(**kwargs)
    # Store name, end_id and position of connectors attached to each ipMT
    connector_data = []
    for sublist in data:
        if mtname == sublist[0]:
            connector_data.append([a.instances[sublist[1]].vertices,
                                   sublist[3],
                                   a.instances[sublist[1]].vertices[sublist[3]].pointOn,
                                   sublist[1]])

    # Sort connectors by position
    if int(mtname[-1]) % 2 == 0:  # if ipMT is right
        connector_data.sort(key=lambda tup: tup[2])
    else:  # else ipMT is left
        connector_data.sort(key=lambda tup: tup[2], reverse=True)

    return connector_data


def attach_connectors(data, **kwargs):
    """
    Create attachment wire between two vertexes and define the properties of connection

    :param data: list containing all ipMTs and connectors of the model

    :param kwargs: dictionary -> keyword arguments

    :return: Null
    """
    a = return_assembly(**kwargs)
    a.regenerate()
    # Create connector section
    modelname = kwargs['modelname']

    # Iterate through unique ipMT names
    mt_names = [sublist[0] for sublist in data]
    mt_names = set(mt_names)
    for mtname in mt_names:
        for sublist in data:
            if mtname == sublist[0]:
                connector_end = sublist[3]
                connector_name = sublist[1]

                # Link points of ipMT with connector ends
                end = a.instances[connector_name].vertices[connector_end].pointOn
                masterPoint = a.instances[connector_name].vertices.findAt(end,)
                masterName = 'master-'+connector_name+'-'+str(connector_end)
                masterRegion = a.Set(vertices=masterPoint,
                                     name=masterName)
                slaveEdge = a.instances[mtname].edges
                slaveName = 'slave-'+mtname
                slaveRegion = a.Set(edges=slaveEdge, name=slaveName)
                mdb.models[modelname].Coupling(name=masterName+slaveName,
                                               controlPoint=masterRegion,
                                               surface=slaveRegion,
                                               influenceRadius=0.2,
                                               couplingType=DISTRIBUTING,
                                               weightingMethod=LINEAR,
                                               localCsys=None, u1=ON, u2=ON, ur3=ON,
                                               adjust=False)
