# -*- coding: utf-8 -*-
"""
Created on Fri Oct 30 00:18:12 2020

We want to assess the loss of a credit portfolio under stressed conditions.
The portfolio is made of a large number of loans from a large number of
borrowers.

@author: anneg
"""

import numpy as np
import pandas as pds
from scipy.stats import norm

from cerm.settings import D_N_D, F, G, K, N_MC, SEED, T, W
from cerm import data


def cumulate_matrix(matrix):
    """ Cumulates a matrix (for QA).
        Shape (K, K)"""
    cum_matrix = np.flipud(np.transpose(matrix))
    cum_matrix = np.cumsum(cum_matrix, axis=0)
    cum_matrix = np.transpose(np.flipud(cum_matrix))
    return cum_matrix


def estimate_matrix(matrix, t):
    """ Computes a migration matrix at t h.
        Shape (K, K)"""
    m_m = np.linalg.matrix_power(matrix, t)
    return m_m


def cum_mat_t(matrix):
    """ Computes a migration matrix up to t h.
        Shape (G, T, K, K)"""
    cum_mat = np.zeros([G, T, K, K])
    for g in range(G):
        for h in range(T):
            mat = matrix[g, h, :, :]
            cum_mat[g, h, :, :] = np.linalg.matrix_power(mat, h+1)
    return cum_mat


def validate_matrix(matrix):
    """ Validates that a transition matrix has the right properties.
        Shape (K,K)"""
    if matrix.shape[0] != matrix.shape[1]:
        print("Matrix is not square.", matrix.shape)
    if np.amin(matrix) < 0:
        print("Negative Probabilities")
    if np.amax(matrix) > 1:
        print("Probabilities Larger than.")
        # checking row sums of matrix
    if np.amax(np.sum(matrix)-1) > 0:
        print("Row sum not equal to one", np.amax(np.sum(matrix)-1))


# Unconditional risk measures.

def expected_loss_g_i_t(u_m, lgd, ead, new):
    """ Expected loss that occur at time t for an initial rating i.
        Shape (G, K, T+1)"""
    el = np.zeros([G, K, T+1])
    e_l = np.zeros([K, K])
    i_m = u_m[:, 0, :, :]  # initial state as matrix is constant
    m_m_new = data.mig_matrix_new_loans(new)
    m_m_new = m_m_new[:, 0, :, :]
    for group in range(G):
        for t in range(1, T+1):
            if D_N_D == 0:  # with migration matrices
                if t == 1:
                    el[group, :, 1] = i_m[group, :, K-1] * ead[group, :] * lgd[group, :]
                elif t >= 2:
                    i_m_w = (1-W) * i_m + W * m_m_new
                    m = estimate_matrix(i_m_w[group, :, :], t-1)
                    e_l = np.einsum('ij,j->ij', m, i_m_w[group, :, K-1])
                    e_l[:, K-1] = 0
                    el[group, :, t] = np.sum(e_l, axis=1) * ead[group, :] * lgd[group, :]
            elif D_N_D == 1:  # simplified default/no default model
                el[group, :, t] = i_m[group, :, 0] * ead[group, :] * lgd[group, :]
    return el


def expected_loss_t(u_m, lgd, ead, new):
    """ Total expected loss.
        Shape (K,)"""
    el = np.sum(expected_loss_g_i_t(u_m, lgd, ead, new), axis=0)
    el = np.sum(el, axis=0)
    return el


def expected_loss_i(u_m, lgd, ead, new):
    """ Total expected loss.
        Shape (K,)"""
    el = np.sum(expected_loss_g_i_t(u_m, lgd, ead, new), axis=2)
    el = np.sum(el, axis=0)
    return el


def expected_loss(u_m, average_lgd, ead, new):
    """ Total expected loss.
        Shape 1"""
    el = np.sum(expected_loss_g_i_t(u_m, average_lgd, ead, new))
    return el


# Conditional risk measures.


