assemble.py

#!/usr/bin/env python3
"""
Assemble 2 Sanger ABI reads (forward + reverse) from a folder into a consensus FASTA.

Requires: biopython
  pip install biopython

Usage:
  python assemble.py /path/to/folder
  python assemble.py --use-paired-mixture /path/to/folder
  python assemble.py --use-single-mixture /path/to/folder
  python assemble.py --min-phred-score-per-base 20 --min-consecutive-high-quality-bases 10 /path/to/folder

Input folder:
  - forward .ab1 file stem ends with 'F'
  - reverse .ab1 file stem ends with 'R'

Outputs are written into the same folder:
  - <folder_name>.fasta
  - forward_trimmed.fasta
  - reverse_rc_trimmed.fasta
  - <folder_name>_alignment.html
  - <folder_name>_QA.html

"""

import argparse
from collections import Counter
import json
from pathlib import Path
import sys
from dataclasses import dataclass
from typing import Any, List, Tuple

from Bio import SeqIO
from Bio.Seq import Seq
from Bio.Align import PairwiseAligner
from jinja2 import Environment, FileSystemLoader, select_autoescape
from align_reference import (
    align_consensus_to_reference_nucleotide,
    read_reference_nucleotide_fasta,
    write_amino_acid_alignment_fasta,
)
from mixture_detector import resolve_consensus_base
from single_strand_mixture_detector import (
    build_single_strand_mixtures_header,
    detect_single_strand_mixtures,
)
from quality_checks import run_quality_checks


TEMPLATE_ENV = Environment(
    loader=FileSystemLoader(Path(__file__).parent),
    autoescape=select_autoescape(["html"]),
)


def parse_min_phred_score_per_base(value: str | tuple[int, int]) -> tuple[int, int]:
    if isinstance(value, tuple):
        if len(value) != 2:
            raise ValueError("min_phred_score_per_base must contain exactly two values.")
        return value

    parts = value.split(":")
    if len(parts) != 2:
        raise ValueError(
            "min_phred_score_per_base must use format forward_phred:reverse_phred, for example 20:20."
        )
    try:
        forward_min_phred_score_per_base = int(parts[0])
        reverse_min_phred_score_per_base = int(parts[1])
    except ValueError as exc:
        raise ValueError(
            "min_phred_score_per_base must use integer values like 20:20."
        ) from exc
    return forward_min_phred_score_per_base, reverse_min_phred_score_per_base


def format_min_phred_score_per_base(value: tuple[int, int]) -> str:
    return f"{value[0]}:{value[1]}"


@dataclass
class Read:
    name: str
    seq: str
    qual: List[int]
    start: int
    end: int


@dataclass
class LowQualityPosition:
    consensus_pos: int
    forward_pos: int
    forward_base: str
    forward_qual: str
    reverse_pos: int
    reverse_base: str
    reverse_qual: str
    reason: str


@dataclass
class AbiTrace:
    name: str
    seq: str
    base_positions: List[int]
    channels: dict[str, List[int]]
    metadata: dict[str, str]
    key_sections: list[str]
    hierarchy_text: str


@dataclass
class AssemblyArtifacts:
    sample_name: str
    report_files: list[Path]
    consensus_files: list[Path]
    warning_file: Path | None

    @property
    def primary_report(self) -> Path | None:
        for suffix in ("_QA.html", "_alignment.html"):
            for path in self.report_files:
                if path.name.endswith(suffix):
                    return path
        if self.warning_file is not None:
            return self.warning_file
        return self.report_files[0] if self.report_files else None


def read_ab1(path: str, name: str) -> Read:
    rec = SeqIO.read(path, "abi")
    seq = str(rec.seq).upper()
    qual = rec.letter_annotations.get("phred_quality")
    if qual is None:
        raise ValueError(f"No phred qualities found in {path}.")
    if len(qual) != len(seq):
        raise ValueError(f"Quality length != sequence length in {path}.")
    return Read(name=name, seq=seq, qual=list(qual), start=0, end=len(seq))


def read_ab1_trace(path: str, name: str) -> AbiTrace:
    rec = SeqIO.read(path, "abi")
    raw = rec.annotations.get("abif_raw")
    if raw is None:
        raise ValueError(f"No ABI raw trace data found in {path}.")

    order = raw.get("FWO_1")
    if isinstance(order, bytes):
        order = order.decode("ascii")
    if not isinstance(order, str) or len(order) != 4:
        raise ValueError(f"Could not determine channel order from {path}.")

    base_positions = raw.get("PLOC2")
    if base_positions is None:
        base_positions = raw.get("PLOC1")
    if base_positions is None:
        raise ValueError(f"Could not determine base positions from {path}.")

    trace_keys = ["DATA9", "DATA10", "DATA11", "DATA12"]
    channels: dict[str, List[int]] = {}
    for base, key in zip(order, trace_keys):
        values = raw.get(key)
        if values is None:
            raise ValueError(f"Missing trace channel {key} in {path}.")
        channels[base] = list(values)

    for base in "ACGT":
        if base not in channels:
            raise ValueError(f"Missing trace channel for base {base} in {path}.")

    metadata = {}
    metadata_keys = {
        "SMPL1": "Sample Name",
        "MCHN1": "Instrument",
        "MODL1": "Instrument Model",
        "LANE1": "Lane",
        "RUNN1": "Run Name",
        "RUND1": "Run Date",
        "RUNT1": "Run Time",
        "DYEP1": "Dye Set",
        "CTID1": "Container ID",
    }
    for key, label in metadata_keys.items():
        value = raw.get(key)
        if isinstance(value, bytes):
            value = value.decode("ascii", errors="ignore").strip("\x00 ")
        if value not in (None, ""):
            metadata[label] = str(value)

    key_section_specs = [
        ("PBAS", ["PBAS2", "PBAS1"], "Base calls for the read sequence."),
        ("PCON", ["PCON2", "PCON1"], "Per-base confidence or quality values."),
        ("PLOC", ["PLOC2", "PLOC1"], "Trace indices for each called base."),
        ("FWO_", ["FWO_1"], "Channel order that maps trace data to A/C/G/T bases."),
        (
            "DATA",
            ["DATA9", "DATA10", "DATA11", "DATA12"],
            "Raw fluorescence trace channels used to build the chromatogram.",
        ),
        ("SMPL", ["SMPL1"], "Sample identifier or sample name."),
        ("MCHN", ["MCHN1"], "Instrument identifier."),
        ("MODL", ["MODL1"], "Instrument model."),
        ("RUN", ["RUNN1", "RUND1", "RUNT1"], "Run name, date, and time metadata."),
        ("DYEP", ["DYEP1"], "Dye set or chemistry information."),
        ("CTID", ["CTID1"], "Container or plate/tube identifier."),
    ]
    key_sections = []
    hierarchy_lines = ["ABIF"]
    for section_name, keys, purpose in key_section_specs:
        present_keys = [key for key in keys if raw.get(key) is not None]
        if present_keys:
            key_sections.append(
                f"{section_name}: present via {', '.join(present_keys)}. {purpose}"
            )
            hierarchy_lines.append(f"+- {section_name}: {purpose}")
            for key in present_keys:
                hierarchy_lines.append(f"|  +- {key}")

    return AbiTrace(
        name=name,
        seq=str(rec.seq).upper(),
        base_positions=list(base_positions),
        channels=channels,
        metadata=metadata,
        key_sections=key_sections,
        hierarchy_text="\n".join(hierarchy_lines),
    )


