#!/usr/bin/env python3
# encoding: utf-8
# Copyright (C) 2020 Space Science and Engineering Center (SSEC),
#  University of Wisconsin-Madison.
#
#     This program is free software: you can redistribute it and/or modify
#     it under the terms of the GNU General Public License as published by
#     the Free Software Foundation, either version 3 of the License, or
#     (at your option) any later version.
#
#     This program is distributed in the hope that it will be useful,
#     but WITHOUT ANY WARRANTY; without even the implied warranty of
#     MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#     GNU General Public License for more details.
#
#     You should have received a copy of the GNU General Public License
#     along with this program.  If not, see <http://www.gnu.org/licenses/>.
#
# This file is part of the polar2grid software package. Polar2grid takes
# satellite observation data, remaps it, and writes it to a file format for
# input into another program.
# Documentation: http://www.ssec.wisc.edu/software/polar2grid/
"""The VIIRS SDR Reader operates on Sensor Data Record (SDR) HDF5 files from
the Suomi National Polar-orbiting Partnership (SNPP), NOAA21 and  NOAA20
Visible/Infrared Imager Radiometer Suite (VIIRS) instruments. The VIIRS
SDR reader requires filenames to match one of a couple different standard
filename schemes used for official products. SDR files are typically named
as below and have corresponding geolocation files::

    SVI01_j02_d20260205_t1911178_e1912424_b16791_c20260205194106784585_cspp_dev.h5

The VIIRS SDR reader supports all instrument spectral bands, identified as
the products shown below. It supports terrain corrected or non-terrain corrected
navigation files. Geolocation files must be included when specifying filepaths to
readers and ``polar2grid.sh``. The VIIRS reader can be specified to the ``polar2grid.sh`` script
with the reader name ``viirs_sdr``.

This reader's default remapping algorithm is ``ewa`` for Elliptical Weighted
Averaging resampling. The ``--weight-delta-max`` parameter set to 40 and the
``--weight-distance-max`` parameter set to 2.

+---------------------------+-----------------------------------------------------+
| Product Name              | Description                                         |
+===========================+=====================================================+
| i01                       | I01 Reflectance Band                                |
+---------------------------+-----------------------------------------------------+
| i02                       | I02 Reflectance Band                                |
+---------------------------+-----------------------------------------------------+
| i03                       | I03 Reflectance Band                                |
+---------------------------+-----------------------------------------------------+
| i04                       | I04 Brightness Temperature Band                     |
+---------------------------+-----------------------------------------------------+
| i05                       | I05 Brightness Temperature Band                     |
+---------------------------+-----------------------------------------------------+
| i01_rad                   | I01 Radiance Band                                   |
+---------------------------+-----------------------------------------------------+
| i02_rad                   | I02 Radiance Band                                   |
+---------------------------+-----------------------------------------------------+
| i03_rad                   | I03 Radiance Band                                   |
+---------------------------+-----------------------------------------------------+
| i04_rad                   | I04 Radiance Band                                   |
+---------------------------+-----------------------------------------------------+
| i05_rad                   | I05 Radiance Band                                   |
+---------------------------+-----------------------------------------------------+
| m01                       | M01 Reflectance Band                                |
+---------------------------+-----------------------------------------------------+
| m02                       | M02 Reflectance Band                                |
+---------------------------+-----------------------------------------------------+
| m03                       | M03 Reflectance Band                                |
+---------------------------+-----------------------------------------------------+
| m04                       | M04 Reflectance Band                                |
+---------------------------+-----------------------------------------------------+
| m05                       | M05 Reflectance Band                                |
+---------------------------+-----------------------------------------------------+
| m06                       | M06 Reflectance Band                                |
+---------------------------+-----------------------------------------------------+
| m07                       | M07 Reflectance Band                                |
+---------------------------+-----------------------------------------------------+
| m08                       | M08 Reflectance Band                                |
+---------------------------+-----------------------------------------------------+
| m09                       | M09 Reflectance Band                                |
+---------------------------+-----------------------------------------------------+
| m10                       | M10 Reflectance Band                                |
+---------------------------+-----------------------------------------------------+
| m11                       | M11 Reflectance Band                                |
+---------------------------+-----------------------------------------------------+
| m12                       | M12 Brightness Temperature Band                     |
+---------------------------+-----------------------------------------------------+
| m13                       | M13 Brightness Temperature Band                     |
+---------------------------+-----------------------------------------------------+
| m14                       | M14 Brightness Temperature Band                     |
+---------------------------+-----------------------------------------------------+
| m15                       | M15 Brightness Temperature Band                     |
+---------------------------+-----------------------------------------------------+
| m16                       | M16 Brightness Temperature Band                     |
+---------------------------+-----------------------------------------------------+
| m01_rad                   | M01 Radiance Band                                   |
+---------------------------+-----------------------------------------------------+
| m02_rad                   | M02 Radiance Band                                   |
+---------------------------+-----------------------------------------------------+
| m03_rad                   | M03 Radiance Band                                   |
+---------------------------+-----------------------------------------------------+
| m04_rad                   | M04 Radiance Band                                   |
+---------------------------+-----------------------------------------------------+
| m05_rad                   | M05 Radiance Band                                   |
+---------------------------+-----------------------------------------------------+
| m06_rad                   | M06 Radiance Band                                   |
+---------------------------+-----------------------------------------------------+
| m07_rad                   | M07 Radiance Band                                   |
+---------------------------+-----------------------------------------------------+
| m08_rad                   | M08 Radiance Band                                   |
+---------------------------+-----------------------------------------------------+
| m09_rad                   | M09 Radiance Band                                   |
+---------------------------+-----------------------------------------------------+
| m10_rad                   | M10 Radiance Band                                   |
+---------------------------+-----------------------------------------------------+
| m11_rad                   | M11 Radiance Band                                   |
+---------------------------+-----------------------------------------------------+
| m12_rad                   | M12 Radiance Band                                   |
+---------------------------+-----------------------------------------------------+
| m13_rad                   | M13 Radiance Band                                   |
+---------------------------+-----------------------------------------------------+
| m14_rad                   | M14 Radiance Band                                   |
+---------------------------+-----------------------------------------------------+
| m15_rad                   | M15 Radiance Band                                   |
+---------------------------+-----------------------------------------------------+
| m16_rad                   | M16 Radiance Band                                   |
+---------------------------+-----------------------------------------------------+
| dnb                       | Raw DNB Band (not useful for images)                |
+---------------------------+-----------------------------------------------------+
| histogram_dnb             | Histogram Equalized DNB Band                        |
+---------------------------+-----------------------------------------------------+
| adaptive_dnb              | Adaptive Histogram Equalized DNB Band               |
+---------------------------+-----------------------------------------------------+
| dynamic_dnb               | Dynamic DNB Band from Steve Miller and              |
|                           | Curtis Seaman. Uses erf to scale the data.          |
+---------------------------+-----------------------------------------------------+
| hncc_dnb                  | Simplified High and Near-Constant Contrast          |
|                           | Approach from Stephan Zinke                         |
+---------------------------+-----------------------------------------------------+
| ifog                      | Temperature difference between I05 and I04          |
+---------------------------+-----------------------------------------------------+
| i_solar_zenith_angle      | I Band Solar Zenith Angle                           |
+---------------------------+-----------------------------------------------------+
| i_solar_azimuth_angle     | I Band Solar Azimuth Angle                          |
+---------------------------+-----------------------------------------------------+
| i_sat_zenith_angle        | I Band Satellite Zenith Angle                       |
+---------------------------+-----------------------------------------------------+
| i_sat_azimuth_angle       | I Band Satellite Azimuth Angle                      |
+---------------------------+-----------------------------------------------------+
| m_solar_zenith_angle      | M Band Solar Zenith Angle                           |
+---------------------------+-----------------------------------------------------+
| m_solar_azimuth_angle     | M Band Solar Azimuth Angle                          |
+---------------------------+-----------------------------------------------------+
| m_sat_zenith_angle        | M Band Satellite Zenith Angle                       |
+---------------------------+-----------------------------------------------------+
| m_sat_azimuth_angle       | M Band Satellite Azimuth Angle                      |
+---------------------------+-----------------------------------------------------+
| dnb_solar_zenith_angle    | DNB Band Solar Zenith Angle                         |
+---------------------------+-----------------------------------------------------+
| dnb_solar_azimuth_angle   | DNB Band Solar Azimuth Angle                        |
+---------------------------+-----------------------------------------------------+
| dnb_sat_zenith_angle      | DNB Band Satellite Zenith Angle                     |
+---------------------------+-----------------------------------------------------+
| dnb_sat_azimuth_angle     | DNB Band Satellite Azimuth Angle                    |
+---------------------------+-----------------------------------------------------+
| dnb_lunar_zenith_angle    | DNB Band Lunar Zenith Angle                         |
+---------------------------+-----------------------------------------------------+
| dnb_lunar_azimuth_angle   | DNB Band Lunar Azimuth Angle                        |
+---------------------------+-----------------------------------------------------+
| true_color                | Ratio sharpened rayleigh corrected true color       |
+---------------------------+-----------------------------------------------------+
| false_color               | Ratio sharpened rayleigh corrected false color      |
+---------------------------+-----------------------------------------------------+

"""

