from datasets import load_dataset, Audio
from transformers import WhisperProcessor,WhisperForConditionalGeneration
from transformers.models.whisper.tokenization_whisper import TO_LANGUAGE_CODE  #可查看whisper支持语言
import soundfile as sf
import io
from typing import Any, Dict, List, Union
from dataclasses import dataclass
import torch
from torch.utils.data import DataLoader
from torch.optim import AdamW
from transformers import get_scheduler
from tqdm import tqdm


processor = WhisperProcessor.from_pretrained(
    "openai/whisper-small", language="albanian", task="transcribe")

alb_dataset = load_dataset("parquet", data_files={'train': 'train.parquet', 'test': 'test.parquet'})
alb_dataset = alb_dataset.cast_column("audio", Audio(decode=False))

def prepare_dataset(example):
    audio_bytes = example["audio"]['bytes']
    audio, samplerate = sf.read(io.BytesIO(audio_bytes))

    example = processor(audio=audio,
                        sampling_rate=16000,
                        text=example["sentence"])
    example["input_length"] = len(audio) / samplerate
    return example
#此处alb_dataset.column_names["train"]是返回的所有column的名字，指删去原datasetdict所有同名列，不是指删去原训练集
dataset = alb_dataset.map(prepare_dataset,remove_columns=alb_dataset.column_names["train"],
                            num_proc=1,
                            batch_size=100)
# sample = dataset["train"][1] #feature形状:(1,80,3000)的list

''' 
@dataclass等同于
def __init__(self, processor):
self.processor = processor
'''
@dataclass
class DataCollatorSpeechSeq2SeqWithPadding:
    processor: Any  #不强制约束类型

    #令对象可像函数一样直接调用
    def __call__(
            self, features: List[Dict[str, Union[List[int], torch.Tensor]]] #输入为一个list，内容为一个字典，键为str，值为list或tensor之一
    ) -> Dict[str, torch.Tensor]: #返回值为一个字典，类似C语言的函数定义

        input_features = [
            {"input_features": feature["input_features"][0]} for feature in features #feature["input_features"]形状为(1,80,3000),因此要取[0]
        ]
        batch = self.processor.feature_extractor.pad(input_features, return_tensors="pt") #processor已经填充过，此处仅为转tensor
        #feature_extractor.pad接口协议是List[{"input_features": (feature_dim, time)}]

        #label同理，但是label是没有提前填充的
        label_features = [{"input_ids": feature["labels"]} for feature in features]
        labels_batch = self.processor.tokenizer.pad(label_features, return_tensors="pt")
        #返回{"input_ids": Tensor(batch, max_len),"attention_mask": Tensor(batch, max_len)}
        labels = labels_batch["input_ids"].masked_fill(
            labels_batch.attention_mask.ne(1), -100     #pytorch 交叉熵默认-100不参与loss计算
        )#用-100代替填充词元，labels_batch.attention_mask.ne(1)可以根据attention_mask生成一个bool mask，pad token为True

        # 如果在之前分词时添加了 bos 词元，那就剪切掉，因为之后还会加上的
        if (labels[:, 0] == self.processor.tokenizer.bos_token_id).all().cpu().item():
            labels = labels[:, 1:]

        batch["labels"] = labels

        return batch

data_collator = DataCollatorSpeechSeq2SeqWithPadding(processor=processor)
# 定义 DataLoader
train_dataloader = DataLoader(
    dataset["train"],
    batch_size=2,
    shuffle=True,
    collate_fn=data_collator
)
eval_dataloader = DataLoader(
    dataset["test"],
    batch_size=2,
    collate_fn=data_collator
)

device = torch.device("cuda")
model = WhisperForConditionalGeneration.from_pretrained("openai/whisper-small")
model.to(device)

optimizer = AdamW(model.parameters(), lr=5e-5)
num_epochs = 3
num_training_steps = num_epochs * len(train_dataloader)
lr_scheduler = get_scheduler(
    "linear",
    optimizer=optimizer,
    num_warmup_steps=0,
    num_training_steps=num_training_steps
)

