from tidyfaf.core import FAFRegionData, FAFNetworkData
from tidyfaf.visualization import FlowMap
import pandas as pd
import geopandas as gpd
import matplotlib.pyplot as plt

class FlowQuery:
    def __init__(self, df: pd.DataFrame, metadata, region_data: FAFRegionData):
        self._df = df
        self._metadata = metadata
        self._region_data = region_data

    def filter_origin(self, origin):
        """Filter by origin state (name or code)."""
        if isinstance(origin, str) and not origin.isdigit():
            origin_code = self._metadata.lookup_state(origin)
        else:
            origin_code = int(origin)
        
        # FAF5 Regional data uses 'dms_orig' (FAF Zone).
        # We assume dms_orig // 10 corresponds to State FIPS.
        if 'dms_origst' in self._df.columns:
             self._df = self._df[self._df['dms_origst'] == origin_code]
        else:
             # Filter where derived state code matches
             self._df = self._df[self._df['dms_orig'] // 10 == origin_code]
             
        return self

        return self
        
    def filter_origin_zone(self, zone):
        """Filter by origin FAF Zone (name or code)."""
        if isinstance(zone, str) and not zone.isdigit():
            zone_code = self._metadata.lookup_zone(zone)
        else:
            zone_code = int(zone)
            
        self._df = self._df[self._df['dms_orig'] == zone_code]
        return self

    def filter_destination(self, destination):
        """Filter by destination state (name or code)."""
        if isinstance(destination, str) and not destination.isdigit():
            dest_code = self._metadata.lookup_state(destination)
        else:
            dest_code = int(destination)
            
        if 'dms_destst' in self._df.columns:
             self._df = self._df[self._df['dms_destst'] == dest_code]
        else:
             self._df = self._df[self._df['dms_dest'] // 10 == dest_code]
             
        return self

    def filter_destination_zone(self, zone):
        """Filter by destination FAF Zone (name or code)."""
        if isinstance(zone, str) and not zone.isdigit():
            zone_code = self._metadata.lookup_zone(zone)
        else:
            zone_code = int(zone)
            
        self._df = self._df[self._df['dms_dest'] == zone_code]
        return self

    def filter_commodity(self, commodity):
        """Filter by commodity (name or code)."""
        if isinstance(commodity, str) and not commodity.isdigit():
            comm_code = self._metadata.lookup_commodity(commodity)
        else:
            comm_code = int(commodity)
            
        self._df = self._df[self._df['sctg2'] == comm_code]
        return self
        
    def filter_mode(self, mode):
        """Filter by mode (name or code)."""
        if isinstance(mode, str) and not mode.isdigit():
            mode_code = self._metadata.lookup_mode(mode)
        else:
            mode_code = int(mode)
            
        self._df = self._df[self._df['dms_mode'] == mode_code]
        return self

    def filter_year(self, year):
        """Filter by year (if applicable, though FAF5 often has columns for years)."""
        # FAF5 often has columns like 'tons_2020', 'value_2020'.
        # This might be more about selecting the column to visualize.
        # But if the dataframe is long-format, we filter.
        # Assuming wide format for now, so this might not filter rows but set a context?
        # Actually, let's assume this is for future use or if data is reshaped.
        # For now, we'll just return self as FAF5 is usually wide.
        return self

    def head(self, n=5):
        return self._df.head(n)
        
    def get_data(self):
        return self._df

    def to_gdf(self):
        """
        Convert the current filtered data to a GeoDataFrame with flow lines.
        Returns:
            gpd.GeoDataFrame: Data with 'geometry' column containing LineString flows.
        """
        # Create a copy to avoid modifying the original query df
        plot_df = self._df.copy()
        
        # Get Zone Centroids from official shapefile
        try:
            zones_gdf = self._region_data.get_zones_gdf()
        except Exception as e:
            print(f"Error loading zones: {e}")
            return None

        # Calculate centroids
        centroids = zones_gdf.geometry.centroid
        zone_map = pd.DataFrame({
            'FAFZONE': zones_gdf['FAFZONE'],
            'lat': centroids.y,
            'lon': centroids.x
        }).set_index('FAFZONE')
        
        # Map Origin
        plot_df = plot_df.join(zone_map, on='dms_orig', rsuffix='_orig')
        plot_df.rename(columns={'lat': 'orig_lat', 'lon': 'orig_lon'}, inplace=True)
        
        # Map Destination
        plot_df = plot_df.join(zone_map, on='dms_dest', rsuffix='_dest')
        plot_df.rename(columns={'lat': 'dest_lat', 'lon': 'dest_lon'}, inplace=True)
        
        # Filter out invalid coordinates (NaNs)
        plot_df = plot_df.dropna(subset=['orig_lat', 'orig_lon', 'dest_lat', 'dest_lon'])
        
        if len(plot_df) == 0:
             print("Warning: No flows to map (check filters or coordinate mapping).")
             return gpd.GeoDataFrame(plot_df, geometry=[])

        # Create LineString geometry
        from shapely.geometry import LineString
        plot_df['geometry'] = plot_df.apply(
            lambda row: LineString([(row['orig_lon'], row['orig_lat']), (row['dest_lon'], row['dest_lat'])]), 
            axis=1
        )
        
        return gpd.GeoDataFrame(plot_df, geometry='geometry', crs="EPSG:4326")

    def plot(self, output_path="faf_map.html", flow_column="tons_2020", max_flows=1000):
        """Generate a Deck.gl map."""
        gdf = self.to_gdf()
        if gdf is None or len(gdf) == 0:
            return self

        # For Deck.gl, we need lat/lon columns, which to_gdf provides
        # But FlowMap expects a DataFrame with specific columns, which gdf is.
        
        fm = FlowMap(gdf)
        fm.generate_map(output_path, flow_column, max_flows)
        return self

class NetworkQuery:
    def __init__(self, gdf: gpd.GeoDataFrame):
        self._gdf = gdf

    def filter_state(self, state_code):
        """Filter network by state code (e.g. 'CA', 'TX' or integer FIPS)."""
        # Network data usually has 'STFIPS' or similar.
        # Based on inspection, we saw 'STFIPS' in columns? 
        # Wait, inspection showed: ['ID', 'LENGTH', 'DIR', 'DATA', 'VERSION', 'geometry']
        # It seems the basic GDB layer might NOT have state codes directly in the main table?
        # Or maybe they are in 'DATA' column?
        # If 'STFIPS' is not there, we might need to spatial join or use another attribute.
        # Let's check columns again. 
        # If we can't filter by state easily, we might skip this or use a bounding box if we had one.
        # For now, let's assume we can filter if the column exists, else warn.
        
        if 'STFIPS' in self._gdf.columns:
             # Handle string vs int
             # STFIPS is often 2-digit string "06"
             if self._gdf['STFIPS'].dtype == 'object':
                 state_code = str(state_code).zfill(2)
             
             self._gdf = self._gdf[self._gdf['STFIPS'] == state_code]
        else:
             print("Warning: 'STFIPS' column not found. Filtering by state not supported on this layer.")
        return self

    def filter_route(self, route_name):
        """Filter by route name/number if available."""
        # Often 'SIGN1', 'SIGN2', 'LNAME' etc.
        # Let's try 'LNAME' or 'SIGN1'
        if 'LNAME' in self._gdf.columns:
             self._gdf = self._gdf[self._gdf['LNAME'].str.contains(route_name, case=False, na=False)]
        elif 'SIGN1' in self._gdf.columns:
             self._gdf = self._gdf[self._gdf['SIGN1'].str.contains(route_name, case=False, na=False)]
        return self

    def plot(self, output_path="network_map.html"):
        """Simple plot using geopandas or pydeck."""
        # For quick viz, geopandas plot is easier but static.
        # For consistency, maybe pydeck? But pydeck for lines needs more setup.
        # Let's stick to simple static plot for now or just save to file?
        # The user asked for .plot(), usually implies showing it.
        # But we are in a headless env mostly.
        # Let's save a static image using matplotlib backend of geopandas.
        
        try:
            ax = self._gdf.plot(figsize=(10, 10))
            fig = ax.get_figure()
            fig.savefig(output_path.replace('.html', '.png'))
            print(f"Network map saved to {output_path.replace('.html', '.png')}")
        except Exception as e:
            print(f"Error plotting network: {e}")
        return self
        
    def get_data(self):
        return self._gdf

# --- Main Data Access Functions ---

def get_flows(origin=None, destination=None, origin_zone=None, destination_zone=None, commodity=None, mode=None, year=None, geometry=False):
    """
    Get FAF flow data as a DataFrame or GeoDataFrame.
    
    Args:
        origin (str or int): Origin state name or code.
        destination (str or int): Destination state name or code.
        origin_zone (str or int): Origin FAF Zone name or code.
        destination_zone (str or int): Destination FAF Zone name or code.
        commodity (str or int): Commodity name or code.
        mode (str or int): Mode name or code.
        year (int): Year to filter.
        geometry (bool): If True, returns a GeoDataFrame with flow lines.
        
    Returns:
        pd.DataFrame or gpd.GeoDataFrame
    """
    # Initialize Data
    data = FAFRegionData()
    query = FlowQuery(data.df, data.metadata, data)
    
    # Apply Filters
    if origin:
        query.filter_origin(origin)
    if destination:
        query.filter_destination(destination)
    if origin_zone:
        query.filter_origin_zone(origin_zone)
    if destination_zone:
        query.filter_destination_zone(destination_zone)
    if commodity:
        query.filter_commodity(commodity)
    if mode:
        query.filter_mode(mode)
    if year:
        query.filter_year(year)
        
    # Return Data
    if geometry:
        return query.to_gdf()
    else:
        return query.get_data()

def get_network(state=None, route_name=None):
    """
    Get FAF network links as a GeoDataFrame.
    
    Args:
        state (str or int): State filter (name or code).
        route_name (str): Route name filter (e.g., 'I-5').
        
    Returns:
        gpd.GeoDataFrame
    """
    data = FAFNetworkData()
    gdf = data.get_network_gdf(layer='links')
    query = NetworkQuery(gdf)
    
    if state:
        query.filter_state(state)
    if route_name:
        query.filter_route(route_name)
        
    return query.get_data()

# --- Metadata Search Helpers ---

def _get_metadata():
    return FAFRegionData().metadata

def search_commodity(query):
    """Search for commodities matching the query string."""
    md = _get_metadata()
    df = md.commodities
    return df[df.astype(str).apply(lambda x: x.str.contains(query, case=False)).any(axis=1)]

def search_zone(query):
    """Search for FAF Zones matching the query string."""
    md = _get_metadata()
    df = md.zones
    return df[df.astype(str).apply(lambda x: x.str.contains(query, case=False)).any(axis=1)]

def search_state(query):
    """Search for States matching the query string."""
    md = _get_metadata()
    df = md.states
    return df[df.astype(str).apply(lambda x: x.str.contains(query, case=False)).any(axis=1)]

def search_mode(query):
    """Search for Modes matching the query string."""
    md = _get_metadata()
    df = md.modes
    return df[df.astype(str).apply(lambda x: x.str.contains(query, case=False)).any(axis=1)]
