import random

import numpy as np
import torch

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


class RandomSparseGenerator:

    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, density: float) -> TopologySkeleton:
        nnz_offdiag = int(n * n * density)

        if nnz_offdiag == 0:
            rows = np.empty(0, dtype=np.int64)
            cols = np.empty(0, dtype=np.int64)
        else:
            rows = np.random.randint(0, n, size=nnz_offdiag, dtype=np.int64)
            cols = np.random.randint(0, n, size=nnz_offdiag, dtype=np.int64)

            non_diag = rows != cols
            rows = rows[non_diag]
            cols = cols[non_diag]

            coords = np.unique(np.stack([rows, cols], axis=1), axis=0)
            rows = coords[:, 0]
            cols = coords[:, 1]

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

        return TopologySkeleton(
            rows=all_rows,
            cols=all_cols,
            n=n,
            domain=MatrixDomain.RANDOM_SPARSE,
            metadata={"num_offdiag": len(rows)},
        )

    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)
        num_offdiag = skeleton.metadata["num_offdiag"]

        log_scale = random.uniform(1.0, 6.0)

        if num_offdiag > 0:
            offdiag_vals = torch.randn(num_offdiag, dtype=torch.float64, device=self.device)
            offdiag_vals = offdiag_vals * (10.0 ** (torch.rand(num_offdiag, dtype=torch.float64, device=self.device) * log_scale))
        else:
            offdiag_vals = torch.empty(0, dtype=torch.float64, device=self.device)

        row_abs_sums = torch.zeros(n, dtype=torch.float64, device=self.device)
        if num_offdiag > 0:
            offdiag_rows = rows_t[:num_offdiag]
            row_abs_sums.scatter_add_(0, offdiag_rows, offdiag_vals.abs())

        dominance_factor = random.uniform(1.1, 5.0)
        diag_vals = row_abs_sums * dominance_factor + 10.0 ** random.uniform(0.0, log_scale)

        all_vals = torch.cat([offdiag_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.RANDOM_SPARSE,
        )

    def generate_batch(self, batch_size: int, num_context_pairs: int) -> BatchMatrixData:
        n = random.choice(self.ns)
        density = 10 ** random.uniform(-3, -1.5)
        skeleton = self.generate_topology_cpu(n, density)
        return self.fill_values_gpu(skeleton, num_context_pairs)
