""" Functions for generating multi-state life tables """

import numpy as np
import pandas as pd
import METER.wrangler as wr


def atrisk_and_transitions(data, transition_names, states):
    """
    Get the number of individuals in a given state at a each time, and number of transitions at each time.

    Parameters
    ----------
    data : pandas dataframe
        the data in wide format as generated by :py:func:`.wide_format`

    transition_names : list
        a list of the names of the columns that contain the transition times

    states : list
        the names of the states in the model, the entry into each of which will correspond to the columns in transition_names

    Returns
    ----------
    list
        a two element list with:

        - first element a dataframe containing the risk sets for each time point
        - second element a dataframe containing the number of transitions of each type at each time point

    """
    t_statuses = ['{}_status'.format(a) for a in transition_names]
    t_ages = ['{}_age'.format(a) for a in transition_names]
    column_positions = {}
    for i in range(0, len(data.columns)):
        column_positions[data.columns[i]] = i
    for i in range(0, len(transition_names)):
        t_statuses[i] = column_positions[t_statuses[i]]
        t_ages[i] = column_positions[t_ages[i]]
    max_age = int(data['final age'].max()) # need the final age so we know how many age specific arrays needed
    row_list = data.values.tolist()
    allowable_transitions = wr.allowed_transitions(states)
    transitions = []
    at_risk = []
    for i in range(0, max_age+1): # create an array containing the number of people in each state for every age in the study
        at_risk.append([])
        transitions.append([[0 for x in range(len(states))] for y in range(len(states))])
    for i in at_risk:
        for z in states:
            i.append(0)
    at_risk[0][0] = len(data) # initially, everyone is in state 1
    for i in row_list:
        current_state = 0
        perf_max_age = i[column_positions["final age"]]
        n = 1
        while n <= perf_max_age:
            at_risk[n][current_state] += 1
            for z in range(current_state, len(states)):
                if i[t_ages[z]] == n and i[t_statuses[z]] == 1:
                    transitions[n][current_state][z] += 1
                    current_state = z
            n += 1
    riskdf = pd.DataFrame(at_risk, columns=states)
    riskdf.index.name = 'age'
    transitionsdf = []
    for z in range(0, max_age+1):
        translst = []
        for i in allowable_transitions:
            translst.append(transitions[z][i[0]][i[1]])
        transitionsdf.append(translst)
    names = []
    for i in allowable_transitions:
        names.append(states[i[0]] + "->" + states[i[1]])
    transdf = pd.DataFrame(transitionsdf, columns=names)
    transdf.index.name = 'age'
    return([riskdf, transdf])


def transitionprobs_and_samplesizes(riskdf, transdf, states):
    """
    Obtain transition matrices governing transitions from each time point, and associated sample sizes.

    Parameters
    ----------

    riskdf : pandas dataframe
        a dataframe with the number of individuals in each state, indexed at each time point. This is the first output from :py:func:`.atrisk_and_transitions`

    transdf : pandas dataframe
        a dataframe containing the number of transitions of each type at each time point. The second output from :py:func:`.atrisk_and_transitions`

    states : list
        the names of the states in the model

    Returns
    ----------
    list
        a two element list with:

        - first element a list containing the transition matrix (as a numpy array) for each time point
        - second element a list containing a matrix of the associated sample sizes

    """
    numstates = len(states)
    raw_transitions = np.full((numstates, numstates, len(riskdf)), 0)
    raw_SS = np.full((numstates, numstates, len(riskdf)), 0)
    allowable_transitions = wr.allowed_transitions(states)
    for n in range(0, len(allowable_transitions)):
        fromstate = states[allowable_transitions[n][0]]
        tostate = states[allowable_transitions[n][1]]
        transition = fromstate + "->" + tostate
        raw_transitions[allowable_transitions[n][0], allowable_transitions[n][1]] = transdf[transition].to_numpy()
        raw_SS[allowable_transitions[n][0], allowable_transitions[n][1]] = riskdf[fromstate].to_numpy()
    samplesizes = raw_SS.transpose((2, 1, 0))
    transmat = np.divide(raw_transitions.transpose((2, 1, 0)), samplesizes,out=np.zeros(samplesizes.shape, dtype=float), where=samplesizes!=0)
    for i in range(0, numstates):
        transitions_out = []
        for z in allowable_transitions:
            if i == z[0]:
                transitions_out.append(z)
        for x in range(0, len(transmat)):
            stay_prob = 1
            for q in transitions_out:
                stay_prob -= transmat[x][q[1], q[0]]
            transmat[x][i,i] = stay_prob
    return [transmat, samplesizes]


