from abaqus import *
from abaqusConstants import *
import sys
from numpy import ones
from SpindleAssembly.PositionComponents import assign_aMT
from SpindleAssembly.PositionConnectors import attach_connectors
# sys.path.append("J:\ECS Research\Public\VLC\Spindle")
from SpindleAssembly import AddComponents as ac
from SpindleAssembly import PositionIpMTs
from SpindleAssembly import PositionConnectors
from Parts.GeometryBuilder import model


def generate_assembly(**kwargs):
    """
    Call modules to generate assembly

    :type kwargs: object

    :param kwargs:

    :return: kwargs -> object
    """
    create_model_assembly(**kwargs)
    add_and_position_centrosomes(**kwargs)
    aMTnames, kwargs = add_and_position_astral_mts(**kwargs)
    MTdata, kwargs, data = add_and_position_interpolar_mts_and_connectors(**kwargs)
    return MTdata, aMTnames, kwargs, data


def create_model_assembly(**kwargs):
    """
    Create Abaqus model and assembly

    :param kwargs: object

    :return: None
    """
    model(**kwargs)
    a = ac.return_assembly(**kwargs)
    a.DatumCsysByDefault(CARTESIAN)
    kwargs.update({'assembly': a})


def add_and_position_centrosomes(**kwargs):
    """
    Create centrosomes and position them at the right or left pole

    :param kwargs: model marameters. object

    :return: None
    """
    # Add right centrosome
    ac.add_centrosome('right', **kwargs)
    # Add left centrosome
    ac.add_centrosome('left', **kwargs)


def add_and_position_interpolar_mts_and_connectors(**kwargs):
    """
    Create and position interpolar microtubules and connects them with connectors

    :param kwargs: model parameters. object

    :return: MTdata -> (dictionary) Contains ipMT names, positions and connector positions
    """
    MTdata, ConnectorData, data = assign_ipMTs(**kwargs)
    attach_connectors(data, **kwargs)
    kwargs.update(dict(MTdata=MTdata, ConnectorData=ConnectorData))
    return MTdata, kwargs, data


def add_and_position_astral_mts(**kwargs):
    """
    Create and position astral microtubules

    :param kwargs: object

    :return: aMTnames -> (list) list of astral microtubule names
             kwargs -> object
    """
    aMTnames, aMTposition = assign_aMT('aMT', **kwargs)
    kwargs.update({'aMTnames': aMTnames, 'aMTposition': aMTposition})

    return aMTnames, kwargs


