import sys
import random
import itertools
import numpy as np
import cv2 as cv

MAP_FILE = 'cape_python.png'

# Przypisuje współrzędne rogów obszarów poszukiwań do pikseli na mapie.
SA1_CORNERS = (130, 265, 180, 315) # (LG X, LG Y, PD X, PD Y)
SA2_CORNERS = (80, 255, 130, 305)  # (LG X, LG Y, PD X, PD Y)
SA3_CORNERS = (105, 205, 155, 255) # (LG X, LG Y, PD X, PD Y)


class Search():
    """
    Bayesowska gra do symulacji misji poszukiwawczo-ratunkowych
    z trzema obszarami poszukiwań.
    """

    def __init__(self, name):
        self.name = name
        self.img = cv.imread(MAP_FILE, cv.IMREAD_COLOR)
        if self.img is None:
            print('Nie można załadować pliku z mapą {}.'.format(MAP_FILE),
                  file=sys.stderr)
            sys.exit(1)

        # Tworzy atrybuty do przechowywania rzeczywistej 
        # lokalizacji zaginionego.
        self.area_actual = 0
        # Lokalne współrzędne w obrębie obszaru poszukiwań
        self.sailor_actual = [0, 0]

        # Tworzy tablicę numpy dla każdego obszaru, wydzielając zakresy z mapy.
        self.sa1 = self.img[SA1_CORNERS[1] : SA1_CORNERS[3],
                            SA1_CORNERS[0] : SA1_CORNERS[2]]

        self.sa2 = self.img[SA2_CORNERS[1] : SA2_CORNERS[3],
                            SA2_CORNERS[0] : SA2_CORNERS[2]]

        self.sa3 = self.img[SA3_CORNERS[1] : SA3_CORNERS[3], 
                            SA3_CORNERS[0] : SA3_CORNERS[2]]

        # Określa początkowe oszacowanie prawdopodobieństwa 
        # znalezienia żeglarza dla każdego obszaru.
        self.p1 = 0.2
        self.p2 = 0.5
        self.p3 = 0.3

        # Inicjalizuje atrybuty przechowujące prawdopodobieństwo 
        # skuteczności poszukiwań.
        self.sep1 = 0
        self.sep2 = 0
        self.sep3 = 0

    def draw_map(self, last_known):
        """
        Wyświetla mapę regionu wraz ze skalą, ostatnią znaną
        lokalizacją oraz obszarami poszukiwań.
        """
        # Rysuje pasek skali.
        cv.line(self.img, (20, 370), (70, 370), (0, 0, 0), 2)
        cv.putText(self.img, '0', (8, 370), cv.FONT_HERSHEY_PLAIN, 1,
                   (0, 0, 0))
        cv.putText(self.img, '50 mil morskich', (71, 370),
                   cv.FONT_HERSHEY_PLAIN, 1, (0, 0, 0))

        # Rysuje i numeruje obszary poszukiwań.
        cv.rectangle(self.img, (SA1_CORNERS[0], SA1_CORNERS[1]),
                     (SA1_CORNERS[2], SA1_CORNERS[3]), (0, 0, 0), 1)
        cv.putText(self.img, '1',
                   (SA1_CORNERS[0] + 3, SA1_CORNERS[1] + 15),
                   cv.FONT_HERSHEY_PLAIN, 1, 0)
        cv.rectangle(self.img, (SA2_CORNERS[0], SA2_CORNERS[1]),
                     (SA2_CORNERS[2], SA2_CORNERS[3]), (0, 0, 0), 1)
        cv.putText(self.img, '2',
                   (SA2_CORNERS[0] + 3, SA2_CORNERS[1] + 15),
                   cv.FONT_HERSHEY_PLAIN, 1, 0)
        cv.rectangle(self.img, (SA3_CORNERS[0], SA3_CORNERS[1]),
                     (SA3_CORNERS[2], SA3_CORNERS[3]), (0, 0, 0), 1)
        cv.putText(self.img, '3',
                   (SA3_CORNERS[0] + 3, SA3_CORNERS[1] + 15),
                   cv.FONT_HERSHEY_PLAIN, 1, 0)

        # Zaznacza na mapie ostatnią znaną lokalizację zaginionego.
        cv.putText(self.img, '+', last_known,
                   cv.FONT_HERSHEY_PLAIN, 1, (0, 0, 255))
        cv.putText(self.img, '+ = ostatnia znana lokalizacja', (240, 355),
                   cv.FONT_HERSHEY_PLAIN, 1, (0, 0, 255))
        cv.putText(self.img, '* = rzeczywista lokalizacja', (242, 370),
                   cv.FONT_HERSHEY_PLAIN, 1, (255, 0, 0))

        cv.imshow('Obszary do przeszukania', self.img)
        cv.moveWindow('Obszary do przeszukania', 750, 10)
        cv.waitKey(500)

    def sailor_final_location(self, num_search_areas):
        """Zwraca współrzędne x i y rzeczywistej lokalizacji zaginionego."""
        # Znajduje współrzędne żeglarza względem podtablicy obszaru poszukiwań.
        self.sailor_actual[0] = np.random.choice(self.sa1.shape[1], 1)
        self.sailor_actual[1] = np.random.choice(self.sa1.shape[0], 1)

        # Losuje obszar poszukiwań.
        area = int(random.triangular(1, num_search_areas + 1))

        # Przekształca lokalne współrzędne obszaru poszukiwań 
        # na współrzędne mapy regionu.
        if area == 1:
            x = self.sailor_actual[0] + SA1_CORNERS[0]
            y = self.sailor_actual[1] + SA1_CORNERS[1]
            self.area_actual = 1
        elif area == 2:
            x = self.sailor_actual[0] + SA2_CORNERS[0]
            y = self.sailor_actual[1] + SA2_CORNERS[1]
            self.area_actual = 2
        elif area == 3:
            x = self.sailor_actual[0] + SA3_CORNERS[0]
            y = self.sailor_actual[1] + SA3_CORNERS[1]
            self.area_actual = 3
        return x, y

    def calc_search_effectiveness(self):
        """
        Wyznacza wartość dziesiętną reprezentującą skuteczność poszukiwań
        dla każdego obszaru poszukiwań.
        """
        self.sep1 = random.uniform(0.2, 0.9)
        self.sep2 = random.uniform(0.2, 0.9)
        self.sep3 = random.uniform(0.2, 0.9)

    def conduct_search(self, area_num, area_array, effectiveness_prob):
        """Zwraca wynik poszukiwań oraz listę przeszukanych współrzędnych."""
        local_y_range = range(area_array.shape[0])
        local_x_range = range(area_array.shape[1])
        coords = list(itertools.product(local_x_range, local_y_range))
        random.shuffle(coords)
        coords = coords[:int((len(coords) * effectiveness_prob))]
        loc_actual = (self.sailor_actual[0], self.sailor_actual[1])
        if area_num == self.area_actual and loc_actual in coords:
            return 'Znaleziono w obszarze nr {}.'.format(area_num), coords
        else:
            return 'Nie znaleziono.', coords

    def revise_target_probs(self):
        """
        Aktualizuje prawdopodobieństwo dla każdego obszaru
        na podstawie skuteczności poszukiwań.
        """
        denom = self.p1 * (1 - self.sep1) + self.p2 * (1 - self.sep2) \
                + self.p3 * (1 - self.sep3)
        self.p1 = self.p1 * (1 - self.sep1) / denom
        self.p2 = self.p2 * (1 - self.sep2) / denom
        self.p3 = self.p3 * (1 - self.sep3) / denom


