import torch
import numpy as np
from dataclasses import dataclass, field
from enum import Enum
from typing import Protocol


class MatrixDomain(Enum):
    DIFFUSION = "diffusion"
    DIFFUSION_ADVECTION = "diffusion_advection"
    GRAPH_LAPLACIAN = "graph_laplacian"
    ELASTICITY = "elasticity"
    STOKES = "stokes"
    SBM = "sbm"
    SPECTRAL_STRESS = "spectral_stress"
    VARIABLE_DIFFUSION = "variable_diffusion"
    VARIABLE_ADVECTION = "variable_advection"
    ENHANCED_DIFFUSION = "enhanced_diffusion"
    ENHANCED_ADVECTION = "enhanced_advection"
    FAST_GRAPH_LAPLACIAN = "fast_graph_laplacian"
    RANDOM_SPARSE = "random_sparse"
    DIRECTED_POWER_LAW = "directed_power_law"
    SUITESPARSE = "suitesparse"


@dataclass
class BatchMatrixData:
    indices: torch.Tensor
    values: torch.Tensor
    diagonals: torch.Tensor
    context_pairs: torch.Tensor
    n: int
    batch_size: int
    domain: MatrixDomain


@dataclass
class TopologySkeleton:
    rows: np.ndarray
    cols: np.ndarray
    n: int
    domain: MatrixDomain
    metadata: dict = field(default_factory=dict)


@dataclass
class GeneratorConfig:
    grid_sizes: tuple[int, ...] = (16, 24, 32, 48, 64, 96, 128)
    grid_size: int = 32
    n: int = 32 * 32
    material_jump_prob: float = 0.3
    material_jump_range: tuple[float, float] = (10.0, 1000.0)


@dataclass
class GridStencil:
    rows: torch.Tensor
    cols: torch.Tensor
    edge_src: torch.Tensor
    edge_dst: torch.Tensor
    offdiag_mask: torch.Tensor
    diag_mask: torch.Tensor
    x_edge_mask: torch.Tensor
    y_edge_mask: torch.Tensor
    advection_x_signs: torch.Tensor
    advection_y_signs: torch.Tensor
    n: int
    nnz: int
    num_edges: int
    grid_size: int


def build_grid_stencil(grid_size: int, device: torch.device) -> GridStencil:
    gs = grid_size
    n = gs * gs

    rows, cols = [], []
    edge_directions = []

    for i in range(gs):
        for j in range(gs):
            idx = i * gs + j

            if i > 0:
                rows.append(idx); cols.append(idx - gs)
                edge_directions.append('up')
            if i < gs - 1:
                rows.append(idx); cols.append(idx + gs)
                edge_directions.append('down')
            if j > 0:
                rows.append(idx); cols.append(idx - 1)
                edge_directions.append('left')
            if j < gs - 1:
                rows.append(idx); cols.append(idx + 1)
                edge_directions.append('right')

            rows.append(idx); cols.append(idx)
            edge_directions.append('diag')

    rows_t = torch.tensor(rows, dtype=torch.long, device=device)
    cols_t = torch.tensor(cols, dtype=torch.long, device=device)
    nnz = len(rows)

    offdiag_mask = rows_t != cols_t
    diag_mask = ~offdiag_mask

    offdiag_dirs = [d for d in edge_directions if d != 'diag']
    offdiag_rows = rows_t[offdiag_mask]
    offdiag_cols = cols_t[offdiag_mask]
    num_edges = len(offdiag_dirs)

    x_edge = torch.tensor([d in ('left', 'right') for d in offdiag_dirs], dtype=torch.bool, device=device)
    y_edge = ~x_edge

    adv_x = torch.zeros(num_edges, dtype=torch.float64, device=device)
    adv_y = torch.zeros(num_edges, dtype=torch.float64, device=device)
    for idx, d in enumerate(offdiag_dirs):
        if d == 'right':
            adv_x[idx] = 1.0
        elif d == 'left':
            adv_x[idx] = -1.0
        elif d == 'up':
            adv_y[idx] = 1.0
        elif d == 'down':
            adv_y[idx] = -1.0

    return GridStencil(
        rows=rows_t,
        cols=cols_t,
        edge_src=offdiag_rows,
        edge_dst=offdiag_cols,
        offdiag_mask=offdiag_mask,
        diag_mask=diag_mask,
        x_edge_mask=x_edge,
        y_edge_mask=y_edge,
        advection_x_signs=adv_x,
        advection_y_signs=adv_y,
        n=n,
        nnz=nnz,
        num_edges=num_edges,
        grid_size=grid_size,
    )


