import unittest

import numpy

from solver import OptimizationSolver, DirectSolver
from surveyladder.geo import distance, chainage_offset_to_xy, poly_area


class TestDirectSolver(unittest.TestCase):

    def test_sanity_check_no_distances(self):
        solver = DirectSolver(['A', 'B'])
        with self.assertRaises(ValueError):
            solver.sanity_check()

    def test_sanity_check_not_enough_neighbors(self):
        solver = DirectSolver(['A', 'B', 'C'])
        solver.add_distance('A', 'B', 3)
        with self.assertRaises(ValueError):
            solver.sanity_check()
    
    def test_simple_triangle_only_distances(self):
        solver = DirectSolver(['A', 'B', 'C'])
        solver.add_distance('A', 'B', 3)
        solver.add_distance('B', 'C', 4)
        solver.add_distance('A', 'C', 5)

        points = solver.solve()

        numpy.testing.assert_almost_equal(distance(points['A'], points['B']), 3, decimal=2)
        numpy.testing.assert_almost_equal(distance(points['B'], points['C']), 4, decimal=2)
        numpy.testing.assert_almost_equal(distance(points['A'], points['C']), 5, decimal=2)


class TestSolver(unittest.TestCase):

    def test_distances_cost_optimal_solution(self):
        solver = OptimizationSolver(['A', 'B', 'C'])
        solver.add_distance('A', 'B', 3)
        solver.add_distance('B', 'C', 4)
        solver.add_distance('A', 'C', 5)
        optimal_solution = numpy.array([[0, 0], [3, 0], [3, 4]]).flatten()
        cost = solver.distances_cost(optimal_solution)
        self.assertAlmostEqual(cost, 0.0, places=3)

    def test_distances_cost_suboptimal_solution(self):
        solver = OptimizationSolver(['A', 'B', 'C'])
        solver.add_distance('A', 'B', 3)
        solver.add_distance('B', 'C', 4)
        solver.add_distance('A', 'C', 5)
        suboptimal_solution = numpy.array([[0, 0], [6, 0], [6, 8]]).flatten()
        # distance(A, B) = 6 instead of 3 -> (6-3)^2 = 9
        # distance(B, C) = 8 instead of 4 -> (8-4)^2 = 16
        # distance(A, C) = 10 instead of 5 -> (10-5)^2 = 25
        # total cost = 9 + 16 + 25 = 50
        cost = solver.distances_cost(suboptimal_solution)
        self.assertAlmostEqual(cost, 50.0, places=3)

    def test_offsets_cost_optimal_solution(self):
        solver = OptimizationSolver(['A', 'B', 'C'])
        solver.add_distance('A', 'B', 10)
        solver.add_offset('C', 'A', 'B', 5, 2)
        optimal_solution = numpy.array([[0, 0], [10, 0], [5, 2]]).flatten()
        cost = solver.offsets_cost(optimal_solution)
        self.assertAlmostEqual(cost, 0.0, places=3)

    def test_offsets_cost_suboptimal_solution(self):
        solver = OptimizationSolver(['A', 'B', 'C'])
        solver.add_distance('A', 'B', 10)
        solver.add_offset('C', 'A', 'B', 5, 2)
        suboptimal_solution = numpy.array([[0, 0], [10, 0], [5, 4]]).flatten()
        cost = solver.offsets_cost(suboptimal_solution)
        # The expected position of C is (5, 2), but the actual position is (5, 4),
        # so the cost should be (distance((5, 4), (5, 2)))^2 = 2^2 = 4
        self.assertAlmostEqual(cost, 4.0, places=3)


    def test_simple_triangle_only_distances_ccw(self):
        solver = OptimizationSolver(['A', 'B', 'C'])
        solver.add_distance('A', 'B', 3)
        solver.add_distance('B', 'C', 4)
        solver.add_distance('A', 'C', 5)

        points = solver.solve()

        numpy.testing.assert_almost_equal(distance(points['A'], points['B']), 3, decimal=2)
        numpy.testing.assert_almost_equal(distance(points['B'], points['C']), 4, decimal=2)
        numpy.testing.assert_almost_equal(distance(points['A'], points['C']), 5, decimal=2)

        triangle = numpy.array([points['A'], points['B'], points['C']])
        self.assertAlmostEqual(poly_area(triangle), 6.0, places=3)


    def test_simple_triangle_only_distances_cw(self):
        solver = OptimizationSolver(['A', 'C', 'B'])
        solver.add_distance('A', 'C', 3)
        solver.add_distance('B', 'C', 4)
        solver.add_distance('A', 'B', 5)

        points = solver.solve()

        numpy.testing.assert_almost_equal(distance(points['A'], points['C']), 3, decimal=2)
        numpy.testing.assert_almost_equal(distance(points['B'], points['C']), 4, decimal=2)
        numpy.testing.assert_almost_equal(distance(points['A'], points['B']), 5, decimal=2)

        triangle = numpy.array([points['A'], points['C'], points['B']])
        self.assertAlmostEqual(poly_area(triangle), -6.0, places=3)


    def test_simple_triangle_with_offset(self):
        solver = OptimizationSolver(['A', 'B', 'C'])
        solver.add_distance('A', 'B', 10)
        solver.add_offset('C', 'A', 'B', 5, 2)

        points = solver.solve()

        numpy.testing.assert_almost_equal(distance(points['A'], points['B']), 10, decimal=2)
        numpy.testing.assert_almost_equal(points['C'], chainage_offset_to_xy(points['A'], points['B'], 5, 2), decimal=2)


unittest.main()