from abaqus import *
from abaqusConstants import *
import math
from SpindleAssembly.AddComponents import return_assembly
from SpindleAssembly.AddComponents import create_MT_instance


def assign_ipMTs_right(pos, i, **kwargs):
    """
    Position interpolar microtubules to start from the right centrosome

    :param pos: (x, y) -> coordinates of MT

    :param i: number of MT

    :param kwargs: object

    :return: MTname -> (string) name of ipMT,
             zBeginPos -> (float) z coordinate of the starting end of ipMT in global
             coordinates
             zEndPos -> (float) z coordinate of the growing end of ipMT in global
             coordinates
             l -> (float) length of ipMT
    """
    # Call an assembly where ipMTs are to be placed
    a = return_assembly(**kwargs)

    # Create MT instance
    l, p, MTname = create_MT_instance('ipMT', i, **kwargs)

    # Rotate ipMT to coincide with global z axis
    a.rotate(instanceList=(MTname,),
             axisPoint=(0.0, 0.0, 0.0),
             axisDirection=(0.0, 1.0, 0.0), angle=270.0)

    # Translate ipMT and position their starting ends on the
    # surfaces of the right centrosome sphere
    centrosomeRadius = kwargs['CentrosomeRadius']
    x = pos[0]
    y = pos[1]
    zpos = math.sqrt(abs(centrosomeRadius**2 - x**2 - y**2))
    zEndPos = zpos + l
    zBeginPos = zpos
    a.translate(instanceList=(MTname,), vector=(x, y, zpos))

    return MTname, zBeginPos, zEndPos, l


def assign_ipMTs_left(pos, i, **kwargs):
    """
    Position interpolar microtubules to start from the left centrosome

    :param pos: (x, y) -> coordinates of MT

    :param i: number of MT

    :param kwargs: object

    :return: MTname -> (string) name of ipMT,
             zBeginPos -> (float) z coordinate of the starting end of ipMT in global
             coordinates
             zEndPos -> (float) z coordinate of the growing end of ipMT in global
             coordinates
             l -> (float) length of ipMT
    """

    # Call an assembly where ipMTs are to be placed
    a = return_assembly(**kwargs)

    # Create MT instance
    l, p, MTname = create_MT_instance('ipMT', i, **kwargs)

    # Rotate ipMT to coincide with global z axis
    a.rotate(instanceList=(MTname,),
             axisPoint=(0.0, 0.0, 0.0),
             axisDirection=(0.0, 1.0, 0.0), angle=90.0)

    # Translate ipMT and position their starting ends on the
    # surface of the left centrosome sphere
    x = pos[0]
    y = pos[1]
    i = kwargs['index']
    spindleLength = kwargs['spindleLength']
    centrosomeRadius = kwargs['CentrosomeRadius']
    zpos = spindleLength - math.sqrt(abs(centrosomeRadius**2 - x**2 - y**2))
    zBeginPos = zpos
    zEndPos = zpos - l
    a.translate(instanceList=(MTname,), vector=(x, y, zpos))

    return MTname, zBeginPos, zEndPos, l


def generate_aMT_position(type, **kwargs):
    """
    Calculate the position of astral microtubules within the spindle

    :param type: 'right', 'left' -> either left or right centrosome bound aMTs specified

    :param kwargs: object

    :return: pos -> (x, y, z) global coordinates of the starting end of aMT
             theta -> orientation angle 1 in spherical coordinates
             phi -> orientation angle 2 in spherical coordinates
    """
    import random

    # Return main parameters of the Spindle
    spindleLength = kwargs['spindleLength']
    centrosomeRadius = kwargs['CentrosomeRadius']

    # Generate random angles that define position on the sphere
    theta = math.radians(180 * random.random() - 90)
    phi = math.radians(360 * random.random() - 180)

    # Calculate Cartesian position
    x = centrosomeRadius * math.cos(theta) * math.cos(phi)
    y = centrosomeRadius * math.sin(theta)

    # Check to which pole the aMT should be assigned
    if type == 'right':
        z = -centrosomeRadius * math.cos(theta) * math.sin(phi)
    elif type == 'left':
        z = spindleLength - centrosomeRadius * math.cos(theta) * math.sin(phi)
    else:
        raise ValueError('Invalid position type')

    pos = (x, y, z)
    return pos, theta, phi


def position_aMT(a, MTname, type, **kwargs):
    """
    Position astral microtubule within the spindle

    :param a: current assembly

    :param MTname: name of aMT

    :param type: 'left' or 'right' -> specifies to which pole the aMT is attached

    :param kwargs: object

    :return: pos -> (x, y, z) global cartesian coordinates of the starting end of aMT
    """

    # Generate random position at the right or left pole
    pos, theta, phi = generate_aMT_position(type, **kwargs)

    # Place aMT with respect to the calculated position
    a.rotate(
        instanceList=(MTname,),
        axisPoint=(0.0, 0.0, 0.0),
        axisDirection=(0.0, 0.0, 1.0),
        angle=math.degrees(theta))
    a.rotate(
        instanceList=(MTname,),
        axisPoint=(0.0, 0.0, 0.0),
        axisDirection=(0.0, 1.0, 0.0),
        angle=math.degrees(phi))
    a.translate(instanceList=(MTname,), vector=pos)

    return pos


def assign_aMT(type, **kwargs):
    """
    Add and position astral microtubules

    :param type: 'right' or 'left' -> specify the pole that aMTs should be attached to

    :param kwargs: object

    :return: aMTnames -> (list) list of aMT names
             positions -> (list of tuples) global coordinates of starting end of each aMT
    """
    aMTnames = []
    positions = []
    a = return_assembly(**kwargs)
    # Create aMTs and connect either to the right or to the left pole
    for i in range(kwargs['aMTnumber']):
        l, p, MTname = create_MT_instance(type, i, **kwargs)
        aMTnames.append(MTname)
        if i % 2 == 0:  # Right pole
            pos = position_aMT(a, MTname, 'right', **kwargs)
            positions.append(pos)
        else:  # Left pole
            pos = position_aMT(a, MTname, 'left', **kwargs)
            positions.append(pos)

    return aMTnames, positions