import unittest
from unittest.mock import Mock, MagicMock, PropertyMock
import pytest

from osgeo import ogr
import numpy as np
import pandas as pd
from shapely.geometry import Polygon as Poly

from autocnet.matcher import cpu_outlier_detector as od
from autocnet.examples import get_path
from autocnet.graph.network import CandidateGraph
from autocnet.io.geodataset import AGeoDataset
from .. import edge
from .. import node

import logging
log = logging.getLogger(__name__)

class TestEdge(unittest.TestCase):

    def setUp(self):
        source = Mock(node.Node)
        destination = Mock(node.Node)
        self.edge = edge.Edge(source=source, destination=destination)

    def test_masks(self):
        self.assertIsInstance(self.edge.masks, pd.DataFrame)

    def test_edge_overlap(self):
        e = edge.Edge()
        source = Mock(spec = node.Node)
        destination = Mock(spec = node.Node)
        e.destination = destination
        e.source = source
        geodata_s = Mock(spec = AGeoDataset)
        geodata_d = Mock(spec = AGeoDataset)
        source.geodata = geodata_s
        destination.geodata = geodata_d

        wkt1 = "POLYGON ((0 40, 40 40, 40 0, 0 0, 0 40))"
        wkt2 = "POLYGON ((20 60, 60 60, 60 20, 20 20, 20 60))"

        poly1 = ogr.CreateGeometryFromWkt(wkt1)
        poly2 = ogr.CreateGeometryFromWkt(wkt2)

        source.geodata.footprint = poly1
        destination.geodata.footprint = poly2

        e.overlap()
        self.assertEqual(e.weights['overlap_area'], 400)
        self.assertAlmostEqual(e.weights['overlap_percn'], 14.285714285)

    def test_weight(self):
        e = edge.Edge()
        assert e.weights == {}
        e.weights = ('foo', 1)
        assert e.weights['foo']  == 1
        e.weights = ('foo', 2)
        assert e.weights['foo']  == 2

    @pytest.mark.xfail
    def test_coverage(self):
        adjacency = get_path('two_image_adjacency.json')
        basepath = get_path('Apollo15')
        cg = CandidateGraph.from_adjacency(adjacency, basepath=basepath)
        keypoint_df = pd.DataFrame({'x': (15, 18, 18, 12, 12), 'y': (5, 10, 15, 15, 10)})
        keypoint_matches = [[0, 0, 1, 0],
                            [0, 1, 1, 1],
                            [0, 2, 1, 2],
                            [0, 3, 1, 3],
                            [0, 4, 1, 4]]

        matches_df = pd.DataFrame(data=keypoint_matches, columns=['source_image', 'source_idx', 'destination_image', 'destination_idx'])

        e = edge.Edge()
        source_node = MagicMock(spec=node.Node())
        destination_node = MagicMock(spec=node.Node())

        source_node.get_keypoint_coordinates = MagicMock(return_value=keypoint_df)
        destination_node.get_keypoint_coordinates = MagicMock(return_value=keypoint_df)

        e.source = source_node
        e.destination = destination_node

        source_geodata = PropertyMock(spec=AGeoDataset)
        destination_geodata = PropertyMock(spec=AGeoDataset)

        e.source.geodata = source_geodata
        e.destination.geodata = destination_geodata

        source_corners = [(0, 0),
                          (20, 0),
                          (20, 20),
                          (0, 20)]

        destination_corners = [(10, 5),
                               (30, 5),
                               (30, 25),
                               (10, 25)]

        e.source.geodata.latlon_corners = source_corners
        e.destination.geodata.latlon_corners = destination_corners

        vals = {(15, 5): (15, 5), (18, 10): (18, 10), (18, 15): (18, 15), (12, 15): (12, 15), (12, 10): (12, 10)}

        def pixel_to_latlon(i, j):
            return vals[(i, j)]

        e.source.geodata.pixel_to_latlon = MagicMock(side_effect=pixel_to_latlon)
        e.destination.geodata.pixel_to_latlon = MagicMock(side_effect=pixel_to_latlon)

        e.matches = matches_df

        #self.assertRaises(AttributeError, cg.edge[0][1].coverage)
        self.assertEqual(e.coverage(), 0.3)

    def test_get_keypoints(self):
        src_keypoint_df = pd.DataFrame({'x': (0, 1, 2, 3, 4), 'y': (5, 6, 7, 8, 9),
                                        'response': (10, 11, 12, 13, 14), 'size': (15, 16, 17, 18, 19),
                                        'angle': (20, 21, 22, 23, 24), 'octave': (25, 26, 27, 28, 29),
                                        'layer': (30, 31, 32, 33, 34)})

        dst_keypoint_df = pd.DataFrame({'x': (34, 33, 32, 31, 30), 'y': (29, 28, 27, 26, 25),
                                        'response': (24, 23, 22, 21, 20), 'size': (19, 18, 17, 16, 15),
                                        'angle': (14, 13, 12, 11, 10), 'octave': (9, 8, 7, 6, 5),
                                        'layer': (4, 3, 2, 1, 0)})

        keypoint_matches = [[0, 0, 1, 4],
                            [0, 1, 1, 3],
                            [0, 2, 1, 2],
                            [0, 3, 1, 1],
                            [0, 4, 1, 0]]

        matches_df = pd.DataFrame(data=keypoint_matches, columns=['source_image', 'source_idx',
                                                                  'destination_image', 'destination_idx'])
        source_node = node.Node(node_id=0)
        destination_node = node.Node(node_id=1)
        e = edge.Edge(source_node, destination_node)

        source_node.get_keypoint_coordinates = MagicMock(return_value=src_keypoint_df)
        destination_node.get_keypoint_coordinates = MagicMock(return_value=dst_keypoint_df)

        e.clean = MagicMock(return_value=(matches_df, None))
        e.matches = matches_df

        # Test all uses for edge.get_keypoints()
        src_matched_keypts = e.get_keypoints("source")
        src_matched_keypts2 = e.get_keypoints(e.source)
        dst_matched_keypts = e.get_keypoints("destination")
        dst_matched_keypts2 = e.get_keypoints(e.destination)

        # [output df to test] [name of node] [df to test against]
        to_test = [[src_matched_keypts, "source", src_keypoint_df],
                   [src_matched_keypts2, "source", src_keypoint_df],
                   [dst_matched_keypts, "destination", dst_keypoint_df],
                   [dst_matched_keypts2, "destination", dst_keypoint_df]]

        for out_df in to_test:
            # For each row index in the appropriate column of the matches_df,
            # assert that row index exists in the function's returned df
            [self.assertIn(row_idx, out_df[0].index.values)
             for row_idx in matches_df[out_df[1] + '_idx']]
            # For each row index in the returned df
            for row_idx in out_df[0].index.values:
                # Assert that row_idx exists in the matches_df's appropriate
                # column
                self.assertIn(row_idx, matches_df[out_df[1] + '_idx'])
                # Assert that all row_idx[column] vals returned by function
                # match their counterpart in orig df
                for column in out_df[0].columns:
                    self.assertTrue(out_df[0].iloc[row_idx][column] ==
                                            out_df[2].iloc[row_idx][column])

        # Test when overlap=True
        # edge["source_mbr"] and edge["destin_mbr"] haven't been calculated
        # yet, so return val should be unmasked df of node's keypts
        s_no_overlap = e.get_keypoints(e.source, overlap=True)
        d_no_overlap = e.get_keypoints(e.destination, overlap=True)
        self.assertTrue(s_no_overlap.equals(src_keypoint_df))
        self.assertTrue(d_no_overlap.equals(dst_keypoint_df))

        # Define the MBRs
        e.overlap_latlon_coords = 0, 0

        source_node.reproject_geom = MagicMock(return_value=Poly([(1, 6), (1, 8), (3, 6), (3, 8)]))
        source_node.keypoints = src_keypoint_df

        destination_node.reproject_geom = MagicMock(return_value=Poly([(31, 26), (31, 28), (33, 26), (33, 28)]))
        destination_node.keypoints = dst_keypoint_df

        # Only keep keypt vals w/I these bounding rects
        e["source_mbr"] = (0, 2, 5, 8)
        e["destin_mbr"] = (31, 33, 26, 28)

        # Grab the keypoints on our MBR overlaps
        s_overlap = e.get_keypoints(e.source, overlap=True)
        d_overlap = e.get_keypoints(e.destination, overlap=True)

        # Assert masked src keypts coords are equal
        s_expected = pd.DataFrame({'x': (0, 1, 2), 'y': (5, 6, 7)})
        self.assertTrue(np.array_equal(s_expected['x'], s_overlap['x'].values))
        self.assertTrue(np.array_equal(s_expected['y'], s_overlap['y'].values))

        # Assert masked dst keypt coords are equal
        d_expected = pd.DataFrame({'x': (33, 32, 31), 'y': (28, 27, 26)})
        self.assertTrue(np.array_equal(d_expected['x'], d_overlap['x'].values))
        self.assertTrue(np.array_equal(d_expected['y'], d_overlap['y'].values))

        # Assert type-checking in method throws proper errors
        with self.assertRaises(TypeError):
            e.get_keypoints("source", index = 456)
        with self.assertRaises(AttributeError):
            e.get_keypoints(1)
        # Check key error thrown when string arg != "source" or "destination"
        with self.assertRaises(AttributeError):
            e.get_keypoints("string")

    def test_eq(self):
        edge1 = edge.Edge()
        edge2 = edge.Edge()
        edge3 = edge.Edge()

        # Test edges w/ different keys are not equal, ones with same keys are
        edge1.__dict__["key"] = 1
        edge2.__dict__["key"] = 1
        edge3.__dict__["not_key"] = 1

        self.assertTrue(edge1 == edge2)
        self.assertFalse(edge1 == edge3)

        # Test edges with same keys, but diff df values
        edge1.__dict__["key"] = pd.DataFrame({'x': (0, 1, 2, 3, 4)})
        edge2.__dict__["key"] = pd.DataFrame({'x': (0, 1, 2, 3, 4)})
        edge3.__dict__["key"] = pd.DataFrame({'x': (0, 1, 2, 3, 5)})

        self.assertTrue(edge1 == edge2)
        self.assertFalse(edge1 == edge3)

        # Test edges with same keys, but diff np array vals
        # edge.__eq__ calls ndarray.all(), which checks that
        # all values in an array eval to true
        edge1.__dict__["key"] = np.array([True, True, True], dtype=bool)
        edge2.__dict__["key"] = np.array([True, True, True], dtype=bool)
        edge3.__dict__["key"] = np.array([True, True, False], dtype=bool)

        self.assertTrue(edge1 == edge2)
        self.assertFalse(edge1 == edge3)

    def test_repr(self):
        src = node.Node()
        dst = node.Node()
        masks = pd.DataFrame()

        e = edge.Edge()
        e.source = src
        e.destination = dst

        expected = """
        Source Image Index: {}
        Destination Image Index: {}
        Available Masks: {}
        """.format(src, dst, masks)

        self.assertEqual(expected, e.__repr__())

    @pytest.mark.filterwarnings('ignore::UserWarning')
    def test_ratio_check(self):
        """
        A pretty basic test that simply tests pass through from the edge to the
        proper distance func.
        """
        # Matches is init to None
        e = edge.Edge()

        # If there are matches...
        keypoint_matches = [[0, 0, 1, 4, 5],
                            [0, 1, 1, 3, 5],
                            [0, 2, 1, 2, 5],
                            [0, 3, 1, 1, 5],
                            [0, 4, 1, 0, 5]]

        matches_df = pd.DataFrame(data=keypoint_matches, columns=['source_image', 'source_idx',
                                                                  'destination_image', 'destination_idx', 'distance'])
        e.matches = matches_df
        expected = od.distance_ratio(None, matches_df)
        e.ratio_check()
        assert expected.equals(e.masks["ratio"])

    def test_overlap_check(self):
        s = node.Node(node_id=0)
        d = node.Node(node_id=1)

        e = edge.Edge(s, d)

        src_keypoint_df = pd.DataFrame({'x': (0, 1, 2, 3, 4), 'y': (5, 6, 7, 8, 9)})
        dst_keypoint_df = pd.DataFrame({'x': (34, 33, 32, 31, 30), 'y': (29, 28, 27, 26, 25)})

        # Create keypt matches
        keypoint_matches = [[0, 0, 1, 4, 5],
                            [0, 1, 1, 3, 5],
                            [0, 2, 1, 2, 5],
                            [0, 3, 1, 1, 5],
                            [0, 4, 1, 0, 5]]

        matches_df = pd.DataFrame(data=keypoint_matches,
                                  columns=['source_image', 'source_idx',
                                           'destination_image', 'destination_idx', 'distance'])
        s.keypoints = src_keypoint_df
        d.keypoints = dst_keypoint_df
        e.matches = matches_df

        s_overlap_keypts = pd.DataFrame({'x': (0, 1), 'y': (5, 6)})
        d_overlap_keypts = pd.DataFrame({'x': (31, 30), 'y': (26, 25)}, index=[3, 4])
        expected_mask = pd.Series(data=[True, True, False, False, False])

        # Mockup of the Edge.get_keypoints() method when overlap=True
        def mock_get_keypts(node, overlap=False):
            if node == s and overlap:
                return s_overlap_keypts
            elif node == d and overlap:
                return d_overlap_keypts
            else:
                return None

        e.get_keypoints = MagicMock(side_effect=mock_get_keypts)

        # Should fail if no src & dst mbrs on edge; Warns user & mask isn't
        # populated
        with self.assertLogs() as captured:
            e.overlap_check()
        self.assertTrue("Overlap between" not in captured.records[0].getMessage()) 

        # Should work after MBRs are set
        e["source_mbr"] = (1, 1, 1, 1)
        e["destin_mbr"] = (1, 1, 1, 1)
        e.overlap_check()
        overlap_matches, overlap_mask = e.clean(clean_keys=['overlap'])
        self.assertTrue(expected_mask.equals(overlap_mask))
        self.assertTrue(overlap_matches.equals(e.matches[overlap_mask]))

    def test_bad_matches_type(self):
        with self.assertRaises(TypeError):
            s = node.Node(node_id=0)
            d = node.Node(node_id=1)

            e = edge.Edge(s, d)
            e.matches = ['a', 'b', 'c']

    def test_prep_subpixel_1(self):
        nmatches = 10
        nstrengths = 1
        e = edge.Edge()
        arrs = e._prep_subpixel(nmatches, nstrengths=nstrengths)
        assert len(arrs) == 5
        assert arrs[2].shape == (nmatches, nstrengths)
        assert arrs[0][0] == 0

    def test_prep_subpixel_2(self):
        nmatches = 10
        nstrengths = 2
        e = edge.Edge()
        arrs = e._prep_subpixel(nmatches, nstrengths=nstrengths)
        assert len(arrs) == 5
        assert arrs[2].shape == (nmatches, nstrengths)
        assert arrs[0][0] == 0