def trim_by_quality(
    r: Read,
    min_phred_score_per_base: int,
    min_consecutive_high_quality_bases: int = 10,
) -> Read:
    """
    Trim low-quality ends:
    - find first position where there's a run of >= min_consecutive_high_quality_bases
      bases with Q>=min_phred_score_per_base
    - find last position similarly (from the end)
    """
    n = len(r.seq)
    good = [q >= min_phred_score_per_base for q in r.qual]

    # find left trim
    left = 0
    found = False
    for i in range(0, n - min_consecutive_high_quality_bases + 1):
        if all(good[i : i + min_consecutive_high_quality_bases]):
            left = i
            found = True
            break
    if not found:
        # nothing good: return empty
        return Read(r.name, "", [], r.start, r.start)

    # find right trim
    right = n
    found = False
    for i in range(n, min_consecutive_high_quality_bases - 1, -1):
        # window ends at i, starts at i-min_consecutive_high_quality_bases
        if all(good[i - min_consecutive_high_quality_bases : i]):
            right = i
            found = True
            break
    if not found or right <= left:
        return Read(r.name, "", [], r.start, r.start)

    return Read(
        r.name,
        r.seq[left:right],
        r.qual[left:right],
        r.start + left,
        r.start + right,
    )


def revcomp_read(r: Read) -> Read:
    """
    Reverse-complement sequence; reverse qualities (no complement needed for qualities).
    """
    rc_seq = str(Seq(r.seq).reverse_complement())
    rc_qual = list(reversed(r.qual))
    return Read(r.name + "_rc", rc_seq, rc_qual, r.start, r.end)


def build_gapped_alignment(
    fwd_seq: str,
    rev_seq: str,
    f_blocks: List[Tuple[int, int]],
    r_blocks: List[Tuple[int, int]],
) -> Tuple[str, str]:
    gf = []
    gr = []
    f_i = 0
    r_i = 0

    for (fs, fe), (rs, re) in zip(f_blocks, r_blocks):
        if fs > f_i or rs > r_i:
            f_gap = fs - f_i
            r_gap = rs - r_i

            if f_gap > 0 and r_gap == 0:
                gf.append(fwd_seq[f_i:fs])
                gr.append("-" * f_gap)
            elif r_gap > 0 and f_gap == 0:
                gf.append("-" * r_gap)
                gr.append(rev_seq[r_i:rs])
            else:
                chunk_length = max(f_gap, r_gap)
                gf.append(fwd_seq[f_i:fs] + "-" * (chunk_length - f_gap))
                gr.append(rev_seq[r_i:rs] + "-" * (chunk_length - r_gap))

            f_i = fs
            r_i = rs

        gf.append(fwd_seq[fs:fe])
        gr.append(rev_seq[rs:re])
        f_i = fe
        r_i = re

    if f_i < len(fwd_seq) or r_i < len(rev_seq):
        f_suf = len(fwd_seq) - f_i
        r_suf = len(rev_seq) - r_i
        if f_suf > 0 and r_suf == 0:
            gf.append(fwd_seq[f_i:])
            gr.append("-" * f_suf)
        elif r_suf > 0 and f_suf == 0:
            gf.append("-" * r_suf)
            gr.append(rev_seq[r_i:])
        else:
            chunk_length = max(f_suf, r_suf)
            gf.append(fwd_seq[f_i:] + "-" * (chunk_length - f_suf))
            gr.append(rev_seq[r_i:] + "-" * (chunk_length - r_suf))

    gapped_f = "".join(gf)
    gapped_r = "".join(gr)
    if len(gapped_f) != len(gapped_r):
        raise RuntimeError("Internal error: gapped strings have different lengths.")
    return gapped_f, gapped_r


def configure_overlap_aligner(
    aligner: PairwiseAligner,
    match: float,
    mismatch: float,
    gap_open: float,
    gap_extend: float,
) -> None:
    aligner.mode = "global"
    aligner.match_score = match
    aligner.mismatch_score = mismatch
    aligner.open_gap_score = gap_open
    aligner.extend_gap_score = gap_extend

    available_attrs = set(dir(aligner))
    for attr in [
        "target_end_open_gap_score",
        "target_end_extend_gap_score",
        "query_end_open_gap_score",
        "query_end_extend_gap_score",
        "target_left_open_gap_score",
        "target_left_extend_gap_score",
        "target_right_open_gap_score",
        "target_right_extend_gap_score",
        "query_left_open_gap_score",
        "query_left_extend_gap_score",
        "query_right_open_gap_score",
        "query_right_extend_gap_score",
        "target_end_gap_score",
        "query_end_gap_score",
        "open_end_gap_score",
        "extend_end_gap_score",
        "end_gap_score",
        "end_open_gap_score",
        "end_extend_gap_score",
    ]:
        if attr in available_attrs:
            setattr(aligner, attr, 0.0)


