import pydeck as pdk
import pandas as pd

class FlowMap:
    def __init__(self, data: pd.DataFrame):
        """
        Initialize the FlowMap with FAF data.
        
        Args:
            data (pd.DataFrame): The FAF data containing origin and destination information.
                                 Expected columns: 'dms_orig', 'dms_dest', 'tons_2020' (or similar flow column)
        """
        self.data = data

    def generate_map(self, output_path="faf_flow_map.html", flow_column="tons_2020", max_flows=1000):
        """
        Generate a Deck.gl ArcLayer map visualizing freight flows.
        
        Args:
            output_path (str): Path to save the HTML map.
            flow_column (str): Column name representing the flow magnitude (e.g., tons, value).
            max_flows (int): Maximum number of top flows to visualize to avoid clutter.
        """
        # Aggregate data if necessary (assuming data might be disaggregated by commodity)
        # For simplicity, we'll assume the user might want to aggregate by OD pair first
        # But here we'll just take the top flows
        
        # Check if we have lat/lon coordinates. 
        # FAF data usually provides region codes. We need a mapping to lat/lon.
        # For this initial version, we will assume the data has been joined with coordinates
        # OR we can provide a simple lookup for major FAF zones if needed.
        # Let's assume the user will provide data with 'orig_lat', 'orig_lon', 'dest_lat', 'dest_lon'
        # If not, we might need a helper to join with centroids.
        
        required_cols = ['orig_lat', 'orig_lon', 'dest_lat', 'dest_lon', flow_column]
        if not all(col in self.data.columns for col in required_cols):
             raise ValueError(f"Data must contain columns: {required_cols}")

        # Sort by flow and take top N
        df_plot = self.data.sort_values(by=flow_column, ascending=False).head(max_flows)

        # Normalize flow for width
        max_val = df_plot[flow_column].max()
        df_plot['width'] = df_plot[flow_column] / max_val * 10

        layer = pdk.Layer(
            "ArcLayer",
            df_plot,
            get_source_position=["orig_lon", "orig_lat"],
            get_target_position=["dest_lon", "dest_lat"],
            get_source_color=[0, 255, 0, 160],
            get_target_color=[255, 0, 0, 160],
            get_width="width",
            pickable=True,
            auto_highlight=True,
        )

        view_state = pdk.ViewState(
            latitude=39.8283,
            longitude=-98.5795,
            zoom=3,
            pitch=45,
            bearing=0
        )

        r = pdk.Deck(
            layers=[layer],
            initial_view_state=view_state,
            tooltip={"text": "{dms_orig} to {dms_dest}\nFlow: {" + flow_column + "}"}
        )

        r.to_html(output_path)
        print(f"Map saved to {output_path}")
