import torch
import numpy as np
from scipy.linalg import hadamard
import transformer_engine.pytorch as te
from transformers import LlamaForCausalLM, LlamaTokenizer
from torch.utils.data import DataLoader
from datasets import load_dataset
from torch.cuda.amp import autocast

# FP4 양자화 클래스
class FP4Quantizer:
    def __init__(self, exponent_bits=2, mantissa_bits=1):
        self.exponent_bits = exponent_bits
        self.mantissa_bits = mantissa_bits
        self.max_value = 6.0
        self.levels = 2 ** (exponent_bits + mantissa_bits)
        self.lookup_table = self._create_lookup_table()

    def _create_lookup_table(self):
        values = []
        for sign in [-1, 1]:
            for e in range(2 ** self.exponent_bits):
                for m in range(2 ** self.mantissa_bits):
                    value = sign * (1 + m / 2 ** self.mantissa_bits) * 2 ** (e - 1)
                    if abs(value) <= self.max_value:
                        values.append(value)
        return torch.tensor(values, dtype=torch.float32).cuda()

    def quantize(self, tensor):
        n = tensor.shape[-1]
        if n % 2 != 0:
            n = 2 ** int(np.ceil(np.log2(n)))
            tensor = torch.nn.functional.pad(tensor, (0, n - tensor.shape[-1]))
        H = torch.tensor(hadamard(n), dtype=torch.float32).cuda()
        transformed = torch.matmul(tensor, H)
        quantized = torch.zeros_like(transformed)
        for i in range(transformed.shape[0]):
            for j in range(transformed.shape[1]):
                value = transformed[i, j]
                closest = self.lookup_table[torch.argmin(torch.abs(self.lookup_table - value))]
                quantized[i, j] = closest
        return torch.matmul(quantized, H.t())[:, :tensor.shape[-1]]

# 확률적 반올림
def stochastic_round(tensor, quantizer):
    quantized = quantizer.quantize(tensor)
    floor = torch.floor(tensor / quantizer.max_value * quantizer.levels) / quantizer.levels * quantizer.max_value
    ceil = torch.ceil(tensor / quantizer.max_value * quantizer.levels) / quantizer.levels * quantizer.max_value
    prob = (tensor - floor) / (ceil - floor + 1e-10)
    mask = torch.rand_like(tensor) < prob
    return torch.where(mask, ceil, floor)

# 스케일링 법칙
def scaling_law_loss(model_size, data_size, loss):
    eff_N = model_size / loss
    eff_D = data_size / loss
    return eff_N, eff_D

# FP4 선형 레이어
class FP4Linear(te.Linear):
    def __init__(self, in_features, out_features):
        super().__init__(in_features, out_features, precision="fp4")
        self.quantizer = FP4Quantizer()

    def forward(self, input):
        with autocast(dtype=torch.float4_e2m1):
            weight = self.quantizer.quantize(self.weight)
            output = torch.matmul(input, weight.t())
            if self.bias is not None:
                output += self.bias
            return output

# 모델 및 데이터 준비
model = LlamaForCausalLM.from_pretrained("meta-llama/Llama-2-7b").cuda()
tokenizer = LlamaTokenizer.from_pretrained("meta-llama/Llama-2-7b")
dataset = load_dataset("c4", split="train")
dataloader = DataLoader(dataset, batch_size=8)

# FP4 선형 레이어로 교체
for name, module in model.named_modules():
    if isinstance(module, torch.nn.Linear):
        setattr(model, name, FP4Linear(module.in_features, module.out_features))

# 훈련 설정
optimizer = torch.optim.Adam(model.parameters(), lr=1e-4)
quantizer = FP4Quantizer()

# 훈련 루프
for epoch in range(3):
    for batch in dataloader:
        inputs = tokenizer(batch["text"], return_tensors="pt", truncation=True, padding=True).to("cuda")
        outputs = model(**inputs)
        loss = outputs.loss
        optimizer.zero_grad()
        loss.backward()
        for param in model.parameters():
            if param.grad is not None:
                param.grad = stochastic_round(param.grad, quantizer)
        optimizer.step()
        eff_N, eff_D = scaling_law_loss(7e9, len(dataset), loss.item())
        print(f"Epoch {epoch}, Loss: {loss.item()}, eff_N: {eff_N}, eff_D: {eff_D}")