#!/usr/bin/env python3
"""
Visualize numeric variables from NetCDF files.

Supported plotting cases include:
    - scalar output
    - 1D time series
    - 1D vertical profiles (value on X, depth on Y)
    - 2D latitude/longitude maps, drawn in a real projection
    - 2D cross-sections and Hovmoller diagrams
    - optional polar stereographic views for 2D maps
    - optional 3D globe views for 2D maps

How a plot is chosen:
    Each dimension is given a role - X, Y, Z, T, unstructured, or none - from the
    file's own metadata (axis, standard_name, positive, units), falling back on
    dimension names only for files that carry no usable metadata. The roles left
    after slicing decide the plot type. Run with --explain to see the reasoning.

Usage:
    1) Command-line mode:
             python display_netcdf.py /path/to/file.nc --variable VAR_NAME [options]

         Main options:
             -v, --variable      Name of the variable to visualize.
             -c, --cmap          Matplotlib colormap, or "auto" (default).
             -o, --output        Save figure to file instead of displaying.
             -e, --extra-indices JSON map selecting an index per dimension.
                                 Indices are 0-based for every dimension.
                                 A value of "avg" averages over that dimension.
                                 Example: '{"Time": 0, "nslope": 1}'
             -x, --x-dim         Dimension to put on the X axis for 2D plots.
             --plot-kind         Force a renderer instead of inferring one.
             --no-remap          Keep the raw physics grid instead of building a map.
             --explain           Report how the file was read and the plot chosen.
             --list-vars         List the variables in the file and exit.
             --no-interpolate    Draw the file's own cells as flat quads instead
                                 of interpolating between their centres.
             --no-show-topo      Do not overlay topography on lat/lon maps.
             --show-polar        Also display polar stereographic map view(s).
             --show-3d           Also display 3D globe view (requires vedo).

         Dimensions left out of --extra-indices are kept in full; the plot type
         follows from what remains.

    2) Interactive mode:
             python display_netcdf.py
         The script prompts for file, variable, and plotting options.

Environment/setup notes and dependency installation are documented in
toolbox/README.md.
"""

import os
import sys
import readline
import argparse

import matplotlib.pyplot as plt

from .conventions import normalize
from .coords import resolve_variable_axes
from .io import open_source
from .interactive import complete_filename, make_varname_completer
from .pipeline import plot_variable, plot_vector_field
from .render import GlobeOptions
from .selection import (autoselect_singletons, get_dimension_indices,
                        parse_extra_indices)

PLOT_KINDS = ('auto', 'scalar', 'line', 'timeseries', 'profile', 'section', 'geomap', 'globe')


def _explain(src, report, varname=None):
    """
    Report what normalize() concluded and, for one variable, how each dimension
    was resolved. This is the window into the metadata layer.
    """
    print(f"\nFile     : {src.path}")
    print(f"Families : {', '.join(report.families)}")
    if report.cell_area:
        print(f"Cell area: {report.cell_area}")
    if report.unstructured:
        print(f"Grid     : unstructured on '{report.unstructured['dim']}'")
    if report.fixes:
        print("Normalization:")
        for fix in report.fixes:
            print(f"  + {fix}")
    if report.warnings:
        print("Guesses (verify these):")
        for warning in report.warnings:
            print(f"  ! {warning}")
    if varname and varname in src:
        print(f"Axes of '{varname}':")
        for axis in resolve_variable_axes(src.ds, varname, report.unstructured):
            coord = axis.coord or '(no coordinate)'
            print(f"  {axis.dim:<20} role={axis.role}  size={axis.size:<6} "
                  f"from={axis.source} ({axis.confidence})  coord={coord}"
                  + (f"  units={axis.units}" if axis.units else ''))
    print()


def _open_normalized(path):
    """
    Open a file and bring its metadata into one vocabulary.
    """
    src = open_source(path)
    _, report = normalize(src.ds)
    return src, report


