# AUTOGENERATED! DO NOT EDIT! File to edit: ../../nbs/core/02_protein.ipynb.

# %% auto 0
__all__ = ['get_uniprot_seq', 'get_uniprot_features', 'get_uniprot_kd', 'get_uniprot_type', 'apply_mut_single',
           'apply_mut_complex', 'compare_seq']

# %% ../../nbs/core/02_protein.ipynb 3
import pandas as pd
import requests,re
from functools import lru_cache

# for compare seq
from Bio.Align import PairwiseAligner

# %% ../../nbs/core/02_protein.ipynb 5
@lru_cache()
def get_uniprot_seq(uniprot_id):
    "Queries the UniProt database to retrieve the protein sequence for a given UniProt ID."
    
    url = f"https://www.uniprot.org/uniprot/{uniprot_id}.fasta"
    response = requests.get(url)

    # Check if the request was successful (status code 200)
    if response.status_code == 200:
        data = response.text
        # The sequence starts after the first line, which is a description
        sequence = ''.join(data.split('\n')[1:]).strip()
        return sequence
    else:
        return f"Error: Unable to retrieve sequence for UniProt ID {uniprot_id}. Status code: {response.status_code}"

# %% ../../nbs/core/02_protein.ipynb 7
@lru_cache()
def get_uniprot_features(uniprot_id):
    "Given uniprot_id, get specific region for uniprot features."
    # uniprot REST API
    url = f"https://rest.uniprot.org/uniprotkb/{uniprot_id}.json"
    response = requests.get(url)

    if response.status_code == 200:
        data = response.json()
        # Extract the "features" section which contains information
        features = data.get('features', [])
        
        protein_name = (
            data.get("proteinDescription", {})
                .get("recommendedName", {})
                .get("fullName", {})
                .get("value")
        )

        gene_name = (
            data.get("genes", [{}])[0]
                .get("geneName", {})
                .get("value")
        )
        return {
            "uniprot_id": uniprot_id,
            "protein_name": protein_name,
            "gene_name": gene_name,
            "features": features
        }
    else:
        raise ValueError(f"Failed to retrieve UniProt features for {uniprot_id}")

# %% ../../nbs/core/02_protein.ipynb 9
def get_uniprot_kd(uniprot_id):
    "Query 'Domain: Protein kinase' based on UniProt ID and get its sequence info."
    data = get_uniprot_features(uniprot_id)
    seq = get_uniprot_seq(uniprot_id)
    out = []

    for feature in data['features']:
        if feature.get("type") == "Domain" and "Protein kinase" in feature.get("description", ""):
            start = feature['location']['start']['value']
            end = feature['location']['end']['value']
            out.append({
                "uniprot_id": uniprot_id,
                "protein_name": data["protein_name"],
                "gene_name": data["gene_name"],
                "start": start,
                "end": end,
                "description": feature.get("description", ""),
                "sequence": seq[start-1:end]
            })

    return out

# %% ../../nbs/core/02_protein.ipynb 11
def get_uniprot_type(uniprot_id,type_='Signal'):
    "Get region sequences based on UniProt ID features."
    data = get_uniprot_features(uniprot_id)
    seq = get_uniprot_seq(uniprot_id)

    out = []

    for feature in data['features']:
        if feature.get("type") == type_:
            start = feature['location']['start']['value']
            end = feature['location']['end']['value']
            region = {
                'uniprot_id': uniprot_id,
                'type': feature['type'],
                "protein_name": data["protein_name"],
                "gene_name": data["gene_name"],
                'start': start,
                'end': end,
                'description': feature['description'],
                'sequence': seq[start-1:end]
            }
            out.append(region)

    if not out:
        available = sorted({f.get("type") for f in data['features'] if f.get("type")})
        print(f"No feature of type '{type_}' found for {uniprot_id}.")
        print(f"Available feature types: {', '.join(available)}")
        return available

    return out

# %% ../../nbs/core/02_protein.ipynb 16
def apply_mut_single(seq, # protein sequence
           *mutations, # e.g., E709A
           start_pos=1, # if the protein sequence does not start from index 1, indicate the start index to match the mutations
           ):
    "Apply mutations to a protein sequence."
    seq_list = list(seq)  # convert to list for mutability
    
    for mut in mutations:
        # check mutation format
        if len(mut) < 3: raise ValueError(f"Invalid mutation format: {mut}")
        
        from_aa,pos,to_aa = mut[0],int(mut[1:-1])-start_pos,mut[-1]

        # make sure position is within the sequence length
        if pos < 0 or pos >= len(seq_list): raise IndexError(f"Position {pos + 1} out of range for sequence length {len(seq_list)}")
        # make sure aa from mutations matches the residue on the sequence
        if seq_list[pos] != from_aa: raise ValueError(f"Expected {from_aa} at position {pos + 1}, found {seq_list[pos]}")
        
        seq_list[pos] = to_aa
        print('Converted:', mut)
        
    return ''.join(seq_list)

