import os.path

import numpy as np
import pandas as pd

from .information import AA

AA_IDX = {AA[i]: i for i in range(len(AA))}

BLOSUM = np.array(
    [
        [
            3.9029,
            0.6127,
            0.5883,
            0.5446,
            0.8680,
            0.7568,
            0.7413,
            1.0569,
            0.5694,
            0.6325,
            0.6019,
            0.7754,
            0.7232,
            0.4649,
            0.7541,
            1.4721,
            0.9844,
            0.4165,
            0.5426,
            0.9365,
        ],
        [
            0.6127,
            6.6656,
            0.8586,
            0.5732,
            0.3089,
            1.4058,
            0.9608,
            0.4500,
            0.9170,
            0.3548,
            0.4739,
            2.0768,
            0.6226,
            0.3807,
            0.4815,
            0.7672,
            0.6778,
            0.3951,
            0.5560,
            0.4201,
        ],
        [
            0.5883,
            0.8586,
            7.0941,
            1.5539,
            0.3978,
            1.0006,
            0.9113,
            0.8637,
            1.2220,
            0.3279,
            0.3100,
            0.9398,
            0.4745,
            0.3543,
            0.4999,
            1.2315,
            0.9842,
            0.2778,
            0.4860,
            0.3690,
        ],
        [
            0.5446,
            0.5732,
            1.5539,
            7.3979,
            0.3015,
            0.8971,
            1.6878,
            0.6343,
            0.6786,
            0.3390,
            0.2866,
            0.7841,
            0.3465,
            0.2990,
            0.5987,
            0.9135,
            0.6948,
            0.2321,
            0.3457,
            0.3365,
        ],
        [
            0.8680,
            0.3089,
            0.3978,
            0.3015,
            19.5766,
            0.3658,
            0.2859,
            0.4204,
            0.3550,
            0.6535,
            0.6423,
            0.3491,
            0.6114,
            0.4390,
            0.3796,
            0.7384,
            0.7406,
            0.4500,
            0.4342,
            0.7558,
        ],
        [
            0.7568,
            1.4058,
            1.0006,
            0.8971,
            0.3658,
            6.2444,
            1.9017,
            0.5386,
            1.1680,
            0.3829,
            0.4773,
            1.5543,
            0.8643,
            0.3340,
            0.6413,
            0.9656,
            0.7913,
            0.5094,
            0.6111,
            0.4668,
        ],
        [
            0.7413,
            0.9608,
            0.9113,
            1.6878,
            0.2859,
            1.9017,
            5.4695,
            0.4813,
            0.9600,
            0.3305,
            0.3729,
            1.3083,
            0.5003,
            0.3307,
            0.6792,
            0.9504,
            0.7414,
            0.3743,
            0.4965,
            0.4289,
        ],
        [
            1.0569,
            0.4500,
            0.8637,
            0.6343,
            0.4204,
            0.5386,
            0.4813,
            6.8763,
            0.4930,
            0.2750,
            0.2845,
            0.5889,
            0.3955,
            0.3406,
            0.4774,
            0.9036,
            0.5793,
            0.4217,
            0.3487,
            0.3370,
        ],
        [
            0.5694,
            0.9170,
            1.2220,
            0.6786,
            0.3550,
            1.1680,
            0.9600,
            0.4930,
            13.5060,
            0.3263,
            0.3807,
            0.7789,
            0.5841,
            0.6520,
            0.4729,
            0.7367,
            0.5575,
            0.4441,
            1.7979,
            0.3394,
        ],
        [
            0.6325,
            0.3548,
            0.3279,
            0.3390,
            0.6535,
            0.3829,
            0.3305,
            0.2750,
            0.3263,
            3.9979,
            1.6944,
            0.3964,
            1.4777,
            0.9458,
            0.3847,
            0.4432,
            0.7798,
            0.4089,
            0.6304,
            2.4175,
        ],
        [
            0.6019,
            0.4739,
            0.3100,
            0.2866,
            0.6423,
            0.4773,
            0.3729,
            0.2845,
            0.3807,
            1.6944,
            3.7966,
            0.4283,
            1.9943,
            1.1546,
            0.3711,
            0.4289,
            0.6603,
            0.5680,
            0.6921,
            1.3142,
        ],
        [
            0.7754,
            2.0768,
            0.9398,
            0.7841,
            0.3491,
            1.5543,
            1.3083,
            0.5889,
            0.7789,
            0.3964,
            0.4283,
            4.7643,
            0.6253,
            0.3440,
            0.7038,
            0.9319,
            0.7929,
            0.3589,
            0.5322,
            0.4565,
        ],
        [
            0.7232,
            0.6226,
            0.4745,
            0.3465,
            0.6114,
            0.8643,
            0.5003,
            0.3955,
            0.5841,
            1.4777,
            1.9943,
            0.6253,
            6.4815,
            1.0044,
            0.4239,
            0.5986,
            0.7938,
            0.6103,
            0.7084,
            1.2689,
        ],
        [
            0.4649,
            0.3807,
            0.3543,
            0.2990,
            0.4390,
            0.3340,
            0.3307,
            0.3406,
            0.6520,
            0.9458,
            1.1546,
            0.3440,
            1.0044,
            8.1288,
            0.2874,
            0.4400,
            0.4817,
            1.3744,
            2.7694,
            0.7451,
        ],
        [
            0.7541,
            0.4815,
            0.4999,
            0.5987,
            0.3796,
            0.6413,
            0.6792,
            0.4774,
            0.4729,
            0.3847,
            0.3711,
            0.7038,
            0.4239,
            0.2874,
            12.8375,
            0.7555,
            0.6889,
            0.2818,
            0.3635,
            0.4431,
        ],
        [
            1.4721,
            0.7672,
            1.2315,
            0.9135,
            0.7384,
            0.9656,
            0.9504,
            0.9036,
            0.7367,
            0.4432,
            0.4289,
            0.9319,
            0.5986,
            0.4400,
            0.7555,
            3.8428,
            1.6139,
            0.3853,
            0.5575,
            0.5652,
        ],
        [
            0.9844,
            0.6778,
            0.9842,
            0.6948,
            0.7406,
            0.7913,
            0.7414,
            0.5793,
            0.5575,
            0.7798,
            0.6603,
            0.7929,
            0.7938,
            0.4817,
            0.6889,
            1.6139,
            4.8321,
            0.4309,
            0.5732,
            0.9809,
        ],
        [
            0.4165,
            0.3951,
            0.2778,
            0.2321,
            0.4500,
            0.5094,
            0.3743,
            0.4217,
            0.4441,
            0.4089,
            0.5680,
            0.3589,
            0.6103,
            1.3744,
            0.2818,
            0.3853,
            0.4309,
            38.1078,
            2.1098,
            0.3745,
        ],
        [
            0.5426,
            0.5560,
            0.4860,
            0.3457,
            0.4342,
            0.6111,
            0.4965,
            0.3487,
            1.7979,
            0.6304,
            0.6921,
            0.5322,
            0.7084,
            2.7694,
            0.3635,
            0.5575,
            0.5732,
            2.1098,
            9.8322,
            0.6580,
        ],
        [
            0.9365,
            0.4201,
            0.3690,
            0.3365,
            0.7558,
            0.4668,
            0.4289,
            0.3370,
            0.3394,
            2.4175,
            1.3142,
            0.4565,
            1.2689,
            0.7451,
            0.4431,
            0.5652,
            0.9809,
            0.3745,
            0.6580,
            3.6922,
        ],
    ]
)


