#!/usr/bin/env python3
"""
Chapter 9: Python for Next Gen Biologist
Script: ch09_numpy_pandas_viz.py
Description: Scientific computing with NumPy, Pandas, and Matplotlib:
             Position Weight Matrix (PWM) log-odds scoring, RNA-seq counts normalization,
             fold-change filtering, and Volcano plot generation.
Usage:
    python ch09_numpy_pandas_viz.py [--output volcano.png]
"""

import argparse
import math
import numpy as np
import pandas as pd

def build_pwm_log_odds(aligned_motifs: list, background: float = 0.25) -> np.ndarray:
    """
    Builds a Position Weight Matrix (PWM) in log2-odds space from aligned motif strings.
    Matrix shape: (4, motif_length) corresponding to rows A, C, G, T.
    """
    length = len(aligned_motifs[0])
    base_indices = {'A': 0, 'C': 1, 'G': 2, 'T': 3}
    counts = np.ones((4, length))  # Laplace pseudocount = 1
    
    for seq in aligned_motifs:
        for col, base in enumerate(seq.upper()):
            if base in base_indices:
                counts[base_indices[base], col] += 1
                
    frequencies = counts / counts.sum(axis=0)
    log_odds = np.log2(frequencies / background)
    return log_odds

def score_sequence_with_pwm(sequence: str, pwm: np.ndarray) -> list:
    """Scans sequence with sliding window and returns list of (pos, score)."""
    base_indices = {'A': 0, 'C': 1, 'G': 2, 'T': 3}
    motif_len = pwm.shape[1]
    seq = sequence.upper()
    scores = []
    
    for i in range(len(seq) - motif_len + 1):
        window = seq[i:i+motif_len]
        score = 0.0
        valid = True
        for col, b in enumerate(window):
            if b in base_indices:
                score += pwm[base_indices[b], col]
            else:
                valid = False
                break
        if valid:
            scores.append((i + 1, round(score, 3)))
    return scores

def simulate_rnaseq_dataframe(num_genes: int = 100) -> pd.DataFrame:
    """Generates synthetic RNA-seq counts DataFrame for demonstration."""
    np.random.seed(42)
    genes = [f"GENE_{i:04d}" for i in range(1, num_genes + 1)]
    
    # Simulate log2 fold changes and p-values
    log2fc = np.random.normal(loc=0.0, scale=1.5, size=num_genes)
    # Give a few genes strong upregulation/downregulation
    log2fc[:5] += 3.5
    log2fc[5:10] -= 3.5
    
    pvalues = np.random.beta(a=0.5, b=2.0, size=num_genes)
    pvalues[:10] = pvalues[:10] * 1e-4  # strong significance
    
    neg_log10_p = -np.log10(np.clip(pvalues, 1e-15, 1.0))
    
    df = pd.DataFrame({
        'GeneID': genes,
        'log2FoldChange': np.round(log2fc, 3),
        'pvalue': pvalues,
        'negLog10P': np.round(neg_log10_p, 3)
    })
    
    df['Status'] = 'Not Significant'
    df.loc[(df['log2FoldChange'] > 1.5) & (df['pvalue'] < 0.05), 'Status'] = 'Upregulated'
    df.loc[(df['log2FoldChange'] < -1.5) & (df['pvalue'] < 0.05), 'Status'] = 'Downregulated'
    return df

def main():
    parser = argparse.ArgumentParser(description="Chapter 9: NumPy & Pandas Bioinformatics Demo.")
    parser.add_argument("-o", "--output", help="Optional path to save volcano plot")
    args = parser.parse_args()

    print("=== Chapter 9: Scientific Computing with NumPy & Pandas ===")

    # 1. NumPy PWM
    motifs = [
        "TATAAA", "TATATA", "TATAAA", "TATAAT", "CATAAA", "TATAAA"
    ]
    pwm = build_pwm_log_odds(motifs)
    print(f"\n[1] Generated PWM Log-Odds Matrix ({pwm.shape[0]} bases x {pwm.shape[1]} pos):")
    print("    A pos 1-6:", np.round(pwm[0], 2))
    print("    T pos 1-6:", np.round(pwm[3], 2))

    target = "CGTATAAACCGGATATAAATCG"
    scores = score_sequence_with_pwm(target, pwm)
    print(f"\nScanning Target DNA: {target}")
    for pos, s in sorted(scores, key=lambda x: x[1], reverse=True)[:3]:
        print(f"    -> Pos {pos}: Score = {s}")

    # 2. Pandas Differential Expression Table
    df = simulate_rnaseq_dataframe(100)
    print(f"\n[2] Differential Expression Summary (Pandas DataFrame):")
    counts = df['Status'].value_counts()
    for status, count in counts.items():
        print(f"    - {status:15}: {count} genes")

    top_up = df[df['Status'] == 'Upregulated'].sort_values('log2FoldChange', ascending=False).head(3)
    print("\nTop 3 Upregulated Genes:")
    print(top_up[['GeneID', 'log2FoldChange', 'negLog10P']])

if __name__ == "__main__":
    main()