import random

import torch

from ..base import (
    BatchMatrixData,
    GeneratorConfig,
    GridStencil,
    MatrixDomain,
    build_grid_stencil,
    make_context_pairs,
    scale_and_shift,
)


class VariableGridDiffusionGenerator:

    def __init__(self, config: GeneratorConfig, device: torch.device):
        self.config = config
        self.device = device
        self.stencils: dict[int, GridStencil] = {
            gs: build_grid_stencil(gs, device) for gs in config.grid_sizes
        }

    def _generate_kappa(self, stencil: GridStencil) -> torch.Tensor:
        n = stencil.n
        gs = stencil.grid_size
        kappa = torch.empty(n, dtype=torch.float64, device=self.device).uniform_(0.8, 1.2)

        if random.random() < self.config.material_jump_prob:
            jump_factor = random.uniform(*self.config.material_jump_range)
            split = random.randint(gs // 4, 3 * gs // 4)
            indices = torch.arange(n, device=self.device)
            if random.random() < 0.5:
                mask = (indices % gs) >= split
            else:
                mask = (indices // gs) >= split
            kappa[mask] *= jump_factor

        return kappa

    def generate_batch(self, batch_size: int, num_context_pairs: int) -> BatchMatrixData:
        grid_size = random.choice(self.config.grid_sizes)
        stencil = self.stencils[grid_size]

        kappa = self._generate_kappa(stencil)

        kappa_src = kappa[stencil.edge_src]
        kappa_dst = kappa[stencil.edge_dst]
        offdiag_values = -2.0 * kappa_src * kappa_dst / (kappa_src + kappa_dst)

        diag_values = torch.zeros(stencil.n, dtype=torch.float64, device=self.device)
        diag_values.scatter_add_(0, stencil.edge_src, -offdiag_values)

        all_values = torch.zeros(stencil.nnz, dtype=torch.float64, device=self.device)
        all_values[stencil.offdiag_mask] = offdiag_values
        all_values[stencil.diag_mask] = diag_values

        indices = torch.stack([stencil.rows, stencil.cols])
        scaled_values, diagonals = scale_and_shift(indices, all_values, stencil.n)
        context_pairs = make_context_pairs(indices, scaled_values, stencil.n, num_context_pairs, self.device)

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