def globe_options_from(args):
    """
    Collect the 3D-globe flags, or None when there are no args to read.

    Built unconditionally rather than only under --show-3d, because the
    interactive prompt can ask for the globe after the map is drawn.
    """
    if args is None:
        return None

    sun = None
    if getattr(args, 'sun', None):
        parts = [p.strip() for p in args.sun.split(',')]
        try:
            if len(parts) != 2:
                raise ValueError("expected LON,LAT")
            sun = (float(parts[0]), float(parts[1]))
        except ValueError as err:
            print(f"Warning: ignoring --sun '{args.sun}' ({err}).")

    return GlobeOptions(
        relief_mode     = getattr(args, 'globe_relief', 'topo') or 'topo',
        sun             = sun,
        sun_ls          = getattr(args, 'sun_ls', None),
        export          = getattr(args, 'globe_export', None),
        lighting        = getattr(args, 'light', None),
        surface_mode    = getattr(args, 'globe_surface', None),
        shell_mode      = getattr(args, 'shell_mode', 'cloud'),
        shell_threshold = getattr(args, 'shell_threshold', None),
        shell_color     = getattr(args, 'shell_color', 'value'),
        shell_level     = getattr(args, 'shell_level', None),
        shell_opacity   = getattr(args, 'shell_opacity', 0.35),
        shell_top_km    = getattr(args, 'shell_top', None),
        shell_resolution= getattr(args, 'shell_resolution', None),
        shell_cut       = getattr(args, 'shell_cut', None),
        shell_exaggeration = getattr(args, 'shell_exaggeration', None),
        fps             = getattr(args, 'fps', 12),
        spin            = getattr(args, 'spin', None),
    )


def _plot_options(args):
    """
    Every presentation option `plot_variable` takes, read off the parsed
    arguments.

    Shared by the two modes because they had already drifted apart: --norm,
    --vmin, --vmax, --x-dim, --show-polar, --show-3d, --anomaly, --reduce,
    --stats and --stats-only reached the command line and not the interactive
    loop, which the README said they did. What genuinely differs between the two
    is where the figure goes, which slice it shows, and whether it may prompt -
    and that is now all the call sites spell out.

    Tolerates `args is None`, as `globe_options_from` does, so the interactive
    mode still works when it was reached with nothing parsed at all.
    """
    return dict(
        colormap      = getattr(args, 'cmap', None) or 'auto',
        norm          = getattr(args, 'norm', None),
        vmin          = getattr(args, 'vmin', None),
        vmax          = getattr(args, 'vmax', None),
        x_dim         = getattr(args, 'x_dim', None),
        show_topo     = getattr(args, 'show_topo', True),
        show_polar    = getattr(args, 'show_polar', False),
        show_3d       = getattr(args, 'show_3d', False),
        plot_kind     = getattr(args, 'plot_kind', None),
        remap         = not getattr(args, 'no_remap', False),
        explain       = getattr(args, 'explain', False),
        anomaly       = getattr(args, 'anomaly', None),
        reduce        = getattr(args, 'reduce', None),
        stats         = getattr(args, 'stats', False),
        stats_only    = getattr(args, 'stats_only', False),
        globe_options = globe_options_from(args),
        animate       = getattr(args, 'animate', None),
        overlay       = getattr(args, 'overlay', None),
        figsize       = getattr(args, 'figsize', None),
        dpi           = getattr(args, 'dpi', None),
        title_override= getattr(args, 'title_override', None),
        overlay_style = getattr(args, 'overlay_style', 'blend'),
        overlay_threshold = getattr(args, 'overlay_threshold', None),
        overlay_scale = getattr(args, 'overlay_scale', 'own'),
        diff          = getattr(args, 'diff', None),
        interpolate   = getattr(args, 'interpolate', True),
    )


