import pandas as pd
import pytest

import woodwork as ww
from woodwork.logical_types import (
    Boolean,
    Categorical,
    CountryCode,
    Datetime,
    Double,
    Integer,
    LogicalType,
    NaturalLanguage,
    Ordinal,
    PersonFullName,
    SubRegionCode,
    Unknown
)
from woodwork.type_sys.inference_functions import (
    categorical_func,
    double_func,
    integer_func
)
from woodwork.type_sys.type_system import TypeSystem


def test_type_system_init(default_inference_functions, default_relationships):
    type_sys = TypeSystem(inference_functions=default_inference_functions,
                          relationships=default_relationships)
    assert len(type_sys.inference_functions) == 5
    assert type_sys.inference_functions[Double] is double_func
    assert type_sys.inference_functions[Integer] is integer_func
    assert type_sys.inference_functions[Categorical] is categorical_func
    assert len(type_sys.relationships) == 2
    assert type_sys.relationships[0] == (Double, Integer)
    assert type_sys.relationships[1] == (Categorical, CountryCode)


def test_add_type_validation_errors(type_sys):
    error_msg = 'logical_type must be a valid LogicalType'
    with pytest.raises(TypeError, match=error_msg):
        type_sys.add_type(logical_type=1)

    with pytest.raises(TypeError, match=error_msg):
        type_sys.add_type(logical_type=Double())

    error_msg = 'inference_function must be a function'
    with pytest.raises(TypeError, match=error_msg):
        type_sys.add_type(logical_type=Ordinal, inference_function='not a function')

    error_msg = 'parent must be a valid LogicalType'
    with pytest.raises(ValueError, match=error_msg):
        type_sys.add_type(logical_type=Ordinal, parent=1)

    with pytest.raises(ValueError, match=error_msg):
        type_sys.add_type(logical_type=Ordinal, parent=Double())


def test_remove_type_validation_errors(type_sys):
    error_msg = 'logical_type must be a valid LogicalType'
    with pytest.raises(TypeError, match=error_msg):
        type_sys.add_type(logical_type=1)


def test_update_inference_function_validation_errors(type_sys):
    error_msg = 'logical_type must be a valid LogicalType'
    with pytest.raises(TypeError, match=error_msg):
        type_sys.update_inference_function(logical_type=1, inference_function=None)

    error_msg = 'inference_function must be a function'
    with pytest.raises(TypeError, match=error_msg):
        type_sys.update_inference_function(logical_type=Ordinal, inference_function='not a function')


def test_update_relationship_validation_errors(type_sys):
    error_msg = 'logical_type must be a valid LogicalType'
    with pytest.raises(TypeError, match=error_msg):
        type_sys.update_relationship(logical_type=1, parent=Ordinal)

    error_msg = 'parent must be a valid LogicalType'
    with pytest.raises(ValueError, match=error_msg):
        type_sys.update_relationship(logical_type=Ordinal, parent=1)


def test_type_system_default_type(default_inference_functions, default_relationships):
    type_sys = TypeSystem(inference_functions=default_inference_functions,
                          relationships=default_relationships,
                          default_type=SubRegionCode)
    assert type_sys.default_type == SubRegionCode
    type_sys.update_inference_function(Categorical, None)
    test_series = pd.Series(['a', 'b', 'c'])
    assert isinstance(type_sys.infer_logical_type(test_series), SubRegionCode)
    assert SubRegionCode in type_sys.registered_types


def test_type_system_default_type_remove_error(default_inference_functions, default_relationships):
    type_sys = TypeSystem(inference_functions=default_inference_functions,
                          relationships=default_relationships,
                          default_type=SubRegionCode)
    error_msg = "Default LogicalType cannot be removed"
    with pytest.raises(ValueError, match=error_msg):
        type_sys.remove_type(SubRegionCode)
    with pytest.raises(ValueError, match=error_msg):
        type_sys.remove_type('SubRegionCode')
    with pytest.raises(ValueError, match=error_msg):
        type_sys.remove_type('sub_region_code')


def test_type_system_registered_types(type_sys):
    assert isinstance(type_sys.registered_types, list)
    assert set(type_sys.registered_types) == {Double, Integer, Categorical, CountryCode, Unknown}


def test_type_system_root_types(type_sys):
    assert isinstance(type_sys.root_types, list)
    assert set(type_sys.root_types) == {Double, Categorical, Unknown}