def build_consensus_from_gapped_alignment(
    fwd: Read,
    rev: Read,
    gapped_f: str,
    gapped_r: str,
    use_paired_mixture: bool,
    forward_min_phred_score_per_base: int,
    reverse_min_phred_score_per_base: int,
    min_phred_score_for_paired_base: int,
    min_overlap: int,
    forward_single_mixture_map: dict[int, str] | None = None,
    reverse_single_mixture_map: dict[int, str] | None = None,
) -> Tuple[str, int, List[LowQualityPosition], List[str], List[str]]:
    fi = 0
    ri = 0
    overlap = 0
    consensus_chars = []
    consensus_line = []
    low_quality_positions: List[LowQualityPosition] = []
    consensus_pos = 0
    fwd_qual_line = []
    rev_qual_line = []

    for cf, cr in zip(gapped_f, gapped_r):
        bf = cf != "-"
        br = cr != "-"

        if bf and br:
            overlap += 1
            qf = fwd.qual[fi]
            qr = rev.qual[ri]
            consensus_pos += 1
            decision = resolve_consensus_base(
                cf,
                qf,
                cr,
                qr,
                forward_min_phred_score_per_base=forward_min_phred_score_per_base,
                reverse_min_phred_score_per_base=reverse_min_phred_score_per_base,
                min_phred_score_for_paired_base=min_phred_score_for_paired_base,
                allow_mixture=use_paired_mixture,
                forward_trimmed_position=fi + 1,
                reverse_trimmed_position=ri + 1,
                forward_single_mixture_map=forward_single_mixture_map,
                reverse_single_mixture_map=reverse_single_mixture_map,
            )
            consensus_chars.append(decision.consensus_base)
            consensus_line.append(decision.consensus_base)
            if decision.low_quality:
                low_quality_positions.append(
                    LowQualityPosition(
                        consensus_pos=consensus_pos,
                        forward_pos=fwd.start + fi + 1,
                        forward_base=cf,
                        forward_qual=str(qf),
                        reverse_pos=rev.start + ri + 1,
                        reverse_base=cr,
                        reverse_qual=str(qr),
                        reason=decision.reason or "low_quality",
                    )
                )
            fwd_qual_line.append(f"{qf:02d}")
            rev_qual_line.append(f"{qr:02d}")
            fi += 1
            ri += 1
        elif bf and not br:
            consensus_pos += 1
            qf = fwd.qual[fi]
            decision = resolve_consensus_base(
                cf,
                qf,
                "-",
                None,
                forward_min_phred_score_per_base=forward_min_phred_score_per_base,
                reverse_min_phred_score_per_base=reverse_min_phred_score_per_base,
                min_phred_score_for_paired_base=min_phred_score_for_paired_base,
                allow_mixture=use_paired_mixture,
                forward_trimmed_position=fi + 1,
                forward_single_mixture_map=forward_single_mixture_map,
                reverse_single_mixture_map=reverse_single_mixture_map,
            )
            consensus_chars.append(decision.consensus_base)
            consensus_line.append(decision.consensus_base)
            if decision.low_quality:
                low_quality_positions.append(
                    LowQualityPosition(
                        consensus_pos=consensus_pos,
                        forward_pos=fwd.start + fi + 1,
                        forward_base=cf,
                        forward_qual=str(qf),
                        reverse_pos=0,
                        reverse_base="-",
                        reverse_qual="-",
                        reason=decision.reason or "low_quality",
                    )
                )
            fwd_qual_line.append(f"{qf:02d}")
            rev_qual_line.append("--")
            fi += 1
        elif br and not bf:
            consensus_pos += 1
            qr = rev.qual[ri]
            decision = resolve_consensus_base(
                "-",
                None,
                cr,
                qr,
                forward_min_phred_score_per_base=forward_min_phred_score_per_base,
                reverse_min_phred_score_per_base=reverse_min_phred_score_per_base,
                min_phred_score_for_paired_base=min_phred_score_for_paired_base,
                allow_mixture=use_paired_mixture,
                reverse_trimmed_position=ri + 1,
                forward_single_mixture_map=forward_single_mixture_map,
                reverse_single_mixture_map=reverse_single_mixture_map,
            )
            consensus_chars.append(decision.consensus_base)
            consensus_line.append(decision.consensus_base)
            if decision.low_quality:
                low_quality_positions.append(
                    LowQualityPosition(
                        consensus_pos=consensus_pos,
                        forward_pos=0,
                        forward_base="-",
                        forward_qual="-",
                        reverse_pos=rev.start + ri + 1,
                        reverse_base=cr,
                        reverse_qual=str(qr),
                        reason=decision.reason or "low_quality",
                    )
                )
            fwd_qual_line.append("--")
            rev_qual_line.append(f"{qr:02d}")
            ri += 1
        else:
            consensus_pos += 1
            decision = resolve_consensus_base(
                "-",
                None,
                "-",
                None,
                forward_min_phred_score_per_base=forward_min_phred_score_per_base,
                reverse_min_phred_score_per_base=reverse_min_phred_score_per_base,
                min_phred_score_for_paired_base=min_phred_score_for_paired_base,
                allow_mixture=use_paired_mixture,
                forward_single_mixture_map=forward_single_mixture_map,
                reverse_single_mixture_map=reverse_single_mixture_map,
            )
            consensus_chars.append(decision.consensus_base)
            consensus_line.append(decision.consensus_base)
            low_quality_positions.append(
                LowQualityPosition(
                    consensus_pos=consensus_pos,
                    forward_pos=0,
                    forward_base="-",
                    forward_qual="-",
                    reverse_pos=0,
                    reverse_base="-",
                    reverse_qual="-",
                    reason=decision.reason or "double_gap",
                )
            )
            fwd_qual_line.append("--")
            rev_qual_line.append("--")

    if overlap < min_overlap:
        raise ValueError(
            f"Overlap too short ({overlap} bp). "
            f"Try lowering min_overlap or check trimming / correct reverse-complement."
        )

    consensus = "".join(consensus_chars).replace("-", "")
    return consensus, overlap, low_quality_positions, consensus_line, fwd_qual_line, rev_qual_line


