import sys
from pathlib import Path

import netCDF4
import ruamel.yaml as yaml
from docopt import docopt

__all__ = ["main", "build"]

__doc__ = """
Generate ncml based on a yaml file.

Usage:
    yaml2ncml INFILE [--output=OUTFILE]

    yaml2ncml (-h | --help | --version)

Examples:
    yaml2ncml roms.yaml
    yaml2ncml roms.yaml --output=roms.ncml

Arguments:
  file      yaml file.

Options:
  -h --help     Show this screen.
  -v --version     Show version.
"""


def str_att(name, value):
    if isinstance(value, list):
        try:
            value = ",".join(value)
        except TypeError as e:
            raise ValueError(f"Expected `str` got {value}") from e
    msg = '  <attribute name="{:s}" type="String" value="{:s}"/>\n'
    return msg.format(name, value)


def header():
    text = '<?xml version="1.0" encoding="UTF-8"?>\n<netcdf xmlns='
    text += '"http://www.unidata.ucar.edu/namespaces/netcdf/ncml-2.2">\n'
    text += str_att("Conventions", "CF-1.6, SGRID-0.1, ACDD-1.3")
    text += str_att("cdm_data_type", "Grid")
    return text


def footer(text):
    text += "</netcdf>\n"
    return text


def add_global_atts(text, a):
    d = a["dataset"]
    for key, value in d.items():
        # Handle simple attribute pairs first.
        if key in [
            "id",
            "license",
            "summary",
            "title",
            "project",
            "naming_authority",
            "references",
            "acknowledgment",
            "keywords_vocabulary",
        ]:
            text += str_att(key, value)
        elif key in ["creator", "publisher"]:
            email = value.get("email", None)
            if email:
                text += str_att("_".join([key, "email"]), email)
            url = value.get("url", None)
            if url:
                text += str_att("_".join([key, "url"]), url)
            name = value.get("name", None)
            if name:
                text += str_att("_".join([key, "name"]), name)
            institution = value.get("institution", None)
            if institution:
                text += str_att("_".join([key, "institution"]), institution)
        elif key in ["contributor"]:
            role = value.get("role", None)
            text += str_att("_".join([key, "role"]), role)
            email = value.get("email", None)
            if email:
                text += str_att("_".join([key, "email"]), email)
            url = value.get("url", None)
            if url:
                text += str_att("_".join([key, "url"]), url)
            name = value.get("name", None)
            if name:
                text += str_att("_".join([key, "name"]), name)
            institution = value.get("institution", None)
            if institution:
                text += str_att("_".join([key, "institution"]), institution)
        elif key in ["date"]:
            created = value.get("created", None)
            if created:
                text += str_att("_".join([key, "creation"]), created)
            modified = value.get("modified", None)
            if modified:
                text += str_att("_".join([key, "revision"]), modified)
            issued = value.get("issued", None)
            if issued:
                text += str_att("_".join([key, "publication"]), issued)
    return text


def add_bed_coord(text, a):
    ncfile = Path(a["aggregation"]["dir"], a["aggregation"]["sample_file"])
    nc = netCDF4.Dataset(ncfile)
    bed_coord_var = """<variable name="Nbed" shape="Nbed" type="double">
      <attribute name="long_name" value="pseudo coordinate at seabed points"/>
      <attribute name="standard_name" value="ocean_sigma_coordinate"/>
      <attribute name="positive" value="up"/>
      <attribute name="formula_terms" value="sigma: Nbed eta: zeta depth: h"/>
      <values start="-1.0" increment="-0.01"/>
    </variable>\n """
    if "Nbed" in nc.dimensions.keys():
        text += bed_coord_var
    return text


