from solving import bellmansolving
import numpy as np
import pytest

from util import constants as cs


@pytest.fixture(scope='module')
def init_mock():
    class Mock():
        pass
    mock = Mock()
    mock.vep_stay = np.ones(10, dtype=np.float64)
    mock.vfp_stay = np.ones(10, dtype=np.float64)/1.3
    mock.vup = np.zeros(10, dtype=np.float64)
    mock.vep_outside = np.zeros((10, cs.NJOBS))
    for i in range(10):
        mock.vep_outside[i,:] = i/10 + np.linspace(0,1, num=cs.NJOBS)

    mock.lambda_u = .4
    mock.offer_probs = np.ones(cs.NJOBS)/cs.NJOBS
    return mock

@pytest.mark.usefixtures("init_mock")
class TestBellmanSolving(object):

    def test_emax_outsideoffer(self, init_mock):
        ecv_e, ecv_f, stayprobs = bellmansolving.emax_outsideoffer(init_mock.vep_stay,
                                                   init_mock.vfp_stay,
                                                   init_mock.vup,
                                                   init_mock.vep_outside,
                                                   init_mock.offer_probs)
        assert np.all(ecv_e>=(init_mock.vep_stay - 1e-8))
        assert np.allclose(ecv_e, np.array([ 1.,  1.01111111, 1.03055556, 1.05833333, 1.09444444,
                                             1.13888889, 1.19444444, 1.25833333, 1.33055556, 1.41111111]))


    def test_emax_no_outside_offer(self, init_mock):
        ecv_e, ecv_f = bellmansolving.emax_if_no_outside_offer(init_mock.vep_stay,
                                                               init_mock.vfp_stay,
                                                               init_mock.vup)
        assert np.allclose(ecv_e, np.ones(10))

    def test_emax_unemployed(self, init_mock):
        ecv_unemp_worker = bellmansolving.emax_unemployed(init_mock.vup,
                                                          init_mock.vep_outside,
                                                          init_mock.lambda_u,
                                                          init_mock.offer_probs)
        assert np.allclose(ecv_unemp_worker, np.array([0.19 , 0.228, 0.266, 0.304,
                                                       0.342, 0.38 , 0.418, 0.456, 0.494,
                                                       0.532]))