import os
import zipfile
import tempfile

def updateZip(zipname, filename, data):
    # generate a temp file
    tmpfd, tmpname = tempfile.mkstemp(dir=os.path.dirname(zipname))
    os.close(tmpfd)

    # create a temp copy of the archive without filename            
    with zipfile.ZipFile(zipname, 'r') as zin:
        with zipfile.ZipFile(tmpname, 'w') as zout:
            zout.comment = zin.comment # preserve the comment
            for item in zin.infolist():
                if item.filename != filename:
                    zout.writestr(item, zin.read(item.filename))

    # replace with the temp archive
    os.remove(zipname)
    os.rename(tmpname, zipname)

    # now add filename with its new data
    with zipfile.ZipFile(zipname, mode='a', compression=zipfile.ZIP_DEFLATED) as zf:
        zf.writestr(filename, data)

ppo_mapping = {
      'action_net.0.weight': 'mlp_extractor.policy_net.0.weight',
      'action_net.0.bias': 'mlp_extractor.policy_net.0.bias',
      'action_net.2.weight': 'mlp_extractor.policy_net.2.weight',
      'action_net.2.bias': 'mlp_extractor.policy_net.2.bias',
      'action_net.4.weight': 'action_net.weight',
      'action_net.4.bias': 'action_net.bias',
}

sac_mapping = {
      'action_net.0.weight': 'actor.latent_pi.0.weight',
      'action_net.0.bias': 'actor.latent_pi.0.bias',
      'action_net.2.weight': 'actor.latent_pi.2.weight',
      'action_net.2.bias': 'actor.latent_pi.2.bias',
      'action_net.4.weight': 'actor.mu.weight',
      'action_net.4.bias': 'actor.mu.bias',
}

td3_mapping = {
      'action_net.0.weight': 'actor.mu.0.weight',
      'action_net.0.bias': 'actor.mu.0.bias',
      'action_net.2.weight': 'actor.mu.2.weight',
      'action_net.2.bias': 'actor.mu.2.bias',
      'action_net.4.weight': 'actor.mu.4.weight',
      'action_net.4.bias': 'actor.mu.4.bias',
}