repliconf.py

# Copyright (C) 2026 Fufu Fang
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <https://www.gnu.org/licenses/>.

"""Replication configuration-related classes for AmplifyP."""

from __future__ import annotations

import builtins
from dataclasses import dataclass

from .dna import DNA, DNADirection, DNAType, Nucleotides, Primer
from .origin import ReplicationOrigin
from .settings import (
    GLOBAL_REPLICATION_SETTINGS,
    LengthWiseWeightTbl,
    ReplicationSettings,
)

COMPLEMENT_TABLE = str.maketrans(
    "ACGTMKRYBDHVacgtmkrybdhv", "TGCAKMYRVHDBtgcakmyrvhdb"
)


@dataclass(slots=True, frozen=True)
class DirIdx:
    """A class representing a directed index.

    This class encapsulates a specific position (index) on a DNA strand,
    along with the direction (forward or reverse) of that strand.

    Attributes:
        direction (DNADirection): The direction of the DNA strand.
        index (int): The integer index position.
    """

    direction: DNADirection
    index: int

    def __int__(self) -> int:
        """Return the index value as an integer.

        Returns:
            int: The index.
        """
        return self.index

    def __index__(self) -> int:
        """Return the index value for slicing/indexing usage.

        Returns:
            int: The index.
        """
        return self.index

    def __add__(self, other: object) -> DirIdx:
        """Add an integer or another DirIdx to this index.

        Args:
            other (int | DirIdx): The value to add.

        Returns:
            DirIdx: A new DirIdx with the updated index value. The direction
                remains unchanged.
        """
        if isinstance(other, int):
            return DirIdx(self.direction, self.index + other)
        if not isinstance(other, DirIdx):
            return NotImplemented
        return DirIdx(self.direction, self.index + other.index)

    def __sub__(self, other: object) -> DirIdx:
        """Subtract an integer or another DirIdx from this index.

        Args:
            other (int | DirIdx): The value to subtract.

        Returns:
            DirIdx: A new DirIdx with the updated index value. The direction
                remains unchanged.
        """
        if isinstance(other, int):
            return DirIdx(self.direction, self.index - other)
        if not isinstance(other, DirIdx):
            return NotImplemented
        return DirIdx(self.direction, self.index - other.index)

    def __eq__(self, other: object) -> bool:
        """Check equality with an integer or another DirIdx.

        Args:
            other (object): The object to compare.

        Returns:
            bool: True if equal, False otherwise. If comparing with int, checks
                index only.
        """
        if isinstance(other, int):
            return self.index == other
        if not isinstance(other, DirIdx):
            return NotImplemented
        return self.direction == other.direction and self.index == other.index

    def __lt__(self, other: object) -> bool:
        """Check if index is less than another.

        Args:
            other (object): Value to compare.

        Returns:
            bool: True if less than.
        """
        if isinstance(other, int):
            return self.index < other
        if not isinstance(other, DirIdx):
            return NotImplemented
        return self.index < other.index

    def __gt__(self, other: object) -> bool:
        """Check if index is greater than another.

        Args:
            other (object): Value to compare.

        Returns:
            bool: True if greater than.
        """
        if isinstance(other, int):
            return self.index > other
        if not isinstance(other, DirIdx):
            return NotImplemented
        return self.index > other.index

    def __le__(self, other: object) -> bool:
        """Check if index is less than or equal to another.

        Args:
            other (object): Value to compare.

        Returns:
            bool: True if less than or equal.
        """
        if isinstance(other, int):
            return self.index <= other
        if not isinstance(other, DirIdx):
            return NotImplemented
        return self.index <= other.index

    def __ge__(self, other: object) -> bool:
        """Check if index is greater than or equal to another.

        Args:
            other (object): Value to compare.

        Returns:
            bool: True if greater than or equal.
        """
        if isinstance(other, int):
            return self.index >= other
        if not isinstance(other, DirIdx):
            return NotImplemented
        return self.index >= other.index

    def __str__(self) -> str:
        """Return the string representation of the index.

        Returns:
            str: The index value as a string.
        """
        return f"{self.index}"


