#!/usr/bin/env python
# Copyright (c) 2024, openradar developers.
# Distributed under the MIT License. See LICENSE for more info.

"""

Datamet
=======

This sub-module contains the Datamet xarray backend for reading Datamet-based radar
data into Xarray structures.

.. autosummary::
   :nosignatures:
   :toctree: generated/

   {}

"""

__all__ = ["DataMetBackendEntrypoint", "open_datamet_datatree"]

__doc__ = __doc__.format("\n   ".join(__all__))

import gzip
import io
import os
import tarfile
from collections import defaultdict
from datetime import datetime, timedelta

import numpy as np
import xarray as xr
from datatree import DataTree
from xarray.backends.common import AbstractDataStore, BackendArray, BackendEntrypoint
from xarray.backends.file_manager import CachingFileManager
from xarray.backends.store import StoreBackendEntrypoint
from xarray.core import indexing
from xarray.core.utils import FrozenDict
from xarray.core.variable import Variable

from ... import util
from ...model import (
    get_altitude_attrs,
    get_azimuth_attrs,
    get_elevation_attrs,
    get_latitude_attrs,
    get_longitude_attrs,
    get_range_attrs,
    get_time_attrs,
    moment_attrs,
    sweep_vars_mapping,
)
from .common import _assign_root, _attach_sweep_groups

#: mapping from DataMet names to CfRadial2/ODIM
datamet_mapping = {
    "UZ": "DBTH",
    "CZ": "DBZH",
    "V": "VRADH",
    "W": "WRADH",
    "ZDR": "ZDR",
    "PHIDP": "PHIDP",
    "RHOHV": "RHOHV",
    "KDP": "KDP",
}


def convert_value(value):
    """
    Try to convert the value to an int or float. If conversion is not possible,
    return the original value.
    """
    try:
        return int(value)
    except (ValueError, TypeError):
        pass

    try:
        return float(value)
    except (ValueError, TypeError):
        pass

    return value


class DataMetFile:
    def __init__(self, filename):
        self.filename = filename
        # Handle .tar.gz case
        if self.filename.endswith(".gz"):
            with gzip.open(self.filename, "rb") as gzip_file:
                tar_bytes = io.BytesIO(gzip_file.read())
            self.tarfile = tarfile.open(fileobj=tar_bytes)
        else:
            self.tarfile = tarfile.open(self.filename, "r")
        self.first_dimension = "azimuth"  # No idea if other scan types are available
        self.scan_metadata = self.get_scan_metadata()
        self.moments = self.scan_metadata["measure"]
        self.data = dict()

    def extract_parameters(self, parameter_path):
        # Extract set of parameters from a file in the tarball
        member = self.tarfile.getmember(parameter_path)
        file = self.tarfile.extractfile(member)
        labels = np.loadtxt(file, delimiter="=", dtype=str, usecols=[0])
        file.seek(0)  # Reset file pointer
        values = np.loadtxt(file, delimiter="=", dtype=str, usecols=[1])
        values = np.array(values, ndmin=1)
        labels = np.array(labels, ndmin=1)
        # Iterate over the labels and values, appending values to the appropriate label key
        parameters = defaultdict(list)
        for label, value in zip(labels, values):
            parameters[label.strip()].append(value.strip())
        parameters = dict(parameters)
        # Convert lists with a single element to individual values
        parameters = {
            k: convert_value(v[0]) if len(v) == 1 else v for k, v in parameters.items()
        }
        return parameters

    def extract_data(self, data_path, dtype):
        # Extract moment data from a file in the tarball
        member = self.tarfile.getmember(data_path + "/SCAN.dat")
        file = self.tarfile.extractfile(member)
        data = np.frombuffer(file.read(), dtype=dtype)
        return data

    def get_scan_metadata(self):
        # Get all metadata at scan level (valid for all sweeps/moments)
        navigation = self.extract_parameters(os.path.join(".", "navigation.txt"))
        archiviation = self.extract_parameters(os.path.join(".", "archiviation.txt"))
        return {**navigation, **archiviation}

    def get_mom_metadata(self, mom, sweep):
        # Get all metadata that is moment and/or sweep dependent
        # Note that DataMet uses 1-indexed but we switch to 0-indexed
        # as is done in other backends
        mom_path = os.path.join(".", mom)
        sweep_path = os.path.join(".", mom, str(sweep + 1))

        generic = self.extract_parameters(os.path.join(sweep_path, "generic.txt"))
        calibration_momlvl = self.extract_parameters(
            os.path.join(mom_path, "calibration.txt")
        )
        calibration_sweeplvl = self.extract_parameters(
            os.path.join(sweep_path, "calibration.txt")
        )
        navigation_var = self.extract_parameters(
            os.path.join(sweep_path, "navigation.txt")
        )
        return {
            **navigation_var,
            **generic,
            **calibration_sweeplvl,
            **calibration_momlvl,
        }

    def get_moment(self, mom, sweep):
        # Get the data for a moment and apply byte to float conversion
        # Note that DataMet uses 1-indexed but we switch to 0-indexed
        # as is done in other backends
        mom_path = os.path.join(".", mom, str(sweep + 1))
        mom_medata = self.get_mom_metadata(mom, sweep)

        bitplanes = 16 if mom == "PHIDP" else int(mom_medata["bitplanes"] or 8)
        nazim = int(mom_medata.get("nlines"))
        nrange = int(mom_medata.get("ncols"))

        dtype = np.uint16 if bitplanes == 16 else np.uint8
        data = self.extract_data(mom_path, dtype)
        data = np.reshape(data, (nazim, nrange))

        return data

    def get_sweep(self, sweep):
        # Get data for all moments and sweeps
        moment_names = self.moments
        if sweep not in self.data:
            self.data[sweep] = dict()

        for moment in moment_names:
            if moment not in self.data[sweep]:
                self.data[sweep][moment] = self.get_moment(moment, sweep)

    def close(self):
        # Clase tarfile
        if self.tarfile is not None:
            self.tarfile.close()


