import os
import stat as stat_module
import torch
import signal
import random
import numpy as np
import pandas as pd
from typing import List
from collections import defaultdict
from concurrent.futures import ThreadPoolExecutor
from torch_geometric.data import Dataset, Data
from torch_geometric.transforms import BaseTransform


def _is_valid_graph_path(p):
    """A single os.stat (regular file, non-empty). One syscall instead of two."""
    try:
        st = os.stat(p)
        return stat_module.S_ISREG(st.st_mode) and st.st_size > 0
    except OSError:
        return False


def filter_valid_graph_paths(graph_paths: list, num_workers: int = 32) -> list:
    """Filter out graph paths that don't exist or have zero size (truncated).

    Parallelized across threads: this is I/O-bound stat() traffic on a network
    filesystem, so a thread pool gives near-linear speedup. On the full dataset
    (~8.9M paths) the serial two-stat version took tens of minutes; this is
    ~1-2 min. Order is preserved (map keeps input order) and the result is
    identical to the serial check.
    """
    with ThreadPoolExecutor(max_workers=num_workers) as ex:
        flags = list(ex.map(_is_valid_graph_path, graph_paths, chunksize=2000))
    valid = [p for p, ok in zip(graph_paths, flags) if ok]
    skipped = len(graph_paths) - len(valid)
    if skipped:
        print(f"WARNING: Filtered out {skipped} missing/empty graph files")
    return valid


class GraphDataset(Dataset):
    def __init__(self, graph_paths: list, transform=None, pre_transform=None):
        super().__init__(None, transform, pre_transform)
        self.graph_paths = filter_valid_graph_paths(graph_paths)

    @property
    def processed_file_names(self):
        return self.graph_paths

    def len(self):
        return len(self.graph_paths)

    def get(self, idx):
        graph_path = self.graph_paths[idx]
        graph = torch.load(graph_path, weights_only=False)
        graph.graph_path = graph_path
        return graph


class NormalizeData(BaseTransform):
    def __init__(
        self,
        scale_dict: dict,
        attrs: List[str] = ["x", "edge_attr"],
        edge_attr_skip_cols: List[int] | None = None,
        min_std: float | None = None,
        clip: float | None = None,
    ):
        tensor_dict = defaultdict(lambda: defaultdict())  # convert the data to tensors
        for key, value in scale_dict.items():
            tensor_dict[key]["mean"] = torch.tensor(value["mean"])
            tensor_dict[key]["std"] = torch.tensor(value["std"])

        # Optional std floor: clamp tiny stds upward to prevent near-constant
        # dims (which have intrinsically tiny variance from the feature pipeline,
        # e.g. high-freq Fourier descriptors, GLCM moments) from inflating
        # outlier raw values by 1/std factors of 10^3-10^4.
        if min_std is not None:
            for key in tensor_dict:
                tensor_dict[key]["std"] = torch.clamp(
                    tensor_dict[key]["std"], min=min_std
                )

        # For specified edge_attr columns (e.g. column 0 = spatial distance),
        # set mean=0, std=1 so they pass through untouched.
        if edge_attr_skip_cols and "edge_attr" in tensor_dict:
            for col in edge_attr_skip_cols:
                tensor_dict["edge_attr"]["mean"][col] = 0.0
                tensor_dict["edge_attr"]["std"][col] = 1.0

        self.tensor_dict = tensor_dict
        self.attrs = attrs
        self.clip = clip
        self.edge_attr_skip_cols = edge_attr_skip_cols or []

    def forward(self, data: Data) -> Data:
        for store in data.stores:
            for key, value in store.items(*self.attrs):
                if value.numel() > 0:
                    mean = self.tensor_dict[key]["mean"]
                    std = self.tensor_dict[key]["std"]
                    value = (value - mean) / std
                    value = torch.nan_to_num(value, nan=0.0)
                    # Hard bound on standardized values to stop residual
                    # cross-dataset explosion (a feature far from the TCGA mean).
                    # Skip edge_attr columns left raw (e.g. col 0 = distance).
                    if self.clip is not None:
                        if key == "edge_attr" and self.edge_attr_skip_cols:
                            keep = value[:, self.edge_attr_skip_cols].clone()
                            value = value.clamp(-self.clip, self.clip)
                            value[:, self.edge_attr_skip_cols] = keep
                        else:
                            value = value.clamp(-self.clip, self.clip)
                    store[key] = value
        return data

    def __repr__(self) -> str:
        return f"{self.__class__.__name__}()"


