#!/usr/bin/python3
 
########### Grammar.py ##############
from lark import Lark

notam_grammar = r"""
start: area+

area: area_intro? clause+

area_intro: "AREA" ("DEFINED" "AS")?
          | "WI" "AREA"

clause: coord_clause
      | arc_clause
      | circle_clause
      | bearing_clause

coord_clause: COORD+

arc_clause: dir? "ARC" "OF" "A" "CIRCLE"
            "RADIUS" NUMBER "NM"
            "CENTRED" "ON" COORD
            ("TO" COORD)?

circle_clause: "RADIUS" NUMBER "NM"
               "CENTRED" "ON" COORD

bearing_clause: "BEARING" NUMBER "TO" NUMBER

dir: "CLOCKWISE" | "ANTICLOCKWISE"

COORD: /[0-9]{6,7}[NS]\s*[0-9]{7}[EW]/

NUMBER: /[0-9]+(\.[0-9]+)?/

%import common.WS
%ignore WS

// Ignore filler words globally (critical for real NOTAMs)
%ignore "THEN"
%ignore "THENCE"
%ignore "FROM"
"""
parser = Lark(notam_grammar, start="start", parser="lalr")


############### AST.py ##################
from dataclasses import dataclass
from typing import List, Optional, Tuple

@dataclass
class Coord:
    lon: float
    lat: float

@dataclass
class LineSeg:
    start: Coord
    end: Coord

@dataclass
class ArcSeg:
    start: Coord
    end: Coord
    center: Coord
    radius_nm: float
    clockwise: bool

@dataclass
class CircleSeg:
    center: Coord
    radius_nm: float

@dataclass
class BearingArcSeg:
    center: Coord
    radius_nm: float
    start_bearing: float
    end_bearing: float

@dataclass
class AreaAST:
    segments: List

@dataclass
class NotamAST:
    areas: List[AreaAST]






######################### Transformer.py ####################
from lark import Transformer
import re

def dms_to_dd(dms):
    m = re.match(r"(\d+)([NSEW])", dms)
    val, hemi = m.groups()

    deg = int(val[:-4])
    minutes = int(val[-4:-2])
    seconds = int(val[-2:])

    dd = deg + minutes/60 + seconds/3600
    if hemi in ("S", "W"):
        dd *= -1
    return dd


def parse_coord_token(token):
    lat, lon = token.split()
    return Coord(dms_to_dd(lon), dms_to_dd(lat))


class NotamTransformer(Transformer):

    def COORD(self, token):
        return parse_coord_token(str(token))

    def NUMBER(self, token):
        return float(token)

    def coord_clause(self, items):
        return ("coords", items)

    def arc_clause(self, items):
        clockwise = True
        idx = 0

        if isinstance(items[0], str):
            clockwise = items[0] == "CLOCKWISE"
            idx += 1

        radius = items[idx]; idx += 1
        center = items[idx]; idx += 1

        end = items[idx] if idx < len(items) else None

        return ("arc", {
            "radius": radius,
            "center": center,
            "end": end,
            "clockwise": clockwise
        })

    def circle_clause(self, items):
        return ("circle", {
            "radius": items[0],
            "center": items[1]
        })

    def bearing_clause(self, items):
        return ("bearing", {
            "start": items[0],
            "end": items[1]
        })

    def area(self, items):
        segments = []
        coords = []

        for item in items:
            t, v = item

            if t == "coords":
                coords.extend(v)

            elif t == "arc":
                start = coords[-1]
                end = v["end"] or start

                segments.append(ArcSeg(
                    start=start,
                    end=end,
                    center=v["center"],
                    radius_nm=v["radius"],
                    clockwise=v["clockwise"]
                ))

                coords.append(end)

            elif t == "circle":
                segments.append(CircleSeg(
                    center=v["center"],
                    radius_nm=v["radius"]
                ))

            elif t == "bearing":
                segments.append(BearingArcSeg(
                    center=coords[-1],
                    radius_nm=5,
                    start_bearing=v["start"],
                    end_bearing=v["end"]
                ))

        # implicit lines
        for i in range(len(coords)-1):
            segments.append(LineSeg(coords[i], coords[i+1]))

        return AreaAST(segments)

    def start(self, items):
        return NotamAST(items)







###################### GeometryBuilder #########################
from shapely.geometry import Polygon, MultiPolygon
from pyproj import Transformer
import numpy as np
import math

NM_TO_M = 1852

def get_aeqd(lat, lon):
    proj = f"+proj=aeqd +lat_0={lat} +lon_0={lon}"
    fwd = Transformer.from_crs("EPSG:4326", proj, always_xy=True).transform
    inv = Transformer.from_crs(proj, "EPSG:4326", always_xy=True).transform
    return fwd, inv


def arc_points(center, start, end, radius_nm, clockwise):
    fwd, inv = get_aeqd(center.lat, center.lon)

    cx, cy = fwd(center.lon, center.lat)
    sx, sy = fwd(start.lon, start.lat)
    ex, ey = fwd(end.lon, end.lat)

    a1 = math.atan2(sy-cy, sx-cx)
    a2 = math.atan2(ey-cy, ex-cx)

    if clockwise:
        if a2 > a1: a2 -= 2*math.pi
    else:
        if a2 < a1: a2 += 2*math.pi

    r = radius_nm * NM_TO_M
    angles = np.linspace(a1, a2, 64)

    pts = []
    for a in angles:
        x = cx + r*math.cos(a)
        y = cy + r*math.sin(a)
        pts.append(inv(x, y))
    return pts


def circle_points(center, radius_nm):
    fwd, inv = get_aeqd(center.lat, center.lon)
    cx, cy = fwd(center.lon, center.lat)

    r = radius_nm * NM_TO_M
    angles = np.linspace(0, 2*math.pi, 128)

    return [inv(cx + r*math.cos(a), cy + r*math.sin(a)) for a in angles]


def build_geometry(ast: NotamAST):
    polys = []

    for area in ast.areas:
        pts = []

        for seg in area.segments:
            if isinstance(seg, LineSeg):
                pts.append((seg.start.lon, seg.start.lat))
                pts.append((seg.end.lon, seg.end.lat))

            elif isinstance(seg, ArcSeg):
                pts.extend(arc_points(
                    seg.center, seg.start, seg.end,
                    seg.radius_nm, seg.clockwise
                ))

            elif isinstance(seg, CircleSeg):
                pts.extend(circle_points(seg.center, seg.radius_nm))

        if pts and pts[0] != pts[-1]:
            pts.append(pts[0])

        if pts:
            polys.append(Polygon(pts))

    return polys[0] if len(polys) == 1 else MultiPolygon(polys)




######################## Usage.py #################
#text = """
#AREA DEFINED AS
#512000N 0001000W THEN 513000N 0002000W
#THEN CLOCKWISE ARC OF A CIRCLE RADIUS 10 NM CENTRED ON 512500N 0001500W TO 512000N 0001000W

text = """
AREA DEFINED AS
RADIUS 5 NM CENTRED ON 512000N 0000000W
"""

tree = parser.parse(text)
ast = NotamTransformer().transform(tree)

geom = build_geometry(ast)

print(type(geom), geom.is_valid)