import torch
import torchvision
import torchvision.transforms as t
from torchvision import models
import torchvision.ops as ops

def get_available_classification_models():
    return [k for k, v in models.__dict__.items() if callable(v) and k[0].lower() == k[0] and k[0] != "_"]

all_classification_models = get_available_classification_models()
print(all_classification_models)


for model_name in all_classification_models:
    if model_name in ('mnasnet0_75', 'mnasnet1_3', 'shufflenet_v2_x1_5', 'shufflenet_v2_x2_0'):
        continue
    getattr(models, model_name)(pretrained=True)




from glob import glob
import io
from pickle import UnpicklingError

all_files = glob('/data/home/nicolashug/.cache/torch/hub/checkpoints/*.pth')
baddies = []
goodies = []

for filename in all_files:
    try:
        state_dict = torch.load(filename, map_location='cuda')
    except RuntimeError:
        baddies.append(filename)
    else:
        goodies.append(filename)

print("these files are OK:")
print(goodies)
print("these ones failed:")
print(baddies)


import hashlib
import os
import uuid

output_dir = '/data/home/nicolashug/new_models'
tmp_path = os.path.join(output_dir, 'blah')

for filename in baddies:
    model_name = '-'.join(filename.split('/')[-1].split('-')[:-1])
    model = torch.load(filename)
    torch.save(model, tmp_path)

    sha256_hash = hashlib.sha256()
    with open(tmp_path, "rb") as f:
        # Read and update hash string value in blocks of 4K
        for byte_block in iter(lambda: f.read(4096), b""):
            sha256_hash.update(byte_block)
        hh = sha256_hash.hexdigest()


    output_path = os.path.join(output_dir, model_name + "-" + str(hh[:8]) + ".pth")
    os.replace(tmp_path, output_path)


# Make sure they can be loaded now
for filename in glob(output_dir + '/*'):
    state_dict = torch.load(filename, map_location='cuda')
    
print("The newly updated weights don't fail anymore")
print(f"They are located in {output_dir}")

# Also make sure the new dicts are the same as the old ones
for d1, d2 in zip(sorted(glob(output_dir + '/*')), sorted(baddies)):
    d1 = torch.load(d1)
    d2 = torch.load(d2)
    assert d1.keys() == d2.keys()
    for v1, v2 in zip(d1.values(), d2.values()):
        torch.testing.assert_allclose(v1, v2)

print("newly updated weights seem to be the same as the previous ones.")
print("all good")