def thresholds_reg(u_m):
    """ Unconditional thresholds values.
        Shape (K,K)
        Can be simplified using th[th > 9.00] = 9.00
        """
    
    th = u_m.copy()[0, 0, :, :]
    nine = np.ones([K]) * 9.00
    cum_u_m = cumulate_matrix(u_m[0, 0, :, :])
    for n in range(0, K):
        th[n, :] = np.fmin(norm.ppf(cum_u_m[n, :]), nine)
        th[n, :] = np.fmax(th[n, :], -nine)
    return th


def thresholds_t(u_m, c, a_ti):
    """ Generates time dependent thresholds.
        Shape (G, T, K, K)"""
    th = thresholds_reg(u_m)
    prism = np.repeat(th[np.newaxis, :, :], T, axis=0)
    prism = np.repeat(prism[np.newaxis, :, :, :], G, axis=0)
    a_reg = data.loadings_a1(u_m, c, a_ti)
    b = data.c_a3(u_m, c, a_ti)
    rect1 = data.norm_c_2D(b, c)
    vect2 = data.norm_c_2D(a_reg, c)[:, :, 0]
    rect2 = np.repeat(vect2[:, :, np.newaxis], T, axis=-1)
    rect = 1 / np.sqrt(1 + rect1 - rect2)
    th_t = np.einsum('gtij, git-> gtij', prism, rect)
    return th_t


def uncond_mig_matrix_t(u_m, c, a_ti, new):
    """ Conditional migration matrices.
        Shape (G, T, K, K) for u_m_t and (G, T, K, K) for u_m_t_new."""
    u_m_t = u_m.copy()
    a3_th_t = thresholds_t(u_m, c, a_ti)
    for g in range(G):
        for j in range(0, K):
            if D_N_D == 0:
                if j == 0:
                    u_m_t[g, :, :, j] = 1-norm.cdf(a3_th_t[g, :, :, 1])
                elif (j >= 1) & (j <= K-2):
                    u_m_t[g, :, :, j] = norm.cdf(a3_th_t[g, :, :, j])- norm.cdf(a3_th_t[g, :, :, j+1])
                elif j == K-1:
                    u_m_t[g, :, :, j] = norm.cdf(a3_th_t[g, :, :, K-1])
            elif D_N_D == 1:
                u_m_t[g, :, :, 0] = norm.cdf(a3_th_t[g, :, :, 0])
    return u_m_t


def expected_loss_a3_g_i_t(u_m_t, lgd, ead, new):
    """ Expected loss at time t with approach 3.
        Shape (G, K, T+1)"""
    el = np.zeros([G, K, T+1])
    e_l = np.zeros([K, K])
    m_m_new = data.mig_matrix_new_loans(new)
    for g in range(G):
        i_m = u_m_t[g, 0, :, :]
        m = np.eye(K)
        for t in range(1, T+1):
            if D_N_D == 0:
                if t == 1:
                    el[g, :, 1] = i_m[:, K-1] * ead[g, :] * lgd[g, :]
                elif t >= 2:
                    u_m_t_new = (1-W) * u_m_t + W * m_m_new
                    m = np.einsum('ij, jk->ik', m, u_m_t_new[g, t-2, :, :])
                    e_l = np.einsum('ij, j->ij', m, u_m_t_new[g, t-1, :, K-1])
                    e_l[:, K-1] = 0
                    el[g, :, t] = np.sum(e_l, axis=1) * ead[g, :] * lgd[g, :]
            elif D_N_D == 1:
                el[g, :, t] = u_m_t[g, t-1, :, 0] * ead[g, :] * lgd[g, :]
    return el

def pandas_expected_loss_a3_g_i_t(u_m_t, lgd, ead, new):
    """ Distribution of losses at time t.
        long format (G, K, T)."""
    mat_3D = expected_loss_a3_g_i_t(u_m_t, lgd, ead, new)
    names = ['group_ID', 'credit_rating_rank', 'time_steps']
    index = pds.MultiIndex.from_product([range(s)for s in mat_3D.shape], names=names)
    df = pds.DataFrame({'expected_loss_g_i_t': mat_3D.flatten()}, index=index)['expected_loss_g_i_t']
    df = pds.Series.to_frame(df)
    df['expected_loss_g_i_t_ID'] = np.arange(len(df))
    df.insert(0, 'portfolio_ID', 'pilot 1 Bank A', allow_duplicates=False)
    return df