from __future__ import annotations

from argparse import ArgumentParser, _ArgumentGroup
from typing import Optional

from satpy import DataQuery, Scene

from polar2grid.core.script_utils import ExtendConstAction

from ._base import ReaderProxyBase

PREFERRED_CHUNK_SIZE: int = 6400

I_VIS_PRODUCTS = [
    "I01",
    "I02",
    "I03",
]
I_IR_PRODUCTS = [
    "I04",
    "I05",
]
I_PRODUCTS = I_VIS_PRODUCTS + I_IR_PRODUCTS
I_ANGLE_PRODUCTS = [
    "i_solar_zenith_angle",
    "i_solar_azimuth_angle",
    "i_sat_zenith_angle",
    "i_sat_azimuth_angle",
]
M_VIS_PRODUCTS = [
    "M01",
    "M02",
    "M03",
    "M04",
    "M05",
    "M06",
    "M07",
    "M08",
    "M09",
    "M10",
    "M11",
]
M_IR_PRODUCTS = [
    "M12",
    "M13",
    "M14",
    "M15",
    "M16",
]
M_PRODUCTS = M_VIS_PRODUCTS + M_IR_PRODUCTS
M_ANGLE_PRODUCTS = [
    "m_solar_zenith_angle",
    "m_solar_azimuth_angle",
    "m_sat_zenith_angle",
    "m_sat_azimuth_angle",
]
DNB_PRODUCTS = [
    "histogram_dnb",
    "adaptive_dnb",
    "dynamic_dnb",
    "hncc_dnb",
]
DNB_ANGLE_PRODUCTS = [
    "dnb_solar_zenith_angle",
    "dnb_solar_azimuth_angle",
    "dnb_sat_zenith_angle",
    "dnb_sat_azimuth_angle",
    "dnb_lunar_zenith_angle",
    "dnb_lunar_azimuth_angle",
]
TRUE_COLOR_PRODUCTS = ["true_color"]
FALSE_COLOR_PRODUCTS = ["false_color"]
OTHER_COMPS = [
    "ifog",
]


