import random
from types import SimpleNamespace

import torch

from ..base import BatchMatrixData, MatrixDomain, batch_arnoldi_context


class StokesSaddlePointGenerator:

    def __init__(self, grid_sizes: tuple[int, ...] | int, device: torch.device):
        if isinstance(grid_sizes, int):
            grid_sizes = (grid_sizes,)
        self.grid_sizes = grid_sizes
        self.device = device
        self._data: dict[int, SimpleNamespace] = {
            gs: self._precompute(gs) for gs in grid_sizes
        }

    def _precompute(self, grid_size: int) -> SimpleNamespace:
        gs = grid_size
        h = 1.0 / (gs + 1)
        h2 = h * h
        inv_h = 1.0 / h
        neg_inv_h2 = -1.0 / h2

        n_u = (gs + 1) * gs
        n_v = gs * (gs + 1)
        n_vel = n_u + n_v
        n_p = gs * gs
        n = n_vel + n_p

        rows, cols, base_vals = [], [], []

        for idx in range(n_vel):
            rows.append(idx)
            cols.append(idx)
            base_vals.append(4.0 / h2)

        u_stride = gs
        for i in range(gs + 1):
            for j in range(gs):
                idx = i * u_stride + j
                if i > 0:
                    rows.append(idx); cols.append(idx - u_stride)
                    base_vals.append(neg_inv_h2)
                if i < gs:
                    rows.append(idx); cols.append(idx + u_stride)
                    base_vals.append(neg_inv_h2)
                if j > 0:
                    rows.append(idx); cols.append(idx - 1)
                    base_vals.append(neg_inv_h2)
                if j < gs - 1:
                    rows.append(idx); cols.append(idx + 1)
                    base_vals.append(neg_inv_h2)

        v_stride = gs + 1
        for i in range(gs):
            for j in range(gs + 1):
                idx = n_u + i * v_stride + j
                if i > 0:
                    rows.append(idx); cols.append(idx - v_stride)
                    base_vals.append(neg_inv_h2)
                if i < gs - 1:
                    rows.append(idx); cols.append(idx + v_stride)
                    base_vals.append(neg_inv_h2)
                if j > 0:
                    rows.append(idx); cols.append(idx - 1)
                    base_vals.append(neg_inv_h2)
                if j < gs:
                    rows.append(idx); cols.append(idx + 1)
                    base_vals.append(neg_inv_h2)

        vel_nnz = len(rows)

        div_rows, div_cols, div_vals = [], [], []
        for i in range(gs):
            for j in range(gs):
                p_idx = i * gs + j

                u_left = i * u_stride + j
                u_right = (i + 1) * u_stride + j
                div_rows.extend([p_idx, p_idx])
                div_cols.extend([u_left, u_right])
                div_vals.extend([-inv_h, inv_h])

                v_bottom = n_u + i * v_stride + j
                v_top = n_u + i * v_stride + (j + 1)
                div_rows.extend([p_idx, p_idx])
                div_cols.extend([v_bottom, v_top])
                div_vals.extend([-inv_h, inv_h])

        rows.extend(p + n_vel for p in div_rows)
        cols.extend(div_cols)
        base_vals.extend(div_vals)

        rows.extend(div_cols)
        cols.extend(p + n_vel for p in div_rows)
        base_vals.extend(div_vals)

        for p_idx in range(n_p):
            rows.append(n_vel + p_idx)
            cols.append(n_vel + p_idx)
            base_vals.append(1e-4)

        def _t(lst, dtype=torch.long):
            return torch.tensor(lst, dtype=dtype, device=self.device)

        base_rows = _t(rows)
        base_cols = _t(cols)
        base_vals_t = _t(base_vals, dtype=torch.float64)
        nnz = len(rows)

        viscosity_mask = torch.zeros(nnz, dtype=torch.bool, device=self.device)
        viscosity_mask[:vel_nnz] = True
        pressure_reg_mask = torch.zeros(nnz, dtype=torch.bool, device=self.device)
        pressure_reg_mask[-n_p:] = True
        diag_mask = base_rows == base_cols

        return SimpleNamespace(
            base_rows=base_rows, base_cols=base_cols, base_vals=base_vals_t,
            nnz=nnz, n=n, n_p=n_p,
            viscosity_mask=viscosity_mask, pressure_reg_mask=pressure_reg_mask,
            diag_mask=diag_mask,
        )

    def generate_batch(self, batch_size: int, num_context_pairs: int) -> BatchMatrixData:
        gs = random.choice(self.grid_sizes)
        d = self._data[gs]

        viscosity = torch.rand(batch_size, 1, dtype=torch.float64, device=self.device) * 0.09 + 0.01

        values = d.base_vals.unsqueeze(0).expand(batch_size, -1).clone()
        values[:, d.viscosity_mask] *= viscosity
        values[:, d.pressure_reg_mask] *= viscosity

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

        scaled_values = values / gamma
        scaled_diag = scaled_values[:, d.diag_mask]

        context_pairs = batch_arnoldi_context(
            d.base_rows, d.base_cols, scaled_values, d.n, batch_size, num_context_pairs, self.device
        )

        return BatchMatrixData(
            indices=torch.stack([d.base_rows, d.base_cols]),
            values=scaled_values,
            diagonals=scaled_diag,
            context_pairs=context_pairs,
            n=d.n,
            batch_size=batch_size,
            domain=MatrixDomain.STOKES,
        )
