"""Loading precomputed LLR matrices written by scripts/compute_llr_matrices.py."""

from __future__ import annotations

from pathlib import Path

import numpy as np
import pandas as pd

from src.config import MATRICES_DIR
from src.constants import AMINO_ACIDS


def list_available_accessions(directory: Path = MATRICES_DIR) -> list[str]:
    """Accessions with a precomputed `<accession>.npz` in `directory`, sorted."""
    if not directory.exists():
        return []
    return sorted(path.stem for path in directory.glob("*.npz"))


def load_precomputed_matrix(accession: str, directory: Path = MATRICES_DIR) -> tuple[str, pd.DataFrame]:
    """Load `<directory>/<accession>.npz`.

    Returns `(sequence, matrix)`, where `matrix` has the same shape and row
    (`AMINO_ACIDS`) / column (`{residue}{position}`) labeling as
    `src.model.calculate_llr_matrix`'s output.
    """
    data = np.load(directory / f"{accession}.npz")
    sequence = str(data["sequence"])
    columns = [f"{residue}{i + 1}" for i, residue in enumerate(sequence)]
    matrix = pd.DataFrame(data["matrix"], index=AMINO_ACIDS, columns=columns)
    return sequence, matrix