# AUTOGENERATED! DO NOT EDIT! File to edit: nbs/04_geo.ipynb (unless otherwise specified).

__all__ = ['dict2json', 'Region', 'RegionST']

# Cell
import rasterio
import numpy as np
import pandas as pd
import json

# Cell
def dict2json(data:dict, file):
    "Writes json file from dict"
    with open(file, 'w') as f:
        f.write(json.dumps(data))

class Region():
    """Defines a geographical region with a name, a bounding box and the pixel size"""
    def __init__(self, name:str, bbox:list, pixel_size:float):
        self.name = name
        self.bbox = rasterio.coords.BoundingBox(*bbox) # left, bottom, right, top
        self.pixel_size = pixel_size

    @property
    def width(self):
        "Width of the region"
        return int(np.round(np.abs(self.bbox.left-self.bbox.right)/self.pixel_size))

    @property
    def height(self):
        "Height of the region"
        return int(np.round(np.abs(self.bbox.top-self.bbox.bottom)/self.pixel_size))

    @property
    def transform(self):
        "Rasterio Affine transform of the region"
        return rasterio.transform.from_bounds(*self.bbox, self.width, self.height)

    @property
    def shape(self):
        "Shape of the region (height, width)"
        return (self.height, self.width)

    def coords(self, offset='ul'):
        "Computes longitude and latitude arrays given a shape and a rasterio Affine transform"
        rxy = rasterio.transform.xy
        ys, xs = map(range, self.shape)
        return (np.array(rxy(self.transform, [0]*len(xs), xs, offset=offset)[0]),
                np.array(rxy(self.transform, ys, [0]*len(ys), offset=offset)[1]))

    @classmethod
    def load(cls, file):
        "Loads region information from json file"
        with open(file, 'r') as f:
            args = json.load(f)
        return cls(args['name'], args['bbox'], args['pixel_size'])

    def export(self, file):
        """Exports region information to json file"""
        dict2json(self.__dict__, file)

    def __repr__(self):
        return '\n'.join([f'{i}: {o}' for i, o in self.__dict__.items()]) + '\n'


class RegionST(Region):
    "Defines a region in space and time with a name, a bounding box and the pixel size."
    def __init__(self, name:str, bbox:list, pixel_size:float, time_start:str=None,
                 time_end:str=None, time_freq:str='D', time_margin:int=0):
        self.name = name
        self.bbox = rasterio.coords.BoundingBox(*bbox) # left, bottom, right, top
        self.pixel_size = pixel_size
        self.time_start = pd.Timestamp(str(time_start))
        self.time_end = pd.Timestamp(str(time_end))
        self.time_margin = time_margin
        self.time_freq = time_freq

    @property
    def times(self):
        "Property that computes the date_range for the region."
        tstart = self.time_start - pd.Timedelta(days=self.time_margin)
        tend = self.time_end + pd.Timedelta(days=self.time_margin)
        return pd.date_range(tstart, tend, freq=self.time_freq)

    @classmethod
    def load(cls, file, time_start=None, time_end=None):
        "Loads region information from json file"
        with open(file, 'r') as f:
            args = json.load(f)
        if time_start is None:
            time_start = args['time_start']
        if time_end is None:
            time_end = args['time_end']
        return cls(args['name'], args['bbox'], args['pixel_size'],
                   time_start, time_end)