@dataclass(slots=True)
class DirIdxDb:
    """A database for storing valid replication origin locations.

    This container holds lists of indices where valid replication origins were
    found, segregated by strand direction (forward and reverse). It also tracks
    whether a search operation has been performed.

    Attributes:
        fwd (list[DirIdx]): A list of locations (DirIdx) for origins on the
            forward strand.
        rev (list[DirIdx]): A list of locations (DirIdx) for origins on the
            reverse strand.
        searched (bool): Indicates if the search has been executed. Defaults to
            False.
    """

    fwd: list[DirIdx]
    rev: list[DirIdx]
    searched: bool = False

    def clear(self) -> None:
        """Clear the database.

        Removes stored indices and resets search flag.
        """
        self.fwd.clear()
        self.rev.clear()
        self.searched = False

    def __getitem__(self, key: tuple[DNADirection, int]) -> DirIdx:
        """Retrieve a stored DirIdx by direction and list index.

        Args:
            key (tuple[DNADirection, int]): A tuple specifying the direction
                (FWD/REV) and the integer index within that list.

        Returns:
            DirIdx: The requested directed index.
        """
        direction, i = key
        if direction == DNADirection.FWD:
            return self.fwd[i]
        else:
            return self.rev[i]


class Repliconf:
    """A class representing a Replication Configuration.

    A 'Repliconf' encapsulates the setup required to search for replication
    origins given a specific Primer and Template DNA. It prepares the template
    sequences (including padding and complementary strands) and manages the
    search process.

    Attributes:
        padding_len (int): The length of padding added to the template,
            typically equal to the primer length.
        primer (Primer): The primer sequence being analyzed.
        template (DNA): The original template DNA sequence.
        template_seq (dict[DNADirection, str]): A dictionary mapping direction
            (FWD/REV) to the processed (padded/complemented) template strings.
        settings (Settings): The configuration settings (weights, cutoffs) used
            for scoring.
        origin_db (DirIdxDb): The results database storing found origin indices.
    """

    def __init__(
        self,
        template: DNA,
        primer: Primer,
        settings: ReplicationSettings = GLOBAL_REPLICATION_SETTINGS,
    ) -> None:
        """Initialize a Repliconf object.

        Args:
            template (DNA): The template DNA sequence to be searched.
            primer (Primer): The primer sequence to search for.
            settings (Settings): The settings used for evaluating origins.
        """
        self.padding_len = len(primer)

        self.primer = primer
        self.template = template

        self.template_seq: dict[DNADirection, str] = {}
        # Add padding the 5' end of the template
        if template.type == DNAType.LINEAR or template.type == DNAType.PRIMER:
            self.template_seq[DNADirection.FWD] = (
                Nucleotides.GAP * self.padding_len + template.seq
            )
            self.template_seq[DNADirection.REV] = (
                template.seq + Nucleotides.GAP * self.padding_len
            ).translate(COMPLEMENT_TABLE)
        elif template.type == DNAType.CIRCULAR:
            # Matches existing behavior including 0 case where padding is empty
            padding = (
                template.seq[-self.padding_len :]
                if self.padding_len > 0
                else ""
            )
            self.template_seq[DNADirection.FWD] = padding + template.seq
            self.template_seq[DNADirection.REV] = (
                template.seq + template.seq[: self.padding_len]
            ).translate(COMPLEMENT_TABLE)

        self.settings = settings
        self.origin_db = DirIdxDb([], [], False)

        # Pre-calculate reversed primer sequence
        self._rev_primer_seq = self.primer.seq[::-1]

        # Pre-calculate reverse complement primer for amplicon generation
        self.rev_comp_primer = self.primer.reverse_complement()

    def range(self) -> range:
        """Return the range of valid starting indices for search.

        Returns:
            range: A range object covering all valid start positions in the
                padded template.
        """
        return range(
            len(self.template_seq[DNADirection.FWD]) - len(self.primer) + 1
        )

    def slice(self, i: int) -> slice:
        """Create a slice object for extracting a target segment.

        Args:
            i (int): The start index.

        Returns:
            slice: A slice of length equal to the primer length.
        """
        return slice(i, i + len(self.primer))

    def origin(self, var: DirIdx) -> ReplicationOrigin:
        """Construct a ReplicationOrigin object for a specific location.

        This method extracts the target sequence from the pre-processed template
        at the specified location and direction, and creates a ReplicationOrigin
        object to evaluate it.

        Args:
            var (DirIdx): A `DirIdx` object containing both direction and index.

        Returns:
            ReplicationOrigin: An evaluated replication origin object.
        """
        direction = var.direction
        i = var.index

        if direction:
            # Optimized reverse slicing: template_seq[end-1:start-1:-1]
            # This creates only 1 string object instead of 2 (slice then
            # reverse)
            end = i + len(self.primer)
            target = (
                self.template_seq[direction][end - 1 : i - 1 : -1]
                if i > 0
                else self.template_seq[direction][end - 1 :: -1]
            )
        else:
            target = self.template_seq[direction][self.slice(i)]

        return ReplicationOrigin(
            target,
            self._rev_primer_seq,
            self.settings,
        )

    def origin_from_db(
        self, direction: DNADirection, i: int
    ) -> ReplicationOrigin:
        """Retrieve a ReplicationOrigin based on database index.

        Args:
            direction (DNADirection): The direction to look up in the database.
            i (int): The index *within the list of found origins* (not the
                genomic index).

        Returns:
            ReplicationOrigin: The replication origin object.
        """
        return self.origin(self.origin_db[direction, i])

    def _scan_direction(
        self,
        direction: DNADirection,
        seq: str,
        search_range: builtins.range,
        L: int,
        prim_score_lookup: list[dict[str, float]],
        stab_score_lookup: list[dict[str, float]],
        constants: tuple[float, float, float, float, LengthWiseWeightTbl],
    ) -> None:
        """Scan a specific direction for valid replication origins.

        Args:
            direction (DNADirection): The direction to scan.
            seq (str): The sequence to scan.
            search_range (range): The range of indices to search.
            L (int): The length of the primer.
            prim_score_lookup (list[dict[str, float]]): Primability scores.
            stab_score_lookup (list[dict[str, float]]): Stability scores.
            constants (tuple): Tuple containing (prim_denom, stab_denom,
                prim_cutoff, stab_cutoff, r).
        """
        prim_denom, stab_denom, prim_cutoff, stab_cutoff, r = constants
        origin_list = (
            self.origin_db.fwd
            if direction == DNADirection.FWD
            else self.origin_db.rev
        )

        # Pre-compute index generator to avoid conditional inside inner loop
        if direction == DNADirection.FWD:
            # FWD: target[k] = seq[i + L - 1 - k]
            # map k (0..L-1) to indices (i+L-1 .. i)
            # Range for seq_indices: start=i+L-1, stop=i-1, step=-1
            def range_args(i: int) -> tuple[int, int, int]:
                return (i + L - 1, i - 1, -1)
        else:
            # REV: target[k] = seq[i + k]
            # map k (0..L-1) to indices (i .. i+L-1)
            def range_args(i: int) -> tuple[int, int, int]:
                return (i, i + L, 1)

        for i in search_range:
            prim_num = 0.0
            stab_num = 0.0
            run_len = 0
            run_score = 0.0

            seq_indices = range(*range_args(i))

            for k, seq_idx in enumerate(seq_indices):
                base_t = seq[seq_idx]

                # Primability
                prim_num += prim_score_lookup[k][base_t]

                # Stability
                val = stab_score_lookup[k][base_t]
                if val > 0:
                    run_len += 1
                    run_score += val
                else:
                    if run_len > 0:
                        idx = run_len - 1
                        stab_num += r[idx] * run_score
                        run_len = 0
                        run_score = 0.0

            # Finish stability run
            if run_len > 0:
                idx = run_len - 1
                stab_num += r[idx] * run_score

            primability = prim_num / prim_denom if prim_denom != 0 else 0.0
            stability = stab_num / stab_denom if stab_denom != 0 else 0.0

            if primability > prim_cutoff and stability > stab_cutoff:
                origin_list.append(DirIdx(direction, i))

    def search(self) -> None:
        """Perform a search for valid replication origins.

        Scans the entire template in both forward and reverse directions.
        Locations where the primability and stability scores exceed the
        configured cutoffs are stored in `origin_db`.
        """
        self.origin_db.clear()

        # Optimization: Pre-calculate constants and lookup tables
        m = self.settings.match_weight
        S = self.settings.base_pair_scores
        r = self.settings.run_weights

        primer_rev = self.primer.seq[::-1]

        prim_denom = 0.0
        stab_denom_base = 0.0

        prim_score_lookup: list[dict[str, float]] = []
        stab_score_lookup: list[dict[str, float]] = []

        # Prepare set of characters for lookup keys
        lookup_keys = set(Nucleotides.TEMPLATE)
        lookup_keys.update(c.lower() for c in list(lookup_keys))

        for k, base_p in enumerate(primer_rev):
            row_max = S.row_max(base_p)
            prim_denom += m[k] * row_max
            stab_denom_base += row_max

            p_dict = {}
            s_dict = {}

            for base_t in lookup_keys:
                score = S[base_p, base_t]
                p_dict[base_t] = m[k] * score
                s_dict[base_t] = score

            prim_score_lookup.append(p_dict)
            stab_score_lookup.append(s_dict)

        stab_denom = stab_denom_base * r[int(max(0, len(primer_rev) - 1))]

        prim_cutoff = self.settings.primability_cutoff
        stab_cutoff = self.settings.stability_cutoff

        constants = (prim_denom, stab_denom, prim_cutoff, stab_cutoff, r)

        for direction in [DNADirection.FWD, DNADirection.REV]:
            # FWD direction:
            # - Searches `template_seq[FWD]`, which is the Sense strand (5'-3').
            # - Matches the primer sequence against the Sense strand (reversed).
            # - Effectively finds binding sites for Forward Primers (which are
            #   identical to the Sense strand) by checking for identity.
            #
            # REV direction:
            # - Searches `template_seq[REV]`, which is the Antisense strand
            #   (complement of Sense, read 3'-5' in physical space, but stored
            #   as a string).
            # - Matches the primer sequence against the Antisense strand.
            # - Effectively finds binding sites for Reverse Primers (which are
            #   complementary to the Sense strand) by checking for identity
            #   against the Antisense strand.

            self._scan_direction(
                direction,
                self.template_seq[direction],
                self.range(),
                len(primer_rev),
                prim_score_lookup,
                stab_score_lookup,
                constants,
            )

        self.origin_db.searched = True

    @property
    def searched(self) -> bool:
        """Check if a search has been performed.

        Returns:
            bool: True if `search()` has been called, False otherwise.
        """
        return self.origin_db.searched

    def __eq__(self, other: object) -> bool:
        """Check equality with another Repliconf.

        Equality is based on having the same primer and template.

        Args:
            other (object): The object to compare.

        Returns:
            bool: True if equal, False otherwise.
        """
        if not isinstance(other, Repliconf):
            return NotImplemented
        return self.primer == other.primer and self.template == other.template

    def __hash__(self) -> int:
        """Compute the hash of the Repliconf.

        Returns:
            int: The hash based on primer and template.
        """
        return hash((self.primer, self.template))

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

        Returns:
            str: Description of the configuration.
        """
        return (
            f"ReplicationConfig: Primer: {self.primer}, Target: {self.template}"
        )