BioSight/python/phi/output.py

"""
phi.output — FASTQ, Adleman graph, and PCR primer protocol

Transforms Φ encoding records (from phi.embed.encode_phi) into
standard bioinformatics formats consumable by external tools.

Formats:
  FASTQ          — pipe through Jellyfish/KMC/BLAST
  Adleman graph  — feed into BioComputing's connectNodes()
  PCR protocol   — wet-lab filtering with allele-specific primers

Dependencies: phi.embed (for encode_phi output dict format).
"""

from __future__ import annotations

import math
from typing import Dict, List

# ── Primer sequences ─────────────────────────────────────────────────────

PRIMER_DESIGN = {
    "pas_primer": {
        "sequence": "CCCCCC",
        "target": "consistency_dna = GGGGGG (all 6 rules pass)",
        "tm": 36.0,
        "purpose": "Capture PASS strands — amplify consistent equations only",
    },
    "fail_primer": {
        "sequence": "AAAAAA",
        "target": "consistency_dna = TTTTTT (all 6 rules fail)",
        "tm": 12.0,
        "purpose": "Capture FAIL strands — quarantine inconsistent equations",
    },
    "universal_primer": {
        "sequence": "ABCABC",
        "target": "F_dna region (first 8 bases)",
        "tm": 15.0,
        "purpose": "Amplify all encoded equations regardless of consistency",
    },
}


# ── FASTQ ────────────────────────────────────────────────────────────────

