"""Load the released encoder and extract features from a synthetic HSI batch."""
import argparse
import torch
from transformers import AutoModel


def load_encoder(model_id="GFM-Bench/LESSViT-Base", device="cpu", revision=None):
    model = AutoModel.from_pretrained(
        model_id, trust_remote_code=True, revision=revision,
    )
    return model.to(device).eval()


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('--model-id', default='GFM-Bench/LESSViT-Base',
                        help='Hugging Face model ID or local snapshot directory')
    parser.add_argument('--revision', default=None, help='Optional commit SHA')
    parser.add_argument('--device', default='cpu')
    args = parser.parse_args()
    torch.set_num_threads(4)
    encoder = load_encoder(args.model_id, args.device, args.revision)
    # Synthetic values demonstrate tensor shapes, not preprocessing for real imagery.
    image = torch.randn(1, 10, 32, 32, device=args.device)
    wavelengths_nm = torch.linspace(450, 2300, 10, device=args.device)[None]
    # Example: fixed full-EnMAP wavelength bounds. Use the same reference bounds
    # for train/eval in your downstream pipeline; do not renormalize each subset.
    wavelengths = (wavelengths_nm - 418.24) / (2445.53 - 418.24)
    with torch.no_grad():
        output = encoder(optical=image, optical_channel_wv=wavelengths,
                         spatial_resolution=30)
    grid = output.last_hidden_state
    assert torch.isfinite(grid).all()
    print('Token grid:', tuple(grid.shape))
    print('Global features:', tuple(grid[:, 0, 0].shape))
    print('Spatial features:', tuple(grid[:, 0, 1:].shape))
    print('Per-band patch features:', tuple(grid[:, 1:, 1:].shape))


if __name__ == '__main__':
    main()