def survivorship_vector(transmat, radix, initial_time, states):
    """
    Predict the expected proportion of a population being in a given state at each time point based on some
    initial time and number of individuals currently in each state. To get probabilities for one individual,
    simply proceed for a 1-individual population.

    Parameters
    ----------
    transmat : list
        a list of numpy matrices that represent the transition probabilities at each
        time point. This is the first output from :py:func:`.transitionprobs_and_samplesizes`.

    states : list
        the names of the states in the model

    radix : numpy array
        an initial condition for the number of subjects in each state we want to model, specified
        as a numpy vector. For example, if we wanted the probabilities of being in each state for
        1 individual starting at the first state in a 6-state model the radix would be generated by
        np.array([[1],[0],[0],[0],[0],[0]], dtype=float).

    initial_time : int
        the initial time to model the survivorship outcomes from

    Returns
    ----------
    pandas dataframe
        a dataframe with the expected proportion of the population in each state at each time point past the
        initial time given
    """
    rows = []
    l = radix
    for i in range(initial_time, len(transmat)):
        l = np.matmul(transmat[i], l)
        rows.append((l.transpose().flatten().tolist()))
    survdf = pd.DataFrame(rows, columns=states)
    survdf.index.name = "time"
    survdf.index += initial_time
    return survdf


def life_expectancy(transmat, initial_state, initial_time, states):
    """
    Predict the life expectancy of an individual based on current state and age

    Parameters
    ----------
    transmat : list
        a list of numpy matrices that represent the transition probabilities at each
        age. This is the first output from :py:func:`.transitionprobs_and_samplesizes`.

    states : list
        the names of the states in the model.

    initial_state : string
        the initial state to model the survivorship ourcomes from.

    initial_age : int
        the initial time to model the survivorship outcomes from

    Returns
    ----------
    float
        the life expectancy of an individual given the parameters specified
    """
    radix = np.zeros((len(states), 1), dtype=float)
    radix[states.index(initial_state), 0] = 1
    survdf = survivorship_vector(transmat, radix, initial_time, states)
    surv = survdf.values.tolist()
    life_ex = 0
    for i in surv:
        life_ex += 1 - i[-1]
    return life_ex


def KaplanMeier(data, termination):
    """
    get Kaplan Meier survival curve for a given set of data

    Parameters
    ----------
    data : pandas dataframe
        the dataframe in wide format as created by :py:func:`.wide_format`

    termination : string
        the termination transition (most commonly death) that we want to base the curve on

    Returns
    ----------
    list
        a two-element list with:

        - first element a list of age-specific survival proportions (points on a Kaplan-Meier curve)
        - second element the area under the curve (ie life expectancy)

    """
    curve = []
    s = 1
    for n in range(0, int(data['final age'].max() + 1)):
        curve.append(s)
        D = len(data[(data[termination + "_age"] == n) & (data[termination + '_status'] == 1)])
        N = len(data[(data['final age'] >= n - 1)])
        s *= (N - D) / N
    total = 0
    for i in range(0, len(curve)):
        total += curve[i]
    return [curve, total]


def censorLE(data, transition_names, states, initial_state, initial_time=0, censor_state='default', conditions=None):
    """
    Get the life expectancy of an individual based on initial state, initial time, and any censor states or covariate conditions.

    Parameters
    ----------
    data : pandas dataframe
        the data in wide format as generated by :py:func:`.wide_format`

    transition_names : list
        a list of the names of the columns that contain the transition times

    states: list
        the names of the states in the model

    initial_state : string
        the initial state that you want to estimate life expectancy from

    initial_time : int
        optional input if you want to estimate life expectancy after a given time (by default 0)

    censor_state : string
        a particular state that you want to restrict movement beyond (by default none)

    conditions: dictionary
        a set of conditions you want the group to be subject to (by default none).
        ex. {'Smoking': 'Yes', 'Race': 'White'}

    Returns
    ----------
    int
        the life expectancy of an individual starting in the initial state given the conditions provided

    """
    if conditions is None:
        conditions = {}
    if censor_state == 'default':
        censor_state = states[-2]
    if censor_state == states[-2]:
        cens_data = data
        cens_states = states
        cens_transitions = transition_names
    else:
        g_ind = states.index(censor_state)
        cens_data = wr.censor(data, transition_names[g_ind + 1], transition_names)
        cens_states = states[0:g_ind+1]
        cens_transitions = transition_names[0:g_ind+1]
        cens_states.append(states[-1])
        cens_transitions.append(transition_names[-1])
    for i in conditions:
        cens_data = cens_data[cens_data[i] == conditions[i]]
    atrisk = atrisk_and_transitions(cens_data, cens_transitions, cens_states)
    transmat = transitionprobs_and_samplesizes(atrisk[0], atrisk[1], cens_states)
    LE = life_expectancy(transmat[0], initial_state, initial_time, cens_states)
    return LE
