import pytest
import pandas as pd
import numpy as np
from sandbox import get_production_func, get_optimal_allocation
from pandas.testing import assert_frame_equal

@pytest.fixture
def math_nan_df():
    df = pd.DataFrame(np.random.randint(0, 80, size=(10, 8)), columns=['math_t_1', 'math_t_2', 'math_t_3', 'math_t_4', 'math_t_5', 'math_t_6', 'math_t_7', 'math_t_8'])
    df = df.apply(lambda row: row.sort_values().values, axis=1, result_type='broadcast')
    df = df.applymap(lambda x: x if x <= 60 else None)
    return df

@pytest.fixture
def mem_nan_df():
    df = pd.DataFrame(np.random.randint(0, 80, size=(10, 10)), columns=['mem_t_1', 'mem_t_2', 'mem_t_3', 'mem_t_4', 'mem_t_5', 'mem_t_6', 'mem_t_7', 'mem_t_8','mem_t_9','mem_t_10'])
    df = df.apply(lambda row: row.sort_values().values, axis=1, result_type='broadcast')
    df = df.applymap(lambda x: x if x <= 60 else None)
    return df

def test_get_production_func_math(math_nan_df):
    df_res = get_production_func(math_nan_df, 'math')
    # Simple asserts
    assert isinstance(df_res, pd.Series)
    assert df_res.index[0] == 0
    assert df_res.index[-1] == 60

def test_get_production_func_mem(mem_nan_df):
    df_res = get_production_func(mem_nan_df, 'mem')
    assert isinstance(df_res, pd.Series)
    assert df_res.index[0] == 0
    assert df_res.index[-1] == 60

def test_get_production_func_nan_test(math_nan_df):
    df_res = get_production_func(math_nan_df, 'math')
    # This should give us all test nans
    nans = math_nan_df.notna()
    # By construction this should give us all
    vals = math_nan_df<=60
    assert_frame_equal(nans,vals)

def test_get_optimal_allocation_last_sec_max():
    mock_df = pd.DataFrame({
        'mem': [1, 2, 4, 6, 8]+list(np.zeros(56)),
        'math': list(np.zeros(56))+[3, 3, 4, 6, 12]
    })
    p_vars = {}
    get_optimal_allocation(mock_df, p_vars)
    assert p_vars['switch_point'] == 60
    assert p_vars['opt_math_score'] == 12
    assert p_vars['opt_mem_score'] == 1

def test_get_optimal_allocation_first_sec_max():
    mock_df = pd.DataFrame({
        'mem': list(np.zeros(56))+[1, 2, 4, 6, 8],
        'math': [14, 3, 4, 6, 12]+list(np.zeros(56))
    })
    p_vars = {}
    get_optimal_allocation(mock_df, p_vars)
    assert p_vars['switch_point'] == 0
    assert p_vars['opt_math_score'] == 14
    assert p_vars['opt_mem_score'] == 8

def test_get_optimal_allocation_multiple_max():
    mock_df = pd.DataFrame({
        'mem': list(np.zeros(56))+[2, 4, 4, 6, 8],
        'math': [2, 3, 4, 12, 14]+list(np.zeros(56))
    })
    p_vars = {}
    get_optimal_allocation(mock_df, p_vars)
    assert p_vars['switch_point'] == 3
    assert p_vars['opt_math_score'] == 12
    assert p_vars['opt_mem_score'] == 4