def visualize_variable_interactive(nc_path=None, args=None):
    """
    Interactive loop: keep prompting for variables to plot until user quits.
    Returns a process exit status.
    """
    if not sys.stdin.isatty():
        print("Error: interactive mode needs a terminal. "
              "Pass a file and --variable to run non-interactively.")
        return 2

    if nc_path:
        path = nc_path
    else:
        readline.set_completer(complete_filename)
        readline.parse_and_bind("tab: complete")
        path = input("Enter path to NetCDF file: ").strip()

    if not os.path.isfile(path):
        print(f"Error: '{path}' not found.")
        return 1

    try:
        src, report = _open_normalized(path)
    except Exception as err:
        print(f"Error: cannot read '{path}': {err}")
        return 1

    try:
        var_list = src.variables
        if not var_list:
            print("No variables found in file.")
            return 1

        if args is not None and args.explain:
            _explain(src, report)

        plt.ion()

        # Every figure stays on screen until it is closed, so two variables can
        # be looked at side by side - which is most of what this loop is for.
        # matplotlib's warning at twenty open figures is about a script leaking
        # them; here they are being kept on purpose, and `close` clears the
        # screen when there are too many.
        plt.rcParams['figure.max_open_warning'] = 0

        while True:
            readline.set_completer(make_varname_completer(list(var_list) + ['close']))
            readline.parse_and_bind("tab: complete")

            print("\nAvailable variables:")
            for name in var_list:
                print(f"  > {name}")
            try:
                varname = input("\nEnter variable name to plot "
                                "('close' to close all figures, Enter to quit): ").strip()
            except EOFError:
                print("\nExiting.")
                break
            if varname == "":
                print("Exiting.")
                break
            # Guarded on the file rather than on the word, so a variable that is
            # really called 'close' is still plotted rather than shadowed.
            if varname.lower() in ('close', 'clear') and varname not in src:
                plt.close('all')
                continue
            if varname not in src:
                print(f"Variable '{varname}' not found. Try again.")
                continue

            dims, shape = src.dims_and_shape(varname)
            print(f"\nVariable '{varname}' has dimensions:")
            for dim, size in zip(dims, shape):
                print(f"  - {dim} (size {size})")
            print()

            if args is not None and args.explain:
                _explain(src, report, varname)

            selection = get_dimension_indices(src, varname)

            # Style options come from the command line even in interactive mode,
            # so `-c RdBu_r` is honoured here instead of being hardcoded to jet.
            # `-o` deliberately stays out of the shared options: one path cannot
            # serve a loop over however many variables the user asks for.
            plot_variable(
                src, report, varname,
                output_path   = None,
                extra_indices = selection,
                interactive   = True,
                **_plot_options(args),
            )
        return 0
    finally:
        src.close()


def visualize_variable_cli(args):
    """
    Command-line mode: visualize one variable directly, without prompting.
    Returns a process exit status.
    """
    if not os.path.isfile(args.nc_file):
        print(f"Error: '{args.nc_file}' not found.")
        return 1

    try:
        src, report = _open_normalized(args.nc_file)
    except Exception as err:
        print(f"Error: cannot read '{args.nc_file}': {err}")
        return 1

    try:
        if args.list_vars:
            for name in src.variables:
                dims, shape = src.dims_and_shape(name)
                sizes = ', '.join(f"{d}={n}" for d, n in zip(dims, shape))
                print(f"{name}\t{sizes}")
            return 0

        if args.vector:
            parts = [p.strip() for p in args.vector.split(',') if p.strip()]
            if len(parts) != 2:
                print("Error: --vector expects exactly two variable names, e.g. --vector u,v")
                return 2
            u_name, v_name = parts
            if args.overlay:
                # Not a stack: a translucent wash under a quiver hides the
                # arrows it is supposed to explain, and --vector-background
                # already puts one scalar under them.
                print("Note: --overlay is ignored with --vector; use "
                      "--vector-background VAR for a second field.")
            dims, shape = src.dims_and_shape(u_name) if u_name in src else ([], [])
            try:
                extra = parse_extra_indices(args.extra_indices, dims, shape)
            except ValueError as err:
                print(f"Error: {err}")
                return 2
            return plot_vector_field(
                src, report, u_name, v_name,
                colormap    = args.cmap,
                output_path = args.output,
                extra_indices = extra,
                show_topo   = args.show_topo,
                style       = args.vector_style,
                density     = args.vector_density,
                background  = args.vector_background,
                norm        = args.norm,
                vmin        = args.vmin,
                vmax        = args.vmax,
                show_polar  = args.show_polar,
                show_3d     = args.show_3d,
                interpolate = args.interpolate,
                globe_options = globe_options_from(args),
            )

        if not args.variable:
            print("Error: --variable is required (or use --list-vars).")
            return 2

        if args.variable not in src:
            print(f"Error: variable '{args.variable}' not in file.")
            return 1

        dims, shape = src.dims_and_shape(args.variable)

        print(f"\nVariable '{args.variable}' has {len(dims)} dimensions:")
        for name, size in zip(dims, shape):
            print(f"  - {name}: size {size}")
        print()

        if args.explain:
            _explain(src, report, args.variable)

        try:
            extra = parse_extra_indices(args.extra_indices, dims, shape)
        except ValueError as err:
            print(f"Error: {err}")
            return 2

        # Singleton dimensions are selected automatically, as in interactive mode
        autoselect_singletons(extra, dims, shape)

        return plot_variable(
            src, report, args.variable,
            output_path   = args.output,
            extra_indices = extra,
            interactive   = False,
            **_plot_options(args),
        )
    finally:
        src.close()