def assign_ipMTs(**kwargs):
    """
    Position the microtubules within the spindle along with the associated connectors

    :param kwargs: model parameters

    :type kwargs: dict

    :return: MTdata, ConnectorData, data

    :rtype: list, list, list
    """
    # Create dictionary to store MT metadata
    MTdata = {'MTnames': [],
              'MTpositions': [],
              'Points': {}}
    ConnectorData = {}
    data = []
    # Create ipMTs:
    NipMTs = kwargs['ipMTnumber']

    if NipMTs >= 1:
        MTname1, pos1, alpha1, kwargs = PositionIpMTs.add_first_ipMT(i=0, **kwargs)

    if NipMTs >= 2:
        MTname2, pos2, alpha2, kwargs = PositionIpMTs.add_second_ipMT(
            pos1, alpha1, i=1, **kwargs)
        # Generate points of connector attachment sites

        z12, z21, z12_global, z21_global = PositionConnectors.generate_connector_positions(
            pos1, pos2, MTtype='antiparallel', **kwargs)

        # Store all the metadata in the dictionary
        MTdata['MTnames'].append(MTname1)
        MTdata['MTnames'].append(MTname2)
        MTdata['MTpositions'].append(pos1)
        MTdata['MTpositions'].append(pos2)
        MTdata['Points'].update({'12': z12})
        MTdata['Points'].update({'21': z21})

        connectornames, alpha = ac.add_connectors(pos1, pos2, MTname1, MTname2, z12, **kwargs)
        PositionConnectors.position_connectors(pos1, pos2, z12, connectornames, alpha, **kwargs)
        for index, name in enumerate(connectornames):
            ConnectorData.update({name: [MTname1, (pos1[0], pos1[1], z12[index], z12_global[index]),
                                         MTname2, (pos2[0], pos2[1], z21[index], z21_global[index])]})
            data.append([MTname1, name, (pos1[0], pos1[1], z12_global[index]), 0])
            # data.append({'MT_name': MTname1, 'Con_name': name, 'Con_pointId': 0, 'MT_pointId': MTids[0],
            #              'pointCoordinate' : (pos1[0], pos1[1], z12_global[index])})
            data.append([MTname2, name, (pos2[0], pos2[1], z21_global[index]), 1])
            # MTids[0] += 1
            # MTids[1] += 1

    if NipMTs >= 3:
        MTname3, pos3, alpha3, kwargs = PositionIpMTs.add_third_ipMT(
            pos2, alpha2, i=2, **kwargs)
        z32, z23, z32_global, z23_global = PositionConnectors.generate_connector_positions(
            pos3, pos2, MTtype='antiparallel', **kwargs)

        MTdata['MTnames'].append(MTname3)
        MTdata['MTpositions'].append(pos3)
        MTdata['Points'].update({'32': z32})
        MTdata['Points'].update({'23': z23})

        connectornames, alpha = ac.add_connectors(pos3, pos2, MTname3, MTname2, z32, **kwargs)
        PositionConnectors.position_connectors(pos3, pos2, z32, connectornames, alpha, **kwargs)
        for index, name in enumerate(connectornames):
            ConnectorData.update({name: [MTname3, (pos3[0], pos3[1], z32[index], z32_global[index]),
                                         MTname2, (pos2[0], pos2[1], z23[index], z23_global[index])]})
            data.append([MTname3, name, (pos3[0], pos3[1], z32_global[index]), 0])
            data.append([MTname2, name, (pos2[0], pos2[1], z23_global[index]), 1])

    if NipMTs >= 4:
        MTname4, pos4, alpha4, kwargs = PositionIpMTs.add_fourth_ipMT(
            pos3, alpha2, alpha3, i=3, **kwargs)
        z34, z43, z34_global, z43_global = PositionConnectors.generate_connector_positions(
            pos3, pos4, MTtype='antiparallel', **kwargs)
        z14, z41, z14_global, z41_global = PositionConnectors.generate_connector_positions(
            pos1, pos4, MTtype='antiparallel', **kwargs)

        MTdata['MTnames'].append(MTname4)
        MTdata['MTpositions'].append(pos4)
        MTdata['Points'].update({'34': z34})
        MTdata['Points'].update({'43': z43})
        MTdata['Points'].update({'14': z14})
        MTdata['Points'].update({'41': z41})

        connectornames, alpha = ac.add_connectors(pos3, pos4, MTname3, MTname4, z34, **kwargs)
        PositionConnectors.position_connectors(pos3, pos4, z34, connectornames, alpha, **kwargs)
        for index, name in enumerate(connectornames):
            ConnectorData.update({name: [MTname3, (pos3[0], pos3[1], z34[index], z34_global[index]),
                                         MTname4, (pos4[0], pos4[1], z43[index], z43_global[index])]})
            data.append([MTname3, name, (pos3[0], pos3[1], z34_global[index]), 0])
            data.append([MTname4, name, (pos4[0], pos4[1], z43_global[index]), 1])

        connectornames, alpha = ac.add_connectors(pos1, pos4, MTname1, MTname4, z14, **kwargs)
        PositionConnectors.position_connectors(pos1, pos4, z14, connectornames, alpha, **kwargs)
        for index, name in enumerate(connectornames):
            ConnectorData.update({name: [MTname1, (pos1[0], pos1[1], z14[index], z14_global[index]),
                                         MTname4, (pos4[0], pos4[1], z41[index], z41_global[index])]})
            data.append([MTname1, name, (pos1[0], pos1[1], z14_global[index]), 0])
            data.append([MTname4, name, (pos4[0], pos4[1], z41_global[index]), 1])

    if NipMTs >= 5:
        MTname5, pos5, alpha5, kwargs = PositionIpMTs.add_fifth_ipMT(
            pos4, alpha4, i=4, **kwargs)

        z54, z45, z54_global, z45_global = PositionConnectors.generate_connector_positions(
            pos5, pos4, MTtype='antiparallel', **kwargs)

        MTdata['MTnames'].append(MTname5)
        MTdata['MTpositions'].append(pos5)
        MTdata['Points'].update({'45': z45})
        MTdata['Points'].update({'54': z54})

        connectornames, alpha = ac.add_connectors(pos5, pos4, MTname5, MTname4, z54, **kwargs)
        PositionConnectors.position_connectors(pos5, pos4, z54, connectornames, alpha, **kwargs)
        for index, name in enumerate(connectornames):
            ConnectorData.update({name: [MTname5, (pos5[0], pos5[1], z54[index], z54_global[index]),
                                         MTname4, (pos4[0], pos4[1], z45[index], z45_global[index])]})
            data.append([MTname5, name, (pos5[0], pos5[1], z54_global[index]), 0])
            data.append([MTname4, name, (pos4[0], pos4[1], z45_global[index]), 1])

    if NipMTs >= 6:
        MTname6, pos6, alpha6, kwargsAddMT = PositionIpMTs.add_sixth_ipMT(
            pos5, alpha4, alpha5, i=5, **kwargs)
        z56, z65, z56_global, z65_global = PositionConnectors.generate_connector_positions(
            pos5, pos6, MTtype='antiparallel', **kwargs)

        MTdata['MTnames'].append(MTname6)
        MTdata['MTpositions'].append(pos6)
        MTdata['Points'].update({'56': z56})
        MTdata['Points'].update({'65': z65})

        connectornames, alpha = ac.add_connectors(pos5, pos6, MTname5, MTname6, z56, **kwargs)
        PositionConnectors.position_connectors(pos5, pos6, z56, connectornames, alpha, **kwargs)
        for index, name in enumerate(connectornames):
            ConnectorData.update({name: [MTname5, (pos5[0], pos5[1], z56[index], z56_global[index]),
                                         MTname6, (pos6[0], pos6[1], z65[index], z65_global[index])]})
            data.append([MTname5, name, (pos5[0], pos5[1], z56_global[index]), 0])
            data.append([MTname6, name, (pos6[0], pos6[1], z65_global[index]), 1])

    # Create partitions that correspond to connector attachment sites
    MTpoints = {}
    for i in range(len(MTdata['MTnames'])):
        MTpoints.update({str(i+1): []})
        for key, value in MTdata['Points'].iteritems():
            if int(key[0]) == i+1:
                MTpoints[str(i + 1)] = sorted(MTpoints[str(i+1)] + value)
        PositionConnectors.generate_partition_points(
            MTpoints[str(i + 1)],
            MTdata['MTnames'][i], **kwargs)
    return MTdata, ConnectorData, data
