import numpy as np
from numba import njit

from solving import tech
from util import constants as cs


def solve_for_values_on_grid(params, value_functions):
    output = tech.production((params.beta0,
                              params.fixed_gamma,
                              params.mincer_gamma,
                              params.thetatypes,
                              params.ften0,
                              params.oten0),
                             value_functions.time,
                             value_functions.employed_worker_grid_points)
    ecv_u_pts = expected_continuation_value_unemployed_at_emp_grid_points(params, value_functions)
    val_unemp_grid = expected_continuation_value_unemployed_at_unemp_grid_points(params, value_functions)
    (ecv_e_pts, ecv_f_pts) = expected_continuation_values_emp_and_firm(params, value_functions)
    divided_surpluses = tech.nashbargain(output, ecv_e_pts, ecv_f_pts, ecv_u_pts)
    return [divided_surpluses['values_emp_on_grid'],
            divided_surpluses['values_firm_on_grid'],
            val_unemp_grid,
            divided_surpluses['wage_on_grid']]


def expected_continuation_value_unemployed_at_emp_grid_points(params, value_functions):
    return expected_continuation_value_unemployed(params, value_functions, "employed")

def expected_continuation_value_unemployed_at_unemp_grid_points(params, value_functions):
    return expected_continuation_value_unemployed(params, value_functions, "unemployed")

def expected_continuation_value_unemployed(params, value_functions, unemployed_or_employed_grid):
    value_function_prime_unemployment = value_functions[0]['value_function_unemployed']
    value_function_prime_employment = value_functions[0]['value_function_employed']

    if unemployed_or_employed_grid == 'employed':
        states = value_functions.employed_worker_grid_points[:,3:]
    elif unemployed_or_employed_grid == 'unemployed':
        states = value_functions.unemployed_worker_grid_points
    else:
        raise ValueError
    vup_stay = value_function_prime_unemployment(states)
    vep_outside = np.zeros((states.shape[0], cs.NJOBS))
    newstates = np.column_stack((np.zeros((states.shape[0], 3)), states))
    for jprime in range(cs.NJOBS):
        newstates[:, 0] = jprime
        vep_outside[:, jprime] = value_function_prime_employment(newstates)
    maxval = emax_unemployed(vup_stay, vep_outside, params.lambda_u, params.offer_probs)
    """ makes optimizing weakly better than always staying"""
    maxval += 1e-8
    assert np.alltrue(maxval >= (.95 * vup_stay))
    return maxval

def expected_continuation_values_emp_and_firm(params, value_functions):
    values_of_possible_outcomes = get_values_for_all_next_period_outcomes(params, value_functions)
    offer_specific_emaxs = calculate_emaxs_for_each_offer_set(params, values_of_possible_outcomes)
    (ecve, ecvf) = expectation_across_possible_offer_sets(params, offer_specific_emaxs)
    return ecve, ecvf

def get_values_for_all_next_period_outcomes(params, value_functions):
    states = np.copy(value_functions.employed_worker_grid_points)
    add_mincer_hc_to_states(params, states)
    staystates = state_vector_if_staying(states)
    (workers_next_period_values_in_same_match,
     firms_next_period_values_in_same_match) = same_match_values(value_functions, staystates)
    workers_next_period_value_in_unemp = unemp_match_values(value_functions, staystates)
    workers_next_period_value_at_outside = values_if_moving_outside(value_functions, states)
    return (workers_next_period_values_in_same_match,
            firms_next_period_values_in_same_match,
            workers_next_period_value_in_unemp,
            workers_next_period_value_at_outside)

def calculate_emaxs_for_each_offer_set(params, values_of_possible_outcomes):
    (workers_next_period_values_in_same_match,
     firms_next_period_values_in_same_match,
     workers_next_period_value_in_unemp,
     workers_next_period_value_at_outside) = values_of_possible_outcomes
    (expected_cv_employed_worker_nooffer,
     expected_cv_firm_nooffer) = emax_if_no_outside_offer(
                             workers_next_period_values_in_same_match,
                             firms_next_period_values_in_same_match,
                             workers_next_period_value_in_unemp)
    (expected_cv_employed_worker_offer,
     expected_cv_firm_offer,
     stayprobs) = emax_outsideoffer(
                                     workers_next_period_values_in_same_match,
                                     firms_next_period_values_in_same_match,
                                     workers_next_period_value_in_unemp,
                                     workers_next_period_value_at_outside,
                                     params.offer_probs
                                    )
    return (expected_cv_employed_worker_nooffer, expected_cv_firm_nooffer,
            expected_cv_employed_worker_offer, expected_cv_firm_offer,
            workers_next_period_value_in_unemp)

