import random

import igraph as ig
import numpy as np
import torch

from ..base import BatchMatrixData, MatrixDomain, TopologySkeleton, make_context_pairs, scale_and_shift


class FastGraphLaplacianGenerator:

    def __init__(self, ns: tuple[int, ...] | int, device: torch.device):
        if isinstance(ns, int):
            ns = (ns,)
        self.ns = ns
        self.device = device

    def generate_topology_cpu(self, n: int, m: int = 3) -> TopologySkeleton:
        g = ig.Graph.Barabasi(n, m, directed=False)
        edges = g.get_edgelist()

        rows = np.array([e[0] for e in edges] + [e[1] for e in edges], dtype=np.int64)
        cols = np.array([e[1] for e in edges] + [e[0] for e in edges], dtype=np.int64)

        diag_idx = np.arange(n, dtype=np.int64)
        all_rows = np.concatenate([rows, diag_idx])
        all_cols = np.concatenate([cols, diag_idx])

        degrees = np.array(g.degree(), dtype=np.float64)

        return TopologySkeleton(
            rows=all_rows,
            cols=all_cols,
            n=n,
            domain=MatrixDomain.FAST_GRAPH_LAPLACIAN,
            metadata={"degrees": degrees, "num_edges": len(edges)},
        )

    def fill_values_gpu(self, skeleton: TopologySkeleton, num_context_pairs: int) -> BatchMatrixData:
        n = skeleton.n
        rows_t = torch.from_numpy(skeleton.rows).to(self.device)
        cols_t = torch.from_numpy(skeleton.cols).to(self.device)
        degrees = torch.from_numpy(skeleton.metadata["degrees"]).to(self.device)

        num_edges_doubled = len(skeleton.rows) - n
        edge_vals = -torch.ones(num_edges_doubled, dtype=torch.float64, device=self.device)
        diag_vals = degrees.clone()
        all_vals = torch.cat([edge_vals, diag_vals])

        indices = torch.stack([rows_t, cols_t])
        scaled_vals, diagonals = scale_and_shift(indices, all_vals, n)

        context_pairs = make_context_pairs(indices, scaled_vals, n, num_context_pairs, self.device)

        return BatchMatrixData(
            indices=indices,
            values=scaled_vals.unsqueeze(0),
            diagonals=diagonals.unsqueeze(0),
            context_pairs=context_pairs.unsqueeze(0),
            n=n,
            batch_size=1,
            domain=MatrixDomain.FAST_GRAPH_LAPLACIAN,
        )

    def generate_batch(self, batch_size: int, num_context_pairs: int) -> BatchMatrixData:
        n = random.choice(self.ns)
        m = random.randint(2, 5)
        skeleton = self.generate_topology_cpu(n, m)
        return self.fill_values_gpu(skeleton, num_context_pairs)