class AddVirtualNode(BaseTransform):
    """Add a virtual node connected to every real node.

    Args:
        mean_edge_distance: Raw (un-normalised) mean spatial distance from
            the training set.  Used as column 0 of the virtual-node edge
            features so that the degree-normalised message weighting remains
            strictly positive.  Remaining columns (visual + relational
            features) are set to 0.0 because those columns *are* normalised,
            and 0 represents their mean.
    """

    def __init__(self, mean_edge_distance: float):
        self.mean_edge_distance = mean_edge_distance

    def forward(self, data: Data) -> Data:
        num_nodes = data.num_nodes
        device = data.x.device if data.x is not None else "cpu"

        virtual_node_feat = torch.zeros((1, data.x.size(-1)), device=device)
        data.x = torch.cat([data.x, virtual_node_feat], dim=0)

        # Keep any node-level label(s) aligned with the added virtual node: the
        # cell-level embedding path slices batch.label with the VN-inclusive ptr,
        # so a per-node label must gain one entry for the VN (placeholder -1).
        # Scalar / graph-level labels (size != num_nodes) are left untouched.
        for _lab in ("label", "labels"):
            _v = getattr(data, _lab, None)
            if torch.is_tensor(_v) and _v.dim() >= 1 and _v.size(0) == num_nodes:
                _pad = torch.full(
                    (1, *_v.shape[1:]), -1, dtype=_v.dtype, device=_v.device
                )
                setattr(data, _lab, torch.cat([_v, _pad], dim=0))

        row = torch.arange(num_nodes, device=device)
        col = torch.full((num_nodes,), num_nodes, device=device)
        new_edges = torch.stack([torch.cat([row, col]), torch.cat([col, row])], dim=0)
        data.edge_index = torch.cat([data.edge_index, new_edges], dim=1)

        data.virtual_node_index = torch.tensor([num_nodes], device=device)

        if data.edge_attr is not None:
            edge_dim = data.edge_attr.size(-1)
            new_edge_attr = torch.zeros(
                (2 * num_nodes, edge_dim), device=device
            )
            # Column 0 = raw spatial distance (not normalised); use training-
            # set mean so virtual-node edges have a typical positive weight.
            new_edge_attr[:, 0] = self.mean_edge_distance
            data.edge_attr = torch.cat([data.edge_attr, new_edge_attr], dim=0)

        return data

    def __repr__(self) -> str:
        return f"{self.__class__.__name__}(mean_edge_distance={self.mean_edge_distance})"


class GracefulKiller:
    kill_now = False

    def __init__(self):
        signal.signal(signal.SIGINT, self.exit_gracefully)
        signal.signal(signal.SIGTERM, self.exit_gracefully)

    def exit_gracefully(self, signum, frame):
        self.kill_now = True


def set_random_seed(seed):
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)
    torch.backends.cudnn.deterministic = True


def seed_worker(worker_id):
    worker_seed = torch.initial_seed() % 2**32
    np.random.seed(worker_seed)
    random.seed(worker_seed)


def split_data(data_df: pd.DataFrame, split_df: pd.DataFrame):

    train_samples = split_df.loc[split_df["split"] == "train", "sample_id"].unique()
    val_samples = split_df.loc[split_df["split"] == "val", "sample_id"].unique()
    test_samples = split_df.loc[split_df["split"] == "test", "sample_id"].unique()

    train_mask = data_df["sample_id"].isin(train_samples)
    val_mask = data_df["sample_id"].isin(val_samples)
    test_mask = data_df["sample_id"].isin(test_samples)

    return train_mask, val_mask, test_mask


def get_current_lr(optimizer):
    return optimizer.state_dict()["param_groups"][0]["lr"]


def create_optimizer(
    opt, model, lr, weight_decay, get_num_layer=None, get_layer_scale=None
):
    opt_lower = opt.lower()

    parameters = model.parameters()
    opt_args = dict(lr=lr, weight_decay=weight_decay)

    opt_split = opt_lower.split("_")
    opt_lower = opt_split[-1]
    if opt_lower == "adam":
        optimizer = torch.optim.Adam(parameters, **opt_args)
    elif opt_lower == "adamw":
        optimizer = torch.optim.AdamW(parameters, **opt_args)
    elif opt_lower == "adadelta":
        optimizer = torch.optim.Adadelta(parameters, **opt_args)
    elif opt_lower == "radam":
        optimizer = torch.optim.RAdam(parameters, **opt_args)
    elif opt_lower == "sgd":
        opt_args["momentum"] = 0.9
        return torch.optim.SGD(parameters, **opt_args)
    else:
        assert False and "Invalid optimizer"

    return optimizer


def load_checkpoint(checkpoint_fpath, model, optimizer, just_model=False):
    checkpoint = torch.load(checkpoint_fpath, weights_only=False)
    model.load_state_dict(checkpoint["model_state_dict"])
    if just_model:
        return model
    else:
        optimizer.load_state_dict(checkpoint["optimizer_state_dict"])
        return (
            model,
            optimizer,
            checkpoint["epoch"],
            checkpoint["best_loss"],
            checkpoint["run_id"],
        )


def save_checkpoint(checkpoint_fpath, model, optimizer, epoch, best_loss, run_id):
    torch.save(
        {
            "epoch": epoch,
            "model_state_dict": model.state_dict(),
            "optimizer_state_dict": optimizer.state_dict(),
            "best_loss": best_loss,
            "run_id": run_id,
        },
        checkpoint_fpath,
    )
    return