"""
CLI entry point for pygndc.
"""

import argparse
import sys
import os

from pygndc import __version__
from pygndc.reader import GNDCReader

MODE_CHOICES = ["auto", "tcnn_cuda", "torch_gpu", "torch_cpu"]


def _add_mode_arg(parser):
    """Add --mode argument to a subparser."""
    parser.add_argument(
        "--mode", choices=MODE_CHOICES, default="auto",
        help="Inference mode (default: auto)"
    )


def main():
    parser = argparse.ArgumentParser(
        description="pygndc - Geographic Neural Data Compression",
        prog="pygndc"
    )
    parser.add_argument("--version", action="version", version=f"pygndc {__version__}")

    subparsers = parser.add_subparsers(dest="command", help="Available commands")

    # Decompress command
    decompress_parser = subparsers.add_parser("decompress", help="Decompress/reconstruct data")
    decompress_parser.add_argument("-i", "--input", required=True, help="Input .gndc model file")
    decompress_parser.add_argument("-o", "--output", required=True, help="Output path (file or directory)")
    decompress_parser.add_argument("--timestamp", help="Specific timestamp (YYYY-MM-DD)")
    decompress_parser.add_argument("--start", dest="start_date", help="Start date (YYYY-MM-DD)")
    decompress_parser.add_argument("--end", dest="end_date", help="End date (YYYY-MM-DD)")
    decompress_parser.add_argument("--interval", type=int, default=5, help="Interval in days")
    decompress_parser.add_argument("--format", choices=["tif", "npy"], default="tif", help="Output format")
    _add_mode_arg(decompress_parser)

    # Info command
    info_parser = subparsers.add_parser("info", help="Show model information")
    info_parser.add_argument("model", help="Path to .gndc model file")
    info_parser.add_argument("--json", action="store_true", help="Output as JSON")

    # Viewer command
    viewer_parser = subparsers.add_parser("viewer", help="Launch interactive viewer")
    viewer_parser.add_argument("-i", "--input", required=True, help="Input .gndc model file")
    _add_mode_arg(viewer_parser)

    # NDVI command
    ndvi_parser = subparsers.add_parser("ndvi", help="Compute NDVI and export as TIF")
    ndvi_parser.add_argument("-i", "--input", required=True, help="Input .gndc model file")
    ndvi_parser.add_argument("-o", "--output", required=True, help="Output NDVI TIF file")
    ndvi_parser.add_argument("--timestamp", type=int, default=0, help="Time index (default: 0)")
    ndvi_parser.add_argument("--red-band", type=int, default=0, help="0-based red band index")
    ndvi_parser.add_argument("--nir-band", type=int, default=1, help="0-based NIR band index")
    _add_mode_arg(ndvi_parser)

    # Derivative command
    deriv_parser = subparsers.add_parser("derivative", help="Compute spatial gradients")
    deriv_parser.add_argument("-i", "--input", required=True, help="Input .gndc model file")
    deriv_parser.add_argument("-o", "--output", required=True, help="Output gradient TIF file")
    deriv_parser.add_argument("--grad-mode", choices=["dx", "dy", "mag", "dir"], default="mag",
                              help="Derivative mode (default: mag)")
    deriv_parser.add_argument("--timestamp", type=int, default=0, help="Time index (default: 0)")
    _add_mode_arg(deriv_parser)

    # Sample command
    sample_parser = subparsers.add_parser("sample", help="Point query by coordinates")
    sample_parser.add_argument("-i", "--input", required=True, help="Input .gndc model file")
    sample_parser.add_argument("--lon", type=float, required=True, help="Longitude")
    sample_parser.add_argument("--lat", type=float, required=True, help="Latitude")
    sample_parser.add_argument("--timestamp", type=int, default=None,
                               help="Time index (omit for full series)")
    _add_mode_arg(sample_parser)

    # Timeseries command
    ts_parser = subparsers.add_parser("timeseries", help="Export pixel time series to CSV")
    ts_parser.add_argument("-i", "--input", required=True, help="Input .gndc model file")
    ts_parser.add_argument("--row", type=int, required=True, help="Pixel row")
    ts_parser.add_argument("--col", type=int, required=True, help="Pixel column")
    ts_parser.add_argument("-o", "--output", required=True, help="Output CSV file")
    _add_mode_arg(ts_parser)

    args = parser.parse_args()

    if args.command == "decompress":
        reader = GNDCReader(mode=args.mode)
        reader.load(args.input)

        if args.timestamp:
            from datetime import datetime
            from pygndc.utils.date_utils import normalize_time
            import numpy as np
            meta = reader.dataset_meta
            g_start = datetime.strptime(str(meta['global_start_date']), "%Y-%m-%d %H:%M:%S")
            g_end = datetime.strptime(str(meta['global_end_date']), "%Y-%m-%d %H:%M:%S")
            t_norms = meta.get('t_norms', [])
            target_dt = datetime.strptime(args.timestamp, "%Y-%m-%d")
            target_tnorm = normalize_time(target_dt, g_start, g_end)
            t_idx = min(range(len(t_norms)), key=lambda i: abs(t_norms[i] - target_tnorm))
            frame = reader.reconstruct_single_frame(t_idx)
            np.save(args.output, frame)
            print(f"Saved frame {t_idx} ({args.timestamp}) to {args.output}")
        else:
            recons_dates = None
            if args.start_date and args.end_date:
                recons_dates = (args.start_date, args.end_date, args.interval)
            reader.reconstruct_and_save_all(args.output, recons_dates=recons_dates)

    elif args.command == "info":
        import json
        import pygndc
        ds = pygndc.open(args.model)

        if args.json:
            print(json.dumps(ds.meta, indent=2, default=str))
        else:
            print(ds.info())
        ds.close()

    elif args.command == "viewer":
        from pygndc.viewer.interactive import GNDCInteractiveViewer
        app = GNDCInteractiveViewer()
        app.root.mainloop()

    elif args.command == "ndvi":
        import pygndc
        import numpy as np
        ds = pygndc.open(args.input, mode=args.mode)
        ndvi = ds.ndvi(t=args.timestamp, red_band=args.red_band, nir_band=args.nir_band)

        from pygndc.interop.interop import save_frame_as_tif
        save_frame_as_tif(
            ndvi[:, :, np.newaxis], args.output, ds.crs, ds.transform, "float32"
        )
        print(f"NDVI saved to {args.output}")
        ds.close()

    elif args.command == "derivative":
        import pygndc
        ds = pygndc.open(args.input, mode=args.mode)
        result = ds.gradient(t=args.timestamp, mode=args.grad_mode)
        if isinstance(result, tuple):
            import numpy as np
            result = np.sqrt(result[0] ** 2 + result[1] ** 2)

        from pygndc.interop.interop import save_frame_as_tif
        if result.ndim == 2:
            import numpy as np
            result = result[:, :, np.newaxis]
        save_frame_as_tif(result, args.output, ds.crs, ds.transform, "float32")
        print(f"Gradient ({args.grad_mode}) saved to {args.output}")
        ds.close()

    elif args.command == "sample":
        import pygndc
        import numpy as np
        ds = pygndc.open(args.input, mode=args.mode)

        if args.timestamp is not None:
            vals = ds.sample(lon=args.lon, lat=args.lat, t=args.timestamp)
            print(f"Values at ({args.lon}, {args.lat}), t={args.timestamp}: {vals}")
        else:
            series = ds.sample(lon=args.lon, lat=args.lat)
            print(f"Time series at ({args.lon}, {args.lat}): shape={series.shape}")
            for t_idx in range(series.shape[0]):
                ts_str = ds.timestamps[t_idx] if t_idx < len(ds.timestamps) else str(t_idx)
                vals_str = ", ".join(f"{v:.4f}" for v in series[t_idx])
                print(f"  [{t_idx}] {ts_str}: {vals_str}")
        ds.close()

    elif args.command == "timeseries":
        import pygndc
        import numpy as np
        ds = pygndc.open(args.input, mode=args.mode)
        series = ds.pixel_series(row=args.row, col=args.col)

        # Write CSV
        os.makedirs(os.path.dirname(os.path.abspath(args.output)), exist_ok=True)
        n_bands = series.shape[1]
        band_headers = [f"band_{i}" for i in range(n_bands)]

        with open(args.output, "w") as f:
            f.write("t_idx,timestamp," + ",".join(band_headers) + "\n")
            for t_idx in range(series.shape[0]):
                ts_str = str(ds.timestamps[t_idx]) if t_idx < len(ds.timestamps) else str(t_idx)
                vals_str = ",".join(f"{v:.6f}" for v in series[t_idx])
                f.write(f"{t_idx},{ts_str},{vals_str}\n")

        print(f"Time series ({args.row}, {args.col}) saved to {args.output}")
        ds.close()

    else:
        parser.print_help()


if __name__ == "__main__":
    main()
