import pandas as pd
import pytest

from finn._tests.utils import assert_colors_equal
from finn.layers.utils.color_encoding import (
    ColorEncoding,
    ConstantColorEncoding,
    DirectColorEncoding,
    ManualColorEncoding,
    NominalColorEncoding,
    QuantitativeColorEncoding,
)


def make_features_with_no_columns(*, num_rows) -> pd.DataFrame:
    return pd.DataFrame({}, index=range(num_rows))


@pytest.fixture
def features() -> pd.DataFrame:
    return pd.DataFrame(
        {
            "class": ["a", "b", "c"],
            "confidence": [0.5, 1, 0],
            "custom_colors": ["red", "green", "cyan"],
        }
    )


def test_constant_call_with_no_rows():
    features = make_features_with_no_columns(num_rows=0)
    encoding = ConstantColorEncoding(constant="red")

    values = encoding(features)

    assert_colors_equal(values, "red")


def test_constant_call_with_some_rows():
    features = make_features_with_no_columns(num_rows=3)
    encoding = ConstantColorEncoding(constant="red")

    values = encoding(features)

    assert_colors_equal(values, "red")


def test_manual_call_with_no_rows():
    features = make_features_with_no_columns(num_rows=0)
    array = ["red", "green", "cyan"]
    default = "yellow"
    encoding = ManualColorEncoding(array=array, default=default)

    values = encoding(features)

    assert_colors_equal(values, [])


def test_manual_call_with_fewer_rows():
    features = make_features_with_no_columns(num_rows=2)
    array = ["red", "green", "cyan"]
    default = "yellow"
    encoding = ManualColorEncoding(array=array, default=default)

    values = encoding(features)

    assert_colors_equal(values, ["red", "green"])


def test_manual_call_with_same_rows():
    features = make_features_with_no_columns(num_rows=3)
    array = ["red", "green", "cyan"]
    default = "yellow"
    encoding = ManualColorEncoding(array=array, default=default)

    values = encoding(features)

    assert_colors_equal(values, ["red", "green", "cyan"])


def test_manual_with_more_rows():
    features = make_features_with_no_columns(num_rows=4)
    array = ["red", "green", "cyan"]
    default = "yellow"
    encoding = ManualColorEncoding(array=array, default=default)

    values = encoding(features)

    assert_colors_equal(values, ["red", "green", "cyan", "yellow"])


def test_direct(features):
    encoding = DirectColorEncoding(feature="custom_colors")
    values = encoding(features)
    assert_colors_equal(values, list(features["custom_colors"]))


def test_direct_with_missing_feature(features):
    encoding = DirectColorEncoding(feature="not_class")
    with pytest.raises(KeyError):
        encoding(features)


def test_nominal_with_dict_colormap(features):
    colormap = {"a": "red", "b": "yellow", "c": "green"}
    encoding = NominalColorEncoding(
        feature="class",
        colormap=colormap,
    )

    values = encoding(features)

    assert_colors_equal(values, ["red", "yellow", "green"])


def test_nominal_with_dict_cycle(features):
    colormap = ["red", "yellow", "green"]
    encoding = NominalColorEncoding(
        feature="class",
        colormap=colormap,
    )

    values = encoding(features)

    assert_colors_equal(values, ["red", "yellow", "green"])


def test_nominal_with_missing_feature(features):
    colormap = {"a": "red", "b": "yellow", "c": "green"}
    encoding = NominalColorEncoding(feature="not_class", colormap=colormap)
    with pytest.raises(KeyError):
        encoding(features)


def test_quantitative_with_colormap_name(features):
    colormap = "gray"
    encoding = QuantitativeColorEncoding(feature="confidence", colormap=colormap)

    values = encoding(features)

    assert_colors_equal(values, [[c] * 3 for c in features["confidence"]])


def test_quantitative_with_colormap_values(features):
    colormap = ["black", "red"]
    encoding = QuantitativeColorEncoding(feature="confidence", colormap=colormap)
    values = encoding(features)
    assert_colors_equal(values, [[c, 0, 0] for c in features["confidence"]])


def test_quantitative_with_contrast_limits(features):
    colormap = "gray"
    encoding = QuantitativeColorEncoding(
        feature="confidence",
        colormap=colormap,
        contrast_limits=(0, 2),
    )

    values = encoding(features)

    assert encoding.contrast_limits == (0, 2)
    assert_colors_equal(values, [[c / 2] * 3 for c in features["confidence"]])


def test_quantitative_with_missing_feature(features):
    colormap = "gray"
    encoding = QuantitativeColorEncoding(feature="not_confidence", colormap=colormap)

    with pytest.raises(KeyError):
        encoding(features)


def test_validate_from_named_color():
    argument = "red"
    expected = ConstantColorEncoding(constant=argument)

    actual = ColorEncoding.validate(argument)

    assert actual == expected


def test_validate_from_sequence():
    argument = ["red", "green", "cyan"]
    expected = ManualColorEncoding(array=argument)

    actual = ColorEncoding.validate(argument)

    assert actual == expected


def test_validate_from_constant_dict():
    constant = "yellow"
    argument = {"constant": constant}
    expected = ConstantColorEncoding(constant=constant)

    actual = ColorEncoding.validate(argument)

    assert actual == expected


def test_validate_from_manual_dict():
    array = ["red", "green", "cyan"]
    default = "yellow"
    argument = {"array": array, "default": default}
    expected = ManualColorEncoding(array=array, default=default)

    actual = ColorEncoding.validate(argument)

    assert actual == expected


def test_validate_from_direct_dict():
    feature = "class"
    argument = {"feature": feature}
    expected = DirectColorEncoding(feature=feature)

    actual = ColorEncoding.validate(argument)

    assert actual == expected


def test_validate_from_nominal_dict():
    feature = "class"
    colormap = ["red", "green", "cyan"]
    argument = {"feature": feature, "colormap": colormap}
    expected = NominalColorEncoding(
        feature=feature,
        colormap=colormap,
    )

    actual = ColorEncoding.validate(argument)

    assert actual == expected


def test_validate_from_quantitative_dict(features):
    feature = "confidence"
    colormap = "gray"
    contrast_limits = (0, 2)
    argument = {
        "feature": feature,
        "colormap": colormap,
        "contrast_limits": contrast_limits,
    }
    expected = QuantitativeColorEncoding(
        feature=feature,
        colormap=colormap,
        contrast_limits=contrast_limits,
    )

    actual = ColorEncoding.validate(argument)

    assert actual == expected
