import pandas as pd
import geopandas as gpd
import matplotlib.pyplot as plt

class GeospatialVisualizer:
    def __init__(self):
        self.world = gpd.read_file(gpd.datasets.get_path('naturalearth_lowres'))

    def visualize(self, data, location_column, value_column):
        # Merge data with world map
        merged = self.world.merge(data, how='left', left_on=['name'], right_on=[location_column])

        # Create the plot
        fig, ax = plt.subplots(figsize=(15, 10))
        merged.plot(column=value_column, ax=ax, legend=True, 
                    legend_kwds={'label': value_column, 'orientation': 'horizontal'},
                    missing_kwds={'color': 'lightgrey'})

        # Customize the plot
        ax.set_title(f'{value_column} by Country')
        ax.axis('off')

        return fig

    def create_choropleth(self, data, location_column, value_column):
        # Merge data with world map
        merged = self.world.merge(data, how='left', left_on=['name'], right_on=[location_column])

        # Create the choropleth map
        fig, ax = plt.subplots(figsize=(15, 10))
        merged.plot(column=value_column, ax=ax, legend=True, 
                    legend_kwds={'label': value_column, 'orientation': 'horizontal'},
                    cmap='YlOrRd', missing_kwds={'color': 'lightgrey'})

        # Customize the plot
        ax.set_title(f'Choropleth Map: {value_column} by Country')
        ax.axis('off')

        return fig

    def create_bubble_map(self, data, lat_column, lon_column, size_column):
        # Create a GeoDataFrame from the data
        gdf = gpd.GeoDataFrame(data, geometry=gpd.points_from_xy(data[lon_column], data[lat_column]))

        # Create the bubble map
        fig, ax = plt.subplots(figsize=(15, 10))
        self.world.plot(ax=ax, color='lightgrey')
        gdf.plot(ax=ax, markersize=data[size_column]/data[size_column].max()*100, alpha=0.5)

        # Customize the plot
        ax.set_title(f'Bubble Map: {size_column} by Location')
        ax.axis('off')

        return fig