import numpy as np
import pandas as pd

from Task import SimpleCounterEncoder

def test_imports():
    with open('Task.py', 'r') as file:
        lines = ' '.join(file.readlines())
        assert 'import numpy' in lines
        assert lines.count('import') == 1
        assert 'sklearn' not in lines
        assert 'get_dummies' not in lines

def test_simple_counters_small():
    data = {'col_1': [0, 0, 0, 1, 1, 1], 'col_2': ['a', 'b', 'c', 'a', 'b', 'c'], 'col_3': [0, 1, 2, 3, 4, 5]}
    df_test = pd.DataFrame.from_dict(data)
    enc = SimpleCounterEncoder()
    enc.fit(df_test[['col_1', 'col_2']], df_test['col_3'])
    counts = enc.transform(df_test[['col_1', 'col_2']], a=1, b=1)

    print("Counters small")
    print(counts)

    ans = np.array([[1, 0.5, 4/3, 1.5, 1/3, 1.875],\
                    [1, 0.5, 4/3, 2.5, 1/3, 2.625],\
                    [1, 0.5, 4/3, 3.5, 1/3, 3.375],\
                    [4, 0.5, 10/3, 1.5, 1/3, 1.875],\
                    [4, 0.5, 10/3, 2.5, 1/3, 2.625],\
                    [4, 0.5, 10/3, 3.5, 1/3, 3.375]])
    assert len(counts.shape) == 2
    assert counts.shape[0] == 6
    assert counts.shape[1] == 6
    assert np.allclose(counts, ans, atol=1e-8)
    assert type(counts) == np.ndarray

def test_simple_counters_diff_shape_small():
    data_enc = {'col_1': [0, 0, 0, 1, 1, 1], 'col_2': ['a', 'b', 'c', 'a', 'b', 'c'], 'col_3': [0, 1, 2, 3, 4, 5]}
    data_trans = {'col_1': [0, 1, 0, 1, 0], 'col_2': ['c', 'a', 'b', 'b', 'c']}
    df_test_enc = pd.DataFrame.from_dict(data_enc)
    df_test_trans = pd.DataFrame.from_dict(data_trans)
    enc = SimpleCounterEncoder()
    enc.fit(df_test_enc[['col_1', 'col_2']], df_test_enc['col_3'])
    counts = enc.transform(df_test_trans[['col_1', 'col_2']], a=1, b=1)
    ans = np.array([[1, 0.5, 4/3, 3.5, 1/3, 3.375],\
                    [4, 0.5, 10/3, 1.5, 1/3, 1.875],\
                    [1, 0.5, 4/3, 2.5, 1/3, 2.625],\
                    [4, 0.5, 10/3, 2.5, 1/3, 2.625],\
                    [1, 0.5, 4/3, 3.5, 1/3, 3.375]])
    assert len(counts.shape) == 2
    assert counts.shape[0] == 5
    assert counts.shape[1] == 6
    assert np.allclose(counts, ans, atol=1e-8)
    assert type(counts) == np.ndarray

def test_simple_counters_big():
    data = {'col_1': [1, 0, 0, 1, 1, 0, 0, 1, 0, 1, 0, 2, 1, 2, 0, 2, 1, 2, 0, 0, 2, 0, 1, 2, 2, 0, 1, 1, 2, 0], 'col_2': [1, 1, 1, 1, 0, 4, 1, 0, 0, 3, 2, 1, 0, 3, 1, 1, 3, 4, 0, 1, 3, 4, 2, 4, 0, 3, 1, 2, 0, 4], 'col_3': [1, 1, 0, 1, 1, 0, 1, 0, 0, 1, 1, 1, 0, 1, 1, 0, 1, 1, 1, 1, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1], "target": [1, 1, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 0, 0, 1, 1, 0, 0, 1, 0, 0, 0, 1, 1, 1, 1, 1, 0]}
    df_test = pd.DataFrame.from_dict(data)
    enc = SimpleCounterEncoder()
    enc.fit(df_test[['col_1', 'col_2', 'col_3']], df_test['target'])
    counts = enc.transform(df_test[['col_1', 'col_2', 'col_3']], a=1, b=2)
    ans = np.array([[0.6, 0.3333333333333333, 0.6857142857142857, 0.4, 0.3333333333333333, 0.6, 0.631578947368421, 0.6333333333333333, 0.619586942038641 ],\
[0.3333333333333333, 0.4, 0.5555555555555556, 0.4, 0.3333333333333333, 0.6, 0.631578947368421, 0.6333333333333333, 0.619586942038641 ],\
[0.3333333333333333, 0.4, 0.5555555555555556, 0.4, 0.3333333333333333, 0.6, 0.36363636363636365, 0.36666666666666664, 0.5761843790012805 ],\
[0.6, 0.3333333333333333, 0.6857142857142857, 0.4, 0.3333333333333333, 0.6, 0.631578947368421, 0.6333333333333333, 0.619586942038641 ],\
[0.6, 0.3333333333333333, 0.6857142857142857, 0.5714285714285714, 0.23333333333333334, 0.7036247334754797, 0.631578947368421, 0.6333333333333333, 0.619586942038641 ],\
[0.3333333333333333, 0.4, 0.5555555555555556, 0.2, 0.16666666666666666, 0.5538461538461539, 0.36363636363636365, 0.36666666666666664, 0.5761843790012805 ],\
[0.3333333333333333, 0.4, 0.5555555555555556, 0.4, 0.3333333333333333, 0.6, 0.631578947368421, 0.6333333333333333, 0.619586942038641 ],\
[0.6, 0.3333333333333333, 0.6857142857142857, 0.5714285714285714, 0.23333333333333334, 0.7036247334754797, 0.36363636363636365, 0.36666666666666664, 0.5761843790012805 ],\
[0.3333333333333333, 0.4, 0.5555555555555556, 0.5714285714285714, 0.23333333333333334, 0.7036247334754797, 0.36363636363636365, 0.36666666666666664, 0.5761843790012805 ],\
[0.6, 0.3333333333333333, 0.6857142857142857, 1.0, 0.16666666666666666, 0.9230769230769231, 0.631578947368421, 0.6333333333333333, 0.619586942038641 ],\
[0.3333333333333333, 0.4, 0.5555555555555556, 0.6666666666666666, 0.1, 0.7936507936507935, 0.631578947368421, 0.6333333333333333, 0.619586942038641 ],\
[0.75, 0.26666666666666666, 0.7720588235294118, 0.4, 0.3333333333333333, 0.6, 0.631578947368421, 0.6333333333333333, 0.619586942038641 ],\
[0.6, 0.3333333333333333, 0.6857142857142857, 0.5714285714285714, 0.23333333333333334, 0.7036247334754797, 0.36363636363636365, 0.36666666666666664, 0.5761843790012805 ],\
[0.75, 0.26666666666666666, 0.7720588235294118, 1.0, 0.16666666666666666, 0.9230769230769231, 0.631578947368421, 0.6333333333333333, 0.619586942038641 ],\
[0.3333333333333333, 0.4, 0.5555555555555556, 0.4, 0.3333333333333333, 0.6, 0.631578947368421, 0.6333333333333333, 0.619586942038641 ],\
[0.75, 0.26666666666666666, 0.7720588235294118, 0.4, 0.3333333333333333, 0.6, 0.36363636363636365, 0.36666666666666664, 0.5761843790012805 ],\
[0.6, 0.3333333333333333, 0.6857142857142857, 1.0, 0.16666666666666666, 0.9230769230769231, 0.631578947368421, 0.6333333333333333, 0.619586942038641 ],\
[0.75, 0.26666666666666666, 0.7720588235294118, 0.2, 0.16666666666666666, 0.5538461538461539, 0.631578947368421, 0.6333333333333333, 0.619586942038641 ],\
[0.3333333333333333, 0.4, 0.5555555555555556, 0.5714285714285714, 0.23333333333333334, 0.7036247334754797, 0.631578947368421, 0.6333333333333333, 0.619586942038641 ],\
[0.3333333333333333, 0.4, 0.5555555555555556, 0.4, 0.3333333333333333, 0.6, 0.631578947368421, 0.6333333333333333, 0.619586942038641 ],\
[0.75, 0.26666666666666666, 0.7720588235294118, 1.0, 0.16666666666666666, 0.9230769230769231, 0.36363636363636365, 0.36666666666666664, 0.5761843790012805 ],\
[0.3333333333333333, 0.4, 0.5555555555555556, 0.2, 0.16666666666666666, 0.5538461538461539, 0.36363636363636365, 0.36666666666666664, 0.5761843790012805 ],\
[0.6, 0.3333333333333333, 0.6857142857142857, 0.6666666666666666, 0.1, 0.7936507936507935, 0.36363636363636365, 0.36666666666666664, 0.5761843790012805 ],\
[0.75, 0.26666666666666666, 0.7720588235294118, 0.2, 0.16666666666666666, 0.5538461538461539, 0.36363636363636365, 0.36666666666666664, 0.5761843790012805 ],\
[0.75, 0.26666666666666666, 0.7720588235294118, 0.5714285714285714, 0.23333333333333334, 0.7036247334754797, 0.36363636363636365, 0.36666666666666664, 0.5761843790012805 ],\
[0.3333333333333333, 0.4, 0.5555555555555556, 1.0, 0.16666666666666666, 0.9230769230769231, 0.631578947368421, 0.6333333333333333, 0.619586942038641 ],\
[0.6, 0.3333333333333333, 0.6857142857142857, 0.4, 0.3333333333333333, 0.6, 0.631578947368421, 0.6333333333333333, 0.619586942038641 ],\
[0.6, 0.3333333333333333, 0.6857142857142857, 0.6666666666666666, 0.1, 0.7936507936507935, 0.631578947368421, 0.6333333333333333, 0.619586942038641 ],\
[0.75, 0.26666666666666666, 0.7720588235294118, 0.5714285714285714, 0.23333333333333334, 0.7036247334754797, 0.631578947368421, 0.6333333333333333, 0.619586942038641 ],\
[0.3333333333333333, 0.4, 0.5555555555555556, 0.2, 0.16666666666666666, 0.5538461538461539, 0.631578947368421, 0.6333333333333333, 0.619586942038641 ]])
    assert len(counts.shape) == 2
    assert counts.shape[0] == 30
    assert counts.shape[1] == 9
    assert np.allclose(counts, ans, atol=1e-8)
    assert type(counts) == np.ndarray