PRODUCT_ALIASES = {}


def _process_legacy_and_rad_products(satpy_names, band_aliases, rad_aliases):
    """Map all lowercase band names to uppercase names and add radiance product."""
    for band in satpy_names:
        # P2G name is lowercase, Satpy is uppercase
        PRODUCT_ALIASES[band.lower()] = band
        band_aliases.append(band.lower())

        # radiance products for M and I bands
        rad_name = band.lower() + "_rad"
        dq = DataQuery(name=band, calibration="radiance")
        PRODUCT_ALIASES[rad_name] = dq
        rad_aliases.append(rad_name)


I_ALIASES = []
I_RAD_PRODUCTS = []
_process_legacy_and_rad_products(I_PRODUCTS, I_ALIASES, I_RAD_PRODUCTS)
M_ALIASES = []
M_RAD_PRODUCTS = []
_process_legacy_and_rad_products(M_PRODUCTS, M_ALIASES, M_RAD_PRODUCTS)

_AWIPS_TRUE_COLOR = ["viirs_crefl08", "viirs_crefl04", "viirs_crefl03"]
_AWIPS_FALSE_COLOR = ["viirs_crefl07", "viirs_crefl09", "viirs_crefl08"]

PRODUCT_ALIASES["dnb_solar_zenith_angle"] = DataQuery(name="dnb_solar_zenith_angle")
PRODUCT_ALIASES["dnb_solar_azimuth_angle"] = DataQuery(name="dnb_solar_azimuth_angle")
PRODUCT_ALIASES["dnb_sat_zenith_angle"] = DataQuery(name="dnb_satellite_zenith_angle")
PRODUCT_ALIASES["dnb_sat_azimuth_angle"] = DataQuery(name="dnb_satellite_azimuth_angle")
PRODUCT_ALIASES["dnb_lunar_zenith_angle"] = DataQuery(name="dnb_lunar_zenith_angle")
PRODUCT_ALIASES["dnb_lunar_azimuth_angle"] = DataQuery(name="dnb_lunar_azimuth_angle")
PRODUCT_ALIASES["m_solar_zenith_angle"] = DataQuery(name="solar_zenith_angle", resolution=742)
PRODUCT_ALIASES["m_solar_azimuth_angle"] = DataQuery(name="solar_azimuth_angle", resolution=742)
PRODUCT_ALIASES["m_sat_zenith_angle"] = DataQuery(name="satellite_zenith_angle", resolution=742)
PRODUCT_ALIASES["m_sat_azimuth_angle"] = DataQuery(name="satellite_azimuth_angle", resolution=742)
PRODUCT_ALIASES["i_solar_zenith_angle"] = DataQuery(name="solar_zenith_angle", resolution=371)
PRODUCT_ALIASES["i_solar_azimuth_angle"] = DataQuery(name="solar_azimuth_angle", resolution=371)
PRODUCT_ALIASES["i_sat_zenith_angle"] = DataQuery(name="satellite_zenith_angle", resolution=371)
PRODUCT_ALIASES["i_sat_azimuth_angle"] = DataQuery(name="satellite_azimuth_angle", resolution=371)

