import tensorflow as tf
import os
import pickle
import numpy as np

from constant_params import input_feature_dim, window_size


def build_dataset(input_tfrecord_files, batch_size):
    drop_remainder = False

    feature_description = {
        'label': tf.io.FixedLenFeature([], tf.int64),
        'ref_aa': tf.io.FixedLenFeature([], tf.int64),
        'alt_aa': tf.io.FixedLenFeature([], tf.int64),
        'feature': tf.io.FixedLenFeature([], tf.string),
        'mask': tf.io.FixedLenFeature([], tf.string),
        'var_id': tf.io.FixedLenFeature([], tf.string),
    }

    def _parser(example_proto):
        parsed = tf.io.parse_single_example(example_proto, feature_description)
        label, ref_aa, alt_aa = parsed['label'], parsed['ref_aa'], parsed[
            'alt_aa']
        var_id = parsed['var_id']

        ref_aa, alt_aa, label = tf.cast(ref_aa, tf.int32), tf.cast(
            alt_aa, tf.int32), tf.cast(label, tf.float32)

        feature = tf.io.decode_raw(parsed['feature'], tf.float32)
        feature = tf.reshape(feature, (window_size, input_feature_dim))

        mask = tf.io.decode_raw(parsed['mask'], tf.float32)
        mask = tf.reshape(mask, (window_size, ))
        h = window_size // 2

        #mask the postion of interest
        mask = tf.concat(
            [mask[:h],
             tf.cast([
                 1,
             ], dtype=tf.float32), mask[h + 1:]],
            axis=-1)
        '''
        pos_encoding = 1.0 + tf.cast(
            tf.math.abs(window_size // 2 - tf.range(window_size)),
            dtype=tf.float32)
        #pos_encoding = tf.math.log() / tf.math.log(2.0)
        feature = tf.concat([feature, pos_encoding[:, tf.newaxis]], axis=-1)
        '''

        return var_id, ref_aa, alt_aa, feature, label, mask

    dataset = tf.data.TFRecordDataset(input_tfrecord_files)

    options = tf.data.Options()
    options.experimental_threading.max_intra_op_parallelism = 1
    dataset = dataset.with_options(options)

    dataset = dataset.shuffle(2048)
    dataset = dataset.map(_parser, num_parallel_calls=8)

    dataset = dataset.batch(batch_size)

    #dataset = dataset.prefetch(4)

    return dataset


def build_all_possible_missenses_dataset(tr_list, feature_dir, batch_size):
    amino_acid_order = 'ACDEFGHIKLMNPQRSTVWY*'

    def _gen_data():
        for transcript_id in tr_list:
            feature_path = f'{feature_dir}/{transcript_id}.pickle'
            if not os.path.exists(feature_path):
                continue
            print(feature_path, flush=True)

            with open(feature_path, 'rb') as fr:
                feature = pickle.load(fr)

            L = feature.shape[0]

            w = window_size // 2

            for aa_pos in range(L):
                ref_aa = int(feature[aa_pos, 0])

                start = max(aa_pos - w, 0)
                end = min(L, aa_pos + 1 + w)
                var_start = start - (aa_pos - w)
                var_end = var_start + (end - start)
                var_feature = np.zeros([w * 2 + 1, feature.shape[1]])
                var_feature[var_start:var_end] = feature[start:end]

                mask = np.ones((w * 2 + 1, ), dtype=np.float32)
                mask[var_start:var_end] = 0.0
                mask[w] = 1.0

                for alt_aa in range(20):
                    var_id = f'{transcript_id}_{str(aa_pos+1)}_{amino_acid_order[ref_aa]}_{amino_acid_order[alt_aa]}'.encode(
                        'utf-8')
                    yield var_id, np.int32(ref_aa), np.int32(
                        alt_aa), np.float32(var_feature), np.float32(mask)

    dataset = tf.data.Dataset.from_generator(
        _gen_data, (tf.string, tf.int32, tf.int32, tf.float32, tf.float32),
        (tf.TensorShape(()), tf.TensorShape(()), tf.TensorShape(
            ()), tf.TensorShape((window_size, input_feature_dim)),
         tf.TensorShape((window_size, ))))

    options = tf.data.Options()
    options.experimental_threading.max_intra_op_parallelism = 1
    dataset = dataset.with_options(options)

    #dataset = dataset.map(_parser, num_parallel_calls=8)

    dataset = dataset.batch(batch_size)

    dataset = dataset.prefetch(4)

    return dataset