for epoch in range(num_epochs):
    model.train()
    train_loop = tqdm(enumerate(train_dataloader), total=len(train_dataloader),
                 desc=f'Train_Epoch {epoch + 1}/{num_epochs}', unit='batch')
    for idx, batch in train_loop:
        batch = {k: v.to(device) for k, v in batch.items()}


        outputs = model(**batch)
        loss = outputs.loss

        loss.backward()

        optimizer.step()
        lr_scheduler.step()
        optimizer.zero_grad()
        train_loop.set_postfix({'loss': '{0:1.5f}'.format(loss.item())})

    model.eval()
    total_loss = 0
    eval_loop = tqdm(enumerate(eval_dataloader), total=len(eval_dataloader),
                 desc=f'Eval_Epoch {epoch + 1}/{num_epochs}', unit='batch')
    for idx, batch in eval_loop:
        batch = {k: v.to(device) for k, v in batch.items()}
        with torch.no_grad():
            outputs = model(**batch)

        total_loss += outputs.loss.item()

    avg_loss = total_loss / len(eval_dataloader)
    print(f"Evaluation Loss: {avg_loss}")

from transformers.models.whisper.english_normalizer import BasicTextNormalizer

normalizer = BasicTextNormalizer()#文本标准化，重要


def compute_metrics(pred):
    pred_ids = pred.predictions
    label_ids = pred.label_ids

    # 用 pad_token_id 替换 -100
    label_ids[label_ids == -100] = processor.tokenizer.pad_token_id

    # 我们希望在计算指标时不要组合起词元（解码结果为词元(subword)，而不是完整单词）
    pred_str = processor.batch_decode(pred_ids, skip_special_tokens=True)
    label_str = processor.batch_decode(label_ids, skip_special_tokens=True)

    # 计算普通的 WER
    wer_ortho = 100 * metric.compute(predictions=pred_str, references=label_str)

    # 计算标准化的 WER
    pred_str_norm = [normalizer(pred) for pred in pred_str]
    label_str_norm = [normalizer(label) for label in label_str]
    # 过滤，从而在评估时只计算 reference 非空的样本
    pred_str_norm = [
        pred_str_norm[i] for i in range(len(pred_str_norm)) if len(label_str_norm[i]) > 0
    ]
    label_str_norm = [
        label_str_norm[i]
        for i in range(len(label_str_norm))
        if len(label_str_norm[i]) > 0
    ]

    wer = 100 * metric.compute(predictions=pred_str_norm, references=label_str_norm)

    return {"wer_ortho": wer_ortho, "wer": wer}

#对于transformers库加载模型：
model = WhisperForConditionalGeneration.from_pretrained("openai/whisper-small",
                                                       ,torch_dype=torch.bfloat16)
#对于nn.Moudle对象：
model = MyModel() # 默认 FP32
model.to(torch.float16) # 整体转为 FP16

model = WhisperForConditionalGeneration.from_pretrained("openai/whisper-small").to("cuda")
optimizer = torch.optim.AdamW(model.parameters(), lr=1e-5)

# 初始化梯度缩放器，防止 FP16 梯度下溢
scaler = GradScaler()

for batch in dataloader:
    optimizer.zero_grad()
    
    # 1. 在 autocast 上下文中运行前向传播
    with autocast(dtype=torch.float16):
        outputs = model(input_features=batch["input_features"], labels=batch["labels"])
        loss = outputs.loss

    # 2. 使用 scaler 缩放 loss 并反向传播
    scaler.scale(loss).backward()
    
    # 3. 更新参数
    scaler.step(optimizer)
    scaler.update()

from peft import LoraConfig
config = LoraConfig(r=32, 
                    lora_alpha=64, 
                    target_modules=["q_proj", "v_proj"],
                    lora_dropout=0.05)

from peft import LoraConfig, get_peft_model

model = WhisperForConditionalGeneration.from_pretrained("openai/whisper-small")
config = LoraConfig(r=32, 
                    lora_alpha=64, 
                    target_modules=["q_proj", "v_proj"],
                    lora_dropout=0.05)

model = get_peft_model(model, config)
# 打印训练参数
model.print_trainable_parameters()

'''
结果
trainable params: 3,538,944 || all params: 245,273,856 || trainable%: 1.4429
即在微调中只有1.44%的参数参与训练
'''

from transformers import BitsAndBytesConfig
from peft import  get_peft_model, prepare_model_for_kbit_training

