import re
import os
import sys
import RNA
import gzip
import pickle
import random
import subprocess
import numpy as np
import scipy.sparse as sp
from functools import partial
import forgi.graph.bulge_graph as fgb

basedir = os.path.split(os.path.dirname(os.path.abspath(__file__)))[0]
sys.path.append(basedir)
sys.path.append(os.path.dirname(os.path.abspath(__file__)))

from lib.general_utils import Pool

def adj_to_bias(adj, nhood=1):
    # [batch_size, nb_nodes, nb_nodes]
    mt = np.stack([np.eye(adj.shape[1])] * adj.shape[0], axis=0)  # self-connection
    for _ in range(nhood):
        mt = np.matmul(mt, (adj + np.stack([np.eye(adj.shape[1])] * adj.shape[0], axis=0)))
    mt = np.greater(mt, 0.).astype(np.float32)
    return -1e9 * (1.0 - mt)

# >>>> equilibrium probability using RNAplfold >>>>>>>

def fold_seq_rnashapes(seq, winsize, iterations=100):
    stride = winsize // 4
    cmd = 'echo %s | RNAshapes -w %d -W %d -i %d -A -t 1 -c %d -M 0 -o 1' % (seq, winsize, stride, iterations, 10)
    ret = subprocess.check_output(cmd, shell=True)
    lines = re.sub(' +', ' ', ret.decode('utf-8')).rstrip().split('\n')

    # assemble adjacency matrix
    row_col, link, prob, norm = [], [], [], []
    length = len(seq)
    for i in range(length):
        if i != length - 1:
            row_col.append((i, i + 1))
            link.append(1)
            prob.append(1.)
            norm.append(1)
        if i != 0:
            row_col.append((i, i - 1))
            link.append(2)
            prob.append(1.)
            norm.append(1)
    for line in lines:
        if len(line) > 0:
            if line[0] >= '0' and line[0] <= '9':
                # indices
                start_idx = int(line.split(' ')[0]) - 1
                local_idx = []
            elif line[0] in ['(', '.', ')']:
                # secondary structures
                tokens = line.split(' ')
                struct = tokens[0]
                probability = float(tokens[2])

                bg = fgb.BulgeGraph.from_dotbracket(struct)
                for i, ele in enumerate(struct):
                    if ele == '(':
                        pair_from = i + start_idx
                        pair_to = bg.pairing_partner(i + 1) - 1 + start_idx
                        if not (pair_from, pair_to) in row_col:
                            row_col.append((pair_from, pair_to))
                            link.append(3)
                            prob.append(probability)
                            norm.append(1)
                            local_idx.append((pair_from, pair_to))
                            # symmetric
                            row_col.append((pair_to, pair_from))
                            link.append(4)
                            prob.append(probability)
                            norm.append(1)
                        else:
                            idx = row_col.index((pair_from, pair_to))
                            prob[idx] += probability
                            prob[idx + 1] += probability
                            if not (pair_from, pair_to) in local_idx:
                                local_idx.append((pair_from, pair_to))
                                norm[idx] += 1
                                norm[idx + 1] += 1

    prob = np.array(prob) / np.array(norm)
    print(norm)
    return (sp.csr_matrix((link, (np.array(row_col)[:, 0], np.array(row_col)[:, 1])), shape=(length, length)),
            sp.csr_matrix((prob, (np.array(row_col)[:, 0], np.array(row_col)[:, 1])), shape=(length, length)),)

# >>>> equilibrium probability using RNAplfold >>>>>>>
# when on compute canada, make sure this is happening on the compute nodes

def fold_seq_rnaplfold(seq, w, l, cutoff, no_lonely_bps):
    np.random.seed(random.seed())
    name = str(np.random.rand())
    # Call RNAplfold on command line.
    no_lonely_bps_str = ""
    if no_lonely_bps:
        no_lonely_bps_str = "--noLP"
    cmd = 'echo %s | RNAplfold -W %d -L %d -c %.4f --id-prefix %s %s' % (seq, w, l, cutoff, name, no_lonely_bps_str)
    ret = subprocess.call(cmd, shell=True)

    # assemble adjacency matrix
    row_col, link, prob = [], [], []
    length = len(seq)
    for i in range(length):
        if i != length - 1:
            row_col.append((i, i + 1))
            link.append(1)
            prob.append(1.)
        if i != 0:
            row_col.append((i, i - 1))
            link.append(2)
            prob.append(1.)
    # Extract base pair information.
    name += '_0001_dp.ps'
    start_flag = False
    with open(name) as f:
        for line in f:
            if start_flag:
                values = line.split()
                if len(values) == 4:
                    source_id = int(values[0]) - 1
                    dest_id = int(values[1]) - 1
                    avg_prob = float(values[2])
                    # source_id < dest_id
                    row_col.append((source_id, dest_id))
                    link.append(3)
                    prob.append(avg_prob ** 2)
                    row_col.append((dest_id, source_id))
                    link.append(4)
                    prob.append(avg_prob ** 2)
            if 'start of base pair probability data' in line:
                start_flag = True
    # delete RNAplfold output file.
    os.remove(name)
    # placeholder for dot-bracket structure
    return (sp.csr_matrix((link, (np.array(row_col)[:, 0], np.array(row_col)[:, 1])), shape=(length, length)),
            sp.csr_matrix((prob, (np.array(row_col)[:, 0], np.array(row_col)[:, 1])), shape=(length, length)),)

# >>>> Boltzmann sampling using RNAsubopt >>>>>>>

def sample_one_seq(seq, passes=10):
    max_len = len(seq)
    adj_mat_tmp = np.zeros((max_len, max_len), dtype=np.int32)
    for i in range(max_len):
        if i != max_len - 1:
            adj_mat_tmp[i, i + 1] = 1
        if i != 0:
            adj_mat_tmp[i, i - 1] = 2
    cmd = 'echo "%s" | RNAsubopt --stochBT=%d' % (seq, passes)
    all_struct = subprocess.check_output(cmd, shell=True). \
                     decode('utf-8').rstrip().split('\n')[1:]
    all_adj_mat = []
    for struct in all_struct:
        adj_mat = adj_mat_tmp.copy()
        bg = fgb.BulgeGraph.from_dotbracket(struct)
        for i, ele in enumerate(struct):
            if ele == '(':
                adj_mat[i, bg.pairing_partner(i + 1) - 1] = 3
            elif ele == ')':
                adj_mat[i, bg.pairing_partner(i + 1) - 1] = 4
        all_adj_mat.append(adj_mat)
    all_adj_mat = np.stack(all_adj_mat, axis=0)
    return all_adj_mat

def fold_seq_subopt(seq, probabilistic=False, sampling_amount=1000):
    # RNAfold is only suitable for short RNA sequences within 100 nucleotides
    if probabilistic:
        # sampling from a boltzmann ensemble
        cmd = 'echo "%s" | RNAsubopt --stochBT=%d' % (seq, sampling_amount)
        struct_list = subprocess.check_output(cmd, shell=True). \
                          decode('utf-8').rstrip().split('\n')[1:]
    else:
        struct_list, energy_list = [], []

        def collect_subopt_result(structure, energy, *args):
            if not structure == None:
                struct_list.append(structure)
                energy_list.append(energy)

        # Enumerate all structures 100 dacal/mol = 1 kcal/mol around
        # default deltaEnergy is the MFE
        RNA.fold_compound(seq).subopt_cb(100, collect_subopt_result, None)

        # sort
        struct_list = list(np.array(struct_list)[np.argsort(energy_list)])

    # merging all structures into a single adjacency matrix
    # probability returning two matrices
    matrix = adj_mat_subopt(struct_list, probabilistic)
    # process the structures
    return structural_content(struct_list), matrix

def structural_content(struct_list):
    size = len(struct_list)
    length = len(struct_list[0])
    content = np.zeros((length, 3), dtype=np.int32)
    for i in range(length):
        for j in range(size):
            idx = '.()'.index(struct_list[j][i])
            content[i][idx] += 1
    return content.astype(np.float32) / size

def adj_mat_subopt(struct_list, probabilistic):
    # create sparse matrix
    row_col, data = [], []
    length = len(struct_list[0])
    counts = []
    for i in range(length):
        if i != length - 1:
            row_col.append((i, i + 1))
            data.append(1)
            counts.append(0)
        if i != 0:
            row_col.append((i, i - 1))
            data.append(2)
            counts.append(0)
    if probabilistic:
        for struct in struct_list:
            bg = fgb.BulgeGraph.from_dotbracket(struct)
            for i, ele in enumerate(struct):
                if ele == '(':
                    if not (i, bg.pairing_partner(i + 1) - 1) in row_col:
                        row_col.append((i, bg.pairing_partner(i + 1) - 1))
                        data.append(3)
                        counts.append(1)
                    else:
                        idx = row_col.index((i, bg.pairing_partner(i + 1) - 1))
                        counts[idx] += 1
                elif ele == ')':
                    if not (i, bg.pairing_partner(i + 1) - 1) in row_col:
                        row_col.append((i, bg.pairing_partner(i + 1) - 1))
                        data.append(4)
                        counts.append(1)
                    else:
                        idx = row_col.index((i, bg.pairing_partner(i + 1) - 1))
                        counts[idx] += 1
        # normalize each row into probabilities
        for i in range(len(row_col)):
            if counts[i] > 0:
                # have to be a hydrogen bond
                counts[i] /= len(struct_list)
            else:
                # covalent bond that forms the stem
                counts[i] = 1.
        return (sp.csr_matrix((data, (np.array(row_col)[:, 0], np.array(row_col)[:, 1])), shape=(length, length)),
                sp.csr_matrix((counts, (np.array(row_col)[:, 0], np.array(row_col)[:, 1])), shape=(length, length)),)
    else:
        for struct in struct_list:
            bg = fgb.BulgeGraph.from_dotbracket(struct)
            for i, ele in enumerate(struct):
                if ele == '(':
                    if not (i, bg.pairing_partner(i + 1) - 1) in row_col:
                        row_col.append((i, bg.pairing_partner(i + 1) - 1))
                        data.append(3)
                elif ele == ')':
                    if not (i, bg.pairing_partner(i + 1) - 1) in row_col:
                        row_col.append((i, bg.pairing_partner(i + 1) - 1))
                        data.append(4)
        return sp.csr_matrix((data, (np.array(row_col)[:, 0], np.array(row_col)[:, 1])),
                             shape=(length, length))

# >>>> MFE structure using RNAfold >>>>>>>

def fold_seq_rnafold(seq):
    '''fold sequence using RNAfold'''
    struct = RNA.fold(seq)[0]
    matrix = adj_mat(struct)
    return structural_content([struct]), matrix

def adj_mat(struct):
    # create sparse matrix
    row_col, data = [], []
    length = len(struct)
    for i in range(length):
        if i != length - 1:
            row_col.append((i, i + 1))
            data.append(1)
        if i != 0:
            row_col.append((i, i - 1))
            data.append(2)
    bg = fgb.BulgeGraph.from_dotbracket(struct)
    for i, ele in enumerate(struct):
        if ele == '(':
            row_col.append((i, bg.pairing_partner(i + 1) - 1))
            data.append(3)
        elif ele == ')':
            row_col.append((i, bg.pairing_partner(i + 1) - 1))
            data.append(4)
    return sp.csr_matrix((data, (np.array(row_col)[:, 0], np.array(row_col)[:, 1])),
                         shape=(length, length))

