# geney/mirna/seed.py — miRNA seed match scanning.
"""
Efficient seed-region scanning for canonical miRNA binding sites.

Canonical miRNA-mRNA interactions involve complementary base pairing at
the 5' end of the miRNA (positions 2-8), called the seed region. Four
canonical seed types are recognized (TargetScan conventions):

  - 8mer:    positions 2-8 match + A at position 1 (strongest)
  - 7mer-m8: positions 2-8 match (no A1 requirement)
  - 7mer-A1: positions 2-7 match + A at position 1
  - 6mer:    positions 2-7 match (weakest canonical)

Only sites in the 3'UTR and last third of the ORF are considered
functionally relevant for miRNA-mediated repression.
"""
from __future__ import annotations

import re
from dataclasses import dataclass
from typing import Optional


# Canonical seed types ordered by strength
SEED_TYPES = ("8mer", "7mer-m8", "7mer-A1", "6mer")

# Watson-Crick complement (RNA context: miRNA G pairs with mRNA C)
_COMPLEMENT = str.maketrans("ACGT", "TGCA")


def _reverse_complement(seq: str) -> str:
    """Reverse complement of a DNA/RNA sequence."""
    return seq.upper().replace("U", "T").translate(_COMPLEMENT)[::-1]


@dataclass
class SeedMatch:
    """A single miRNA seed match on an mRNA sequence."""
    position: int           # 0-based start position on the mRNA
    end: int                # 0-based end (exclusive)
    seed_type: str          # '8mer', '7mer-m8', '7mer-A1', '6mer'
    mirna_id: str           # miRNA identifier (if known)
    seed_seq: str           # The matched seed sequence on the mRNA
    region: str             # 'utr3', 'orf_last_third', 'orf'
    strength_rank: int      # 0=8mer (strongest) to 3=6mer (weakest)

    @property
    def strength(self) -> float:
        """Relative binding strength: 8mer=1.0, 7mer-m8=0.8, 7mer-A1=0.6, 6mer=0.4."""
        return {0: 1.0, 1: 0.8, 2: 0.6, 3: 0.4}.get(self.strength_rank, 0.2)


def find_seed_matches(
    mrna_seq: str,
    mirna_seeds: dict[str, str],
    utr3_start: Optional[int] = None,
    orf_start: int = 0,
    orf_end: Optional[int] = None,
) -> list[SeedMatch]:
    """Scan an mRNA for canonical miRNA seed matches.

    Efficiently finds all canonical binding sites by searching for the
    reverse complement of each miRNA's seed region (positions 2-8).

    Args:
        mrna_seq: Full mRNA sequence (5'UTR + CDS + 3'UTR).
        mirna_seeds: Dict of {mirna_id: seed_sequence} where seed_sequence
            is positions 2-8 of the mature miRNA (7 nt). Can also provide
            the full mature miRNA (21-24 nt) and the seed will be extracted.
        utr3_start: 0-based position where the 3'UTR begins. If None,
            estimated as the last third of the sequence.
        orf_start: 0-based position where the ORF begins (default 0).
        orf_end: 0-based position where the ORF ends. If None, uses utr3_start.

    Returns:
        List of SeedMatch objects sorted by position.
    """
    seq = mrna_seq.upper().replace("U", "T")
    n = len(seq)

    # Estimate UTR boundaries if not provided
    if utr3_start is None:
        utr3_start = int(n * 0.67)  # rough: last third is 3'UTR-ish
    if orf_end is None:
        orf_end = utr3_start

    orf_last_third_start = orf_start + int((orf_end - orf_start) * 0.67)

    matches: list[SeedMatch] = []

    for mirna_id, seed_input in mirna_seeds.items():
        seed_input = seed_input.upper().replace("U", "T")

        # Extract seed if full miRNA provided
        if len(seed_input) > 10:
            seed_7 = seed_input[1:8]   # positions 2-8
            seed_6 = seed_input[1:7]   # positions 2-7
            pos1_base = seed_input[0]  # position 1
        elif len(seed_input) >= 7:
            seed_7 = seed_input[:7]
            seed_6 = seed_input[:6]
            pos1_base = None
        else:
            seed_6 = seed_input[:6] if len(seed_input) >= 6 else seed_input
            seed_7 = seed_input[:7] if len(seed_input) >= 7 else None
            pos1_base = None

        # Target site = reverse complement of seed on the mRNA
        target_7 = _reverse_complement(seed_7) if seed_7 else None
        target_6 = _reverse_complement(seed_6)

        # Scan for 7nt matches first (covers 8mer and 7mer-m8)
        if target_7:
            for m in re.finditer(re.escape(target_7), seq):
                pos = m.start()
                end = m.end()

                # Classify region
                if pos >= utr3_start:
                    region = "utr3"
                elif pos >= orf_last_third_start:
                    region = "orf_last_third"
                else:
                    region = "orf"

                # Check for 8mer: A at position immediately downstream of target
                # (mRNA position 1 opposite = A downstream of the 7mer match)
                if end < n and seq[end] == "A":
                    matches.append(SeedMatch(
                        position=pos, end=end + 1, seed_type="8mer",
                        mirna_id=mirna_id, seed_seq=seq[pos:end + 1],
                        region=region, strength_rank=0,
                    ))
                else:
                    matches.append(SeedMatch(
                        position=pos, end=end, seed_type="7mer-m8",
                        mirna_id=mirna_id, seed_seq=seq[pos:end],
                        region=region, strength_rank=1,
                    ))

        # Scan for 6nt matches (7mer-A1 and 6mer)
        for m in re.finditer(re.escape(target_6), seq):
            pos = m.start()
            end = m.end()

            # Skip if this is already covered by a 7nt match at the same position
            if target_7 and seq[pos:pos + len(target_7)] == target_7:
                continue

            if pos >= utr3_start:
                region = "utr3"
            elif pos >= orf_last_third_start:
                region = "orf_last_third"
            else:
                region = "orf"

            # Check for 7mer-A1: A downstream of the 6mer
            if end < n and seq[end] == "A":
                matches.append(SeedMatch(
                    position=pos, end=end + 1, seed_type="7mer-A1",
                    mirna_id=mirna_id, seed_seq=seq[pos:end + 1],
                    region=region, strength_rank=2,
                ))
            else:
                matches.append(SeedMatch(
                    position=pos, end=end, seed_type="6mer",
                    mirna_id=mirna_id, seed_seq=seq[pos:end],
                    region=region, strength_rank=3,
                ))

    # Sort by position
    matches.sort(key=lambda m: (m.position, m.strength_rank))
    return matches