def assemble_pair(
    fwd: Read,
    rev: Read,
    forward_min_phred_score_per_base: int,
    reverse_min_phred_score_per_base: int,
    min_phred_score_for_paired_base: int,
    use_paired_mixture: bool = True,
    use_single_mixture: bool = False,
    forward_single_mixture_map: dict[int, str] | None = None,
    reverse_single_mixture_map: dict[int, str] | None = None,
    match: float = 2.0,
    mismatch: float = -1.0,
    gap_open: float = -5.0,
    gap_extend: float = -1.0,
    min_overlap: int = 40,
) -> Tuple[str, dict, List[LowQualityPosition]]:
    """
    Align fwd vs rev (rev should already be reverse-complemented) and build a consensus.

    Returns:
      consensus_seq, alignment_payload, low_quality_positions
    """
    if not fwd.seq or not rev.seq:
        raise ValueError("One of the reads is empty after trimming.")

    aligner = PairwiseAligner()
    configure_overlap_aligner(
        aligner=aligner,
        match=match,
        mismatch=mismatch,
        gap_open=gap_open,
        gap_extend=gap_extend,
    )

    best = next(iter(aligner.align(fwd.seq, rev.seq)))
    gapped_f, gapped_r = build_gapped_alignment(
        fwd.seq,
        rev.seq,
        best.aligned[0],
        best.aligned[1],
    )
    consensus, overlap, low_quality_positions, consensus_line, fwd_qual_line, rev_qual_line = (
        build_consensus_from_gapped_alignment(
            fwd=fwd,
            rev=rev,
            gapped_f=gapped_f,
            gapped_r=gapped_r,
            use_paired_mixture=use_paired_mixture,
            forward_min_phred_score_per_base=forward_min_phred_score_per_base,
            reverse_min_phred_score_per_base=reverse_min_phred_score_per_base,
            min_phred_score_for_paired_base=min_phred_score_for_paired_base,
            min_overlap=min_overlap,
            forward_single_mixture_map=(
                forward_single_mixture_map if use_single_mixture else None
            ),
            reverse_single_mixture_map=(
                reverse_single_mixture_map if use_single_mixture else None
            ),
        )
    )
    fwd_readpos_line = []
    rev_readpos_line = []
    fwd_pos = fwd.start + 1
    rev_pos = rev.start + 1
    for cf, cr in zip(gapped_f, gapped_r):
        if cf != "-":
            fwd_readpos_line.append(str(fwd_pos))
            fwd_pos += 1
        else:
            fwd_readpos_line.append("")
        if cr != "-":
            rev_readpos_line.append(str(rev_pos))
            rev_pos += 1
        else:
            rev_readpos_line.append("")
    mid = []
    for cf, cr in zip(gapped_f, gapped_r):
        mid.append("|" if (cf == cr and cf != "-" and cr != "-") else " ")
    alignment_payload = {
        "fwd_positions": f"{fwd.start + 1}-{fwd.end}",
        "rev_positions": f"{rev.start + 1}-{rev.end}",
        "overlap": overlap,
        "alignment_data_json": json.dumps(
            {
                "gapped_f": list(gapped_f),
                "gapped_r": list(gapped_r),
                "consensus_line": consensus_line,
                "match_line": list("".join(mid)),
                "fwd_readpos_line": fwd_readpos_line,
                "fwd_qual_line": fwd_qual_line,
                "rev_readpos_line": rev_readpos_line,
                "rev_qual_line": rev_qual_line,
            }
        ),
    }
    return consensus, alignment_payload, low_quality_positions


def write_fasta(path: str, header: str, seq: str) -> None:
    with open(path, "w") as f:
        f.write(f">{header}\n")
        for i in range(0, len(seq), 80):
            f.write(seq[i : i + 80] + "\n")


def write_text(path: str, content: str) -> None:
    with open(path, "w") as f:
        f.write(content)


def write_warning_html(
    folder: Path,
    message: str,
    title: str = "Assembly Warning",
    banner_message: str | None = None,
) -> None:
    template = TEMPLATE_ENV.get_template("warning_template.html")
    warning_html = template.render(
        title=title,
        banner_message=banner_message,
        intro_message=message,
        min_consecutive_high_quality_bases=None,
        forward_min_phred_score_per_base=None,
        reverse_min_phred_score_per_base=None,
        empty_reads=[],
        footer_message=None,
    )
    write_text(str(folder / "Warning.html"), warning_html)


def cleanup_folder_outputs(folder: Path) -> None:
    for path in sorted(
        folder.rglob("*"), key=lambda item: len(item.parts), reverse=True
    ):
        if path.is_file() and path.suffix.lower() != ".ab1":
            path.unlink()
        elif path.is_dir() and path != folder:
            try:
                path.rmdir()
            except OSError:
                pass


def collect_assembly_artifacts(folder: Path) -> AssemblyArtifacts:
    folder = Path(folder).resolve()
    sample_name = folder.name
    report_files = sorted(
        [
            path
            for path in (
                folder / f"{sample_name}_QA.html",
                folder / f"{sample_name}_alignment.html",
            )
            if path.is_file()
        ]
    )
    consensus_files = sorted(
        [
            path
            for path in (
                folder / f"{sample_name}.fasta",
                folder / f"{sample_name}_aa_alignment.fasta",
                folder / "forward_trimmed.fasta",
                folder / "reverse_rc_trimmed.fasta",
            )
            if path.is_file()
        ]
    )
    warning_file = folder / "Warning.html"
    return AssemblyArtifacts(
        sample_name=sample_name,
        report_files=report_files,
        consensus_files=consensus_files,
        warning_file=warning_file if warning_file.is_file() else None,
    )


def build_low_quality_table_rows(
    positions: List[LowQualityPosition],
) -> List[List[str | int]]:
    return [
        [
            pos.consensus_pos,
            pos.forward_pos,
            pos.forward_base,
            pos.forward_qual,
            pos.reverse_pos,
            pos.reverse_base,
            pos.reverse_qual,
            pos.reason,
        ]
        for pos in positions
    ]


def build_single_strand_mixture_rows(
    calls,
    trimmed_read: Read,
) -> List[List[str | int | float]]:
    return [
        [
            call.read_position,
            (
                call.read_position - trimmed_read.start
                if trimmed_read.start < call.read_position <= trimmed_read.end
                else ""
            ),
            call.trace_index,
            call.primary_base,
            call.secondary_base,
            call.iupac_code,
            call.phred_quality,
            call.primary_height,
            call.secondary_height,
            f"{call.peak_ratio:.3f}",
            f"{call.local_noise:.3f}",
            f"{call.secondary_snr:.3f}",
            call.confidence,
        ]
        for call in calls
    ]


