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

__all__ = ['calculate_x_offset', 'get_seviri_area_def', 'load_scene', 'construct_area_def', 'construct_TM_area_def',
           'reproj_with_manual_grid', 'reproj_to_xarray', 'full_scene_pyresample', 'full_scene_pyinterp', 'Reprojector']

# Cell
import pandas as pd
import xarray as xr
import numpy as np
import numpy.ma as ma

import matplotlib as mpl
import matplotlib.pyplot as plt
from matplotlib import colors
import seaborn as sns

import os
import time
from itertools import product
from collections import OrderedDict
from datetime import datetime
from ipypb import track
import FEAutils as hlp

import satpy
from satpy import Scene
from pyresample.geometry import AreaDefinition
import geopandas as gpd
from shapely.geometry import Point
import cartopy.crs as ccrs
import pyinterp
import pyinterp.backends.xarray

import rasterio
from rasterio import Affine as A
from rasterio.warp import reproject, Resampling, calculate_default_transform, transform
from rasterio.control import GroundControlPoint
from rasterio.transform import xy

# Cell
def calculate_x_offset(native_fp):
    handler = satpy.readers.seviri_l1b_native.NativeMSGFileHandler(native_fp, {}, None)
    lower_east_column_planned = handler.header['15_DATA_HEADER']['ImageDescription']['PlannedCoverageHRV']['LowerEastColumnPlanned']
    x_offset = 32500 + ((2733 - lower_east_column_planned) * 1000)

    return x_offset

def get_seviri_area_def(native_fp, num_x_pixels=5568, num_y_pixels=4176) -> AreaDefinition:
    """
    The HRV channel on Meteosat Second Generation satellites doesn't scan the full number of columns.
    The east boundary of the HRV channel changes (e.g. to maximise the amount of the image which
    is illuminated by sunlight.

    Parameters:
        native_fp: Data filepath


    """

    x_offset = calculate_x_offset(native_fp)

    # The EUMETSAT docs say "The distance between spacecraft and centre of earth is 42,164 km. The idealized earth
    # is a perfect ellipsoid with an equator radius of 6378.1690 km and a polar radius of 6356.5838 km."
    # The projection used by SatPy expresses height as height above the Earth's surface (not distance
    # to the centre of the Earth).

    projection = {
        'proj': 'geos',
        'lon_0': 9.5,
        'a': 6378169.0,
        'b': 6356583.8,
        'h': 35785831.00,
        'units': 'm'}

    seviri = AreaDefinition(
        area_id='seviri',
        description='SEVIRI RSS HRV',
        proj_id='seviri',
        projection=projection,
        width=num_x_pixels,
        height=num_y_pixels,
        area_extent=[
            -2768872.0236 + x_offset, # left
             1394687.3495,            # bottom (from scene['HRV'].area)
             2799876.1893 + x_offset, # right
             5570248.4773]            # top (from scene['HRV'].area)
    )

    return seviri

# Cell
def load_scene(native_fp):
    # Reading scene and loading HRV
    scene = Scene(filenames=[native_fp], reader='seviri_l1b_native')

    # Identifying and recording lower_east_column_planned
    handler = satpy.readers.seviri_l1b_native.NativeMSGFileHandler(native_fp, {}, None)
    scene.attrs['lower_east_column_planned'] = handler.header['15_DATA_HEADER']['ImageDescription']['PlannedCoverageHRV']['LowerEastColumnPlanned']

    return scene

# Cell
def construct_area_def(scene, area_id, description,
                       proj_id, projection,
                       west, south, east, north,
                       pixel_size=None):

    # If None then will use same number of x and y points
    # HRV's resolution will be more like 4km for Europe
    if meters_per_pixel is not None:
        width = int((east - west) / pixel_size)
        height = int((north - south) / pixel_size)
    else:
        width = scene[list(scene.keys())[0]['name']].x.values.shape[0]
        height = scene[list(scene.keys())[0]['name']].y.values.shape[0]

    area_extent = (west, south, east, north)

    area_def = AreaDefinition(area_id, description,
                              proj_id, projection,
                              width, height, area_extent)

    return area_def