def get_gfp_base_seq():
    """Returns the wild type GFP sequence"""
    f = open(
        os.path.join(
            os.path.dirname(__file__), "ref", "data", "avGFP_reference_sequence.fa"
        )
    )
    lines = f.readlines()
    f.close()
    seq = lines[1].strip()
    return seq


def read_gfp_data(path=None, df_save_file=None):
    """Reads the GFP brightness data in a pandas DataFrame"""
    if path is None:
        path = "../data/nucleotide_genotypes_to_brightness.tsv"
    f = open(path)
    base_seq = get_base_seq()
    mod = list(base_seq)
    mod[191] = "A"
    base_seq = "".join(mod)
    data = []
    i = 0
    cols = None
    for line in f:
        if i == 0:
            cols = line.strip().split("\t")
            cols = ["nucSequence", "numNucMutations", "numAAMutations"] + cols[2:]
            i += 1
            continue
        else:
            split = line.split("\t")
            mutations = split[0].split(":")
            seq, n_nuc_mut = convert_mutations_to_sequence(base_seq, mutations)
            n_aa_mut = len([m for m in split[1].split(":") if m != ""])
            if split[-1].strip() == "":
                std = 0
            else:
                std = float(split[-1].strip())
            data.append([seq, n_nuc_mut, n_aa_mut, int(split[2]), float(split[3]), std])
            i += 1
    df = pd.DataFrame(data, columns=cols)
    if df_save_file is not None:
        df.to_csv(df_save_file)
    return df