def build_test_dataset(input_tfrecord_files, batch_size):
    drop_remainder = False

    feature_description = {
        'ref_aa': tf.io.FixedLenFeature([], tf.int64),
        'alt_aa': tf.io.FixedLenFeature([], tf.int64),
        'feature': tf.io.FixedLenFeature([], tf.string),
        'mask': tf.io.FixedLenFeature([], tf.string),
        'var_id': tf.io.FixedLenFeature([], tf.string),
    }

    def _parser(example_proto):
        parsed = tf.io.parse_single_example(example_proto, feature_description)
        ref_aa, alt_aa = parsed['ref_aa'], parsed['alt_aa']
        var_id = parsed['var_id']

        ref_aa, alt_aa = tf.cast(ref_aa, tf.int32), tf.cast(alt_aa, tf.int32)

        feature = tf.io.decode_raw(parsed['feature'], tf.float32)
        feature = tf.reshape(feature, (window_size, input_feature_dim))

        mask = tf.io.decode_raw(parsed['mask'], tf.float32)
        mask = tf.reshape(mask, (window_size, ))
        h = window_size // 2

        #mask the postion of interest
        mask = tf.concat(
            [mask[:h],
             tf.cast([
                 1,
             ], dtype=tf.float32), mask[h + 1:]],
            axis=-1)

        return var_id, ref_aa, alt_aa, feature, mask

    dataset = tf.data.TFRecordDataset(input_tfrecord_files)

    options = tf.data.Options()
    options.experimental_threading.max_intra_op_parallelism = 1
    dataset = dataset.with_options(options)

    dataset = dataset.map(_parser, num_parallel_calls=8)

    dataset = dataset.batch(batch_size)

    #dataset = dataset.prefetch(4)

    return dataset