def summarize_non_acgt_symbols(consensus_line: list[str]) -> str:
    counts = Counter(
        base for base in consensus_line if base not in {"A", "C", "G", "T", "-"}
    )
    if not counts:
        return "none"
    return ", ".join(
        f"{symbol}={count}"
        for symbol, count in sorted(counts.items(), key=lambda item: (-item[1], item[0]))
    )


def build_single_mixture_map(
    rows: List[List[str | int | float]],
    use_single_mixture: str,
) -> dict[int, str]:
    mixture_map: dict[int, str] = {}
    allowed_confidence = {"high"} if use_single_mixture == "high" else {"high", "medium"}
    for row in rows:
        trimmed_position = row[1]
        iupac_code = row[5]
        confidence = row[12]
        if (
            isinstance(trimmed_position, int)
            and isinstance(iupac_code, str)
            and isinstance(confidence, str)
            and confidence in allowed_confidence
        ):
            mixture_map[trimmed_position] = iupac_code
    return mixture_map


def render_alignment_html(
    folder_name: str,
    alignment_payload: dict,
    reference_aa_alignment_payload: dict | None,
    forward_summary: str,
    reverse_summary: str,
    non_acgt_summary: str,
    alignment_parameters: list[tuple[str, str]],
    resolve_consensus_base_rules: list[str],
    low_quality_rows: List[List[str | int]],
    forward_mixture_rows: List[List[str | int | float]] | None = None,
    reverse_mixture_rows: List[List[str | int | float]] | None = None,
    forward_mixture_header_text: str | None = None,
    reverse_mixture_header_text: str | None = None,
    forward_mixture_name: str | None = None,
    reverse_mixture_name: str | None = None,
    single_mixture_parameters: list[tuple[str, str]] | None = None,
    banner_message: str | None = None,
) -> str:
    template = TEMPLATE_ENV.get_template("alignment_template.html")
    return template.render(
        folder_name=folder_name,
        banner_message=banner_message,
        forward_summary=forward_summary,
        reverse_summary=reverse_summary,
        non_acgt_summary=non_acgt_summary,
        overlap=alignment_payload["overlap"],
        alignment_data_json=alignment_payload["alignment_data_json"],
        reference_aa_alignment_payload=reference_aa_alignment_payload,
        alignment_parameters=alignment_parameters,
        resolve_consensus_base_rules=resolve_consensus_base_rules,
        low_quality_headers=[
            "ConsensusPos",
            "ForwardPos",
            "ForwardBase",
            "ForwardQ",
            "ReversePos",
            "ReverseBase",
            "ReverseQ",
            "Reason",
        ],
        low_quality_rows=low_quality_rows,
        mixture_headers=[
            "ReadPos",
            "TrimmedPos",
            "TraceIndex",
            "PrimaryBase",
            "SecondaryBase",
            "SuggestedIUPAC",
            "Phred",
            "PrimaryHeight",
            "SecondaryHeight",
            "PeakRatio",
            "LocalNoise",
            "SecondarySNR",
            "Confidence",
        ],
        forward_mixture_rows=forward_mixture_rows,
        reverse_mixture_rows=reverse_mixture_rows,
        forward_mixture_header_text=forward_mixture_header_text,
        reverse_mixture_header_text=reverse_mixture_header_text,
        forward_mixture_name=forward_mixture_name,
        reverse_mixture_name=reverse_mixture_name,
        single_mixture_parameters=single_mixture_parameters,
    )


def build_resolve_consensus_base_rules(
    forward_min_phred_score_per_base: int,
    reverse_min_phred_score_per_base: int,
    min_phred_score_for_paired_base: int,
) -> list[str]:
    template = TEMPLATE_ENV.get_template("resolve_consensus_base_rules.txt")
    return [
        line
        for line in template.render(
            forward_min_phred_score_per_base=forward_min_phred_score_per_base,
            reverse_min_phred_score_per_base=reverse_min_phred_score_per_base,
            min_phred_score_for_paired_base=min_phred_score_for_paired_base,
        ).splitlines()
        if line.strip()
    ]


def reverse_complement_trace(trace: AbiTrace) -> AbiTrace:
    complement = {"A": "T", "T": "A", "C": "G", "G": "C"}
    rc_channels = {
        complement[base]: list(reversed(values))
        for base, values in trace.channels.items()
    }
    trace_length = len(next(iter(trace.channels.values())))
    rc_base_positions = [
        trace_length - 1 - pos for pos in reversed(trace.base_positions)
    ]
    rc_seq = str(Seq(trace.seq).reverse_complement())
    return AbiTrace(
        name=trace.name + "_rc",
        seq=rc_seq,
        base_positions=rc_base_positions,
        channels=rc_channels,
        metadata=dict(trace.metadata),
        key_sections=list(trace.key_sections),
        hierarchy_text=trace.hierarchy_text,
    )


def find_ab1_pair(
    folder: Path,
    allow_single_strand: bool = False,
) -> Tuple[Path | None, Path | None]:
    ab1_files = [
        p for p in folder.iterdir() if p.is_file() and p.suffix.lower() == ".ab1"
    ]
    forward = [p for p in ab1_files if p.stem.endswith("F")]
    reverse = [p for p in ab1_files if p.stem.endswith("R")]

    if len(forward) == 1 and len(reverse) == 0:
        if allow_single_strand:
            return forward[0], None
        raise ValueError(
            f"Found only a forward .ab1 file ending with 'F' in {folder}; missing reverse .ab1 file ending with 'R'."
        )
    if len(forward) == 0 and len(reverse) == 1:
        if allow_single_strand:
            return None, reverse[0]
        raise ValueError(
            f"Found only a reverse .ab1 file ending with 'R' in {folder}; missing forward .ab1 file ending with 'F'."
        )
    if len(forward) != 1 or len(reverse) != 1:
        raise ValueError(
            "Expected exactly one forward .ab1 file ending with 'F' and one reverse .ab1 file ending with 'R' "
            f"in {folder}."
        )

    return forward[0], reverse[0]