def convert_mutations_to_sequence(base_seq, mutations):
    """
    Given the wild type sequence and a formatted mtuation string, returns
    the mutated sequence
    """
    new_seq = list(base_seq)
    n = 0
    for m in mutations:
        if len(m) == 0:
            continue
        style = m[0]
        assert style == "S"
        nuc1 = m[1]
        nuc2 = m[-1]
        pos = int(m[2:-1])
        assert new_seq[pos] == nuc1
        new_seq[pos] = nuc2
        n += 1
    new_seq = "".join(new_seq)
    return new_seq, n


def one_hot_encode_dna(dna_str, pad=None, base_order="ATCG"):
    """Convert length M string into M x 4 tokenized array"""
    dna_str = dna_str.upper()
    if pad is not None:
        M = pad
    else:
        M = len(dna_str)
    dna_arr = np.zeros((M, 4))
    for i in range(len(dna_str)):
        idx = base_order.index(dna_str[i])
        dna_arr[i, idx] = 1
    return dna_arr


def one_hot_encode_aa_array(X_aa: np.ndarray) -> np.ndarray:
    """
    OneHot encodes array: (batch_size, L) -> (batch_size, L, alphabet_size)
    """
    oh_arr = np.zeros((X_aa.shape[0], X_aa.shape[1], len(AA)), dtype=int)
    for i in range(X_aa.shape[0]):
        for j in range(X_aa.shape[1]):
            oh_arr[i, j, AA_IDX[X_aa[i, j]]] = 1
    # test if first and last entry match
    np.testing.assert_array_equal(
        np.array([AA[idx] for idx in np.argmax(oh_arr[0], axis=1)]), X_aa[0]
    )
    np.testing.assert_array_equal(
        np.array([AA[idx] for idx in np.argmax(oh_arr[-1], axis=1)]), X_aa[-1]
    )
    return oh_arr


def one_hot_encode_aa(aa_str, pad=None):
    """Returns a one hot encoded amino acid sequence"""
    M = len(aa_str)
    aa_arr = np.zeros((M, 20), dtype=int)
    for i in range(M):
        aa_arr[i, AA_IDX[aa_str[i]]] = 1
    return aa_arr


def convert_aas_to_idx_array(X_aa):
    """
    Converts a list of amino acid sequences into an array of amino acid indices from AA_IDX.
    E.g. the amino sequence ARN becomes a row containing [1, 2, 3]
    """
    N = len(X_aa)
    M = len(X_aa[0])
    X_aa_idx = np.zeros((N, M), dtype=int)
    for i in range(N):
        for j in range(M):
            X_aa_idx[i, j] = AA_IDX[X_aa[i][j]]

    return X_aa_idx


def convert_idx_array_to_aas(X_aa):
    """Converts an array containing indices of amino acids into the
    corresponding string amino acid sequences. E.g. a row with [1, 2, 3]
    becomes the sequence ARN."""
    N = len(X_aa)
    M = len(X_aa[0])
    X_aa_str = [["A"] * M] * N
    for i in range(N):
        for j in range(M):
            X_aa_str[i][j] = AA[X_aa[i, j]]
        X_aa_str[i] = "".join(X_aa_str[i])
    return X_aa_str


def get_argmax(Xt_p):
    """Given a categorical probability distribution specifying the probability
    of amino acids at each position in a sequence, returns the most probable sequence"""
    Xt_argmax = np.zeros_like(Xt_p)
    Xt_argmax[
        np.arange(Xt_p.shape[0]).reshape(Xt_p.shape[0], 1),
        np.arange(Xt_p.shape[1]).reshape(1, Xt_p.shape[1]),
        np.argmax(Xt_p, axis=-1),
    ] = 1
    return Xt_argmax


