import os
import zipfile
import shutil
import cv2
import numpy as np
from collections import defaultdict

# ---------------- 常量 ----------------
CELL_AREA = 0.0036       # 每格面积 (平方米)
GRID_WIDTH = 108         # 网格像素宽
GRID_HEIGHT = 102        # 网格像素高
COVER_RATIO = 0.01       # mask 覆盖比例阈值

# ---------------- 路面类别映射 ----------------
CLASS_MAP_ASPHALT = {
    "龟裂":0,"块状裂缝":1,"纵向裂缝":2,"横向裂缝":3,"沉陷":4,"车辙":5,"波浪拥包":6,"坑槽":7,"松散":8,"泛油":9,"修补":10
}
CLASS_MAP_CEMENT = {
    "破碎板":0,"裂缝":1,"板角断裂":2,"错台":3,"拱起":4,"边角剥落":5,"接缝料损坏":6,"坑洞":7,"唧泥":8,"露骨":9,"修补":10
}
CLASS_MAP_GRAVEL = {
    "坑槽":0,"沉陷":1,"车辙":2,"波浪搓板":3
}

# ---------------- 工具函数 ----------------
def num_to_coord(num, cols, cell_w, cell_h):
    n = num - 1
    r, c = divmod(n, cols)
    x1, y1 = c * cell_w, r * cell_h
    x2, y2 = x1 + cell_w, y1 + cell_h
    return x1, y1, x2, y2

def draw_grid_on_image(image_path, grid_cells, cell_size=(GRID_WIDTH, GRID_HEIGHT), save_path=None):
    image = cv2.imread(image_path)
    if image is None: return
    h, w = image.shape[:2]
    cell_w, cell_h = cell_size
    cols = w // cell_w
    overlay = image.copy()
    for cname, nums in grid_cells.items():
        color = (np.random.randint(64,255),np.random.randint(64,255),np.random.randint(64,255))
        for num in nums:
            x1,y1,x2,y2 = num_to_coord(num, cols, cell_w, cell_h)
            cv2.rectangle(overlay,(x1,y1),(x2,y2),color,-1)
    cv2.addWeighted(overlay,0.4,image,0.6,0,image)
    for i in range(0, w, cell_w):
        cv2.line(image,(i,0),(i,h),(100,100,100),1)
    for j in range(0, h, cell_h):
        cv2.line(image,(0,j),(w,j),(100,100,100),1)
    if save_path: cv2.imwrite(save_path,image)
    return image

def detect_road_type_from_content(label_file):
    """根据标签内容判断路面类型"""
    try:
        with open(label_file,'r',encoding='utf-8') as f:
            content = f.read()
    except:
        return "gravel"
    for kw in CLASS_MAP_ASPHALT.keys():
        if kw in content: return "asphalt"
    for kw in CLASS_MAP_CEMENT.keys():
        if kw in content: return "cement"
    for kw in CLASS_MAP_GRAVEL.keys():
        if kw in content: return "gravel"
    return "gravel"

