from Parts import standard_parts as sparts
from Parts import GeometryBuilder as gb
from abaqus import *
from abaqusConstants import *
import math


def return_assembly(**kwargs):
    """
    Return the current assembly

    :param kwargs: object

    :return: a -> (object) current assembly
    """
    modelname = kwargs['modelname']
    a = mdb.models[modelname].rootAssembly
    return a


def add_centrosome(position, **kwargs):
    """
    Generate and add centrosome

    :param position: 'right' or 'left' -> positioning of centrosome in the spindle

    :param kwargs: object

    :return: Null
    """
    modelname = kwargs['modelname']
    if position == 'right':
        sparts.centrosome(**kwargs)
        p = mdb.models[modelname].parts['centrosome']
        a = return_assembly(**kwargs)
        a.Instance(name='centrosome-right', part=p, dependent=ON)
        a = mdb.models[modelname].rootAssembly
        a.rotate(
            instanceList=('centrosome-right',),
            axisPoint=(0.0, 0.0, 0.0),
            axisDirection=(1.0, 0.0, 0.0),
            angle=90)
    elif position == 'left':
        sparts.centrosome(**kwargs)
        p = mdb.models[modelname].parts['centrosome']
        a = return_assembly(**kwargs)
        a.Instance(name='centrosome-left', part=p, dependent=ON)
        a = mdb.models[modelname].rootAssembly
        # Translate left centrosome and position it at the left pole
        a.translate(
            instanceList=('centrosome-left',),
            vector=(0, 0, kwargs['spindleLength']))
        a.rotate(
            instanceList=('centrosome-left',),
            axisPoint=(0.0, 0.0, kwargs['spindleLength']),
            axisDirection=(1.0, 0.0, kwargs['spindleLength']),
            angle=90)


def add_microtubule(l, type, i, **kwargs):
    """
    Generate and add microtubules to the assembly

    :param l: length of the microtubule

    :param type: 'ipMT' or 'aMT' -> Either interpolar or astral microtubule

    :param i: number of MT

    :param kwargs: object

    :return: p -> microtubule part, MTname -> (string) name of the MT
    """

    # generate name for microtubule:
    MTname = gb.create_mt_name(l, type, i)

    # generate microtubule
    modelname = kwargs['modelname']
    sparts.microtubule(type, l, i, **kwargs)
    p = mdb.models[modelname].parts[MTname]
    return p, MTname


def add_connectors(pos1, pos2, MTname1, MTname2, z, **kwargs):
    """
    Generate and add connectors to couple ipMT pair

    :param pos1: (x, y, z) of the 1-st ipMT of the pair to be connected

    :param pos2: (x, y, z) of the 2-nd ipMT of the pair to be connected

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

    :param kwargs: dictionary -> keyword arguments

    :return: list -> list of connector names, float ->  angle of connector orientation
    """

    # Calculate orientation and length of connector
    alpha = math.degrees(math.atan((pos2[1] - pos1[1]) / (pos2[0] - pos1[0])))
    length = math.sqrt((pos2[0] - pos1[0])**2 + (pos2[1] - pos1[1])**2)

    # Add connectors
    modelname = kwargs['modelname']
    connectornames = []
    for i in range(len(z)):
        connectorname = 'connector-' + MTname1 + '-' + MTname2 + '-' + str(i)
        sparts.connector(i, length, connectorname, **kwargs)
        p = mdb.models[modelname].parts[connectorname]
        a = return_assembly(**kwargs)
        a.Instance(name=connectorname, part=p, dependent=OFF)
        connectornames.append(connectorname)
    return connectornames, alpha


def generate_MT_length( type, **kwargs ):
    """
    calculate length of an MT based on its type

    :param type: 'ipMT' or 'aMT' -> interpolar or astral microtubules

    :param kwargs: object

    :return: l -> (float) length of the MT
    """
    import random
    if type == 'ipMT':
        l = kwargs['lengthInterval'][0] * random.random() + kwargs['lengthInterval'][1]
    elif type == 'aMT':
        l = kwargs['aMTlength']
    else:
        raise ValueError('Only ipMT and aMT are valid types')

    return l


def create_MT_instance(type, i, **kwargs):
    """
    create abaqus instance that contain a single microtubule

    :param type: 'ipMT' or 'aMT' -> interpolar or astral microtubule

    :param i: number of microtubule

    :param kwargs: object

    :return: l -> (float) length of MT,
             p -> (object) MT part,
             MTname -> (string) name of MT
    """

    # Calculate a random length of MT
    l = generate_MT_length(type, **kwargs)

    # Create MT part
    p, MTname = add_microtubule(l, type, i, **kwargs)

    # Create MT instance
    a = return_assembly(**kwargs)
    a.Instance(name=MTname, part=p, dependent=ON)
    return l, p, MTname