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.arange(0, pool_size+pool_size/gridsize, step=pool_size/gridsize)
    df = pd.DataFrame(S.flatten(), columns=['Probability'], index=raised)
    return df

def plot_scenarios(df, pool_size):
    ''' takes run_scenario df '''
    bins = [0,0.01,50,100,150,200,250,300,350,400,450,500,pool_size]
    df1 = df.groupby(pd.cut(df.index,bins,right=False)).sum()
    ax = df1.plot(kind='bar')
    plt.pyplot.tight_layout()
    return ax

def prob_over(df, amount):
    return df[amount:].sum()

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