from analytics.tranche_functions import BCloss_recov_dist, GHquad
import numpy as np
import pandas as pd
import matplotlib as plt

gridsize = 500

def run_scenario(pool_size, rho, successprob, issuerweights, amount):
    #settings and running
    Z, w = GHquad(gridsize)
    S, _ = BCloss_recov_dist(successprob, issuerweights, amount,
                             rho, Z, w, Ngrid=gridsize+1, defaultflag=False)
    raised = np.linspace(0, pool_size, gridsize+1)
    df = pd.DataFrame(S.flatten(), columns=['Probability'], index=raised)
    return df

def plot_scenarios(df, bins = [0,0.1, 50, 100, 150, 200, 250, 300, 350, 400, 450, 500, 2000]):
    ''' takes run_scenario df '''
    df1 = df.groupby(pd.cut(df.index, bins, right=False)).sum()
    #have to manually label or else we lose ordering
    bins_label = []
    for this, next in zip(bins, bins[1:]):
        bins_label.append(str(this) + " to " + str(next))
    bins_label[0] = 'none'
    label = pd.DataFrame(bins_label, index=df1.index, columns=['bin'])
    df1 = df1.join(label).set_index('bin')
    ax = df1.plot(kind='bar')

    ax.legend().remove()
    ax.set_xlabel('Raised AUM (mm)')
    ax.set_ylabel('Probability (%)')
    y_ticks = ax.get_yticks()
    ax.set_yticklabels(['{:.0f}%'.format(y*100) for y in y_ticks])
    plt.pyplot.tight_layout()
    return ax

def stats(df):
    output = {}
    output['Expected Raised'] = (df.index.values * df.Probability).sum()
    output['Probability over 50mm'] = df[50:].sum().iloc[0]
    output['Raising No Capital'] = df.iloc[0][0]
    return output

def plot_prob_over(df):
    ax = (1-df.cumsum()).plot()
    ax.set_title('Probability of raising X amount of AUM')
    ax.set_xlabel('Amount Raised')
    ax.set_ylabel('Probability')
    return ax

def add_to_plot(df, ax):
    ax.plot(1-df.cumsum())
    return ax