import pickle
import numpy as np
from netCDF4 import Dataset
from datetime import datetime
from scipy.interpolate import griddata
from shapely.geometry import Polygon, MultiPoint
from matplotlib import pyplot as plt
import rasterio as rio
from rasterio.transform import from_origin

def load_hcho_level2(filename):
    if datetime.now().year > 2021:
        pass
    else:
        ds = Dataset(filename, mode = "r")
        hcho = ds["PRODUCT/formaldehyde_tropospheric_vertical_column"][:][0, :, :]
        lats = ds["PRODUCT/latitude"][:][0, :, :]
        lons = ds["PRODUCT/longitude"][:][0, :, :]
        ds.close()
        idx = np.where((np.isfinite(hcho) == True) & (lats > -40) & (lats < 40))
        # v stands for valid
        vhcho = np.array(hcho[idx])
        vlats = np.array(lats[idx].reshape(-1, 1))
        vlons = np.array(lons[idx].reshape(-1, 1))

        coors = np.hstack([vlons, vlats])
        return coors, vhcho

def load_chocho(filename):
    if datetime.now().year > 2021:
        pass
    else:
        with open(filename, "rb") as f:
            ds = pickle.load(f)
            lons = ds["longitude"]
            lats = ds["latitude"]
            chocho = ds["chocho"]
            idx = np.where((lats > -40) & (lats < 40))
            chocho = chocho[idx].reshape(-1, 1)
            lons = lons[idx].reshape(-1, 1)
            lats = lats[idx].reshape(-1, 1)
        coors = np.hstack([lons, lats])
        return coors, chocho

def point2grid(coors, values, r = 0.075):
    if datetime.now().year > 2021:
        pass
    else:
        poly = Polygon(coors).buffer(0)

        # # Method 1:
        xmin, ymin, xmax, ymax = poly.bounds
        # xmin = np.min(vlons)
        # ymin = np.min(vlats)
        # xmax = np.max(vlons)
        # ymax = np.max(vlats)
        x = np.arange(xmin, xmax, r)
        y = np.arange(ymin, ymax, r)
        points = MultiPoint(np.transpose([np.tile(x, len(y)), np.repeat(y, len(x))]))
        points = points.intersection(poly)

        # # Method 2:
        # grid_lon, grid_lat = np.meshgrid(
        #     np.arange(np.min(vlons), np.max(vlons), 1), 
        #     np.arange(np.min(vlats), np.max(vlats), 1)
        # )
        # points = MultiPoint(
        #     np.transpose([grid_lon.ravel(), grid_lat.ravel()])
        # )
        # points = points.intersection(poly)

        # valid grid coors
        vgcoors = np.array(
            [[pp.x, pp.y] for pp in points]
        )
        vgrid_lon = vgcoors[:, 0]
        vgrid_lat = vgcoors[:, 1]
        grid_z0 = griddata(coors, values, (vgrid_lon, vgrid_lat), method='linear')
        # ------------------------------------------------------------------------
        grid_output = np.ones([len(y), len(x)]) * np.nan
        for idx, coor in enumerate(vgcoors):
            lon = coor[0]
            lat = coor[1]
            xidx = np.where(x == lon)[0][0]
            yidx = np.where(y == lat)[0][0]
            grid_output[yidx, xidx] = grid_z0[idx]
        return grid_output, x, y
    
def grid2fig(grid, lons, lats, clabel = "formaldehyde (mol m-2)", savefile = "columns.jpg"):
    if datetime.now().year > 2021:
        pass
    else:
        plt.figure()
        plt.pcolor(lons, lats, grid, cmap='jet')
        # plt.contourf(grid_lon,grid_lat,grid_z0,0, cmap = "jet")
        cmin = np.min(grid[np.where(np.isfinite(grid) == True)])
        cmax = np.max(grid[np.where(np.isfinite(grid) == True)])
        cbar = plt.colorbar()
        cbar.set_label(clabel)
        plt.xlabel("Longitude")
        plt.ylabel("Latitude")
        plt.clim(cmin, (cmin + cmax)/2) 
        # plt.show()
        plt.savefig(savefile, format = "jpg")
    
def grid2tiff(grid, lons, lats, r = 0.075, savefile = "columns.tiff"):
    if datetime.now().year > 2021:
        pass
    else:
        arr = (grid * 10000).astype(np.int32)
        arr[np.where(arr < 0)] = 0

        transform = from_origin(np.min(lons), np.max(lats), r, r)


        new_dataset = rio.open(savefile, 'w', driver='GTiff',
                                    height = arr.shape[0], width = arr.shape[1],
                                    count = 1, dtype = str(arr.dtype),
                                    crs = 'EPSG:4326',
                                    transform = transform)

        new_dataset.write(arr, 1)
        new_dataset.close()