# QuatOrtho.py
# Andrew Davison, ad@coe.psu.ac.th, August 2025
# Input: New York, Cape Town
'''
  Draw a Great Circle arc between two cities
  on an Orthographic map, and add Quaternion
  LERP and SLERP points along the arc

  The arc is drawn using Cartopy and using my 
  nav.greatCircleCoords().

  The Orthographic extent is adjusted to make 
  both cities visible on the map.

  Possible cities
         New York  Plymouth  Tokyo Berkeley
         Cape Town  Panama City  Seattle   Melbourne
         Reykjavik  Helsinki Lima  Nord  Sydney

  Optional command line argument:  s  or  l
  for SLERP-only or LERP-only markings
'''
import math, sys

import cartopy.crs as ccrs
import cartopy.feature as cf

import matplotlib.pyplot as plt
from matplotlib.lines import Line2D

import nav
from Quat import Quat


def plotMarkers(ax, coords, mark, col):
  for coord in coords:
    ax.plot(coord[1], coord[0], marker=mark, 
           color=col, alpha=0.5,
           transform=ccrs.PlateCarree())


def plotCity(ax, city, lat, lon):
  # Add marker
  ax.plot(lon, lat, 'o', color='green', 
         transform=ccrs.PlateCarree())
  # label city
  ax.text(lon, lat, " "+city, color='green', 
         transform=ccrs.PlateCarree())


interp = 'SL'   # show both markers
if len(sys.argv)  == 2:
  interp = sys.argv[1].upper()
print("Visible Markers:", interp)

# get coordinates for the cities
# startCity = "Sydney"  #  "New York"   
# endCity =  "Berkeley"  #  "Cape Town"
startCity = input("Start city? ").capitalize()
endCity = input("End City? ").capitalize()
cities = nav.loadCityData()
lat1, lon1 = nav.readCoord(cities, startCity) 
lat2, lon2 = nav.readCoord(cities, endCity)
print(f"{startCity} (lat,lon): ({lat1:.2f}, {lon1:.2f})")
print(f"{endCity} (lat,lon): ({lat2:.2f}, {lon2:.2f})")

qStart = Quat.fromLatLon(lat1, lon1)
qEnd = Quat.fromLatLon(lat2, lon2)
print("\nStart quat:", qStart)
print("End quat:", qEnd)

usesDL = nav.usesDateline(lon1, lon2)
if usesDL:
  print("Crosses the dateline")
print()

plt.figure(f"{startCity} to {endCity}", figsize=(8,6))
if usesDL:
  proj = ccrs.Orthographic(180, 30)  # on dateline
else:
  proj = ccrs.Orthographic((lon1+lon2)/2, 30)

proj.threshold = proj.threshold/20  # finer threshold
ax = plt.axes(projection=proj)

ax.set_global()  # show the entire globe
lineSep = 10
ax.gridlines(draw_labels=True, ls='--',
             xlocs=range(-180, 181, lineSep),
             ylocs=range(-90, 91, lineSep))
ax.coastlines()
ax.add_feature(cf.BORDERS, ls=':')


# add LERP and SLERP markers
dotHandles = []
markStep = 10

if 'L' in interp:
  quats = Quat.lerps(qStart, qEnd, markStep)
  coords = [ q.toLatLon() for q in quats]
  plotMarkers(ax, coords, "s", "red")
  dotHandles.append(Line2D([0], [0], marker='s', 
                 markerfacecolor='red', color='none', 
                 label='LERP'))

if 'S' in interp:
  quats = Quat.slerpsNInv(qStart, qEnd, markStep)
  coords = [ q.toLatLon() for q in quats]
  plotMarkers(ax, coords, "o", "green")
  print("SLERPNInv (lat,lon)s:")
  for coord in coords:
    cLat, cLon = coord
    print(f"  ({cLat:.2f}, {cLon:.2f})")
  dotHandles.append( Line2D([0], [0], marker='o', 
                markerfacecolor='green', color='none', 
                label='SLERP'))

# use Cartopy to draw the Great Circle
handle1, = plt.plot([lon1, lon2], [lat1, lat2],
     color='blue', lw=5, alpha=0.2,
     label="Great Circle", transform=ccrs.Geodetic())  
                 # spherical; great circle

# draw the Great Circle trip as a series of coordinates
GCStep = 20
coords = nav.greatCircleCoords(lat1, lon1, lat2, lon2, GCStep)
'''
print("\nGreat Circle (lat,lon)s and deg Bearing:")
for coord in coords:
  cLat, cLon, cBear = coord
  print(f"  ({cLat:.2f}, {cLon:.2f}); {cBear:.2f}")
'''
lats, lons, _ = zip(*coords)
if usesDL:
  lons = [ lon+360 if lon < 0 else lon for lon in lons]
handle2, = plt.plot(lons, lats, color='blue', linestyle='--',
         label="Nav Great Circle",
         transform=ccrs.PlateCarree())
                # straight lines link the coords

plotCity(ax, startCity, lat1, lon1)
plotCity(ax, endCity, lat2, lon2)

plt.title(f"{startCity} to {endCity}\n")
allHandles = [handle1, handle2] + dotHandles
plt.legend(handles=allHandles)
plt.show()