def augment_features(path):
    '''
    try to avoid this as much as possible.
    we are mostly interested in an end-to-end learning scenario
    '''
    # region type: 101 x 5
    region_types = np.loadtxt(gzip.open(os.path.join(path, "matrix_RegionType.tab.gz")), skiprows=1)
    assert (region_types.shape[1] == 505)  # 4 region types
    region_types = np.transpose(region_types.reshape((region_types.shape[0], 5, 101)), [0, 2, 1])

    coclip = np.loadtxt(gzip.open(os.path.join(path, "matrix_Cobinding.tab.gz")), skiprows=1)
    assert (coclip.shape[1] % 101 == 0)
    nb_exprs = coclip.shape[1] // 101
    coclip = np.transpose(coclip.reshape((coclip.shape[0], nb_exprs, 101)), [0, 2, 1])

    return np.concatenate([region_types, coclip], axis=-1)

def load_fasta_format(file):
    all_id = []
    all_seq = []
    seq = ''
    for row in file:
        if type(row) is bytes:
            row = row.decode('utf-8')
        row = row.rstrip()
        if row.startswith('>'):
            all_id.append(row)
            if seq != '':
                all_seq.append(seq)
                seq = ''
        else:
            seq += row
    all_seq.append(seq)
    return all_id, all_seq

def load_dotbracket(filepath, pool=None, fold_algo='rnafold', probabilistic=False, **kwargs):
    prefix = '%s_%s_' % (fold_algo, probabilistic)
    if fold_algo == 'rnaplfold' or fold_algo == 'rnashapes':
        prefix += '%d_' % (kwargs.get('w', 150))
    full_path = os.path.join(os.path.dirname(filepath), '{}structures.npy'.format(prefix))
    if not os.path.exists(full_path):
        print(full_path, 'is missing. Begin folding from scratch.')
        fold_rna_from_file(filepath, pool, fold_algo, probabilistic, **kwargs)
    # load secondary structures
    all_struct = np.load(
        os.path.join(os.path.dirname(filepath), '{}structures.npy'.format(prefix)))
    return all_struct

def load_mat(filepath, pool=None, fold_algo='rnafold', probabilistic=False, **kwargs):
    prefix = '%s_%s_' % (fold_algo, probabilistic)
    # folding length is crucial hyperparam for local RNA folding, therefore should be marked
    if fold_algo == 'rnaplfold' or fold_algo == 'rnashapes':
        prefix += '%d_' % (kwargs.get('w', 150))
    if kwargs.get('modify_leaks', False):
        # this will make sure we load the modified sequences
        # incorrect secondary structure may still give away information / data statistics
        prefix = 'modified_' + prefix

    if not os.path.exists(
            os.path.join(os.path.dirname(filepath), '{}rel_mat.obj'.format(prefix))) or probabilistic and \
            not os.path.exists(os.path.join(os.path.dirname(filepath), '{}prob_mat.obj'.format(prefix))):
        print('adj mat or prob mat is missing. Begin folding from scratch.')
        fold_rna_from_file(filepath, pool, fold_algo, probabilistic, **kwargs)

    load_dense = kwargs.get('load_dense', True)
    sp_rel_matrix = pickle.load(open(os.path.join(os.path.dirname(filepath), '{}rel_mat.obj'.format(prefix)), 'rb'))
    if load_dense:
        adjacency_matrix = np.array([mat.toarray() for mat in sp_rel_matrix])
    else:
        adjacency_matrix = np.array(sp_rel_matrix)

    if probabilistic:
        sp_prob_matrix = pickle.load(
            open(os.path.join(os.path.dirname(filepath), '{}prob_mat.obj'.format(prefix)), 'rb'))
        if load_dense:
            probability_matrix = np.array([mat.toarray() for mat in sp_prob_matrix])
        else:
            probability_matrix = np.array(sp_prob_matrix)
        matrix = (adjacency_matrix, probability_matrix)
    else:
        matrix = adjacency_matrix

    return matrix

def load_seq(filepath):
    if filepath.endswith('.fa'):
        file = open(filepath, 'r')
    else:
        file = gzip.open(filepath, 'rb')

    all_id, all_seq = load_fasta_format(file)
    return all_id, all_seq

def fold_rna_from_file(filepath, p=None, fold_algo='rnafold', probabilistic=False, **kwargs):
    assert (fold_algo in ['rnafold', 'rnasubopt', 'rnaplfold'])
    if fold_algo == 'rnafold':
        assert (probabilistic is False)
    if fold_algo == 'rnaplfold':
        assert (probabilistic is True)
    print('Parsing', filepath)
    _, all_seq = load_seq(filepath)

    # compatible with already computed structures with RNAfold
    prefix = '%s_%s_' % (fold_algo, probabilistic)
    if fold_algo == 'rnaplfold' or fold_algo == 'rnashapes':
        prefix += '%d_' % (kwargs.get('w', 150))
    if kwargs.get('modify_leaks', False):
        prefix = 'modified_' + prefix

    if p is None:
        pool = Pool(int(os.cpu_count() * 2 / 3))
    else:
        pool = p

    if fold_algo == 'rnafold':
        fold_func = fold_seq_rnafold
        res = list(pool.imap(fold_func, all_seq))
        sp_rel_matrix = []
        structural_content = []
        for struct, matrix in res:
            structural_content.append(struct)
            sp_rel_matrix.append(matrix)
        np.save(os.path.join(os.path.dirname(filepath), '{}structures.npy'.format(prefix)),
                np.array(structural_content))  # [size, length, 3]
        pickle.dump(sp_rel_matrix,
                    open(os.path.join(os.path.dirname(filepath), '{}rel_mat.obj'.format(prefix)), 'wb'))
    elif fold_algo == 'rnasubopt':
        fold_func = partial(fold_seq_subopt, fold_algo=fold_algo, probabilistic=probabilistic)
        res = list(pool.imap(fold_func, all_seq))
        sp_rel_matrix = []
        sp_prob_matrix = []
        structural_content = []
        for struct, matrix in res:
            structural_content.append(struct)
            if probabilistic:
                rel_mat, prob_mat = matrix
                sp_prob_matrix.append(prob_mat)
            else:
                rel_mat = matrix
            sp_rel_matrix.append(rel_mat)
        np.save(os.path.join(os.path.dirname(filepath), '{}structures.npy'.format(prefix)),
                np.array(structural_content))  # [size, length, 3]
        pickle.dump(sp_rel_matrix,
                    open(os.path.join(os.path.dirname(filepath), '{}rel_mat.obj'.format(prefix)), 'wb'))
        if probabilistic:
            pickle.dump(sp_prob_matrix,
                        open(os.path.join(os.path.dirname(filepath), '{}prob_mat.obj'.format(prefix)), 'wb'))
    elif fold_algo == 'rnaplfold':
        winsize = kwargs.get('w', 150)
        print('running rnaplfold with winsize %d' % (winsize))
        fold_func = partial(fold_seq_rnaplfold, w=winsize, l=min(winsize, 150), cutoff=1e-4, no_lonely_bps=True)
        res = list(pool.imap(fold_func, all_seq))
        sp_rel_matrix = []
        sp_prob_matrix = []
        for rel_mat, prob_mat in res:
            sp_rel_matrix.append(rel_mat)
            sp_prob_matrix.append(prob_mat)
        pickle.dump(sp_rel_matrix,
                    open(os.path.join(os.path.dirname(filepath), '{}rel_mat.obj'.format(prefix)), 'wb'))
        pickle.dump(sp_prob_matrix,
                    open(os.path.join(os.path.dirname(filepath), '{}prob_mat.obj'.format(prefix)), 'wb'))
    else:
        raise ValueError('Supported folding algorithms are ' + ', '.join(['rnafold', 'rnasubopt', 'rnaplfold']))

    if p is None:
        pool.close()
        pool.join()

    print('Parsing', filepath, 'finished')

def fold_and_check_hairpin(seq, return_label=True):
    regex = r'^\(\(\(\.\.\.\)\)\)[\.\(]|[\.\)]\(\(\(\.\.\.\)\)\)$|[\.\)]\(\(\(\.\.\.\)\)\)|\(\(\(\.\.\.\)\)\)[\.\(]'
    '''return label, or an annotation over the entire seq'''
    '''fold rna and check if the structure contains a hairpin of 3 loose nucleotide connected by a stem of 3 basepairs'''
    struct = RNA.fold(seq)[0]
    mat = adj_mat(struct)
    if return_label:
        match = re.search(regex, struct)
        return struct, mat, int(match is not None)
    else:
        annotation = [0] * len(seq)
        for match in re.finditer(regex, struct):
            start_idx = struct[match.start(): match.end()].index('(((...)))') + match.start()
            annotation[start_idx:start_idx + 9] = [1] * 9
        return struct, mat, annotation

def fold_and_check_element(seq, element_symbol, return_label=True):
    '''simply use forgi to annotate the whole string and check if it contains the element'''
    '''return label, or an annotation over the entire seq'''
    '''fold rna and check if the structure contains a hairpin of 3 loose nucleotide connected by a stem of 3 basepairs'''
    struct = RNA.fold(seq)[0]
    mat = adj_mat(struct)
    bg = fgb.BulgeGraph.from_dotbracket(struct)
    annotation = bg.to_element_string(())
    if return_label:
        return struct, mat, int(element_symbol in annotation)
    else:
        return struct, mat, [int(element_symbol == c) for c in annotation]

def generate_hairpin_dataset(n, length, p=None, return_label=True):
    '''
    generate toy dataset
    positive examples: RNA sequences that contain specific structural motifs:
        1. a hairpin of three nucleotides connected by a stem of 3 base-pairs
        2. nucleotidal composition does not matter.
    negative examples: RNA sequences that do not contain this specific motifs
    '''
    data_path = os.path.join(basedir, 'Data/toy-data/hairpin/%s' % ('label' if return_label else 'annotation'))
    if not os.path.exists(data_path):
        os.makedirs(data_path)

    if os.path.exists(os.path.join(data_path, 'seq-and-struct.fa')) and \
            os.path.exists(os.path.join(data_path, 'adj_mat.obj')):
        all_labels = []
        all_seqs = []
        all_struct = []
        with open(os.path.join(data_path, 'seq-and-struct.fa'), 'r') as file:
            for line in file:
                if line[0] == '>':
                    label = line.rstrip().split(' ')[-1].split(':')[-1]
                    if return_label:
                        all_labels.append(int(label))
                    else:
                        all_labels.append([int(c) for c in label.split(',')])
                elif line[0] in 'ACGT':
                    all_seqs.append(['ACGT'.index(c) for c in line.rstrip()])
                elif line[0] in '.()':
                    all_struct.append(['.()'.index(c) for c in line.rstrip()])

        all_seqs = np.array(all_seqs)
        sp_adj_matrix = pickle.load(open(os.path.join(data_path, 'adj_mat.obj'), 'rb'))
    else:
        all_seqs = np.zeros((n, length), dtype=int)
        for j in range(length):
            all_seqs[:, j] = np.random.choice([0, 1, 2, 3], n, p=[0.25, 0.25, 0.25, 0.25])
        seqs_str = [''.join(['ACGT'[c] for c in seq]) for seq in all_seqs]

        if p is None:
            pool = Pool(8)
        else:
            pool = p

        res = list(pool.imap(partial(fold_and_check_hairpin, return_label=return_label), seqs_str))

        sp_adj_matrix = []
        all_labels = []
        all_struct = []
        with open(os.path.join(data_path, 'seq-and-struct.fa'), 'w') as file:
            for seq, (struct, mat, label) in zip(seqs_str, res):
                file.writelines('> label:%s\n%s\n%s\n' %
                                (label if return_label else ','.join([str(c) for c in label]), seq, struct))
                sp_adj_matrix.append(mat)
                all_labels.append(label)
                all_struct.append(['.()'.index(c) for c in struct])

        pickle.dump(sp_adj_matrix, open(os.path.join(data_path, 'adj_mat.obj'), 'wb'))
        if p is None:
            pool.close()
            pool.join()

    all_labels = np.array(all_labels)
    adjacency_matrix = np.stack([mat.toarray() for mat in sp_adj_matrix], axis=0)
    all_struct = np.array(all_struct)
    return all_seqs, adjacency_matrix, all_labels, all_struct