[((1795, 1840), 'tensorflow.data.TFRecordDataset', 'tf.data.TFRecordDataset', (['input_tfrecord_files'], {}), '(input_tfrecord_files)\n', (1818, 1840), True, 'import tensorflow as tf\n'), ((1856, 1873), 'tensorflow.data.Options', 'tf.data.Options', ([], {}), '()\n', (1871, 1873), True, 'import tensorflow as tf\n'), ((3909, 3926), 'tensorflow.data.Options', 'tf.data.Options', ([], {}), '()\n', (3924, 3926), True, 'import tensorflow as tf\n'), ((5424, 5469), 'tensorflow.data.TFRecordDataset', 'tf.data.TFRecordDataset', (['input_tfrecord_files'], {}), '(input_tfrecord_files)\n', (5447, 5469), True, 'import tensorflow as tf\n'), ((5485, 5502), 'tensorflow.data.Options', 'tf.data.Options', ([], {}), '()\n', (5500, 5502), True, 'import tensorflow as tf\n'), ((255, 290), 'tensorflow.io.FixedLenFeature', 'tf.io.FixedLenFeature', (['[]', 'tf.int64'], {}), '([], tf.int64)\n', (276, 290), True, 'import tensorflow as tf\n'), ((310, 345), 'tensorflow.io.FixedLenFeature', 'tf.io.FixedLenFeature', (['[]', 'tf.int64'], {}), '([], tf.int64)\n', (331, 345), True, 'import tensorflow as tf\n'), ((365, 400), 'tensorflow.io.FixedLenFeature', 'tf.io.FixedLenFeature', (['[]', 'tf.int64'], {}), '([], tf.int64)\n', (386, 400), True, 'import tensorflow as tf\n'), ((421, 457), 'tensorflow.io.FixedLenFeature', 'tf.io.FixedLenFeature', (['[]', 'tf.string'], {}), '([], tf.string)\n', (442, 457), True, 'import tensorflow as tf\n'), ((475, 511), 'tensorflow.io.FixedLenFeature', 'tf.io.FixedLenFeature', (['[]', 'tf.string'], {}), '([], tf.string)\n', (496, 511), True, 'import tensorflow as tf\n'), ((531, 567), 'tensorflow.io.FixedLenFeature', 'tf.io.FixedLenFeature', (['[]', 'tf.string'], {}), '([], tf.string)\n', (552, 567), True, 'import tensorflow as tf\n'), ((625, 687), 'tensorflow.io.parse_single_example', 'tf.io.parse_single_example', (['example_proto', 'feature_description'], {}), '(example_proto, feature_description)\n', (651, 687), True, 'import tensorflow as tf\n'), ((965, 1012), 'tensorflow.io.decode_raw', 'tf.io.decode_raw', (["parsed['feature']", 'tf.float32'], {}), "(parsed['feature'], tf.float32)\n", (981, 1012), True, 'import tensorflow as tf\n'), ((1031, 1084), 'tensorflow.reshape', 'tf.reshape', (['feature', '(window_size, input_feature_dim)'], {}), '(feature, (window_size, input_feature_dim))\n', (1041, 1084), True, 'import tensorflow as tf\n'), ((1101, 1145), 'tensorflow.io.decode_raw', 'tf.io.decode_raw', (["parsed['mask']", 'tf.float32'], {}), "(parsed['mask'], tf.float32)\n", (1117, 1145), True, 'import tensorflow as tf\n'), ((1161, 1193), 'tensorflow.reshape', 'tf.reshape', (['mask', '(window_size,)'], {}), '(mask, (window_size,))\n', (1171, 1193), True, 'import tensorflow as tf\n'), ((4324, 4359), 'tensorflow.io.FixedLenFeature', 'tf.io.FixedLenFeature', (['[]', 'tf.int64'], {}), '([], tf.int64)\n', (4345, 4359), True, 'import tensorflow as tf\n'), ((4379, 4414), 'tensorflow.io.FixedLenFeature', 'tf.io.FixedLenFeature', (['[]', 'tf.int64'], {}), '([], tf.int64)\n', (4400, 4414), True, 'import tensorflow as tf\n'), ((4435, 4471), 'tensorflow.io.FixedLenFeature', 'tf.io.FixedLenFeature', (['[]', 'tf.string'], {}), '([], tf.string)\n', (4456, 4471), True, 'import tensorflow as tf\n'), ((4489, 4525), 'tensorflow.io.FixedLenFeature', 'tf.io.FixedLenFeature', (['[]', 'tf.string'], {}), '([], tf.string)\n', (4510, 4525), True, 'import tensorflow as tf\n'), ((4545, 4581), 'tensorflow.io.FixedLenFeature', 'tf.io.FixedLenFeature', (['[]', 'tf.string'], {}), '([], tf.string)\n', (4566, 4581), True, 'import tensorflow as tf\n'), ((4639, 4701), 'tensorflow.io.parse_single_example', 'tf.io.parse_single_example', (['example_proto', 'feature_description'], {}), '(example_proto, feature_description)\n', (4665, 4701), True, 'import tensorflow as tf\n'), ((4894, 4941), 'tensorflow.io.decode_raw', 'tf.io.decode_raw', (["parsed['feature']", 'tf.float32'], {}), "(parsed['feature'], tf.float32)\n", (4910, 4941), True, 'import tensorflow as tf\n'), ((4960, 5013), 'tensorflow.reshape', 'tf.reshape', (['feature', '(window_size, input_feature_dim)'], {}), '(feature, (window_size, input_feature_dim))\n', (4970, 5013), True, 'import tensorflow as tf\n'), ((5030, 5074), 'tensorflow.io.decode_raw', 'tf.io.decode_raw', (["parsed['mask']", 'tf.float32'], {}), "(parsed['mask'], tf.float32)\n", (5046, 5074), True, 'import tensorflow as tf\n'), ((5090, 5122), 'tensorflow.reshape', 'tf.reshape', (['mask', '(window_size,)'], {}), '(mask, (window_size,))\n', (5100, 5122), True, 'import tensorflow as tf\n'), ((852, 877), 'tensorflow.cast', 'tf.cast', (['ref_aa', 'tf.int32'], {}), '(ref_aa, tf.int32)\n', (859, 877), True, 'import tensorflow as tf\n'), ((879, 904), 'tensorflow.cast', 'tf.cast', (['alt_aa', 'tf.int32'], {}), '(alt_aa, tf.int32)\n', (886, 904), True, 'import tensorflow as tf\n'), ((919, 945), 'tensorflow.cast', 'tf.cast', (['label', 'tf.float32'], {}), '(label, tf.float32)\n', (926, 945), True, 'import tensorflow as tf\n'), ((3728, 3746), 'tensorflow.TensorShape', 'tf.TensorShape', (['()'], {}), '(())\n', (3742, 3746), True, 'import tensorflow as tf\n'), ((3748, 3766), 'tensorflow.TensorShape', 'tf.TensorShape', (['()'], {}), '(())\n', (3762, 3766), True, 'import tensorflow as tf\n'), ((3768, 3786), 'tensorflow.TensorShape', 'tf.TensorShape', (['()'], {}), '(())\n', (3782, 3786), True, 'import tensorflow as tf\n'), ((3801, 3849), 'tensorflow.TensorShape', 'tf.TensorShape', (['(window_size, input_feature_dim)'], {}), '((window_size, input_feature_dim))\n', (3815, 3849), True, 'import tensorflow as tf\n'), ((3860, 3890), 'tensorflow.TensorShape', 'tf.TensorShape', (['(window_size,)'], {}), '((window_size,))\n', (3874, 3890), True, 'import tensorflow as tf\n'), ((4822, 4847), 'tensorflow.cast', 'tf.cast', (['ref_aa', 'tf.int32'], {}), '(ref_aa, tf.int32)\n', (4829, 4847), True, 'import tensorflow as tf\n'), ((4849, 4874), 'tensorflow.cast', 'tf.cast', (['alt_aa', 'tf.int32'], {}), '(alt_aa, tf.int32)\n', (4856, 4874), True, 'import tensorflow as tf\n'), ((1325, 1355), 'tensorflow.cast', 'tf.cast', (['[1]'], {'dtype': 'tf.float32'}), '([1], dtype=tf.float32)\n', (1332, 1355), True, 'import tensorflow as tf\n'), ((2444, 2472), 'os.path.exists', 'os.path.exists', (['feature_path'], {}), '(feature_path)\n', (2458, 2472), False, 'import os\n'), ((2619, 2634), 'pickle.load', 'pickle.load', (['fr'], {}), '(fr)\n', (2630, 2634), False, 'import pickle\n'), ((3009, 3048), 'numpy.zeros', 'np.zeros', (['[w * 2 + 1, feature.shape[1]]'], {}), '([w * 2 + 1, feature.shape[1]])\n', (3017, 3048), True, 'import numpy as np\n'), ((3141, 3180), 'numpy.ones', 'np.ones', (['(w * 2 + 1,)'], {'dtype': 'np.float32'}), '((w * 2 + 1,), dtype=np.float32)\n', (3148, 3180), True, 'import numpy as np\n'), ((5254, 5284), 'tensorflow.cast', 'tf.cast', (['[1]'], {'dtype': 'tf.float32'}), '([1], dtype=tf.float32)\n', (5261, 5284), True, 'import tensorflow as tf\n'), ((3493, 3509), 'numpy.int32', 'np.int32', (['ref_aa'], {}), '(ref_aa)\n', (3501, 3509), True, 'import numpy as np\n'), ((3511, 3527), 'numpy.int32', 'np.int32', (['alt_aa'], {}), '(alt_aa)\n', (3519, 3527), True, 'import numpy as np\n'), ((3554, 3577), 'numpy.float32', 'np.float32', (['var_feature'], {}), '(var_feature)\n', (3564, 3577), True, 'import numpy as np\n'), ((3579, 3595), 'numpy.float32', 'np.float32', (['mask'], {}), '(mask)\n', (3589, 3595), True, 'import numpy as np\n')]