class DataMetArrayWrapper(BackendArray):
    """Wraps array of DataMet RAW data."""

    def __init__(self, datastore, variable_name):
        self.datastore = datastore
        self.group = datastore._group
        self.variable_name = variable_name

        array = self.get_array()
        self.dtype = array.dtype
        self.shape = array.shape

    def _getitem(self, key):
        # read the data and put it into dict
        key = tuple(list(k) if isinstance(k, np.ndarray) else k for k in key)
        array = self.get_array()
        return array[key]

    def __getitem__(self, key):
        return indexing.explicit_indexing_adapter(
            key,
            self.shape,
            indexing.IndexingSupport.BASIC,
            self._getitem,
        )

    def get_array(self):
        ds = self.datastore._acquire()
        ds.get_sweep(self.group)
        return ds.data[self.group][self.variable_name]


class DataMetStore(AbstractDataStore):
    """Store for reading DataMet sweeps."""

    def __init__(self, manager, group=None):
        self._manager = manager
        self._group = int(group[6:])
        self._filename = self.filename
        self._need_time_recalc = False

    @classmethod
    def open(cls, filename, mode="r", group=None, **kwargs):
        manager = CachingFileManager(DataMetFile, filename, kwargs=kwargs)
        return cls(manager, group=group)

    @property
    def filename(self):
        with self._manager.acquire_context(False) as root:
            return root.filename

    @property
    def root(self):
        with self._manager.acquire_context(False) as root:
            return root

    def _acquire(self):
        # Does it need a lock as other backends ?
        ds = self.root
        return ds

    @property
    def ds(self):
        return self._acquire()

    def open_store_variable(self, mom):
        dim = self.root.first_dimension

        data = indexing.LazilyOuterIndexedArray(DataMetArrayWrapper(self, mom))
        encoding = {"group": self._group, "source": self._filename}

        mom_metadata = self.root.get_mom_metadata(mom, self._group)
        add_offset = float(mom_metadata.get("offset") or 0.0)
        scale_factor = float(mom_metadata.get("slope") or 1.0)
        maxval = mom_metadata.get("maxval", None)

        if maxval is not None:
            maxval = float(maxval)
            top = 255
            bottom = float(mom_metadata["bottom"])
            scale_factor = (maxval + add_offset) / (top - bottom)

        mname = datamet_mapping.get(mom, mom)
        mapping = sweep_vars_mapping.get(mname, {})
        attrs = {key: mapping[key] for key in moment_attrs if key in mapping}
        attrs["add_offset"] = add_offset
        attrs["scale_factor"] = scale_factor
        attrs["_FillValue"] = 0
        attrs["coordinates"] = (
            "elevation azimuth range latitude longitude altitude time"
        )
        return {mname: Variable((dim, "range"), data, attrs, encoding)}

    def open_store_coordinates(self, mom):
        scan_mdata = self.root.scan_metadata
        mom_mdata = self.root.get_mom_metadata(mom, self._group)
        dim = self.root.first_dimension

        # Radar coordinates
        lat = scan_mdata["orig_lat"]
        lon = scan_mdata["orig_lon"]
        alt = scan_mdata["orig_alt"]

        rng = (
            mom_mdata["Rangeoff"]
            + np.arange(mom_mdata["ncols"]) * mom_mdata["Rangeres"]
        )
        azimuth = (
            mom_mdata["Azoff"] + np.arange(mom_mdata["nlines"]) * mom_mdata["Azres"]
        )
        # Unravel azimuth
        azimuth[azimuth > 360] -= 360
        elevation = mom_mdata["Eloff"] + np.arange(mom_mdata["nlines"]) * 0

        sweep_mode = "azimuth_surveillance" if dim == "azimuth" else "rhi"
        sweep_number = self._group
        prt_mode = "not_set"
        follow_mode = "not_set"

        raytime = 0  # TODO find out raytime
        raytimes = np.array(
            [
                timedelta(seconds=x * raytime).total_seconds()
                for x in range(azimuth.shape[0] + 1)
            ]
        )
        diff = np.diff(raytimes) / 2.0
        rtime = raytimes[:-1] + diff

        timestr = scan_mdata["dt_acq"]
        time = datetime.strptime(timestr, "%Y-%m-%d-%H%M")

        time_attrs = get_time_attrs(f"{time.isoformat()}Z")

        encoding = {"group": self._group}
        rng = Variable(("range",), rng, get_range_attrs())
        azimuth = Variable((dim,), azimuth, get_azimuth_attrs(), encoding)
        elevation = Variable((dim,), elevation, get_elevation_attrs(), encoding)
        time = Variable((dim,), rtime, time_attrs, encoding)

        coords = {
            "azimuth": azimuth,
            "elevation": elevation,
            "range": rng,
            "time": time,
            "sweep_mode": Variable((), sweep_mode),
            "sweep_number": Variable((), sweep_number),
            "prt_mode": Variable((), prt_mode),
            "follow_mode": Variable((), follow_mode),
            "sweep_fixed_angle": Variable((), float(elevation[0])),
            "longitude": Variable((), lon, get_longitude_attrs()),
            "latitude": Variable((), lat, get_latitude_attrs()),
            "altitude": Variable((), alt, get_altitude_attrs()),
        }

        return coords

    def get_variables(self):
        return FrozenDict(
            (k1, v1)
            for k1, v1 in {
                **{
                    k: v
                    for list_item in [
                        self.open_store_variable(k) for k in self.ds.moments
                    ]
                    for (k, v) in list_item.items()
                },
                **self.open_store_coordinates(self.ds.moments[0]),
            }.items()
        )

    def get_attrs(self):
        return FrozenDict()


