import random

import torch

from ..base import BatchMatrixData, MatrixDomain, make_context_pairs


class SpectralStressGenerator:

    def __init__(self, base_generator, device: torch.device):
        self.base_generator = base_generator
        self.device = device

    def generate_batch(self, batch_size: int, num_context_pairs: int) -> BatchMatrixData:
        assert batch_size == 1, "SpectralStressGenerator only supports batch_size=1"

        base_data = self.base_generator.generate_batch(1, num_context_pairs)

        indices = base_data.indices
        values = base_data.values[0].clone()
        n = base_data.n

        target_kappa = 10 ** random.uniform(2.0, 4.0)

        diag_mask = indices[0] == indices[1]
        offdiag_mask = ~diag_mask

        scale_factor = target_kappa ** 0.3
        values[offdiag_mask] /= scale_factor

        row_sums = torch.zeros(n, dtype=torch.float64, device=values.device)
        row_sums.scatter_add_(0, indices[0][offdiag_mask], values[offdiag_mask].abs())

        shift = row_sums.max().item() / target_kappa
        values[diag_mask] = row_sums[indices[0][diag_mask]] + shift

        diagonals = torch.zeros(n, dtype=torch.float64, device=values.device)
        diagonals[indices[0][diag_mask]] = values[diag_mask]

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

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