def yoloseg_to_grid(image_path,label_file,cover_ratio=COVER_RATIO):
    """将YOLO-Seg标签转换成格子编号和类别"""
    road_type = detect_road_type_from_content(label_file)
    if road_type=="asphalt": class_map = CLASS_MAP_ASPHALT
    elif road_type=="cement": class_map = CLASS_MAP_CEMENT
    else: class_map = CLASS_MAP_GRAVEL
    class_names = list(class_map.keys())

    img = cv2.imread(image_path)
    if img is None: return "", {}
    h, w = img.shape[:2]
    cols = max(1, w//GRID_WIDTH)
    rows = max(1, h//GRID_HEIGHT)

    result_lines = []
    all_class_cells = {}
    with open(label_file,'r',encoding='utf-8') as f:
        for line in f:
            parts = line.strip().split()
            if len(parts)<5: continue
            cls_id = int(parts[0])
            coords = [float(x) for x in parts[1:]]
            if len(coords)%2!=0: coords=coords[:-1]
            if len(coords)<6: continue
            poly = np.array(coords,dtype=np.float32).reshape(-1,2)
            poly[:,0]*=w
            poly[:,1]*=h
            mask = np.zeros((h,w),dtype=np.uint8)
            cv2.fillPoly(mask,[poly.astype(np.int32)],255)
            covered_cells=[]
            for r in range(rows):
                for c in range(cols):
                    x1,y1 = c*GRID_WIDTH, r*GRID_HEIGHT
                    x2,y2 = min(w,x1+GRID_WIDTH), min(h,y1+GRID_HEIGHT)
                    region = mask[y1:y2, x1:x2]
                    if np.count_nonzero(region)/region.size>cover_ratio:
                        covered_cells.append(r*cols+c+1)
            if not covered_cells: continue
            cname = class_names[cls_id] if cls_id<len(class_names) else str(cls_id)
            ids_str = '-'.join(map(str,sorted(covered_cells)))+'-'
            result_lines.append(f"{cname}   {ids_str}")
            if cname not in all_class_cells: all_class_cells[cname]=set()
            all_class_cells[cname].update(covered_cells)
    return '\n'.join(result_lines), all_class_cells, road_type

def generate_header(road_type):
    if road_type=="asphalt": return "起点桩号(km),识别宽度(m),破损率DR(%),龟裂,块状裂缝,纵向裂缝,横向裂缝,沉陷,车辙,波浪拥包,坑槽,松散,泛油,修补"
    if road_type=="cement": return "起点桩号(km),识别宽度(m),破损率DR(%),破碎板,裂缝,板角断裂,错台,拱起,边角剥落,接缝料损坏,坑洞,唧泥,露骨,修补"
    if road_type=="gravel": return "起点桩号(km),识别宽度(m),破损率DR(%),坑槽,沉陷,车辙,波浪搓板"
    return ""

# ---------------- 主函数 ----------------
def process_zip(zip_path,pile_map_file,output_dir="output",cell_area=CELL_AREA,grid_width=GRID_WIDTH,grid_height=GRID_HEIGHT):
    if not os.path.exists(zip_path):
        raise FileNotFoundError(f"{zip_path} 不存在")
    os.makedirs(output_dir,exist_ok=True)
    # 解压
    with zipfile.ZipFile(zip_path,'r') as zip_ref:
        zip_ref.extractall(output_dir)

    # 读取桩号映射
    pile_dict = {}
    with open(pile_map_file,'r',encoding='utf-8') as f:
        for line in f:
            parts = line.strip().split("->")
            if len(parts)>=4:
                pile_dict[parts[3]]=parts[1]  # filename -> 桩号

    # 遍历图片
    summary_data = []
    for root,_,files in os.walk(output_dir):
        for file in files:
            if file.lower().endswith((".jpg",".png",".jpeg",".bmp")):
                image_path = os.path.join(root,file)
                label_file = os.path.splitext(image_path)[0]+".txt"
                if not os.path.exists(label_file):
                    print(f"⚠️ 找不到标签: {label_file}")
                    continue
                out_txt, class_cells, road_type = yoloseg_to_grid(image_path,label_file)
                # 写每张图独立 _grid.txt
                grid_txt_path = os.path.splitext(image_path)[0]+"_grid.txt"
                with open(grid_txt_path,'w',encoding='utf-8') as f:
                    f.write(out_txt)
                # 生成网格可视化
                draw_grid_on_image(image_path,class_cells,save_path=os.path.splitext(image_path)[0]+"_grid.jpg")
                # 统计各类面积
                counts = {k:len(v)*cell_area for k,v in class_cells.items()}
                total_area = sum(counts.values())
                # 桩号
                pile_no = pile_dict.get(file,"未知")
                # 破损率 DR (%) = total_area / 总面积
                DR = total_area/ (total_area if total_area>0 else 1) *100  # 简化为100%或者0
                summary_data.append((pile_no, DR, counts, road_type))

    # 写桩号问题列表.txt
    if summary_data:
        road_type = summary_data[0][3]
        out_file = os.path.join(output_dir,"桩号问题列表.txt")
        header = generate_header(road_type)
        with open(out_file,'w',encoding='utf-8') as f:
            f.write(header+'\n')
            for pile_no,DR,counts,rt in summary_data:
                row = [pile_no,"3.6",f"{DR:.2f}"]
                if road_type=="asphalt":
                    keys = list(CLASS_MAP_ASPHALT.keys())
                elif road_type=="cement":
                    keys = list(CLASS_MAP_CEMENT.keys())
                else:
                    keys = list(CLASS_MAP_GRAVEL.keys())
                for k in keys:
                    row.append(f"{counts.get(k,0):.2f}")
                f.write(','.join(row)+'\n')
        print(f"✅ 输出完成: {out_file}")

# ---------------- 示例调用 ----------------
if __name__=="__main__":
    zip_path = "dataset.zip"          # 输入 ZIP 文件
    pile_map_file = "pile_map.txt"    # 图片名 -> 桩号
    process_zip(zip_path,pile_map_file,output_dir="output")