PAF I/O¶

@dataclass
class PafRecord:
    """A single PAF alignment record.

    The twelve required PAF columns are represented as typed attributes.
    Optional SAM-like tags (e.g. ``tp:A:P``, ``cg:Z:10M``) are stored in
    :attr:`tags`.  If a ``cg:Z:`` tag is present, CIGAR-derived alignment
    statistics are populated automatically.

    Parameters
    ----------
    query_name : str
        Query sequence name (column 1).
    query_len : int
        Query sequence length (column 2).
    query_start : int
        Query start position, 0-based (column 3).
    query_end : int
        Query end position, exclusive (column 4).
    strand : str
        Relative strand: ``"+"`` or ``"-"`` (column 5).
    target_name : str
        Target sequence name (column 6).
    target_len : int
        Target sequence length (column 7).
    target_start : int
        Target start position, 0-based (column 8).
    target_end : int
        Target end position, exclusive (column 9).
    residue_matches : int
        Number of residue matches (column 10).
    alignment_block_len : int
        Number of bases in the alignment block (column 11).
    mapping_quality : int
        Mapping quality (0–255; 255 = missing) (column 12).
    tags : dict[str, Any]
        Optional SAM-like tags decoded as ``{tag_name: value}``.
    cigar : str or None
        CIGAR string from ``cg:Z:`` tag, or ``None`` if absent.
    alignment_length : int or None
        Target-span alignment length derived from CIGAR, or ``None``.
    n_matches : int or None
        Count of exact-match bases (``=`` ops) from CIGAR; falls back to
        ``residue_matches`` when only ``M`` ops are present.
    n_mismatches : int or None
        Count of mismatch bases (``X`` ops) from CIGAR, or ``None``.
    n_gaps : int or None
        Total number of gap bases (``I`` + ``D`` bases) from CIGAR.
    n_gap_bases : int or None
        Same as ``n_gaps`` (alias kept for clarity).
    tag_types : dict[str, str]
        SAM tag type characters (e.g. ``{'tp': 'A'}``) recorded at parse
        time so :meth:`to_line` can round-trip tags exactly.
    """

    query_name: str
    query_len: int
    query_start: int
    query_end: int
    strand: str
    target_name: str
    target_len: int
    target_start: int
    target_end: int
    residue_matches: int
    alignment_block_len: int
    mapping_quality: int
    tags: dict[str, Any] = field(default_factory=dict)
    cigar: str | None = None
    alignment_length: int | None = None
    n_matches: int | None = None
    n_mismatches: int | None = None
    n_gaps: int | None = None
    n_gap_bases: int | None = None
    tag_types: dict[str, str] = field(default_factory=dict)

    @property
    def query_aligned_len(self) -> int:
        """Return the aligned length on the query sequence.

        Returns
        -------
        int
            ``query_end - query_start``.
        """
        return self.query_end - self.query_start

    @property
    def target_aligned_len(self) -> int:
        """Return the aligned length on the target sequence.

        Returns
        -------
        int
            ``target_end - target_start``.
        """
        return self.target_end - self.target_start

    @property
    def identity(self) -> float:
        """Return the fraction identity of this alignment in ``[0, 1]``.

        The most accurate available metric is used, in order of preference:

        1. ``1 - de`` from the ``de:f`` tag — minimap2's gap-compressed
           per-base divergence, emitted with base-level alignment (``-c``
           or ``--cs``).  A gap run counts as a single difference.
        2. Gap-compressed identity derived from the CIGAR:
           ``residue_matches / (aligned columns + gap openings)``.  Column
           10 is used as the match count because minimap2's default
           ``M``-style CIGAR does not distinguish mismatches.
        3. BLAST-style identity from the required PAF columns:
           ``residue_matches / alignment_block_len``.  For output without
           base-level alignment this is an approximation.

        Returns
        -------
        float
            Fraction identity, clamped to ``[0, 1]``.
        """
        de = self.tags.get('de')
        if isinstance(de, float):
            return max(0.0, min(1.0, 1.0 - de))
        if self.cigar is not None:
            ops = _parse_cigar(self.cigar)
            aligned_cols = ops.get('M', 0) + ops.get('=', 0) + ops.get('X', 0)
            denom = aligned_cols + (self.n_gaps or 0)
            if denom > 0:
                return min(1.0, self.residue_matches / denom)
        if self.alignment_block_len > 0:
            return min(1.0, self.residue_matches / self.alignment_block_len)
        return 1.0

    @classmethod
    def from_line(cls, line: str) -> 'PafRecord':
        """Parse a single PAF text line into a :class:`PafRecord`.

        Parameters
        ----------
        line : str
            A single PAF record line (tab-separated, trailing newline optional).

        Returns
        -------
        PafRecord
            The parsed record.

        Raises
        ------
        ValueError
            If the line has fewer than 12 tab-separated fields.
        """
        fields = line.rstrip('\n').split('\t')
        if len(fields) < 12:
            raise ValueError(
                f'PAF line has {len(fields)} fields; expected at least 12: {line!r}'
            )
        tags: dict[str, Any] = {}
        tag_types: dict[str, str] = {}
        cigar: str | None = None
        for tag_field in fields[12:]:
            parts = tag_field.split(':', 2)
            if len(parts) == 3:
                tag_name, tag_type, tag_value = parts
                if tag_type == 'i':
                    tags[tag_name] = int(tag_value)
                elif tag_type == 'f':
                    tags[tag_name] = float(tag_value)
                else:
                    tags[tag_name] = tag_value
                tag_types[tag_name] = tag_type
                if tag_name == 'cg' and tag_type == 'Z':
                    cigar = tag_value

        residue_matches = int(fields[9])
        stats: dict[str, int] = {}
        if cigar is not None:
            stats = _cigar_stats(cigar, residue_matches)

        return cls(
            query_name=fields[0],
            query_len=int(fields[1]),
            query_start=int(fields[2]),
            query_end=int(fields[3]),
            strand=fields[4],
            target_name=fields[5],
            target_len=int(fields[6]),
            target_start=int(fields[7]),
            target_end=int(fields[8]),
            residue_matches=residue_matches,
            alignment_block_len=int(fields[10]),
            mapping_quality=int(fields[11]),
            tags=tags,
            cigar=cigar,
            alignment_length=stats.get('alignment_length'),
            n_matches=stats.get('n_matches'),
            n_mismatches=stats.get('n_mismatches'),
            n_gaps=stats.get('n_gaps'),
            n_gap_bases=stats.get('n_gap_bases'),
            tag_types=tag_types,
        )

    def to_line(self) -> str:
        """Serialise this record back to a PAF-format string (no trailing newline).

        Returns
        -------
        str
            Tab-separated PAF line with the 12 required columns followed by
            any optional tags in insertion order.  Tag type characters come
            from :attr:`tag_types` when recorded at parse time, otherwise
            they are inferred from the Python value (``int`` -> ``i``,
            ``float`` -> ``f``, else ``Z``).
        """
        parts = [
            str(v)
            for v in [
                self.query_name,
                self.query_len,
                self.query_start,
                self.query_end,
                self.strand,
                self.target_name,
                self.target_len,
                self.target_start,
                self.target_end,
                self.residue_matches,
                self.alignment_block_len,
                self.mapping_quality,
            ]
        ]
        for name, value in self.tags.items():
            tag_type = self.tag_types.get(name)
            if tag_type is None:
                if isinstance(value, int):
                    tag_type = 'i'
                elif isinstance(value, float):
                    tag_type = 'f'
                else:
                    tag_type = 'Z'
            text = f'{value:g}' if isinstance(value, float) else str(value)
            parts.append(f'{name}:{tag_type}:{text}')
        return '\t'.join(parts)