def build_parser():
    parser = argparse.ArgumentParser(
        description=__doc__,
        formatter_class=argparse.RawDescriptionHelpFormatter,
    )
    parser.add_argument('nc_file', nargs='?', help='NetCDF file (omit for interactive mode)')
    parser.add_argument('-v', '--variable', help='Variable name')
    parser.add_argument('-c', '--cmap', default='auto',
                        help='Matplotlib colormap, or "auto" to choose one from the '
                             'data (default: auto). Pass "jet" for the historical look.')
    parser.add_argument('-o', '--output', help='Save the figure to this path instead of displaying it')
    parser.add_argument('--figsize', metavar='W,H',
                        help='Figure size in inches, e.g. 12,5. Each plot kind '
                             'keeps its own default otherwise - a profile is '
                             'portrait and a map is landscape for a reason')
    parser.add_argument('--dpi', type=int,
                        help='Resolution of saved output. The figure on screen '
                             'is left alone, so this is for the printed page')
    parser.add_argument('--title', dest='title_override', metavar='TEXT',
                        help='Title the figure exactly like this, instead of '
                             'building one from the variable and the slice')
    parser.add_argument('-e', '--extra-indices',
                        help='JSON object selecting a 0-based index per dimension, '
                             'or "avg" to average over it. Example: \'{"Time": 0, "nslope": 1}\'')
    parser.add_argument('-x', '--x-dim', help='Dimension to place on the X axis of 2D plots')
    parser.add_argument('--plot-kind', choices=PLOT_KINDS, default='auto',
                        help='Force a renderer instead of inferring one (default: auto)')
    parser.add_argument('--no-remap', action='store_true',
                        help='Keep the raw unstructured grid instead of rebuilding a map')
    parser.add_argument('--explain', action='store_true',
                        help='Report how the file was interpreted and the plot chosen')
    parser.add_argument('--list-vars', action='store_true',
                        help='List the variables in the file and exit')
    parser.add_argument('--norm', choices=('linear', 'log', 'symlog', 'centered'),
                        help='Colour scale (default: chosen from the data)')
    parser.add_argument('--vmin', type=float, help='Lower limit of the colour scale')
    parser.add_argument('--vmax', type=float, help='Upper limit of the colour scale')
    parser.add_argument('--interpolate', action=argparse.BooleanOptionalAction,
                        default=True,
                        help='Draw the field interpolated between cell centres '
                             '(default). --no-interpolate draws the file\'s own '
                             'cells as flat quads instead, which is what to '
                             'reach for when the question is about the grid')
    parser.add_argument('--anomaly', choices=('time', 'zonal'),
                        help='Plot the departure from the time mean or the zonal mean')
    parser.add_argument('--reduce', action='append', metavar='OP:DIM[,DIM]',
                        help='Reduce dimensions before plotting, e.g. mean:lon. '
                             'Repeatable. Operations: mean, std, min, max, sum')
    parser.add_argument('--stats', action='store_true',
                        help='Print summary statistics alongside the plot')
    parser.add_argument('--stats-only', action='store_true',
                        help='Print summary statistics and produce no figure')
    parser.add_argument('--overlay', action='append', metavar='VAR[:CMAP]',
                        help='Draw another variable over this one, in its own '
                             'colours. Repeatable; later ones go on top. Each '
                             'layer is translucent where its own field is weak, '
                             'so an overlap reads as an overlap - CO2 and H2O '
                             'ice on one map, or two cloud decks on one globe')
    parser.add_argument('--diff', metavar='VAR',
                        help='Plot this variable subtracted from the one being '
                             'shown, on a scale centred on zero. Often the '
                             'question two overlaid layers were standing in '
                             'for: where one exceeds the other is one diverging '
                             'map, not two stacked fields')
    parser.add_argument('--overlay-style',
                        choices=('blend', 'contour', 'hatch', 'panels',
                                 'bivariate', 'glyph'),
                        default='blend', metavar='STYLE',
                        help='How the --overlay variables share the figure: '
                             '"blend" stacks them as translucent colour '
                             '(default), "contour" draws them as labelled lines '
                             'over the base - the one to reach for when both '
                             'fields have to stay readable - "hatch" marks where '
                             'they exceed a threshold, "panels" gives each its '
                             'own panel and works for every plot kind, '
                             '"bivariate" puts exactly two in one 2-D colour '
                             'scheme with a square key, "glyph" draws them as '
                             'sized symbols')
    parser.add_argument('--overlay-scale', choices=('own', 'shared'), default='own',
                        help='Under --overlay-style panels, whether every panel '
                             'uses one colour scale. "shared" is right for one '
                             'variable at several times; "own" (default) for '
                             'different variables, where a shared scale is set '
                             'by whichever is largest and flattens the rest')
    parser.add_argument('--overlay-threshold', metavar='VALUES',
                        help='Contour levels, or the level a hatch is drawn '
                             'above: comma-separated values, or "pNN" for a '
                             'percentile of the layer (default: six even levels '
                             'for contours, p90 for a hatch)')
    parser.add_argument('--vector', metavar='U,V',
                        help='Plot a vector field from two named components')
    parser.add_argument('--vector-style', choices=('quiver', 'stream'), default='quiver',
                        help='How to draw the vectors (default: quiver)')
    parser.add_argument('--vector-density', type=float,
                        help='Arrow count per axis (quiver) or streamline density')
    parser.add_argument('--vector-background', default='magnitude',
                        help='Scalar drawn under the vectors: "magnitude", "none", '
                             'or a variable name (default: magnitude)')
    parser.add_argument('--show-topo', action=argparse.BooleanOptionalAction, default=True,
                        help='Overlay topography on lat/lon maps (default: enabled)')
    parser.add_argument('--show-polar', action='store_true', help='Also display polar-stereo views')
    parser.add_argument('--show-3d', action='store_true', help='Also display the 3D globe view')
    parser.add_argument('--globe-relief', default='topo', metavar='SOURCE',
                        help='What shapes the 3D globe: "topo" (MOLA, the default), '
                             '"none" for a bare sphere, "data" for the plotted field, '
                             'or a variable name. A field carries no metric height and '
                             'is normalized into a fraction of the radius.')
    parser.add_argument('--globe-surface', choices=('field', 'terrain', 'none'),
                        help='What the sphere itself shows: "field" (the plotted '
                             'variable), "terrain" (MOLA relief in greys) or "none". '
                             'Defaults to terrain under an atmospheric shell, so the '
                             'data colours belong to the shell, and to field otherwise')
    parser.add_argument('--light', action=argparse.BooleanOptionalAction, default=None,
                        help='Shade the 3D globe. Off by default, because shading '
                             'makes the same value read differently across the '
                             'sphere; --sun turns it back on')
    parser.add_argument('--sun', metavar='LON,LAT',
                        help='Light the 3D globe from this sub-solar point and draw '
                             'the day/night terminator')
    parser.add_argument('--sun-ls', type=float, metavar='LS',
                        help='Sub-solar latitude from a solar longitude, in degrees. '
                             'Ls does not fix the longitude; pass --sun for that.')
    parser.add_argument('--globe-export', metavar='FILE',
                        help='Write the 3D globe mesh to a 3D file '
                             '(.vtk, .vtp, .ply, .obj, .stl)')
    parser.add_argument('--shell-mode', choices=('cloud', 'iso', 'column', 'layers'),
                        default='cloud',
                        help='How --plot-kind globe draws an atmospheric field: '
                             '"cloud" ray-casts it as a translucent volume (default), '
                             '"iso" contours it into cloud and dust bodies, '
                             '"column" integrates it over the vertical onto one '
                             'translucent shell, "layers" stacks one surface per level')
    parser.add_argument('--shell-threshold', metavar='VALUES',
                        help='Isosurface level(s), comma-separated: a value, or "pNN" '
                             'for a percentile of the field (default: p99). Several '
                             'nest one surface inside another, e.g. p90,p99,p99.9')
    parser.add_argument('--shell-color', choices=('value', 'height'), default='value',
                        help='What colours an isosurface: "value" draws it in its own '
                             'threshold\'s colour on the field scale (default), '
                             '"height" colours it by altitude instead')
    parser.add_argument('--shell-level', type=int, metavar='N',
                        help='Draw only vertical level N as a solid shell')
    parser.add_argument('--shell-opacity', type=float, default=0.35, metavar='A',
                        help='Opacity of the shell, cloud or surfaces (default: 0.35)')
    parser.add_argument('--shell-cut', action=argparse.BooleanOptionalAction, default=None,
                        help='Slice the shell open with a plane through the '
                             'planet, draggable by mouse, so the view gets '
                             'inside it. On by default for --shell-mode layers, '
                             'which cannot be read any other way; worth trying '
                             'on a cloud too')
    parser.add_argument('--shell-resolution', type=int, default=None, metavar='N',
                        help='Voxels per side of the box --shell-mode cloud is '
                             'ray-cast from (default: 256). Higher is sharper and '
                             'costs memory as the cube of N; the cloud keeps the '
                             'same density either way')
    parser.add_argument('--animate', metavar='DIM',
                        help='Step a lat/lon map, its polar views and the 3D globe '
                             'through this dimension. Interactively it becomes a '
                             'slider; with -o and a movie extension (.mp4, .gif) it '
                             'is written as a video.')
    parser.add_argument('--spin', type=float, nargs='?', const=1.0, default=None,
                        metavar='TURNS',
                        help='Rotate the 3D globe camera. On screen it turns by '
                             'itself and stays draggable; with -o and a movie '
                             'extension it is written as a turning video, on its '
                             'own or on top of --animate (default: 1 turn)')
    parser.add_argument('--fps', type=int, default=12,
                        help='Frames per second for --animate video output (default: 12)')
    parser.add_argument('--shell-top', type=float, metavar='KM',
                        help='Where the shell tops out when the vertical coordinate '
                             'is not a length, such as pressure or sigma levels '
                             '(default: 50 km)')
    parser.add_argument('--shell-exaggeration', type=float, metavar='N',
                        help='Vertical exaggeration of the atmosphere (default: 30). '
                             'The relief keeps its own x10, so the weather in the '
                             'bottom scale height is visible instead of being a film '
                             'on the ground; pass 10 for one shared scale')
    return parser


