import random
from pathlib import Path

import torch
from scipy.io import mmread

from ..base import BatchMatrixData, MatrixDomain, make_context_pairs, scale_and_shift


class SuiteSparseLoader:

    def __init__(self, data_dir: Path, device: torch.device):
        self.device = device
        self.matrices = self._load_all(data_dir)

    def _load_all(self, data_dir: Path) -> list[dict]:
        mtx_files = sorted(data_dir.rglob("*.mtx"))
        if not mtx_files:
            raise FileNotFoundError(f"No .mtx files found in {data_dir}")

        matrices = []
        for path in mtx_files:
            scipy_matrix = mmread(str(path)).tocoo()

            n = scipy_matrix.shape[0]
            if scipy_matrix.shape[0] != scipy_matrix.shape[1]:
                raise ValueError(f"Non-square matrix: {path}")

            indices = torch.tensor(
                [scipy_matrix.row, scipy_matrix.col],
                dtype=torch.long,
                device=self.device,
            )
            values = torch.tensor(
                scipy_matrix.data,
                dtype=torch.float64,
                device=self.device,
            )

            scaled_values, diagonals = scale_and_shift(indices, values, n)

            matrices.append({
                "indices": indices,
                "values": scaled_values,
                "diagonals": diagonals,
                "n": n,
                "name": path.parent.name + "/" + path.stem,
            })

        return matrices

    def generate_batch(self, batch_size: int, num_context_pairs: int) -> BatchMatrixData:
        entry = random.choice(self.matrices)

        context_pairs = make_context_pairs(
            entry["indices"], entry["values"], entry["n"], num_context_pairs, self.device
        )

        return BatchMatrixData(
            indices=entry["indices"],
            values=entry["values"].unsqueeze(0),
            diagonals=entry["diagonals"].unsqueeze(0),
            context_pairs=context_pairs.unsqueeze(0),
            n=entry["n"],
            batch_size=1,
            domain=MatrixDomain.SUITESPARSE,
        )