def generate_consensus_outputs(
    folder: Path,
    folder_name: str,
    fwd_trim: Read,
    rev_rc_trim: Read,
    fwd: Read,
    rev_rc_full: Read,
    fwd_trace: AbiTrace,
    rev_rc_trace: AbiTrace,
    fwd_path: Path,
    rev_path: Path,
    min_overlap: int,
    use_paired_mixture: bool,
    use_single_mixture: str,
    detect_single_strand_mixture: bool,
    min_phred_score_per_base: tuple[int, int],
    min_phred_score_for_paired_base: int,
    min_peak_ratio: float,
    min_secondary_snr: float,
    noise_window_radius: int,
    noise_window_exclude_radius: int,
    high_phred_quality: int,
    high_peak_ratio: float,
    high_secondary_snr: float,
    reference_nucleotide_fasta: Path | None = None,
    banner_message: str | None = None,
) -> Tuple[Path, Path, Path, Path]:
    (
        forward_min_phred_score_per_base,
        reverse_min_phred_score_per_base,
    ) = min_phred_score_per_base
    forward_out = folder / "forward_trimmed.fasta"
    reverse_out = folder / "reverse_rc_trimmed.fasta"
    consensus_out = folder / f"{folder_name}.fasta"
    amino_acid_alignment_out = folder / f"{folder_name}_aa_alignment.fasta"
    alignment_html_out = folder / f"{folder_name}_alignment.html"
    forward_mixture_rows = None
    reverse_mixture_rows = None
    forward_mixture_header_text = None
    reverse_mixture_header_text = None
    forward_mixture_name = None
    reverse_mixture_name = None
    single_mixture_parameters = None
    reference_aa_alignment_payload = None
    forward_single_mixture_map = None
    reverse_single_mixture_map = None
    should_detect_single_mixture = detect_single_strand_mixture or bool(use_single_mixture)

    if should_detect_single_mixture:
        forward_mixture_calls = detect_single_strand_mixtures(
            seq=fwd_trace.seq,
            qual=fwd.qual,
            base_positions=fwd_trace.base_positions,
            channels=fwd_trace.channels,
            min_phred_score_per_base=forward_min_phred_score_per_base,
            min_peak_ratio=min_peak_ratio,
            min_secondary_snr=min_secondary_snr,
            noise_window_radius=noise_window_radius,
            noise_window_exclude_radius=noise_window_exclude_radius,
            high_phred_quality=high_phred_quality,
            high_peak_ratio=high_peak_ratio,
            high_secondary_snr=high_secondary_snr,
        )
        reverse_rc_mixture_calls = detect_single_strand_mixtures(
            seq=rev_rc_trace.seq,
            qual=rev_rc_full.qual,
            base_positions=rev_rc_trace.base_positions,
            channels=rev_rc_trace.channels,
            min_phred_score_per_base=reverse_min_phred_score_per_base,
            min_peak_ratio=min_peak_ratio,
            min_secondary_snr=min_secondary_snr,
            noise_window_radius=noise_window_radius,
            noise_window_exclude_radius=noise_window_exclude_radius,
            high_phred_quality=high_phred_quality,
            high_peak_ratio=high_peak_ratio,
            high_secondary_snr=high_secondary_snr,
        )
        forward_mixture_rows = build_single_strand_mixture_rows(
            forward_mixture_calls,
            fwd_trim,
        )
        reverse_mixture_rows = build_single_strand_mixture_rows(
            reverse_rc_mixture_calls,
            rev_rc_trim,
        )
        if use_single_mixture:
            forward_single_mixture_map = build_single_mixture_map(
                forward_mixture_rows,
                use_single_mixture,
            )
            reverse_single_mixture_map = build_single_mixture_map(
                reverse_mixture_rows,
                use_single_mixture,
            )
        forward_mixture_header_text = "\n".join(
            build_single_strand_mixtures_header(fwd_trace.name)[:3]
        )
        reverse_mixture_header_text = "\n".join(
            build_single_strand_mixtures_header(rev_rc_trace.name)[:3]
        )
        forward_mixture_name = f"{fwd_path.stem}_candidate_mixtures"
        reverse_mixture_name = f"{rev_path.stem}_rev_rc_candidate_mixtures"
        single_mixture_parameters = [
            ("forward_min_phred_score", str(forward_min_phred_score_per_base)),
            ("reverse_min_phred_score", str(reverse_min_phred_score_per_base)),
            ("min_peak_ratio", str(min_peak_ratio)),
            ("min_secondary_snr", str(min_secondary_snr)),
            ("noise_window_radius", str(noise_window_radius)),
            (
                "noise_window_exclude_radius",
                str(noise_window_exclude_radius),
            ),
            ("high_phred_quality", str(high_phred_quality)),
            ("high_peak_ratio", str(high_peak_ratio)),
            ("high_secondary_snr", str(high_secondary_snr)),
        ]

    write_fasta(str(forward_out), fwd_trim.name, fwd_trim.seq)
    write_fasta(str(reverse_out), rev_rc_trim.name, rev_rc_trim.seq)

    consensus, alignment_payload, low_quality_positions = assemble_pair(
        fwd_trim,
        rev_rc_trim,
        forward_min_phred_score_per_base=forward_min_phred_score_per_base,
        reverse_min_phred_score_per_base=reverse_min_phred_score_per_base,
        min_phred_score_for_paired_base=min_phred_score_for_paired_base,
        use_paired_mixture=use_paired_mixture,
        use_single_mixture=use_single_mixture,
        forward_single_mixture_map=forward_single_mixture_map,
        reverse_single_mixture_map=reverse_single_mixture_map,
        min_overlap=min_overlap,
    )

    low_quality_rows = build_low_quality_table_rows(low_quality_positions)
    forward_summary = (
        f"original positions {fwd.start + 1}-{fwd.end}, "
        f"original length {len(fwd.seq)}, "
        f"trimmed positions {fwd_trim.start + 1}-{fwd_trim.end}, "
        f"trimmed length {len(fwd_trim.seq)}"
    )
    reverse_summary = (
        f"original positions {rev_rc_full.start + 1}-{rev_rc_full.end}, "
        f"original length {len(rev_rc_full.seq)}, "
        f"trimmed positions {rev_rc_trim.start + 1}-{rev_rc_trim.end}, "
        f"trimmed length {len(rev_rc_trim.seq)}"
    )
    non_acgt_summary = summarize_non_acgt_symbols(
        json.loads(alignment_payload["alignment_data_json"])["consensus_line"]
    )
    alignment_parameters = [
        ("use_paired_mixture", str(use_paired_mixture)),
        ("min_overlap", str(min_overlap)),
        (
            "min_phred_score_per_base",
            format_min_phred_score_per_base(min_phred_score_per_base),
        ),
        ("min_phred_score_for_paired_base", str(min_phred_score_for_paired_base)),
        ("use_single_mixture", use_single_mixture or ""),
        ("detect_single_strand_mixture", str(detect_single_strand_mixture)),
        (
            "reference_nucleotide_fasta",
            (
                read_reference_nucleotide_fasta(reference_nucleotide_fasta)[0]
                if reference_nucleotide_fasta
                else ""
            ),
        ),
        ("match_score", "2.0"),
        ("mismatch_score", "-1.0"),
        ("gap_open_score", "-5.0"),
        ("gap_extend_score", "-1.0"),
    ]
    resolve_consensus_base_rules = build_resolve_consensus_base_rules(
        forward_min_phred_score_per_base=forward_min_phred_score_per_base,
        reverse_min_phred_score_per_base=reverse_min_phred_score_per_base,
        min_phred_score_for_paired_base=min_phred_score_for_paired_base,
    )
    if not consensus:
        return forward_out, reverse_out, consensus_out, alignment_html_out
    write_fasta(str(consensus_out), folder_name, consensus)
    if reference_nucleotide_fasta is not None:
        aa_alignment = align_consensus_to_reference_nucleotide(
            consensus_nt=consensus,
            sample_name=folder_name,
            reference_nt_path=reference_nucleotide_fasta,
        )
        write_amino_acid_alignment_fasta(amino_acid_alignment_out, aa_alignment)
        reference_aa_alignment_payload = {
            "reference_header": aa_alignment.reference_header,
            "consensus_header": aa_alignment.consensus_header,
            "strand": aa_alignment.strand,
            "strand_label": (
                "consensus"
                if aa_alignment.strand == "forward"
                else "reverse complement"
            ),
            "frame": aa_alignment.frame,
            "score": f"{aa_alignment.score:.2f}",
            "translated_length": aa_alignment.translated_length,
            "alignment_data_json": json.dumps(
                {
                    "reference_aligned": list(aa_alignment.reference_aligned),
                    "reference_na_aligned": aa_alignment.reference_na_aligned,
                    "consensus_aligned": list(aa_alignment.consensus_aligned),
                    "consensus_na_aligned": aa_alignment.consensus_na_aligned,
                }
            ),
        }
    alignment_html = render_alignment_html(
        folder_name=folder_name,
        alignment_payload=alignment_payload,
        reference_aa_alignment_payload=reference_aa_alignment_payload,
        banner_message=banner_message,
        forward_summary=forward_summary,
        reverse_summary=reverse_summary,
        non_acgt_summary=non_acgt_summary,
        alignment_parameters=alignment_parameters,
        resolve_consensus_base_rules=resolve_consensus_base_rules,
        low_quality_rows=low_quality_rows,
        forward_mixture_rows=forward_mixture_rows,
        reverse_mixture_rows=reverse_mixture_rows,
        forward_mixture_header_text=forward_mixture_header_text,
        reverse_mixture_header_text=reverse_mixture_header_text,
        forward_mixture_name=forward_mixture_name,
        reverse_mixture_name=reverse_mixture_name,
        single_mixture_parameters=single_mixture_parameters,
    )
    write_text(str(alignment_html_out), alignment_html)

    return forward_out, reverse_out, consensus_out, alignment_html_out