DEFAULT_PRODUCTS = I_ALIASES + M_ALIASES + DNB_PRODUCTS[1:] + TRUE_COLOR_PRODUCTS + FALSE_COLOR_PRODUCTS + OTHER_COMPS
P2G_PRODUCTS = I_ALIASES + M_ALIASES + DNB_PRODUCTS + I_RAD_PRODUCTS + M_RAD_PRODUCTS
P2G_PRODUCTS += I_ANGLE_PRODUCTS + M_ANGLE_PRODUCTS + DNB_ANGLE_PRODUCTS + OTHER_COMPS
P2G_PRODUCTS += TRUE_COLOR_PRODUCTS + FALSE_COLOR_PRODUCTS
P2G_PRODUCTS += [
    "viirs_crefl01",
    "viirs_crefl02",
    "viirs_crefl03",
    "viirs_crefl04",
    "viirs_crefl07",
    "viirs_crefl08",
    "viirs_crefl09",
]
P2G_PRODUCTS += ["dynamic_dnb_saturation", "unsharp_true_color"]

FILTERS = {
    "day_only": {
        "standard_name": [
            "toa_bidirectional_reflectance",
            "true_color",
            "false_color",
            "natural_color",
            "corrected_reflectance",
        ],
    },
    "night_only": {
        "standard_name": ["temperature_difference"],
    },
}


