import torch
import os

# --- Configuration ---
# Pfad zur Eingabe-BIN-Datei
input_checkpoint_path = "./pytorch_model_oldLayerNames.bin"
# Pfad zur Ausgabe-BIN-Datei (wo die geänderte Version gespeichert wird)
output_checkpoint_path = "./pytorch_model_renamed_final.bin"

# Define the layers that *must* be named 'segmentation_head'
# These are the specific layers you identified
specific_segmentation_head_layers = [
    'decode_head.conv_seg.bias',
    'decode_head.conv_seg.weight',
    'decode_head.convs.0.conv.bias',
    'decode_head.convs.0.conv.weight',
    'decode_head.convs.1.conv.bias',
    'decode_head.convs.1.conv.weight',
    'decode_head.convs.2.conv.bias',
    'decode_head.convs.2.conv.weight',
    'decode_head.convs.3.conv.bias',
    'decode_head.convs.3.conv.weight',
    'decode_head.fusion_conv.conv.bias',
    'decode_head.fusion_conv.conv.weight'
]

# --- Check, ob die Eingabedatei existiert ---
if not os.path.exists(input_checkpoint_path):
    print(f"Fehler: Eingabedatei nicht gefunden unter {input_checkpoint_path}. Bitte den Pfad korrigieren. ❌")
else:
    # --- Checkpoint laden ---
    state_dict = torch.load(input_checkpoint_path, map_location="cpu")

    # --- Layer-Namen ändern und neues State Dict erstellen ---
    new_state_dict = {}
    renamed_count_segmentation = 0
    renamed_count_segformer = 0
    skipped_count = 0

    for old_key, value in state_dict.items():
        if old_key in specific_segmentation_head_layers:
            # These specific layers get 'segmentation_head.' prefix
            new_key = old_key.replace('decode_head.', 'segmentation_head.', 1)
            new_state_dict[new_key] = value
            renamed_count_segmentation += 1
        elif old_key.startswith('decode_head.'):
            # All other layers starting with 'decode_head.' get 'segformer_head.' prefix
            new_key = old_key.replace('decode_head.', 'segformer_head.', 1)
            new_state_dict[new_key] = value
            renamed_count_segformer += 1
        else:
            # Keep other layers as they are (e.g., backbone layers)
            new_state_dict[old_key] = value
            skipped_count += 1

    # --- Geändertes State Dict speichern ---
    torch.save(new_state_dict, output_checkpoint_path)

    print(f"✅ Fertig! Die umbenannte Datei wurde gespeichert unter: {output_checkpoint_path}")
    print(f"Zusammenfassung der Umbenennungen:")
    print(f"  - '{renamed_count_segmentation}' Layer von 'decode_head.' zu 'segmentation_head.' umbenannt.")
    print(f"  - '{renamed_count_segformer}' Layer von 'decode_head.' zu 'segformer_head.' umbenannt.")
    print(f"  - '{skipped_count}' Layer behielten ihren ursprünglichen Namen (z.B. Backbone).")