def add_var_atts(text, a):
    ncfile = Path(a["aggregation"]["dir"], a["aggregation"]["sample_file"])
    nc = netCDF4.Dataset(ncfile)
    ncv = nc.variables

    # Get a list of all variables more than 1D.
    vars = [var for var, vart in ncv.items() if vart.ndim > 1]

    #   identify all the rho, u and v vars

    rho_vars = [
        var
        for var in vars
        if "eta_rho" in ncv[var].dimensions and "xi_rho" in ncv[var].dimensions
    ]
    u_vars = [
        var
        for var in vars
        if "eta_u" in ncv[var].dimensions and "xi_u" in ncv[var].dimensions
    ]
    v_vars = [
        var
        for var in vars
        if "eta_v" in ncv[var].dimensions and "xi_v" in ncv[var].dimensions
    ]

    vars_all = set(vars)
    vars_include = set(a["variables"]["include"])
    vars_exclude = set(a["variables"]["exclude"])

    #   include/exclude only variables that actually occur in variable list

    vars_include = vars_all.intersection(vars_include)
    vars_exclude = vars_all.intersection(vars_exclude)

    #   If there are variables excluded, exclude them and keep all rest.
    #   If no variables are excluded, take just the included variables
    #   If no variables are included or excluded, take all variables (leave
    #     list of variables unchanged)

    if vars_exclude:
        vars_display = vars_all - vars_all.intersection(vars_exclude)
    else:
        if vars_include:
            vars_display = vars_all.intersection(vars_include)
        else:
            vars_display = vars_all

    #   remove some variables we never want (if they exist)
    Tobc = {"Tobc_in", "Tobc_out"}
    vars_display = vars_display - vars_display.intersection(Tobc)

    vars_display = list(vars_display)

    # add the variable attributes: S-grid stuff, display=T|F, ...
    for var in vars:
        text += f'<variable name="{var:s}">\n'
        try:
            text += str_att("standard_name", cf[var])
        except Exception:
            pass
        text += str_att("grid", "grid")

        if "Nbed" in ncv[var].dimensions:
            text += str_att("coordinates", ncv[var].coordinates + " Nbed")

        if var in vars_display:
            text += str_att("display", "True")
        else:
            text += str_att("display", "False")

        text += str_att("coverage_content_type", "modelResult")
        if var in rho_vars:
            text += str_att("location", "face")
        elif var in u_vars:
            text += str_att("location", "edge1")
        elif var in v_vars:
            text += str_att("location", "edge2")
        text += "</variable>\n\n"

    # write standard_name for time coordinate variable
    var = "ocean_time"
    if var in ncv.keys():
        try:
            text += f'\n<variable name="{var:s}">\n'
            text += str_att("standard_name", cf[var])
            text += "</variable>\n\n"
        except Exception:
            pass

    nc.close()
    return text


def write_grid_var(text):
    grid_var = """<variable name="grid" type="int">
        <attribute name="cf_role" value="grid_topology"/>
        <attribute name="topology_dimension" type="int" value="2"/>
        <attribute name="node_dimensions" value="xi_psi eta_psi"/>
        <attribute name="face_dimensions"
            value="xi_rho: xi_psi (padding: both) eta_rho: eta_psi (padding: both)"/>
        <attribute name="edge1_dimensions" value="xi_u: xi_psi eta_u: eta_psi (padding: both)"/>
        <attribute name="edge2_dimensions" value="xi_v: xi_psi (padding: both) eta_v: eta_psi"/>
        <attribute name="node_coordinates" value="lon_psi lat_psi"/>
        <attribute name="face_coordinates" value="lon_rho lat_rho"/>
        <attribute name="edge1_coordinates" value="lon_u lat_u"/>
        <attribute name="edge2_coordinates" value="lon_v lat_v"/>
        <attribute name="vertical_dimensions" value="s_rho: s_w (padding: none)"/>
    </variable>\n"""  # noqa
    text += grid_var
    return text


def add_aggregation_scan(text, a):
    agg = a["aggregation"]
    text += '<aggregation dimName="{:s}" type="joinExisting">\n'.format(agg["time_var"])
    text += '<scan location="{:s}" regExp="{:s}" subdirs="false"/>\n</aggregation>\n'.format(
        agg["dir"],
        agg["pattern"],
    )  # noqa
    return text


# Map ROMS variables to CF standard_names.
cf = {
    "ocean_time": "time",
    "zeta": "sea_surface_height_above_geopotential_datum",
    "temp": "sea_water_potential_temperature",
    "salt": "sea_water_salinity",
    "u": "sea_water_x_velocity",
    "v": "sea_water_y_velocity",
    "ubar": "barotropic_sea_water_x_velocity",
    "vbar": "barotropic_sea_water_y_velocity",
    "Hwave": "sea_surface_wave_significant_height",
    "bed_thickness": "sediment_bed_thickness",
}  # sediment_bed_thickness not in CF standard_names


def build(yml):
    text = header()
    text = add_global_atts(text, yml)
    text = add_var_atts(text, yml)
    text = write_grid_var(text)
    text = add_bed_coord(text, yml)
    text = add_aggregation_scan(text, yml)
    text = footer(text)
    return text


def main():
    args = docopt(__doc__, version="0.6.0")
    fname = args.get("INFILE")
    fout = args.get("--output", None)

    with open(fname) as stream:
        yml = yaml.load(stream, Loader=yaml.RoundTripLoader)

    text = build(yml)

    if fout:
        with open(fout, "w") as f:
            f.write(f"{text:s}")
    else:
        sys.stdout.write(text)
