
import os

import xarray as xr
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sn
import pandas as pd
import matplotlib as mpl
import logging

logger = logging.getLogger(__name__)

try:
    from rasterio import features
    from affine import Affine
    import geopandas as gpd
except:
    pass

from pathlib import Path

import tinysunset as ts

def transform_from_latlon(lat, lon):
    """ input 1D array of lat / lon and output an Affine transformation
    """
    lat = np.asarray(lat)
    lon = np.asarray(lon)
    trans = Affine.translation(lon[0], lat[0])
    scale = Affine.scale(lon[1] - lon[0], lat[1] - lat[0])
    return trans * scale

def rasterize(shapes, coords, latitude='lat', longitude='lon',
              fill=np.nan, **kwargs):

    transform = transform_from_latlon(coords[latitude], coords[longitude])
    out_shape = (len(coords[latitude]), len(coords[longitude]))
    
    raster = features.rasterize(shapes, out_shape=out_shape,
                                fill=fill, transform=transform,
                                dtype=float, **kwargs)
    
    spatial_coords = {latitude: coords[latitude], longitude: coords[longitude]}
    
    return xr.DataArray(raster, coords=spatial_coords, dims=(latitude, longitude))


def _add_shape_coord_from_data_array(xr_da, shp_path):

    # 1. read in shapefile
    shp_gpd = gpd.read_file(shp_path)

    # 2. create a list of tuples (shapely.geometry, id)
    #    this allows for many different polygons within a .shp file (e.g. States of US)
    shapes = [(shape, n) for n, shape in enumerate(shp_gpd.geometry)]
    
    ds = xr.Dataset(coords=xr_da.coords)
    
    for shape in shapes:

        region_name = shp_gpd.loc[shape[1], "Acronym"]

        # add as new variable to the dataset
        da = rasterize([shape], xr_da.coords)
        ds[region_name] = da.rename(region_name)
        
    return ds

def _create_mask_from_shapefile(
    da: xr.DataArray,
    shp_path: str
    ) -> xr.Dataset:


    mask = _add_shape_coord_from_data_array(
        da.to_dataset(name="mask"), 
        shp_path
    )
    try:
        mask = mask.drop(["height"])
    except:
        pass
    try:
        mask = mask.drop(["time"])
    except:
        pass

    return mask

def apply_mask_to_fcst(
    ds: xr.Dataset, 
    regions_shp_path: str
    ) -> xr.Dataset:
    """This function applies a lat lon mask to a dataset. 
    Computes the mean of the region and stores the result in a new variable."""
    
    lonlat_res = f"{ds.lon.size}x{ds.lat.size}"
    mask_dir = Path(regions_shp_path).parent
    mask_filename = Path(regions_shp_path).stem
    nc_mask_filename = f"{mask_dir}/{mask_filename}_{lonlat_res}.nc"

    if not os.path.exists(nc_mask_filename):
    
        mask = _create_mask_from_shapefile(
            ds.isel(time=0,ensemble=0),
            regions_shp_path)#.isel(month=0)
        
        mask = mask.drop(["month"], errors="ignore")
        
        logger.info(f"Saving mask to {nc_mask_filename}")
        mask.to_netcdf(nc_mask_filename)

        
    else: 
        mask = xr.open_dataset(nc_mask_filename)

    new_ds = []
    for region in mask.data_vars:
        # apply mask to each region
        new_ds.append(ds.where(~np.isnan(mask[region])).mean(dim=["lat","lon"]))
        
    new_ds = xr.concat(new_ds, dim="region")
    new_ds = new_ds.assign_coords({"region": mask.data_vars})
            
    return new_ds    

def compute_scorecard(
    fcst: xr.DataArray,
    obs: xr.Dataset,
    verification_config,
    clim_fcst: xr.DataArray = None,
    ):
    
    SHPS_DIR = os.path.dirname(os.path.realpath(__file__)) + "/../../shapefiles/"
    IPCC_REGIONS_SHP_PATH = SHPS_DIR + "/IPCC-WGI-reference-regions-v4_shapefile/IPCC-WGI-reference-regions-v4.shp"
    
    regions_shp_path = IPCC_REGIONS_SHP_PATH
    
    fcst_agg = apply_mask_to_fcst(fcst, regions_shp_path)
    obs_agg = apply_mask_to_fcst(obs, regions_shp_path)

    if clim_fcst is not None:
        clim_fcst_agg = apply_mask_to_fcst(clim_fcst, regions_shp_path)
    else:
        clim_fcst_agg = None

    verification = ts.verification.Verification(
        obs = obs_agg,
        fcst = fcst_agg,
        verification_config = verification_config
        )
    
    del fcst_agg
    del obs_agg

    verification.compute(clim_fcst=clim_fcst_agg)

    scorecard_ds = verification.fcst_metrics
    del verification

    try:
    
        scorecard_ds = scorecard_ds.drop(["month"])
        scorecard_ds = scorecard_ds.isel(month=0)
        
    except:
        pass

    return scorecard_ds


def plot_scorecard(
    scorecard_ds: xr.Dataset,
    metrics_list = None,
    title = None,
    fileout: Path = None,
    heatmap_kws = None
    ):

    scorecard_ds = scorecard_ds.drop(["height", "observations_category_edge", "forecasts_category_edge"], errors="ignore")

    
    if metrics_list is not None:
        scorecard_ds = scorecard_ds[metrics_list]
    
    
    if "model" in scorecard_ds.dims:

        df_list = pd.DataFrame()      

        for model in scorecard_ds.model.values:

            _ds = scorecard_ds.sel(model=model)
            _ds = _ds.drop("model")

            df_list = pd.concat([df_list, _ds.to_dataframe().transpose()])
            
        
        num_models = scorecard_ds.dims["model"]
        models = scorecard_ds.model.values

    else:

        df = scorecard_ds.to_dataframe().transpose()    
        df_list = [df]
        del df
        num_models = 1
        models = [""]
        
        df_list = pd.concat(df_list)

    num_metrics = len(scorecard_ds.data_vars)
    num_regions = scorecard_ds.dims["region"]

    fig, ax = plt.subplots(1,1, figsize=(30,2*num_models))

    sn.heatmap(df_list, ax=ax, **heatmap_kws)
        
    ax.collections[0].colorbar.remove()
    ax.hlines([num_metrics*i for i in range(num_models+1)], 
              *ax.get_xlim(), linewidths=2, colors="black")

    for i, model in enumerate(models):

        model = models[-i-1]     
        
        plt.figtext(0.91, 0+i*0.18 + 0.2,
                    model, fontsize=20, rotation=0)

    ax.set_xticks(np.arange(num_regions)+0.5)
    ax.set_xticklabels(scorecard_ds.region.values, rotation=60, fontsize=18)
    
    if title is not None:
        plt.title(title, fontsize=25)
    
    if fileout is not None:
        plt.savefig(f"{fileout}")
        plt.close()