import random

import torch

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


class VariableAdvectionGenerator:

    def __init__(self, config: GeneratorConfig, device: torch.device):
        self.config = config
        self.device = device
        self.base = VariableGridDiffusionGenerator(config, device)

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

        kappa = self.base._generate_kappa(stencil)
        advection_strength = random.uniform(0.3, 2.0)

        kappa_src = kappa[stencil.edge_src]
        kappa_dst = kappa[stencil.edge_dst]
        kappa_harmonic = 2.0 * kappa_src * kappa_dst / (kappa_src + kappa_dst)
        offdiag_values = -kappa_harmonic
        offdiag_values += stencil.advection_x_signs * advection_strength * 0.5 / grid_size

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

        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_ADVECTION,
        )
