import random

import torch

from ..base import BatchMatrixData, GeneratorConfig, MatrixDomain
from .enhanced_diffusion import EnhancedDiffusionGenerator


class EnhancedAdvectionGenerator:

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

    def generate_batch(self, batch_size: int, num_context_pairs: int) -> BatchMatrixData:
        grid_size = random.choice(self.base.config.grid_sizes)
        offdiag_values, diag_values, stencil, _ = self.base._assemble_values(grid_size)

        advection_strength = random.uniform(0.1, 3.0)
        offdiag_values += stencil.advection_x_signs * advection_strength * 0.5 / grid_size

        return self.base._finalize(offdiag_values, diag_values, stencil,
                                   MatrixDomain.ENHANCED_ADVECTION, num_context_pairs)
