#!/usr/bin/env python3
import json
import surfex
from datetime import datetime
import argparse
import sys


def parse_args_set_geo_from_obs_set(argv):
    parser = argparse.ArgumentParser()
    parser.add_argument("-v", type=str, dest="variable", help="Variable name", required=True)
    parser.add_argument("-t", dest="validtime", help="Validtime (YYYYMMDDHH)", required=True)
    parser.add_argument("-i", type=str, dest="inputfile", help="Input file", required=False)
    parser.add_argument("-it", type=str, dest="inputtype", help="Input type", required=True,
                        choices=["netatmo", "frost", "bufr", "json"])
    parser.add_argument("--lonrange", type=str,  dest="lonrange", help="Longitude range", default=None, required=False)
    parser.add_argument("--latrange", type=str, dest="latrange", help="Latitude range", default=None, required=False)
    parser.add_argument("-o", type=str, dest="output", help="Output file", required=True)

    if len(argv) == 0:
        parser.print_help()
        sys.exit(1)

    return parser.parse_args(argv)


def set_geo_from_obs_set(args):

    obs_time = datetime.strptime(args.validtime, "%Y%m%d%H")
    varname = args.variable
    inputtype = args.inputtype
    inputfile = args.inputtype
    output = args.output
    lonrange = args.lonrange
    latrange = args.latrange

    settings = {
        "obs": {
            "varname": varname,
            "filetype": inputtype,
            "inputfile": inputfile,
            "filenames": [inputfile]
        }
    }
    if lonrange is None:
        lonrange = [-180, 180]
    if latrange is None:
        latrange = [-90, 90]

    print("Get data source")
    times, lons, lats, stids, elevs, values, varnames = surfex.get_datasources(obs_time, settings)[0].get_obs()

    selected_lons = []
    selected_lats = []
    for i in range(0, len(lons)):
        lon = lons[i]
        lat = lats[i]

        if lonrange[0] <= lon <= lonrange[1] and latrange[0] <= lat <= latrange[1]:
            lon = round(lon, 5)
            lat = round(lat, 5)
            # print(i, lon, lat)
            selected_lons.append(lon)
            selected_lats.append(lat)

    dx = ["0.3"] * len(selected_lons)
    geo_json = {
        "nam_pgd_grid": {
            "cgrid": "LONLATVAL"
        },
        "nam_lonlatval": {
            "xx": selected_lons,
            "xy": selected_lats,
            "xdx": dx,
            "xdy": dx
        }
    }
    # geo = surfex.LonLatVal(geo_json)
    json.dump(geo_json, open(output, "w"))


if __name__ == "__main__":
    my_args = parse_args_set_geo_from_obs_set(sys.argv[1:])
    set_geo_from_obs_set(my_args)
