import os
import pprint
import random
import re
from collections import defaultdict

import numpy as np
import torch
import torch.nn as nn
import tqdm
from PIL import Image
from datasets import Dataset as HFDataset
from sklearn.metrics import classification_report, f1_score
from torch.nn import BCELoss
from torch.functional import F
from torch.utils.data import Dataset
# 加载 SwinTransformer 模型和处理器
from transformers import (AutoFeatureExtractor, SwinModel, Trainer,
                          TrainingArguments)
from transformers.models.swin.modeling_swin import SwinImageClassifierOutput
from torchvision import transforms

os.environ["WANDB_DISABLED"] = "true"
MODEL_NAME_OR_PATH = './swin-large-patch4-window12-384-in22k/'


transform = transforms.Compose([
    transforms.ColorJitter(brightness=0.2, contrast=0.2,
                           saturation=0.2, hue=0.1),
    transforms.RandomResizedCrop(384),
])

def load_dataset_from_path():
    result = []
    image_dirs = ['./images', ]
    for image_dir in image_dirs:
        for image_name in tqdm.tqdm(os.listdir(image_dir), total=len(os.listdir(image_dir)), postfix=image_dir):
            label = re.search(r"l\-(是|否)\.", image_name).group(1)
            image_path = os.path.join(image_dir, image_name)
            try:
                transform(Image.open(image_path).convert("RGB"))
            except Exception as e:
                print(f"invalid image: {image_path}, err:{e}")
                continue

            result.append({"image_path": image_path, "label": label})

    # #################### split
    random.shuffle(result)

    counter = defaultdict(int)
    for item in result:
        counter[item['label']] += 1
    pprint.pprint(counter)

    trains, tests = [], []
    train_counter = defaultdict(int)
    for item in result:
        if train_counter[item['label']] / counter[item['label']] > 0.85:
            tests.append(item)
        else:
            train_counter[item['label']] += 1
            trains.append(item)
    return HFDataset.from_list(trains), HFDataset.from_list(tests)
    # all_dataset = HFDataset.from_list(result)
    #
    # split = all_dataset.train_test_split(0.15, seed=1)
    # return split['train'], split['test']


train_dataset, eval_dataset = load_dataset_from_path()
print(
    f'total train size:{train_dataset.shape[0]},eval size:{eval_dataset.shape[0]}')


class SwinForClassify(nn.Module):
    def __init__(self, model_name="microsoft/swin-large-patch4-window12-384-in22k", num_labels=1):
        super().__init__()
        self.backbone = SwinModel.from_pretrained(model_name)  # 只加载特征提取部分
        hidden_size = self.backbone.config.hidden_size  # Swin的最后一层特征维度
        self.classify = nn.Linear(hidden_size, num_labels)  # 线性回归层

    def forward(self, pixel_values, labels=None):
        outputs = self.backbone(pixel_values)  # 提取特征
        pooled_output = outputs.pooler_output  # (batch_size, hidden_size)
        logits = torch.sigmoid(self.classify(pooled_output))  # 线性映射到目标维度
        loss = None
        if labels is not None:
            cri = BCELoss()
            loss = cri(logits, labels)

        return SwinImageClassifierOutput(
            loss=loss,
            logits=logits,
            hidden_states=outputs.hidden_states,
            attentions=outputs.attentions,
            reshaped_hidden_states=outputs.reshaped_hidden_states,

        )


# 示例
model = SwinForClassify(model_name=MODEL_NAME_OR_PATH, num_labels=1)

processor = AutoFeatureExtractor.from_pretrained(MODEL_NAME_OR_PATH)


class ClassifyDataset(Dataset):
    def __init__(self, dataset, processor, transform=None):
        self.dataset = dataset
        self.processor = processor
        self.transform = transform

    def __len__(self):
        return len(self.dataset)

    def __getitem__(self, idx):
        image = self.dataset[idx]['image_path']
        image = Image.open(image).convert('RGB')
        image = self.transform(image)

        label = self.dataset[idx]['label']
        if label == '是':
            label = 1
        elif label == '否':
            label = 0
        else:
            raise ValueError(f'{label} not valid.')
        return {"image": image, "label": label}


