#!/usr/bin/env python3
"""
Chapter 8: Python for Next Gen Biologist
Script: ch08_regex_data_cleaning.py
Description: Regular expressions for biological sequence patterns:
             Zinc finger C2H2 motifs, bacterial Pribnow box, restriction cut sites,
             and clinical metadata cleaning.
Usage:
    python ch08_regex_data_cleaning.py
"""

import re
import csv
import io

def find_c2h2_zinc_fingers(protein_seq: str) -> list:
    """
    Finds classic C2H2 zinc finger motifs: C-x(2,4)-C-x(12)-H-x(3,5)-H
    """
    pattern = re.compile(r'C.{2,4}C.{12}H.{3,5}H')
    matches = []
    for match in pattern.finditer(protein_seq):
        matches.append({
            'start': match.start() + 1,  # 1-based coordinate
            'end': match.end(),
            'motif': match.group()
        })
    return matches

def find_promoter_pribnow_box(dna_seq: str) -> list:
    """
    Scans for bacterial -10 Pribnow box consensus (TATAAT) with up to 1 mismatch.
    Regex uses alternation for core conservation.
    """
    # Pattern allowing variation in central positions
    pattern = re.compile(r'TA[TA][AA]T', re.IGNORECASE)
    matches = []
    for m in pattern.finditer(dna_seq):
        matches.append((m.start() + 1, m.end(), m.group().upper()))
    return matches

def clean_clinical_metadata(raw_csv_data: str) -> list:
    """
    Cleans messy patient/sample metadata: standardizes sample IDs, cleans ages,
    and formats diagnosis stages.
    """
    reader = csv.DictReader(io.StringIO(raw_csv_data.strip()))
    cleaned = []
    for row in reader:
        # Standardize Sample ID: e.g. " patient_001 " -> "PT-001"
        sid = re.sub(r'[^0-9]', '', row['sample_id'])
        clean_sid = f"PT-{int(sid):03d}" if sid else "UNKNOWN"
        
        # Clean Age: "54 yrs" / "54yo" -> 54
        age_match = re.search(r'\d+', row['age'])
        age = int(age_match.group()) if age_match else None
        
        # Clean Stage: "stage IIIa" -> "Stage IIIA"
        stage = row['stage'].strip().upper()
        stage = re.sub(r'\s+', ' ', stage)
        
        cleaned.append({
            'sample_id': clean_sid,
            'age': age,
            'gender': row['gender'].strip().upper()[:1],
            'stage': stage
        })
    return cleaned

def main():
    print("=== Chapter 8: Genomic Regular Expressions & Data Cleaning ===")

    # Demo 1: C2H2 Zinc Finger Search
    protein = "MKGEELFTGVVPILVELDGDVNGHKFSVSGEGEGDATYGKLTLKFICTYCKTFVLDDNLQAHVTHMGPVI"
    zf_matches = find_c2h2_zinc_fingers(protein)
    print(f"\n[1] Scanning Protein ({len(protein)} AA) for C2H2 Zinc Fingers:")
    if zf_matches:
        for m in zf_matches:
            print(f"    -> Found Motif at {m['start']}-{m['end']}: {m['motif']}")
    else:
        print("    -> No standard C2H2 motif detected.")

    # Demo 2: Bacterial Pribnow Box (-10 Promoter)
    promoter_dna = "AGCTTTTCATTCTGACTGCAACGGGCAATATGTCTCTGTGTGGATTAAAAAAAGAGTGTCTGATAGCAGCTTCTGAACTGGTTACCTGCCGTGAGTAAATTAAA"
    pribnow = find_promoter_pribnow_box(promoter_dna)
    print(f"\n[2] Scanning DNA for Pribnow Box (-10 Promoter):")
    for start, end, seq in pribnow:
        print(f"    -> Consensus Match at bp {start}-{end}: {seq}")

    # Demo 3: Cleaning Clinical Metadata Table
    dirty_csv = """sample_id,age,gender,stage
patient_1, 58 yrs ,Female,stage IIa
pt-02, 63yo ,M,STAGE 3B
sample 003,45 years,F,stage i
"""
    cleaned = clean_clinical_metadata(dirty_csv)
    print(f"\n[3] Cleaned Clinical Metadata Table:")
    for row in cleaned:
        print(f"    Sample: {row['sample_id']} | Age: {row['age']} | Sex: {row['gender']} | Stage: {row['stage']}")

if __name__ == "__main__":
    main()