import tkinter as tk
from tkinter import filedialog, messagebox, ttk
import pandas as pd
import jieba
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.metrics.pairwise import cosine_similarity
from scipy.cluster.hierarchy import linkage, fcluster, leaves_list
from scipy.spatial.distance import squareform
import openpyxl
from openpyxl.styles import PatternFill
import threading
import os
import numpy as np

# --- 核心后端逻辑 (部分更新) ---

def chinese_tokenizer(text): return " ".join(jieba.cut(text))

# calculate_similarity现在接收一个已经是副本的DataFrame
def calculate_similarity(df_processed, selected_columns, method='concatenate'):
    # 注意：传入的df_processed已经是副本，并且相关的列已经转为str
    if method == 'concatenate':
        corpus = df_processed[selected_columns].apply(lambda row: ' '.join(row), axis=1).tolist()
        if not any(corpus): return np.identity(len(df_processed))
        vectorizer = TfidfVectorizer(tokenizer=jieba.lcut); tfidf_matrix = vectorizer.fit_transform(corpus)
        return cosine_similarity(tfidf_matrix)
    elif method == 'weighted':
        num_rows, final_similarity_matrix, valid_cols_count = len(df_processed), np.zeros((len(df_processed), len(df_processed))), 0
        for col in selected_columns:
            corpus = df_processed[col].tolist()
            if not any(c for c in corpus if c.strip()): continue # 检查列是否真的有文本内容
            valid_cols_count += 1
            vectorizer = TfidfVectorizer(tokenizer=jieba.lcut); tfidf_matrix = vectorizer.fit_transform(corpus)
            final_similarity_matrix += cosine_similarity(tfidf_matrix)
        return final_similarity_matrix / valid_cols_count if valid_cols_count > 0 else np.identity(num_rows)

def perform_hierarchical_analysis(similarity_matrix):
    # ... (此函数无变化)
    distance_matrix = 1 - np.clip(similarity_matrix, 0, 1); np.fill_diagonal(distance_matrix, 0)
    condensed_distance = squareform(distance_matrix)
    Z = linkage(condensed_distance, method='ward'); leaf_order = leaves_list(Z)
    return Z, leaf_order