def get_samples(Xt_p):
    """Samples from a categorical probability distribution specifying the probability
    of amino acids at each position in a sequence"""
    Xt_sampled = np.zeros_like(Xt_p)
    for i in range(Xt_p.shape[0]):
        for j in range(Xt_p.shape[1]):
            p = Xt_p[i, j]
            k = np.random.choice(range(len(p)), p=p)
            Xt_sampled[i, j, k] = 1.0
    return Xt_sampled


def get_balaji_predictions(preds, Xt):
    """Given a set of predictors built according to the methods in
    the Balaji Lakshminarayanan paper 'Simple and scalable predictive
    uncertainty estimation using deep ensembles' (2017), returns the mean and
    variance of the total prediction."""
    M = len(preds)
    N = Xt.shape[0]
    means = np.zeros((M, N))
    variances = np.zeros((M, N))
    for m in range(M):
        y_pred = preds[m].predict(Xt)
        means[m, :] = y_pred[:, 0]
        variances[m, :] = np.log(1 + np.exp(y_pred[:, 1])) + 1e-6
    mu_star = np.mean(means, axis=0)
    var_star = (1 / M) * (
        np.sum(variances, axis=0) + np.sum(means**2, axis=0)
    ) - mu_star**2
    return mu_star, var_star


def partition_data(
    X, y, percentile=40, train_size=1000, random_state=1, return_test=False
):
    """Partition a (X, y) data set by a percentile of the y values"""
    np.random.seed(random_state)
    assert percentile * 0.01 * len(y) >= train_size
    y_percentile = np.percentile(y, percentile)
    idx = np.where(y < y_percentile)[0]
    rand_idx = np.random.choice(idx, size=train_size, replace=False)
    X_train = X[rand_idx]
    y_train = y[rand_idx]
    if return_test:
        test_idx = [i for i in idx if i not in rand_idx]
        X_test = X[test_idx]
        y_test = y[test_idx]
        return X_train, y_train, X_test, y_test
    else:
        return X_train, y_train


def get_experimental_X_y(
    random_state=1, train_size=5000, return_test=False, return_all=False
):
    """For the GFP testing experiments. Loads the ground truth data (i.e. predictions from the GP model
    we use as the ground truth), partitions it such that we only observe values below the 20th percentile
    and adds measurement noise"""
    prefix = os.path.join(os.path.dirname(os.path.abspath(__file__)), "ref")
    df = pd.read_csv(prefix + "/data/gfp_data.csv")
    X, _ = get_gfp_X_y_aa(df, large_only=True, ignore_stops=True)
    y_gt = np.load(prefix + "/data/gfp_gt_evals.npy")
    if return_test:
        X_train, gt_train, X_test, gt_test = partition_data(
            X,
            y_gt,
            percentile=20,
            train_size=train_size,
            random_state=random_state,
            return_test=return_test,
        )
        np.random.seed(random_state)
        gt_var = 0.01
        y_train = gt_train + np.random.randn(*gt_train.shape) * gt_var
        y_test = gt_test + np.random.randn(*gt_test.shape) * gt_var
        return X_train, y_train, gt_train, X_test, y_test, gt_test
    else:
        X_train, gt_train = partition_data(
            X,
            y_gt,
            percentile=20,
            train_size=train_size,
            random_state=random_state,
            return_test=return_test,
        )
        np.random.seed(random_state)
        gt_var = 0.01
        y_train = gt_train + np.random.randn(*gt_train.shape) * gt_var
        return X_train, y_train, gt_train


def get_gfp_X_y_aa(
    data_df: pd.DataFrame,
    functional_only: bool = False,
    ignore_stops: bool = True,
    return_str: bool = False,
):
    """
    Converts the raw GFP data to a set of X and y values that are ready to use
    in a model
    """
    if functional_only:
        idx = data_df.loc[
            (data_df["medianBrightness"] > data_df["medianBrightness"].mean())
        ].index
    else:
        idx = data_df.index
    data_df = data_df.loc[idx]

    if ignore_stops:
        idx = data_df.loc[~data_df["aaSequence"].str.contains("!")].index
    data_df = data_df.loc[idx]
    seqs = data_df["aaSequence"]

    M = len(seqs[0])
    N = len(seqs)
    X = np.zeros((N, M, 20))
    j = 0
    for i in idx:
        X[j] = one_hot_encode_aa(seqs[i])
        j += 1
    y_raw = np.array(data_df["medianBrightness"][idx])
    y = y_raw
    if return_str:
        return X, y, list(data_df["aaSequence"])
    else:
        return X, y