# 创建数据集对象
train_dataset = ClassifyDataset(train_dataset, processor)
eval_dataset = ClassifyDataset(eval_dataset, processor)

# 定义训练参数
training_args = TrainingArguments(
    output_dir="./results",  # 输出结果目录
    overwrite_output_dir=True,
    num_train_epochs=10,  # 训练轮数
    per_device_train_batch_size=4,  # 每设备训练的batch size
    per_device_eval_batch_size=8,  # 每设备评估的batch size
    evaluation_strategy="epoch",  # 每个epoch评估一次
    save_strategy="epoch",  # 每个epoch保存一次
    logging_dir='./logs',  # 日志目录
    logging_steps=10,
    learning_rate=2e-5,  # 学习率
    weight_decay=0.01,  # 权重衰减
    load_best_model_at_end=True,  # 训练结束时加载最好的模型
    remove_unused_columns=False,
    metric_for_best_model='f1',
)


def data_collator(batch):
    images = [item["image"] for item in batch]
    labels = [item["label"] for item in batch]
    inputs = processor(images, return_tensors="pt")
    inputs["labels"] = torch.tensor(labels, dtype=torch.float).reshape(-1,1)
    return inputs


def compute_metrics(eval_preds):
    preds, labels = [], []
    acc, total = 0, 0

    for (pred, label) in zip(eval_preds.predictions, eval_preds.label_ids):
        if pred > 0.5:
            pred = 1
        else:
            pred = 0
        preds.append(pred)
        labels.append(label)

        total += 1
        if pred == label:
            acc += 1
    print(f'ACC: {acc/(total+1e-5)}')
    print(classification_report(y_true=labels, y_pred=preds, ))
    return {"f1": f1_score(y_true=labels, y_pred=preds, average='macro')}


# 定义 Trainer
trainer = Trainer(
    model=model,  # 预训练模型
    args=training_args,  # 训练参数
    train_dataset=train_dataset,  # 训练数据集
    eval_dataset=eval_dataset,  # 验证数据集
    tokenizer=processor,  # 使用Swin图像处理器
    compute_metrics=compute_metrics,
    data_collator=data_collator
)

# 训练模型
#trainer.train(resume_from_checkpoint=True)
trainer.train()

# 保存模型
trainer.save_model("./final_model")

import torch
import torch.nn as nn
import tqdm
from PIL import Image
from datasets import Dataset as HFDataset
from sklearn.metrics import classification_report, f1_score
from torch.nn import BCELoss
from torch.functional import F
from torch.utils.data import Dataset
# 加载 SwinTransformer 模型和处理器
from transformers import (AutoFeatureExtractor, SwinModel, Trainer,
                                  TrainingArguments, set_seed)
from transformers.models.swin.modeling_swin import SwinImageClassifierOutput
from torchvision import transforms
from safetensors.torch import load_file

set_seed(1)

transform = transforms.Compose([
    transforms.RandomResizedCrop(384),
])


class SwinForClassify(nn.Module):
    def __init__(self, model_name="microsoft/swin-large-patch4-window12-384-in22k", num_labels=1):
        super().__init__()
        self.backbone = SwinModel.from_pretrained(model_name)  # 只加载特征提取部分
        hidden_size = self.backbone.config.hidden_size  # Swin的最后一层特征维度
        self.classify = nn.Linear(hidden_size, num_labels)  # 线性回归层

    @torch.no_grad()
    def forward(self, pixel_values, labels=None):
        outputs = self.backbone(pixel_values)  # 提取特征
        pooled_output = outputs.pooler_output  # (batch_size, hidden_size)
        logits = torch.sigmoid(self.classify(pooled_output))  # 线性映射到目标维度
        loss = None
        if labels is not None:
            cri = BCELoss()
            loss = cri(logits, labels)

        return SwinImageClassifierOutput(
            loss=loss,
            logits=logits,
            hidden_states=outputs.hidden_states,
            attentions=outputs.attentions,
            reshaped_hidden_states=outputs.reshaped_hidden_states,

        )
      