def test_add_type_without_parent():
    type_sys = TypeSystem()
    type_sys.add_type(Double, inference_function=double_func)
    assert Double in type_sys.inference_functions.keys()
    assert type_sys.inference_functions[Double] is double_func

    type_sys.add_type(Integer)
    assert Integer in type_sys.inference_functions.keys()
    assert type_sys.inference_functions[Integer] is None


def test_add_type_with_parent():
    type_sys = TypeSystem()
    type_sys.add_type(Double, inference_function=double_func)
    type_sys.add_type(Integer, inference_function=integer_func, parent=Double)
    assert Integer in type_sys.inference_functions.keys()
    assert type_sys.inference_functions[Integer] is integer_func
    assert len(type_sys.relationships) == 1
    assert type_sys.relationships[0] == (Double, Integer)


def test_add_duplicate_ltype(type_sys):
    inference_fn = type_sys.inference_functions[ww.logical_types.Integer]

    assert ww.type_system.str_to_logical_type('Integer') == ww.logical_types.Integer

    class Integer(LogicalType):
        primary_dtype = 'string'

    error_msg = 'Logical Type with name Integer already present in the Type System. Please rename the LogicalType or remove existing one.'
    with pytest.raises(ValueError, match=error_msg):
        type_sys.add_type(Integer, inference_function=inference_fn)

    type_sys.remove_type(ww.logical_types.Integer)
    type_sys.add_type(Integer, inference_function=inference_fn)

    ltype = type_sys.str_to_logical_type('Integer')
    assert ltype.primary_dtype == 'string'
    assert ltype == Integer

    type_sys.reset_defaults()
    ltype = type_sys.str_to_logical_type('Integer')
    assert ltype.primary_dtype == 'int64'
    assert ltype == ww.logical_types.Integer


def test_remove_type_no_children(type_sys):
    type_sys.remove_type(Integer)
    assert len(type_sys.inference_functions) == 4
    assert Integer not in type_sys.inference_functions.keys()
    assert len(type_sys.relationships) == 1
    type_sys.remove_type('CountryCode')
    assert len(type_sys.inference_functions) == 3
    assert CountryCode not in type_sys.inference_functions.keys()
    assert len(type_sys.relationships) == 0


def test_remove_type_with_children(type_sys):
    type_sys.remove_type(Double)
    assert len(type_sys.inference_functions) == 4
    assert Double not in type_sys.inference_functions.keys()
    assert len(type_sys.relationships) == 2
    assert (None, Integer) in type_sys.relationships
    assert Integer in type_sys.root_types


def test_get_children(type_sys):
    type_sys.add_type(Ordinal, parent=Categorical)
    children = type_sys._get_children(Categorical)
    assert isinstance(children, list)
    assert len(children) == 2
    assert set(children) == {CountryCode, Ordinal}


def test_get_parent(type_sys):
    assert type_sys._get_parent(CountryCode) == Categorical
    assert type_sys._get_parent(Integer) == Double


def test_get_depth(type_sys):
    type_sys.add_type(Ordinal, parent=CountryCode)
    assert type_sys._get_depth(Categorical) == 0
    assert type_sys._get_depth(CountryCode) == 1
    assert type_sys._get_depth(Ordinal) == 2


def test_update_inference_function(type_sys):
    assert type_sys.inference_functions[Double] is double_func
    type_sys.update_inference_function(Double, integer_func)
    assert type_sys.inference_functions[Double] is integer_func
    type_sys.update_inference_function('CountryCode', integer_func)
    assert type_sys.inference_functions[CountryCode] is integer_func


def test_update_relationship_no_children(type_sys):
    type_sys.update_relationship(CountryCode, Integer)
    assert len(type_sys.relationships) == 2
    assert (Integer, CountryCode) in type_sys.relationships
    assert type_sys._get_parent(CountryCode) == Integer


def test_update_relationship_string_input(type_sys):
    type_sys.update_relationship('CountryCode', 'Integer')
    assert len(type_sys.relationships) == 2
    assert (Integer, CountryCode) in type_sys.relationships
    assert type_sys._get_parent(CountryCode) == Integer


