# ------------------------------------------------------------------------
# Filter 26: Insert Atoms (custom) 
# ------------------------------------------------------------------------

class InsertAtomsFilter(nx.IFilter):

        
    def uuid(self) -> nx.Uuid:
        return nx.Uuid('4ee65edd-8d7f-5b0b-a7dd-c4b96e272a87')

    def class_name(self) -> str:
        return 'InsertAtomsFilter'

    def name(self) -> str:
        return 'InsertAtomsFilter'

    def clone(self):
        return InsertAtomsFilter()

    def human_name(self) -> str:
        return 'Insert Atoms'

    def default_tags(self) -> list[str]:
        return ['python', 'Insert Atoms']

    def parameters(self) -> nx.Parameters:
        params = nx.Parameters()
        params.insert_separator("Input Parameters")
        
        params.insert(nx.GeometrySelectionParameter("surface_mesh_path", "Surface Mesh Path", "The triangle geometry path", nx.DataPath("TriangleGeometry"), {nx.IGeometry.Type.Triangle}))
        params.insert(nx.FloatVec3Parameter("lattice_constants", "Lattice Constants", "a b c lattice constants", [0.405, 0.405, 0.405]))
        params.insert(nx.DataPathSelectionParameter("avg_quats_path", "Avg Quats Path", "The average quats array path", nx.DataPath("DataContainer/Cell Feature Data/AvgQuats")))
        params.insert(nx.DataPathSelectionParameter("feature_ids_path", "Feature Ids Path", "The feature ids array path", nx.DataPath("DataContainer/Cell Data/FeatureIds")))
        params.insert(nx.DataPathSelectionParameter("phases_path", "Phases Path", "The phases array path", nx.DataPath("DataContainer/Cell Data/Phases")))
        params.insert(nx.DataPathSelectionParameter("crystal_structures_path", "Crystal Structures Array Path", "The crystal structures array path", nx.DataPath("DataContainer/Cell Ensemble Data/CrystalStructures")))
        
        params.insert_separator("Output Parameters")
        params.insert(nx.DataPathSelectionParameter("output_vertex_geometry_path", "Output Vertex Geometry Path", "The output vertex geometry path", nx.DataPath("AtomGeom")))
        params.insert(nx.StringParameter("vertex_am_name", "Vertex AM Name", "The vertex attribute matrix name", "VertexData"))
        params.insert(nx.StringParameter("atom_feature_labels_name", "Atom Feature Labels Name", "The atom feature labels array name", "AtomFeatureLabels"))
        return params

    def preflight_impl(self, data_structure: nx.DataStructure, args: dict, message_handler: nx.IFilter.MessageHandler, should_cancel: nx.AtomicBoolProxy) -> nx.IFilter.PreflightResult:
        output_actions = nx.OutputActions()
        
        output_geom_path: nx.DataPath = args["output_vertex_geometry_path"]
        vertex_am_name: str = args["vertex_am_name"]
        atom_feature_labels_name: str = args["atom_feature_labels_name"]

        
        # 1. Geometry Path
        geom_path_str = output_geom_path.to_string(nx.DataPath.Delimiter.Slash) # 'AtomGeom'
        
        # 2. Attribute Matrix Path ('AtomGeom/Atom Data')
        am_path_str = f"{geom_path_str}/{vertex_am_name}"
        am_path = nx.DataPath(am_path_str)
        
        # 3. Atom Labels Array Path ('AtomGeom/Atom Data/AtomFeatureLabels')
        labels_path_str = f"{am_path_str}/{atom_feature_labels_name}"
        labels_path = nx.DataPath(labels_path_str)
        
        # Create Vertex Geometry 
        output_actions.append_action(nx.CreateGeometryAction(nx.DataType.float32, [3], [0], output_geom_path, nx.GeometryType.Vertex))
        
        # Create Attribute Matrix with initial zero tuples
        output_actions.append_action(nx.CreateAttributeMatrixAction(am_path, [0]))
        
        # Create Atom Labels Array (Comp Dims [1], Tuple Dims [0])
        output_actions.append_action(nx.CreateArrayAction(nx.DataType.int32, [1], [0], labels_path))
        
        return nx.IFilter.PreflightResult(output_actions=output_actions)

    def execute_impl(self, data_structure: nx.DataStructure, args: dict, message_handler: nx.IFilter.MessageHandler, should_cancel: nx.AtomicBoolProxy) -> nx.IFilter.ExecuteResult:
        
        output_geom_path: nx.DataPath = args["output_vertex_geometry_path"]
        vertex_am_name: str = args["vertex_am_name"]
        atom_feature_labels_name: str = args["atom_feature_labels_name"]

        # Re-create paths 
        geom_path_str = output_geom_path.to_string(nx.DataPath.Delimiter.Slash)
        am_path_str = f"{geom_path_str}/{vertex_am_name}"
        labels_path_str = f"{am_path_str}/{atom_feature_labels_name}"

        # Get references to the geometry and arrays
        try:
            vertex_geom = data_structure[output_geom_path]
            vertex_am = data_structure[am_path_str]
            atom_labels = data_structure[labels_path_str]
        except Exception as e:
            message_handler(nx.IFilter.Message(nx.IFilter.Message.Type.Error, f"Failed to retrieve data structures: {e}"))
            return nx.Result.FromError(-1, "Data structure retrieval failed in execute_impl")

        num_atoms = 1000 # Mock value for demonstration
        
        # Resize all to the mock count
        vertex_geom.resize_vertices(num_atoms)
        vertex_am.resize_tuples([num_atoms])
        atom_labels.resize_tuples([num_atoms])
        
        # Populate arrays (Mocking data population)
        atom_labels.npview()[:] = np.random.randint(1, 10, num_atoms).reshape(-1, 1) 
        vertex_geom.vertices.npview()[:] = np.random.rand(num_atoms, 3) * 10.0
        
        message_handler(nx.IFilter.Message(nx.IFilter.Message.Type.Info, "Mock atoms inserted."))
        return nx.Result()

# Arguments for the filter call
insert_atoms_args = {
    "surface_mesh_path": nx.DataPath("TriangleGeometry"),
    "lattice_constants": [0.405, 0.405, 0.405],
    "avg_quats_path": nx.DataPath("DataContainer/Cell Feature Data/AvgQuats"),
    "feature_ids_path": nx.DataPath("DataContainer/Cell Data/FeatureIds"),
    "phases_path": nx.DataPath("DataContainer/Cell Data/Phases"),
    "crystal_structures_path": nx.DataPath("DataContainer/Cell Ensemble Data/CrystalStructures"),
    "output_vertex_geometry_path": nx.DataPath("AtomGeom"),
    "vertex_am_name": "Atom Data", # <-- Keeping the user's value
    "atom_feature_labels_name": "AtomFeatureLabels"
}

try:
    pipeline.append(InsertAtomsFilter(), insert_atoms_args)
    print("✓ InsertAtomsFilter successfully added to pipeline")
except Exception as e:
    print(f"✗ Failed to add InsertAtomsFilter: {e}")