#!/usr/bin/env python3
"""
Chapter 6: Python for Next Gen Biologist
Script: ch06_dna_oop.py
Description: Object-Oriented Programming (OOP) in Python for Biologists:
             DnaSequence, ProteinSequence, and GenomicVariant classes with inheritance,
             transcription, translation, and mutation application.
Usage:
    python ch06_dna_oop.py
"""

import math

class BiologicalSequence:
    """Abstract base class representing an arbitrary biological sequence."""
    
    def __init__(self, seq_id: str, sequence: str, organism: str = "Unknown"):
        self.seq_id = str(seq_id)
        self.sequence = sequence.upper().strip()
        self.organism = organism
        
    def __len__(self):
        return len(self.sequence)
        
    def __repr__(self):
        return f"<{self.__class__.__name__} id='{self.seq_id}' len={len(self)} organism='{self.organism}'>"
        
    def __getitem__(self, item):
        return self.sequence[item]


class DnaSequence(BiologicalSequence):
    """Represents a double-stranded genomic DNA sequence."""
    
    VALID_BASES = set("ACGTN")
    CODON_TABLE = {
        'ATA':'I', 'ATC':'I', 'ATT':'I', 'ATG':'M',
        'ACA':'T', 'ACC':'T', 'ACG':'T', 'ACT':'T',
        'AAC':'N', 'AAT':'N', 'AAA':'K', 'AAG':'K',
        'AGC':'S', 'AGT':'S', 'AGA':'R', 'AGG':'R',
        'CTA':'L', 'CTC':'L', 'CTG':'L', 'CTT':'L',
        'CCA':'P', 'CCC':'P', 'CCG':'P', 'CCT':'P',
        'CAC':'H', 'CAT':'H', 'CAA':'Q', 'CAG':'Q',
        'CGA':'R', 'CGC':'R', 'CGG':'R', 'CGT':'R',
        'GTA':'V', 'GTC':'V', 'GTG':'V', 'GTT':'V',
        'GCA':'A', 'GCC':'A', 'GCG':'A', 'GCT':'A',
        'GAC':'D', 'GAT':'D', 'GAA':'E', 'GAG':'E',
        'GGA':'G', 'GGC':'G', 'GGG':'G', 'GGT':'G',
        'TCA':'S', 'TCC':'S', 'TCG':'S', 'TCT':'S',
        'TTC':'F', 'TTT':'F', 'TTA':'L', 'TTG':'L',
        'TAC':'Y', 'TAT':'Y', 'TAA':'*', 'TAG':'*',
        'TGC':'C', 'TGT':'C', 'TGA':'*', 'TGG':'W',
    }

    def __init__(self, seq_id: str, sequence: str, organism: str = "Homo sapiens"):
        super().__init__(seq_id, sequence, organism)
        if not set(self.sequence).issubset(self.VALID_BASES):
            raise ValueError(f"Sequence {seq_id} contains invalid DNA bases!")

    def gc_content(self) -> float:
        """Returns the GC percentage of the sequence."""
        if len(self) == 0:
            return 0.0
        gc = self.sequence.count('G') + self.sequence.count('C')
        return round((gc / len(self)) * 100, 2)

    def reverse_complement(self) -> "DnaSequence":
        """Returns a new DnaSequence instance with the reverse complement."""
        trans = str.maketrans("ACGTN", "TGCAN")
        rev_seq = self.sequence.translate(trans)[::-1]
        return DnaSequence(f"{self.seq_id}_revcomp", rev_seq, self.organism)

    def transcribe(self) -> str:
        """Returns the single-stranded mRNA transcript string (T -> U)."""
        return self.sequence.replace("T", "U")

    def translate(self) -> "ProteinSequence":
        """Translates coding DNA into a ProteinSequence object."""
        peptide = []
        for i in range(0, len(self) - 2, 3):
            codon = self.sequence[i:i+3]
            aa = self.CODON_TABLE.get(codon, 'X')
            if aa == '*':
                break
            peptide.append(aa)
        return ProteinSequence(f"{self.seq_id}_protein", "".join(peptide), self.organism)

    def apply_variant(self, variant: "GenomicVariant") -> "DnaSequence":
        """Applies a single-nucleotide variant (SNV) to create a mutated sequence."""
        pos = variant.pos - 1  # Convert 1-based to 0-based
        if pos < 0 or pos >= len(self):
            raise IndexError("Variant position out of bounds.")
        if self.sequence[pos] != variant.ref:
            raise ValueError(f"Reference mismatch at pos {variant.pos}: expected {variant.ref}, found {self.sequence[pos]}")
            
        mutated = self.sequence[:pos] + variant.alt + self.sequence[pos+1:]
        return DnaSequence(f"{self.seq_id}_{variant.ref}{variant.pos}{variant.alt}", mutated, self.organism)


