# Copyright 2026 Angshul Majumdar. Licensed under the Apache License, Version 2.0.
"""TensorPEFTConfig."""
import json
import os
from dataclasses import asdict, dataclass, field
from typing import List, Optional, Union

CONFIG_NAME = "adapter_config.json"
WEIGHTS_NAME = "adapter_model.safetensors"


@dataclass
class TensorPEFTConfig:
    """Configuration of a TensorPEFT adapter.

    Args:
        atoms_per_group: rank-one CP atoms per tensor group (all chosen at initialisation).
            Trainable parameters = sum over groups of atoms_per_group * (1 + sum of the group's mode sizes).
        target_modules: ``"all-linear"`` (every Linear / Conv1D inside the transformer blocks) or a
            list of module-name suffixes, e.g. ``["q_proj", "v_proj"]`` (``"c_attn.q"`` etc. for GPT-2).
        exclude_modules: substrings of module names that are never adapted.
        split_fused_qkv: split GPT-2-style fused ``c_attn`` = [q|k|v] into q, k, v.
        hopm_iters: higher-order power iterations used to compute each atom.
        alpha_scale: coefficients are parametrised as ``alpha_scale * alpha_raw`` so that a single
            optimiser learning rate (about 3e-3 for Adam/AdamW) suits both coefficients and factors.
        init_batches: default number of batches used by ``init_atoms``.
    """
    atoms_per_group: int = 4
    target_modules: Union[str, List[str]] = "all-linear"
    exclude_modules: List[str] = field(default_factory=lambda: ["lm_head", "embed", "wte", "wpe"])
    split_fused_qkv: bool = True
    hopm_iters: int = 30
    alpha_scale: float = 100.0 / 3.0
    init_batches: int = 4
    base_model_name_or_path: Optional[str] = None
    peft_type: str = "TENSORPEFT"
    task_type: str = "CAUSAL_LM"
    tensorpeft_version: Optional[str] = None

    def to_dict(self):
        return asdict(self)

    def save_pretrained(self, save_directory, extra=None):
        os.makedirs(save_directory, exist_ok=True)
        d = self.to_dict()
        if extra:
            d.update(extra)
        with open(os.path.join(save_directory, CONFIG_NAME), "w") as f:
            json.dump(d, f, indent=2)

    @classmethod
    def from_dict(cls, d):
        names = {f for f in cls.__dataclass_fields__}
        return cls(**{k: v for k, v in d.items() if k in names})

    @classmethod
    def from_pretrained(cls, model_id, **hub_kwargs):
        path = _resolve(model_id, CONFIG_NAME, **hub_kwargs)
        with open(path) as f:
            d = json.load(f)
        return cls.from_dict(d), d


def _resolve(model_id, filename, **hub_kwargs):
    local = os.path.join(model_id, filename)
    if os.path.isfile(local):
        return local
    from huggingface_hub import hf_hub_download
    return hf_hub_download(repo_id=model_id, filename=filename, **hub_kwargs)