def _validate(args, parser):
    """
    The argument combinations argparse cannot express on its own.

    Refused rather than corrected. Silently swapping an inverted pair of limits
    would be the quiet substitution this tool avoids everywhere else - and an
    inverted scale is a plausible typo, not a plausible intent. `parser.error`
    exits 2, the documented status for bad arguments.
    """
    if args.vmin is not None and args.vmax is not None and args.vmin >= args.vmax:
        parser.error(f"--vmin {args.vmin:g} is not below --vmax {args.vmax:g}; "
                     f"an inverted or empty colour scale cannot be read")

    # A comma, matching --sun LON,LAT. Deliberately not also accepting 8x6:
    # one spelling per idea is what keeps the options readable.
    if args.figsize is not None:
        parts = [p.strip() for p in str(args.figsize).split(',')]
        try:
            if len(parts) != 2:
                raise ValueError('expected two numbers')
            width, height = (float(p) for p in parts)
            if width <= 0 or height <= 0:
                raise ValueError('both must be positive')
        except ValueError as err:
            parser.error(f"--figsize '{args.figsize}' is not W,H in inches ({err})")
        args.figsize = (width, height)

    if args.dpi is not None and args.dpi <= 0:
        parser.error(f"--dpi {args.dpi} is not a positive number of dots per inch")


def main():
    parser = build_parser()
    args = parser.parse_args()
    _validate(args, parser)

    if args.nc_file and (args.variable or args.list_vars or args.vector):
        return visualize_variable_cli(args)

    if (args.variable or args.vector) and not args.nc_file:
        parser.error("--variable requires a NetCDF file argument")
    return visualize_variable_interactive(args.nc_file, args)