# %% ../../nbs/core/02_protein.ipynb 18
def apply_mut_complex(seq, # protein sequence
                      mut, # mutation (e.g., G776delinsVC/S783C, G778dupGSP)
                      start_pos=1, # if truncated protein sequence, indicate where it starts to match the position of mutation
                      ):
    """
    Apply a composite mutation like 'G776delinsVC/S783C' to `seq`,
    assuming `seq[0]` corresponds to residue number `start_pos`.

    * At most one delins **or** dup is allowed.
    * Point substitutions are executed first; the indel/dup is done last.
    """
    _sub_pat    = re.compile(r'^([A-Z])(\d+)([A-Z])$')          # e.g. S783C
    _delins_pat = re.compile(r'^([A-Z])(\d+)delins([A-Z]+)$')   # e.g. G776delinsVC
    _dup_pat    = re.compile(r'^([A-Z])(\d+)dup([A-Z]+)$')      # e.g. G778dupGSP

    seq = list(seq)
    tokens = mut.split('/')

    # ---------- 1) substitutions (length-neutral) ----------
    for m in tokens:
        if _sub_pat.match(m):
            orig, pos, new = _sub_pat.match(m).groups()
            idx = int(pos) - start_pos
            if seq[idx] != orig:
                raise ValueError(
                    f"Mismatch at position {pos}: expected {orig}, got {seq[idx]}"
                )
            seq[idx] = new

    # ---------- 2) the single length-changing event ----------
    for m in tokens:
        if _delins_pat.match(m):
            orig, pos, ins = _delins_pat.match(m).groups()
            idx = int(pos) - start_pos
            if seq[idx] != orig:
                raise ValueError(
                    f"Mismatch at position {pos}: expected {orig}, got {seq[idx]}"
                )
            seq[idx : idx + 1] = list(ins)         # replace 1 residue with many
            break

        if _dup_pat.match(m):
            orig, pos, dup = _dup_pat.match(m).groups()
            idx = int(pos) - start_pos
            if seq[idx] != orig:
                raise ValueError(
                    f"Mismatch at position {pos}: expected {orig}, got {seq[idx]}"
                )
            seq[idx + 1 : idx + 1] = list(dup)     # insert right after the residue
            break

    return ''.join(seq)

# %% ../../nbs/core/02_protein.ipynb 22
def compare_seq(
    seq1: str,  # original
    seq2: str,  # mutant
    *,
    start_pos: int = 1,
    label1: str = "Original",
    label2: str = "Mutant",
    visualize: bool = True
):
    """
    Align two protein sequences and summarise differences.
    """

    # ----- global alignment using PairwiseAligner -----
    aligner = PairwiseAligner()
    aligner.mode = "global"
    aligner.match_score = 2
    aligner.mismatch_score = -1
    aligner.open_gap_score = -5
    aligner.extend_gap_score = -0.5

    alignment = aligner.align(seq1, seq2)[0]
    aln1 = alignment.aligned[0]
    aln2 = alignment.aligned[1]

    # Reconstruct aligned strings from aligned segments
    aligned_seq1 = []
    aligned_seq2 = []
    i1, i2 = 0, 0

    for (start1, end1), (start2, end2) in zip(aln1, aln2):
        # Add gaps to make sequences align
        while i1 < start1:
            aligned_seq1.append(seq1[i1])
            aligned_seq2.append('-')
            i1 += 1
        while i2 < start2:
            aligned_seq1.append('-')
            aligned_seq2.append(seq2[i2])
            i2 += 1

        # Add aligned part
        for j in range(end1 - start1):
            aligned_seq1.append(seq1[i1])
            aligned_seq2.append(seq2[i2])
            i1 += 1
            i2 += 1

    # Remaining tails
    while i1 < len(seq1):
        aligned_seq1.append(seq1[i1])
        aligned_seq2.append('-')
        i1 += 1
    while i2 < len(seq2):
        aligned_seq1.append('-')
        aligned_seq2.append(seq2[i2])
        i2 += 1

    aln1_str = ''.join(aligned_seq1)
    aln2_str = ''.join(aligned_seq2)

    # ----- find differences -----
    diffs, raw_i1, raw_i2 = [], 0, 0
    for a1, a2 in zip(aln1_str, aln2_str):
        if a1 != '-' and a2 != '-':
            if a1 != a2:
                diffs.append((start_pos + raw_i1, a1, a2, 'substitution'))
            raw_i1 += 1
            raw_i2 += 1
        elif a1 == '-' and a2 != '-':
            diffs.append((start_pos + raw_i1, '-', a2, 'insertion'))
            raw_i2 += 1
        elif a1 != '-' and a2 == '-':
            diffs.append((start_pos + raw_i1, a1, '-', 'deletion'))
            raw_i1 += 1

    # ----- visualization -----
    if visualize:
        for block in range(0, len(aln1_str), 80):
            s1_block = aln1_str[block:block + 80]
            s2_block = aln2_str[block:block + 80]
            marker = ''.join(' ' if x == y else '^' for x, y in zip(s1_block, s2_block))
            left_idx = start_pos + aln1_str[:block].replace('-', '').__len__()
            right_idx = left_idx + s1_block.replace('-', '').__len__() - 1
            print(f"{label1:<10} {left_idx:>5}-{right_idx:<5}: {s1_block}")
            print(f"{label2:<10} {'':>11}: {s2_block}")
            print(f"{'':>22}  {marker}\n")

    # ----- summary list -----
    print("Differences:")
    for pos, ref, new, kind in diffs:
        print(f"  {kind:<12} at {pos:>4}: {ref} → {new}")

    # return diffs
