# ActivitySim
# See full license in LICENSE.txt.
from __future__ import print_function

import os.path

import numpy.testing as npt
import numpy as np
import pandas as pd
import pandas.util.testing as pdt
import pytest

from .. import inject

from .. import simulate


@pytest.fixture(scope='module')
def data_dir():
    return os.path.join(os.path.dirname(__file__), 'data')


@pytest.fixture(scope='module')
def spec_name(data_dir):
    return 'sample_spec.csv'


@pytest.fixture(scope='module')
def spec(data_dir, spec_name):
    return simulate.read_model_spec(
        file_name=spec_name,
        spec_dir=data_dir,
        description_name='description',
        expression_name='expression')


@pytest.fixture(scope='module')
def data(data_dir):
    return pd.read_csv(os.path.join(data_dir, 'data.csv'))


def test_read_model_spec(data_dir, spec_name):

    spec = simulate.read_model_spec(
        file_name=spec_name,
        spec_dir=data_dir,
        description_name='description', expression_name='expression')

    assert len(spec) == 4
    assert spec.index.name == 'expression'
    assert list(spec.columns) == ['alt0', 'alt1']
    npt.assert_array_equal(
        spec.values,
        [[1.1, 11], [2.2, 22], [3.3, 33], [4.4, 44]])


def test_eval_variables(spec, data):

    result = simulate.eval_variables(spec.index, data)

    expected = pd.DataFrame([
            [1, 0, 4, 1],
            [0, 1, 4, 1],
            [0, 1, 5, 1]],
            index=data.index, columns=spec.index)

    expected[expected.columns[0]] = expected[expected.columns[0]].astype(np.int8)
    expected[expected.columns[1]] = expected[expected.columns[1]].astype(np.int8)
    expected[expected.columns[2]] = expected[expected.columns[2]].astype(np.int64)
    expected[expected.columns[3]] = expected[expected.columns[3]].astype(int)

    print("\nexpected\n", expected.dtypes)
    print("\nresult\n", result.dtypes)

    pdt.assert_frame_equal(result, expected, check_names=False)


def test_simple_simulate(data, spec):

    inject.add_injectable("settings", {'check_for_variability': False})

    choices = simulate.simple_simulate(choosers=data, spec=spec, nest_spec=None)
    expected = pd.Series([1, 1, 1], index=data.index)
    pdt.assert_series_equal(choices, expected)


def test_simple_simulate_chunked(data, spec):

    inject.add_injectable("settings", {'check_for_variability': False})

    choices = simulate.simple_simulate(choosers=data, spec=spec, nest_spec=None, chunk_size=2)
    expected = pd.Series([1, 1, 1], index=data.index)
    pdt.assert_series_equal(choices, expected)