class ProteinSequence(BiologicalSequence):
    """Represents an amino acid peptide sequence."""
    
    # Molecular weights of amino acids in Daltons (Da)
    AA_WEIGHTS = {
        'A': 71.08, 'R': 156.20, 'N': 114.11, 'D': 115.09, 'C': 103.14,
        'E': 129.12, 'Q': 128.13, 'G': 57.05,  'H': 137.14, 'I': 113.17,
        'L': 113.17, 'K': 128.18, 'M': 131.21, 'F': 147.18, 'P': 97.12,
        'S': 87.08,  'T': 101.11, 'W': 186.21, 'Y': 163.18, 'V': 99.13
    }

    def molecular_weight_kda(self) -> float:
        """Calculates molecular weight in kiloDaltons including terminal water molecule."""
        weight = 18.015  # H2O terminal
        for aa in self.sequence:
            weight += self.AA_WEIGHTS.get(aa, 110.0)
        return round(weight / 1000.0, 2)


class GenomicVariant:
    """Represents a genomic single-nucleotide variant (SNV)."""
    
    def __init__(self, chrom: str, pos: int, ref: str, alt: str):
        self.chrom = str(chrom)
        self.pos = int(pos)
        self.ref = str(ref).upper()
        self.alt = str(alt).upper()
        
    def __repr__(self):
        return f"<GenomicVariant {self.chrom}:{self.pos} {self.ref}>{self.alt}>"


def main():
    print("=== Chapter 6: Object-Oriented Biological Sequences ===")
    
    # 1. Instantiate Wild-Type DNA Sequence (Human Beta-Globin Exon 1 fragment)
    wt_dna = DnaSequence(
        seq_id="HBB_WT",
        sequence="ATGGTGCACCTGACTCCTGAGGAGAAGTCTGCCGTTACTGCCCTGTGGGGCAAGGTGAACGTGGATGAAGTTGGTGGTGAGGCCCTGGGCAGGCTGCTGG",
        organism="Homo sapiens"
    )
    print(f"Wild-Type DNA     : {wt_dna}")
    print(f"GC Content        : {wt_dna.gc_content()}%")
    
    wt_protein = wt_dna.translate()
    print(f"Translated Protein: {wt_protein} ({wt_protein.sequence})")
    print(f"Molecular Weight  : {wt_protein.molecular_weight_kda()} kDa")

    # 2. Apply Sickle Cell Mutation: HBB c.20A>T (p.Glu7Val / HbS)
    # Codon 6: GAG -> GTG at nucleotide position 20
    scd_var = GenomicVariant(chrom="chr11", pos=20, ref="A", alt="T")
    print(f"\nApplying Variant  : {scd_var}")
    
    mut_dna = wt_dna.apply_variant(scd_var)
    mut_protein = mut_dna.translate()
    
    print(f"Mutated DNA       : {mut_dna}")
    print(f"Mutated Protein   : {mut_protein} ({mut_protein.sequence})")
    print(f"Amino Acid Delta  : Pos 7 changed from '{wt_protein[6]}' to '{mut_protein[6]}'")

if __name__ == "__main__":
    main()
