"""Rigid alignment for structure dataclasses."""

from __future__ import annotations

from dataclasses import Field, replace
from typing import Any, ClassVar, Protocol, TypeVar

import numpy as np
import torch
from torch import Tensor

from .esmfold2_protein_structure import compute_affine_and_rmsd


class Alignable(Protocol):
    """Minimum structure interface accepted by :class:`Aligner`."""

    __dataclass_fields__: ClassVar[dict[str, Field[Any]]]

    @property
    def atom37_positions(self) -> np.ndarray: ...

    @property
    def atom37_mask(self) -> np.ndarray: ...

    def __len__(self) -> int: ...


AlignableT = TypeVar("AlignableT", bound=Alignable)


def _coordinate_batch(structure: Alignable) -> Tensor:
    return torch.as_tensor(structure.atom37_positions, dtype=torch.double).unsqueeze(0)


def _shared_atom_mask(mobile: Alignable, target: Alignable, backbone_only: bool) -> Tensor:
    shared = np.asarray(mobile.atom37_mask, dtype=bool) & np.asarray(
        target.atom37_mask,
        dtype=bool,
    )
    if backbone_only:
        shared = shared.copy()
        shared[:, 3:] = False
    return torch.from_numpy(shared).unsqueeze(0)


class Aligner:
    """Fit a mobile structure onto a target with masked Kabsch alignment."""

    def __init__(
        self,
        mobile: Alignable,
        target: Alignable,
        only_use_backbone: bool = False,
        use_reflection: bool = False,
    ) -> None:
        if len(mobile) != len(target):
            raise AssertionError("mobile and target must contain the same residue count")

        mobile_coordinates = _coordinate_batch(mobile)
        target_coordinates = _coordinate_batch(target)
        if use_reflection:
            target_coordinates = -target_coordinates
        atom_mask = _shared_atom_mask(mobile, target, only_use_backbone)
        self._affine3D, rmsd = compute_affine_and_rmsd(
            mobile_coordinates,
            target_coordinates,
            atom_exists_mask=atom_mask,
        )
        self._rmsd = rmsd.item()

    @property
    def rmsd(self) -> float:
        return self._rmsd

    def apply(self, mobile: AlignableT) -> AlignableT:
        """Return a dataclass copy with all present atom coordinates aligned."""

        present = np.asarray(mobile.atom37_mask, dtype=bool)
        packed = torch.as_tensor(
            mobile.atom37_positions[present],
            dtype=torch.float32,
        ).unsqueeze(0)
        aligned = self._affine3D.apply(packed).squeeze(0).cpu().numpy()
        atom37_positions = np.full_like(mobile.atom37_positions, np.nan)
        atom37_positions[present] = aligned
        return replace(mobile, atom37_positions=atom37_positions)