# 配置 4-bit 量化
bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_compute_dtype=torch.bfloat16, 
    bnb_4bit_quant_type="nf4",            # 使用正态浮点 4 位，精度更好
    bnb_4bit_use_double_quant=True        # 嵌套量化，进一步省显存
)

model = WhisperForConditionalGeneration.from_pretrained(
    "openai/whisper-small", 
    quantization_config=bnb_config
)

#量化后的权重是不可训练的。如果你要微调，必须配合 LoRA 使用！！
model = prepare_model_for_kbit_training(model)
...
lora_model = get_peft_model(model, lora_config)

from transformers import pipeline

classifier = pipeline("sentiment-analysis")
classifier(
    [
        "I've been waiting for a HuggingFace course my whole life.",
        "I hate this so much!",
    ]
)

>>>[{'label': 'POSITIVE', 'score': 0.9598047137260437},
>>> {'label': 'NEGATIVE', 'score': 0.9994558095932007}]

from transformers import AutoTokenizer, AutoModelForSequenceClassification

checkpoint = "distilbert-base-uncased-finetuned-sst-2-english"
tokenizer = AutoTokenizer.from_pretrained(checkpoint)
model = AutoModelForSequenceClassification.from_pretrained(checkpoint)
sequences = ["I've been waiting for a HuggingFace course my whole life.", "So have I!"]

tokens = tokenizer(sequences, padding=True, truncation=True, return_tensors="pt")

#仅分词
tokenizer = AutoTokenizer.from_pretrained("bert-base-cased")
sequence = "Using a Transformer network is simple"
tokens = tokenizer.tokenize(sequence)
>>>['Using', 'a', 'transform', '##er', 'network', 'is', 'simple']

#将分词转为ID
ids = tokenizer.convert_tokens_to_ids(tokens)
print(ids)
>>>[7993, 170, 11303, 1200, 2443, 1110, 3014]

#解码
decoded_string = tokenizer.decode([7993, 170, 11303, 1200, 2443, 1110, 3014])
print(decoded_string)
>>>'Using a Transformer network is simple'

# 将句子序列填充到最长句子的长度
model_inputs = tokenizer(sequences, padding="longest")

# 将句子序列填充到模型的最大长度
# (512 for BERT or DistilBERT)
model_inputs = tokenizer(sequences, padding="max_length")

# 将句子序列填充到指定的最大长度
model_inputs = tokenizer(sequences, padding="max_length", max_length=8)

# 将截断比模型最大长度长的句子序列
# (512 for BERT or DistilBERT)
model_inputs = tokenizer(sequences, truncation=True)

# 将截断长于指定最大长度的句子序列
model_inputs = tokenizer(sequences, max_length=8, truncation=True)

# 返回 PyTorch tensors
model_inputs = tokenizer(sequences, padding=True, return_tensors="pt")

# 返回 TensorFlow tensors
model_inputs = tokenizer(sequences, padding=True, return_tensors="tf")

# 返回 NumPy arrays
model_inputs = tokenizer(sequences, padding=True, return_tensors="np")

from transformers import BertConfig, BertModel

config = BertConfig()
model = BertModel(config)

model = BertModel.from_pretrained("bert-base-cased")

#或是直接通过AutoModel类
checkpoint = "bert-base-cased"
model = AutoModel.from_pretrained(checkpoint)

model(**tokens) 
"""
python的字典解包机制，等同于model(
    input_ids=tensor(...),
    attention_mask=tensor(...)
"""

from datasets import load_dataset

raw_datasets = load_dataset("glue", "mrpc")
raw_train_dataset = raw_datasets["train"]
sentence1 = raw_train_dataset['sentence1']
sample = raw_train_dataset[0]

def tokenize_function(examples):
    return tokenizer(
        examples["sentence1"],
        examples["sentence2"],
        truncation=True
    )
tokenized_datasets = raw_datasets.map(tokenize_function, batched=True)

samples = tokenized_datasets["train"][:]
samples = {k: v for k, v in samples.items() if k not in ["idx", "sentence1", "sentence2"]}

tokenized_datasets = tokenized_datasets.remove_columns(["sentence1", "sentence2", "idx"])

