import pandas as pd
import numpy as np
import random


class Disease:
    def __init__(self, disease):
        self.disease = disease
        self.disease_data = pd.read_csv('../data/' + disease + '-genedatalist.csv')
        self.allbc = pd.read_csv('../data/betweenness.csv')
        self.seeds = pd.DataFrame({'gene': pd.read_csv('../data/' + disease + 'seeds.txt',
                                                       header=None, squeeze=True),
                                   'seed': True})
        self.score_method_name = ''

        # join disease dataframe with all disease betweenness
        # centrality, and drop the extra name column
        d = self.disease_data.merge(self.allbc, how='left',
                                    left_on='gene', right_on='name')
        del d['name']
        d.rename(columns={'betweenness': 'allbc'}, inplace=True)
        d.allbc = 1000 * d.allbc
        d = d.merge(self.seeds, how='left', on='gene').fillna(False)
        self.df = d

    def sort_gene(self):
        self.df.sort_index(by='gene', inplace=True)

    def sort_score_up(self):
        self.df.sort_index(by='score', inplace=True)

    def sort_score_down(self):
        self.df.sort_index(by='score', inplace=True, ascending=False)

    def compute_score(self, score_fn):
        self.score_method_name = score_fn.func_name.upper()
        self.df['score'] = score_fn(self.df.disbc, self.df.allbc, self.df.gcounts)
        self.sort_score_down()
        self.df['score_index'] = np.arange(len(self.df.score))

    def compute_filter(self, disbcf=0, allbcf=500, countf=0):
        d = self.df
        self.df['selected'] = ((d.disbc >= disbcf) & (d.allbc <= allbcf) &
                               (d.gcounts >= countf))

    def compute_seed_ranks(self):
        flag = self.df.seed & self.df.selected
        return float(sum(flag * self.df.score_index)) / sum(flag)

    def compute_seed_ranks_percentiles(self):
        return self.compute_seed_ranks() / sum(self.df.selected)

    def run_expt(self, score_fn, disbcf=0, allbcf=500, countf=0):
        self.compute_score(score_fn)
        self.compute_filter(disbcf, allbcf, countf)
        return self.compute_seed_ranks_percentiles()


def exponent(disbc, allbc, gcounts):
    return ((disbc / allbc)**2 * gcounts).fillna(0)


def subtract(disbc, allbc, gcounts):
    return disbc - allbc


def divide(disbc, allbc, gcounts):
    # divide and replace NaN with 0.
    return (disbc / allbc).fillna(0)


def count(disbc, allbc, gcounts):
    return gcounts


def random_shuffle(disbc, allbc, gcounts):
    n = len(gcounts)
    return random.sample(range(n), n)