from src.model import tdc_prompts, txgemma_predict

def predict_kiba_score(drug_smile, amino_acid):
    TDC_PROMPT = tdc_prompts["KIBA"].replace("{Drug SMILES}", drug_smile).replace("{Target amino acid sequence}", amino_acid)
    response = txgemma_predict(TDC_PROMPT)
    return response.split("Answer:")[1].strip()

def predict(task, drug_smile, amino_acid=None):
    if task == "KIBA Score":
        if amino_acid is None:
            raise ValueError("amino_acid parameter is required for KIBA task")
        kiba_score = predict_kiba_score(drug_smile, amino_acid)
        return f"{kiba_score} Binding Affinity On Scale of 0-1000"


    if task == "Skin Reaction":
        TDC_PROMPT = tdc_prompts["Skin_Reaction"].replace("{Drug SMILES}", drug_smile)
        response = txgemma_predict(TDC_PROMPT).split("Answer:")[1].strip()

        if "(A)" in response: response = f"{drug_smile} does not cause a skin reaction!"
        elif "(B)" in response: response = f"{drug_smile} causes a skin reaction!"

        return response
    
    if task == "Liver Safety":
        TDC_PROMPT = tdc_prompts["DILI"].replace("{Drug SMILES}", drug_smile)
        response = txgemma_predict(TDC_PROMPT).split("Answer:")[1].strip()

        if "(A)" in response: response = f"{drug_smile} does not damage a liver!"
        elif "(B)" in response: response = f"{drug_smile} can damage a liver!"

        return response