def generate_element_dataset(n, length, element_symbol, p=None, return_label=True):
    '''
    generate toy dataset
    positive examples: RNA sequences that contain specific structural motifs:
        1. a hairpin of three nucleotides connected by a stem of 3 base-pairs
        2. nucleotidal composition does not matter.
    negative examples: RNA sequences that do not contain this specific motifs
    '''
    assert (len(element_symbol) == 1 and str.isalpha(element_symbol))
    data_path = os.path.join(basedir,
                             'Data/toy-data/%s/%s' % (element_symbol, 'label' if return_label else 'annotation'))
    if not os.path.exists(data_path):
        os.makedirs(data_path)

    if os.path.exists(os.path.join(data_path, 'seq-and-struct.fa')) and \
            os.path.exists(os.path.join(data_path, 'adj_mat.obj')):
        all_labels = []
        all_seqs = []
        all_struct = []
        with open(os.path.join(data_path, 'seq-and-struct.fa'), 'r') as file:
            for line in file:
                if line[0] == '>':
                    label = line.rstrip().split(' ')[-1].split(':')[-1]
                    if return_label:
                        all_labels.append(int(label))
                    else:
                        all_labels.append([int(c) for c in label.split(',')])
                elif line[0] in 'ACGT':
                    all_seqs.append(['ACGT'.index(c) for c in line.rstrip()])
                elif line[0] in '.()':
                    all_struct.append(['.()'.index(c) for c in line.rstrip()])

        all_seqs = np.array(all_seqs)
        sp_adj_matrix = pickle.load(open(os.path.join(data_path, 'adj_mat.obj'), 'rb'))
    else:
        all_seqs = np.zeros((n, length), dtype=int)
        for j in range(length):
            all_seqs[:, j] = np.random.choice([0, 1, 2, 3], n, p=[0.25, 0.25, 0.25, 0.25])
        seqs_str = [''.join(['ACGT'[c] for c in seq]) for seq in all_seqs]

        if p is None:
            pool = Pool(8)
        else:
            pool = p

        res = list(pool.imap(partial(fold_and_check_element,
                                     element_symbol=element_symbol,
                                     return_label=return_label), seqs_str))

        sp_adj_matrix = []
        all_labels = []
        all_struct = []
        with open(os.path.join(data_path, 'seq-and-struct.fa'), 'w') as file:
            for seq, (struct, mat, label) in zip(seqs_str, res):
                file.writelines('> label:%s\n%s\n%s\n' %
                                (label if return_label else ','.join([str(c) for c in label]), seq, struct))
                sp_adj_matrix.append(mat)
                all_labels.append(label)
                all_struct.append(['.()'.index(c) for c in struct])

        pickle.dump(sp_adj_matrix, open(os.path.join(data_path, 'adj_mat.obj'), 'wb'))
        if p is None:
            pool.close()
            pool.join()

    all_labels = np.array(all_labels)
    adjacency_matrix = np.stack([mat.toarray() for mat in sp_adj_matrix], axis=0)
    all_struct = np.array(all_struct)
    return all_seqs, adjacency_matrix, all_labels, all_struct

if __name__ == "__main__":
    np.set_printoptions(threshold=np.inf, edgeitems=30, linewidth=100000, )

    # res = fold_seq_rnaplfold(
    #     'cgcgggacgcggcccgaggccgtgcgcgagccggggcaccgggcggcggcggcggcggcgcgcgccatgtcgttcagtgaaatgaaccgcaggacgctggcgttccgaggaggcgggttggtcaccgctagcggcggcggctccacgaacAATAACGCTGGCGGGGAGGCCTCAGcttggcctccgcagccccagccgagacagcccccgccgccagcgccgcccgcgcttcagccgcctaatgggcggggggccgacgaggaagtggaattggagggcctggagccccaagacctggaggcctccgccgggccggccgccggcg',
    #     150, 150, 0.0001, True)
    # print(res[1].todense().sum(axis=-1))
    res = fold_seq_rnashapes(
        'cgcgggacgcggcccgaggccgtgcgcgagccggggcaccgggcggcggcggcggcggcgcgcgccatgtcgttcagtgaaatgaaccgcaggacgctggcgttccgaggaggcgggttggtcaccgctagcggcggcggctccacgaacAATAACGCTGGCGGGGAGGCCTCAGcttggcctccgcagccccagccgagacagcccccgccgccagcgccgcccgcgcttcagccgcctaatgggcggggggccgacgaggaagtggaattggagggcctggagccccaagacctggaggcctccgccgggccggccgccggcg',
        150, iterations=100)
    # print(res[0].todense())
    print(res[1].todense().sum(axis=-1))

    # with open('boltzmann-sampling-acc.txt', 'w') as file:
    #     for amount in [5, 10, 100, 1000, 5000, 10000]:
    #         rel_diff, prob_diff = [], []
    #         for replcate in range(100):
    #             _, res = fold_seq_subopt(
    #                 'TGTGAAGCGCGGCTAGCTGCCGGGGTTCGAGGTGGGTCCCAGGGTTAAAATCCCTTGTTGTCTTACTGGTGGCAGCAAGCTAGGACTATACTCCTCGGTCG',
    #                 'rnafold', True, amount)
    #             _, new_res = fold_seq_subopt(
    #                 'TGTGAAGCGCGGCTAGCTGCCGGGGTTCGAGGTGGGTCCCAGGGTTAAAATCCCTTGTTGTCTTACTGGTGGCAGCAAGCTAGGACTATACTCCTCGGTCG',
    #                 'rnafold', True, amount)
    #
    #             diff = (res[0].todense() != new_res[0].todense()).astype(np.int32)
    #             rel_diff.append(np.sum(diff))
    #
    #             diff = np.abs(res[1].todense() - new_res[1].todense())
    #             prob_diff.append(np.mean(np.max(diff, axis=-1)))
    #         file.writelines('sampling amount %d, relation difference: %.4f\u00b1%.4f, probability difference: %.4f\u00b1%.4f' %
    #               (amount, np.mean(rel_diff), np.std(rel_diff), np.mean(prob_diff), np.std(prob_diff)))
    #         print('sampling amount %d, relation difference: %.4f\u00b1%.4f, probability difference: %.4f\u00b1%.4f' %
    #               (amount, np.mean(rel_diff), np.std(rel_diff), np.mean(prob_diff), np.std(prob_diff)))

    # annotation for the multiloop elements
    # all_seqs, adjacency_matrix, all_labels, _ = generate_element_dataset(80000, 101, 'i', return_label=False)
    # print(all_labels.shape)
    # print(np.where(np.count_nonzero(all_labels, axis=-1) > 0)[0].__len__())
    #
    # all_seqs, adjacency_matrix, all_labels, _ = generate_hairpin_dataset(80000, 101, 'm', return_label=False)
    # print(all_labels.shape)
    # print(np.where(np.count_nonzero(all_labels, axis=-1) > 0)[0].__len__())


import numpy as np
from skimage.future.graph import RAG
import heapq
from skimage.segmentation import find_boundaries
from skimage.morphology import skeletonize
from sklearn.linear_model import LogisticRegression
from lib.utils import get_xymaps, shuffler
from .nmf import NMFDataCompressor

def run_lrc_mrm(dataset_slice, cps, sample_size):
    compressor = NMFDataCompressor(cps)
    compressor.fit(dataset_slice, sample_size)
    compressed_dataset = compressor.transform(dataset_slice['data'])
    dataset_slice['data'] = compressed_dataset
    dataset_slice = _fit_lrc_model(
        dataset_slice,
        model=LogisticRegression(penalty='none'),
        training_set_size=sample_size,
    )
    dataset_slice = _lrc_mrm_segmentation(dataset_slice)
    return dataset_slice

def _lrc_mrm_segmentation(dataset):
    """
    Implementation of the multi-region merging segmentation controlled by
    a trained classifier model. Design of the function was originally inspired
    by the merge_hierarchical function of Skimage:
    (https://github.com/scikit-image/scikit-image/blob/master/skimage/)
    """
    # Collect spatial resolution and data from dataset
    rx, ry = dataset.get('spatial_resol')
    data = dataset.get('data')
    model = dataset.get('lrc_model')

    # Define merging decision function
    def mdf(v):
        return model.predict_proba(np.atleast_2d(v))[0, 1]

    # Initialize region adjacency graph (RAG)
    rag, edge_heap, segments = _initialize_graph(rx, ry, data, model)

    # Start the region-merging algorithm
    while (len(edge_heap) > 0) and (edge_heap[0][0] < 0.5):
        # Pop the smallest edge from the heap if weight < 0.5
        smallest_weight, n1, n2, valid = heapq.heappop(edge_heap)

        # Check that the edge is valid
        if valid:
            # Make sure that n1 is the smallest regiom
            if rag.nodes[n1]['count'] > rag.nodes[n2]['count']:
                n1, n2 = n2, n1

            # Update properties of n2
            rag.nodes[n2]['labels'] = (rag.nodes[n1]['labels']
                                       + rag.nodes[n2]['labels'])
            rag.nodes[n2]['count'] = (rag.nodes[n1]['count']
                                      + rag.nodes[n2]['count'])

            # Get new neighbors of n2
            n1_numbers = set(rag.neighbors(n1))
            n2_numbers = set(rag.neighbors(n2))
            new_neighbors = (n1_numbers | n2_numbers) - n2_numbers - {n1, n2}

            # Disable edges of n1 in the heap list
            for nbr in rag.neighbors(n1):
                edge = rag[n1][nbr]
                edge['heap item'][3] = False

            # Remove n1 from the graph (edges are still in the heap list)
            rag.remove_node(n1)

            # Update new edges of n2
            for nbr in new_neighbors:
                rag.add_edge(n2, nbr)
                edge = rag[n2][nbr]
                master_n2 = rag.nodes[n2]['master']
                master_nbr = rag.nodes[nbr]['master']
                weight = mdf(_vector_similarity(master_n2, master_nbr))
                heap_item = [weight, n2, nbr, (weight < 0.5)]
                edge['heap item'] = heap_item
                # Push edges to the heap
                heapq.heappush(edge_heap, heap_item)

    # Compute grain segmentation map
    label_map = np.arange(segments.max() + 1)
    for ix, (n, d) in enumerate(rag.nodes(data=True)):
        for lab in d['labels']:
            label_map[lab] = ix
    segmentation = label_map[segments]

    # Compute grain boundary map
    gbs = skeletonize(find_boundaries(segmentation, mode='inner'))

    # Return updated dataset
    dataset['segmentation'] = segmentation.ravel()
    dataset['boundaries'] = gbs.ravel()

    return dataset

def _fit_lrc_model(dataset, model, training_set_size):
    """Fits a model, computes precision, recall and accuracy"""
    training_set = _get_sample_set(dataset, training_set_size)
    model.fit(training_set['x'], training_set['y'])
    dataset['lrc_model'] = model
    return dataset

