import os
os.environ['OMP_NUM_THREADS'] = '1'
import sys

project_base_dir = os.path.abspath(os.path.join(os.path.dirname(os.path.abspath(__file__)), './')) + '/assembly/'


sys.path.append(project_base_dir)
sys.path.append(project_base_dir + '/baselines/')
sys.path.append(project_base_dir + '/assets/')

import json
import numpy as np
from assembly.assets.subdivide import subdivide_to_size
from assembly.baselines.run_joint_plan import PyPlanner
from assembly.examples.run_joint_plan import get_planner
from assembly.assets.process_mesh import process_mesh

from spatialmath.base import *
from spatialmath import *
import shutil
from scipy.spatial.transform import Rotation

class FS:
    def readJSON(path: str):
        return json.loads((open(path)).read())

    def writeFile(data, filePath, fileName):

        file_to_open = filePath + fileName

        f = open(file_to_open, 'w', )

        f.write(data)

    def readFile(path: str):
        return open(path).read()

    def readFilesTypeFolder(pathFolder: str, fileType='.json'):
        return os.listdir(pathFolder)
    def readFolder(pathFolder: str):
        return list( map(lambda el: pathFolder + '/' + el,  os.listdir(pathFolder) ))
        

def listGetFirstValue(iterable, default=False, pred=None):
    return next(filter(pred, iterable), default)


def filterModels(filterModels, filterModelsDescription):
    models = []
    for el in filterModelsDescription:
        models.append(listGetFirstValue(
            filterModels, None, lambda x:  x.name == el))
    return models
def meshTransformation():
    from argparse import ArgumentParser
    parser = ArgumentParser()
    parser.add_argument('--asp-path', type=str, required=True)
    args = parser.parse_args()
    aspDir = args.asp_dir
    if(aspDir == None):
        args.print_helper()
    if(aspDir[aspDir.__len__() - 1] != '/'):
        aspDir += '/'
    assemblys = FS.readFolder(aspDir + 'assembly')
    assemblyDirNormalize = []
    for el in assemblys:
        try:
            
            process_mesh(source_dir=el, target_dir=el+'/process/',subdivide=el, verbose=True)
            assemblyDirNormalize.append(el + '/process/')
        except Exception as e:
            print('ERRROR:')
            print(e)
    for el in assemblyDirNormalize:
        asset_folder = os.path.join(project_base_dir, aspDir)
        assembly_dir = os.path.join(asset_folder, el)
        planner = get_planner('bfs')(assembly_dir,assembly_dir, 0, [1], False, 'sdf', 0.05, 0.01, 100, 100, True)
        
         
        status, t_plan, path = planner.plan(
            120, seed=1, return_path=True, render=False, record_path=None
        )
        coords = []
        
        for k in path:
            seMatrix = SE3(k)
            euler = seMatrix.eul()
            coord = seMatrix.A[0:3,3]
            rot = Rotation.from_euler('xyz', euler, degrees=True).as_quat()
            coords.append({ 'quadrelion': [rot[0], rot[1], rot[2], rot[3]], 'xyz':[coord[0], coord[1], coord[2]],'euler': [euler[0], euler[1], euler[2]]})           
        planingObject = {
            "time": t_plan,
            "insertion_path": coords,
            "status":status,
        } 
        FS.writeFile(json.dumps(planingObject),el[0:el.__len__() - 8], 'insertion_path.json' )
         
        try:
            planner = PyPlanner(assembly_dir, 'process', still_ids=[1],   )
            status, t_plan, path = planner.plan(
                planner_name=args.planner, 
                step_size=args.step_size, 
                max_time=args.max_time, 
                seed=args.seed, 
                return_path=True, 
                simplify=args.simplify, 
                render=args.render
            )

            print(f'Status: {status}, planning time: {t_plan}')

            if args.save_dir is not None:
                planner.save_path(path, args.save_dir, args.n_save_state)    
        except Exception as e:
            print(e)
       
    
meshTransformation()