class ReaderProxy(ReaderProxyBase):
    """Provide Polar2Grid-specific information about this reader's products."""

    is_polar2grid_reader = True

    def __init__(self, scn: Scene, user_products: list[str]):
        self._modified_aliases = PRODUCT_ALIASES.copy()
        if "dynamic_dnb_saturation" in user_products:
            # they specified --dnb-saturation-correction
            # let's modify the aliases so dynamic_dnb points to this product
            user_products.remove("dynamic_dnb_saturation")
            self._modified_aliases["dynamic_dnb"] = DataQuery(name="dynamic_dnb_saturation")
        super().__init__(scn, user_products)

    def get_default_products(self) -> list[str]:
        """Get products to load if users hasn't specified any others."""
        return DEFAULT_PRODUCTS

    def get_all_products(self):
        """Get all polar2grid products that could be loaded."""
        return P2G_PRODUCTS

    @property
    def _aliases(self):
        return self._modified_aliases


def add_reader_argument_groups(
    parser: ArgumentParser, group: Optional[_ArgumentGroup] = None
) -> tuple[Optional[_ArgumentGroup], Optional[_ArgumentGroup]]:
    """Add reader-specific command line arguments to an existing argument parser.

    If ``group`` is provided then arguments are added to this group. If not,
    a new group is added to the parser and arguments added to this new group.

    """
    if group is None:
        group = parser.add_argument_group(title="VIIRS SDR Reader")

    group.add_argument(
        "--i-bands",
        dest="products",
        action=ExtendConstAction,
        const=I_ALIASES,
        help="Add all I-band raw products to list of products",
    )
    group.add_argument(
        "--m-bands",
        dest="products",
        action=ExtendConstAction,
        const=M_ALIASES,
        help="Add all M-band raw products to list of products",
    )
    group.add_argument(
        "--dnb-angle-products",
        dest="products",
        action=ExtendConstAction,
        const=DNB_ANGLE_PRODUCTS,
        help="Add DNB-band geolocation 'angle' products to list of products",
    )
    group.add_argument(
        "--dnb-saturation-correction",
        dest="products",
        action=ExtendConstAction,
        const=["dynamic_dnb_saturation"],
        help="Enable dynamic DNB saturation correction (normally used for aurora scenes)",
    )
    group.add_argument(
        "--i-angle-products",
        dest="products",
        action=ExtendConstAction,
        const=I_ANGLE_PRODUCTS,
        help="Add I-band geolocation 'angle' products to list of products",
    )
    group.add_argument(
        "--m-angle-products",
        dest="products",
        action=ExtendConstAction,
        const=M_ANGLE_PRODUCTS,
        help="Add M-band geolocation 'angle' products to list of products",
    )
    group.add_argument(
        "--m-rad-products",
        dest="products",
        action=ExtendConstAction,
        const=M_RAD_PRODUCTS,
        help="Add M-band geolocation radiance products to list of products",
    )
    group.add_argument(
        "--i-rad-products",
        dest="products",
        action=ExtendConstAction,
        const=I_RAD_PRODUCTS,
        help="Add I-band geolocation radiance products to list of products",
    )
    group.add_argument(
        "--awips-true-color",
        dest="products",
        action=ExtendConstAction,
        const=_AWIPS_TRUE_COLOR,
        help="Add individual CREFL corrected products to create the 'true_color' composite in AWIPS.",
    )
    group.add_argument(
        "--awips-false-color",
        dest="products",
        action=ExtendConstAction,
        const=_AWIPS_FALSE_COLOR,
        help="Add individual CREFL corrected products to create the 'false_color' composite in AWIPS.",
    )
    return group, None
