import _bootstrap
import torch
from model.build_vocab import WordVocab
from model.pretrain_trfm import TrfmSeq2seq
from model.utils import split

from transformers import T5EncoderModel, T5Tokenizer

import re
import gc
import numpy as np
import pandas as pd
import pickle
import math
from project_paths import (

    DEGREE_SMILES_PATH, KCAT_DATASET_PATH, KCAT_KM_SAMPLES_PATH, KM_TEST_PATH,

    PH_SMILES_PATH, PROT_T5_MODEL, TRFM_PATH, UNIKP_MODEL_DIR, VOCAB_PATH,

)


# ============================
# SMILES -> vector
# ============================

def smiles_to_vec(Smiles):

    pad_index = 0
    unk_index = 1
    eos_index = 2
    sos_index = 3

    vocab = WordVocab.load_vocab(
        VOCAB_PATH
    )


    def get_inputs(sm):

        seq_len = 220

        sm = sm.split()

        if len(sm) > 218:
            print(
                "SMILES too long:",
                len(sm)
            )
            sm = sm[:109] + sm[-109:]


        ids = [
            vocab.stoi.get(
                token,
                unk_index
            )
            for token in sm
        ]


        ids = (
            [sos_index]
            +
            ids
            +
            [eos_index]
        )


        seg = [1] * len(ids)


        padding = [
            pad_index
        ] * (
            seq_len - len(ids)
        )


        ids.extend(padding)
        seg.extend(padding)


        return ids, seg



    def get_array(smiles):

        x_id = []
        x_seg = []

        for sm in smiles:

            a, b = get_inputs(sm)

            x_id.append(a)
            x_seg.append(b)


        return (
            torch.tensor(x_id),
            torch.tensor(x_seg)
        )



    # load SMILES Transformer

    trfm = TrfmSeq2seq(
        len(vocab),
        256,
        len(vocab),
        4
    )


    trfm.load_state_dict(
        torch.load(
            TRFM_PATH,
            map_location="cpu"
        )
    )


    trfm.eval()



    x_split = [
        split(sm)
        for sm in Smiles
    ]


    xid, xseg = get_array(
        x_split
    )


    X = trfm.encode(
        torch.t(xid)
    )


    return X



# ============================
# Protein sequence -> vector
# ============================


def Seq_to_vec(Sequence):


    for i in range(len(Sequence)):

        if len(Sequence[i]) > 1000:

            Sequence[i] = (
                Sequence[i][:500]
                +
                Sequence[i][-500:]
            )



    sequences_Example = []


    for seq in Sequence:

        spaced = ""

        for aa in seq[:-1]:

            spaced += aa + " "

        spaced += seq[-1]


        sequences_Example.append(
            spaced
        )



    # load ProtT5

    tokenizer = T5Tokenizer.from_pretrained(
        PROT_T5_MODEL,
        do_lower_case=False
    )


    model = T5EncoderModel.from_pretrained(
        PROT_T5_MODEL
    )


    gc.collect()


    device = torch.device(
        "cuda:0"
        if torch.cuda.is_available()
        else "cpu"
    )


    print(
        "Using device:",
        device
    )


    model = model.to(device)

    model.eval()



    features = []



    for i, seq in enumerate(sequences_Example):

        print(
            "Processing protein:",
            i + 1
        )


        seq = [
            re.sub(
                r"[UZOB]",
                "X",
                seq
            )
        ]


        ids = tokenizer.batch_encode_plus(
            seq,
            add_special_tokens=True,
            padding=True
        )


        input_ids = torch.tensor(
            ids["input_ids"]
        ).to(device)


        attention_mask = torch.tensor(
            ids["attention_mask"]
        ).to(device)



        with torch.no_grad():

            embedding = model(
                input_ids=input_ids,
                attention_mask=attention_mask
            )



        embedding = (
            embedding
            .last_hidden_state
            .cpu()
            .numpy()
        )



        seq_len = (
            attention_mask[0] == 1
        ).sum()


        seq_embedding = (
            embedding[0]
            [:seq_len-1]
        )


        features.append(
            seq_embedding
        )



    # mean pooling

    features_normalize = np.zeros(
        [
            len(features),
            len(features[0][0])
        ],
        dtype=float
    )



    for i in range(len(features)):

        for k in range(
            len(features[0][0])
        ):

            for j in range(
                len(features[i])
            ):

                features_normalize[i][k] += (
                    features[i][j][k]
                )


            features_normalize[i][k] /= (
                len(features[i])
            )


    return features_normalize




# ============================
# Main
# ============================


if __name__ == "__main__":


    # 示例蛋白序列
    sequences = [
        "MEDIPDTSRPPLKYVKGIPLIKYFAEALESLQDFQAQPDDLLISTYPKSGTTWVSEILDMIYQDGDVEKCRRAPVFIRVPFLEFKAPGIPTGLEVLKDTPAPRLIKTHLPLALLPQTLLDQKVKVVYVARNAKDVAVSYYHFYRMAKVHPDPDTWDSFLEKFMAGEVSYGSWYQHVQEWWELSHTHPVLYLFYEDMKENPKREIQKILKFVGRSLPEETVDLIVQHTSFKEMKNNSMANYTTLSPDIMDHSISAFMRKGISGDWKTTFTVAQNERFDADYAKKMEGCGLSFRTQL"
    ]


    # 示例底物 SMILES

    Smiles = [
        "OC1=CC=C(C[C@@H](C(O)=O)N)C=C1"
    ]



    print(
        "Extracting protein embedding..."
    )

    seq_vec = Seq_to_vec(
        sequences
    )


    print(
        "Extracting SMILES embedding..."
    )

    smiles_vec = smiles_to_vec(
        Smiles
    )


    print(
        "Protein vector:",
        seq_vec.shape
    )


    print(
        "SMILES vector:",
        smiles_vec.shape
    )



    # concatenate

    fused_vector = np.concatenate(
        (
            smiles_vec,
            seq_vec
        ),
        axis=1
    )



    print(
        "Fused vector:",
        fused_vector.shape
    )



    # load UniKP model

    with open(
        UNIKP_MODEL_DIR / "UniKP for kcat.pkl",
        "rb"
    ) as f:

        model = pickle.load(f)



    # prediction

    pred = model.predict(
        fused_vector
    )


    # log10 inverse transform

    pred_value = [
        math.pow(
            10,
            x
        )
        for x in pred
    ]



    print("\n========== Result ==========")

    print(
        "Protein:",
        sequences[0][:50],
        "..."
    )

    print(
        "SMILES:",
        Smiles[0]
    )


    print(
        "Predicted kcat:",
        pred_value[0],
        "s^-1"
    )



    # save

    result = pd.DataFrame(
        {
            "sequence": sequences,
            "SMILES": Smiles,
            "kcat": pred_value
        }
    )


    result.to_excel(
        "UniKP_kcat_prediction.xlsx",
        index=False
    )


    print(
        "\nSaved:"
        " UniKP_kcat_prediction.xlsx"
    )