import torch
import torch.nn as nn
from contextlib import contextmanager
from collections import OrderedDict

@contextmanager
def use_pretrained_weights(builder, **kwargs):
  bIsFineTuning = builder.pretrained_weights is not None
  
  bIsFrozenStem = kwargs.get("is_frozen_stem", False)
  nCreationClassCount = builder.hyperparams.output_neurons
  bMustChangeClassifier = False
  if bIsFineTuning:
    nPretrainedClassCount = len(builder.pretrained_weights.meta["categories"])
    bMustChangeClassifier = (builder.hyperparams.output_neurons != nPretrainedClassCount)
    nCreationClassCount = nPretrainedClassCount
  
  try:
    yield nCreationClassCount  # Your code runs here
  finally:
    if bIsFineTuning:
      bIsModifiedStructure = kwargs.get("is_modified", False)
      if bIsModifiedStructure:
        remap = kwargs["remap"]
        # 2) Load pretrained state dict (torchvision weights or your checkpoint)
        loaded = builder.pretrained_weights.get_state_dict(progress=kwargs.get("progress", False), check_hash=True)
        
        # 3) Merge: keep current params where missing/mismatch (clip params, new layers, classifier)
        merged_sd = merge_state_dict(
          builder.model,
          loaded,
          key_remap_rules=remap,
          #drop_prefixes=("classifier.",),  # prevents 1000->314 mismatch
          verbose=True,
        )
        
        # 4) Now load merged. This should be clean because merged has EXACT model keys.
        builder.model.load_state_dict(merged_sd, strict=True)
      else:
        if builder.pretrained_weights is not None:
          builder.model.load_state_dict(
            builder.pretrained_weights.get_state_dict(progress=kwargs.get("progress", False), check_hash=True))
      
      # Replace the last Linear layer in the classifier
      if bMustChangeClassifier:
        adjust_classifier_to_class_count(builder.model, builder.hyperparams.output_neurons)
      
    if builder.import_stem_weights() or bIsFrozenStem:
      freeze_stem(builder.model)
      print_frozen(builder.model)
  
def adjust_classifier_to_class_count(model: nn.Module, class_count):
  if hasattr(model, "classifier"):
    if isinstance(model.classifier, nn.Sequential):
      # TODO: Test
      # Efficient, ConvNext
      for nIndex,oLayer in enumerate(model.classifier):
        if isinstance(oLayer, nn.Linear):
          nLastFeatures = model.classifier[nIndex].in_features
          model.classifier[nIndex] = nn.Linear(nLastFeatures, class_count)
    else:
      # DenseNet
      nLastFeatures = model.classifier.in_features
      model.classifier = nn.Linear(nLastFeatures, class_count)
  elif hasattr(model, "fc"):
    # ResNet, WideResNet
    nLastFeatures = model.fc.in_features
    model.fc = nn.Linear(nLastFeatures, class_count)


  

def merge_state_dict(
    model: torch.nn.Module,
    loaded_sd: dict,
    *,
    key_remap_rules=None,
    drop_prefixes=("classifier.",),  # drop imagenet head by default
    verbose=True,
):
    """
    Returns a merged state_dict:
      - Starts from model.state_dict() (keeps your new params)
      - Overwrites with loaded weights where (key exists AND shape matches)
      - Optionally remaps loaded keys to match your new module structure
    """
    model_sd = model.state_dict()
    model_keys = set(model_sd.keys())

    # Normalize incoming dict (handles checkpoints that store under "state_dict")
    if "state_dict" in loaded_sd and isinstance(loaded_sd["state_dict"], dict):
        loaded_sd = loaded_sd["state_dict"]

    # Apply key remapping to loaded sd (if you changed module indexing/names)
    remapped_loaded = OrderedDict()
    for k, v in loaded_sd.items():
        # drop some prefixes (e.g. classifier from imagenet)
        if any(k.startswith(p) for p in drop_prefixes):
            continue

        k2 = k
        if key_remap_rules:
            for old, new in key_remap_rules:
                k2 = k2.replace(old, new)

        remapped_loaded[k2] = v

    merged = OrderedDict(model_sd)  # start from current model sd (keeps new params)
    loaded_ok = []
    skipped = []  # (reason, key)

    for k, v in remapped_loaded.items():
        if k not in model_keys:
            skipped.append(("no_such_key_in_model", k))
            continue

        if merged[k].shape != v.shape:
            skipped.append((f"shape_mismatch model={tuple(merged[k].shape)} ckpt={tuple(v.shape)}", k))
            continue

        merged[k] = v
        loaded_ok.append(k)

    if verbose:
        print(f"Loaded {len(loaded_ok)} tensors into current model.")
        print(f"Skipped {len(skipped)} tensors.")
        if skipped[:10]:
            print("First skips:")
            for r, k in skipped[:10]:
                print("  -", r, ":", k)

    return merged




def freeze_stem(model):
    for module in model.modules():
        if isinstance(module, nn.Conv2d):
            for param in module.parameters():
                param.requires_grad = False
            break  # stop after first conv layer

def print_frozen(model):
  for name, param in model.named_parameters():
      if not param.requires_grad:
        print("Frozen:", name)
        
        
def fix_optimizer(model):
  # TODO: Experimental for on-the-fly freezing
  optimizer = torch.optim.Adam(
      filter(lambda p: p.requires_grad, model.parameters()),
      lr=1e-4
  )
  return optimizer

