import torch
import yaml
import timm
from config.utils import (
    get_device,
    get_dataloader,
    get_resolution,
    get_model,
    get_optimizer,
    get_scheduler,
)
from trainer.trainer_IC import Trainer

torch.set_float32_matmul_precision("high")
path_cfg = "config/cfg_IC.yaml"
with open(path_cfg, "r") as file:
    cfg = yaml.safe_load(file)

device = get_device(cfg)
train_loader, test_loader, input_shape, output_shape = get_dataloader(cfg)
resolution = get_resolution(cfg)
model = get_model(cfg, device, output_shape, resolution)

if cfg["finetuning"]:
    # We download a Swin V2 with pretrained weights, num_classes = 0 means that there is no classification head
    pretrained = timm.create_model(
        "swinv2_tiny_window8_256", pretrained=True, num_classes=0
    )

    pretrained_dict = pretrained.state_dict()

    missing_keys, unexpected_keys = model.load_state_dict(pretrained_dict, strict=False)
    print(f"Loaded weights, missing: {missing_keys}, unexpected: {unexpected_keys}")

    # Let's freeze all the weights
    for param in model.parameters():
        param.requires_grad = False  # Freeze all parameters

    # Unfreeze parameters for each missing key
    for key in missing_keys:
        # Split key into parts (e.g., 'head.fc.weight' → ['head', 'fc', 'weight'])
        parts = key.split(".")

        # Traverse to the module containing the parameter
        module = model
        for part in parts[:-1]:
            module = getattr(module, part)

        # The last part is the parameter name (weight or bias)
        param_name = parts[-1]
        param = getattr(module, param_name)

        # Unfreeze (requires_grad=True)
        param.requires_grad = True


optimizer = get_optimizer(cfg, model)
scheduler = get_scheduler(cfg, optimizer)

trainer = Trainer(
    model,
    train_loader,
    test_loader,
    optimizer,
    scheduler,
    cfg["loss_fn"],
    cfg["epochs"],
    cfg["freq_print"],
    cfg["device"],
    cfg["state_dict_name"],
    cfg["mixed_precision"],
    cfg["torch_dot_compile"],
    cfg["before_training"],
)
trainer.train()
