# (C) Copyright 2005-2021 Enthought, Inc., Austin, TX
# All rights reserved.
#
# This software is provided without warranty under the terms of the BSD
# license included in LICENSE.txt and may be redistributed only under
# the conditions described in the aforementioned license. The license
# is also available online at http://www.enthought.com/licenses/BSD.txt
#
# Thanks for using Enthought open source!

import unittest

from numpy import alltrue, arange, array
from numpy.testing import assert_array_equal

# Chaco imports
from chaco.api import (
    ArrayDataSource,
    DataRange1D,
    LinearMapper,
    PlotGraphicsContext,
    SegmentPlot,
    viridis,
)
from chaco.base import rgba_dtype


class SegmentPlotTest(unittest.TestCase):
    def setUp(self):
        self.size = (250, 250)
        index_data_source = ArrayDataSource(arange(10))
        index_range = DataRange1D()
        index_range.add(index_data_source)
        index_mapper = LinearMapper(range=index_range)

        value_data_source = ArrayDataSource(arange(1, 11))
        value_range = DataRange1D()
        value_range.add(value_data_source)
        value_mapper = LinearMapper(range=value_range)

        self.segment_plot = SegmentPlot(
            index=index_data_source,
            index_mapper=index_mapper,
            value=value_data_source,
            value_mapper=value_mapper,
            border_visible=False,
        )
        self.segment_plot.outer_bounds = list(self.size)

    def set_color_data(self):
        color_data_source = ArrayDataSource(arange(2, 7))
        color_range = DataRange1D()
        color_range.add(color_data_source)
        color_mapper = viridis(range=color_range)

        self.segment_plot.color_by_data = True
        self.segment_plot.color_data = color_data_source
        self.segment_plot.color_mapper = color_mapper

    def set_width_data(self):
        width_data_source = ArrayDataSource(arange(3, 8))
        width_range = DataRange1D()
        width_range.add(width_data_source)
        width_mapper = LinearMapper(low_pos=1, high_pos=10, range=width_range)

        self.segment_plot.width_by_data = True
        self.segment_plot.width_data = width_data_source
        self.segment_plot.width_mapper = width_mapper

    def test_segment(self):
        self.assertEqual(self.segment_plot.origin, "bottom left")
        self.assertIs(
            self.segment_plot.x_mapper, self.segment_plot.index_mapper
        )
        self.assertIs(
            self.segment_plot.y_mapper, self.segment_plot.value_mapper
        )

        gc = PlotGraphicsContext(self.size)
        gc.render_component(self.segment_plot)
        actual = gc.bmp_array[:, :, :]
        self.assertFalse(alltrue(actual == 255))

    def test_segment_orthogonal(self):
        self.segment_plot.render_style = "orthogonal"

        self.assertEqual(self.segment_plot.origin, "bottom left")
        self.assertIs(
            self.segment_plot.x_mapper, self.segment_plot.index_mapper
        )
        self.assertIs(
            self.segment_plot.y_mapper, self.segment_plot.value_mapper
        )

        gc = PlotGraphicsContext(self.size)
        gc.render_component(self.segment_plot)
        actual = gc.bmp_array[:, :, :]
        self.assertFalse(alltrue(actual == 255))

    def test_segment_quad(self):
        self.segment_plot.render_style = "quad"

        self.assertEqual(self.segment_plot.origin, "bottom left")
        self.assertIs(
            self.segment_plot.x_mapper, self.segment_plot.index_mapper
        )
        self.assertIs(
            self.segment_plot.y_mapper, self.segment_plot.value_mapper
        )

        gc = PlotGraphicsContext(self.size)
        gc.render_component(self.segment_plot)
        actual = gc.bmp_array[:, :, :]
        self.assertFalse(alltrue(actual == 255))

    def test_segment_cubic(self):
        self.segment_plot.render_style = "cubic"

        self.assertEqual(self.segment_plot.origin, "bottom left")
        self.assertIs(
            self.segment_plot.x_mapper, self.segment_plot.index_mapper
        )
        self.assertIs(
            self.segment_plot.y_mapper, self.segment_plot.value_mapper
        )

        gc = PlotGraphicsContext(self.size)
        gc.render_component(self.segment_plot)
        actual = gc.bmp_array[:, :, :]
        self.assertFalse(alltrue(actual == 255))

    def test_segment_color(self):
        self.set_color_data()

        gc = PlotGraphicsContext(self.size)
        gc.render_component(self.segment_plot)
        actual = gc.bmp_array[:, :, :]
        self.assertFalse(alltrue(actual == 255))

    def test_segment_color_orthogonal(self):
        self.segment_plot.render_style = "orthogonal"
        self.set_color_data()

        gc = PlotGraphicsContext(self.size)
        gc.render_component(self.segment_plot)
        actual = gc.bmp_array[:, :, :]
        self.assertFalse(alltrue(actual == 255))

    def test_segment_color_quad(self):
        self.segment_plot.render_style = "quad"
        self.set_color_data()

        gc = PlotGraphicsContext(self.size)
        gc.render_component(self.segment_plot)
        actual = gc.bmp_array[:, :, :]
        self.assertFalse(alltrue(actual == 255))

    def test_segment_color_cubic(self):
        self.segment_plot.render_style = "cubic"
        self.set_color_data()

        gc = PlotGraphicsContext(self.size)
        gc.render_component(self.segment_plot)
        actual = gc.bmp_array[:, :, :]
        self.assertFalse(alltrue(actual == 255))

    def test_segment_width(self):
        self.set_width_data()

        gc = PlotGraphicsContext(self.size)
        gc.render_component(self.segment_plot)
        actual = gc.bmp_array[:, :, :]
        self.assertFalse(alltrue(actual == 255))

    def test_segment_width_orthogonal(self):
        self.segment_plot.render_style = "orthogonal"
        self.set_width_data()

        gc = PlotGraphicsContext(self.size)
        gc.render_component(self.segment_plot)
        actual = gc.bmp_array[:, :, :]
        self.assertFalse(alltrue(actual == 255))

    def test_segment_width_quad(self):
        self.segment_plot.render_style = "quad"
        self.set_width_data()

        gc = PlotGraphicsContext(self.size)
        gc.render_component(self.segment_plot)
        actual = gc.bmp_array[:, :, :]
        self.assertFalse(alltrue(actual == 255))

    def test_segment_width_cubic(self):
        self.segment_plot.render_style = "cubic"
        self.set_width_data()

        gc = PlotGraphicsContext(self.size)
        gc.render_component(self.segment_plot)
        actual = gc.bmp_array[:, :, :]
        self.assertFalse(alltrue(actual == 255))

    def test_segment_width_color(self):
        self.set_width_data()
        self.set_color_data()

        gc = PlotGraphicsContext(self.size)
        gc.render_component(self.segment_plot)
        actual = gc.bmp_array[:, :, :]
        self.assertFalse(alltrue(actual == 255))

    def test_segment_width_orthogonal_color(self):
        self.segment_plot.render_style = "orthogonal"
        self.set_width_data()
        self.set_color_data()

        gc = PlotGraphicsContext(self.size)
        gc.render_component(self.segment_plot)
        actual = gc.bmp_array[:, :, :]
        self.assertFalse(alltrue(actual == 255))

    def test_segment_width_quad_color(self):
        self.segment_plot.render_style = "quad"
        self.set_width_data()
        self.set_color_data()

        gc = PlotGraphicsContext(self.size)
        gc.render_component(self.segment_plot)
        actual = gc.bmp_array[:, :, :]
        self.assertFalse(alltrue(actual == 255))

    def test_segment_width_cubic_color(self):
        self.segment_plot.render_style = "cubic"
        self.set_width_data()
        self.set_color_data()

        gc = PlotGraphicsContext(self.size)
        gc.render_component(self.segment_plot)
        actual = gc.bmp_array[:, :, :]
        self.assertFalse(alltrue(actual == 255))

    def test_segment_alpha(self):
        self.segment_plot.alpha = 0.5

        gc = PlotGraphicsContext(self.size)
        gc.render_component(self.segment_plot)
        actual = gc.bmp_array[:, :, :]
        self.assertFalse(alltrue(actual == 255))

    def test_segment_orthogonal_alpha(self):
        self.segment_plot.render_style = "orthogonal"
        self.segment_plot.alpha = 0.5

        gc = PlotGraphicsContext(self.size)
        gc.render_component(self.segment_plot)
        actual = gc.bmp_array[:, :, :]
        self.assertFalse(alltrue(actual == 255))

    def test_segment_quad_alpha(self):
        self.segment_plot.render_style = "quad"
        self.segment_plot.alpha = 0.5

        gc = PlotGraphicsContext(self.size)
        gc.render_component(self.segment_plot)
        actual = gc.bmp_array[:, :, :]
        self.assertFalse(alltrue(actual == 255))

    def test_segment_cubic_alpha(self):
        self.segment_plot.render_style = "cubic"
        self.segment_plot.alpha = 0.5

        gc = PlotGraphicsContext(self.size)
        gc.render_component(self.segment_plot)
        actual = gc.bmp_array[:, :, :]
        self.assertFalse(alltrue(actual == 255))

    def test_segment_selection(self):
        mask = array([True, True, False, False, True])
        self.segment_plot.index.metadata["selections"] = [mask]

        black = self.segment_plot.color_
        yellow = self.segment_plot.selection_color_
        expected_colors = array([yellow, yellow, black, black, yellow])
        expected_colors = expected_colors.astype("float32").view(rgba_dtype)
        expected_colors.shape = (5,)
        expected_colors["a"][~mask] *= 0.3

        assert_array_equal(mask, self.segment_plot.selected_mask)
        assert_array_equal(expected_colors, self.segment_plot.effective_colors)

        gc = PlotGraphicsContext(self.size)
        gc.render_component(self.segment_plot)
        actual = gc.bmp_array[:, :, :]
        self.assertFalse(alltrue(actual == 255))

    def test_segment_selection_color(self):
        mask = array([True, True, False, False, True])
        self.segment_plot.index.metadata["selections"] = [mask]
        self.set_color_data()

        color_data = self.segment_plot.color_data.get_data()
        colors = self.segment_plot.color_mapper.map_screen(color_data)
        expected_colors = colors.astype("float32").view(rgba_dtype)
        expected_colors.shape = (5,)
        expected_colors["a"][~mask] *= 0.3

        assert_array_equal(mask, self.segment_plot.selected_mask)
        assert_array_equal(expected_colors, self.segment_plot.effective_colors)

        gc = PlotGraphicsContext(self.size)
        gc.render_component(self.segment_plot)
        actual = gc.bmp_array[:, :, :]
        self.assertFalse(alltrue(actual == 255))