import numpy as np
import pandas as pd

# choose buildings for extrusion
# import gml_par
# only massive construction types
# only residential buildings

class choose_building:
    def __init__(self, gml_par, bldg_extension, residential_code, extrusion_percentage):
        self.gml_par = gml_par
        self.bldg_extension = bldg_extension
        self.residential_code = residential_code
        self.extrusion_percentage = extrusion_percentage

    def choose_building(self):
        gml_par = self.gml_par
        residential_code = self.residential_code
        extrusion_percentage = self.extrusion_percentage
        storeys_to_add = self.bldg_extension

        # Determine if buildings can be extended
        gml_par['extruded_bool'] = (gml_par['function'] == residential_code) & (gml_par['construction_type'] == 'massive')
        
        # Filter possible buildings
        possible_buildings = gml_par[gml_par['extruded_bool']]
        
        # Calculate the number of buildings to be extended
        num_buildings_to_extend = int(len(possible_buildings) * extrusion_percentage)
        
        gml_par['extruded_bool_percentage'] = False

        # Randomly select buildings to extend
        possible_buildings['extruded_bool_percentage'] = False
        if num_buildings_to_extend > 0:
            selected_indices = np.random.choice(possible_buildings.index, num_buildings_to_extend, replace=False)
            gml_par.loc[selected_indices, 'extruded_bool_percentage'] = True
        if num_buildings_to_extend == 0:
            gml_par['extruded_bool_percentage'] = False
        
        # Set the number of storeys that will be added to the building
        #gml_par['extruded'] = np.where(gml_par['extruded_bool_percentage'], storeys_to_add, 0)

        cond1 = ((gml_par['roofType'].astype(int) == 1000) & (gml_par['extruded_bool_percentage'] == True))
        cond2 = ((gml_par['roofType'].astype(int) > 1000) & (gml_par['extruded_bool_percentage'] == True))
        gml_par['extruded'] = np.where(cond1, storeys_to_add,
                                           np.where(cond2, storeys_to_add - 1, 0))
    
        return gml_par