from PIL import Image
import torch
import os
import json
import torch.distributed as dist
from tqdm import tqdm

def is_main_process():
    """
    检查当前进程是否为主进程。
    """
    return not dist.is_initialized() or dist.get_rank() == 0

def configure_trainable_params(model, modules, freeze = False,save_to_file=None):
    """
    筛选并设置模型中哪些模块的参数需要计算梯度，同时可以保存筛选结果到文件。

    Args:
        model (torch.nn.Module): 待操作的模型对象。
        modules (list): 模块名列表。
        save_to_file (str, optional): 如果提供文件路径，筛选的参数信息会保存到文件。
    
    Returns:
        dict: 包含目标模块参数的字典 {参数名: 参数对象}。
    """
    target_params = {}

    if freeze:
        for name, param in model.named_parameters():
            # 如果参数名属于被冻结的模块，则设置 requires_grad=False，否则为 True
            if any(module in name for module in modules):
                param.requires_grad = False  # 冻结模块的参数
            else:
                param.requires_grad = True  # 其它模块的参数可训练
                target_params[name] = param
    else:
        # 遍历所有参数
        for name, param in model.named_parameters():
            # 如果参数名属于目标模块，则设置 requires_grad=True，否则为 False
            if any(module in name for module in modules):
                param.requires_grad = True
                target_params[name] = param
            else:
                param.requires_grad = False

    # 保存目标参数到文件（如果提供了路径）
    if save_to_file:
        with open(f"{save_to_file}/filtered_trainable_parameters.txt", "w") as f:
            for name, param in target_params.items():
                f.write(f"{name}: requires_grad={param.requires_grad}, shape={param.shape}\n")
        with open(f"{save_to_file}/trainable_parameters_with_shapes.txt", "w") as f:
            for name, param in model.named_parameters():
                f.write(f"{name}: requires_grad={param.requires_grad}, shape={param.shape}\n")

    return target_params

def count_model_parameters(model):
    """
    统计模型的参数数量
    """
    total_params = sum(p.numel() for p in model.parameters())
    trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
    return total_params, trainable_params


def save_dataset_and_categories(train_loader, output_file="dataset_samples.txt"):
    """
    从数据集中提取前五个样本示例并提取第一个样本的类别名称保存到文件。

    Args:
        train_loader: 数据加载器。
        output_file (str): 目标文件路径。
    """
    print("#########################  datasampling  #################################")
    print(f"Preparing to write to {output_file}, existing content will be cleared.")
    
    with open(output_file, "w", encoding="utf-8") as f:
        # 解封装数据集
        dataset = train_loader.dataset
        if hasattr(dataset, 'dataset'):
            dataset = dataset.dataset  # 解封装内部数据集

        all_samples = []
        category_names = set()
        
        # 提取前五个样本
        for idx, sample in enumerate(dataset):
            all_samples.append(sample)
            # 仅从第一个样本提取类别名称
            if idx == 0 and "category_names" in sample:
                category_names.update(sample["category_names"])
            if idx >= 4:
                break

        # 写入前五个样本到文件
        for idx, sample in enumerate(all_samples):
            f.write(f"Sample {idx + 1}:\n")
            f.write(f"{sample}\n\n")

        # 排序和索引类别名称
        sorted_categories = sorted(category_names)
        indexed_categories = {index: category for index, category in enumerate(sorted_categories)}

        # 写入类别信息到文件
        f.write("Sorted categories with indices:\n")
        for index, category in indexed_categories.items():
            f.write(f"{index}: {category}\n")

    print(f"Dataset samples and categories written to {output_file}")


def check_cropping_status(output_dir):
    """
    检查裁剪结果是否已经存在且完整。

    Args:
        output_dir (str): 裁剪结果保存目录。

    Returns:
        bool: 如果裁剪结果已经存在且完整，则返回 True；否则返回 False。
    """
    # 检查标志文件是否存在
    flag_file = os.path.join(output_dir, "cropped_done.json")
    if not os.path.exists(flag_file):
        return False

    # 检查裁剪文件目录是否存在
    if not os.path.exists(output_dir):
        return False

    # 加载标志文件内容
    try:
        with open(flag_file, "r") as f:
            cropped_info = json.load(f)
    except Exception as e:
        print(f"Error reading flag file: {e}")
        return False

    # 验证裁剪文件是否存在
    for file_name in cropped_info["files"]:
        if not os.path.exists(os.path.join(output_dir, file_name)):
            print(f"Missing cropped file: {file_name}")
            return False

    print("Cropping status verified: All files are complete.")
    return True


def save_cropping_status(output_dir, file_list):
    """
    保存裁剪结果的状态到标志文件。

    Args:
        output_dir (str): 裁剪结果保存目录。
        file_list (list): 保存的裁剪文件列表。
    """
    flag_file = os.path.join(output_dir, "cropped_done.json")
    cropped_info = {"files": file_list}

    with open(flag_file, "w") as f:
        json.dump(cropped_info, f, indent=4)
    print(f"Cropping status saved to {flag_file}")

def save_cropped_instances_with_annotations(train_loader, output_dir="cropped_instances"):
    """
    裁剪数据加载器中的每个实例，并保存为单独的图片，同时生成新的标注文件。

    Args:
        train_loader: 数据加载器。
        output_dir (str): 保存裁剪图像和标注文件的目录。
    """
    if not is_main_process():
        return  # 非主进程直接跳过

    if check_cropping_status(output_dir):
        print("Cropping already done. Skipping...")
        return

    os.makedirs(output_dir, exist_ok=True)
    cropped_annotations = []  # 用于存储裁剪后的标注信息
    file_list = []

    # 解封装数据集
    dataset = train_loader.dataset
    if hasattr(dataset, 'dataset'):
        dataset = dataset.dataset

    # 初始化进度条
    with tqdm(total=len(dataset), desc="Cropping Instances", unit="sample") as pbar:
        for idx, sample in enumerate(dataset):
            file_name = sample["file_name"]
            image = Image.open(file_name).convert("RGB")
            annotations = sample["annotations"]  # 每个样本的标注列表

            for ann_idx, annotation in enumerate(annotations):
                bbox = annotation["bbox"]  # [x, y, w, h]
                category_id = annotation["category_id"]

                x1, y1, w, h = bbox
                x2, y2 = x1 + w, y1 + h

                # 裁剪图像
                cropped_img = image.crop((int(x1), int(y1), int(x2), int(y2)))
                cropped_file_name = f"{idx}_{ann_idx}_class_{category_id}.jpg"
                cropped_img_path = os.path.join(output_dir, cropped_file_name)
                cropped_img.save(cropped_img_path)
                file_list.append(cropped_file_name)

                # 保存裁剪后的标注信息
                cropped_annotations.append({
                    "file_name": cropped_file_name,
                    "width": int(x2 - x1),
                    "height": int(y2 - y1),
                    "category_id": category_id,
                    "original_file": file_name,
                    "original_bbox": bbox
                })

            pbar.update(1)

    # 保存裁剪状态
    save_cropping_status(output_dir, file_list)

    # 保存新的标注文件
    annotation_file = os.path.join(output_dir, "annotations.json")
    with open(annotation_file, "w") as f:
        json.dump(cropped_annotations, f, indent=4)
    print(f"Annotations saved to {annotation_file}")

    print("All cropped instances and annotations saved.")