def run_folder_assembly(
    folder: Path,
    use_paired_mixture: bool = True,
    use_single_mixture: str = "",
    allow_single_strand: bool = False,
    min_phred_score_per_base: str | tuple[int, int] = "20:20",
    min_phred_score_for_paired_base: int = 10,
    min_consecutive_high_quality_bases: int = 10,
    min_overlap: int = 40,
    detect_single_strand_mixture: bool = True,
    min_peak_ratio: float = 0.25,
    min_secondary_snr: float = 5.0,
    noise_window_radius: int = 20,
    noise_window_exclude_radius: int = 4,
    high_phred_quality: int = 30,
    high_peak_ratio: float = 0.33,
    high_secondary_snr: float = 8.0,
    reference_nucleotide_fasta: Path | None = None,
) -> None:
    folder = Path(folder).resolve()
    if not folder.is_dir():
        raise ValueError(f"Input must be a folder: {folder}")
    if reference_nucleotide_fasta is not None:
        reference_nucleotide_fasta = Path(reference_nucleotide_fasta).expanduser().resolve()
        if not reference_nucleotide_fasta.is_file():
            raise ValueError(
                f"Reference nucleotide FASTA is not a file: {reference_nucleotide_fasta}"
            )
    (
        forward_min_phred_score_per_base,
        reverse_min_phred_score_per_base,
    ) = parse_min_phred_score_per_base(min_phred_score_per_base)

    try:
        fwd_path, rev_path = find_ab1_pair(
            folder,
            allow_single_strand=allow_single_strand,
        )
    except ValueError as exc:
        write_warning_html(folder, str(exc))
        return

    banner_message = None
    if allow_single_strand and (fwd_path is None or rev_path is None):
        min_phred_score_for_paired_base = 20
        banner_message = (
            "Single-strand mode is active. The available read was mirrored into the missing strand for QA and consensus generation."
        )

    if fwd_path is not None and rev_path is not None:
        fwd = read_ab1(str(fwd_path), "forward")
        rev = read_ab1(str(rev_path), "reverse")
        fwd_trace = read_ab1_trace(str(fwd_path), "forward")
        rev_trace = read_ab1_trace(str(rev_path), "reverse")
        rev_rc_full = revcomp_read(rev)
        rev_rc_trace = reverse_complement_trace(rev_trace)
    elif fwd_path is not None:
        fwd = read_ab1(str(fwd_path), "forward")
        fwd_trace = read_ab1_trace(str(fwd_path), "forward")
        # Mirror the available forward read into reverse orientation so downstream logic still sees a pair.
        rev = revcomp_read(fwd)
        rev.name = "reverse"
        rev_trace = reverse_complement_trace(fwd_trace)
        rev_trace.name = "reverse"
        rev_rc_full = fwd
        rev_rc_trace = fwd_trace
        rev_path = fwd_path
    else:
        rev = read_ab1(str(rev_path), "reverse")
        rev_trace = read_ab1_trace(str(rev_path), "reverse")
        rev_rc_full = revcomp_read(rev)
        rev_rc_trace = reverse_complement_trace(rev_trace)
        # Mirror the available reverse read into forward orientation so downstream logic still sees a pair.
        fwd = rev_rc_full
        fwd.name = "forward"
        fwd_trace = rev_rc_trace
        fwd_trace.name = "forward"
        fwd_path = rev_path

    fwd_trim = trim_by_quality(
        fwd,
        min_phred_score_per_base=forward_min_phred_score_per_base,
        min_consecutive_high_quality_bases=min_consecutive_high_quality_bases,
    )
    rev_trim = trim_by_quality(
        rev,
        min_phred_score_per_base=reverse_min_phred_score_per_base,
        min_consecutive_high_quality_bases=min_consecutive_high_quality_bases,
    )

    if not run_quality_checks(
        folder_name=folder.name,
        output_folder=folder,
        fwd_path=fwd_path,
        rev_path=rev_path,
        fwd=fwd,
        rev=rev,
        rev_rc_full=rev_rc_full,
        fwd_trace=fwd_trace,
        rev_trace=rev_trace,
        rev_rc_trace=rev_rc_trace,
        fwd_trim=fwd_trim,
        rev_trim=rev_trim,
        forward_min_phred_score_per_base=forward_min_phred_score_per_base,
        reverse_min_phred_score_per_base=reverse_min_phred_score_per_base,
        min_consecutive_high_quality_bases=min_consecutive_high_quality_bases,
        min_peak_ratio=min_peak_ratio,
        min_secondary_snr=min_secondary_snr,
        noise_window_radius=noise_window_radius,
        noise_window_exclude_radius=noise_window_exclude_radius,
        high_phred_quality=high_phred_quality,
        high_peak_ratio=high_peak_ratio,
        high_secondary_snr=high_secondary_snr,
        banner_message=banner_message,
    ):
        return

    rev_rc_trim = revcomp_read(rev_trim)

    try:
        generate_consensus_outputs(
            folder=folder,
            folder_name=folder.name,
            fwd_trim=fwd_trim,
            rev_rc_trim=rev_rc_trim,
            fwd=fwd,
            rev_rc_full=rev_rc_full,
            fwd_trace=fwd_trace,
            rev_rc_trace=rev_rc_trace,
            fwd_path=fwd_path,
            rev_path=rev_path,
            min_overlap=min_overlap,
            use_paired_mixture=use_paired_mixture,
            use_single_mixture=use_single_mixture,
            detect_single_strand_mixture=detect_single_strand_mixture,
            min_phred_score_per_base=(
                forward_min_phred_score_per_base,
                reverse_min_phred_score_per_base,
            ),
            min_phred_score_for_paired_base=min_phred_score_for_paired_base,
            min_peak_ratio=min_peak_ratio,
            min_secondary_snr=min_secondary_snr,
            noise_window_radius=noise_window_radius,
            noise_window_exclude_radius=noise_window_exclude_radius,
            high_phred_quality=high_phred_quality,
            high_peak_ratio=high_peak_ratio,
            high_secondary_snr=high_secondary_snr,
            reference_nucleotide_fasta=reference_nucleotide_fasta,
            banner_message=banner_message,
        )
    except ValueError as exc:
        write_warning_html(folder, str(exc), banner_message=banner_message)
        return


