from part import *
from material import *
from section import *
from assembly import *
from step import *
from interaction import *
from load import *
from mesh import *
from optimization import *
from job import *
from sketch import *
from visualization import *
from connectorBehavior import *
import regionToolset

import numpy as np
import os

model_name = "aba1"  # enter model name
part_name = "PART-1"  # enter part name

radius = 1e-03  # radius of sphere to pick nodes within (default value for point selection = 1e-3)
tol = 1e-6  # default tolerance value for finding the nodes 1e-6


# CSV FILE INFO CONTAINING CO-ORDINATES OF RPs & MASS value
directory = r"C:\Suvopam\ABA\08_vPumps_CD"
file_name = "cog_mass"
full_path = os.path.join(directory, file_name + ".csv")
arr = np.genfromtxt(full_path, unpack=False, skip_header=1, filling_values=np.NaN, delimiter=";")

m = mdb.models[model_name]
p = m.parts[part_name]
a = m.rootAssembly


def rp_coord_n_mass_creator(arr_):
    '''to create list of co-ordinates of RP and center of the sphere in the form of tuple'''
    arr_ = arr_[~np.isnan(arr_).any(axis=1)]  # dropna() for numpy
    n_row, n_col = arr_.shape
    p_cog_ = []
    mass_ = []
    for x in range(n_row):
        temp = []
        for y in range(3):
            temp.append(arr_[x][y])
        p_cog_.append(tuple(temp))
        mass_.append(arr_[x][3])
    return (p_cog_, mass_)


def rp_dict_creation(f_):
    '''creating RP dictionary'''
    f_dict = {}
    for i in f_.keys():
        try:
            redundant = f_[i].xValue  # to clean PART number key from selection
            f_dict[f_[i].id] = []
            f_dict[f_[i].id].append(f_[i].name)
            f_dict[f_[i].id].append(f_[i].xValue)
            f_dict[f_[i].id].append(f_[i].yValue)
            f_dict[f_[i].id].append(f_[i].zValue)
        except AttributeError:
            pass
    return f_dict


def sph_extrm(radius_, rp_coord_):
    '''to create extremes of the sphere '''
    p_extr_x_ = (rp_coord_[0] + radius_, rp_coord_[0] - radius_)
    p_extr_y_ = (rp_coord_[1] + radius_, rp_coord_[1] - radius_)
    p_extr_z_ = (rp_coord_[2] + radius_, rp_coord_[2] - radius_)
    return p_extr_x_, p_extr_y_, p_extr_z_


def search_inside_vol_n_create_mass(rp_dict_, a_, p_mass_, p_extr_x_, p_extr_y_, p_extr_z_):
    '''searches if rp co-ordinate is inside the sphere and creates mass if true'''
    rp_key_list = rp_dict_.keys()
    for i in rp_key_list:
        x_flag = False
        y_flag = False
        z_flag = False
        if p_extr_x_[0] + tol >= rp_dict_[i][1] >= p_extr_x_[1] - tol or p_extr_x_[1] + tol >= rp_dict_[i][1] >= p_extr_x_[0] - tol:
            x_flag = True
        if p_extr_y_[0] + tol >= rp_dict_[i][2] >= p_extr_y_[1] - tol or p_extr_y_[1] + tol >= rp_dict_[i][2] >= p_extr_y_[0] - tol:
            y_flag = True
        if p_extr_z_[0] + tol >= rp_dict_[i][3] >= p_extr_z_[1] - tol or p_extr_z_[1] + tol >= rp_dict_[i][3] >= p_extr_z_[0] - tol:
            z_flag = True
        if x_flag and y_flag and z_flag:
            # creating region
            r1 = a.referencePoints
            refPoints1 = (r1[i],)
            region = regionToolset.Region(referencePoints=refPoints1)
            # creating inertia mass
            a_.engineeringFeatures.PointMassInertia(name='Inertia-x_{}'.format(i), region=region, mass=p_mass_, alpha=0.0, composite=0.0)
        else:
            pass


def master_run():
    rp_dict = rp_dict_creation(f)
    p_cog, mass = rp_coord_n_mass_creator(arr)
    for x in range(len(mass)):
        p_extr_x, p_extr_y, p_extr_z = sph_extrm(radius, p_cog[x])
        search_inside_vol_n_create_mass(rp_dict, a, mass[x], p_extr_x, p_extr_y, p_extr_z)  # multi run


master_run()