from pathlib import Path
import os
import sys

ROOT = Path(__file__).resolve().parents[1]
BACKBONE = ROOT/'assets/esm2_t33_650M_UR50D'
PREDICTOR = ROOT/'checkpoints/metalatte_predictor.ckpt'
DIFFUSION = ROOT/'checkpoints/metalorian_diffusion.pt'
TRAIN = ROOT/'data/metalatte/train'
VALIDATION = ROOT/'data/metalatte/validation'
os.environ['USE_TF'] = '0'
os.environ.setdefault('HF_HUB_OFFLINE', '1')
os.environ.setdefault('WANDB_MODE', 'disabled')
os.environ.setdefault('MPLCONFIGDIR', str(ROOT/'.cache/matplotlib'))
sys.dont_write_bytecode = True
_configured = False


def configure():
    global _configured
    if _configured:
        return
    import torch
    original = torch.load

    def load(file, *args, **kwargs):
        candidate = file if isinstance(file, (str, Path)) else getattr(file, 'name', getattr(file, 'path', None))
        if isinstance(candidate, (str, Path)) and Path(candidate).resolve() in {PREDICTOR.resolve(), DIFFUSION.resolve()}:
            file = str(Path(candidate).resolve())
            kwargs.setdefault('weights_only', False)
            kwargs.setdefault('mmap', True)
        return original(file, *args, **kwargs)

    torch.load = load
    _configured = True