
"""
PKTron Biology — Simplified Protein Folding Model
===================================================
HP lattice model (Hydrophobic-Polar) for protein folding simulation.
"""
import numpy as np
_R = np.float32

HYDROPHOBIC = set(["Ala","Val","Ile","Leu","Met","Phe","Trp","Pro","Gly"])
POLAR       = set(["Ser","Thr","Cys","Tyr","Asn","Gln","Asp","Glu","Lys","Arg","His"])


def amino_acid_properties(aa_name):
    """Return basic physicochemical properties of an amino acid."""
    props = {
        "Ala":{"mw":89.1, "hydrophobic":True,  "charge": 0},
        "Arg":{"mw":174.2,"hydrophobic":False, "charge":+1},
        "Asn":{"mw":132.1,"hydrophobic":False, "charge": 0},
        "Asp":{"mw":133.1,"hydrophobic":False, "charge":-1},
        "Cys":{"mw":121.2,"hydrophobic":False, "charge": 0},
        "Gln":{"mw":146.2,"hydrophobic":False, "charge": 0},
        "Glu":{"mw":147.1,"hydrophobic":False, "charge":-1},
        "Gly":{"mw":75.0, "hydrophobic":True,  "charge": 0},
        "His":{"mw":155.2,"hydrophobic":False, "charge":+0.5},
        "Ile":{"mw":131.2,"hydrophobic":True,  "charge": 0},
        "Leu":{"mw":131.2,"hydrophobic":True,  "charge": 0},
        "Lys":{"mw":146.2,"hydrophobic":False, "charge":+1},
        "Met":{"mw":149.2,"hydrophobic":True,  "charge": 0},
        "Phe":{"mw":165.2,"hydrophobic":True,  "charge": 0},
        "Pro":{"mw":115.1,"hydrophobic":True,  "charge": 0},
        "Ser":{"mw":105.1,"hydrophobic":False, "charge": 0},
        "Thr":{"mw":119.1,"hydrophobic":False, "charge": 0},
        "Trp":{"mw":204.2,"hydrophobic":True,  "charge": 0},
        "Tyr":{"mw":181.2,"hydrophobic":False, "charge": 0},
        "Val":{"mw":117.1,"hydrophobic":True,  "charge": 0},
    }
    return props.get(aa_name, {"mw":110.0,"hydrophobic":False,"charge":0})


def hp_encode(amino_acids):
    """
    Convert list of amino acid names to HP string.
    H=hydrophobic, P=polar.
    """
    return "".join("H" if aa in HYDROPHOBIC else "P" for aa in amino_acids)


def fold_hp_2d(hp_string, seed=0):
    """
    Greedy HP lattice folding on 2D grid.
    Returns (positions, energy) where energy < 0 is better.
    """
    rng  = np.random.default_rng(seed)
    n    = len(hp_string)
    pos  = [(0,0)]
    dirs = [(1,0),(-1,0),(0,1),(0,-1)]
    occupied = {(0,0)}

    for i in range(1,n):
        x0,y0 = pos[-1]
        rng.shuffle(dirs)
        placed = False
        for dx,dy in dirs:
            nx,ny = x0+dx,y0+dy
            if (nx,ny) not in occupied:
                pos.append((nx,ny)); occupied.add((nx,ny))
                placed = True; break
        if not placed:
            pos.append((x0+i,y0))   # extend linearly if stuck
            occupied.add((x0+i,y0))

    # Energy = number of H-H non-bonded contacts
    energy = 0
    for i,p in enumerate(pos):
        if hp_string[i]!="H": continue
        x,y = p
        for dx,dy in [(1,0),(-1,0),(0,1),(0,-1)]:
            nb = (x+dx,y+dy)
            if nb in occupied:
                j = pos.index(nb)
                if abs(j-i)>1 and hp_string[j]=="H":
                    energy -= 1
    return pos, energy//2   # each contact counted twice


def protein_isoelectric_point(amino_acids):
    """
    Approximate pI of a protein from its amino acid composition.
    """
    pKa_table = {"Asp":3.9,"Glu":4.1,"His":6.0,"Cys":8.3,
                 "Tyr":10.1,"Lys":10.5,"Arg":12.5}
    total_charge = sum(amino_acid_properties(aa)["charge"]
                       for aa in amino_acids)
    return round(7.0 + total_charge*0.1, 2)


def protein_molecular_weight(amino_acids):
    """Approximate MW in Da (subtracting water per peptide bond)."""
    mw = sum(amino_acid_properties(aa)["mw"] for aa in amino_acids)
    mw -= 18.015*(len(amino_acids)-1)
    return round(mw,2)


def secondary_structure_predict(hp_string):
    """
    Trivial secondary structure prediction:
    runs of 4+ H -> helix, runs of P -> coil.
    """
    struct = []
    i = 0
    while i < len(hp_string):
        if hp_string[i]=="H" and i+3<len(hp_string) and hp_string[i:i+4]=="HHHH":
            j = i
            while j<len(hp_string) and hp_string[j]=="H": j+=1
            struct.extend(["helix"]*(j-i)); i=j
        elif hp_string[i:i+3]=="PPP":
            j=i
            while j<len(hp_string) and hp_string[j]=="P": j+=1
            struct.extend(["sheet"]*(j-i)); i=j
        else:
            struct.append("coil"); i+=1
    return struct
