import random
from types import SimpleNamespace

import torch

from ..base import BatchMatrixData, MatrixDomain, batch_arnoldi_context


class ElasticityGenerator:

    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._data: dict[int, SimpleNamespace] = {
            gs: self._precompute(gs) for gs in grid_sizes
        }

    def _precompute(self, grid_size: int) -> SimpleNamespace:
        gs = grid_size
        n_nodes = gs * gs
        n = 2 * n_nodes
        h = 1.0 / (gs + 1)

        rows, cols = [], []
        groups = {name: {'idx': [], 'nodes': [], 'nbrs': []}
                  for name in ('ux_ux_strong', 'ux_ux_weak',
                               'uy_uy_strong', 'uy_uy_weak',
                               'coupling_neg', 'coupling_pos')}
        diag_ux_idx, diag_ux_nodes = [], []
        diag_uy_idx, diag_uy_nodes = [], []

        entry_idx = 0

        for i in range(gs):
            for j in range(gs):
                node = i * gs + j
                ux = 2 * node
                uy = 2 * node + 1

                axis_nbrs = []
                if i > 0:
                    axis_nbrs.append(('y', node - gs))
                if i < gs - 1:
                    axis_nbrs.append(('y', node + gs))
                if j > 0:
                    axis_nbrs.append(('x', node - 1))
                if j < gs - 1:
                    axis_nbrs.append(('x', node + 1))

                for axis, nbr in axis_nbrs:
                    rows.append(ux); cols.append(2 * nbr)
                    g = groups['ux_ux_strong' if axis == 'x' else 'ux_ux_weak']
                    g['idx'].append(entry_idx); g['nodes'].append(node); g['nbrs'].append(nbr)
                    entry_idx += 1

                    rows.append(uy); cols.append(2 * nbr + 1)
                    g = groups['uy_uy_strong' if axis == 'y' else 'uy_uy_weak']
                    g['idx'].append(entry_idx); g['nodes'].append(node); g['nbrs'].append(nbr)
                    entry_idx += 1

                corners = []
                if i > 0 and j > 0:
                    corners.append(('neg', node - gs - 1))
                if i > 0 and j < gs - 1:
                    corners.append(('pos', node - gs + 1))
                if i < gs - 1 and j > 0:
                    corners.append(('pos', node + gs - 1))
                if i < gs - 1 and j < gs - 1:
                    corners.append(('neg', node + gs + 1))

                for sign, corner in corners:
                    g = groups[f'coupling_{sign}']
                    rows.append(ux); cols.append(2 * corner + 1)
                    g['idx'].append(entry_idx); g['nodes'].append(node); g['nbrs'].append(corner)
                    entry_idx += 1
                    rows.append(uy); cols.append(2 * corner)
                    g['idx'].append(entry_idx); g['nodes'].append(node); g['nbrs'].append(corner)
                    entry_idx += 1

                rows.append(ux); cols.append(ux)
                diag_ux_idx.append(entry_idx); diag_ux_nodes.append(node)
                entry_idx += 1
                rows.append(uy); cols.append(uy)
                diag_uy_idx.append(entry_idx); diag_uy_nodes.append(node)
                entry_idx += 1

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

        d = SimpleNamespace(
            base_rows=_t(rows), base_cols=_t(cols), nnz=entry_idx,
            n_nodes=n_nodes, n=n, h2=h * h,
            diag_ux_idx=_t(diag_ux_idx), diag_ux_nodes=_t(diag_ux_nodes),
            diag_uy_idx=_t(diag_uy_idx), diag_uy_nodes=_t(diag_uy_nodes),
        )

        for name, g in groups.items():
            setattr(d, f'{name}_idx', _t(g['idx']))
            setattr(d, f'{name}_nodes', _t(g['nodes']))
            setattr(d, f'{name}_nbrs', _t(g['nbrs']))

        d.diag_mask = torch.zeros(entry_idx, dtype=torch.bool, device=self.device)
        d.diag_mask[d.diag_ux_idx] = True
        d.diag_mask[d.diag_uy_idx] = True

        return d

    def generate_batch(self, batch_size: int, num_context_pairs: int = 10,
                       E_range: tuple = (1e6, 1e9), nu: float = 0.3) -> BatchMatrixData:
        gs = random.choice(self.grid_sizes)
        d = self._data[gs]

        E = torch.rand(batch_size, d.n_nodes, dtype=torch.float64, device=self.device)
        E = E * (E_range[1] - E_range[0]) + E_range[0]

        c1 = 1.0 / (1.0 - nu * nu)
        c2 = (1.0 - nu) / 2.0
        c3 = (1.0 + nu) / 2.0
        c1_h2 = c1 / d.h2

        all_values = torch.zeros(batch_size, d.nnz, dtype=torch.float64, device=self.device)

        diag_coeff = 2.0 * (1.0 + c2) * c1_h2
        all_values[:, d.diag_ux_idx] = diag_coeff * E[:, d.diag_ux_nodes]
        all_values[:, d.diag_uy_idx] = diag_coeff * E[:, d.diag_uy_nodes]

        def _e_avg(nodes, nbrs):
            return 0.5 * (E[:, nodes] + E[:, nbrs]) * c1_h2

        all_values[:, d.ux_ux_strong_idx] = -_e_avg(d.ux_ux_strong_nodes, d.ux_ux_strong_nbrs)
        all_values[:, d.ux_ux_weak_idx] = -_e_avg(d.ux_ux_weak_nodes, d.ux_ux_weak_nbrs) * c2
        all_values[:, d.uy_uy_strong_idx] = -_e_avg(d.uy_uy_strong_nodes, d.uy_uy_strong_nbrs)
        all_values[:, d.uy_uy_weak_idx] = -_e_avg(d.uy_uy_weak_nodes, d.uy_uy_weak_nbrs) * c2

        all_values[:, d.coupling_neg_idx] = -_e_avg(d.coupling_neg_nodes, d.coupling_neg_nbrs) * c3 * 0.25
        all_values[:, d.coupling_pos_idx] = _e_avg(d.coupling_pos_nodes, d.coupling_pos_nbrs) * c3 * 0.25

        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), all_values.abs())
        gamma = row_sums.max(dim=1, keepdim=True)[0]

        scaled_values = all_values / gamma
        scaled_values[:, d.diag_mask] += 0.1
        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.ELASTICITY,
        )