def assembly():
    parser = argparse.ArgumentParser(
        description="Assemble paired AB1 reads from a folder."
    )
    parser.add_argument(
        "folder", help="Folder containing one *F.ab1 and one *R.ab1 file"
    )
    parser.add_argument(
        "--clean",
        action="store_true",
        help="Remove all non-.ab1 files in the sample folder before assembly",
    )
    parser.add_argument(
        "--use-paired-mixture",
        action=argparse.BooleanOptionalAction,
        default=True,
        help="Allow IUPAC mixture codes when both strands disagree but both are above Q20 (default: enabled)",
    )
    parser.add_argument(
        "--use-single-mixture",
        choices=["high", "medium"],
        default="",
        help="Apply single-strand candidate mixtures back onto the trimmed reads before consensus calling: high=use only high confidence, medium=use both high and medium",
    )
    parser.add_argument(
        "--allow-single-strand",
        action="store_true",
        help="Allow a single forward-only or reverse-only .ab1 file and mirror it into the missing strand",
    )
    parser.add_argument(
        "--min-phred-score-per-base",
        default="20:20",
        help="Minimum per-base Phred score in forward_phred:reverse_phred format (default: 20:20)",
    )
    parser.add_argument(
        "--min-phred-score-for-paired-base",
        type=int,
        default=10,
        help="Minimum Phred score accepted when both strands agree on the same base (default: 10)",
    )
    parser.add_argument(
        "--min-consecutive-high-quality-bases",
        type=int,
        default=10,
        help="Minimum consecutive high-quality bases used by trim_by_quality (default: 10)",
    )
    parser.add_argument(
        "--min-overlap",
        type=int,
        default=40,
        help="Minimum overlap required after alignment (default: 40)",
    )
    parser.add_argument(
        "--detect-single-strand-mixture",
        action="store_true",
        default=True,
        help="Enable single-strand candidate mixture detection reports (default: enabled)",
    )
    parser.add_argument(
        "--reference-nucleotide-fasta",
        type=Path,
        help="Optional FASTA file with one nucleotide reference sequence used for codon-aware reference alignment and writing <folder_name>_aa_alignment.fasta",
    )
    args = parser.parse_args()

    try:
        folder = Path(args.folder).expanduser().resolve()
        if not folder.is_dir():
            raise ValueError(f"Input must be a folder: {folder}")
        if args.clean:
            cleanup_folder_outputs(folder)
            return
        run_folder_assembly(
            folder=folder,
            use_paired_mixture=args.use_paired_mixture,
            use_single_mixture=args.use_single_mixture,
            allow_single_strand=args.allow_single_strand,
            min_phred_score_per_base=args.min_phred_score_per_base,
            min_phred_score_for_paired_base=args.min_phred_score_for_paired_base,
            min_consecutive_high_quality_bases=args.min_consecutive_high_quality_bases,
            min_overlap=args.min_overlap,
            detect_single_strand_mixture=args.detect_single_strand_mixture,
            reference_nucleotide_fasta=args.reference_nucleotide_fasta,
        )
    except ValueError as exc:
        write_warning_html(folder, str(exc))
        return
        print(str(exc), file=sys.stderr)
        return


if __name__ == "__main__":
    assembly()