from __future__ import annotations
import pytest

import numpy as np
import pandas as pd

from datashader.bundling import directly_connect_edges, hammer_bundle
from datashader.layout import circular_layout, forceatlas2_layout, random_layout

skimage = pytest.importorskip("skimage")


@pytest.fixture
def nodes():
    # Four nodes arranged at the corners of a 200x200 square with one node
    # at the center
    nodes_df = pd.DataFrame({'id': np.arange(5),
                             'x': [0, -100, 100, -100, 100],
                             'y': [0, 100, 100, -100, -100]})
    nodes_df.set_index('id')
    return nodes_df


@pytest.fixture
def edges():
    # Four edges originating from the center node and connected to each
    # corner
    edges_df = pd.DataFrame({'id': np.arange(4),
                             'source': np.zeros(4, dtype=int),
                             'target': np.arange(1, 5)})
    edges_df.set_index('id')
    return edges_df


@pytest.fixture
def weighted_edges():
    # Four weighted edges originating from the center node and connected
    # to each corner
    edges_df = pd.DataFrame({'id': np.arange(4),
                             'source': np.zeros(4, dtype=int),
                             'target': np.arange(1, 5),
                             'weight': np.ones(4)})
    edges_df.set_index('id')
    return edges_df


def test_immutable_nodes(nodes, edges):
    # Expect nodes to remain immutable after any bundling operation
    original = nodes.copy()
    directly_connect_edges(nodes, edges)
    assert original.equals(nodes)


@pytest.mark.parametrize('bundle', [directly_connect_edges, hammer_bundle])
def test_renamed_columns(nodes, weighted_edges, bundle):
    nodes = nodes.rename(columns={'x': 'xx', 'y': 'yy'})
    edges = weighted_edges.rename(columns={'source': 'src', 'target': 'dst', 'weight': 'w'})

    df = bundle(nodes, edges, x='xx', y='yy', source='src', target='dst', weight='w')

    assert 'xx' in df and 'x' not in df
    assert 'yy' in df and 'y' not in df
    assert 'w' in df and 'weight' not in df


@pytest.mark.parametrize('bundle', [directly_connect_edges, hammer_bundle])
@pytest.mark.parametrize('layout', [random_layout, circular_layout, forceatlas2_layout])
def test_same_path_endpoints(layout, bundle):
    # Expect path endpoints to match original edge source/target
    edges = pd.DataFrame({'id': [0], 'source': [0], 'target': [1]}).set_index('id')
    nodes = pd.DataFrame({'id': np.unique(edges.values)}).set_index('id')

    node_positions = layout(nodes, edges)
    bundled = bundle(node_positions, edges)

    source, target = edges.iloc[0]
    expected_source = node_positions.loc[source]
    expected_target = node_positions.loc[target]

    actual_source = bundled.loc[0]
    actual_target = bundled.loc[len(bundled)-2]

    assert np.allclose(expected_source, actual_source)
    assert np.allclose(expected_target, actual_target)


@pytest.mark.parametrize("include_edge_id", [True, False])
def test_directly_connect_with_weights(nodes, weighted_edges, include_edge_id):
    # Expect four lines starting at center (0.5, 0.5) and terminating
    # at a different corner and NaN
    data = pd.DataFrame({'edge_id':
                            [1.0, 1.0, np.nan, 2.0, 2.0, np.nan,
                             3.0, 3.0, np.nan, 4.0, 4.0, np.nan],
                         'x':
                            [0.0, -100.0, np.nan, 0.0, 100.0, np.nan,
                             0.0, -100.0, np.nan, 0.0, 100.0, np.nan],
                         'y':
                            [0.0, 100.0, np.nan, 0.0, 100.0, np.nan,
                             0.0, -100.0, np.nan, 0.0, -100.0, np.nan]})
    columns = ['edge_id', 'x', 'y'] if include_edge_id else ['x', 'y']
    expected = pd.DataFrame(data, columns=columns)

    given = directly_connect_edges(nodes, weighted_edges, include_edge_id=include_edge_id)
    assert given.equals(expected)


@pytest.mark.parametrize("include_edge_id", [True, False])
def test_directly_connect_without_weights(nodes, edges, include_edge_id):
    # Expect four lines starting at center (0.5, 0.5) and terminating
    # at a different corner and NaN
    data = pd.DataFrame({'edge_id':
                            [1.0, 1.0, np.nan, 2.0, 2.0, np.nan,
                             3.0, 3.0, np.nan, 4.0, 4.0, np.nan],
                         'x':
                            [0.0, -100.0, np.nan, 0.0, 100.0, np.nan,
                             0.0, -100.0, np.nan, 0.0, 100.0, np.nan],
                         'y':
                            [0.0, 100.0, np.nan, 0.0, 100.0, np.nan,
                             0.0, -100.0, np.nan, 0.0, -100.0, np.nan]})
    columns = ['edge_id', 'x', 'y'] if include_edge_id else ['x', 'y']
    expected = pd.DataFrame(data, columns=columns)

    given = directly_connect_edges(nodes, edges, include_edge_id=include_edge_id)
    assert given.equals(expected)


@pytest.mark.parametrize("include_edge_id", [True, False])
def test_hammer_bundle_with_weights(nodes, weighted_edges, include_edge_id):
    # Expect four lines starting at center (0.0, 0.0) and terminating
    # with NaN
    data = pd.DataFrame({'edge_id':
                            [1.0, np.nan, 2.0, np.nan,
                             3.0, np.nan, 4.0, np.nan],
                         'x':
                            [0.0, np.nan, 0.0, np.nan,
                             0.0, np.nan, 0.0, np.nan],
                         'y':
                            [0.0, np.nan, 0.0, np.nan,
                             0.0, np.nan, 0.0, np.nan],
                         'weight':
                            [1.0, np.nan, 1.0, np.nan,
                             1.0, np.nan, 1.0, np.nan]})
    columns = ['edge_id', 'x', 'y', 'weight'] if include_edge_id else ['x', 'y', 'weight']
    expected = pd.DataFrame(data, columns=columns)

    df = hammer_bundle(nodes, weighted_edges, include_edge_id=include_edge_id)

    starts = df[(df.x == 0.0) & (df.y == 0.0)]
    ends = df[df.isnull().any(axis=1)]
    given = pd.concat([starts, ends])
    given.sort_index(inplace=True)
    given.reset_index(drop=True, inplace=True)

    assert given.equals(expected)


@pytest.mark.parametrize("include_edge_id", [True, False])
def test_hammer_bundle_without_weights(nodes, edges, include_edge_id):
    # Expect four lines starting at center (0.0, 0.0) and terminating
    # with NaN
    data = pd.DataFrame({'edge_id':
                            [1.0, np.nan, 2.0, np.nan,
                             3.0, np.nan, 4.0, np.nan],
                         'x':
                            [0.0, np.nan, 0.0, np.nan,
                             0.0, np.nan, 0.0, np.nan],
                         'y':
                            [0.0, np.nan, 0.0, np.nan,
                             0.0, np.nan, 0.0, np.nan]})
    columns = ['edge_id', 'x', 'y'] if include_edge_id else ['x', 'y']
    expected = pd.DataFrame(data, columns=columns)

    df = hammer_bundle(nodes, edges, include_edge_id=include_edge_id)

    starts = df[(df.x == 0.0) & (df.y == 0.0)]
    ends = df[df.isnull().any(axis=1)]
    given = pd.concat([starts, ends])
    given.sort_index(inplace=True)
    given.reset_index(drop=True, inplace=True)

    assert given.equals(expected)