class DataMetBackendEntrypoint(BackendEntrypoint):
    """Xarray BackendEntrypoint for DataMet data."""

    description = "Open DataMet files in Xarray"
    url = "https://xradar.rtfd.io/latest/io.html#datamet-data-i-o"

    def open_dataset(
        self,
        filename_or_obj,
        *,
        mask_and_scale=True,
        decode_times=True,
        concat_characters=True,
        decode_coords=True,
        drop_variables=None,
        use_cftime=None,
        decode_timedelta=None,
        group=None,
        reindex_angle=False,
        first_dim="auto",
        site_coords=True,
    ):
        store = DataMetStore.open(
            filename_or_obj,
            group=group,
        )

        store_entrypoint = StoreBackendEntrypoint()

        ds = store_entrypoint.open_dataset(
            store,
            mask_and_scale=mask_and_scale,
            decode_times=decode_times,
            concat_characters=concat_characters,
            decode_coords=decode_coords,
            drop_variables=drop_variables,
            use_cftime=use_cftime,
            decode_timedelta=decode_timedelta,
        )

        # reassign azimuth/elevation/time coordinates
        ds = ds.assign_coords({"azimuth": ds.azimuth})
        ds = ds.assign_coords({"elevation": ds.elevation})
        ds = ds.assign_coords({"time": ds.time})

        ds.encoding["engine"] = "datamet"

        # handle duplicates and reindex
        if decode_coords and reindex_angle is not False:
            ds = ds.pipe(util.remove_duplicate_rays)
            ds = ds.pipe(util.reindex_angle, **reindex_angle)
            ds = ds.pipe(util.ipol_time, **reindex_angle)

        # handling first dimension
        dim0 = "elevation" if ds.sweep_mode.load() == "rhi" else "azimuth"
        if first_dim == "auto":
            if "time" in ds.dims:
                ds = ds.swap_dims({"time": dim0})
            ds = ds.sortby(dim0)
        else:
            if "time" not in ds.dims:
                ds = ds.swap_dims({dim0: "time"})
            ds = ds.sortby("time")

        # assign geo-coords
        if site_coords:
            ds = ds.assign_coords(
                {
                    "latitude": ds.latitude,
                    "longitude": ds.longitude,
                    "altitude": ds.altitude,
                }
            )

        return ds


def open_datamet_datatree(filename_or_obj, **kwargs):
    """Open DataMet dataset as :py:class:`datatree.DataTree`.

    Parameters
    ----------
    filename_or_obj : str, Path, file-like or DataStore
        Strings and Path objects are interpreted as a path to a local or remote
        radar file

    Keyword Arguments
    -----------------
    sweep : int, list of int, optional
        Sweep number(s) to extract, default to first sweep. If None, all sweeps are
        extracted into a list.
    first_dim : str
        Can be ``time`` or ``auto`` first dimension. If set to ``auto``,
        first dimension will be either ``azimuth`` or ``elevation`` depending on
        type of sweep. Defaults to ``auto``.
    reindex_angle : bool or dict
        Defaults to False, no reindexing. Given dict should contain the kwargs to
        reindex_angle. Only invoked if `decode_coord=True`.
    fix_second_angle : bool
        If True, fixes erroneous second angle data. Defaults to ``False``.
    site_coords : bool
        Attach radar site-coordinates to Dataset, defaults to ``True``.
    kwargs : dict
        Additional kwargs are fed to :py:func:`xarray.open_dataset`.

    Returns
    -------
    dtree: datatree.DataTree
        DataTree
    """
    # handle kwargs, extract first_dim
    backend_kwargs = kwargs.pop("backend_kwargs", {})
    # first_dim = backend_kwargs.pop("first_dim", None)
    kwargs["backend_kwargs"] = backend_kwargs

    sweep = kwargs.pop("sweep", None)
    sweeps = []
    kwargs["backend_kwargs"] = backend_kwargs

    if isinstance(sweep, str):
        sweeps = [sweep]
    elif isinstance(sweep, int):
        sweeps = [f"sweep_{sweep}"]
    elif isinstance(sweep, list):
        if isinstance(sweep[0], int):
            sweeps = [f"sweep_{i}" for i in sweep]
        else:
            sweeps.extend(sweep)
    else:
        # Get number of sweeps from data
        dmet = DataMetFile(filename_or_obj)
        sweeps = [
            f"sweep_{i}" for i in range(0, dmet.scan_metadata["elevation_number"])
        ]

    ds = [
        xr.open_dataset(
            filename_or_obj, group=swp, engine=DataMetBackendEntrypoint, **kwargs
        )
        for swp in sweeps
    ]

    ds.insert(0, xr.Dataset())

    # create datatree root node with required data
    dtree = DataTree(data=_assign_root(ds), name="root")
    # return datatree with attached sweep child nodes
    return _attach_sweep_groups(dtree, ds[1:])
