
"""
PKTron Biology — DNA Sequence Simulation
=========================================
String-based DNA operations: complement, transcription, translation,
GC content, codon analysis, mutation simulation.
"""
import numpy as np
import random

BASES       = ("A","T","C","G")
COMPLEMENT  = {"A":"T","T":"A","C":"G","G":"C"}
CODON_TABLE = {
    "TTT":"Phe","TTC":"Phe","TTA":"Leu","TTG":"Leu",
    "CTT":"Leu","CTC":"Leu","CTA":"Leu","CTG":"Leu",
    "ATT":"Ile","ATC":"Ile","ATA":"Ile","ATG":"Met",
    "GTT":"Val","GTC":"Val","GTA":"Val","GTG":"Val",
    "TCT":"Ser","TCC":"Ser","TCA":"Ser","TCG":"Ser",
    "CCT":"Pro","CCC":"Pro","CCA":"Pro","CCG":"Pro",
    "ACT":"Thr","ACC":"Thr","ACA":"Thr","ACG":"Thr",
    "GCT":"Ala","GCC":"Ala","GCA":"Ala","GCG":"Ala",
    "TAT":"Tyr","TAC":"Tyr","TAA":"Stop","TAG":"Stop",
    "CAT":"His","CAC":"His","CAA":"Gln","CAG":"Gln",
    "AAT":"Asn","AAC":"Asn","AAA":"Lys","AAG":"Lys",
    "GAT":"Asp","GAC":"Asp","GAA":"Glu","GAG":"Glu",
    "TGT":"Cys","TGC":"Cys","TGA":"Stop","TGG":"Trp",
    "CGT":"Arg","CGC":"Arg","CGA":"Arg","CGG":"Arg",
    "AGT":"Ser","AGC":"Ser","AGA":"Arg","AGG":"Arg",
    "GGT":"Gly","GGC":"Gly","GGA":"Gly","GGG":"Gly",
}


def simulate_dna(sequence=None, length=60, seed=None):
    """
    Validate or randomly generate a DNA sequence.

    Parameters
    ----------
    sequence : str or None   if str, validates & returns uppercase
    length   : int           length if sequence is None
    seed     : int           rng seed for random generation

    Returns
    -------
    str  DNA sequence (uppercase, valid ATCG only)
    """
    if sequence is not None:
        seq = sequence.upper().strip()
        invalid = set(seq) - set(BASES)
        if invalid:
            raise ValueError(f"Invalid bases: {invalid}")
        return seq
    rng = np.random.default_rng(seed)
    return "".join(rng.choice(list(BASES), size=length))


def complement(sequence):
    """Return complementary DNA strand (5'->3' antiparallel)."""
    return "".join(COMPLEMENT[b] for b in sequence.upper())


def reverse_complement(sequence):
    return complement(sequence)[::-1]


def gc_content(sequence):
    """Return GC content as fraction [0,1]."""
    seq = sequence.upper()
    return (seq.count("G") + seq.count("C")) / max(1, len(seq))


def transcribe(dna_sequence):
    """DNA -> mRNA: T replaced by U."""
    return dna_sequence.upper().replace("T","U")


def translate(mrna_sequence):
    """mRNA -> amino acid sequence (stops at Stop codon)."""
    mrna   = mrna_sequence.upper().replace("U","T")
    codons = [mrna[i:i+3] for i in range(0, len(mrna)-2, 3)]
    aas    = []
    for codon in codons:
        aa = CODON_TABLE.get(codon, "?")
        if aa == "Stop":
            break
        aas.append(aa)
    return aas


def codon_usage(sequence):
    """Count codons in reading frame 0."""
    seq    = sequence.upper()
    codons = [seq[i:i+3] for i in range(0, len(seq)-2, 3)]
    usage  = {}
    for c in codons:
        if len(c)==3:
            usage[c] = usage.get(c,0)+1
    return usage


def mutate_dna(sequence, rate=0.01, seed=None):
    """
    Apply random point mutations to a DNA sequence.

    Parameters
    ----------
    sequence : str   original DNA
    rate     : float mutation probability per base (0-1)
    seed     : int   rng seed

    Returns
    -------
    str  mutated sequence
    """
    rng = np.random.default_rng(seed)
    seq = list(sequence.upper())
    mutations = 0
    for i in range(len(seq)):
        if rng.random() < rate:
            alternatives = [b for b in BASES if b != seq[i]]
            seq[i] = str(rng.choice(alternatives))
            mutations += 1
    return "".join(seq), mutations


def hamming_distance(seq1, seq2):
    """Count differing positions between two equal-length sequences."""
    if len(seq1) != len(seq2):
        raise ValueError("Sequences must be the same length")
    return sum(a!=b for a,b in zip(seq1.upper(),seq2.upper()))


def find_orfs(sequence, min_length=30):
    """
    Find all Open Reading Frames starting with ATG, ending at Stop.
    Returns list of (start, end, protein_length).
    """
    seq  = sequence.upper()
    orfs = []
    for i in range(len(seq)-2):
        if seq[i:i+3] == "ATG":
            for j in range(i, len(seq)-2, 3):
                codon = seq[j:j+3]
                if len(codon)<3: break
                if CODON_TABLE.get(codon)=="Stop":
                    if (j-i) >= min_length:
                        orfs.append((i,j+3,(j-i)//3))
                    break
    return orfs
