import numpy as np
from numpy.linalg import matrix_power
from dataframe import DataFrame
from rl import *
import pandas as pd
from lib import Lib
import contextily as cx
from matplotlib import pyplot as plt
import geopandas as gpd
from shapely.geometry import Point
import xarray as xr

class GIS:
    """
    GIS class
    """
    def __init__(self):
        self.data_layers = {}
        self.fig, self.ax = plt.subplots(figsize=(17,17))

    def add_data_layer(self, layer_path, layer_name, data_source='sf'):
        if data_source == 'df':
            self.data_layers[layer_name] = gpd.GeoDataFrame(layer_path)
        elif data_source == 'sf':
            self.data_layers[layer_name] = gpd.read_file(layer_path)
        
    def get_data_layer(self, layer_name):
        return self.data_layers.get(layer_name)
    
    def get_shape(self, layer_name):
        return self.data_layers[layer_name].shape
    
    def add_random_series_column(self, layer_name, column_name='random',min=0, max=100, distrubution_type='random', mean=0, sd=1):
        if distrubution_type == 'random':
            series = pd.Series(np.random.randint(min, max, self.get_shape(layer_name)[0]))
        elif distrubution_type == 'standard_normal':
            series = pd.Series(np.random.standard_normal(self.get_shape(layer_name)[0]))
        elif distrubution_type == 'normal':
            series = pd.Series(np.random.normal(mean, sd, self.get_shape(layer_name)[0]))
        else:
            series = pd.Series(np.random.randn(self.get_shape(layer_name)[0]))
        self.add_column(layer_name, series, column_name)
    
    def join_layer(self, layer_name, geo_datframe_to_join, on):
        self.data_layers[layer_name] = self.data_layers.get(layer_name).merge(geo_datframe_to_join, on=on)
        
    def plot(self, layer_name, column4color=None, color=None, alpha=0.5, legend=False, 
             figsize_tuple=(15,10), cmap=None, ):
        """_summary_

        Args:
            layer_name (_type_): _description_
            column4color (_type_, optional): _description_. Defaults to None.
            color (_type_, optional): _description_. Defaults to None.
            alpha (float, optional): _description_. Defaults to 0.5.
            legend (bool, optional): _description_. Defaults to False.
            figsize_tuple (tuple, optional): _description_. Defaults to (15,10).
            cmap (str, optional): exmaple: 'Reds' for heatmaps. Defaults to None.
        """
        layer = self.data_layers.get(layer_name).to_crs(epsg=3857)
        layer.plot(ax=self.ax, alpha=alpha, edgecolor='k', color=color, legend=legend, 
                   column=column4color, figsize=figsize_tuple, cmap=cmap)
        cx.add_basemap(ax=self.ax, source=cx.providers.Esri.WorldImagery)
        
    def show(self, layer_name=None, interactive_mode=False):
        if interactive_mode is True: 
            return self.data_layers.get(layer_name).explore()
        else:
            self.ax.set_aspect('equal')
            plt.show()
        
    def get_crs(self, layer_name):
        """
        Cordonate Reference System
        EPSG: european petroleum survey group
        """
        return self.get_data_layer(layer_name).crs
    
    def reorder_columns(self, layer_name, new_order_as_list):
        self.data_layers[layer_name].reindex_axis(new_order_as_list, axis=1)
    
    def export(self, layer_name, file_name, file_format='geojson'):
        if file_format == 'geojson':
            self.data_layers[layer_name].to_file(file_name + '.geojson', driver='GeoJSON')
        elif file_format == 'shapefile':
            self.data_layers[layer_name].to_file(file_name + '.shp')
            
    def to_crs(self, layer_name, epsg="3857"):
        self.data = self.data_layers[layer_name].to_crs(epsg)
        
    def set_crs(self, layer_name, epsg="3857"):
        self.data = self.data_layers[layer_name].set_crs(epsg)
        
    def show_points(self, x_y_csv_path, crs="3857"):
        pass
    
    def show_point(self, x_y_tuple, crs="3857"):
        pass
    
    def add_point(self, x_y_tuple, layer_name, crs="3857"):
        point = Point(0.0, 0.0)
        #self.__dataframe = self.get_dataframe().append(row_as_dict, ignore_index=True)
        row_as_dict = {'geometry': point}
        self.data_layers[layer_name].append(row_as_dict, ignore_index=True)
    
    def new_data_layer(self, layer_name, crs="EPSG:3857"):
        self.data_layers[layer_name] = gpd.GeoDataFrame(crs=crs)
        self.data_layers[layer_name].crs = crs
        
    def add_column(self, layer_name, column, column_name):
        y = column
        if (not isinstance(column, pd.core.series.Series or not isinstance(column, pd.core.frame.DataFrame))):
            y = np.array(column)
            y = np.reshape(y, (y.shape[0],))
            y = pd.Series(y)
        self.data_layers[layer_name][column_name] = y
        
    def show_data_layer(self, layer_name, number_of_row=None):
        if number_of_row is None:
            print(self.get_data_layer(layer_name))
        elif number_of_row < 0:
            return self.get_data_layer(layer_name).tail(abs(number_of_row)) 
        else:
            return self.get_data_layer(layer_name).head(number_of_row) 
        
    def add_row(self, layer_name, row_as_dict):
        self.data_layers[layer_name] = self.get_data_layer(layer_name).append(row_as_dict, ignore_index=True)
    
    def get_row(self, layer_name, row_index, column=None):
        if column is not None:
            return self.data_layers[layer_name].loc[self.data_layers[layer_name][column] == row_index].reset_index(drop=True)
        return self.data_layers[layer_name].iloc[row_index]
    
    def get_layer_shape(self, layer_name):
        """
        return (Number of lines, number of columns)
        """
        return self.data_layers[layer_name].shape
    
    def get_columns_names(self, layer_name):
        header = list(self.data_layers[layer_name].columns)
        return header 
    
    def drop_column(self, layer_name, column_name):
        """Drop a given column from the dataframe given its name

        Args:
            column (str): name of the column to drop

        Returns:
            [dataframe]: the dataframe with the column dropped
        """
        self.data_layers[layer_name] = self.data_layers[layer_name].drop(column_name, axis=1)
        return self.data_layers[layer_name]
    
    def keep_columns(self, layer_name, columns_names_as_list):
        for p in self.get_columns_names(layer_name):
            if p not in columns_names_as_list:
                self.data_layers[layer_name] = self.data_layers[layer_name].drop(p, axis=1)
                
    def get_area_column(self, layer_name):
        return self.get_data_layer(layer_name).area
    
    def get_perimeter_column(self, layer_name):
        return self.get_data_layer(layer_name).length
    
    def get_row_area(self, layer_name, row_index):
        return self.data_layers[layer_name].area.iloc[row_index]
    
    def get_distance(self, layer_name, index_column, row_index_a, row_index_b):
        if 1 == 1:
            other = self.get_row(layer_name, row_index_b, index_column)
            return self.get_row(layer_name, row_index_a, index_column).distance(other)
    
    def filter_dataframe(self, layer_name, column, func_de_decision, in_place=True, *args):
        if in_place is True:
            if len(args) == 2:
                self.set_dataframe(
                    self.data_layers[layer_name].loc[self.get_column(column).apply(func_de_decision, args=(args[0], args[1]))])
            else:
                self.set_dataframe(
                    self.data_layers[layer_name].loc[self.get_column(column).apply(func_de_decision)])
        else:
            if len(args) == 2:
                return self.data_layers[layer_name].loc[self.get_column(column).apply(func_de_decision, args=(args[0], args[1]))]
            else:
                return self.data_layers[layer_name].loc[self.get_column(column).apply(func_de_decision)]

    def transform_column(self, layer_name, column_to_trsform, column_src, fun_de_trasformation, in_place= True,*args):
        if in_place is True:
            if (len(args) != 0):
                self.set_column(layer_name, column_to_trsform, self.get_column(layer_name, column_src).apply(fun_de_trasformation, args=(args[0],)))
            else:
                self.set_column(layer_name, column_to_trsform, self.get_column(layer_name, column_src).apply(fun_de_trasformation)) 
        else:
            if (len(args) != 0):
                return self.get_column(layer_name, column_src).apply(fun_de_trasformation, args=(args[0],))
            else:
                return self.get_column(layer_name, column_src).apply(fun_de_trasformation)
            
    def set_column(self, layer_name, column_name, new_column):
        self.data_layers[layer_name][column_name] = new_column
    
    def get_column(self, layer_name, column):
        return self.data_layers[layer_name][column]
    
    def reindex_dataframe(self, layer_name, index_as_liste=None, index_as_column_name=None):
        if index_as_liste is not None:
            new_index = new_index = index_as_liste
            self.data_layers[layer_name].index = new_index
        if index_as_column_name is not None:
            self.data_layers[layer_name].set_index(index_as_column_name, inplace=True)
        if index_as_column_name is None and index_as_liste is None:
            new_index = pd.Series(np.arange(self.get_shape()[0]))
            self.data_layers[layer_name].index = new_index
            
    def get_era5_land_grib_as_dataframe(self, file_path, layer_name):
        grip_path = file_path
        ds = xr.load_dataset(grip_path, engine="cfgrib")
        self.data_layers[layer_name] = DataFrame()
        self.data_layers[layer_name].set_dataframe(ds.to_dataframe())
        return ds.to_dataframe()
    
    def rename_columns(self, layer_name, column_dict_or_all_list, all_columns=False):
        if all_columns is True:
            types = {}
            self.data_layers[layer_name].columns = column_dict_or_all_list
            for p in column_dict_or_all_list:
                types[p] = str
            self.data_layers[layer_name] = self.get_dataframe().astype(types)
        else:
            self.data_layers[layer_name].rename(columns=column_dict_or_all_list, inplace=True)
            
    def calculate_area_as_column(self, layer_name):
        self.add_column(layer_name, self.get_area_column(layer_name), 'area')
        
    def calculate_perimeter_as_column(self, layer_name):
        self.add_column(layer_name, self.get_perimeter_column(layer_name), 'perimeter')
        
    def count_occurence_of_each_row(self, layer_name, column_name):
        return self.data_layers[layer_name].pivot_table(index=[column_name], aggfunc='size')
    
    def add_transformed_columns(self, layer_name, dest_column_name="new_column", transformation_rule="okk*2"):
        columns_names = self.get_columns_names(layer_name)
        columns_dict = {}
        for column_name in columns_names:
            if column_name in transformation_rule:
                columns_dict.update({column_name: self.get_column(layer_name, column_name)})
        y_transformed = eval(transformation_rule, columns_dict)
        self.data_layers[layer_name][dest_column_name] = y_transformed
         
    
    @staticmethod
    def new_geodaraframe_from_points():
        map.new_data_layer('valves', crs="ESRI:102191")
        for p in range(map.get_layer_shape('pipelines')[0]):
            #print(map.get_row('pipelines', p))
            vi = Point(map.get_row('pipelines', p)['X_Start'], map.get_row('pipelines', p)['Y_Start'])
            vf = Point(map.get_row('pipelines', p)['X_End'], map.get_row('pipelines', p)['Y_End'])
            id_pipeline = map.get_row('pipelines', p)['Nom_CANAL']
            row_as_dict = {'id_pipeline': id_pipeline,
                        'geometry': vi}
            map.add_row('valves', row_as_dict)
            row_as_dict = {'id_pipeline': id_pipeline,
                        'geometry': vf}
            map.add_row('valves', row_as_dict)