def _initialize_graph(rx, ry, data, model):
    """Initializes the Region Adjacency Graph (RAG)."""

    # Define merging decision function
    def mdf(v):
        return model.predict_proba(np.atleast_2d(v))[0, 1]

    # Initialize RAG
    x_map, y_map = get_xymaps(rx, ry)
    segments = np.arange(rx * ry).reshape((rx, ry))
    rag = RAG(segments)

    # Initialize nodes
    data_reshaped = data.reshape((rx, ry, data.shape[1]))
    for n in rag:
        rag.nodes[n].update({'labels': [n]})
    for index in np.ndindex(segments.shape):
        current = segments[index]
        rag.nodes[current]['count'] = 1
        rag.nodes[current]['master'] = data_reshaped[index]
        rag.nodes[current]['xpos'] = x_map[current]
        rag.nodes[current]['ypos'] = y_map[current]

    # Initialize edges
    edge_heap = []
    for n1, n2, d in rag.edges(data=True):
        master_x = rag.nodes[n1]['master']
        master_y = rag.nodes[n2]['master']
        weight = mdf(_vector_similarity(master_x, master_y))
        # Push the edge into the heap
        heap_item = [weight, n1, n2, (weight < 0.5)]
        d['heap item'] = heap_item
        heapq.heappush(edge_heap, heap_item)

    return rag, edge_heap, segments

def _get_sample_set(dataset, sample_size):
    """
    Randomly extracts a training or test set from the dataset.
    - Class 0: pairs of adjacent voxels
    - Class 1: pairs of non-adjacent voxels
    """
    # Collect data from the dataset
    rx, ry = dataset.get('spatial_resol')
    data = dataset.get('data')
    x_map, y_map = get_xymaps(rx, ry)

    # Extract adjacent sample
    x_close, y_close = _get_adjacent_sample(
        rx, ry, data, sample_size, x_map, y_map)

    # Extract non-adjacent sample
    x_far, y_far = _get_non_adjacent_sample(data, sample_size, x_map, y_map)

    # Stack both samples
    x = np.vstack((x_close, x_far))
    y = np.hstack((y_close, y_far))

    # Shuffle extracted set
    idx = np.arange(0, x.shape[0])
    np.random.shuffle(idx)
    x = x[idx]
    y = y[idx]

    sample_set = {'x': x, 'y': y}

    return sample_set

def _get_adjacent_sample(rx, ry, data, sample_size, xmap, ymap):
    """Samples Sbar, the distribution of adjacent pixel feature vectors."""

    # Get set of random data
    x0, idx = shuffler(data, sample_size)

    # Modify location by 1 pixel
    modified_x = xmap[idx] + (np.random.randint(0, 2, sample_size) * 2 - 1)
    modified_x = np.clip(modified_x.astype('int'), 0, rx - 1)
    modified_y = ymap[idx] + (np.random.randint(0, 2, sample_size) * 2 - 1)
    modified_y = np.clip(modified_y.astype('int'), 0, ry - 1)

    # Find corresponding signal
    x1 = np.empty_like(x0)
    c = 0
    i = np.arange(rx * ry)
    for xc, yc in zip(modified_x, modified_y):
        u = np.zeros((rx, ry))
        u[xc, yc] = 1
        num = i[(u.ravel() == 1)]
        x1[c] = data[num]
        c += 1

    # Compute distance vectors
    x_close = _vector_similarity(x1, x0)

    # Label as 0
    y_close = np.zeros(x_close.shape[0], dtype=np.uint8)

    return x_close, y_close

def _get_non_adjacent_sample(data, sample_size, xmap, ymap):
    """Samples the distribution of non-adjacent pixels."""

    # Get random set of data and location of selected pixel pairs, twice
    x0, idx = shuffler(data, sample_size)
    loc_x_0, loc_y_0 = xmap[idx], ymap[idx]
    x1, idx = shuffler(data, sample_size)
    loc_x_1, loc_y_1 = xmap[idx], ymap[idx]

    # Compute distance vectors
    x_far = _vector_similarity(x1, x0)

    # Filter out adjacent examples
    adjacent_filter = np.abs(loc_x_0 - loc_x_1) + np.abs(loc_y_0 - loc_y_1) < 2
    x_far = x_far[~adjacent_filter]

    # Label as 1
    y_far = np.ones(x_far.shape[0], dtype=np.uint8)

    return x_far, y_far

def _vector_similarity(a, b):
    """Returns distance vector of two input feature vectors"""
    return np.square(np.subtract(a, b))


# -*- coding: utf-8 -*-

"""
scikit-rebate was primarily developed at the University of Pennsylvania by:
    - Randal S. Olson (rso@randalolson.com)
    - Pete Schmitt (pschmitt@upenn.edu)
    - Ryan J. Urbanowicz (ryanurb@upenn.edu)
    - Weixuan Fu (weixuanf@upenn.edu)
    - and many more generous open source contributors

Permission is hereby granted, free of charge, to any person obtaining a copy of this software
and associated documentation files (the "Software"), to deal in the Software without restriction,
including without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense,
and/or sell copies of the Software, and to permit persons to whom the Software is furnished to do so,
subject to the following conditions:

The above copyright notice and this permission notice shall be included in all copies or substantial
portions of the Software.

THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT
LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.
IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY,
WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE
SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
"""

import numpy as np

# (Subset of continuous-valued feature data, Subset of discrete-valued feature data, max/min difference, instance index, boolean mask for continuous, boolean mask for discrete)
def get_row_missing(xc, xd, cdiffs, index, cindices, dindices):
    """ Calculate distance between index instance and all other instances. """
    row = np.empty(0, dtype=np.double)  # initialize empty row
    cinst1 = xc[index]  # continuous-valued features for index instance
    dinst1 = xd[index]  # discrete-valued features for index instance
    # Boolean mask locating missing values for continuous features for index instance
    can = cindices[index]
    # Boolean mask locating missing values for discrete features for index instance
    dan = dindices[index]
    tf = len(cinst1) + len(dinst1)  # total number of features.

    # Progressively compare current instance to all others. Excludes comparison with self indexed instance. (Building the distance matrix triangle).
    for j in range(index):
        dist = 0
        dinst2 = xd[j]  # discrete-valued features for compared instance
        cinst2 = xc[j]  # continuous-valued features for compared instance

        # Manage missing values in discrete features
        # Boolean mask locating missing values for discrete features for compared instance
        dbn = dindices[j]
        # indexes where there is at least one missing value in the feature between an instance pair.
        idx = np.unique(np.append(dan, dbn))
        # Number of features excluded from distance calculation due to one or two missing values within instance pair. Used to normalize distance values for comparison.
        dmc = len(idx)
        d1 = np.delete(dinst1, idx)  # delete unique missing features from index instance
        d2 = np.delete(dinst2, idx)  # delete unique missing features from compared instance

        # Manage missing values in continuous features
        # Boolean mask locating missing values for continuous features for compared instance
        cbn = cindices[j]
        # indexes where there is at least one missing value in the feature between an instance pair.
        idx = np.unique(np.append(can, cbn))
        # Number of features excluded from distance calculation due to one or two missing values within instance pair. Used to normalize distance values for comparison.
        cmc = len(idx)
        c1 = np.delete(cinst1, idx)  # delete unique missing features from index instance
        c2 = np.delete(cinst2, idx)  # delete unique missing features from compared instance
        # delete unique missing features from continuous value difference scores
        cdf = np.delete(cdiffs, idx)

        # Add discrete feature distance contributions (missing values excluded) - Hamming distance
        dist += len(d1[d1 != d2])

        # Add continuous feature distance contributions (missing values excluded) - Manhattan distance (Note that 0-1 continuous value normalization is included ~ subtraction of minimums cancel out)
        dist += np.sum(np.absolute(np.subtract(c1, c2)) / cdf)

        # Normalize distance calculation based on total number of missing values bypassed in either discrete or continuous features.
        tnmc = tf - dmc - cmc  # Total number of unique missing counted
        # Distance normalized by number of features included in distance sum (this seeks to handle missing values neutrally in distance calculation)
        dist = dist/float(tnmc)

        row = np.append(row, dist)

    return row

# For iter relief
def get_row_missing_iter(xc, xd, cdiffs, index, cindices, dindices, weights):
    """ Calculate distance between index instance and all other instances. """
    row = np.empty(0, dtype=np.double)  # initialize empty row
    cinst1 = xc[index]  # continuous-valued features for index instance
    dinst1 = xd[index]  # discrete-valued features for index instance
    # Boolean mask locating missing values for continuous features for index instance
    can = cindices[index]
    # Boolean mask locating missing values for discrete features for index instance
    dan = dindices[index]
    tf = len(cinst1) + len(dinst1)  # total number of features.
    # Progressively compare current instance to all others. Excludes comparison with self indexed instance. (Building the distance matrix triangle).
    for j in range(index):
        dist = 0
        dinst2 = xd[j]  # discrete-valued features for compared instance
        cinst2 = xc[j]  # continuous-valued features for compared instance

        # Manage missing values in discrete features
        # Boolean mask locating missing values for discrete features for compared instance
        dbn = dindices[j]
        # indexes where there is at least one missing value in the feature between an instance pair.
        idx = np.unique(np.append(dan, dbn))
        # Number of features excluded from distance calculation due to one or two missing values within instance pair. Used to normalize distance values for comparison.
        dmc = len(idx)
        d1 = np.delete(dinst1, idx)  # delete unique missing features from index instance
        d2 = np.delete(dinst2, idx)  # delete unique missing features from compared instance

        wd = np.delete(weights, idx)  # delete weights corresponding to missing discrete features
        # Manage missing values in continuous features
        # Boolean mask locating missing values for continuous features for compared instance
        cbn = cindices[j]
        # indexes where there is at least one missing value in the feature between an instance pair.
        idx = np.unique(np.append(can, cbn))
        # Number of features excluded from distance calculation due to one or two missing values within instance pair. Used to normalize distance values for comparison.
        cmc = len(idx)
        c1 = np.delete(cinst1, idx)  # delete unique missing features from index instance
        c2 = np.delete(cinst2, idx)  # delete unique missing features from compared instance
        # delete unique missing features from continuous value difference scores
        cdf = np.delete(cdiffs, idx)
        wc = np.delete(weights, idx)  # delete weights corresponding to missing continuous features

        # Add discrete feature distance contributions (missing values excluded) - Hamming distance
        if len(d1)!=0: #To ensure there is atleast one discrete variable
            hamming_dist = np.not_equal(d1, d2).astype(float)
            weight_hamming_dist = np.dot(hamming_dist, wd)/np.sum(wd)
            dist += weight_hamming_dist

        # Add continuous feature distance contributions (missing values excluded) - Manhattan distance (Note that 0-1 continuous value normalization is included ~ subtraction of minimums cancel out)
        if len(c1)!=0: #To ensure there is atleast one continuous variable
            dist += np.dot((np.absolute(np.subtract(c1, c2)) / cdf), wc)/np.sum(wc)

        # Normalize distance calculation based on total number of missing values bypassed in either discrete or continuous features.
        tnmc = tf - dmc - cmc  # Total number of unique missing counted
        # Distance normalized by number of features included in distance sum (this seeks to handle missing values neutrally in distance calculation)
        dist = dist/float(tnmc)

        row = np.append(row, dist)

    return row

def ramp_function(data_type, attr, fname, xinstfeature, xNNifeature):
    """ Our own user simplified variation of the ramp function suggested by Hong 1994, 1997. Hong's method requires the user to specifiy two thresholds
    that indicate the max difference before a score of 1 is given, as well a min difference before a score of 0 is given, and any in the middle get a
    score that is the normalized difference between the two continuous feature values. This was done because when discrete and continuous features were mixed,
    continuous feature scores were underestimated.  Towards simplicity, automation, and a dataset adaptable approach,
    here we simply check whether the difference is greater than the standard deviation for the given feature; if so we assign a score of 1, otherwise we
    assign the normalized feature score difference.  This should help compensate for the underestimation. """
    diff = 0
    mmdiff = attr[fname][3]  # Max/Min range of values for target feature
    rawfd = abs(xinstfeature - xNNifeature)  # prenormalized feature value difference

    if data_type == 'mixed':  # Ramp function utilized
        # Check whether feature value difference is greater than the standard deviation
        standDev = attr[fname][4]
        if rawfd > standDev:  # feature value difference is is wider than a standard deviation
            diff = 1
        else:
            diff = abs(xinstfeature - xNNifeature) / mmdiff

    else:  # Normal continuous feature scoring
        diff = abs(xinstfeature - xNNifeature) / mmdiff

    return diff

def compute_score(attr, mcmap, NN, feature, inst, nan_entries, headers, class_type, X, y, labels_std, data_type, near=True):
    """Flexible feature scoring method that can be used with any core Relief-based method. Scoring proceeds differently
    based on whether endpoint is binary, multiclass, or continuous. This method is called for a single target instance
    + feature combination and runs over all items in NN. """

    fname = headers[feature]  # feature identifier
    ftype = attr[fname][0]  # feature type
    ctype = class_type  # class type (binary, multiclass, continuous)
    diff_hit = diff_miss = 0.0  # Tracks the score contribution
    # Tracks the number of hits/misses. Used in normalizing scores by 'k' in ReliefF, and by m or h in SURF, SURF*, MultiSURF*, and MultiSURF
    count_hit = count_miss = 0.0
    # Initialize 'diff' (The score contribution for this target instance and feature over all NN)
    diff = 0
    # mmdiff = attr[fname][3] # Max/Min range of values for target feature

    datalen = float(len(X))

    # If target instance is missing, then a 'neutral' score contribution of 0 is returned immediately since all NN comparisons will be against this missing value.
    if nan_entries[inst][feature]:
        return 0.
    # Note missing data normalization below regarding missing NN feature values is accomplished by counting hits and misses (missing values are not counted) (happens in parallel with hit/miss imbalance normalization)

    xinstfeature = X[inst][feature]  # value of target instances target feature.

    #--------------------------------------------------------------------------
    if ctype == 'binary':
        for i in range(len(NN)):
            if nan_entries[NN[i]][feature]:  # skip any NN with a missing value for this feature.
                continue

            xNNifeature = X[NN[i]][feature]

            if near:  # SCORING FOR NEAR INSTANCES
                if y[inst] == y[NN[i]]:   # HIT
                    count_hit += 1
                    if ftype == 'continuous':
                        # diff_hit -= abs(xinstfeature - xNNifeature) / mmdiff #Normalize absolute value of feature value difference by max-min value range for feature (so score update lies between 0 and 1)
                        diff_hit -= ramp_function(data_type, attr, fname, xinstfeature, xNNifeature)
                    else:  # discrete feature
                        if xinstfeature != xNNifeature:  # A difference in feature value is observed
                            # Feature score is reduced when we observe feature difference between 'near' instances with the same class.
                            diff_hit -= 1
                else:  # MISS
                    count_miss += 1
                    if ftype == 'continuous':
                        # diff_miss += abs(xinstfeature - xNNifeature) / mmdiff
                        diff_miss += ramp_function(data_type, attr, fname,
                                                   xinstfeature, xNNifeature)
                    else:  # discrete feature
                        if xinstfeature != xNNifeature:  # A difference in feature value is observed
                            # Feature score is increase when we observe feature difference between 'near' instances with different class values.
                            diff_miss += 1

            else:  # SCORING FOR FAR INSTANCES (ONLY USED BY MULTISURF* BASED ON HOW CODED)
                if y[inst] == y[NN[i]]:   # HIT
                    count_hit += 1
                    if ftype == 'continuous':

                        # diff_hit -= abs(xinstfeature - xNNifeature) / mmdiff  #Hits differently add continuous value differences rather than subtract them
                        # Sameness should yield most negative score
                        diff_hit -= (1-ramp_function(data_type, attr,
                                                     fname, xinstfeature, xNNifeature))
                    else:  # discrete feature
                        # The same feature value is observed (Used for more efficient 'far' scoring, since there should be fewer same values for 'far' instances)
                        if xinstfeature == xNNifeature:
                            # Feature score is reduced when we observe the same feature value between 'far' instances with the same class.
                            diff_hit -= 1
                else:  # MISS
                    count_miss += 1
                    if ftype == 'continuous':
                        # diff_miss += abs(xinstfeature - xNNifeature) / mmdiff #Misses differntly subtract continuous value differences rather than add them
                        # Sameness should yield most negative score
                        diff_miss += (1-ramp_function(data_type, attr,
                                                      fname, xinstfeature, xNNifeature))
                    else:  # discrete feature
                        # The same feature value is observed (Used for more efficient 'far' scoring, since there should be fewer same values for 'far' instances)
                        if xinstfeature == xNNifeature:
                            # Feature score is increased when we observe the same feature value between 'far' instances with different class values.
                            diff_miss += 1

        """ Score Normalizations:
        *'n' normalization dividing by the number of training instances (this helps ensure that all final scores end up in the -1 to 1 range
        *'k','h','m' normalization dividing by the respective number of hits and misses in NN (after ignoring missing values), also helps account for class imbalance within nearest neighbor radius)"""
        if count_hit == 0.0 or count_miss == 0.0:  # Special case, avoid division error
            if count_hit == 0.0 and count_miss == 0.0:
                return 0.0
            elif count_hit == 0.0:
                diff = (diff_miss / count_miss) / datalen
            else:  # count_miss == 0.0
                diff = (diff_hit / count_hit) / datalen
        else:  # Normal diff normalization
            diff = ((diff_hit / count_hit) + (diff_miss / count_miss)) / datalen

    #--------------------------------------------------------------------------
    elif ctype == 'multiclass':
        class_store = dict()  # only 'miss' classes will be stored
        # missClassPSum = 0

        for each in mcmap:
            if(each != y[inst]):  # Identify miss classes for current target instance.
                class_store[each] = [0, 0]
                # missClassPSum += mcmap[each]

        for i in range(len(NN)):
            if nan_entries[NN[i]][feature]:  # skip any NN with a missing value for this feature.
                continue

            xNNifeature = X[NN[i]][feature]

            if near:  # SCORING FOR NEAR INSTANCES
                if(y[inst] == y[NN[i]]):  # HIT
                    count_hit += 1
                    if ftype == 'continuous':
                        # diff_hit -= abs(xinstfeature - xNNifeature) / mmdiff
                        diff_hit -= ramp_function(data_type, attr, fname, xinstfeature, xNNifeature)
                    else:  # discrete feature
                        if xinstfeature != xNNifeature:
                            # Feature score is reduced when we observe feature difference between 'near' instances with the same class.
                            diff_hit -= 1
                else:  # MISS
                    for missClass in class_store:
                        if(y[NN[i]] == missClass):  # Identify which miss class is present
                            class_store[missClass][0] += 1
                            if ftype == 'continuous':
                                # class_store[missClass][1] += abs(xinstfeature - xNNifeature) / mmdiff
                                class_store[missClass][1] += ramp_function(
                                    data_type, attr, fname, xinstfeature, xNNifeature)
                            else:  # discrete feature
                                if xinstfeature != xNNifeature:
                                    # Feature score is increase when we observe feature difference between 'near' instances with different class values.
                                    class_store[missClass][1] += 1

            else:  # SCORING FOR FAR INSTANCES (ONLY USED BY MULTISURF* BASED ON HOW CODED)
                if(y[inst] == y[NN[i]]):  # HIT
                    count_hit += 1
                    if ftype == 'continuous':
                        # diff_hit -= abs(xinstfeature - xNNifeature) / mmdiff  #Hits differently add continuous value differences rather than subtract them
                        # Sameness should yield most negative score
                        diff_hit -= (1-ramp_function(data_type, attr,
                                                     fname, xinstfeature, xNNifeature))
                    else:  # discrete features
                        if xinstfeature == xNNifeature:
                            # Feature score is reduced when we observe the same feature value between 'far' instances with the same class.
                            diff_hit -= 1
                else:  # MISS
                    for missClass in class_store:
                        if(y[NN[i]] == missClass):
                            class_store[missClass][0] += 1
                            if ftype == 'continuous':
                                # class_store[missClass][1] += abs(xinstfeature - xNNifeature) / mmdiff
                                # Sameness should yield most negative score
                                class_store[missClass][1] += (1-ramp_function(data_type,
                                                                              attr, fname, xinstfeature, xNNifeature))
                            else:  # discrete feature
                                if xinstfeature == xNNifeature:
                                    # Feature score is increased when we observe the same feature value between 'far' instances with different class values.
                                    class_store[missClass][1] += 1

        """ Score Normalizations:
        *'n' normalization dividing by the number of training instances (this helps ensure that all final scores end up in the -1 to 1 range
        *'k','h','m' normalization dividing by the respective number of hits and misses in NN (after ignoring missing values), also helps account for class imbalance within nearest neighbor radius)
        * multiclass normalization - accounts for scoring by multiple miss class, so miss scores don't have too much weight in contrast with hit scoring. If a given miss class isn't included in NN
        then this normalization will account for that possibility. """
        # Miss component
        for each in class_store:
            count_miss += class_store[each][0]

        if count_hit == 0.0 and count_miss == 0.0:
            return 0.0
        else:
            if count_miss == 0:
                pass
            else:  # Normal diff normalization
                for each in class_store:  # multiclass normalization
                    # Contribution of given miss class weighted by it's observed frequency within NN set.
                    diff += class_store[each][1] * \
                        (class_store[each][0] / count_miss) * len(class_store)
                diff = diff / count_miss  # 'm' normalization

            # Hit component: with 'h' normalization
            if count_hit == 0:
                pass
            else:
                diff += (diff_hit / count_hit)

        diff = diff / datalen  # 'n' normalization

    #--------------------------------------------------------------------------
    else:  # CONTINUOUS endpoint
        same_class_bound = labels_std

        for i in range(len(NN)):
            if nan_entries[NN[i]][feature]:  # skip any NN with a missing value for this feature.
                continue

            xNNifeature = X[NN[i]][feature]

            if near:  # SCORING FOR NEAR INSTANCES
                if abs(y[inst] - y[NN[i]]) < same_class_bound:  # HIT approximation
                    count_hit += 1
                    if ftype == 'continuous':
                        # diff_hit -= abs(xinstfeature - xNNifeature) / mmdiff
                        diff_hit -= ramp_function(data_type, attr, fname, xinstfeature, xNNifeature)
                    else:  # discrete feature
                        if xinstfeature != xNNifeature:
                            # Feature score is reduced when we observe feature difference between 'near' instances with the same 'class'.
                            diff_hit -= 1
                else:  # MISS approximation
                    count_miss += 1
                    if ftype == 'continuous':
                        # diff_miss += abs(xinstfeature - xNNifeature) / mmdiff
                        diff_miss += ramp_function(data_type, attr, fname,
                                                   xinstfeature, xNNifeature)
                    else:  # discrete feature
                        if xinstfeature != xNNifeature:
                            # Feature score is increase when we observe feature difference between 'near' instances with different class value.
                            diff_miss += 1

            else:  # SCORING FOR FAR INSTANCES (ONLY USED BY MULTISURF* BASED ON HOW CODED)
                if abs(y[inst] - y[NN[i]]) < same_class_bound:  # HIT approximation
                    count_hit += 1
                    if ftype == 'continuous':
                        # diff_hit += abs(xinstfeature - xNNifeature) / mmdiff
                        # Sameness should yield most negative score
                        diff_hit -= (1-ramp_function(data_type, attr,
                                                     fname, xinstfeature, xNNifeature))
                    else:  # discrete feature
                        if xinstfeature == xNNifeature:
                            # Feature score is reduced when we observe the same feature value between 'far' instances with the same class.
                            diff_hit -= 1
                else:  # MISS approximation
                    count_miss += 1
                    if ftype == 'continuous':
                        # diff_miss -= abs(xinstfeature - xNNifeature) / mmdiff
                        # Sameness should yield most negative score
                        diff_miss += (1-ramp_function(data_type, attr,
                                                      fname, xinstfeature, xNNifeature))
                    else:  # discrete feature
                        if xinstfeature == xNNifeature:
                            # Feature score is increased when we observe the same feature value between 'far' instances with different class values.
                            diff_miss += 1

        """ Score Normalizations:
        *'n' normalization dividing by the number of training instances (this helps ensure that all final scores end up in the -1 to 1 range
        *'k','h','m' normalization dividing by the respective number of hits and misses in NN (after ignoring missing values), also helps account for class imbalance within nearest neighbor radius)"""

        if count_hit == 0.0 or count_miss == 0.0:  # Special case, avoid division error
            if count_hit == 0.0 and count_miss == 0.0:
                return 0.0
            elif count_hit == 0.0:
                diff = (diff_miss / count_miss) / datalen
            else:  # count_miss == 0.0
                diff = (diff_hit / count_hit) / datalen
        else:  # Normal diff normalization
            diff = ((diff_hit / count_hit) + (diff_miss / count_miss)) / datalen

    return diff

def ReliefF_compute_scores(inst, attr, nan_entries, num_attributes, mcmap, NN, headers, class_type, X, y, labels_std, data_type, weight_flag=0, weights=None):
    """ Unique scoring procedure for ReliefF algorithm. Scoring based on k nearest hits and misses of current target instance. """
    scores = np.zeros(num_attributes)
    if weight_flag == 2:
        for feature_num in range(num_attributes):
            scores[feature_num] += weights[feature_num]*compute_score(attr, mcmap, NN, feature_num, inst,
                                                                      nan_entries, headers, class_type, X, y, labels_std, data_type)
    else:
        for feature_num in range(num_attributes):
            scores[feature_num] += compute_score(attr, mcmap, NN, feature_num, inst,
                                                 nan_entries, headers, class_type, X, y, labels_std, data_type)
    return scores

def SURF_compute_scores(inst, attr, nan_entries, num_attributes, mcmap, NN, headers, class_type, X, y, labels_std, data_type, weight_flag=0, weights=None):
    """ Unique scoring procedure for SURF algorithm. Scoring based on nearest neighbors within defined radius of current target instance. """
    scores = np.zeros(num_attributes)
    if weight_flag == 2:
        if len(NN) <= 0:
            return scores
        for feature_num in range(num_attributes):
            scores[feature_num] += weights[feature_num]*compute_score(attr, mcmap, NN, feature_num, inst,
                                                                      nan_entries, headers, class_type, X, y, labels_std, data_type)
    else:
        if len(NN) <= 0:
            return scores
        for feature_num in range(num_attributes):
            scores[feature_num] += compute_score(attr, mcmap, NN, feature_num, inst,
                                                 nan_entries, headers, class_type, X, y, labels_std, data_type)
    return scores

def SURFstar_compute_scores(inst, attr, nan_entries, num_attributes, mcmap, NN_near, NN_far, headers, class_type, X, y, labels_std, data_type, weight_flag=0, weights=None):
    """ Unique scoring procedure for SURFstar algorithm. Scoring based on nearest neighbors within defined radius, as well as
    'anti-scoring' of far instances outside of radius of current target instance"""
    scores = np.zeros(num_attributes)
    if weight_flag == 2:
        for feature_num in range(num_attributes):
            if len(NN_near) > 0:
                scores[feature_num] += weights[feature_num]*compute_score(attr, mcmap, NN_near, feature_num, inst,
                                                                          nan_entries, headers, class_type, X, y, labels_std, data_type)
            # Note that we are using the near scoring loop in 'compute_score' and then just subtracting it here, in line with original SURF* paper.
            if len(NN_far) > 0:
                scores[feature_num] -= weights[feature_num]*compute_score(attr, mcmap, NN_far, feature_num, inst,
                                                                          nan_entries, headers, class_type, X, y, labels_std, data_type)
    else:
        for feature_num in range(num_attributes):
            if len(NN_near) > 0:
                scores[feature_num] += compute_score(attr, mcmap, NN_near, feature_num, inst,
                                                     nan_entries, headers, class_type, X, y, labels_std, data_type)
            # Note that we are using the near scoring loop in 'compute_score' and then just subtracting it here, in line with original SURF* paper.
            if len(NN_far) > 0:
                scores[feature_num] -= compute_score(attr, mcmap, NN_far, feature_num, inst,
                                                     nan_entries, headers, class_type, X, y, labels_std, data_type)
    return scores

def MultiSURF_compute_scores(inst, attr, nan_entries, num_attributes, mcmap, NN_near, headers, class_type, X, y, labels_std, data_type, weight_flag=0, weights=None):
    """ Unique scoring procedure for MultiSURF algorithm. Scoring based on 'extreme' nearest neighbors within defined radius of current target instance. """
    scores = np.zeros(num_attributes)
    if weight_flag == 2:
        for feature_num in range(num_attributes):
            if len(NN_near) > 0:
                scores[feature_num] += weights[feature_num]*compute_score(attr, mcmap, NN_near, feature_num, inst,
                                                                          nan_entries, headers, class_type, X, y, labels_std, data_type)
    else:
        for feature_num in range(num_attributes):
            if len(NN_near) > 0:
                scores[feature_num] += compute_score(attr, mcmap, NN_near, feature_num, inst,
                                                     nan_entries, headers, class_type, X, y, labels_std, data_type)

    return scores

def MultiSURFstar_compute_scores(inst, attr, nan_entries, num_attributes, mcmap, NN_near, NN_far, headers, class_type, X, y, labels_std, data_type, weight_flag=0, weights=None):
    """ Unique scoring procedure for MultiSURFstar algorithm. Scoring based on 'extreme' nearest neighbors within defined radius, as
    well as 'anti-scoring' of extreme far instances defined by outer radius of current target instance. """
    scores = np.zeros(num_attributes)
    if weight_flag == 2:
        for feature_num in range(num_attributes):
            if len(NN_near) > 0:
                scores[feature_num] += weights[feature_num]*compute_score(attr, mcmap, NN_near, feature_num, inst,
                                                                          nan_entries, headers, class_type, X, y, labels_std, data_type)
            # Note that we add this term because we used the far scoring above by setting 'near' to False.  This is in line with original MultiSURF* paper.
            if len(NN_far) > 0:
                scores[feature_num] += weights[feature_num]*compute_score(attr, mcmap, NN_far, feature_num, inst,
                                                                          nan_entries, headers, class_type, X, y, labels_std, data_type, near=False)
    else:
        for feature_num in range(num_attributes):
            if len(NN_near) > 0:
                scores[feature_num] += compute_score(attr, mcmap, NN_near, feature_num, inst,
                                                     nan_entries, headers, class_type, X, y, labels_std, data_type)
            # Note that we add this term because we used the far scoring above by setting 'near' to False.  This is in line with original MultiSURF* paper.
            if len(NN_far) > 0:
                scores[feature_num] += compute_score(attr, mcmap, NN_far, feature_num, inst,
                                                     nan_entries, headers, class_type, X, y, labels_std, data_type, near=False)

    return scores


#!/usr/bin/env python3

"""
MIT License

Copyright (c) 2018 Jiun Yen
"""

# Imports
import os
import argparse
import numpy as np
from datetime import datetime
from random import sample
from multiprocessing import Pool, cpu_count
from functools import partial
from time import time

# Constants
_amino_acids_ = 'ARNDCQEGHILKMFPSTWYVXBZUO'
_aa_set_ = set(_amino_acids_)
_blank_ = [0 for _ in range(len(_amino_acids_) + 1)]
_verbose_ = False
_p_test_ = '../test/'
_p_data_ = '../data/'
_cpu_ = cpu_count()

# Functions
def composition(sequence):

    # Compute properties
    sequence = sequence.upper()
    query_length = len(sequence)
    mass = 1 / query_length

    query_unique = list(set(sequence) - _aa_set_)
    if query_unique:
        print('There are special residues in sequence: %s (ignored during comparison)' % ','.join(query_unique))
        sequence = ''.join([i for i in sequence if i not in query_unique])

    # Compute composition
    comp = _blank_.copy()
    comp[0] = query_length
    for e in sequence:
        comp[1 + _amino_acids_.index(e)] += mass

    return comp

def build_db(p_fasta, p_db=''):

    t0 = time()

    if not p_db:
        p_db = p_fasta + '.disorderdb'

    if _verbose_:
        print('Building database from FASTA: %s ...' % p_fasta)

    n_seqs = 0
    db_comps = []
    with open(p_fasta, 'r') as fi, open(p_db, 'w+') as fo:
        seq = ''
        header = ''
        for l in fi:
            if l.startswith('>'):
                if seq:
                    comp = composition(seq)
                    _ = fo.write('%s\t%d\t%s\n' % (header, comp[0], '\t'.join(['%.6f' % i for i in comp[1:]])))
                    db_comps.append([header] + comp)
                    n_seqs += 1

                seq = ''
                header = l[1:].strip().replace('\t', ' ')

            else:
                seq += l.strip()

            if n_seqs % 100 == 0:
                fo.flush()

        if seq:
            comp = composition(seq)
            _ = fo.write('%s\t%d\t%s\n' % (header, comp[0], '\t'.join(['%.6f' % i for i in comp[1:]])))
            db_comps.append([header] + comp)
            n_seqs += 1

    if _verbose_:
        print('\tCompleted in %.1f minutes' % ((time() - t0)/60))
        print('Generated database for %d sequences at %s' % (n_seqs, p_db))

    return p_db, db_comps

def read_db(p_db):

    db_comps = []
    with open(p_db, 'r') as f:
        for l in f:
            tmp = l.split('\t')
            comp = [tmp[0]]
            comp.append(int(tmp[1]))
            comp += [float(i) for i in tmp[2:]]
            db_comps.append(comp)

    return db_comps

