#!/usr/bin/env python3
# Copyleft Rostislav Kouznetsov 11.2023- 

descr="""
    Approximates lcc grid from a given file wth SILAM rll grid
    optionally, calculates grid indices of given lon,lat pair
    and indices of areea of interest
"""


import netCDF4 as nc
import numpy as np
from pyproj import CRS, Transformer
import argparse
import sys

infname="grid_template_2500.nc"
infname="grid_template_0750.nc"


parser = argparse.ArgumentParser(
                    prog=sys.argv[0],
                    description=descr,
                    epilog='')
parser.add_argument('filename')
parser.add_argument('-l', '--location', default=None, help="lon,lat of the point of interest")
parser.add_argument('-s', '--size', default=None, help="nx,ny around the point of interest")


args = parser.parse_args(sys.argv[1:])
infname = args.filename




#infname = "grid_template_shmu_2000.nc"

with nc.Dataset(infname) as inf:
    projvar=inf.variables["Lambert_Conformal"]
    xvals = inf.variables["x"][:]
    yvals = inf.variables["y"][:]
    nx = len(xvals)
    ny = len(yvals)

    try:
        proj4strIn = projvar.proj4
    except AttributeError:
        proj4strIn = "+proj=lcc +a=%(a)f +b=%(b)f +lat_1=%(lat_1)f +lat_2=%(lat_2)f +lon_0=%(lon_0)f +lat_0=%(lat_0)f +k_0=1 +x_0=%(x_0)f +y_0=%(y_0)f +datum=WGS84 +units=m "%dict(
                a = projvar.earth_radius,
                b = projvar.earth_radius,
                lat_1 = projvar.latitude_of_projection_origin,
                lat_2 = projvar.latitude_of_projection_origin,
                lon_0 = projvar.longitude_of_central_meridian,
                lat_0 = projvar.latitude_of_projection_origin,
                x_0 = projvar.false_easting,
                y_0 = projvar.false_northing,
                )
        #proj4strIn = "+proj=lcc +a=6371229 +b=6371229 +lat_1=51.970001 +lat_2=51.970001 +lon_0=0.0 +lat_0=51.97 +k_0=1 +x_0=-324328.303458971 +y_0=17464.56278250541 +datum=WGS84 +units=m " ;

    crsIn=CRS.from_proj4(proj4strIn)

    spolelon = projvar.longitude_of_central_meridian
    spolelat = projvar.standard_parallel - 90.
    proj4rll = "+proj=ob_tran +o_proj=longlat +o_lat_p=%10f +o_lon_p=%10f +lon_0=%10f"%(-spolelat, spolelon,spolelon)
    crsRll = CRS.from_proj4(proj4rll)

    transformer = Transformer.from_crs(crsIn, crsRll)

    xx,yy = np.meshgrid(xvals,yvals)

    lons2d,lats2d = transformer.transform(xx,yy)
    lons = np.mean(lons2d, axis=0)
    lats = np.mean(lats2d, axis=1)

    xstep = (lons[-1] - lons[0]) / (len(lons) - 1)
    ystep = (lats[-1] - lats[0]) / (len(lats) - 1)
    ## Restore
    lo2d,la2d = np.meshgrid(lons,lats)

    dist = np.sqrt((lons2d -lo2d)**2 + (lats2d -la2d)**2)
    maxdist = np.amax(dist)

    lons -= spolelon

    grid ="""
   grid_method = CUSTOM_GRID     # METEO_GRID / EMISSION_GRID / OUTPUT_GRID / CUSTOM_GRID
   grid_type = lon_lat
   grid_title = %(name) s
   # proj string="%(proj)s"
   # maxdist/xstep = %(relerrx)f, maxdist/ystep = %(relerry)f
   resol_flag = 128
   ifReduced = 0
   earth_flag = 0
   wind_component = 0
   reduced_nbr_str = 0
   lon_start = %(x0)f
   lat_start = %(y0)f
   dx = %(dx)f
   dy = %(dy)f
   nx = %(nx)d
   ny = %(ny)d
   # lon_end = %(x1)f
   # lat_end = %(y1)f

    #lon0 to 0th meridian
    #lat0 to equator
   lat_s_pole = %(splat)f
   lon_s_pole = %(splon)f
   lat_pole_stretch = 0.
   lon_pole_stretch = 0.
"""%dict(name=infname, x0=lons[0], y0=lats[0], dx=xstep, dy=ystep, nx=nx, ny=ny, 
            x1 = lons[-1], y1 = lats[-1], proj=proj4rll,
            splat=spolelat, splon=spolelon, relerrx = maxdist/xstep, relerry = maxdist/ystep)
    print(grid)

proj4wgs84 = "+proj=longlat +datum=WGS84 +no_defs +type=crs"
crsll = CRS.from_proj4(proj4wgs84)
transformer = Transformer.from_crs(crsIn, crsll)
lons2d,lats2d = transformer.transform(xx,yy)


transformer_to = Transformer.from_crs(crsll,crsIn)

if args.location is  None:
    ix = len(xvals)//2 
    iy = len(yvals)//2

else:    
    lonc,latc = [ float(v) for v in args.location.split(',') ]
    xc,yc = transformer_to.transform(lonc,latc)
    ix = np.argmin((xvals-xc)**2)
    iy = np.argmin((yvals-yc)**2)

if args.size is None:
    xrang="0,%d"%(len(xvals)-1)
    yrang="0,%d"%(len(yvals)-1)
else:
    if args.location is  None:
        raise ValueError("size specified but not location. Too bad...")
    nxo,nyo = [ int(v) for v in args.size.split(',') ]
    xrang="%d,%d"%(ix - nxo//2, ix + nxo//2 ) ## nco indexing _includes_ last element
    yrang="%d,%d"%(iy - nyo//2, iy + nyo//2 )

print("#minloc for location(%s), ix=%d, iy=%d: lon=%f, lat=%f\n"%(args.location,ix,iy, lons2d[iy,ix],lats2d[iy,ix]))
print("#DEODExrange = "+xrang)
print("#DEODEyrange = "+yrang)
#    print(lons[ix-2:])






WGS84lonrange="#WGS84lonrange = %.1f,%.1f"%(np.amin(lons2d), np.amax(lons2d))
WGS84latrange="#WGS84latrange = %.1f,%.1f"%(np.amin(lats2d), np.amax(lats2d))
print(WGS84lonrange)
print(WGS84latrange)