import math
import random

import torch

from ..base import BatchMatrixData, GridStencil, MatrixDomain, batch_arnoldi_context, build_grid_stencil


class DiffusionGenerator:

    def __init__(self, grid_sizes: tuple[int, ...] | int, device: torch.device = torch.device('cuda')):
        if isinstance(grid_sizes, int):
            grid_sizes = (grid_sizes,)
        self.grid_sizes = grid_sizes
        self.device = device
        self.stencils: dict[int, GridStencil] = {
            gs: build_grid_stencil(gs, device) for gs in grid_sizes
        }

    def _generate_kappa_batch(self, batch_size: int, grid_size: int,
                               num_modes: int = 5,
                               kappa_range: tuple = (0.1, 10.0)) -> torch.Tensor:
        gs = grid_size

        x = torch.linspace(0, 1, gs, dtype=torch.float64, device=self.device)
        y = torch.linspace(0, 1, gs, dtype=torch.float64, device=self.device)
        X, Y = torch.meshgrid(x, y, indexing='ij')
        X = X.unsqueeze(0).expand(batch_size, -1, -1)
        Y = Y.unsqueeze(0).expand(batch_size, -1, -1)

        kappa = torch.zeros(batch_size, gs, gs, dtype=torch.float64, device=self.device)

        for _ in range(num_modes):
            kx = torch.randint(1, 4, (batch_size, 1, 1), device=self.device).double()
            ky = torch.randint(1, 4, (batch_size, 1, 1), device=self.device).double()
            amp = torch.rand(batch_size, 1, 1, dtype=torch.float64, device=self.device)
            phase = torch.rand(batch_size, 1, 1, dtype=torch.float64, device=self.device) * 2 * math.pi
            kappa += amp * torch.sin(kx * math.pi * X + phase) * torch.sin(ky * math.pi * Y + phase)

        kappa_min = kappa.view(batch_size, -1).min(dim=1, keepdim=True)[0].unsqueeze(-1)
        kappa = kappa - kappa_min
        kappa_max = kappa.view(batch_size, -1).max(dim=1, keepdim=True)[0].unsqueeze(-1)
        kappa = kappa / (kappa_max + 1e-8)
        kappa = kappa * (kappa_range[1] - kappa_range[0]) + kappa_range[0]

        return kappa.view(batch_size, -1)

    def _build_matrix(self, batch_size: int, grid_size: int, kappa: torch.Tensor,
                      advection_offdiag: torch.Tensor | None = None) -> tuple[torch.Tensor, GridStencil]:
        stencil = self.stencils[grid_size]
        h = 1.0 / (grid_size + 1)
        h2 = h * h
        n = stencil.n

        kappa_src = kappa[:, stencil.edge_src]
        kappa_dst = kappa[:, stencil.edge_dst]
        kappa_avg = 0.5 * (kappa_src + kappa_dst)

        diffusion_offdiag = -kappa_avg / h2

        if advection_offdiag is not None:
            offdiag_values = diffusion_offdiag + advection_offdiag
        else:
            offdiag_values = diffusion_offdiag

        diag_contributions = torch.zeros(batch_size, n, dtype=torch.float64, device=self.device)
        diag_contributions.scatter_add_(1, stencil.edge_src.unsqueeze(0).expand(batch_size, -1), -diffusion_offdiag)

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

        row_sums = torch.zeros(batch_size, n, dtype=torch.float64, device=self.device)
        row_sums.scatter_add_(1, stencil.rows.unsqueeze(0).expand(batch_size, -1), all_values.abs())
        gamma = row_sums.max(dim=1, keepdim=True)[0]

        scaled_values = all_values / gamma
        return scaled_values, stencil

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

        kappa = self._generate_kappa_batch(batch_size, gs)
        scaled_values, _ = self._build_matrix(batch_size, gs, kappa)

        scaled_diag = scaled_values[:, stencil.diag_mask]

        context_pairs = batch_arnoldi_context(
            stencil.rows, stencil.cols, scaled_values, n, batch_size, num_context_pairs, self.device
        )

        return BatchMatrixData(
            indices=torch.stack([stencil.rows, stencil.cols]),
            values=scaled_values,
            diagonals=scaled_diag,
            context_pairs=context_pairs,
            n=n,
            batch_size=batch_size,
            domain=MatrixDomain.DIFFUSION,
        )

    def generate_advection_batch(self, batch_size: int, num_context_pairs: int = 10,
                                  advection_range: tuple = (0.1, 1.0)) -> BatchMatrixData:
        gs = random.choice(self.grid_sizes)
        stencil = self.stencils[gs]
        n = stencil.n
        h = 1.0 / (gs + 1)

        kappa = self._generate_kappa_batch(batch_size, gs)

        adv_strength = torch.rand(batch_size, 1, dtype=torch.float64, device=self.device)
        adv_strength = adv_strength * (advection_range[1] - advection_range[0]) + advection_range[0]

        vx = (torch.rand(batch_size, 1, dtype=torch.float64, device=self.device) - 0.5) * 2
        vy = (torch.rand(batch_size, 1, dtype=torch.float64, device=self.device) - 0.5) * 2
        norm = torch.sqrt(vx**2 + vy**2) + 1e-8
        vx = vx / norm * adv_strength
        vy = vy / norm * adv_strength

        h2_inv = 1.0 / (2.0 * h)
        advection_offdiag = (vx * stencil.advection_y_signs + vy * (-stencil.advection_x_signs)) * h2_inv

        scaled_values, _ = self._build_matrix(batch_size, gs, kappa, advection_offdiag)
        scaled_diag = scaled_values[:, stencil.diag_mask]

        context_pairs = batch_arnoldi_context(
            stencil.rows, stencil.cols, scaled_values, n, batch_size, num_context_pairs, self.device
        )

        return BatchMatrixData(
            indices=torch.stack([stencil.rows, stencil.cols]),
            values=scaled_values,
            diagonals=scaled_diag,
            context_pairs=context_pairs,
            n=n,
            batch_size=batch_size,
            domain=MatrixDomain.DIFFUSION_ADVECTION,
        )
