"""Data conversion utilities for the GNN router.

Converts raw numpy arrays into PyTorch Geometric Data objects suitable
for the bipartite GNN training pipeline.
"""

from __future__ import annotations

import numpy as np
import torch
from torch_geometric.data import Data


class FormData:
    """Converts numpy arrays to PyTorch Geometric Data objects."""

    def __init__(self, device: torch.device | str) -> None:
        self.device = device

    def formulation(
        self,
        task_features: np.ndarray,
        query_features: np.ndarray,
        llm_features: np.ndarray,
        source_nodes: list[int],
        dest_nodes: list[int],
        edge_features: np.ndarray,
        labels: np.ndarray,
        edge_mask: torch.Tensor,
        combined_edge: np.ndarray,
        train_mask: torch.Tensor,
        valid_mask: torch.Tensor,
        test_mask: torch.Tensor,
        query_task_labels: list[str] | None = None,
    ) -> Data:
        """Create a PyG Data object from raw arrays."""
        query_tensor = torch.as_tensor(query_features, dtype=torch.float, device=self.device)
        llm_tensor = torch.as_tensor(llm_features, dtype=torch.float, device=self.device)
        task_tensor = torch.as_tensor(task_features, dtype=torch.float, device=self.device)

        query_indices = list(range(len(query_tensor)))
        llm_indices = [i + len(query_indices) for i in range(len(llm_tensor))]
        dest_node_ids = [(i + 1 + source_nodes[-1]) for i in dest_nodes]

        edge_index = torch.tensor(
            [source_nodes, dest_node_ids],
            dtype=torch.long,
            device=self.device,
        )
        edge_weight = torch.as_tensor(
            edge_features,
            dtype=torch.float,
            device=self.device,
        ).reshape(-1, 1)
        combined = torch.as_tensor(combined_edge, dtype=torch.float, device=self.device)

        return Data(
            task_id=task_tensor,
            query_features=query_tensor,
            llm_features=llm_tensor,
            edge_index=edge_index,
            edge_attr=edge_weight,
            query_indices=query_indices,
            llm_indices=llm_indices,
            label=torch.as_tensor(labels, dtype=torch.float, device=self.device),
            edge_mask=torch.as_tensor(edge_mask, dtype=torch.bool, device=self.device),
            combined_edge=combined,
            train_mask=torch.as_tensor(train_mask, dtype=torch.bool, device=self.device),
            valid_mask=torch.as_tensor(valid_mask, dtype=torch.bool, device=self.device),
            test_mask=torch.as_tensor(test_mask, dtype=torch.bool, device=self.device),
            query_task_labels=query_task_labels,
        )