class PafAlignment:
    """A collection of PAF alignment records with contig-ordering utilities.

    Can be constructed from a file path or an iterable of :class:`PafRecord`
    objects.  Provides :meth:`reorder_contigs` to sort query and target
    sequence names so that a subsequent dotplot shows maximum collinearity.

    Parameters
    ----------
    records : list of PafRecord
        The alignment records.

    Examples
    --------
    Load from a file and reorder contigs:

    >>> aln = PafAlignment.from_file("alignments.paf")
    >>> q_order, t_order = aln.reorder_contigs(aln.query_names, aln.target_names)
    """

    def __init__(self, records: list[PafRecord]) -> None:
        self.records: list[PafRecord] = records
        # Custom group assignments.  None means use the default (query_names
        # → 'a', target_names → 'b') which is computed lazily from records.
        self._groups: dict[str, list[str]] | None = None
        # Query contigs detected as reverse-oriented by the most recent
        # :meth:`reorder_contigs` call (empty until then).
        self._reversed_query: set[str] = set()

    # ------------------------------------------------------------------
    # Constructors
    # ------------------------------------------------------------------

    @classmethod
    def from_file(cls, path: str | Path) -> 'PafAlignment':
        """Load records from a PAF file.

        Parameters
        ----------
        path : str or Path
            Path to the PAF file.

        Returns
        -------
        PafAlignment
            New instance with all records loaded.
        """
        return cls(list(parse_paf_file(path)))

    @classmethod
    def from_records(cls, records: Iterable[PafRecord]) -> 'PafAlignment':
        """Construct from an iterable of :class:`PafRecord` objects.

        Parameters
        ----------
        records : iterable of PafRecord
            Source records.

        Returns
        -------
        PafAlignment
            New instance.
        """
        return cls(list(records))

    # ------------------------------------------------------------------
    # Properties
    # ------------------------------------------------------------------

    @property
    def query_names(self) -> list[str]:
        """Return a deduplicated list of query sequence names (insertion order).

        Returns
        -------
        list[str]
            Unique query names in the order first seen.
        """
        seen: dict[str, None] = {}
        for rec in self.records:
            seen[rec.query_name] = None
        return list(seen)

    @property
    def target_names(self) -> list[str]:
        """Return a deduplicated list of target sequence names (insertion order).

        Returns
        -------
        list[str]
            Unique target names in the order first seen.
        """
        seen: dict[str, None] = {}
        for rec in self.records:
            seen[rec.target_name] = None
        return list(seen)

    def sequence_names(self) -> list[str]:
        """Return a deduplicated list of all query and target sequence names.

        The list contains each name at most once, in the order it was first
        encountered (queries before targets within each record).  This method
        makes :class:`PafAlignment` compatible with :class:`~dot_explorer.dotplot.DotPlotter`.

        Returns
        -------
        list[str]
            All unique sequence names across query and target fields.
        """
        seen: dict[str, None] = {}
        for rec in self.records:
            seen[rec.query_name] = None
            seen[rec.target_name] = None
        return list(seen)

    def get_sequence_length(self, name: str) -> int:
        """Return the length of a sequence by name as stored in PAF records.

        Looks up *name* in the ``query_name`` and ``target_name`` fields of
        every record and returns the corresponding ``query_len`` or
        ``target_len``.  This method makes :class:`PafAlignment` compatible
        with :class:`~dot_explorer.dotplot.DotPlotter`.

        Parameters
        ----------
        name : str
            Sequence name to look up.

        Returns
        -------
        int
            Length of the sequence.

        Raises
        ------
        KeyError
            If *name* is not found in any record.
        """
        for rec in self.records:
            if rec.query_name == name:
                return rec.query_len
            if rec.target_name == name:
                return rec.target_len
        raise KeyError(f'Sequence {name!r} not found in PAF records.')

    def __len__(self) -> int:
        """Return the number of records.

        Returns
        -------
        int
            Record count.
        """
        return len(self.records)

    def __repr__(self) -> str:
        """Return a concise string representation.

        Returns
        -------
        str
            ``PafAlignment(records=<n>, queries=<q>, targets=<t>)``.
        """
        return (
            f'PafAlignment(records={len(self.records)}, '
            f'queries={len(self.query_names)}, '
            f'targets={len(self.target_names)})'
        )

    # ------------------------------------------------------------------
    # Filtering
    # ------------------------------------------------------------------

    def filter_by_query(self, names: Iterable[str]) -> 'PafAlignment':
        """Return a new :class:`PafAlignment` containing only the given query names.

        Parameters
        ----------
        names : iterable of str
            Query names to keep.

        Returns
        -------
        PafAlignment
            Filtered alignment.
        """
        keep = set(names)
        return PafAlignment([r for r in self.records if r.query_name in keep])

    def filter_by_target(self, names: Iterable[str]) -> 'PafAlignment':
        """Return a new :class:`PafAlignment` containing only the given target names.

        Parameters
        ----------
        names : iterable of str
            Target names to keep.

        Returns
        -------
        PafAlignment
            Filtered alignment.
        """
        keep = set(names)
        return PafAlignment([r for r in self.records if r.target_name in keep])

    def filter_by_min_length(self, min_length: int) -> 'PafAlignment':
        """Return a new :class:`PafAlignment` keeping only records of sufficient length.

        Filters on the query aligned length (``query_end - query_start``), which
        equals the alignment block span for both merged k-mer runs and PAF
        alignments imported from a file.

        Parameters
        ----------
        min_length : int
            Minimum alignment length (inclusive).  Records with a query aligned
            length strictly less than ``min_length`` are discarded.

        Returns
        -------
        PafAlignment
            Filtered alignment containing only records with
            ``query_aligned_len >= min_length``.
        """
        return PafAlignment(
            [r for r in self.records if r.query_aligned_len >= min_length]
        )

    # ------------------------------------------------------------------
    # Group management
    # ------------------------------------------------------------------

    @property
    def groups(self) -> dict[str, list[str]]:
        """Return the current group assignments.

        If groups have not been set explicitly via :meth:`set_groups` or
        :meth:`rename_group`, returns the default: all query sequence names
        in group ``'a'`` and all target sequence names in group ``'b'``.

        Returns
        -------
        dict[str, list[str]]
            Mapping of group label → list of sequence names.
        """
        if self._groups is not None:
            return dict(self._groups)
        return {'a': self.query_names, 'b': self.target_names}

    def set_groups(self, groups: dict[str, list[str]]) -> None:
        """Set custom group assignments for sequence names.

        Parameters
        ----------
        groups : dict[str, list[str]]
            Mapping of group label → list of sequence names belonging to
            that group.

        Warns
        -----
        Logs a warning for every sequence name that appears in more than one
        group.
        """
        seen: dict[str, str] = {}
        for group, names in groups.items():
            for name in names:
                if name in seen:
                    _log.warning(
                        'PafAlignment.set_groups: sequence %r is assigned to '
                        'both group %r and group %r',
                        name,
                        seen[name],
                        group,
                    )
                else:
                    seen[name] = group
        self._groups = {g: list(ns) for g, ns in groups.items()}

    def rename_group(self, old_name: str, new_name: str) -> None:
        """Rename a group label.

        If custom groups have not been set yet, the default assignment
        (``'a'`` → query names, ``'b'`` → target names) is materialised
        first.

        Parameters
        ----------
        old_name : str
            Current group label.
        new_name : str
            New group label.

        Raises
        ------
        KeyError
            If *old_name* is not a known group.
        ValueError
            If *new_name* already exists as a different group label.
        """
        current = self._groups if self._groups is not None else self.groups
        if old_name not in current:
            raise KeyError(f'Group {old_name!r} not found.')
        if new_name in current and new_name != old_name:
            raise ValueError(f'Group {new_name!r} already exists.')
        self._groups = {
            (new_name if k == old_name else k): v for k, v in current.items()
        }
        _log.info('PafAlignment: renamed group %r → %r', old_name, new_name)

    # ------------------------------------------------------------------
    # Contig reordering
    # ------------------------------------------------------------------

    def reorder_contigs(
        self,
        query_names: list[str] | None = None,
        target_names: list[str] | None = None,
        query_group: str | None = None,
        target_group: str | None = None,
    ) -> tuple[list[str], list[str]]:
        """Sort query and target contigs to maximise collinearity in the dotplot.

        Uses the gravity-centre algorithm: each contig is assigned a gravity
        equal to the weighted mean position of its alignment blocks on the
        opposing axis.  Contigs are then sorted by ascending gravity.

        Parameters
        ----------
        query_names : list[str] or None, optional
            Query contigs to reorder.  Ignored when *query_group* is given.
            Defaults to :attr:`query_names`.
        target_names : list[str] or None, optional
            Target contigs to reorder.  Ignored when *target_group* is given.
            Defaults to :attr:`target_names`.
        query_group : str or None, optional
            Group label whose members are used as query contigs.  When
            provided, the corresponding entry in :attr:`groups` is used and
            *query_names* is ignored.
        target_group : str or None, optional
            Group label whose members are used as target contigs.  When
            provided, the corresponding entry in :attr:`groups` is used and
            *target_names* is ignored.

        Returns
        -------
        tuple[list[str], list[str]]
            ``(sorted_query_names, sorted_target_names)``.  The set of query
            contigs detected as reverse-oriented is stored on
            :attr:`reversed_contigs` (not returned, for backward compatibility).

        Raises
        ------
        KeyError
            If a supplied group label is not present in :attr:`groups`.
        """
        current_groups = self.groups
        if query_group is not None:
            if query_group not in current_groups:
                raise KeyError(f'Group {query_group!r} not found.')
            q = current_groups[query_group]
        else:
            q = query_names if query_names is not None else self.query_names

        if target_group is not None:
            if target_group not in current_groups:
                raise KeyError(f'Group {target_group!r} not found.')
            t = current_groups[target_group]
        else:
            t = target_names if target_names is not None else self.target_names

        sorted_q, sorted_t, reversed_q = compute_gravity_contigs(self.records, q, t)
        self._reversed_query = reversed_q
        return sorted_q, sorted_t

    @property
    def reversed_contigs(self) -> set[str]:
        """Query contigs detected as reverse-oriented by :meth:`reorder_contigs`.

        Empty until :meth:`reorder_contigs` has been called.  These contigs
        align in reverse orientation against their best-matching target and can
        be passed to :meth:`~dot_explorer.dotplot.DotPlotter.plot` via
        ``reverse_contigs=`` to be rendered flipped along the main diagonal.

        Returns
        -------
        set[str]
            Names of reverse-oriented query contigs.
        """
        return set(self._reversed_query)