def construct_TM_area_def():
    meters_per_pixel = 4000
    west, south, east, north = (-3090000, 1690000, 4390000, 9014000)

    area_id = 'TM'
    description = 'Transverse Mercator'
    proj_id = 'TM'

    projection = {
        'ellps': 'WGS84',
        'proj': 'tmerc',  # Transverse Mercator
        'units': 'm'  # meters
    }

    tm_area_def = construct_area_def(scene, area_id, description,
                                     proj_id, projection,
                                     west, south, east, north,
                                     meters_per_pixel)

    return tm_area_def

# Cell
def reproj_with_manual_grid(da, new_x_positions, new_y_positions):

    x_axis = pyinterp.Axis(da.x.values)
    y_axis = pyinterp.Axis(da.y.values)

    grid = pyinterp.Grid2D(x_axis, y_axis, da.data.T)

    reproj_data = (pyinterp
                   .bivariate(grid, new_x_positions, new_y_positions)
                   .reshape((da_resampled.x.size, da_resampled.y.size))
                  )

    return reproj_data

def reproj_to_xarray(da, x_coords, y_coords):
    # We'll reproject the data
    reproj_data = reproj_with_manual_grid(da, x_coords, y_coords)

    # Then put it in an XArray DataArray
    da_reproj = xr.DataArray(np.flip(reproj_data.T, axis=(0, 1)),
                             dims=('y', 'x'),
                             coords={
                                 'x': da_resampled.x.values[::-1],
                                 'y': da_resampled.y.values[::-1]
                             },
                             attrs=da_resampled.attrs)

    return da_reproj

# Cell
def full_scene_pyresample(native_fp):
    # Loading scene
    scene = load_scene(native_fp)
    dataset_names = scene.all_dataset_names()
    scene.load(dataset_names)

    # Correcting HRV area definition
    for dataset_name in dataset_names:
        num_y_pixels, num_x_pixels = scene[dataset_name].shape
        seviri_area_def = get_seviri_area_def(native_fp, num_x_pixels=num_x_pixels, num_y_pixels=num_y_pixels)
        scene[dataset_name].attrs['area'] = seviri_area_def

    # Resampling
    tm_area_def = construct_TM_area_def()
    ds_reproj = (scene
                 .resample(tm_area_def, resampler='nearest')
                 .to_xarray_dataset()
                )

    return ds_reproj

def full_scene_pyinterp(native_fp, new_x_coords, new_y_coords):
    # Loading data
    scene = load_scene(native_fp)
    dataset_names = scene.all_dataset_names()
    scene.load(dataset_names)

    # Correcting x coordinates
    seviri_area_def = get_seviri_hrv_area_def(native_fp)
    area_extent = seviri_area_def.area_extent
    x_offset = calculate_x_offset(native_fp)

    width = scene['HRV'].x.size
    corrected_x_coords = np.linspace(area_extent[2], area_extent[0], width)
    scene['HRV'] = scene['HRV'].assign_coords({'x': corrected_x_coords})

    # Reprojecting
    dataset_name_to_da_reproj = dict()

    for dataset_name in dataset_names:
        da_reproj = reproj_to_xarray(scene[dataset_name], new_x_coords, new_y_coords)
        dataset_name_to_da_reproj[dataset_name] = da_reproj

    ds_reproj = xr.Dataset(dataset_name_to_da_reproj)

    return ds_reproj

class Reprojector:
    def __init__(self, new_coords_fp):
        df_new_coords = pd.read_csv(new_coords_fp)

        self.new_x_coords = df_new_coords['x']
        self.new_y_coords = df_new_coords['y']

        return

    def reproject(self, native_fp, reproj_library='pyinterp'):
        if reproj_library == 'pyinterp':
            ds_reproj = full_scene_pyinterp(native_fp, self.new_x_coords, self.new_y_coords)
        elif reproj_library == 'pyresample':
            ds_reproj = full_scene_pyresample(native_fp)
        else:
            raise ValueError(f'`reproj_library` must be one of: pyresample, pyinterp. {reproj_library} can not be passed.')

        return ds_reproj