"""Hugging Face interface for the original LESSViT encoder."""
# Explicit dependencies let Transformers copy the full local/remote module tree.
from .pos_chan_embed import RopePositionChannelEmbedding
from .hyperspectral_patch_embed import HyperspectralPatchEmbed
from .low_rank_attention import LowRankAttention
from .attention_ops import LowRankBlock
from .spatial_spectral_low_rank_vit import (
    SpatialSpectralLowRankViTConfig,
    SpatialSpectralLowRankViTEncoder,
)


class LESSViTConfig(SpatialSpectralLowRankViTConfig):
    model_type = "lessvit"

    def __init__(self, channel_embed_dims_per_head=2, **kwargs):
        # The original config constructor recomputes these derived dimensions.
        kwargs.pop("channel_dim", None)
        kwargs.pop("spatial_dim", None)
        kwargs.pop("pretrain", None)
        kwargs.pop("mask_ratio", None)
        kwargs.pop("channel_mask_ratio", None)
        kwargs.pop("channel_dropout", None)
        super().__init__(channel_embed_dims_per_head=channel_embed_dims_per_head,
                         mask_ratio=0.0, channel_mask_ratio=0.0,
                         channel_dropout=None, **kwargs)
        self.channel_embed_dims_per_head = channel_embed_dims_per_head
        self.pretrain = False


class LESSViTModel(SpatialSpectralLowRankViTEncoder):
    config_class = LESSViTConfig
    main_input_name = "optical"

    def __init__(self, config):
        super().__init__(config)
        self.post_init()