from transformers import DataCollatorWithPadding

data_collator = DataCollatorWithPadding(tokenizer=tokenizer)
samples = tokenized_datasets["train"][:8]
samples = {k: v for k, v in samples.items() if k not in ["idx", "sentence1", "sentence2"]}
batch = data_collator(samples)

from datasets import load_dataset
from transformers import AutoTokenizer, DataCollatorWithPadding, AutoModelForSequenceClassification
from transformers import TrainingArguments, Trainer


raw_datasets = load_dataset("glue", "mrpc")
checkpoint = "bert-base-uncased"
tokenizer = AutoTokenizer.from_pretrained(checkpoint)


def tokenize_function(example):
    return tokenizer(example["sentence1"], example["sentence2"], truncation=True)


tokenized_datasets = raw_datasets.map(tokenize_function, batched=True)
data_collator = DataCollatorWithPadding(tokenizer=tokenizer)
training_args = TrainingArguments("test-trainer")
model = AutoModelForSequenceClassification.from_pretrained(checkpoint, num_labels=2)

trainer = Trainer(
    model,
    training_args,
    train_dataset=tokenized_datasets["train"],
    eval_dataset=tokenized_datasets["validation"],
    data_collator=data_collator,
    tokenizer=tokenizer,
)

trainer.train()

predictions = trainer.predict(tokenized_datasets["validation"])

def compute_metrics(eval_preds):
    metric = evaluate.load("glue", "mrpc")
    logits, labels = eval_preds
    predictions = np.argmax(logits, axis=-1)
    return metric.compute(predictions=predictions, references=labels)

trainer = Trainer(
    model,
    training_args,
    train_dataset=tokenized_datasets["train"],
    eval_dataset=tokenized_datasets["validation"],
    data_collator=data_collator,
    tokenizer=tokenizer,
    compute_metrics=compute_metrics,
)

tokenized_datasets = raw_datasets.map(tokenize_function, batched=True)
data_collator = DataCollatorWithPadding(tokenizer=tokenizer)

#删去不需要列，将label改名为labels（模型默认输入格式）
tokenized_datasets = tokenized_datasets.remove_columns(["sentence1", "sentence2", "idx"])
tokenized_datasets = tokenized_datasets.rename_column("label", "labels")
tokenized_datasets.set_format("torch")
#>>>["attention_mask", "input_ids", "labels", "token_type_ids"]

from torch.utils.data import DataLoader
#Dataloader兼容transformers的DataCollatorWithPadding
train_dataloader = DataLoader(
    tokenized_datasets["train"], shuffle=True, batch_size=8, collate_fn=data_collator
)
eval_dataloader = DataLoader(
    tokenized_datasets["validation"], batch_size=8, collate_fn=data_collator
)

from transformers import get_scheduler
import torch

optimizer = AdamW(model.parameters(), lr=5e-5)
num_epochs = 3
num_training_steps = num_epochs * len(train_dataloader)
lr_scheduler = get_scheduler(	#学习率线性衰减
    "linear",
    optimizer=optimizer,
    num_warmup_steps=0,
    num_training_steps=num_training_steps,
)
model.to(device)

for epoch in range(num_epoch):
    model.train()
    for batch in train_dataloader:
        batch = {k: v.to(device) for k, v in batch.items()}
        outputs = model(**batch)
        loss = outputs.loss
        loss.backward()

        optimizer.step()
        lr_scheduler.step()
        optimizer.zero_grad()
    
    model.eval()
	for batch in eval_loader:
        batch = {k: v.to(device) for k, v in batch.items()}
    	with torch.no_grad():
        	outputs = model(**batch)
        logits = outputs.logits
        ...

joint_states = p.getJointStates(robotId,controllable_joints)
q_actual = np.array([state[0] for state in joint_states])
qd_actual = np.array([state[1] for state in joint_states])


q_e = q[n] - q_actual
qd_e = qd[n] - qd_actual


aq = qdd[n] + 400 * q_e + 40 * qd_e

tau = p.calculateInverseDynamics(robotId,list(q_actual),list(qd_actual),list(aq))
p.setJointMotorControlArray(robotId,controllable_joints,p.TORQUE_CONTROL,
                            forces = tau,
                            )