# -------------------------------------------------------------------------------- NOTEBOOK-CELL: CODE
# -*- coding: utf-8 -*-
import dataiku
import pandas as pd, numpy as np
from dataiku import pandasutils as pdu
from TransitionMatrix import TransitionMatrix

# -------------------------------------------------------------------------------- NOTEBOOK-CELL: CODE
# Read recipe inputs
scenario_transition_matrices_probabilities = dataiku.Dataset("scenario_transition_matrices_probabilities")
scenario_transition_matrices_probabilities_df = scenario_transition_matrices_probabilities.get_dataframe()
last_positions = dataiku.Dataset("EAD_maturity_by_quarter")
last_positions_df = last_positions.get_dataframe()

# -------------------------------------------------------------------------------- NOTEBOOK-CELL: CODE
dates = sorted(last_positions_df['date'].unique())

# -------------------------------------------------------------------------------- NOTEBOOK-CELL: CODE
forecast_scenarios = scenario_transition_matrices_probabilities_df[scenario_transition_matrices_probabilities_df['date']>=min(dates)]
scenarios = forecast_scenarios['scenario'].unique()

# -------------------------------------------------------------------------------- NOTEBOOK-CELL: CODE
credit_statuses = sorted(last_positions_df['credit_status'].unique())

# -------------------------------------------------------------------------------- NOTEBOOK-CELL: CODE
risk_categories = [str(cs) for cs in credit_statuses]
risk_categories.append('D')

# -------------------------------------------------------------------------------- NOTEBOOK-CELL: CODE
long_term_matrix_data = matrix_data = forecast_scenarios[(forecast_scenarios['date']==dates[0]) & (forecast_scenarios['scenario']==scenarios[0])].reset_index(drop=True)
long_term_matrix_data['probability_scenario'] = long_term_matrix_data['probability']

# -------------------------------------------------------------------------------- NOTEBOOK-CELL: CODE
def compute_future_positions():
    future_positions = pd.DataFrame()
    transition_matrix = TransitionMatrix(risk_categories[:-1])

    for credit_status in credit_statuses:
        for scenario in scenarios:
            init_positions = pd.DataFrame({'credit_status': [str(cs) for cs in credit_statuses], 'EAD': [1 if cs==credit_status else 0 for cs in credit_statuses]})
            for date in dates:
                matrix_data = forecast_scenarios[(forecast_scenarios['date']==date) & (forecast_scenarios['scenario']==scenario)].reset_index(drop=True)
                if len(matrix_data) == 0:
                    matrix_data = long_term_matrix_data
                matrix_data['credit_status'] = matrix_data['credit_status'].astype('str')
                transition_matrix.input_matrix_data(matrix_data)
                next_positions = transition_matrix.forecast_next_positions(init_positions)
                saved_positions = pd.DataFrame({'credit_status': risk_categories})
                saved_positions['EAD'] = next_positions
                init_positions = saved_positions.copy()
                saved_positions['date'] = date
                saved_positions['scenario'] = scenario
                saved_positions['start_credit_status'] = credit_status
                future_positions = future_positions.append(saved_positions)
    return future_positions

# -------------------------------------------------------------------------------- NOTEBOOK-CELL: CODE
future_positions = compute_future_positions()

# -------------------------------------------------------------------------------- NOTEBOOK-CELL: CODE
future_positions

# -------------------------------------------------------------------------------- NOTEBOOK-CELL: CODE
# Write recipe outputs
position_forecasts = dataiku.Dataset("probability_default_updates")
position_forecasts.write_with_schema(future_positions)