import re
import numpy as np
import pandas as pd
from sphere.core.schemas.buildings import Buildings
from sphere.core.schemas.field_mapping import FieldMapping
from sphere.core.schemas.abstract_vulnerability_function import AbstractVulnerabilityFunction

import importlib.resources as pkg_resources

class DefaultFloodVulnerability(AbstractVulnerabilityFunction):
    """
    Default flood vulnerability function implementation.
    
    This class implements flood-specific vulnerability calculations
    using HAZUS methodology for riverine and coastal flooding.
    """
    
    def __init__(
        self,
        buildings: Buildings,
        flood_type: str,
    ):
        """
        Initializes a DefaultFloodVulnerability object.

        Args:
            buildings (Buildings): Buildings object.
            flood_type (str): The type of flood to analyze (R, C).
                R = Riverine flooding
                C = Coastal flooding
        """
        
        if flood_type not in ["R", "C"]:
            raise ValueError(f"Invalid flood_type: {flood_type}. Must be one of: R, C")
        
        self.flood_type = flood_type
        self.buildings = buildings

        # Use package resources to load CSV files from the sphere_data package.
        with (
            pkg_resources.files("sphere.data")
            .joinpath("flBldgDmgFn.csv")
            .open("r", encoding="utf-8-sig") as bdf_file
        ):
            self.bdf = pd.read_csv(bdf_file)
        with (
            pkg_resources.files("sphere.data")
            .joinpath("flContDmgFn.csv")
            .open("r", encoding="utf-8-sig") as cdf_file
        ):
            self.cdf = pd.read_csv(cdf_file)
        with (
            pkg_resources.files("sphere.data")
            .joinpath("flInvDmgFn.csv")
            .open("r", encoding="utf-8-sig") as idf_file
        ):
            self.idf = pd.read_csv(idf_file)
        with (
            pkg_resources.files("sphere.data")
            .joinpath("flDmgXRef.csv")
            .open("r", encoding="utf-8-sig") as xdf_file
        ):
            self.xdf = pd.read_csv(xdf_file)

        # index is used later so very important these are set.
        self.bdf = self.bdf.set_index("BldgDmgFnId")
        self.cdf = self.cdf.set_index("ContDmgFnId")
        self.idf = self.idf.set_index("InvDmgFnId")
        self.xdf = self.xdf.set_index(
            [
                "Occupancy",
                "Basement",
                "StoriesMin",
                "StoriesMax",
                "HazardR",
                "HazardCV",
                "HazardCA",
            ]
        )

    def calculate_vulnerability(self):
        """
        Calculate vulnerability percentages for buildings based on flood depth and type.
        This is the main method that implements the abstract vulnerability function.
        """
        
        # Apply damage IDs based on flood type and depth
        self.apply_damage_ids_vectorized()
        
        # Calculate damage percentages using the assigned IDs
        self.apply_damage_percentages()

    def apply_damage_ids_vectorized(self):
        """
        Applies damage IDs to the buildings GeoDataFrame in a vectorized manner.
        """
        fields = self.buildings.fields
        gdf = self.buildings.gdf

        # Handle special case for RES3 occupancy
        gdf['OccupancyClean'] = self.buildings.occupancy_type.str.replace(r'RES3(?=[^\d])', 'RES3', regex=True)

        # Ensure all necessary columns exist
        for col in [fields.get_field_name('occupancy_type'), fields.get_field_name('number_stories'), fields.get_field_name('bddf_id'), fields.get_field_name('cddf_id'), fields.get_field_name('iddf_id'), fields.get_field_name('flood_depth')]:
            if col not in gdf.columns:
                if col == fields.get_field_name('iddf_id'):
                    # Optional field, can be skipped if not present.
                    continue
                raise ValueError(f"Column '{col}' not found in gdf.")

        # Map foundation_type to basement (0 or 1)
        gdf['Basement'] = self.buildings.foundation_type.isin(['B', 4]).astype(int)
        
        # Apply damage IDs in depth-based subsets
        if (self.flood_type == 'C'):
            self._apply_damage_ids_for_subset(gdf, fields, "CV", self.buildings.flood_depth >= 6)
            self._apply_damage_ids_for_subset(gdf, fields, "CA", (self.buildings.flood_depth >= 3) & (self.buildings.flood_depth < 6))
            self._apply_damage_ids_for_subset(gdf, fields, "R", self.buildings.flood_depth < 3)
        else:
            self._apply_damage_ids_for_subset(gdf, fields, "R", self.buildings.flood_depth > -5)
        
        #clean up temp column
        del gdf["OccupancyClean"]
    
    def _apply_damage_ids_for_subset(self, gdf, fields: FieldMapping, hazard_type, subset_mask):
        """
        Applies damage IDs to a subset of the buildings GeoDataFrame based on hazard type.

        Args:
            gdf: The GeoDataFrame.
            fields: BuildingMapping fields.
            hazard_type: "R", "CA", or "CV".
            subset_mask: Boolean mask for the subset of rows.
        """
        # Early Exit if subset is empty.
        if not subset_mask.any():
            return

        # Create a MultiIndex from the GeoDataFrame columns for efficient matching
        gdf_multi_index = pd.MultiIndex.from_arrays(
            [
                gdf.loc[subset_mask, "OccupancyClean"],
                gdf.loc[subset_mask, "Basement"].astype(int),  # Cast boolean to int for matching
                gdf.loc[subset_mask, fields.get_field_name('number_stories')],
            ],
            names=["Occupancy", "Basement", "Stories"],
        )

        # Prepare the hazard match.
        hazard_match = self.xdf.index.get_level_values(f"Hazard{hazard_type}") == 1
        
        # Prepare common vectorized conditions
        stories_match = (
            (gdf_multi_index.get_level_values("Stories").to_numpy()[:, np.newaxis] >= self.xdf.index.get_level_values("StoriesMin").to_numpy()) &
            (gdf_multi_index.get_level_values("Stories").to_numpy()[:, np.newaxis] < self.xdf.index.get_level_values("StoriesMax").to_numpy())
        )

        occupancy_match = gdf_multi_index.get_level_values("Occupancy").to_numpy()[:, np.newaxis] == self.xdf.index.get_level_values("Occupancy").to_numpy()
        basement_match = gdf_multi_index.get_level_values("Basement").to_numpy()[:, np.newaxis] == self.xdf.index.get_level_values("Basement").to_numpy()

        # Iterate over the damage ID fields
        for dmgIdField, output_column in [
            ("BldgDmgFnId", fields.get_field_name('bddf_id')),
            ("ContDmgFnId", fields.get_field_name('cddf_id')) if fields.get_field_name('cddf_id') in gdf.columns else (None, None),
            ("InvDmgFnId", fields.get_field_name('iddf_id')) if fields.get_field_name('iddf_id') in gdf.columns else (None, None),
        ]:
            if dmgIdField is None:
                continue

            # Only process rows where the output column is empty, 0, or NaN
            # First get the indices of the subset
            subset_indices = gdf.loc[subset_mask].index
            
            # Find rows in the subset where the values need updating
            need_update = (gdf.loc[subset_indices, output_column] == 0) | pd.isna(gdf.loc[subset_indices, output_column])
            
            # Skip if nothing needs updating
            if not need_update.any():
                continue
                
            # Get the final indices that need updating
            update_indices = subset_indices[need_update]
            
            # Only run the expensive calculations if we have rows to update
            if len(update_indices) > 0:
                # Get just the subset we need to process
                subset_to_process = gdf.loc[update_indices]
                
                # Create new MultiIndex for just the rows we're updating
                update_multi_index = pd.MultiIndex.from_arrays(
                    [
                        subset_to_process["OccupancyClean"],
                        subset_to_process["Basement"].astype(int),
                        subset_to_process[fields.get_field_name('number_stories')],
                    ],
                    names=["Occupancy", "Basement", "Stories"],
                )
                
                # Recalculate the matching conditions for just the update rows
                update_stories_match = (
                    (update_multi_index.get_level_values("Stories").to_numpy()[:, np.newaxis] >= self.xdf.index.get_level_values("StoriesMin").to_numpy()) &
                    (update_multi_index.get_level_values("Stories").to_numpy()[:, np.newaxis] < self.xdf.index.get_level_values("StoriesMax").to_numpy())
                )
                
                update_occupancy_match = update_multi_index.get_level_values("Occupancy").to_numpy()[:, np.newaxis] == self.xdf.index.get_level_values("Occupancy").to_numpy()
                update_basement_match = update_multi_index.get_level_values("Basement").to_numpy()[:, np.newaxis] == self.xdf.index.get_level_values("Basement").to_numpy()
                
                # Combine conditions for vectorized lookup
                all_conditions = (update_occupancy_match & update_basement_match & update_stories_match & hazard_match)
                
                # Lookup the damage ID values from the cross-reference table.
                lookup_values = self.xdf[dmgIdField].values
                
                # Extract the correct ids.
                matches = lookup_values[all_conditions.argmax(axis=1)]
                
                # Set the value of those who didn't have any matches to 0.
                matches[all_conditions.sum(axis=1) == 0] = 0
                
                # Update directly using the update_indices
                gdf.loc[update_indices, output_column] = matches

    def apply_damage_percentages(self):
         # Add ouptut columns to building_points
        fields = self.buildings.fields

        self.buildings.building_damage_percent = pd.NA
        self.buildings.content_damage_percent = pd.NA
        self.buildings.inventory_damage_percent = pd.NA

        # Already have the three IDs identified.
        self._interpolate_from_lookup(self.bdf, fields.get_field_name('depth_in_structure'), fields.get_field_name('building_damage_percent'), fields.get_field_name('bddf_id'))
        self._interpolate_from_lookup(self.cdf, fields.get_field_name('depth_in_structure'), fields.get_field_name('content_damage_percent'), fields.get_field_name('cddf_id'))
        if fields.get_field_name('iddf_id') in self.buildings.gdf.columns:
            self._interpolate_from_lookup(self.idf, fields.get_field_name('depth_in_structure'), fields.get_field_name('inventory_damage_percent'), fields.get_field_name('iddf_id'))

    def _interpolate_from_lookup(self, lookup_df: pd.DataFrame, flooddepth_col: str | None, result_col: str | None, id_col_gdf: str | None):
        """
        Interpolates values from a lookup table, handling negative columns,
        matching by an ID column in the gdf to a unique ID column in the lookup_df,
        and adds the result to a specified column. Modifies the gdf in-place.

        Args:
            lookup_df: Dataframe with columns like 'ft04m', 'ft00', 'ft01', etc.,
                and a *unique* ID column.  The ID column should be set as the index.
            flooddepth_col: Name of the flood depth column in gdf.
            result_col: Name of the column in gdf to store the interpolated result.
            id_col_gdf: Name of the ID column in the gdf.

        Returns:
            None (Modifies the gdf in-place)
        """
        gdf = self.buildings.gdf

        # Check if the ID columns exist
        if id_col_gdf not in gdf.columns:
            raise ValueError(f"ID column '{id_col_gdf}' not found in gdf.")

        if not lookup_df.index.is_unique:
            raise ValueError("The index of lookup_df must be unique.")
            
        # Pre-compile the regex for efficiency
        col_name_pattern = re.compile(r"^ft(\d+)(m)?$")

        # Create mapping dictionaries *outside* the loop
        def col_name_to_int(col_name):
            match = col_name_pattern.match(col_name)
            if match:
                value = int(match.group(1))
                return -value if match.group(2) == 'm' else value
            return col_name  # Return original if no match

        col_to_int_map = {
            col: col_name_to_int(col)
            for col in lookup_df.columns
            if isinstance(col_name_to_int(col), int) # only map columns we can convert.
        }
        int_to_col_map = {value: key for key, value in col_to_int_map.items()}

        # --- Main Interpolation Logic ---
        gdf[result_col] = np.nan  # Initialize the result column

        # Iterate over the unique IDs in the GDF
        for unique_id in gdf[id_col_gdf].unique():

            # Filter the gdf for the current ID (this is our primary subset)
            gdf_subset = gdf[gdf[id_col_gdf] == unique_id]

            # Efficiently check if the ID exists in the lookup_df
            if unique_id not in lookup_df.index:
                continue  # Skip if ID is not found
            lookup_row = lookup_df.loc[unique_id] # a Series

            # Get the valid lookup ints *for this row*
            lookup_ints = np.array(sorted([col_to_int_map[col] for col in lookup_row.index if col in col_to_int_map]))

            # 2. Find Floor and Ceiling (using the gdf_subset)
            flood_values = gdf_subset[flooddepth_col].to_numpy()
            # Use 'side=right' to handle the case where flood_value is exactly an integer.
            insertion_points = np.searchsorted(lookup_ints, flood_values, side='right')
            #Clip insertion points.  This is very important.
            insertion_points = np.clip(insertion_points, 1, len(lookup_ints) -1)
            
            floor_indices = insertion_points - 1
            ceil_indices = insertion_points

            floor_ints = lookup_ints[floor_indices]
            ceil_ints = lookup_ints[ceil_indices]
            
            # Map integer indices back to column names, handling potential missing mappings
            floor_cols = [int_to_col_map.get(i) for i in floor_ints]
            ceil_cols = [int_to_col_map.get(i) for i in ceil_ints]


            # Handle edge cases:  If the flood depth is exactly on a lookup value, use that value directly.
            floor_lookup_values = np.array([lookup_row[col] if col is not None else np.nan for col in floor_cols])
            ceil_lookup_values = np.array([lookup_row[col] if col is not None else np.nan for col in ceil_cols])


            # 4. Interpolate
            interp_factor = np.where(
                (ceil_ints - floor_ints) != 0,
                (flood_values - floor_ints) / (ceil_ints - floor_ints),
                0  # Avoid division by zero
            )
            interp_result = floor_lookup_values + interp_factor * (ceil_lookup_values - floor_lookup_values)

            # Assign to the *original* gdf using the index from the subset.
            gdf.loc[gdf_subset.index, result_col] = interp_result
