# AUTOGENERATED! DO NOT EDIT! File to edit: 03_l2data.ipynb (unless otherwise specified).

__all__ = ['DIVINER_URL', 'root', 'L2DataManager', 'get_l2_image_paths', 'L2Data', 'LocalTime']

# Cell
import warnings
from pathlib import Path

import numpy as np
import pvl
from yarl import URL

from planetarypy import geotools as gt
from planetarypy.utils import url_retrieve

DIVINER_URL = URL(
    "https://pds-geosciences.wustl.edu/lro/lro-l-dlre-4-rdr-v1/lrodlr_1001/data"
)

root = Path("/luna4/maye/l2_data")

# Cell
class L2DataManager:
    "Small helper to get lists of locally stored data."

    @property
    def labels(self):
        return sorted(list(root.glob("dgdr_*.lbl")))

    @property
    def images(self):
        return sorted(list(root.glob("dgdr_*.tif")))


def get_l2_image_paths():
    return L2DataManager().images

# Cell
class L2Data:
    GDR_L2_URL = DIVINER_URL / "gdr_l2"

    def __init__(
        self,
        cycle=None,
        datatype="ltim",
        map_res=1,  # pix per degrees
        projection="cylindrical",
        format="jp2",
        year=None,
    ):
        self.cycle = cycle
        self.datatype = datatype
        self.map_res = map_res
        self.projection = projection
        self.format = format
        self.year = year

        if cycle is None and year is None:
            warnings.warn("Set `year` for getting a correct folder URL.")

    @property
    def cycle(self):
        return self._cycle

    @cycle.setter
    def cycle(self, value):
        self._cycle = value

    @property
    def year(self):
        if self.cycle is not None:
            return str(self.cycle)[:4]
        else:
            return self._year

    @year.setter
    def year(self, value):
        self._year = str(value)

    @property
    def map_res(self):
        return self._map_res

    @map_res.setter
    def map_res(self, value):
        self._map_res = str(value).zfill(3)

    @property
    def datatype(self):
        return self._datatype

    @datatype.setter
    def datatype(self, value):
        """Set datatype string.

        This is for a Diviner GDR L3 datatype.

        Parameters
        ----------
        value : {"jd", "ltim", "tb3", "tbol"}
            RA = Rock abundance
            RMS = RMS error for RA
            ST = Regolith temperature
            TBOL = Average Bolometric Temperature
        """
        allowed = "jd ltim".split()
        allowed += [f"tb{i}" for i in range(3, 10)]
        allowed += [f"vb{i}" for i in range(1, 3)]
        if not value.lower() in allowed:
            raise ValueError(f"Only {allowed} allowed.")
        else:
            self._datatype = value.lower()

    @property
    def second_token(self):
        # TODO: deal with "err, cnt"
        return "avg"

    @property
    def fname(self):
        "Construct L3 GDR data filename."

        res = 128
        if self.datatype == "tbol":
            res = self.map_res
        return f"dgdr_{self.datatype}_{self.second_token}_cyl_{self.cycle}n_{res}_{self.format}.{self.format}"

    @property
    def label(self):
        return str(Path(self.fname).with_suffix(".lbl"))

    @property
    def folder_url(self):
        return self.GDR_L2_URL / self.year / self.projection / self.format

    @property
    def data_url(self):
        return self.folder_url / self.fname

    @property
    def label_url(self):
        return self.folder_url / self.label

    def download_label(self, subfolder=""):
        if not subfolder:
            p = root
        else:
            p = Path(subfolder)
        p.mkdir(exist_ok=True)
        url_retrieve(self.label_url, p / self.label)

    def download_data(self, subfolder="", overwrite=False):
        if not subfolder:
            p = root
        else:
            p = Path(subfolder)
        p.mkdir(exist_ok=True)
        savepath = p / self.fname
        if savepath.exists() and not overwrite:
            print("File exists, use `overwrite=True` to force download.")
            return
        else:
            url_retrieve(self.data_url, p / self.fname)

# Cell
class LocalTime:
    root = Path("/luna4/maye/l2_data")
    name = "lt"

    @classmethod
    def from_fpath(cls, fpath):
        cycle = int(fpath.name.split("_")[4][:-1])
        return cls(cycle)

    def __init__(self, cycle):
        self.cycle = cycle
        self.l2data = L2Data(cycle=cycle)
        self.img = gt.ImgData(str(self.fname))

    def read_window(self, ul_lon=0, ul_lat=1, width_degrees=1):
        ul = gt.Point.copy_geodata(self.img.center, lon=ul_lon, lat=ul_lat)
        ul.lonlat_to_pixel()
        lr = gt.Point.copy_geodata(
            ul, lon=ul_lon + width_degrees, lat=ul_lat - width_degrees
        )
        lr.lonlat_to_pixel()
        win = gt.Window(ulPoint=ul, lrPoint=lr)
        self.img.read_window(win)

    @property
    def label(self):
        return pvl.load(root / self.l2data.label)

    @property
    def fname(self):
        return root / self.l2data.fname

    @property
    def SCALING_FACTOR(self):
        return self.label["UNCOMPRESSED_FILE"]["IMAGE"]["SCALING_FACTOR"]

    @property
    def OFFSET(self):
        return self.label["UNCOMPRESSED_FILE"]["IMAGE"]["OFFSET"]

    @property
    def NODATA(self):
        return self.label["UNCOMPRESSED_FILE"]["IMAGE"]["MISSING_CONSTANT"]

    @property
    def data(self):
        data = self.img.data.astype("float")
        data[data == self.NODATA] = np.nan
        return data

    @property
    def scaled_data(self):
        return self.data * self.SCALING_FACTOR + self.OFFSET

    def get_pixel(self, xoff, yoff):
        value = np.squeeze(self.img.ds.ReadAsArray(xoff, yoff, 1, 1))
        if value == self.NODATA:
            return np.nan
        else:
            return value * self.SCALING_FACTOR + self.OFFSET

    def plot_window(self):
        plt.figure()
        plt.imshow(self.scaled_data, cmap="plasma")
        plt.colorbar()

    @property
    def window_mean(self):
        return np.nanmean(self.scaled_data)

    @property
    def window_std(self):
        return np.nanstd(self.scaled_data)

    @property
    def n_valid(self):
        return np.count_nonzero(~np.isnan(self.data))