import sys
import matplotlib
matplotlib.use('QtAgg')
import matplotlib.pyplot as plt
from matplotlib.figure import Figure
from matplotlib.axes import Axes
from shapely import Polygon
from shapely.plotting import plot_polygon
from typing import Tuple

class GeometrySetup:
    def __init__(self,current_setup:bool):
        self.current_setup = current_setup
        self.spawn_area = self.spawn()
        self.queue_area = self.queue()
        self.grass_area = self.grass()
        self.crosswalk_area = self.crosswalk()
        self.entry_area = self.entry_polygon()
        self.exit_area = self.exit_polygon()

    def spawn(self)->Polygon:
        spawn_area = Polygon([
            (25.0,16.823),
            (25.0,4.569),
            (26.163,4.569),
            (28.214,5.07),
            (28.214,17.761),
            (25.787,17.761),
            ])
        return spawn_area

    def queue(self)->Polygon:
        if self.current_setup is True:
            queue_area = Polygon([
                (18.767,9.395),
                (16.924,9.395),
                (15.214,9.395),    
                (15.214,0),
                (16.924,0),  
                (16.924,4.569),
                (25.0,4.569),
                (25.0,16.823)
                ])
        else:
            queue_area = Polygon([
                (19.531,10.306),
                (15.214,10.306),
                (13.88,10.306),
                (12.98,9.896),
                (12.98,0),
                (15.214,0),
                (16.924,0),
                (16.924,4.569),
                (25.0,4.569),
                (25.0,16.823)
                ])
        return queue_area

    def grass(self)->Polygon:
        if self.current_setup is True:
            grass_area = Polygon([
                (4,0),
                (0,0),
                (0,17.761),
                (4,17.761)
                ])
        else:
            grass_area = Polygon([
                (6.23,0),
                (0,0),
                (0,17.761),
                (4,17.761),
                (4,10.306),
                (5.33,10.306),
                (6.23,9.896)
                ])
        return grass_area

    def crosswalk(self)->Polygon:
        if self.current_setup is True:
            crosswalk_area = Polygon([
                (4,6.068),
                (4,7.896),
                (15.214,7.896),
                (15.214,6.068)
                ])
        else:
            crosswalk_area = Polygon([
                (6.23,4.982),
                (6.23,8.982),
                (12.98,8.982),
                (12.98,4.982)
                ])
        return crosswalk_area
    
    def entry_polygon(self)->Polygon:
        if self.current_setup is True:
            entry_area = Polygon([ #  x: 2.9m,  y: 3.428m
                (15.314,5.268),    # dx: 0.1m, dy: 0.8m
                (15.314,8.696),
                (18.214,8.696),
                (18.214,5.268)
                ])
        else:
            entry_area = Polygon([ # x: 2.9m, y: 5.6m
                (15.98,9.782),
                (15.98,4.182),
                (13.08,4.182),
                (13.08,9.782)
                ])
        return entry_area

    def exit_polygon(self)->Polygon:
        if self.current_setup is True:
            exit_area = Polygon([ # x: 2.9m, y: 3.428m
                (1,5.268),
                (1,8.696),
                (3.9,8.696),
                (3.9,5.268)
                ])
        else:
            exit_area = Polygon([ # x: 2.9m, y: 5.6m
                (3.23,4.182),
                (3.23,9.782),
                (6.13,9.782),
                (6.13,4.182)
                ])
        return exit_area
    
    def plot_all(self)->Tuple[Figure,Axes]:
        fig,ax = plt.subplots(figsize=(12,8))
        plot_polygon(
            self.spawn_area,
            color="lightblue",
            add_points=False,
            ax=ax,
            alpha=0.4
            )
        plot_polygon(
            self.queue_area,
            color="blue",
            add_points=False,
            ax=ax,
            alpha=0.4
            )
        plot_polygon(
            self.grass_area,
            color="blue",
            add_points=False,
            ax=ax,
            alpha=0.4
            )
        plot_polygon(
            self.crosswalk_area,
            color="red",
            add_points=False,
            ax=ax,
            alpha=0.4
            )
        plot_polygon(
            self.entry_area,
            color="green",
            add_points=False,
            ax=ax,
            alpha=0.4
            )
        plot_polygon(
            self.exit_area,
            color="green",
            add_points=False,
            ax=ax,
            alpha=0.4
            )
        plt.plot(
            *self.spawn_area.exterior.xy,
            color='black',
            alpha=0.7,
            linewidth=1
            )
        plt.plot(
            *self.queue_area.exterior.xy,
            color='black',
            alpha=0.7,
            linewidth=1
            )
        plt.plot(
            *self.grass_area.exterior.xy,
            color='black',
            alpha=0.7,
            linewidth=1
            )
        plt.plot(
            *self.crosswalk_area.exterior.xy,
            color='black',
            alpha=0.7,
            linewidth=1
            )
        plt.plot(
            *self.entry_area.exterior.xy,
            color='black',
            alpha=0.7,
            linewidth=1
            )
        plt.plot(
            *self.exit_area.exterior.xy,
            color='black',
            alpha=0.7,
            linewidth=1
            )
        ax.text(
            self.spawn_area.centroid.x, 
            self.spawn_area.centroid.y,
            'Spawn\nArea',
            ha='center',
            va='center',
            fontweight='bold'
            )
        ax.text(
            self.queue_area.centroid.x,
            self.queue_area.centroid.y,
            'Queue',
            ha='center',
            va='center',
            fontweight='bold'
            )
        ax.text(
            self.crosswalk_area.centroid.x,
            self.crosswalk_area.centroid.y,
            'Crosswalk',
            ha='center',
            va='center',
            fontweight='bold',
            )
        ax.text(
            self.entry_area.centroid.x,
            self.entry_area.centroid.y,
            'Crosswalk\nEntry',
            ha='center',
            va='center',
            fontsize=9,
            fontweight='bold',
            #color="white"
            )
        ax.text(
            self.exit_area.centroid.x,
            self.exit_area.centroid.y,
            'Crosswalk\nExit',
            ha='center',
            va='center',
            fontsize=9,
            fontweight='bold',
            #color="white"
            )
        ax.text(
            (self.grass_area.centroid.x-0.9 if not\
                self.current_setup else self.grass_area.centroid.x),
            self.grass_area.centroid.y+2,
            'Grass',
            ha='center',
            va='center',
            fontweight='bold'
            )
        ax.set_xlabel('X (m)')
        ax.set_ylabel('Y (m)')
        ax.set_title(
            f'{"Current" if self.current_setup else "Smart"} Crosswalk Design',
            fontweight='bold',
            fontsize=15
            )
        ax.set_aspect('equal')
        ax.grid(True, alpha=0.3)
        return fig,ax   

if __name__ == "__main__":
    from PyQt6 import QtWidgets
    app = QtWidgets.QApplication(sys.argv)
    fig,ax = GeometrySetup(True).plot_all()
    plt.show(block=False)
    sys.exit(app.exec())