import random

import torch

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


class EnhancedDiffusionGenerator:

    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_holes(self, n: int, grid_size: int) -> torch.Tensor:
        hole_mask = torch.ones(n, dtype=torch.bool, device=self.device)
        if random.random() < 0.2:
            num_holes = random.randint(1, 3)
            for _ in range(num_holes):
                hole_size = random.randint(2, max(3, grid_size // 8))
                hole_i = random.randint(1, grid_size - hole_size - 1)
                hole_j = random.randint(1, grid_size - hole_size - 1)
                row_indices = torch.arange(hole_i, hole_i + hole_size, device=self.device)
                col_indices = torch.arange(hole_j, hole_j + hole_size, device=self.device)
                flat_indices = (row_indices.unsqueeze(1) * grid_size + col_indices).flatten()
                hole_mask[flat_indices] = False
        return hole_mask

    def _assemble_values(self, grid_size: int) -> tuple[torch.Tensor, torch.Tensor, GridStencil, torch.Tensor]:
        stencil = self.stencils[grid_size]
        n = stencil.n

        anisotropy = random.uniform(0.1, 10.0)
        kappa_x = torch.empty(n, dtype=torch.float64, device=self.device).uniform_(0.5, 2.0)
        kappa_y = kappa_x / anisotropy

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

        hole_mask = self._generate_holes(n, grid_size)

        kappa_combined = torch.zeros(stencil.num_edges, dtype=torch.float64, device=self.device)
        src_kx = kappa_x[stencil.edge_src]
        dst_kx = kappa_x[stencil.edge_dst]
        src_ky = kappa_y[stencil.edge_src]
        dst_ky = kappa_y[stencil.edge_dst]
        x_mask = stencil.x_edge_mask
        y_mask = stencil.y_edge_mask
        kappa_combined[x_mask] = 2.0 * src_kx[x_mask] * dst_kx[x_mask] / (src_kx[x_mask] + dst_kx[x_mask])
        kappa_combined[y_mask] = 2.0 * src_ky[y_mask] * dst_ky[y_mask] / (src_ky[y_mask] + dst_ky[y_mask])

        src_active = hole_mask[stencil.edge_src]
        dst_active = hole_mask[stencil.edge_dst]
        edge_active = src_active & dst_active
        kappa_combined[~edge_active] = 0.0

        offdiag_values = -kappa_combined

        diag_values = torch.zeros(n, dtype=torch.float64, device=self.device)
        diag_values.scatter_add_(0, stencil.edge_src, kappa_combined)
        diag_values[~hole_mask] = 1.0
        diag_values[hole_mask] = diag_values[hole_mask].clamp(min=1.0)

        return offdiag_values, diag_values, stencil, hole_mask

    def _finalize(self, offdiag_values: torch.Tensor, diag_values: torch.Tensor,
                  stencil: GridStencil, domain: MatrixDomain, num_context_pairs: int) -> BatchMatrixData:
        n = stencil.n
        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])

        if random.random() < 0.3:
            perm = torch.randperm(n, device=self.device)
            inv_perm = torch.argsort(perm)
            indices = torch.stack([inv_perm[indices[0]], inv_perm[indices[1]]])

        scaled_values, diagonals = scale_and_shift(indices, all_values, n)
        context_pairs = make_context_pairs(indices, scaled_values, 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=n,
            batch_size=1,
            domain=domain,
        )

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