def expectation_across_possible_offer_sets(params, offer_specific_emaxs):
    (ecve_nooffer, ecvf_nooffer,
     ecve_offer, ecvf_offer,
     vup) = offer_specific_emaxs
    ecve = 0.95 * ((1 - params.lambda_e) * ecve_nooffer +
                   params.lambda_e * ecve_offer) * (1 - params.fire_rate)
    ecve += 0.95 * params.fire_rate * vup
    ecvf = 0.95 * ((1 - params.lambda_e) * ecvf_nooffer +
                   params.lambda_e * ecvf_offer) * (1 - params.fire_rate)
    return (ecve, ecvf)


def same_match_values(value_functions, staystates):
    value_emp_prime = value_functions[0]['value_function_employed']
    value_firm_prime = value_functions[0]['value_function_firm']
    return (value_emp_prime(staystates), value_firm_prime(staystates))


def unemp_match_values(value_functions, staystates):
    value_unemp_prime = value_functions[0]['value_function_unemployed']
    return value_unemp_prime(staystates[:, 3:])

def values_if_moving_outside(value_functions, states):
    value_employed_worker = value_functions[0]['value_function_employed']
    vep_outside = np.zeros((states.shape[0], cs.NJOBS))
    newstates = np.copy(states)
    current_j = states[:, 0]
    new_oten = np.minimum(states[:, 2] + 1, cs.N_TEN_PTS - 1)
    for jprime in range(cs.NJOBS):
        newstates[:, 0] = jprime
        newstates[:, 1] = 0
        newstates[:, 2] = (jprime == current_j) * new_oten
        vep_outside[:, jprime] = value_employed_worker(newstates) - 1e-8
    return vep_outside


def state_vector_if_staying(states):
    staystates = np.copy(states)
    staystates[:, 1] = np.minimum(staystates[:, 1] + 1, cs.N_TEN_PTS - 1)
    staystates[:, 2] = np.minimum(staystates[:, 2] + 1, cs.N_TEN_PTS - 1)
    return staystates


def add_mincer_hc_to_states(params, states):
    states[:,4] = tech.mincer_evo(params.mincer_rates,
                                 states[:, 0].astype(np.intp),
                                 states[:, 4])
    states[:,5] = tech.mincer_evo(params.mincer_rates,
                                 states[:, 0].astype(np.intp),
                                 states[:, 5])


@njit
def emax_unemployed(vup_stay, vep_outside, lambda_u, offerprobs):
    nstates = vep_outside.shape[0]
    ecv_u = np.zeros(nstates)
    for i in range(nstates):
        ecv_u[i] = (1-lambda_u) * vup_stay[i]
        for jobs_p in range(cs.NJOBS):
            ecv_u[i] += (lambda_u
                         * max(vup_stay[i],
                               vep_outside[i, jobs_p])
                         * offerprobs[jobs_p])
    return 0.95 * ecv_u


@njit
def emax_if_no_outside_offer(vep_stay, vfp_stay, vup):
    numpts = vep_stay.size
    ecv_e = np.zeros((numpts))
    ecv_f = np.zeros((numpts))
    for i in range(numpts):
        if vep_stay[i] >= vup[i]:
            ecv_f[i] = vfp_stay[i]
            ecv_e[i] = vep_stay[i]
        else:
            ecv_e[i] = vup[i]
    return ecv_e, ecv_f


@njit
def emax_outsideoffer(vep_stay, vfp_stay, vup, vep_outside, offerprobs):
    numpts = vep_stay.size
    ecv_e = np.zeros((numpts))
    ecv_f = np.zeros((numpts))
    stayprobs = np.zeros(numpts)
    for i in range(numpts):
        for jobs_p in range(cs.NJOBS):
            if ((vep_stay[i] + 1e-8) >= vep_outside[i, jobs_p]
                 and (vep_stay[i] + 1e-8)>=vup[i]):
                ecv_f[i] += offerprobs[jobs_p] * vfp_stay[i]
                ecv_e[i] += offerprobs[jobs_p] * vep_stay[i]
                stayprobs[i] += offerprobs[jobs_p]
            elif (vep_stay[i] < vep_outside[i, jobs_p]
                  and vep_outside[i, jobs_p] >= vup[i]):
                ecv_e[i] += offerprobs[jobs_p] * vep_outside[i, jobs_p]
            elif vup[i] > vep_outside[i, jobs_p] and vup[i] > vep_stay[i]:
                ecv_e[i] += offerprobs[jobs_p] * vup[i]
    return ecv_e, ecv_f, stayprobs