def expected_loss_a3_t(u_m_t, lgd, ead, new):
    """ Total expected loss.
        Shape (T,)"""
    el = np.sum(expected_loss_a3_g_i_t(u_m_t, lgd, ead, new), axis=0)
    el = np.sum(el, axis=0)
    return el


def expected_loss_a3_i(u_m_t, lgd, ead, new):
    """ Distribution of total losses.
        Shape K"""
    el = expected_loss_a3_g_i_t(u_m_t, lgd, ead, new)
    el = np.sum(el, axis=2)
    el = np.sum(el, axis=0)
    return el


def expected_loss_a3(u_m_t, lgd, ead, new):
    """ Distribution of total losses.
        Shape 1"""
    el = np.sum(expected_loss_a3_g_i_t(u_m_t, lgd, ead, new))
    return el


def hazards(c):
    """ Monte Carlo simulations.
        Shape (N_MC, T, F)"""
    mean = np.zeros([F])
    np.random.seed(SEED)
    z = np.random.multivariate_normal(mean, c, (N_MC, T))
    return z


def conditional_mig_matrix(u_m, c, a, new):
    """ Conditional migration matrices."""
    c_m = np.zeros([N_MC, G, T, K, K])
    c_th = thresholds_reg(u_m)
    z = hazards(c)
    div = np.sqrt(1-data.norm_c_2D(a, c))
    div = np.swapaxes(div, 1, 2)

    for p in range(N_MC):
        if p == N_MC-1:
            print('Conditional migration matrix')
        z_path = z[p, :, :]
        a_z = np.einsum('gitf,tf->gti', a, z_path)
        c_th_new = np.repeat(c_th[np.newaxis, :, :], T, axis=0)
        c_th_new = np.repeat(c_th_new[np.newaxis, :, :, :], G, axis=0)
        if D_N_D == 0:
            for j in range(0, K):
                if j == 0:
                    c_m[p, :, :, :, j] = 1-norm.cdf((c_th_new[:, :, :, 1]
                                                  - a_z)/div)
                elif (j >= 1) & (j <= K-2):
                    c_m[p, :, :, :, j] = norm.cdf((c_th_new[:, :, :, j]
                                               - a_z)/div) - norm.cdf((c_th_new[:, :, :, j+1] - a_z)/div)
                elif j == K-1:
                    c_m[p, :, :, :, j] = norm.cdf((c_th_new[:, :, :, K-1]
                                                - a_z)/div)
        elif D_N_D == 1:
            c_m[p, :, :, :, 0] = norm.cdf((c_th_new[:, :, :, 0] - a_z)/div)
    return c_m


def conditional_mig_matrix_t(u_m, c, a_ti, new):
    """ Conditional migration matrices."""
    c_m = np.zeros([N_MC, G, T, K, K])
    b = data.loadings_a3(u_m, c, a_ti)
    c_th_t = thresholds_t(u_m, c, a_ti)
    z = hazards(c)
    div = np.sqrt(1-data.norm_c_2D(b, c))
    div = np.swapaxes(div, 1, 2)
    for p in range(N_MC):
        if p == N_MC-1:
            print('Conditional_migration_matrix_t')
        z_path = z[p, :, :]
        a_z = np.einsum('gitf,tf->gti', b, z_path)
        if D_N_D == 0:
            for j in range(0, K):
                if j == 0:
                    c_m[p, :, :, :, j] = 1-norm.cdf((c_th_t[:, :, :, 1]
                                                  - a_z)/div)
                elif (j >= 1) & (j <= K-2):
                    c_m[p, :, :, :, j] = norm.cdf((c_th_t[:, :, :, j]
                                               - a_z)/div) - norm.cdf((c_th_t[:, :, :, j+1] - a_z)/div)
                elif j == K-1:
                    c_m[p, :, :, :, j] = norm.cdf((c_th_t[:, :, :, K-1]
                                                - a_z)/div)
        elif D_N_D == 1:
            c_m[p, :, :, :, 0] = norm.cdf((c_th_t[:, :, :, 0] - a_z) / div)
    return c_m