def draw_menu(search_num):
    """Drukuje menu z wyborem obszaru do przeszukania."""
    print('\nPodejście nr {}'.format(search_num))
    print(
        """
        Wybierz następne obszary do przeszukania:

        0 - Wyjdź z programu
        1 - Przeszukaj dwukrotnie obszar pierwszy
        2 - Przeszukaj dwukrotnie obszar drugi
        3 - Przeszukaj dwukrotnie obszar trzeci
        4 - Przeszukaj obszary pierwszy i drugi
        5 - Przeszukaj obszary pierwszy i trzeci
        6 - Przeszukaj obszary drugi i trzeci
        7 - Zacznij od początku
        """
    )


def main():
    app = Search('Cape_Python')
    app.draw_map(last_known=(160, 290))
    sailor_x, sailor_y = app.sailor_final_location(num_search_areas=3)
    print("-" * 65)
    print("\nPoczątkowe oszacowanie prawdopodobieństwa (P):")
    print("P1 = {:.3f}, P2 = {:.3f}, P3 = {:.3f}"
          .format(app.p1, app.p2, app.p3))
    search_num = 1

    while True:
        app.calc_search_effectiveness()
        draw_menu(search_num)
        choice = input("Wybierz opcję: ")

        if choice == "0":
            sys.exit()

        elif choice == "1":
            results_1, coords_1 = app.conduct_search(1, app.sa1, app.sep1)
            results_2, coords_2 = app.conduct_search(1, app.sa1, app.sep1)
            app.sep1 = (len(set(coords_1 + coords_2))) / (len(app.sa1)**2)
            app.sep2 = 0
            app.sep3 = 0

        elif choice == "2":
            results_1, coords_1 = app.conduct_search(2, app.sa2, app.sep2)
            results_2, coords_2 = app.conduct_search(2, app.sa2, app.sep2)
            app.sep1 = 0
            app.sep2 = (len(set(coords_1 + coords_2))) / (len(app.sa2)**2)
            app.sep3 = 0

        elif choice == "3":
            results_1, coords_1 = app.conduct_search(3, app.sa3, app.sep3)
            results_2, coords_2 = app.conduct_search(3, app.sa3, app.sep3)
            app.sep1 = 0
            app.sep2 = 0
            app.sep3 = (len(set(coords_1 + coords_2))) / (len(app.sa3)**2)

        elif choice == "4":
            results_1, coords_1 = app.conduct_search(1, app.sa1, app.sep1)
            results_2, coords_2 = app.conduct_search(2, app.sa2, app.sep2)
            app.sep3 = 0

        elif choice == "5":
            results_1, coords_1 = app.conduct_search(1, app.sa1, app.sep1)
            results_2, coords_2 = app.conduct_search(3, app.sa3, app.sep3)
            app.sep2 = 0

        elif choice == "6":
            results_1, coords_1 = app.conduct_search(2, app.sa2, app.sep2)
            results_2, coords_2 = app.conduct_search(3, app.sa3, app.sep3)
            app.sep1 = 0

        elif choice == "7":
            main()

        else:
            print("\nTo nie jest poprawy wybór.", file=sys.stderr)
            continue

        # Używa twierdzenia Bayesa do zaktualizowania prawdopodobieństwa.
        app.revise_target_probs()

        print("\nPodejście nr {} - wynik 1: {}"
              .format(search_num, results_1), file=sys.stderr)
        print("Podejście nr {} - wynik 2: {}\n"
              .format(search_num, results_2), file=sys.stderr)
        print("Skuteczność poszukiwań (E) dla podejścia nr {}:"
              .format(search_num))
        print("E1 = {:.3f}, E2 = {:.3f}, E3 = {:.3f}"
              .format(app.sep1, app.sep2, app.sep3))

        # Drukuje zaktualizowaną wartość prawdopodobieństwa,
        # jeżeli nie znaleziono żeglarza. W przeciwnym przypadku
        # pokazuje pozycję.
        if results_1 == 'Nie znaleziono.' and results_2 == 'Nie znaleziono.':
            print("\nNowe oszacowanie prawdopodobieństwa (P) "
                  "dla podejścia nr {}:".format(search_num + 1))
            print("P1 = {:.3f}, P2 = {:.3f}, P3 = {:.3f}"
                  .format(app.p1, app.p2, app.p3))
        else:
            cv.circle(app.img, (sailor_x[0], sailor_y[0]), 3, (255, 0, 0), -1)
            cv.imshow('Obszary do przeszukania', app.img)
            cv.waitKey(1500)
            main()
        search_num += 1


if __name__ == '__main__':
    main()