def test_update_relationship_with_children(type_sys):
    type_sys.add_type(SubRegionCode, parent=CountryCode)
    type_sys.update_relationship(CountryCode, Integer)
    assert len(type_sys.relationships) == 3
    assert (Integer, CountryCode) in type_sys.relationships
    assert (CountryCode, SubRegionCode) in type_sys.relationships
    assert type_sys._get_children(CountryCode) == [SubRegionCode]


def test_inference_multiple_matches_same_depth(default_relationships):
    def always_true(series):
        return True
    inference_functions = {
        Categorical: always_true,
        Double: always_true,
        Integer: always_true,
        CountryCode: always_true,
    }
    type_sys = TypeSystem(inference_functions=inference_functions,
                          relationships=default_relationships,
                          default_type=NaturalLanguage)
    type_sys.update_inference_function(Integer, always_true)
    type_sys.update_inference_function(CountryCode, always_true)
    inferred_type = type_sys.infer_logical_type(pd.Series([1, 2, 3]))
    # Should match CountryCode - same depth as Integer, but CountryCode parent
    # (Categorical) is tried and found first
    assert isinstance(inferred_type, CountryCode)


def test_inference_multiple_matches_different_depths(default_relationships):
    def always_true(series):
        return True
    inference_functions = {
        Categorical: always_true,
        Double: always_true,
        Integer: always_true,
        CountryCode: always_true,
    }
    type_sys = TypeSystem(inference_functions=inference_functions,
                          relationships=default_relationships,
                          default_type=NaturalLanguage)
    type_sys.update_inference_function(Integer, always_true)
    type_sys.update_inference_function(CountryCode, always_true)
    type_sys.add_type(SubRegionCode, inference_function=always_true, parent=CountryCode)
    inferred_type = type_sys.infer_logical_type(pd.Series([1, 2, 3]))
    # Should match SubRegionCode as it is the deepest match
    assert isinstance(inferred_type, SubRegionCode)


def test_reset_defaults(type_sys, default_inference_functions, default_relationships):
    type_sys.update_inference_function('Integer', None)
    type_sys.update_relationship(CountryCode, parent=Unknown)
    type_sys.default_type = Categorical
    type_sys.reset_defaults()

    assert type_sys.inference_functions == default_inference_functions
    assert type_sys.relationships == default_relationships
    assert type_sys.default_type == Unknown


def test_get_logical_types():
    all_types = ww.type_system.registered_types
    logical_types = ww.type_system._get_logical_types()

    for logical_type in all_types:
        assert logical_types[logical_type.__name__] == logical_type
        assert logical_types[logical_type.type_string] == logical_type

    assert len(logical_types) == 2 * len(all_types)


def test_str_to_logical_type():
    all_types = ww.type_system.registered_types

    with pytest.raises(ValueError, match='String test is not a valid logical type'):
        ww.type_system.str_to_logical_type('test')
    assert ww.type_system.str_to_logical_type('test', raise_error=False) is None

    for logical_type in all_types:
        assert ww.type_system.str_to_logical_type(logical_type.__name__) == logical_type
        assert ww.type_system.str_to_logical_type(logical_type.type_string) == logical_type

    assert ww.type_system.str_to_logical_type('bOoLeAn') == Boolean
    assert ww.type_system.str_to_logical_type('person_full_NAME') == PersonFullName
    assert ww.type_system.str_to_logical_type('PersonFullnamE') == PersonFullName

    ymd = '%Y-%m-%d'
    datetime_with_format = ww.type_system.str_to_logical_type('datetime', params={'datetime_format': ymd})
    assert datetime_with_format.__class__ == Datetime
    assert datetime_with_format.datetime_format == ymd
    assert datetime_with_format == Datetime(datetime_format=ymd)

    datetime_no_format = ww.type_system.str_to_logical_type('datetime', params={'datetime_format': None})
    assert datetime_no_format.__class__ == Datetime
    assert datetime_no_format.datetime_format is None
    assert datetime_no_format == Datetime()

    # When parameters are supplied in a non-empty dictionary, the logical type gets instantiated
    assert ww.type_system.str_to_logical_type('person_full_NAME', params={}) == PersonFullName
    assert datetime_no_format != Datetime

    # Input a different type system
    new_type_sys = TypeSystem()
    with pytest.raises(ValueError, match='String Integer is not a valid logical type'):
        new_type_sys.str_to_logical_type('Integer')
    new_type_sys.add_type(Boolean)
    assert Boolean == new_type_sys.str_to_logical_type('Boolean')
