"""Unit tests for the bin_numeric_column function."""
import copd
import numpy as np
import pandas as pd


def test_binned_ages():
    """Test output is as expected for typical age binning values."""
    age_bins = [0, 50, 60, 70, 80, np.inf]
    labels = ['<50', '50-59', '60-69', '70-79', '80+']
    df = pd.DataFrame({'Age': [10, 49, 50, 55, 59, 60, 65, 69, 70, 75, 79, 80, 85, 100]})
    output = copd.bin_numeric_column(col=df['Age'], bins=age_bins, labels=labels)
    assert list(output.values) == ['<50', '<50', '50-59', '50-59', '50-59', '60-69',
                                   '60-69', '60-69', '70-79', '70-79', '70-79', '80+',
                                   '80+', '80+']


def test_binned_days_since():
    """Test output is as expected for typical days since last exac binning."""
    exac_bins = [-1, 0, 21, 90, 180, np.inf]
    labels = ['None', '<21 days', '21 - 89 days', '90 - 179 days', '>= 180 days']
    df = pd.DataFrame({'DaysSince': [-1, 0, 10, 21, 25, 89, 90, 150, 179, 180, 200]})
    output = copd.bin_numeric_column(col=df['DaysSince'], bins=exac_bins, labels=labels)
    assert list(output) == ['None', '<21 days', '<21 days', '21 - 89 days',
                            '21 - 89 days', '21 - 89 days', '90 - 179 days',
                            '90 - 179 days', '90 - 179 days', '>= 180 days',
                            '>= 180 days']


def test_binned_comorbidities():
    """Test output is as expected for typical comorbidity count binning."""
    comorb_bins = [0, 1, 3, np.inf]
    labels = ['None', '1-2', '3+']
    df = pd.DataFrame({'Comorbs': [0, 1, 2, 3, 4, 5]})
    output = copd.bin_numeric_column(col=df['Comorbs'], bins=comorb_bins, labels=labels)
    assert list(output) == ['None', '1-2', '1-2', '3+', '3+', '3+']