from const import *


def get_sequence_fragment(x: pd.Series, feature_size: int = 200) -> str:
    seq = x["seq"]
    len_seq = len(seq)
    if len_seq >= feature_size * 2:
        return seq[0:feature_size] + seq[-feature_size:]
    s1 = seq[0:int(len_seq / 2)]
    s2 = seq[int(len_seq / 2):]
    n1 = "N" * (feature_size - len(s1))
    n2 = "N" * (feature_size - len(s2))
    s1 = s1 + n1
    s2 = n2 + s2
    return s1 + s2


def feature_encoding(df_in: pd.DataFrame, flag_verbose: bool) -> pd.DataFrame:
    feature_int_encoder = LabelEncoder()
    voc = ["A", "C", "G", "T", "N"]
    num_classes = len(voc)
    feature_int_encoder.fit(voc)

    df = df_in.loc[:, ["id", "seq_frag"]].copy()
    print("  Step 2/7: Label Encoding - Transforming non-numerical labels to numerical labels")
    df["int_enc"] = df.swifter.progress_bar(flag_verbose).apply(
        lambda x: np.array(feature_int_encoder.transform(list(x["seq_frag"]))).reshape(-1, 1), axis=1)
    df = df.drop(columns="seq_frag")

    print("  Step 3/7: One-Hot Encoding - Converting class vectors to binary class matrices")
    df["feature"] = df.swifter.progress_bar(flag_verbose).apply(
        lambda x: keras.utils.to_categorical(x["int_enc"], num_classes=num_classes), axis=1)
    df = df.drop(columns="int_enc")

    return df


def predict(df_in: pd.DataFrame, genome_file: str, path_to_model: str) -> Optional[pd.DataFrame]:
    model = keras.models.load_model(path_to_model)
    pre_feature = df_in["feature"].to_numpy()
    df = df_in.drop(columns="feature")

    if pre_feature.shape[0] == 0:
        print("Info: " + genome_file + " has no candidate to be classified")
        return None

    l_class = ["DTA", "DTC", "DTH", "DTM", "DTT", "NonTIR"]
    target_int_encoder = LabelEncoder()
    target_int_encoder.fit(l_class)
    target_int_encoded = target_int_encoder.transform(l_class)
    d = dict(zip(target_int_encoded, l_class))

    print("  Step 4/7: CNN prediction")
    # try:
    #     predicted_labels = model.predict(np.stack(pre_feature))
    # except InternalError as e:
    #     print(e)
    #     with tf.device("/cpu:0"):
    #         pre_feature_tensor = tf.convert_to_tensor(np.stack(pre_feature), np.float32)
    #         predicted_labels = model.predict(pre_feature_tensor)
    predicted_labels = model.predict(np.stack(pre_feature))

    df["percent"] = pd.Series(predicted_labels.max(axis=-1))
    y_classes = predicted_labels.argmax(axis=-1)
    df["TIR_type"] = pd.Series([d[i] for i in y_classes])
    return df


def postprocessing(df_in: pd.DataFrame, flag_verbose: bool) -> pd.DataFrame:
    df = df_in.loc[:, ["id", "TIR_type"]]
    df = df[df["TIR_type"] != "NonTIR"].reset_index(drop=True)
    print("  Step 5/7: Retrieving sequence ID")
    df["seqid"] = df.swifter.progress_bar(flag_verbose).apply(lambda x: x["id"].split(":")[0], axis=1)
    print("  Step 6/7: Retrieving sequence starting coordinate")
    df["sstart"] = df.swifter.progress_bar(flag_verbose).apply(lambda x: int(x["id"].split(":")[1]), axis=1)
    print("  Step 7/7: Retrieving sequence ending coordinate")
    df["send"] = df.swifter.progress_bar(flag_verbose).apply(lambda x: int(x["id"].split(":")[2]), axis=1)
    df = df.loc[:, ["TIR_type", "id", "seqid", "sstart", "send"]]
    df = df.sort_values(["TIR_type", "seqid", "sstart", "send"], ignore_index=True)
    return df


def execute(TIRLearner_instance) -> pd.DataFrame:
    df = TIRLearner_instance["base"].copy()

    print("  Step 1/7: Getting sequence fragment for prediction")
    df["seq_frag"] = df.swifter.progress_bar(TIRLearner_instance.flag_verbose).apply(get_sequence_fragment, axis=1)
    df = df.drop(columns="seq")

    df = feature_encoding(df, TIRLearner_instance.flag_verbose)

    df = predict(df, TIRLearner_instance.genome_file_path, CNN_MODEL_DIR_PATH)

    return postprocessing(df, TIRLearner_instance.flag_verbose)
