#!/usr/bin/env python3
"""
Tool: pwm_motif_promoter_scanner.py
Category: Regulatory Genomics & Transcription Factors
Description: Converts aligned transcription factor binding sites into Position Weight Matrices (PWM),
             calculates log-odds scores against background nucleotide distributions, and scans
             promoter sequences for significant regulatory binding sites.
Usage:
    python pwm_motif_promoter_scanner.py -m motifs.txt -s promoter.fasta [--threshold 5.0]
"""

import argparse
import math
import sys

def build_pwm(motif_instances, bg=0.25):
    """Constructs log2-odds PWM from list of aligned nucleotide strings."""
    k = len(motif_instances[0])
    counts = {b: [1.0] * k for b in "ACGT"} # Laplace pseudocount = 1
    
    for seq in motif_instances:
        for i, b in enumerate(seq.upper()):
            if b in counts:
                counts[b][i] += 1
                
    total_seqs = len(motif_instances) + 4
    pwm = {}
    for b in "ACGT":
        pwm[b] = [math.log2((counts[b][i] / total_seqs) / bg) for i in range(k)]
        
    return pwm, k

def scan_sequence(seq, pwm, k, threshold=4.0):
    """Scans sequence with sliding window of length k and returns hits above threshold."""
    hits = []
    seq_up = seq.upper()
    for i in range(len(seq_up) - k + 1):
        window = seq_up[i:i+k]
        score = 0.0
        valid = True
        for col, b in enumerate(window):
            if b in pwm:
                score += pwm[b][col]
            else:
                valid = False
                break
        if valid and score >= threshold:
            hits.append((i + 1, window, round(score, 2)))
    return hits

def main():
    print("=== Position Weight Matrix (PWM) Regulatory Motif Scanner ===")
    
    # Example: Human TATA-binding protein (TBP) consensus motifs
    training_sites = [
        "TATAAAAG",
        "TATAAATA",
        "TATAAAAG",
        "TATAATAG",
        "TATAAACG",
        "TATAAAAG"
    ]
    
    pwm, k = build_pwm(training_sites)
    print(f"\n[1] Built PWM for motif length k={k} from {len(training_sites)} verified sites.")
    print(f"    - Base 'T' pos 1 weight : {pwm['T'][0]:.2f}")
    print(f"    - Base 'A' pos 2 weight : {pwm['A'][1]:.2f}")

    promoter = "CGGGCTAGCTAGCTATAAAAGCCCCCGTATAAATAGCCGATCGATC"
    print(f"\n[2] Scanning Promoter Sequence ({len(promoter)} bp):\n    {promoter}")
    
    hits = scan_sequence(promoter, pwm, k, threshold=3.5)
    print(f"\nDetected Motif Matches (Score >= 3.5):")
    for pos, matched_seq, score in hits:
        print(f"    -> Pos {pos:02d}: {matched_seq} | Log-Odds Score: {score}")

if __name__ == "__main__":
    main()