class MatrixGenerator(Protocol):
    def generate_batch(self, batch_size: int, num_context_pairs: int) -> BatchMatrixData: ...


def batch_spmv(
    base_rows: torch.Tensor,
    base_cols: torch.Tensor,
    values: torch.Tensor,
    x: torch.Tensor,
    n: int,
) -> torch.Tensor:
    batch_size = values.shape[0]
    num_vecs = x.shape[2]
    x_gathered = x[:, base_cols, :]
    products = values.unsqueeze(-1) * x_gathered
    result = torch.zeros(batch_size, n, num_vecs, dtype=torch.float64, device=values.device)
    rows_expanded = base_rows.unsqueeze(0).unsqueeze(-1).expand(batch_size, -1, num_vecs)
    result.scatter_add_(1, rows_expanded, products)
    return result


def scale_and_shift(
    indices: torch.Tensor,
    values: torch.Tensor,
    n: int,
    diag_shift: float = 0.1,
) -> tuple[torch.Tensor, torch.Tensor]:
    diag_mask = indices[0] == indices[1]

    if values.dim() == 1:
        row_sums = torch.zeros(n, dtype=torch.float64, device=values.device)
        row_sums.scatter_add_(0, indices[0], values.abs())
        gamma = row_sums.max()
        scaled = values / gamma
        scaled[diag_mask] = scaled[diag_mask] + diag_shift
        diagonals = torch.zeros(n, dtype=torch.float64, device=values.device)
        diagonals.scatter_(0, indices[1][diag_mask], scaled[diag_mask])
        return scaled, diagonals
    else:
        batch_size = values.shape[0]
        row_sums = torch.zeros(batch_size, n, dtype=torch.float64, device=values.device)
        row_sums.scatter_add_(1, indices[0].unsqueeze(0).expand(batch_size, -1), values.abs())
        gamma = row_sums.max(dim=1, keepdim=True)[0]
        scaled = values / gamma
        scaled[:, diag_mask] = scaled[:, diag_mask] + diag_shift
        diag_indices = indices[1][diag_mask]
        diagonals = torch.zeros(batch_size, n, dtype=torch.float64, device=values.device)
        diagonals.scatter_(1, diag_indices.unsqueeze(0).expand(batch_size, -1), scaled[:, diag_mask])
        return scaled, diagonals


def make_context_pairs(
    indices: torch.Tensor,
    values: torch.Tensor,
    n: int,
    num_context_pairs: int,
    device: torch.device,
) -> torch.Tensor:
    A = torch.sparse_coo_tensor(indices, values, (n, n)).coalesce()
    k = num_context_pairs

    V = torch.zeros(n, k + 1, dtype=torch.float64, device=device)
    Av = torch.zeros(n, k, dtype=torch.float64, device=device)

    v0 = torch.randn(n, dtype=torch.float64, device=device)
    V[:, 0] = v0 / torch.linalg.vector_norm(v0)

    for j in range(k):
        w = A @ V[:, j]
        Av[:, j] = w

        for i in range(j + 1):
            w = w - torch.dot(V[:, i], w) * V[:, i]

        h = torch.linalg.vector_norm(w)
        if h < 1e-14:
            V[:, j + 1:] = 0
            Av[:, j + 1:] = 0
            break
        V[:, j + 1] = w / h

    return torch.stack([V[:, :k], Av], dim=-1)


def batch_arnoldi_context(
    base_rows: torch.Tensor,
    base_cols: torch.Tensor,
    values: torch.Tensor,
    n: int,
    batch_size: int,
    num_context_pairs: int,
    device: torch.device,
) -> torch.Tensor:
    k = num_context_pairs

    V = torch.zeros(batch_size, n, k + 1, dtype=torch.float64, device=device)
    Av = torch.zeros(batch_size, n, k, dtype=torch.float64, device=device)

    v0 = torch.randn(batch_size, n, dtype=torch.float64, device=device)
    V[:, :, 0] = v0 / torch.linalg.vector_norm(v0, dim=1, keepdim=True)

    for j in range(k):
        w = batch_spmv(base_rows, base_cols, values, V[:, :, j:j + 1], n).squeeze(-1)
        Av[:, :, j] = w

        for i in range(j + 1):
            coeffs = (V[:, :, i] * w).sum(dim=1, keepdim=True)
            w = w - coeffs * V[:, :, i]

        h = torch.linalg.vector_norm(w, dim=1, keepdim=True)
        V[:, :, j + 1] = torch.where(h > 1e-14, w / h, torch.zeros_like(w))

    return torch.stack([V[:, :, :k], Av], dim=-1)
