""" handle GMF's boundary condition



"""

def cell2block(geo, cell):
    """ returns mulgrid block name (str) of a cell id (int) """
    return geo.block_name_list[geo.num_atmosphere_blocks:][cell]

def update_waiwera_model_heat_bc(geo, wjson, bcjson):
    """ bcjson is loaded from the GMF modifier _boundary_conditions.json
    """
    # heat component
    if wjson['eos']['name'] == 'wae':
        heat_components = [3, 'energy']
    else:
        raise NotImplementedError(f"(gimu)    EOS {wjson['eos']['name']} not implemented")


    if 'Heat' not in bcjson:
        return
    heat_bc = bcjson['Heat']
    default_flux = float(heat_bc['Default']['Heat_Flux'])
    count = 0
    for src in wjson['source']:
        if 'component' in src:
            if src['component'] in heat_components:
                block_name = cell2block(geo, src['cell'])
                column_name = geo.column_name(block_name)
                area = geo.column[column_name].area
                flux = default_flux
                for zone_name, zone in heat_bc.get('Zones', {}).items():
                    if column_name in zone['Columns']:
                        flux = float(zone['Heat_Flux'])
                        break
                src['rate'] = area * flux
                count += 1
                # print(f"Updated {src['name']} heatflow")
    print(f"(gimu)    Updated heatflow in {count} sources")


def update_aut2_model_heat_bc(geo, dat, bcjson):
    if 'Heat' not in bcjson:
        return
    heat_bc = bcjson['Heat']
    default_flux = float(heat_bc['Default']['Heat_Flux'])
    count = 0
    for gen in dat.generatorlist:
        if gen.type == 'HEAT':
            block_name = gen.block
            column_name = geo.column_name(block_name)
            area = geo.column[column_name].area
            flux = default_flux
            for zone_name, zone in heat_bc.get('Zones', {}).items():
                if column_name in zone['Columns']:
                    flux = float(zone['Heat_Flux'])
                    break
            gen.gx = area * flux
            count += 1
            # print(f"Updated {gen.name} heatflow")
    print(f"(gimu)    Updated heatflow in {count} generators")