def search(p_query, db_comps, p_out=''):

    t0 = time()

    if not p_out:
        stamp = 'search-%s-%s.csv' % (datetime.now().strftime('%Y%m%d%H%M%S'), ''.join(sample('ABCDEF', 4)))
        p_file, _  = os.path.splitext(p_query)
        p_out = p_file + '_' + stamp

    with open(p_out, 'w+') as f:
        _ = f.write('Queries,Hits,Distances\n')

    query_seqs = read_fasta(p_query)

    if _verbose_:
        print('Searching on %d CPUs ...' % _cpu_)

    for q_header, q_seq in query_seqs.items():

        print('\tQuery: %s' % q_header)

        query_comp = composition(q_seq)

        with Pool(processes=_cpu_) as pool:
            res = pool.map(partial(_compare_, q_comp=query_comp), db_comps)

        candidates = {h:d for h,d in res if h}

        candidates_sorted = sorted(candidates.items(), key=lambda x: x[1])

        with open(p_out, 'a') as f:

            for h,d in candidates_sorted:
                _ = f.write('%s,%s,%.4f\n' % (q_header, h, d))

    if _verbose_:
        print('\tCompleted in %.1f minutes' % ((time() - t0) / 60))
        print('Search complete, results saved as %s' % p_out)

    return p_out

def _compare_(t_db_comp, q_comp):

    header = ''
    dist = 1000.

    if q_comp[0] == t_db_comp[1]:
        header = t_db_comp[0]

        # Compute Euclidean distance
        dist = np.linalg.norm(np.array(q_comp[1:]) - np.array(t_db_comp[2:]))

    return header, dist

def read_fasta(p_fasta):

    sequences = {}

    n_seqs = 0
    with open(p_fasta, 'r') as f:
        seq = ''
        header = ''
        for l in f:
            if l.startswith('>'):
                if seq:
                    sequences[header] = seq
                    n_seqs += 1

                seq = ''
                header = l[1:].strip().replace('\t',' ')

            else:
                seq += l.strip()

        if seq:
            sequences[header] = seq
            n_seqs += 1

    if _verbose_:
        print('Found %d sequences in %s' % (n_seqs, p_fasta))

    return sequences

# Run
if __name__ == '__main__':

    parser = argparse.ArgumentParser(description='Search similar proteins by composition.')
    parser.add_argument('-i', dest='p_query', default='', help='Query FASTA')
    parser.add_argument('-fb', dest='p_db_fasta', default='', help='FASTA of database')
    parser.add_argument('-o', dest='p_out', default='', help='Output file name/path')
    parser.add_argument('-db', dest='p_db', default='', help='Database name/path')
    parser.add_argument('-v', dest='verbose', action='store_true', help='To verbose')

    args = parser.parse_args()

    _verbose_ = args.verbose

    if args.p_db_fasta:
        _, db_comps = build_db(args.p_db_fasta, args.p_db)
    else:
        db_comps = read_db(args.p_db)

    if args.p_query:
        search(args.p_query, db_comps, args.p_out)


#!/usr/bin/env python
#
# A script topographically correct SAR data using RIOS applier.
#
# Dan Clewley (clewley@usc.edu) - 05/02/2013
#
# The correction uses:
#
# sigma0_norm = sigma0 * (A_flat / A_slope) * (cos(theta_ref) / cos(theta_loc))^n
#
# Where n is a parameter related to optical thicknes. Defaults to 1.
#
# Castel et al. 2001. Sensitivity of space-borne SAR data to forest parameters over sloping terrain.
# Theory and experiment. International journal of remote sensing. 22(12) pp. 2351-2376
#
# Designed to take outputs of GAMMA:
# - Sigma0
# - Normalised pixel area (A_flat / A_slope) '.pix'
# - Local incidence angle '.inc'
#
# Can be run as a stand alone script or from within BatchGamma.py to process
# multiple files (different polarizations).
#
# Copyright 2014 Daniel Clewley. All rights reserved.
#
# Permission is hereby granted, free of charge, to any person
# obtaining a copy of this software and associated documentation
# files (the "Software"), to deal in the Software without restriction,
# including without limitation the rights to use, copy, modify,
# merge, publish, distribute, sublicense, and/or sell copies of the
# Software, and to permit persons to whom the Software is furnished
# to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be
# included in all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
# EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES
# OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.
# IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR
# ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF
# CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION
# WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
#
#

import sys
import argparse
from rios import applier
from rios import cuiprogress
import numpy as np
import os, glob

havescipy = True

try:
    from scipy import ndimage
except ImportError:
    havescipy = False

def getOutDriver(outFileName):

    """ Set output driver type and creation options
        based on file extension.

        Returns driver name and list of creation options
        as dictionary.
    """

    outControls = {}

    gdalFormat = 'ENVI'
    gdalCOOptions = []
    calcStats = False

    extension = os.path.splitext(outFileName)[-1]
    if extension == '.kea':
        gdalFormat = 'KEA'
        calcStats = True
    elif extension == '.tif':
        gdalFormat = 'GTiff'
        gdalCOOptions = ['COMPRESS=DEFLATE']
        calcStats = True
    elif extension == '.img':
        gdalFormat = 'HFA'
        gdalCOOptions = ['COMPRESSED=YES']
        calcStats = True
    elif extension == '.pix':
        gdalFormat = 'PCIDSK'
        gdalCOOptions = ['COMPRESSION=RLE']
        calcStats = True

    outControls['gdalFormat'] = gdalFormat
    outControls['gdalCOOptions'] = gdalCOOptions
    outControls['calcStats'] = gdalCOOptions

    return outControls

def castelCorrection(info, inputs, outputs, otherargs):
    """
    Apply topographic correction of Castel et al (2001)
    """
    theta_ref_deg = otherargs.thetaref
    nFactor = otherargs.nFactor
    filterSize = otherargs.filterSize
    theta_ref = np.deg2rad(theta_ref_deg)
    insigma=inputs.insigma.astype(np.float32)
    inpix=inputs.inpix.astype(np.float32)
    inlinc=inputs.inlinc.astype(np.float32)

    if havescipy and filterSize is not None:
        inpix = ndimage.uniform_filter(inpix,size=filterSize)
        inlinc = ndimage.uniform_filter(inlinc,size=filterSize)

    outputs.outimage = insigma * inpix * (np.cos(theta_ref) / np.cos(inlinc))

def runCorrection(insigma, inlinc, inpix, outsigma, thetaref=39.0, nFactor=1.0, filterSize=None):

    controls = applier.ApplierControls()

    # Set up input images
    infiles = applier.FilenameAssociations()
    infiles.insigma = insigma
    infiles.inlinc = inlinc
    infiles.inpix = inpix

    # Set up output image
    outfiles = applier.FilenameAssociations()
    outfiles.outimage = outsigma

    # Set format for output image
    outControls = getOutDriver(outsigma)

    # Set options
    controls.setOutputDriverName(outControls['gdalFormat'])
    controls.setCreationOptions(outControls['gdalCOOptions'])
    controls.setCalcStats(outControls['calcStats'])

    # Set up parameters
    otherargs = applier.OtherInputs()
    otherargs.thetaref = thetaref
    otherargs.nFactor = nFactor
    otherargs.filterSize = filterSize

    # Run correction
    controls.progress = cuiprogress.CUIProgressBar()
    applier.apply(castelCorrection, infiles, outfiles, otherargs, controls=controls)

def runCorrectionDIR(inDIR, sigmaExt='utm', lincExt='inc', pixExt='pix', outExt='kea', thetaref=39.0, nFactor=1, filterSize=None):
    """ Run correction for directory.
        Finds files based on supplied extension.
    """
    try:
        inlinc = glob.glob(inDIR + '/*' + lincExt)[0]
        inpix = glob.glob(inDIR + '/*' + pixExt)[0]
    except Exception as err:
        print("Couldn't find local incidence angle or pixel area images, is the extension correct?")
        print(err)

    inSigmaList = glob.glob(inDIR + '/*' + sigmaExt)

    for insigma in inSigmaList:
        insigmaBase = os.path.splitext(insigma)[0]
        outsigma = insigmaBase + '_topo.' + outExt
        runCorrection(insigma, inlinc, inpix, outsigma, thetaref, nFactor, filterSize)

if __name__ == '__main__':

    # Read in parameters

    parser = argparse.ArgumentParser()
    parser.add_argument("-i", "--insigma", type=str, required=True, help="Input sigma0 file")
    parser.add_argument("-p", "--inpix", type=str, required=True, help="Input pixel area")
    parser.add_argument("-l", "--inlinc", type=str, required=True, help="Input local incidence angle")
    parser.add_argument("-o", "--outsigma", type=str, required=True, help="Output topographically corrected sigma0")
    parser.add_argument("--thetaref", type=float, required=False, default=39.0, help="Reference incidence angle (default 39)")
    parser.add_argument("--n", type=float, required=False, default=1.0, help="n parameter (default 1)")
    parser.add_argument("--filterSize", type=int, required=False, default=None, help="Size of filter to apply to linc and pix data")
    args = parser.parse_args()

    # Run
    runCorrection(args.insigma, args.inlinc, args.inpix, args.outsigma, args.thetaref, args.n, args.filterSize)


#from phil import *
from sys import exit
import math
from os.path import exists
from pathlib import Path
import pandas as pd
import numpy as np
from . import util
from . import preprocess
from . import imhc_scoring
from . import cd8_scoring
from .tcrdist.all_genes import all_genes

cdr3_score_FG = 'fg'
cdr3_score_CENTER = 'cen'

cdr3_score_modes = [ cdr3_score_FG, cdr3_score_CENTER ]
default_cdr3_score_mode = cdr3_score_FG

fg_trim = 4
center_len = 5

aa_props_file = Path.joinpath( Path(util.path_to_data), 'aa_props.tsv')
assert exists(aa_props_file)
aa_props_df = pd.read_csv(aa_props_file, sep='\t')
aa_props_df.set_index('aa', inplace=True)

# all_tcr_scorenames = ['alphadist', 'cd8', 'cdr3len', 'imhc', 'mait', 'inkt'] +\
#                      [ '{}_{}'.format(x,y) for x in aa_props_df.columns for y in cdr3_score_modes ]

#tmp hacking SIMPLIFY -- dont include info on which version of the loop is used for scoring
all_tcr_scorenames = ['alphadist', 'cd8', 'cdr3len', 'imhc', 'mait', 'inkt', 'nndists_tcr'] + list(aa_props_df.columns)

amino_acids = ['A', 'C', 'D', 'E', 'F', 'G', 'H', 'I', 'K', 'L', \
               'M', 'N', 'P', 'Q', 'R', 'S', 'T', 'V', 'W', 'Y']

def read_cd8_score_params():
    # setup the scoring params
    # made by read_flunica_gene_usage_clustermaps*py
    infofile = Path.joinpath( Path(util.path_to_data), 'cd48_score_params_nomait.txt')

    scoretags = 'cdr3_len cdr3_aa gene'.split()

    all_scorevals = {'A':{}, 'B':{} }
    for tag in scoretags:
        all_scorevals['A'][tag] = {}
        all_scorevals['B'][tag] = {}

    min_l2e = -0.6
    max_l2e = 0.6

    # scores reflect cd8 versus cd4 preference

    for line in open(infofile,'r'):
        l = line.split()
        tag = l[0]
        assert tag.endswith('_cdx_pval:')
        tag = tag[:-10]
        key = l[10]
        if tag[:4] == 'cdr3':
            ab = tag[4]
            tag = tag[:4]+tag[5:]
        else:
            assert key[:2] == 'TR'
            ab = key[2]
        assert tag in scoretags
        assert l[6] == 'cd4:'
        cd4_mean = float(l[7])
        cd8_mean = float(l[9])
        tot = 0.5*(cd4_mean + cd8_mean)
        if not tot:
            l2e = 0.0
        elif cd8_mean==0:
            l2e = min_l2e
        else:
            l2e = max( min_l2e, min( max_l2e, math.log( cd8_mean / tot, 2.0 ) ) )

        all_scorevals[ab][tag][key] = l2e
        #print '{:7.3f} {} {} {}'.format(l2e,ab,tag,key)
    return all_scorevals