if __name__ == "__main__":
    MODEL_NAME_OR_PATH = './swin-large-patch4-window12-384-in22k/'
    # 示例
    model = SwinForClassify(model_name=MODEL_NAME_OR_PATH, num_labels=1)
    state_dict = load_file("./final_model/model.safetensors")
    model.load_state_dict(state_dict)
    model.eval()

    processor = AutoFeatureExtractor.from_pretrained(MODEL_NAME_OR_PATH)

    test_image_dir = "../image-rotate/test_images"
    for image_name in os.listdir(test_image_dir):
        image_path = os.path.join(test_image_dir, image_name)
        image = Image.open(image_path).convert("RGB")
        image = transform(image)
        input = processor([image], return_tensors='pt')
        output = model(**input).logits[0]
        print(image_name,output>0.5, output)

import pathlib

import torch
from datasets import load_dataset
from PIL import Image, ImageDraw
from sympy import im

# 这里路径根据实际路径来
cppe5 = load_dataset(
    'parquet',
    data_files={
        "train": pathlib.Path(__file__).parent.parent.joinpath("train-00000-of-00001.parquet").__str__(),
        "test": pathlib.Path(__file__).parent.parent.joinpath("test-00000-of-00001.parquet").__str__()
    }
)
categories = cppe5["train"].features["objects"].feature["category"].names
id2label = {index: x for index, x in enumerate(categories, start=0)}
label2id = {v: k for k, v in id2label.items()}
# ['Coverall', 'Face_Shield', 'Gloves', 'Goggles', 'Mask']
print(categories)


def cppe5_to_yolo_format(examples, data_type='train'):
    # 注意这里，images/train放的是训练图片，images/test放的是测试图片
    # 同理，labels/train放的是对应图片标签
    pathlib.Path(__file__).parent.joinpath(
        f"images/{data_type}").mkdir(exist_ok=True, parents=True)
    pathlib.Path(__file__).parent.joinpath(
        f"labels/{data_type}").mkdir(exist_ok=True, parents=True)
    for index, sample in enumerate(examples):
        image = sample["image"].convert("RGB")
        yolo_labels = []

        # draw = ImageDraw.ImageDraw(image)
        # for bbox, category in zip(sample["objects"]["bbox"], sample["objects"]["category"]):
        #     x, y, w, h = tuple(bbox)
        #     if max(bbox) > 1.0:
        #         x1, y1 = int(x), int(y)
        #         x2, y2 = int(x + w), int(y + h)
        #     else:
        #         raise NotImplementedError
        #     draw.rectangle((x1, y1, x2, y2), outline="red")
        #     draw.text((x1, y1), str(id2label[category]), fill="red")
        # image.show()

        for bbox, category in zip(sample["objects"]["bbox"], sample["objects"]["category"]):

            x, y, w, h = tuple(bbox)
            if max(bbox) > 1.0:
                x1, y1 = int(x), int(y)
                x2, y2 = int(x + w), int(y + h)
            else:
                raise NotImplementedError
            img_w, img_h = image.size
            x_center = (x1 + x2) / 2.0 / img_w
            y_center = (y1 + y2) / 2.0 / img_h
            width = (x2 - x1) / img_w
            height = (y2 - y1) / img_h
            yolo_labels.append(
                f'{category} {x_center} {y_center} {width} {height}'
            )
        image.save(pathlib.Path(__file__).parent.joinpath(
            f"images/{data_type}/{index}.jpg"))
        with open(pathlib.Path(__file__).parent.joinpath(f"labels/{data_type}/{index}.txt"), "w") as f:
            f.write("\n".join(yolo_labels))