def conditional_lgd(lgd):
    """Conditional LGD.
       Shape (N_MC, K, T), (K,) at this stage"""
    c_lgd = lgd.copy()
    return c_lgd


def loss_g_i_t(c_m, lgd, ead, new):
    """ Distribution of losses at time t.
        Shape (N_MC, G, K, T)."""
    c_m_new = np.zeros([N_MC, G, T, K, K])
    m_m_new = data.mig_matrix_new_loans(new)
    lo_g_t_i = np.zeros([N_MC, G, K, T+1])

    for p in range(N_MC):
        if p == N_MC-1:
            print('Loss')
        for g in range(G):
            i_m = c_m[p, g, 0, :, :]
            m = np.eye(K)
            for t in range(T+1):
                if t == 1:
                    lo_g_t_i[p, g, :, 1] = i_m[:, K-1] * ead[g, :] * conditional_lgd(lgd)[g, :]
                elif t >= 2:
                    c_m_new[p, g, :, :, :] = (1-W) * c_m[p, g, :, :, :] + W * m_m_new[g, :, :, :]
                    m = np.einsum('ij, jk->ik', m, c_m_new[p, g, t-2, :, :])
                    l_o = np.einsum('ij, j->ij', m, c_m_new[p, g, t-1, :, K-1])
                    l_o[:, K-1] = 0
                    lo_g_t_i[p, g, :, t] = np.sum(l_o, axis=1) * ead[g, :] * conditional_lgd(lgd) [g, :]            
    return lo_g_t_i


def pandas_loss_p_g_i_t(c_m, lgd, ead, new):
    """ Distribution of losses at time t.
        long format (N_MC, G, K, T)."""
    mat_4D = loss_g_i_t(c_m, lgd, ead, new)
    names = ['paths', 'group_ID', 'credit_rating_rank', 'time_steps']
    index = pds.MultiIndex.from_product([range(s)for s in mat_4D.shape], names=names)
    df = pds.DataFrame({'loss_p_g_i_t': mat_4D.flatten()}, index=index)['loss_p_g_i_t']
    df = pds.Series.to_frame(df)
    df['loss_p_g_i_t_ID'] = np.arange(len(df))
    df.insert(0, 'portfolio_ID', 'pilot 1 Bank A', allow_duplicates=False)
    return df

def loss_t(c_m, lgd, ead, new):
    """ Distribution of losses at time t.
        Shape (N_MC, T)."""
    lo_t = np.sum(loss_g_i_t(c_m, lgd, ead, new), axis=2)
    lo_t = np.sum(lo_t, axis=1)
    return lo_t

def loss_g_i(c_m, lgd, ead, new):
    """ Distribution of losses at time t.
        Shape (N_MC, G, K)."""
    lo_g_i = np.sum(loss_g_i_t(c_m, lgd, ead, new), axis=2)
    return lo_g_i


def loss_i(c_m, lgd, ead, new):
    """ Distribution of losses for rating i.
        Shape (N_MC, K)."""
    lo_i = np.sum(loss_g_i_t(c_m, lgd, ead, new), axis=3)
    lo_i = np.sum(lo_i, axis=1)
    return lo_i


def loss(c_m, lgd, ead, new):
    """ Distribution of total losses."""
    lo = np.zeros([N_MC])
    lo = np.sum(loss_i(c_m, lgd, ead, new), axis=1)
    return lo

def pandas_loss(c_m, lgd, ead, new):
    """ Distribution of losses at time t.
        long format (N_MC)."""
    vect = loss(c_m, lgd, ead, new)
    df = pds.DataFrame(data=vect, columns=['loss'])
    df['path_ID'] = np.arange(len(df))
    df.insert(0, 'portfolio_ID', 'pilot 1 Bank A', allow_duplicates=False)
    return df