all_scorevals = read_cd8_score_params()

def cd8_score_tcr_chain( ab, v, j, cdr3 ):
    global all_scorevals

    cdr3_len_ranges = { 'A': ( min( int(x) for x in all_scorevals['A']['cdr3_len'] ),
                               max( int(x) for x in all_scorevals['A']['cdr3_len'] ) ),
                        'B': ( min( int(x) for x in all_scorevals['B']['cdr3_len'] ),
                               max( int(x) for x in all_scorevals['B']['cdr3_len'] ) ) }

    assert ab in 'AB'
    if '-' in v:
        v = v[:v.index('-')]
    if '-' in j:
        j = j[:j.index('-')]
    v_score = all_scorevals[ab]['gene'].get(v,0.0)
    j_score = all_scorevals[ab]['gene'].get(j,0.0)
    mn,mx = cdr3_len_ranges[ab]
    L = max( mn, min( mx, len(cdr3) ) )
    cdr3_len_score = all_scorevals[ab]['cdr3_len'][str(L)]
    cdr3_aa_score = 0.0
    if L>8:
        for aa in cdr3[4:-4]:
            cdr3_aa_score += all_scorevals[ab]['cdr3_aa'][aa]
    score = v_score + j_score + cdr3_len_score + cdr3_aa_score

    return ( score, v_score, j_score, cdr3_len_score, cdr3_aa_score )

def old_imhc_score_cdr3( cdr3 ): # cdr3 is untrimmed!
    '''This is the old version of the score, "fit" by staring at the sequences a little bit
    '''
    if len(cdr3) <= 8:
        return 0
    fgloop = cdr3[4:-4]
    return ( len(fgloop) + 3.0 * fgloop.count('C') + 2.0 * fgloop.count('W') + fgloop.count('R') + fgloop.count('K')
             + 0.5*fgloop.count('H') - fgloop.count('D') - fgloop.count('E') )

def old_imhc_score_tcr( tcr ):
    '''This is the old version of the score, "fit" by staring at the sequences a little bit
    '''
    return old_imhc_score_cdr3( tcr[0][2] ) + 2*old_imhc_score_cdr3( tcr[1][2] ) # double-weight the beta CDR3

def cd8_score_tcr( tcr ):
    atcr, btcr = tcr
    return ( cd8_score_tcr_chain( 'A', atcr[0], atcr[1], atcr[2] )[0] +
             cd8_score_tcr_chain( 'B', btcr[0], btcr[1], btcr[2] )[0] )

def is_human_mait_alpha_chain(atcr):
    return ( atcr[0].startswith('TRAV1-2') and
             ( atcr[1].startswith('TRAJ33') or
               atcr[1].startswith('TRAJ20') or
               atcr[1].startswith('TRAJ12') ) and
             len(atcr[2]) == 12 )

def is_mouse_mait_alpha_chain(atcr):
    return ( atcr[0].startswith('TRAV1-') and atcr[1].startswith('TRAJ33') and len(atcr[2]) == 12 )

def is_mouse_inkt_alpha_chain(atcr):
    return ( atcr[0].startswith('TRAV11') and atcr[1].startswith('TRAJ18') and len(atcr[2]) == 15 )

def is_human_inkt_tcr(tcr):
    # could also put some limits on cdr3s?
    return ( tcr[0][0].startswith('TRAV10') and
             tcr[0][1].startswith('TRAJ18') and
             len(tcr[0][2]) in [14,15,16] and  # 15 seems to be the consensus
             tcr[1][0].startswith('TRBV25') )

def mait_score_tcr(tcr, organism):
    if 'human' in organism:
        return float(is_human_mait_alpha_chain(tcr[0]))

    elif 'mouse' in organism:
        return float(is_mouse_mait_alpha_chain(tcr[0]))

    else:
        print('unrecognized organism:', organism)
        exit()
        return 0.

def inkt_score_tcr(tcr, organism):
    if 'human' in organism:
        return float(is_human_inkt_tcr(tcr))

    elif 'mouse' in organism:
        return float(is_mouse_inkt_alpha_chain(tcr[0]))

    else:
        print('unrecognized organism:', organism)
        exit()
        return 0.

def read_locus_order( remove_slashes_from_gene_names= False ):
    ''' returns  all_locus_order:  all_locus_order[ab][gene] = int(l[1])
    '''
    # read the gene order from imgt
    all_locus_order = {'A':{}, 'B':{}}
    for ab in 'AB':
        fn = f'imgt_tr{ab.lower()}_locus_order.txt'
        filename = Path.joinpath( Path( util.path_to_data ), fn)
        assert exists(filename)
        for line in open(filename,'r'):
            l = line.split()
            if l[0] == 'IMGT':continue
            assert len(l) in [2,3]
            gene = l[0]
            if '/' in gene and remove_slashes_from_gene_names:
                print( 'remove /',gene)
                gene = gene.replace('/','')
            all_locus_order[ab][gene] = int(l[1])
    return all_locus_order

all_locus_order = read_locus_order()

trav_list = [ x[1] for x in sorted( (y,x) for x,y in all_locus_order['A'].items() if x[:4] == 'TRAV' ) ]
traj_list = [ x[1] for x in sorted( (y,x) for x,y in all_locus_order['A'].items() if x[:4] == 'TRAJ' ) ]

def alphadist_score_tcr( tcr ):
    global trav_list
    global traj_list
    va, ja = tcr[0][:2]
    # print(va)
    # print(ja)
    if "-" in va :
        va = va[:va.index('-')]
    if "-" in ja :
        ja = ja[:ja.index('-')]
    if va in trav_list:
        va_dist = len(trav_list)-1 -trav_list.index(va)
    else:
        #print('alphadist_score_tcr: unrecognized va:', va)
        va_dist = 0.5*(len(trav_list)-1)

    if ja in traj_list:
        ja_dist = traj_list.index(ja)
    else:
        #print('alphadist_score_tcr: unrecognized ja:', ja)
        ja_dist = 0.5*(len(traj_list)-1)
    return va_dist + ja_dist

def cdr3len_score_tcr(tcr):
    ''' double-weight the beta chain
    '''
    return len(tcr[0][2]) + 2*len(tcr[1][2])

def property_score_cdr3(cdr3, score_name, score_mode):
    ''' Currently averages the score over the number of aas scores
    '''
    global aa_props_df
    global cdr3_score_CENTER
    global cdr3_score_FG
    global fg_trim
    global center_len

    col = aa_props_df[score_name]

    if score_mode == cdr3_score_CENTER:
        cdr3 = cdr3[1:] # trim off the first 'C', makes structurally symmetric
        if len(cdr3)<center_len:
            return np.mean(col)
        else:
            ntrim = (len(cdr3)-center_len)//2
            return sum(col[aa] for aa in cdr3[ntrim:ntrim+center_len])/center_len

    elif score_mode == cdr3_score_FG:
        fgloop = cdr3[fg_trim:-fg_trim]
        if fgloop:
            return sum( col[aa] for aa in fgloop )/len(fgloop)
        else:
            return np.mean(col) # was returning 0 here; prob should use aa-frequency-weighted average...
    else:
        print( 'property_score_cdr3:: unrecognized score_mode:', score_mode)
        exit()
        return 0.0

def property_score_tcr(tcr, score_name, score_mode, alpha_weight=1.0, beta_weight=1.0):
    return ( alpha_weight * property_score_cdr3(tcr[0][2], score_name, score_mode ) +
             beta_weight  * property_score_cdr3(tcr[1][2], score_name, score_mode ) )

def make_tcr_score_table(adata, scorenames):
    ''' Returns an array of the tcr scores of shape: (adata.shape[0], len(scorenames))
    '''
    global aa_props_df
    organism = adata.uns['organism']

    tcrs = preprocess.retrieve_tcrs_from_adata(adata)
    clusters_tcr = np.array(adata.obs['clusters_tcr'])

    organism_genes = all_genes[organism]
    genes = frozenset( organism_genes.keys())
    count_reps = frozenset( [ x.count_rep for x in organism_genes.values() ] )

    cols = []
    for name in scorenames:
        if name == 'cdr3len':
            cols.append( [ cdr3len_score_tcr(x) for x in tcrs ])
        elif name.startswith('tcr_cluster'):
            num = int(name[11:])
            cols.append( [ float(x==num) for x in clusters_tcr])
        elif name == 'alphadist':
            cols.append( [ alphadist_score_tcr(x) for x in tcrs ])
        elif name == 'oldcd8':# the 'old' cd8 score
            cols.append( [ cd8_score_tcr(x) for x in tcrs ])
        elif name == 'cd8': # see comparison between old/new in cd8_scoring.py
            cols.append(cd8_scoring.make_cd8_score_table_column(tcrs))
        elif name == 'old_imhc':
            cols.append( [ old_imhc_score_tcr(x) for x in tcrs ])
        elif name == 'imhc':
            cols.append( imhc_scoring.make_imhc_score_table_column(tcrs, aa_props_df))
        elif name == 'mait':
            organism = adata.uns['organism']
            cols.append( [ mait_score_tcr(x, organism) for x in tcrs ])
        elif name == 'inkt':
            organism = adata.uns['organism']
            cols.append( [ inkt_score_tcr(x, organism) for x in tcrs ])
        elif name == 'nndists_tcr':
            if 'nndists_tcr' not in adata.obs_keys():
                print('WARNING nndists_tcr score requested but not present in adata.obs!!!!')
                cols.append( np.zeros( adata.shape[0] ) )
            else:
                cols.append( np.array(adata.obs['nndists_tcr']) )
        elif name == 'N_ins': # number of N insertions
            if 'N_ins' not in adata.obs_keys():
                print('WARNING N_ins requested but not present in adata.obs!!!!')
                cols.append( np.zeros( adata.shape[0] ) )
            else:
                cols.append( np.array(adata.obs['N_ins']).astype(float) )
        elif name in genes:
            matched = False
            for i_ab,ab in enumerate('AB'):
                for i_vj,vj in enumerate('VJ'):
                    ii_genes = set([x for x,y in organism_genes.items() if y.chain==ab and y.region==vj ])
                    if name in ii_genes:
                        assert not matched
                        matched = True
                        cols.append( [float(x[i_ab][i_vj]==name) for x in tcrs])
            assert matched

        elif name in count_reps:
            matched = False
            for i_ab,ab in enumerate('AB'):
                for i_vj,vj in enumerate('VJ'):
                    ii_count_reps = set([x.count_rep for x in organism_genes.values() if x.chain==ab and x.region==vj ])
                    if name in ii_count_reps:
                        assert not matched
                        matched = True
                        cols.append( [float(organism_genes[x[i_ab][i_vj]].count_rep==name) for x in tcrs])
            assert matched

        else:
            score_mode = name.split('_')[-1]
            score_name = '_'.join( name.split('_')[:-1])
            if score_mode not in cdr3_score_modes:
                score_mode = default_cdr3_score_mode
                score_name = name
            cols.append( [ property_score_tcr(x, score_name, score_mode) for x in tcrs ] )
    table = np.array(cols).transpose()#[:,np.newaxis]
    #print( table.shape, (adata.shape[0], len(scorenames)) )
    assert table.shape == (adata.shape[0], len(scorenames))

    return table