cppe5_train = cppe5_to_yolo_format(cppe5["train"], data_type='train')
cppe5_test = cppe5_to_yolo_format(cppe5["test"], data_type='test')

import os.path

from ultralytics import YOLO

os.environ['WANDB_DISABLED'] = 'true'
os.environ['YOLO_OFFLINE'] = 'true'


model_path = '/object_detection/yolo_impl/yolo11n.pt'
assert os.path.exists(model_path)

model = YOLO(model_path)

model.train(
    data=os.path.join(os.path.dirname(__file__), 'data.yaml'),
    epochs=100,             # 训练轮数
    imgsz=640,              # 输入图像尺寸
    batch=64,               # 批量大小，根据显存可调
    workers=4,              # 数据加载线程数

    # 优化器与学习率
    lr0=0.001,              # 初始学习率
    lrf=0.01,               # 最低学习率（cosine decay 最终值）
    optimizer='SGD',        # 优化器
    weight_decay=0.0005,    # 权重衰减
    warmup_epochs=3,        # 热身轮数
    patience=20,            # 提前停止的容忍度

    # 增强与正则
    augment=True,           # 启用数据增强
    mosaic=1.0,
    mixup=0.1,
    hsv_h=0.015,
    hsv_s=0.7,
    hsv_v=0.4,
    scale=0.5,
    translate=0.1,
    degrees=0.0,            # 不启用旋转

    # 损失函数超参
    box=7.5,                # 边界框损失权重
    cls=0.5,                # 类别损失权重

    iou=0.2,                # IoU 阈值（用于损失计算）

    device=0                # 使用GPU:0

)

import pathlib

from ultralytics import YOLO

model = YOLO(pathlib.Path(__file__).parent.joinpath(
    "runs/detect/train/weights/best.pt"
    )
)

results = model.predict(

    source='/object_detection/yolo_impl/images/test',
    save=True

)
# 处理结果
for result in results:
    print(f"检测到的对象数量: {len(result.boxes)}")
    for box, conf, cls in zip(result.boxes.xyxy, result.boxes.conf, result.boxes.cls):
        print(f"类别: {result.names[int(cls)]}, 置信度: {conf:.2f}, 坐标: {box}")
        print(f"类别: {result.names[int(cls)]}, 置信度: {conf:.2f}, 坐标: {box}")

import os

from langchain_core.messages import AIMessage, AnyMessage
from langchain_deepseek import ChatDeepSeek
from langgraph.constants import END, START
from langgraph.graph import StateGraph
from pydantic import BaseModel, Field
from typing_extensions import TypedDict

INTENTION_PROMPT = """根据下面用户对话信息判断用户意图，意图有：查询天气、讲笑话、闲聊。结果请以json格式输出。
例如：
对话信息：
    User: 今天天气怎么样？
输出：
    {"intention": "查询天气"}

真实对话信息：
{{messages}}
"""


class IntentionOutput(BaseModel):
    intention: str = Field(description='用户意图')


llm = ChatDeepSeek(
    model="deepseek-chat",  # 指定 DeepSeek 的模型名称
    api_key=os.getenv("DEEPSEEK_API_KEY"),  # 替换为您自己的 DeepSeek API 密钥
)


class State(TypedDict):
    messages: list[AnyMessage]
    intention: str


# 初始化状态图
graph_builder = StateGraph(State)


def chatbot(state: State):
    """闲聊"""
    return {"messages": AIMessage(content='哦吼，我也不知道聊啥子')}


def get_weather(state: State):
    """查询天气"""
    return {"messages": AIMessage(content='今天天气晴朗，23度，适合外出打羽毛球。')}


def joke(state: State):
    """讲笑话"""
    history_msgs = ""
    for msg in state['messages']:
        history_msgs += f"{msg['role']}:{msg['content']}\n"
    prompt = '请根据用户最新的对话信息讲个笑话：\n' + history_msgs
    output = llm.invoke(prompt)
    return {"messages": output}


