import torch
import torch.nn as nn
from dataclasses import dataclass
from transformers.quantizers import HfQuantizer, register_quantizer
from transformers.utils.quantization_config import QuantizationConfigMixin, register_quantization_config

# =====================================================================
# ۱. تعریف و ثبت کلاس کانفیگ سفارشی در اکوسیستم هاگینگ فیس
# =====================================================================
@register_quantization_config("adaptive_quant")
@dataclass
class AdaptiveQuantConfig(QuantizationConfigMixin):
    """
    تنظیمات مربوط به فریم‌ورک کوانتیزاسیون پویای زمان اجرا (Adaptive Quantization)
    """
    def __init__(
        self,
        quant_method="adaptive_quant",
        fp32_threshold=0.7,
        int8_threshold=0.3,
        scorer_weight_path=None,
        **kwargs
    ):
        self.quant_method = quant_method
        self.fp32_threshold = fp32_threshold
        self.int8_threshold = int8_threshold
        self.scorer_weight_path = scorer_weight_path
        super().__init__(**kwargs)


# =====================================================================
# ۲. شبیه‌سازی لایه خطی تطبیقی (محیط داینامیک زمان اجرا)
# =====================================================================
class AdaptiveQuantizedLinear(nn.Module):
    """
    لایه‌ای سفارشی که جایگزین nn.Linear معمولی شده و بر اساس امتیاز Scorer،
    دقت محاسبات (FP32, INT8, INT4) را به صورت پویا انتخاب می‌کند.
    """
    def __init__(self, in_features, out_features, bias=True, config=None):
        super().__init__()
        self.in_features = in_features
        self.out_features = out_features
        self.config = config
        
        # تعریف وزن پیش‌فرض (این وزن‌ها هنگام لود مدل با مقادیر اصلی پر می‌شوند)
        self.weight = nn.Parameter(torch.empty(out_features, in_features))
        if bias:
            self.bias = nn.Parameter(torch.empty(out_features))
        else:
            self.register_parameter('bias', None)
            
        # شبکه‌ مینیاتوری امتیازدهی به ورودی (<1% overhead)
        self.input_complexity_scorer = nn.Sequential(
            nn.Linear(in_features, 16),
            nn.ReLU(),
            nn.Linear(16, 1),
            nn.Sigmoid()
        )

    def forward(self, x):
        # ۱. محاسبه امتیاز سختی ورودی (بین 0.0 تا 1.0)
        with torch.no_grad():
            mean_input = x.mean(dim=0) if x.dim() > 2 else x
            score = self.input_complexity_scorer(mean_input).mean().item()
        
        # ۲. مسیریابی پویا بر اساس امتیاز (Dynamic Routing)
        if score > self.config.fp32_threshold:
            # مسیر فرکانس بالا / سخت: پردازش کامل با دقت FP32
            return nn.functional.linear(x, self.weight, self.bias)
            
        elif score > self.config.int8_threshold:
            # مسیر متوسط: کوانتیزاسیون و محاسبه با دقت INT8
            scale = self.weight.abs().max() / 127.0
            quant_weight = torch.clamp(torch.round(self.weight / scale), -128, 127)
            dequant_weight = quant_weight * scale
            return nn.functional.linear(x, dequant_weight.to(x.dtype), self.bias)
            
        else:
            # مسیر ساده: کوانتیزاسیون فشرده با دقت INT4
            scale = self.weight.abs().max() / 7.0
            quant_weight = torch.clamp(torch.round(self.weight / scale), -8, 7)
            dequant_weight = quant_weight * scale
            return nn.functional.linear(x, dequant_weight.to(x.dtype), self.bias)


# =====================================================================
# ۳. تعریف و ثبت کلاس کوانتیزر اصلی در اکوسیستم هاگینگ فیس
# =====================================================================
@register_quantizer("adaptive_quant")
class AdaptiveHfQuantizer(HfQuantizer):
    """
    مدیریت تزریق لایه AdaptiveQuantizedLinear به جای لایه‌های خطی استاندارد مدل
    """
    requires_calibration = False  # پروژه شما کاملاً پویا است و نیاز به کالیبراسیون آفلاین ندارد

    def __init__(self, quantization_config, **kwargs):
        super().__init__(quantization_config, **kwargs)
        self.quantization_config = quantization_config

    def _process_model_before_weight_loading(self, model, **kwargs):
        """
        قبل از پر شدن مدل با وزن‌ها، تمام لایه‌های nn.Linear را پیدا کرده 
        و آن‌ها را با لایه پویای پروژه شما جایگزین می‌کند.
        """
        for name, module in model.named_modules():
            if isinstance(module, nn.Linear):
                # نادیده گرفتن لایه نهایی خروجی (LM Head) برای حفظ کیفیت متن تولیدی
                if "lm_head" in name:
                    continue
                    
                parent_name = ".".join(name.split(".")[:-1])
                child_name = name.split(".")[-1]
                parent = model.get_submodule(parent_name) if parent_name else model
                
                # ساخت لایه تطبیقی جدید
                new_layer = AdaptiveQuantizedLinear(
                    in_features=module.in_features,
                    out_features=module.out_features,
                    bias=module.bias is not None,
                    config=self.quantization_config
                )
                
                # جایگزینی لایه قدیمی با لایه جدید در ساختار مدل
                setattr(parent, child_name, new_layer)

    def _process_model_after_weight_loading(self, model, **kwargs):
        """
        پس از بارگذاری وزن‌ها، در صورت وجود وزن‌های ذخیره‌شده برای Scorer، آن‌ها را اعمال می‌کند.
        """
        if self.quantization_config.scorer_weight_path:
            pass
        return model

# =====================================================================
# ۴. تزریق ایمن به نگاشت داخلی پلتفرم (پشتیبانی از تمام نسخه‌های Transformers)
# =====================================================================
try:
    from transformers.quantizers import AUTO_QUANTIZER_MAPPING
    if "adaptive_quant" not in AUTO_QUANTIZER_MAPPING:
        AUTO_QUANTIZER_MAPPING["adaptive_quant"] = AdaptiveHfQuantizer
except ImportError:
    pass