def auto_generate_levels(num_rows):
    # ... (此函数无变化)
    if num_rows < 10: return list(range(num_rows // 2, 1, -1))
    k_start = max(5, num_rows // 3)
    if k_start >= num_rows: k_start = num_rows - 1
    levels, current_k = [], k_start
    while current_k > 2:
        levels.append(int(current_k))
        next_k = current_k / 1.8
        current_k = current_k - 1 if int(next_k) == int(current_k) else next_k
    fixed_levels = [100, 50, 25, 10, 5, 3, 2]
    for fl in fixed_levels:
        if fl < k_start: levels.append(fl)
    final_levels = sorted(list(set(levels)), reverse=True)
    final_levels = [k for k in final_levels if 1 < k < num_rows]
    return final_levels if final_levels else ([num_rows // 2] if num_rows > 3 else [])

# generate_hierarchical_report 现在接收原始df
def generate_hierarchical_report(df_original, Z, leaf_order, levels, output_path):
    df_report = df_original.copy()

    level_cols = []
    for k in levels:
        col_name = f'层级_{k}_类ID'
        level_cols.append(col_name)
        cluster_ids = fcluster(Z, t=k, criterion='maxclust')
        df_report[col_name] = cluster_ids

    sorted_df = df_report.iloc[leaf_order]

    # ... (后续的上色逻辑无变化)
    writer = pd.ExcelWriter(output_path, engine='openpyxl')
    sorted_df.to_excel(writer, index=False, sheet_name='层级聚类结果')
    workbook = writer.book; worksheet = writer.sheets['层级聚类结果']
    colors = ['FFFF00', '00FF00', '00FFFF', 'FF00FF', 'FFA500', 'ADD8E6', 'F0E68C', '90EE90', 'FFB6C1', 'E6E6FA']
    header_map = {col: i + 1 for i, col in enumerate(sorted_df.columns)}
    for col_name in level_cols:
        col_idx = header_map[col_name]
        unique_ids = sorted_df[col_name].unique()
        color_map = {uid: colors[i % len(colors)] for i, uid in enumerate(unique_ids)}
        for row_idx, uid in enumerate(sorted_df[col_name], start=2):
            fill = PatternFill(start_color=color_map[uid], end_color=color_map[uid], fill_type='solid')
            worksheet.cell(row=row_idx, column=col_idx).fill = fill
    writer.close()


# --- GUI 应用 (逻辑更新) ---

class ClusterToolApp:
    def __init__(self, root):
        # ... (GUI布局代码与上一版完全相同，此处省略以保持简洁)
        self.root = root
        self.root.title("全自动Excel行相似度层级聚类工具 (V3-类型保持)")
        self.root.geometry("1000x800")
        self.df_original = None # 新增，用于存储原始数据
        self.filepath = ""
        # ... (GUI组件创建代码)
        top_frame = ttk.Frame(root, padding="10"); top_frame.pack(side=tk.TOP, fill=tk.X)
        left_frame = ttk.LabelFrame(top_frame, text="1. 输入数据", padding="10"); left_frame.pack(side=tk.LEFT, fill=tk.Y, padx=5, pady=5)
        middle_frame = ttk.LabelFrame(top_frame, text="2. 设置参数", padding="10"); middle_frame.pack(side=tk.LEFT, fill=tk.Y, padx=5, pady=5)
        right_frame = ttk.LabelFrame(top_frame, text="3. 执行与日志", padding="10"); right_frame.pack(side=tk.LEFT, fill=tk.BOTH, expand=True, padx=5, pady=5)
        bottom_frame = ttk.LabelFrame(root, text="数据预览 (前100行)", padding="10"); bottom_frame.pack(side=tk.BOTTOM, fill=tk.BOTH, expand=True, padx=10, pady=10)
        btn_load = ttk.Button(left_frame, text="加载Excel文件", command=self.load_excel); btn_load.pack(fill=tk.X, pady=5)
        self.lbl_filepath = ttk.Label(left_frame, text="未加载文件", wraplength=200); self.lbl_filepath.pack(fill=tk.X, pady=5)
        ttk.Label(left_frame, text="选择用于分析的列:").pack(anchor=tk.W, pady=(10, 2))
        self.listbox_cols = tk.Listbox(left_frame, selectmode=tk.MULTIPLE, exportselection=False, width=30, height=10); self.listbox_cols.pack(fill=tk.BOTH, expand=True)
        self.sim_method_var = tk.StringVar(value="concatenate"); ttk.Label(middle_frame, text="相似度计算方法:").pack(anchor=tk.W)
        ttk.Radiobutton(middle_frame, text="拼接文本计算", variable=self.sim_method_var, value="concatenate").pack(anchor=tk.W)
        ttk.Radiobutton(middle_frame, text="按列加权平均", variable=self.sim_method_var, value="weighted").pack(anchor=tk.W)
        self.btn_run = ttk.Button(right_frame, text="开始分析", command=self.start_analysis_thread, state=tk.DISABLED); self.btn_run.pack(fill=tk.X, ipady=10, pady=5)
        self.log_text = tk.Text(right_frame, wrap=tk.WORD, state=tk.DISABLED, height=10); self.log_text.pack(fill=tk.BOTH, expand=True)
        self.tree = ttk.Treeview(bottom_frame, show='headings'); vsb = ttk.Scrollbar(bottom_frame, orient="vertical", command=self.tree.yview); hsb = ttk.Scrollbar(bottom_frame, orient="horizontal", command=self.tree.xview)
        self.tree.configure(yscrollcommand=vsb.set, xscrollcommand=hsb.set); vsb.pack(side='right', fill='y'); hsb.pack(side='bottom', fill='x'); self.tree.pack(side='left', fill='both', expand=True)


    def log(self, message):
        def _log():
            self.log_text.config(state=tk.NORMAL); self.log_text.insert(tk.END, message + "\n"); self.log_text.see(tk.END); self.log_text.config(state=tk.DISABLED)
        self.root.after(0, _log)

    def update_preview(self):
        for i in self.tree.get_children(): self.tree.delete(i)
        self.tree["columns"] = []
        if self.df_original is None: return
        df_preview = self.df_original.head(100)
        self.tree["columns"] = list(df_preview.columns)
        for col in df_preview.columns:
            self.tree.heading(col, text=col); self.tree.column(col, width=100, anchor='w')
        for index, row in df_preview.iterrows():
            self.tree.insert("", "end", values=[str(v) if pd.notna(v) else "" for v in row])

    # --- load_excel 更新 ---
    def load_excel(self):
        filepath = filedialog.askopenfilename(
            title="选择Excel文件", filetypes=(("Excel 文件", "*.xlsx;*.xls"), ("所有文件", "*.*"))
        )
        if not filepath: return
        self.filepath = filepath
        self.log(f"正在加载: {os.path.basename(filepath)}...")
        try:
            # 读取时不指定dtype，让pandas自动推断
            self.df_original = pd.read_excel(filepath, sheet_name=0)

            if len(self.df_original) > 10000 and not messagebox.askyesno("性能警告", f"文件包含 {len(self.df_original)} 行。\n分析可能会非常缓慢。\n确定要继续吗？"):
                self.df_original = None # 取消加载
                return

            self.lbl_filepath.config(text=os.path.basename(filepath))
            self.listbox_cols.delete(0, tk.END)
            for col in self.df_original.columns: self.listbox_cols.insert(tk.END, col)
            self.update_preview()
            self.btn_run.config(state=tk.NORMAL)
            self.log("文件加载成功，请选择列并点击分析。")
        except Exception as e:
            self.log(f"错误: {e}")
            messagebox.showerror("加载失败", f"无法加载文件: {e}")

    def start_analysis_thread(self):
        self.btn_run.config(state=tk.DISABLED)
        self.log("="*40)
        thread = threading.Thread(target=self.run_analysis, daemon=True)
        thread.start()

    # --- run_analysis 更新 ---
    def run_analysis(self):
        try:
            selected_indices = self.listbox_cols.curselection()
            if not selected_indices: raise ValueError("请至少选择一个列进行分析。")
            if self.df_original is None or self.df_original.shape[0] < 2: raise ValueError("数据行数不足。")

            selected_columns = [self.listbox_cols.get(i) for i in selected_indices]
            sim_method = self.sim_method_var.get()

            self.log(f"选择的列: {', '.join(selected_columns)}")
            self.log(f"相似度方法: {'拼接文本' if sim_method == 'concatenate' else '加权平均'}")

            self.log("\n步骤1/4: 自动生成聚类层级...")
            levels = auto_generate_levels(len(self.df_original))
            if not levels: raise ValueError("数据行数太少，无法生成层级。")
            self.log(f"自动生成的层级为: {levels}")

            # --- 核心修改点：创建处理用的副本 ---
            self.log("\n步骤2/4: 准备数据并计算相似度...")
            df_for_processing = self.df_original.copy()
            # 只对选中的列进行类型转换和填充
            for col in selected_columns:
                df_for_processing[col] = df_for_processing[col].fillna('').astype(str)

            similarity_matrix = calculate_similarity(df_for_processing, selected_columns, sim_method)

            self.log("步骤3/4: 正在执行层次聚类分析...")
            Z, leaf_order = perform_hierarchical_analysis(similarity_matrix)

            self.log("步骤4/4: 正在生成Excel报告 (保持原始类型)...")
            output_dir = os.path.dirname(self.filepath)
            base_name = os.path.splitext(os.path.basename(self.filepath))[0]
            output_path = os.path.join(output_dir, f"{base_name}_自动层级聚类结果_v3.xlsx")

            # 传入原始DataFrame进行报告生成
            generate_hierarchical_report(self.df_original, Z, leaf_order, levels, output_path)

            self.log(f"\n分析完成！结果已保存至:\n{output_path}")
            messagebox.showinfo("完成", f"分析已完成！\n结果已保存至: {output_path}")

        except Exception as e:
            self.log(f"错误: {e}")
            import traceback
            self.log(traceback.format_exc())
            messagebox.showerror("运行错误", f"分析过程中发生错误: {e}")
        finally:
            self.root.after(0, lambda: self.btn_run.config(state=tk.NORMAL))


if __name__ == "__main__":
    root = tk.Tk()
    app = ClusterToolApp(root)
    root.mainloop()