def user_intention(state: State):
    """意图识别"""
    history_msgs = ""
    for msg in state['messages']:
        history_msgs += f"{msg['role']}:{msg['content']}\n"
    prompt = INTENTION_PROMPT.replace("{{messages}}", history_msgs)
    print('意图识别'.center(60, '-'))
    print(prompt)
    structure_llm = llm.with_structured_output(IntentionOutput)
    output = structure_llm.invoke(prompt)
    final_intention = output.intention
    if final_intention not in ("查询天气", "讲笑话", "闲聊"):
        final_intention = '闲聊'
    return {"intention": final_intention}


def intention_conditional_edge(state: State):
    if state['intention'] == '查询天气':
        return 'get_weather'
    elif state['intention'] == '讲笑话':
        return 'joke'
    else:
        return 'chatbot'


# 添加节点到状态图

graph_builder.add_node('user_intention', user_intention)
graph_builder.add_node("chatbot", chatbot)
graph_builder.add_node('get_weather', get_weather)
graph_builder.add_node('joke', joke)

graph_builder.add_edge(START, 'user_intention')
graph_builder.add_conditional_edges('user_intention', intention_conditional_edge, [
    'chatbot', 'get_weather', 'joke'])
graph_builder.add_edge('chatbot', END)
graph_builder.add_edge('get_weather', END)
graph_builder.add_edge('joke', END)

# 编译状态图
graph = graph_builder.compile()

# with open("/tmp/b.png", 'wb') as f:
#     f.write(graph.get_graph().draw_mermaid_png())

if __name__ == "__main__":
    round1 = graph.invoke(
        {"messages": [{"role": "user", "content": "我很开心，讲个好笑的事情吧"}]})
    print(round1)

# Python >= 3.11 is required.

pip install --upgrade "langgraph-cli[inmem]"

DEEPSEEK_API_KEY=<your key>  langgraph dev

from langgraph_sdk import get_sync_client

client = get_sync_client(url="http://localhost:2024")

for chunk in client.runs.stream(
    None,  # Threadless run
    "agent",  # Name of assistant. Defined in langgraph.json.
    input={
        "messages": [
            {
                "role": "user",
                "content": "今天南京天气好吗？适合外出吗？",
            },
            {
                "role": "assistant",
                "content": "今天天气晴朗，23度，适合外出打羽毛球。"
            },
            {
                "role": "user",
                "content": "那你给我讲个天气相关的笑话吧"
            }
        ],
    },
    stream_mode="messages",
):
    print(f"Receiving new event of type: {chunk.event}...")
    print(chunk.data)
    print("\n\n")

import os
os.environ['HF_ENDPOINT'] = 'https://hf-mirror.com'  # noqa

import requests
from PIL import Image
from transformers import AutoProcessor, LlavaForConditionalGeneration, LlavaConfig


config = LlavaConfig.from_pretrained("llava-hf/llava-1.5-7b-hf")


model = LlavaForConditionalGeneration(config=config)

processor = AutoProcessor.from_pretrained("llava-hf/llava-1.5-7b-hf")
prompt = "USER: <image>\nWhat's the content of the image? ASSISTANT:"
url = "https://www.ilankelman.org/stopsigns/australia.jpg"
image = Image.open(requests.get(url, stream=True).raw)
inputs = processor(text=prompt, images=image, return_tensors="pt")

generate_ids = model.generate(**inputs, max_new_tokens=15)
processor.batch_decode(generate_ids, skip_special_tokens=True,
                       clean_up_tokenization_spaces=False)[0]

print()

(336-14)/14+1 = 24
24 * 24 = 596
# 如果考虑CLIP CLS的话就是597。

(final_embedding[:, 5:576+5, :] == image_features[:, :, :]).all()
# > tensor(True)