"""BABAPPA representation construction."""

from __future__ import annotations

from dataclasses import asdict, dataclass

import numpy as np

from babappa.features.codon import codon_to_embedding
from babappa.io.fasta import CodonAlignment


@dataclass
class RepresentationSpec:
    embedding_version: str = "v1_bounded8"
    embed_dim: int = 8
    projection_dim: int = 16
    projection_seed: int = 0

    def to_dict(self) -> dict[str, int | str]:
        return asdict(self)


def projection_matrix(taxa_count: int, spec: RepresentationSpec) -> np.ndarray:
    width = taxa_count * spec.embed_dim
    if spec.projection_dim > width:
        raise ValueError(
            f"projection_dim={spec.projection_dim} exceeds stacked width={width}; reduce projection_dim"
        )
    seed = spec.projection_seed + taxa_count * 1009 + spec.projection_dim * 17
    rng = np.random.default_rng(seed)
    raw = rng.normal(size=(width, spec.projection_dim))
    q, _ = np.linalg.qr(raw)
    return q.T


def alignment_to_codon_matrix(alignment: CodonAlignment, taxa_order: list[str]) -> np.ndarray:
    lengths = {len(alignment.sequences[t]) for t in taxa_order}
    if len(lengths) != 1:
        raise ValueError(f"Inconsistent sequence lengths in gene {alignment.gene}")
    nt_len = next(iter(lengths))
    if nt_len % 3 != 0:
        raise ValueError(f"Sequence length is not a multiple of 3 for gene {alignment.gene}")

    codon_len = nt_len // 3
    mat = np.empty((len(taxa_order), codon_len), dtype=object)
    for i, taxon in enumerate(taxa_order):
        seq = alignment.sequences[taxon]
        mat[i, :] = [seq[j : j + 3].upper() for j in range(0, nt_len, 3)]
    return mat


def site_representation(site_codons: list[str], b_matrix: np.ndarray, embed_dim: int = 8) -> np.ndarray:
    taxa_count = len(site_codons)
    stacked = np.zeros(taxa_count * embed_dim, dtype=float)
    for i, codon in enumerate(site_codons):
        stacked[i * embed_dim : (i + 1) * embed_dim] = codon_to_embedding(codon)
    scaled = stacked / np.sqrt(float(taxa_count))
    return b_matrix @ scaled


def gene_representation(
    alignment: CodonAlignment,
    taxa_order: list[str],
    spec: RepresentationSpec,
) -> tuple[np.ndarray, np.ndarray]:
    codon_mat = alignment_to_codon_matrix(alignment, taxa_order)
    b = projection_matrix(len(taxa_order), spec)

    site_vectors = np.zeros((codon_mat.shape[1], spec.projection_dim), dtype=float)
    for s in range(codon_mat.shape[1]):
        site = [str(x) for x in codon_mat[:, s].tolist()]
        site_vectors[s, :] = site_representation(site, b_matrix=b, embed_dim=spec.embed_dim)
    return site_vectors, codon_mat