def to_fastq(record: Dict) -> str:
    """Format a Φ encoding record as a 4-line FASTQ entry.

    FASTQ is the standard DNA sequencing format.  This lets us pipe
    equation encodings through any existing bioinformatics pipeline
    without tool modification.

    Example:
        >>> record = {
        ...     "equation": "E(x) = x^2",
        ...     "dna_sequence": "ATCG",
        ...     "quality_scores": "IIII",
        ...     "consistency": {"balanced_parens": True, "valid_operator_order": True,
        ...                     "valid_variable_name": True, "no_empty_expression": True,
        ...                     "single_expression": True, "defined_reference": True},
        ...     "sha256_prefix": "abc123",
        ...     "consistency_pass": True,
        ... }
        >>> fas = to_fastq(record)
        >>> fas.startswith("@")
        True
        >>> len(fas.splitlines())
        4
    """
    eq = record["equation"]
    seq = record["dna_sequence"]
    qual = record["quality_scores"]
    n_cons = sum(1 for v in record["consistency"].values() if v)
    eq_hash = record["sha256_prefix"]
    pass_flag = "PASS" if record["consistency_pass"] else "FAIL"

    header = f"@{eq_hash} n_consistency={n_cons}/6 pass={pass_flag} len={len(eq)}"
    qual_line = qual * (len(seq) // len(qual)) if len(seq) > len(qual) else qual

    return f"{header}\n{seq}\n+\n{qual_line}\n"


# ── Adleman integration ──────────────────────────────────────────────────

def to_adleman_graph(records: List[Dict]) -> Dict:
    """Build an Adleman/Lipton-compatible graph from Φ-encoded equations.

    Each equation is a vertex.  Edges represent Fisher-distance similarity
    (d_F < 0.5) between their F(E) byte-class histograms.

    Feed the output into BioComputing's `hampath.connectNodes()` or
    `sattv.createNodes()` to generate wet-lab DNA sequences.

    Example:
        >>> records = [
        ...     {"equation": "a=b", "dna_sequence": "ATCG", "consistency_pass": True,
        ...      "F": [0.125]*8},
        ...     {"equation": "c=d", "dna_sequence": "GCTA", "consistency_pass": False,
        ...      "F": [1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]},
        ...     {"equation": "e=f", "dna_sequence": "TATA", "consistency_pass": True,
        ...      "F": [0.125]*8},
        ... ]
        >>> graph = to_adleman_graph(records)
        >>> graph["n_vertices"]
        3
        >>> graph["n_edges"]
        1
    """
    vertices = {}
    for i, rec in enumerate(records):
        vid = f"E{i}"
        vertices[vid] = {
            "equation": rec["equation"],
            "dna": rec["dna_sequence"],
            "consistency": rec["consistency_pass"],
        }

    edges = []
    vert_ids = list(vertices.keys())
    for i in range(len(vert_ids)):
        for j in range(i + 1, len(vert_ids)):
            Fi = records[i]["F"]
            Fj = records[j]["F"]
            dist = _fisher_distance(Fi, Fj)
            if dist < 0.5:
                edges.append((vert_ids[i], vert_ids[j], round(dist, 4)))

    return {
        "vertices": vertices,
        "edges": edges,
        "format": "adleman_graph_v1",
        "n_vertices": len(vertices),
        "n_edges": len(edges),
    }


# ── PCR protocol ─────────────────────────────────────────────────────────

def design_filtering_protocol(records: List[Dict]) -> Dict:
    """Design a PCR filtering protocol for a batch of encoded equations.

    Uses allele-specific PCR (Newton et al. 1989) — the 24°C Tm gap
    between PASS and FAIL primers provides single-base discrimination.

    Example:
        >>> records = [
        ...     {"equation": "x=1", "consistency_pass": True},
        ...     {"equation": "y=2", "consistency_pass": False},
        ...     {"equation": "z=3", "consistency_pass": True},
        ... ]
        >>> proto = design_filtering_protocol(records)
        >>> proto["pas_primer"]
        'CCCCCC'
        >>> proto["fail_primer"]
        'AAAAAA'
        >>> proto["universal_primer"]
        'ABCABC'
    """
    n_pass = sum(1 for r in records if r["consistency_pass"])
    n_fail = len(records) - n_pass

    return {
        "protocol": "allele_specific_pcr_v1",
        "reference": "Newton et al. (1989) Nucleic Acids Res 17:2503; "
                      "allele-specific PCR for single-base discrimination",
        "pas_primer": PRIMER_DESIGN["pas_primer"]["sequence"],
        "fail_primer": PRIMER_DESIGN["fail_primer"]["sequence"],
        "universal_primer": PRIMER_DESIGN["universal_primer"]["sequence"],
        "annealing_tm_pass": PRIMER_DESIGN["pas_primer"]["tm"],
        "annealing_tm_fail": PRIMER_DESIGN["fail_primer"]["tm"],
        "annealing_tm_universal": PRIMER_DESIGN["universal_primer"]["tm"],
        "expected_amplification_pass": n_pass,
        "expected_amplification_fail": n_fail,
    }


# ── Fisher distance ↓ (internal; no circular dep) ────────────────────────

def _fisher_distance(p: List[float], q: List[float]) -> float:
    """Fisher-Rao distance between two points on the probability simplex.

    The Fisher-Rao distance is the geodesic distance on the statistical
    manifold of multinomial distributions:

        d_F(p, q) = 2 · arccos(Σ √(p_i · q_i))

    where p, q ∈ Δ_{n-1} (the (n-1)-simplex of n-dimensional probability
    vectors).  The result is in [0, π].

    Args:
        p: First probability vector (values must be non-negative and
           sum to 1.0).
        q: Second probability vector (same length as p).

    Returns:
        Fisher-Rao distance in [0, π].

    Example:
        >>> _fisher_distance([1.0, 0.0, 0.0], [0.0, 1.0, 0.0])  # orthogonal
        3.141592653589793
        >>> _fisher_distance([0.5, 0.5], [0.5, 0.5])  # identical
        0.0
    """
    dot = sum(math.sqrt(pi * qi) for pi, qi in zip(p, q))
    dot = max(-1.0, min(1.0, dot))
    return 2.0 * math.acos(dot)


# ── Verification (python3 -m phi.output) ──────────────────────────────────

if __name__ == "__main__":
    import sys

    fail = 0

    # 1. _fisher_distance
    uni = [1 / 12] * 12
    if _fisher_distance(uni, uni) != 0.0:
        fail += 1; print("FAIL: _fisher_distance(p,p) != 0")
    else:
        print("  ✓ _fisher_distance(p,p) = 0")

    p_orth = [1.0] + [0.0] * 11
    q_orth = [0.0, 1.0] + [0.0] * 10
    d_orth = _fisher_distance(p_orth, q_orth)
    if abs(d_orth - math.pi) > 1e-5:
        fail += 1
        print(f"FAIL: _fisher_distance(orthogonal) = {d_orth}, expected π")
    else:
        print("  ✓ _fisher_distance(orthogonal) ≈ 3.14159")

    # 2. to_fastq
    rec_fq = {
        "equation": "a = b",
        "dna_sequence": "ABCGPSTZABCGPSTZABCGPSTZGGGGGG",
        "quality_scores": "IIIIII",
        "consistency": {"balanced_parens": True, "valid_operator_order": True,
                        "valid_variable_name": True, "no_empty_expression": True,
                        "single_expression": True, "defined_reference": True},
        "sha256_prefix": "abc123def456",
        "consistency_pass": True,
    }
    fas = to_fastq(rec_fq)
    if not fas.startswith("@"):
        fail += 1; print("FAIL: to_fastq doesn't start with @")
    else:
        print("  ✓ to_fastq starts with @")
    if len(fas.splitlines()) != 4:
        fail += 1
        print(f"FAIL: to_fastq has {len(fas.splitlines())} lines, expected 4")
    else:
        print("  ✓ to_fastq has 4 lines")

    # 3. to_adleman_graph
    recs_ag = [
        {"equation": "x = 1", "dna_sequence": "A" * 30, "consistency_pass": True,
         "F": [1 / 12] * 12},
        {"equation": "y = 2", "dna_sequence": "G" * 30, "consistency_pass": False,
         "F": [1.0] + [0.0] * 11},
        {"equation": "z = 3", "dna_sequence": "C" * 30, "consistency_pass": True,
         "F": [1 / 12] * 12},
    ]
    graph = to_adleman_graph(recs_ag)
    if graph["n_vertices"] != 3:
        fail += 1
        print(f"FAIL: n_vertices = {graph['n_vertices']}, expected 3")
    else:
        print("  ✓ to_adleman_graph n_vertices = 3")
    if graph["n_edges"] != 1:
        fail += 1
        print(f"FAIL: n_edges = {graph['n_edges']}, expected 1 (E0↔E2)")
    else:
        print("  ✓ to_adleman_graph n_edges = 1 (E0↔E2 via identical F)")

    # 4. design_filtering_protocol
    recs_pcr = [
        {"equation": "x = 1", "consistency_pass": True},
        {"equation": "y = 2", "consistency_pass": False},
        {"equation": "z = 3", "consistency_pass": True},
    ]
    proto = design_filtering_protocol(recs_pcr)
    for key in ["pas_primer", "fail_primer", "universal_primer"]:
        if key not in proto:
            fail += 1; print(f"FAIL: protocol missing {key}")
        else:
            print(f"  ✓ protocol has {key} = {proto[key]}")
    if proto["pas_primer"] != "CCCCCC":
        fail += 1; print(f"FAIL: pas_primer = {proto['pas_primer']}")
    else:
        print("  ✓ pas_primer sequence correct")
    if proto["fail_primer"] != "AAAAAA":
        fail += 1; print(f"FAIL: fail_primer = {proto['fail_primer']}")
    else:
        print("  ✓ fail_primer sequence correct")
    if proto["universal_primer"] != "ABCABC":
        fail += 1; print(f"FAIL: universal_primer = {proto['universal_primer']}")
    else:
        print("  ✓ universal_primer sequence correct")

    # 5. PRIMER_DESIGN structure
    if len(PRIMER_DESIGN) != 3:
        fail += 1
        print(f"FAIL: PRIMER_DESIGN has {len(PRIMER_DESIGN)} entries, expected 3")
    else:
        print("  ✓ PRIMER_DESIGN has exactly 3 entries")
    tm_checks = [("pas_primer", 36.0), ("fail_primer", 12.0), ("universal_primer", 15.0)]
    for name, expected in tm_checks:
        actual = PRIMER_DESIGN[name]["tm"]
        if actual != expected:
            fail += 1
            print(f"FAIL: {name} Tm = {actual}, expected {expected}")
        else:
            print(f"  ✓ {name} Tm = {expected}")

    print(f"\nVerdict: {fail} failure(s)")
    sys.exit(fail)