'''
Created on 04-Jul-2021

@author: bthakur
'''

import inspect
import traceback
import os
from enum import Enum     
import re
import pandas as pd
import random
from xv.util import convertToUtf8

from xv.util import UtilClass

class DataProvider(UtilClass):
    '''
    Class to load GeoSpatial data
    '''
    
    _MAX_SEARCH_RECORDS = 10

    class Type(Enum):
        CONTINENT = "CONTINENT"
        COUNTRY = "COUNTRY"
        CITY = "CITY"
        RIVERS_AND_LAKES = "RIVERS_AND_LAKES"
    
    def __init__(self, **kwargs):
        '''
        '''
        super().__init__(**kwargs)
        self._world_df = None
        self._cities_df = None
        self._ny_df = None
        self._rivers_and_lakes_df = None
        self._elevation_points_df = None
        self._usa_states_df = None
        self._regions_polys_df = None
        self._geographic_lines_df = None
        self._random_names = None
        
        
        #self._superset_df = None
        
        pass


    def get_random_names(self, count = 6):
        if self._random_names is None:
            self._init_random_names()
        return random.sample(self._random_names, count)

    def get_random_continents(self, count = 4):
        dfs = [self.get_continents_df()]
        names = self._parse_names(dfs)
        
        if len(names) < count:
            return names
        return random.sample(names, count)
        pass

    def get_random_countries(self, count = 4):
        dfs = [self.get_countries_df()]
        names = self._parse_names(dfs)
        
        if len(names) < count:
            return names
        return random.sample(names, count)
        pass


    def get_random_cities(self, count = 4):
        dfs = [self.get_cities_df()]
        names = self._parse_names(dfs)
        
        if len(names) < count:
            return names
        return random.sample(names, count)
        pass


    
    def reset_superset_df(self):
        self._superset_df = None
        pass
    
    def _set_superset_df(self, superset_df):
        if superset_df is None:
            self._log(f"superset_df is None: {superset_df}")
            self._superset_df = None
            return
        
        if superset_df.shape[0] <= 0:
            self._log(f"The records in superset_df is zero: {superset_df.shape[0]}")
            self._superset_df = None
            return
        
        self._superset_df = superset_df
        pass

    def get_superset_df(self):
        if self._superset_df is None:
            return self.get_world_df()
        return self._superset_df
        pass

    

    def remove_disjoint_records(self, df, superset_df = None):
        try:
            if superset_df is None:
                superset_df = self.get_superset_df()
                
            if superset_df.shape[0] != 1:
                return df
            superset = superset_df.geometry.unary_union
            
            modified_df = df[~df.geometry.disjoint(superset)]
            return modified_df
        except Exception as e:
            sError = f"""Error {self.__class__.__name__}.{inspect.stack()[0][3]}:{e}
            Failed to remove disjointed records
            """
            self._log(sError)
            self._log(traceback.format_exc())
    
        return df
        pass



    def _restrict_records(self, df):
        if df is None:
            return df
        max_records = DataProvider._MAX_SEARCH_RECORDS
        if df.shape[0] <= max_records:
            return df
        self._log(f"Number of records restricted to {max_records}")
        return df.sample(n = max_records)
    
    
        
    def search_NOT_USED(self, searchStrings, superset = None):
        data = {}
        
        continents = self.search_continents(searchStrings, superset = superset)
        if not continents is None:
            data['continents'] = self._restrict_records(continents)
            
        countries = self.search_countries(searchStrings, superset = superset)
        if not countries is None:
            data['countries'] = self._restrict_records(countries)
            
        rivers_and_lakes = self.search_rivers_and_lakes(searchStrings, superset = superset)
        if not rivers_and_lakes is None:
            data['rivers_and_lakes'] = self._restrict_records(rivers_and_lakes)
            
        cities = self.search_cities(searchStrings, superset = superset)
        if not cities is None:
            data['cities'] = self._restrict_records(cities)
            
        return data
        pass


    def search(self, searchStrings, superset = None):
        data = {}
        
        continents = self.search_continents(searchStrings, superset = superset)
        if not continents is None:
            data['continents'] = continents
            
        countries = self.search_countries(searchStrings, superset = superset)
        if not countries is None:
            data['countries'] = countries
            
        rivers_and_lakes = self.search_rivers_and_lakes(searchStrings, superset = superset)
        if not rivers_and_lakes is None:
            data['rivers_and_lakes'] = rivers_and_lakes
            
        cities = self.search_cities(searchStrings, superset = superset)
        if not cities is None:
            data['cities'] = cities
            
        return data
        pass


    def search_continents(self, searchStrings, superset = None):
        continents = self.get_continents_df()
        found_df = self._search_df(continents, 
                                   column_name = 'name', 
                                   searchStrings = searchStrings, 
                                   superset = superset)
        return self._restrict_records(found_df)


    def get_all_countries(self):
        world = self.get_world_df()
        all_countries = world['name'].to_list()
        return all_countries

    def search_countries(self, searchStrings, superset = None):
        world = self.get_world_df()
        found_df = self._search_df(world, 
                                   column_name = 'name', 
                                   searchStrings = searchStrings, 
                                   superset = superset,
                                   return_columns = ['name', 'geometry'],
                                   )
        return self._restrict_records(found_df)

    def search_rivers_and_lakes(self, searchStrings, superset = None):
        rivers_and_lakes = self.get_rivers_and_lakes_df()
        found_df = self._search_df(rivers_and_lakes, 
                                   column_name = 'name', 
                                   searchStrings = searchStrings, 
                                   superset = superset,
                                   return_columns = ['name', 'geometry'],
                                   )
        return self._restrict_records(found_df)

    def search_usa_states(self, searchStrings, superset = None):
        usa_states = self.get_usa_states_df()
        found_df = self._search_df(usa_states, 
                                   column_name = 'name', 
                                   searchStrings = searchStrings, 
                                   superset = superset,
                                   return_columns = ['name', 'geometry'],
                                   )
        return self._restrict_records(found_df)

    def search_cities(self, searchStrings, superset = None):
        cities = self.get_cities_df()
        
        found_df = self._search_df(cities, 
                                   column_name = 'name', 
                                   searchStrings = searchStrings, 
                                   superset = superset,
                                   return_columns = ['name', 'geometry'],
                                   )
        return self._restrict_records(found_df)

    def _search_df(self, df, 
                   column_name = 'name', 
                   searchStrings = None, 
                   superset = None,
                   return_columns = None
                   ):
        
        if df is None:
            sError = f"""Error {self.__class__.__name__}.{inspect.stack()[0][3]}:
            df is None.
            """
            self._log(sError)
            pass
        
        if not type(df).__name__ == "GeoDataFrame":
            sError = f"""Error {self.__class__.__name__}.{inspect.stack()[0][3]}:
            Invalide df type.
            type(df) = {type(df)}
            """
            self._log(sError)
            pass
            
        
        if not searchStrings:
            return None
        
        pattren = self._get_search_pattern(searchStrings)
        if pattren is None:
            return None
        
        df = df[df[column_name].str.contains(
                    pattren, flags = self._get_re_IGNORECASE(), regex=True, 
                    na=False #added by B Thakur Mar 09, 2024
                    )]
        
        
        if return_columns:
            df = df[return_columns]
        
        if not superset is None:
            df = self.remove_disjoint_records(df, superset)
        
        return df

    def _get_search_pattern(self, searchStrings):
        if not searchStrings:
            return None
        
        if isinstance(searchStrings, str):
            searchStrings = [searchStrings]

        searchStrings = [x.lower() for x in searchStrings]
        pattren = '|'.join(searchStrings)
        return pattren


    def get_world_df(self):
        if self._world_df is None:
            self._init_world_df()
            pass
        return self._world_df

    def get_continents_df(self, column_name = None, aggfunc = 'sum'):
        if column_name is None:
            world = self.get_world_df()[['continent', 'geometry']]
            continents_df = world.dissolve(by='continent')
        else:
            world = self.get_world_df()[['continent', 'geometry', column_name]]
            continents_df = world.dissolve(by='continent', aggfunc=aggfunc)

        continents_df.reset_index(level=0, inplace=True)
        continents_df.rename({'continent': 'name'}, axis=1, inplace=True)
        
        return continents_df
        pass

    def get_countries_df(self):
        return self.get_world_df()

    def get_cities_df(self):
        if self._cities_df is None:
            self._init_cities_df()
            pass
        return self._cities_df

    def get_rivers_and_lakes_df(self):
        if self._rivers_and_lakes_df is None:
            self._init_rivers_and_lakes_df()
            pass
        return self._rivers_and_lakes_df

    def get_elevation_points_df(self):
        if self._elevation_points_df is None:
            self._init_elevation_points_df()
            pass
        return self._elevation_points_df



    def get_filtered_elevation_points(self, data = None):
        '''
        data must be either list or dict of a geopandas dataframes
        '''
        elevation_points = self.get_elevation_points_df()
        if data is None:
            return elevation_points
        
        if isinstance(data, dict):
            data = data.values()
        
        elevation_dfs = []
        for df in data:
            elevation_df = self.remove_disjoint_records(elevation_points, df)
            
            if elevation_df is None:
                continue
            
            if elevation_df.shape[0] == 0:
                continue
            
            elevation_dfs.append(elevation_df)
            pass
        
        filtered_elevation_points = pd.concat(elevation_dfs).drop_duplicates().reset_index(drop=True)
        return filtered_elevation_points



    def get_usa_states_df(self):
        if self._usa_states_df is None:
            self._init_usa_states_df()
            pass
        return self._usa_states_df


    def get_regions_polys_df(self):
        if self._regions_polys_df is None:
            self._init_regions_polys_df()
            pass
        return self._regions_polys_df

    def get_filtered_regions_polys(self, data = None):
        '''
        data must be either list or dict of a geopandas dataframes
        '''
        regions_polys = self.get_regions_polys_df()
        if data is None:
            return regions_polys
        
        if isinstance(data, dict):
            data = data.values()
        
        elevation_dfs = []
        for df in data:
            elevation_df = self.remove_disjoint_records(regions_polys, df)
            
            if elevation_df is None:
                continue
            
            if elevation_df.shape[0] == 0:
                continue
            
            elevation_dfs.append(elevation_df)
            pass
        
        filtered_regions_polys = pd.concat(elevation_dfs).drop_duplicates().reset_index(drop=True)
        return filtered_regions_polys


    def get_ny_df(self):
        if self._ny_df is None:
            self._init_ny_df()
            pass
        return self._ny_df

    def get_geographic_lines_df(self):
        if self._geographic_lines_df is None:
            self._init_geographic_lines_df()
            pass
        return self._geographic_lines_df


    def _init_world_df(self):
        file_path = "geodata/ne_110m_admin_0_countries.zip"

        df = self._init_load_data_file(file_path, 
                                       dropna_columns = ['geometry', 'name']
                                       )
        self._world_df = df
        pass

    def _init_default_world_df(self):
        gpd = self._get_geopandas()
        file_path = gpd.datasets.get_path('naturalearth_lowres')     
        self._world_df = gpd.read_file(file_path)
        pass
    
    def _init_cities_df(self):
        file_path = "geodata/cities.csv"

        df = self._init_load_cities_csv(file_path)
        
        if not df is None:
            df.crs = "EPSG:4326"
            self._cities_df = df
            return
        
        file_path = "geodata/ne_110m_populated_places.zip"

        df = self._init_load_data_file(file_path, 
                                       dropna_columns = ['geometry', 'name']
                                       )
        self._cities_df = df
        pass

    def _init_default_cities_df(self):
        gpd = self._get_geopandas()
        file_path = gpd.datasets.get_path('naturalearth_cities')     
        self._cities_df = gpd.read_file(file_path)
        pass
    
    def _init_geographic_lines_df(self):
        file_path = "geodata/ne_50m_geographic_lines.zip"

        df = self._init_load_data_file(file_path, 
                                       dropna_columns = ['geometry']
                                       )
        
        self._geographic_lines_df = df
        pass
    
    def _init_rivers_and_lakes_df(self):
        file_path = "geodata/ne_50m_rivers_lake_centerlines.zip"

        df = self._init_load_data_file(file_path, 
                                       dropna_columns = ['geometry', 'name']
                                       )
        self._rivers_and_lakes_df = df
        pass

    def _init_elevation_points_df(self):
        file_path = "geodata/ne_50m_geography_regions_elevation_points.zip"

        df = self._init_load_data_file(file_path, 
                                       dropna_columns = ['geometry', 'name', 'elevation']
                                       )
    
        df['name'] = df[['name', 'elevation']].apply(
                            lambda x: f"{x[0]} ({x[1]})", axis = 1)
        self._elevation_points_df = df
        pass


    def _init_usa_states_df(self, ignore_alaska = True):
        file_path = "geodata/cb_2018_us_state_20m.zip"

        df = self._init_load_data_file(file_path, 
                                       dropna_columns = ['geometry', 'name']
                                       )
        if ignore_alaska:
            df = df[~df['name'].isin(['Alaska'])]
            pass
    
        self._usa_states_df = df
        pass

    def _init_regions_polys_df(self):
        file_path = "geodata/ne_10m_geography_regions_polys.zip"

        df = self._init_load_data_file(file_path, 
                                       dropna_columns = ['geometry', 'name']
                                       )
        self._regions_polys_df = df
        pass

    def load_data_file(self, 
                         file_path, 
                         dropna_columns = None,
                         lower_column_names = True,
                         search_file_path = True
                         ):
        
        df = self._init_load_data_file(file_path, 
                                       dropna_columns = dropna_columns,
                                       lower_column_names = lower_column_names,
                                       search_file_path = search_file_path
                                       )
        if df is None:
            self._log(f"Failed to load {file_path}")
            
        return df
        pass



    def _init_load_cities_csv(self, 
                             file_path, 
                             dropna_columns = None,
                             lower_column_names = True,
                             search_file_path = True
                             ):
        if not file_path:
            self._log(f"file_path is None or empty")
            return None
        
        if search_file_path:
            found_file_path = self._init_file_path(file_path)
            if not found_file_path:
                self._log(f"file_path not found: {file_path}")
                return None
            file_path = found_file_path
        else:
            if not os.path.isfile(file_path):
                self._log(f"file_path not found: {file_path}")
                return None
        
        self._log(f"Loading: {file_path}")
        df = pd.read_csv(file_path)
        
        if lower_column_names:
            df.columns = df.columns.str.lower()
        
        #df['geometry'] = df[['latitude','longitude']].apply(
        #                lambda x : f'Point({x[0]},{x[0]})', axis=1)

        df = self._get_geopandas().GeoDataFrame(
            df, geometry = self._get_geopandas().points_from_xy(df.longitude, df.latitude))
        
        
        if dropna_columns:
            df = df.dropna(subset = dropna_columns)
        
        return df
        pass


    def _init_load_data_file(self, 
                             file_path, 
                             dropna_columns = None,
                             lower_column_names = True,
                             search_file_path = True
                             ):
        if not file_path:
            self._log(f"file_path is None or empty")
            return None
        
        if search_file_path:
            found_file_path = self._init_file_path(file_path)
            if not found_file_path:
                self._log(f"file_path not found: {file_path}")
                return None
            file_path = found_file_path
        else:
            if not os.path.isfile(file_path):
                self._log(f"file_path not found: {file_path}")
                return None
        
        self._log(f"Loading: {file_path}")
        df = self._get_geopandas().read_file(file_path)
        
        if lower_column_names:
            df.columns = df.columns.str.lower()
        
        if dropna_columns:
            df = df.dropna(subset = dropna_columns)
        
        return df
        pass


    def _init_ny_df(self):
        gpd = self._get_geopandas()
        file_path = gpd.datasets.get_path('nybb')     
        self._ny_df = gpd.read_file(file_path)
        pass

    def _init_random_names(self):
        dfs = [self.get_world_df(),
               self.get_cities_df(),
               self.get_regions_polys_df(),
               self.get_rivers_and_lakes_df()
               ]
        names = self._parse_names(dfs)
        self._random_names = names
        pass


    def _parse_names(self, dfs):
        names = []
        for df in dfs:
            df_names = [x for x in df.name.to_list() if x]
            df_names = [x for x in df_names if x == convertToUtf8(x)]
            names.extend(df_names)
            pass
        
        names = list(set(names))
        
        names = " ".join(names).split(" ")
        names = [x.strip() for x in names if len(x.strip()) >= 3]
        
        names = list(set(names))
        return names
        pass



    def _init_file_path(self, file_path):
        
        if os.path.isfile(file_path):
            return file_path
        
        from xv.util import Repository
        repository = Repository(verbose = self.verbose)
        found_file_path = repository.get_geo_path(file_path,
                                                search_type = 'file',
                                                part_match = False)

        if found_file_path:
            return found_file_path
        
        return None
        pass

    def _log_DELETED(self, sOut, to_print = None):
        if to_print is None:
            to_print = self.verbose
            
        if to_print:
            print(sOut)
        
    def __str__(self):
        cls = self.__class__
        s = []
        s.append(f"{id(self)}@{cls.__name__}")


        s = "\n".join(s)
        return s

    def _repr_html_(self):
        return self.__str__().replace("\n", "<br />")
    
    def _get_geopandas(self):
        import geopandas
        return geopandas

    def _get_folium(self):
        import folium
        return folium
    
    def _get_re_IGNORECASE(self):
        return re.IGNORECASE


    pass
#class DataProvider:
