import argparse
import json
import re
from pathlib import Path
from collections import defaultdict, Counter
import tkinter as tk
from tkinter import ttk, filedialog, messagebox
import threading

import pandas as pd
import requests
from openpyxl import load_workbook
from openpyxl.chart import BarChart, PieChart, Reference
from openpyxl.chart.label import DataLabelList

def download_results(self):
    """Allow user to save results file"""
    if not self.results_path or not self.results_path.exists():
        messagebox.showerror("Error", "Results file not found")
        return

    save_path = filedialog.asksaveasfilename(
        title="Save Analysis Results",
        defaultextension=".xlsx",
        filetypes=[("Excel files", "*.xlsx"), ("All files", "*.*")],
        initialfile=f"{self.file_path.stem}_topic_analysis.xlsx"
    )

    if save_path:
        try:
            import shutil
            # Copy the temp file to user's chosen location
            shutil.copy2(self.results_path, save_path)

            # Also offer to save CSV
            csv_save = messagebox.askyesno(
                "Save CSV?",
                "Would you also like to save the CSV version?"
            )

            if csv_save and hasattr(self, 'results_csv_path') and self.results_csv_path.exists():
                csv_path = filedialog.asksaveasfilename(
                    title="Save CSV Results",
                    defaultextension=".csv",
                    filetypes=[("CSV files", "*.csv"), ("All files", "*.*")],
                    initialfile=f"{self.file_path.stem}_topic_analysis.csv"
                )
                if csv_path:
                    shutil.copy2(self.results_csv_path, csv_path)

            # Clean up temp files after successful copy
            try:
                self.results_path.unlink()
                if hasattr(self, 'results_csv_path') and self.results_csv_path.exists():
                    self.results_csv_path.unlink()
            except:
                pass  # Ignore cleanup errors

            self.show_download_complete()
        except Exception as e:
            messagebox.showerror("Error",
                                 f"Failed to save file:\n{str(e)}")  # gemma2_topic_analyzer_gui.py - Integrated Topic Analysis with GUI


# ==================== CORE ANALYSIS FUNCTIONS ====================

def load_single_sheet_excel(path: Path, sheet_arg):
    """Load a specific sheet from Excel file"""
    xls = pd.read_excel(path, sheet_name=None)
    if not isinstance(xls, dict) or len(xls) == 0:
        raise SystemExit("The workbook is empty or could not be read.")

    sheet_names = list(xls.keys())

    if sheet_arg is None:
        chosen_name = sheet_names[0]
    else:
        try:
            idx = int(sheet_arg)
            if idx < 0 or idx >= len(sheet_names):
                raise SystemExit(f"Sheet index {idx} out of range. Available sheets: {sheet_names}")
            chosen_name = sheet_names[idx]
        except (ValueError, TypeError):
            name = str(sheet_arg)
            if name not in xls:
                match = [s for s in sheet_names if s.lower() == name.lower()]
                if not match:
                    raise SystemExit(f"Sheet '{name}' not found. Available sheets: {sheet_names}")
                chosen_name = match[0]
            else:
                chosen_name = name

    df = xls[chosen_name]
    if not isinstance(df, pd.DataFrame):
        raise SystemExit(f"Selected sheet '{chosen_name}' did not load as a DataFrame.")

    return df, chosen_name, sheet_names


def call_ollama(prompt: str, model: str, base_url: str) -> str:
    """Call Ollama API for text generation"""
    url = f"{base_url}/api/generate"

    payload = {
        "model": model,
        "prompt": prompt,
        "stream": False,
        "options": {
            "temperature": 0.2,
            "top_p": 0.85,
            "num_predict": 2000,
        }
    }

    try:
        response = requests.post(url, json=payload, timeout=180)
        response.raise_for_status()
        result = response.json()
        return result.get("response", "")
    except requests.exceptions.RequestException as e:
        print(f"Error calling Ollama API: {e}")
        return ""


def analyze_comments_batch(comments: list[dict], model: str, base_url: str) -> list[dict]:
    """Analyze a batch of comments for topics and sentiment"""

    # Prepare comments text
    comments_text = "\n\n".join([
        f"[ID: {c['id']}]\n{c['text'][:500]}"
        for c in comments if c['text'].strip()
    ])

    if not comments_text.strip():
        return []

    prompt = f"""Analiza estos comentarios de hospital en español. Para cada uno extrae:
1. Máximo 2 temas principales
2. Sentimiento: POSITIVO, NEUTRAL o NEGATIVO
3. Puntuación: -10 (muy negativo) a +10 (muy positivo)

TEMAS posibles: Atención médica, Personal, Tiempo de espera, Instalaciones, Comunicación, Trato, Citas, Seguridad, Limpieza, Comida

COMENTARIOS:
{comments_text}

Responde SOLO con JSON válido (sin markdown):
{{
  "analisis": [
    {{"id": "ID", "temas": ["tema1", "tema2"], "sentimiento": "POSITIVO", "puntuacion": 5}}
  ]
}}"""

    response = call_ollama(prompt, model, base_url)

    # Parse JSON response with error recovery
    try:
        # Clean response - remove markdown code blocks
        cleaned = response.strip()
        if cleaned.startswith("```"):
            cleaned = re.sub(r'^```(?:json)?\s*\n', '', cleaned)
            cleaned = re.sub(r'\n```\s*$', '', cleaned)

        # FIX: Add missing commas between JSON objects
        # Pattern: }[whitespace]{ → },{
        cleaned = re.sub(r'}\s*\n\s*{', '},{', cleaned)

        data = json.loads(cleaned)
        return data.get("analisis", [])

    except json.JSONDecodeError as e:
        print(f"⚠ JSON parsing still failed after fix: {e}")
        print(f"Response preview: {response[:300]}")

        # Save failed response for debugging
        with open('failed_batch.json', 'w', encoding='utf-8') as f:
            f.write(response)

        return []


def normalize_topic(topic: str) -> str:
    """Normalize topic names for consistency"""
    topic = topic.lower().strip()

    mappings = {
        "atencion medica": "Atención médica",
        "atencion": "Atención médica",
        "calidad": "Atención médica",
        "cuidado": "Atención médica",

        "personal": "Personal/Enfermería",
        "enfermeria": "Personal/Enfermería",
        "enfermeras": "Personal/Enfermería",
        "medicos": "Personal/Enfermería",
        "doctores": "Personal/Enfermería",

        "tiempo de espera": "Tiempo de espera",
        "espera": "Tiempo de espera",
        "demora": "Tiempo de espera",
        "rapidez": "Tiempo de espera",

        "instalaciones": "Instalaciones",
        "limpieza": "Instalaciones",
        "equipamiento": "Instalaciones",
        "infraestructura": "Instalaciones",

        "comunicacion": "Comunicación",
        "informacion": "Comunicación",

        "trato": "Trato/Amabilidad",
        "amabilidad": "Trato/Amabilidad",
        "atencion personal": "Trato/Amabilidad",

        "citas": "Citas/Programación",
        "programacion": "Citas/Programación",
        "turnos": "Citas/Programación",

        "seguridad": "Seguridad del paciente",
        "seguridad paciente": "Seguridad del paciente",
    }

    for key, value in mappings.items():
        if key in topic:
            return value

    return topic.title()


def create_visualizations(out_xlsx: Path, results_df: pd.DataFrame):
    """Create summary sheet with visualizations"""

    # Calculate statistics
    responded_df = results_df[results_df['sentiment'] != 'NO_RESPONSE']
    total_responses = len(responded_df)  # Only count responses with actual comments

    positive = (responded_df['sentiment'] == 'POSITIVO').sum()
    neutral = (responded_df['sentiment'] == 'NEUTRAL').sum()
    negative = (responded_df['sentiment'] == 'NEGATIVO').sum()

    # Average sentiment score
    avg_score = responded_df['sentiment_score'].mean()

    # Load workbook
    wb = load_workbook(out_xlsx)

    # Create summary sheet
    if 'Summary' in wb.sheetnames:
        del wb['Summary']
    ws = wb.create_sheet('Summary', 0)

    # Title
    ws['A1'] = 'Análisis de Temas y Sentimientos'
    ws['A1'].font = ws['A1'].font.copy(bold=True, size=14)

    # Overall statistics
    ws['A3'] = 'Estadísticas Generales:'
    ws['A3'].font = ws['A3'].font.copy(bold=True)
    ws['A4'] = 'Total de comentarios analizados:';
    ws['B4'] = total_responses
    ws['A5'] = 'Puntuación promedio:';
    ws['B5'] = f"{avg_score:.2f} / 10"

    # Sentiment distribution table
    ws['A10'] = 'Distribución de Sentimientos'
    ws['A10'].font = ws['A10'].font.copy(bold=True)

    sentiment_data = [
        ['Categoría', 'Cantidad', 'Porcentaje'],
        ['Positivo', positive, f"{positive / total_responses * 100:.1f}%"],
        ['Neutral', neutral, f"{neutral / total_responses * 100:.1f}%"],
        ['Negativo', negative, f"{negative / total_responses * 100:.1f}%"]
    ]

    for r_idx, row_data in enumerate(sentiment_data, start=11):
        for c_idx, value in enumerate(row_data, start=1):
            cell = ws.cell(row=r_idx, column=c_idx, value=value)
            if r_idx == 11:
                cell.font = cell.font.copy(bold=True)

    # Bar Chart for sentiment
    bar_chart = BarChart()
    bar_chart.type = "col"
    bar_chart.title = "Distribución de Sentimientos"
    bar_chart.y_axis.title = "Número de comentarios"

    data = Reference(ws, min_col=2, min_row=11, max_row=14)
    cats = Reference(ws, min_col=1, min_row=12, max_row=14)
    bar_chart.add_data(data, titles_from_data=True)
    bar_chart.set_categories(cats)
    bar_chart.height = 12
    bar_chart.width = 20
    bar_chart.dataLabels = DataLabelList()
    bar_chart.dataLabels.showVal = True
    ws.add_chart(bar_chart, "E3")

    # Pie Chart for sentiment
    pie_chart = PieChart()
    pie_chart.title = "Distribución Porcentual"
    data = Reference(ws, min_col=2, min_row=11, max_row=14)
    labels = Reference(ws, min_col=1, min_row=12, max_row=14)
    pie_chart.add_data(data, titles_from_data=True)
    pie_chart.set_categories(labels)
    pie_chart.height = 12
    pie_chart.width = 20
    pie_chart.dataLabels = DataLabelList()
    pie_chart.dataLabels.showPercent = True
    pie_chart.dataLabels.showCatName = True
    ws.add_chart(pie_chart, "E25")

    # Column widths
    ws.column_dimensions['A'].width = 25
    ws.column_dimensions['B'].width = 18
    ws.column_dimensions['C'].width = 15

    wb.save(out_xlsx)
    print(f"✓ Summary sheet created")


def create_topic_analysis(out_xlsx: Path, results_df: pd.DataFrame):
    """Create topic analysis sheet organized by sentiment"""

    # Extract all topics with their sentiments
    topic_sentiment_scores = defaultdict(list)
    topic_counts = Counter()

    for _, row in results_df.iterrows():
        if pd.notna(row['topics']) and row['topics']:
            topics = row['topics'].split('; ')
            for topic in topics:
                topic = normalize_topic(topic)
                topic_counts[topic] += 1
                topic_sentiment_scores[topic].append(row['sentiment_score'])

    # Calculate topic statistics and categorize by sentiment
    positive_topics = []
    neutral_topics = []
    negative_topics = []

    for topic, count in topic_counts.items():
        scores = topic_sentiment_scores[topic]
        avg_score = sum(scores) / len(scores)

        if avg_score >= 3:
            positive_topics.append({
                'Tema': topic,
                'Menciones': count,
                'Puntuación Promedio': round(avg_score, 2)
            })
        elif avg_score <= -3:
            negative_topics.append({
                'Tema': topic,
                'Menciones': count,
                'Puntuación Promedio': round(avg_score, 2)
            })
        else:
            neutral_topics.append({
                'Tema': topic,
                'Menciones': count,
                'Puntuación Promedio': round(avg_score, 2)
            })

    # Sort each category by mentions and take top 10
    positive_topics = sorted(positive_topics, key=lambda x: x['Menciones'], reverse=True)[:10]
    neutral_topics = sorted(neutral_topics, key=lambda x: x['Menciones'], reverse=True)[:10]
    negative_topics = sorted(negative_topics, key=lambda x: x['Menciones'], reverse=True)[:10]

    # Load workbook
    wb = load_workbook(out_xlsx)

    if 'Análisis de Temas' in wb.sheetnames:
        del wb['Análisis de Temas']

    ws = wb.create_sheet('Análisis de Temas')

    # Set column widths
    ws.column_dimensions['A'].width = 30
    ws.column_dimensions['B'].width = 12
    ws.column_dimensions['C'].width = 20
    ws.column_dimensions['D'].width = 5
    ws.column_dimensions['E'].width = 30
    ws.column_dimensions['F'].width = 12
    ws.column_dimensions['G'].width = 20
    ws.column_dimensions['H'].width = 5
    ws.column_dimensions['I'].width = 30
    ws.column_dimensions['J'].width = 12
    ws.column_dimensions['K'].width = 20

    def write_sentiment_section(topics, sentiment_label, start_col):
        row = 1
        ws.cell(row=row, column=start_col, value=f"TEMAS {sentiment_label.upper()}").font = \
            ws.cell(row=row, column=start_col).font.copy(bold=True, size=12)
        row += 1

        ws.cell(row=row, column=start_col, value="Tema").font = \
            ws.cell(row=row, column=start_col).font.copy(bold=True)
        ws.cell(row=row, column=start_col + 1, value="Menciones").font = \
            ws.cell(row=row, column=start_col + 1).font.copy(bold=True)
        ws.cell(row=row, column=start_col + 2, value="Puntuación").font = \
            ws.cell(row=row, column=start_col + 2).font.copy(bold=True)
        row += 1

        for topic_data in topics:
            ws.cell(row=row, column=start_col, value=topic_data['Tema'])
            ws.cell(row=row, column=start_col + 1, value=topic_data['Menciones'])
            ws.cell(row=row, column=start_col + 2, value=topic_data['Puntuación Promedio'])
            row += 1

        return row

    write_sentiment_section(positive_topics, "POSITIVOS", 1)
    write_sentiment_section(neutral_topics, "NEUTRALES", 5)
    write_sentiment_section(negative_topics, "NEGATIVOS", 9)

    chart_start_row = max(len(positive_topics), len(neutral_topics), len(negative_topics)) + 5

    # Charts
    if positive_topics:
        chart = BarChart()
        chart.type = "bar"
        chart.title = "Temas Positivos (Top 10)"
        chart.x_axis.title = "Número de Menciones"
        chart.y_axis.title = ""
        chart.height = max(10, len(positive_topics) * 0.8)
        chart.width = 22
        data = Reference(ws, min_col=2, min_row=2, max_row=2 + len(positive_topics))
        cats = Reference(ws, min_col=1, min_row=3, max_row=2 + len(positive_topics))
        chart.add_data(data, titles_from_data=True)
        chart.set_categories(cats)
        chart.dataLabels = DataLabelList()
        chart.dataLabels.showVal = True
        chart.gapWidth = 50
        ws.add_chart(chart, f"A{chart_start_row}")

    if neutral_topics:
        chart = BarChart()
        chart.type = "bar"
        chart.title = "Temas Neutrales (Top 10)"
        chart.x_axis.title = "Número de Menciones"
        chart.y_axis.title = ""
        chart.height = max(10, len(neutral_topics) * 0.8)
        chart.width = 22
        data = Reference(ws, min_col=6, min_row=2, max_row=2 + len(neutral_topics))
        cats = Reference(ws, min_col=5, min_row=3, max_row=2 + len(neutral_topics))
        chart.add_data(data, titles_from_data=True)
        chart.set_categories(cats)
        chart.dataLabels = DataLabelList()
        chart.dataLabels.showVal = True
        chart.gapWidth = 50
        ws.add_chart(chart, f"A{chart_start_row + 20}")

    if negative_topics:
        chart = BarChart()
        chart.type = "bar"
        chart.title = "Temas Negativos (Top 10)"
        chart.x_axis.title = "Número de Menciones"
        chart.y_axis.title = ""
        chart.height = max(10, len(negative_topics) * 0.8)
        chart.width = 22
        data = Reference(ws, min_col=10, min_row=2, max_row=2 + len(negative_topics))
        cats = Reference(ws, min_col=9, min_row=3, max_row=2 + len(negative_topics))
        chart.add_data(data, titles_from_data=True)
        chart.set_categories(cats)
        chart.dataLabels = DataLabelList()
        chart.dataLabels.showVal = True
        chart.gapWidth = 50
        ws.add_chart(chart, f"A{chart_start_row + 40}")

    wb.save(out_xlsx)
    print(f"✓ Topic analysis sheet created:")
    print(f"  - {len(positive_topics)} positive topics")
    print(f"  - {len(neutral_topics)} neutral topics")
    print(f"  - {len(negative_topics)} negative topics")


# ==================== GUI CLASS ====================

class TopicAnalyzerGUI:
    def __init__(self, root):
        self.root = root
        self.root.title("Patient Safety Comments Topic Analyzer")
        self.root.geometry("1024x768")
        self.root.resizable(True, True)
        self.root.minsize(1024, 768)

        # Variables
        self.file_path = None
        self.sheet_name = None
        self.column_name = None
        self.id_column_name = None
        self.batch_size = tk.IntVar(value=5)
        self.results_path = None
        self.results_csv_path = None  # Add CSV path tracking
        self.model_var = tk.StringVar(value="gemma2:9b")  # Default model
        self.ollama_url = "http://localhost:11434"

        # Styling
        self.style = ttk.Style()
        self.style.theme_use('clam')

        self.bg_color = "#f0f0f0"
        self.primary_color = "#2196F3"
        self.success_color = "#4CAF50"
        self.root.configure(bg=self.bg_color)

        self.show_main_page()

        self.start_time = None
        self.estimated_total_time = None

    def clear_window(self):
        """Clear all widgets from window"""
        for widget in self.root.winfo_children():
            widget.destroy()

    def show_main_page(self):
        """Display the main page for file upload"""
        self.clear_window()

        main_frame = tk.Frame(self.root, bg=self.bg_color)
        main_frame.place(relx=0.5, rely=0.5, anchor="center")

        title = tk.Label(
            main_frame,
            text="Patient Safety Comments Topic Analyzer",
            font=("Arial", 24, "bold"),
            bg=self.bg_color,
            fg="#333"
        )
        title.pack(pady=(0, 10))

        subtitle = tk.Label(
            main_frame,
            text="Spanish patient survey comments analyzed using Ollama",
            font=("Arial", 12),
            bg=self.bg_color,
            fg="#666"
        )
        subtitle.pack(pady=(0, 40))

        upload_label = tk.Label(
            main_frame,
            text="Select Excel File (.xlsx)",
            font=("Arial", 14, "bold"),
            bg=self.bg_color
        )
        upload_label.pack(pady=(0, 15))

        self.file_display = tk.Label(
            main_frame,
            text="No file selected",
            font=("Arial", 10),
            bg="white",
            width=50,
            height=2,
            relief="solid",
            borderwidth=1
        )
        self.file_display.pack(pady=(0, 15))

        browse_btn = tk.Button(
            main_frame,
            text="📁 Browse Files",
            font=("Arial", 12, "bold"),
            bg=self.primary_color,
            fg="white",
            width=20,
            height=2,
            relief="flat",
            cursor="hand2",
            command=self.browse_file
        )
        browse_btn.pack(pady=(0, 20))

        exit_btn = tk.Button(
            main_frame,
            text="Exit",
            font=("Arial", 10),
            bg="#f44336",
            fg="white",
            width=15,
            relief="flat",
            cursor="hand2",
            command=self.exit_app
        )
        exit_btn.pack(pady=(30, 0))

    def browse_file(self):
        """Open file browser to select Excel file"""
        file_path = filedialog.askopenfilename(
            title="Select Excel File",
            filetypes=[("Excel files", "*.xlsx"), ("All files", "*.*")]
        )

        if file_path:
            self.file_path = Path(file_path)
            self.file_display.config(text=self.file_path.name)
            self.show_sheet_selection()

    def show_sheet_selection(self):
        """Display available sheets for selection"""
        self.clear_window()

        try:
            xls = pd.read_excel(self.file_path, sheet_name=None)
            sheet_names = list(xls.keys())
        except Exception as e:
            messagebox.showerror("Error", f"Failed to read Excel file:\n{str(e)}")
            self.show_main_page()
            return

        main_frame = tk.Frame(self.root, bg=self.bg_color)
        main_frame.pack(expand=True, fill="both", padx=40, pady=40)

        title = tk.Label(
            main_frame,
            text="Select Sheet",
            font=("Arial", 20, "bold"),
            bg=self.bg_color
        )
        title.pack(pady=(0, 20))

        sheet_frame = tk.Frame(main_frame, bg=self.bg_color)
        sheet_frame.pack(pady=20)

        tk.Label(
            sheet_frame,
            text="Available Sheets:",
            font=("Arial", 12, "bold"),
            bg=self.bg_color
        ).pack(anchor="w", pady=(0, 10))

        listbox_frame = tk.Frame(sheet_frame, bg="white", relief="solid", borderwidth=1)
        listbox_frame.pack()

        scrollbar = tk.Scrollbar(listbox_frame)
        scrollbar.pack(side="right", fill="y")

        self.sheet_listbox = tk.Listbox(
            listbox_frame,
            font=("Arial", 11),
            width=60,
            height=10,
            yscrollcommand=scrollbar.set
        )
        self.sheet_listbox.pack(side="left", fill="both")
        scrollbar.config(command=self.sheet_listbox.yview)

        for sheet in sheet_names:
            self.sheet_listbox.insert(tk.END, sheet)

        self.sheet_listbox.select_set(0)

        btn_frame = tk.Frame(main_frame, bg=self.bg_color)
        btn_frame.pack(pady=30)

        tk.Button(
            btn_frame,
            text="← Back",
            font=("Arial", 11),
            bg="#757575",
            fg="white",
            width=15,
            relief="flat",
            cursor="hand2",
            command=self.show_main_page
        ).pack(side="left", padx=10)

        tk.Button(
            btn_frame,
            text="Next →",
            font=("Arial", 11, "bold"),
            bg=self.primary_color,
            fg="white",
            width=15,
            relief="flat",
            cursor="hand2",
            command=lambda: self.select_sheet(sheet_names)
        ).pack(side="left", padx=10)

    def select_sheet(self, sheet_names):
        """Process selected sheet"""
        selection = self.sheet_listbox.curselection()
        if not selection:
            messagebox.showwarning("Warning", "Please select a sheet")
            return

        self.sheet_name = sheet_names[selection[0]]
        self.show_column_selection()

    def show_column_selection(self):
        """Display available columns for selection"""
        self.clear_window()

        try:
            df = pd.read_excel(self.file_path, sheet_name=self.sheet_name)
            df.columns = df.columns.astype(str).str.strip()
            columns = list(df.columns)
        except Exception as e:
            messagebox.showerror("Error", f"Failed to read sheet:\n{str(e)}")
            self.show_sheet_selection()
            return

        # Find ID column
        self.id_column_name = None
        for col in columns:
            col_lower = col.lower()
            if ('response' in col_lower and 'id' in col_lower) or col_lower == 'id':
                self.id_column_name = col
                break

        main_frame = tk.Frame(self.root, bg=self.bg_color)
        main_frame.pack(expand=True, fill="both", padx=40, pady=40)

        title = tk.Label(
            main_frame,
            text="Select Column Containing Patient Comments",
            font=("Arial", 20, "bold"),
            bg=self.bg_color
        )
        title.pack(pady=(0, 10))

        subtitle = tk.Label(
            main_frame,
            text=f"Sheet: {self.sheet_name}",
            font=("Arial", 11),
            bg=self.bg_color,
            fg="#666"
        )
        subtitle.pack(pady=(0, 20))

        column_frame = tk.Frame(main_frame, bg=self.bg_color)
        column_frame.pack(pady=20)

        tk.Label(
            column_frame,
            text="Available Columns:",
            font=("Arial", 12, "bold"),
            bg=self.bg_color
        ).pack(anchor="w", pady=(0, 10))

        listbox_frame = tk.Frame(column_frame, bg="white", relief="solid", borderwidth=1)
        listbox_frame.pack()

        scrollbar = tk.Scrollbar(listbox_frame)
        scrollbar.pack(side="right", fill="y")

        self.column_listbox = tk.Listbox(
            listbox_frame,
            font=("Arial", 11),
            width=60,
            height=10,
            yscrollcommand=scrollbar.set
        )
        self.column_listbox.pack(side="left", fill="both")
        scrollbar.config(command=self.column_listbox.yview)

        for col in columns:
            self.column_listbox.insert(tk.END, col)

        # Auto-select Comments column
        for i, col in enumerate(columns):
            if "comment" in col.lower():
                self.column_listbox.select_set(i)
                break

        btn_frame = tk.Frame(main_frame, bg=self.bg_color)
        btn_frame.pack(pady=30)

        tk.Button(
            btn_frame,
            text="← Back",
            font=("Arial", 11),
            bg="#757575",
            fg="white",
            width=15,
            relief="flat",
            cursor="hand2",
            command=self.show_sheet_selection
        ).pack(side="left", padx=10)

        tk.Button(
            btn_frame,
            text="Next →",
            font=("Arial", 11, "bold"),
            bg=self.primary_color,
            fg="white",
            width=15,
            relief="flat",
            cursor="hand2",
            command=lambda: self.select_column(columns)
        ).pack(side="left", padx=10)

    def select_column(self, columns):
        """Process selected column"""
        selection = self.column_listbox.curselection()
        if not selection:
            messagebox.showwarning("Warning", "Please select a column")
            return

        self.column_name = columns[selection[0]]
        self.show_batch_size_selection()

    def show_batch_size_selection(self):
        """Display batch size selection"""
        self.clear_window()

        main_frame = tk.Frame(self.root, bg=self.bg_color)
        main_frame.place(relx=0.5, rely=0.5, anchor="center")

        title = tk.Label(
            main_frame,
            text="Configure Analysis",
            font=("Arial", 20, "bold"),
            bg=self.bg_color
        )
        title.pack(pady=(0, 30))

        # Batch size section
        batch_frame = tk.Frame(main_frame, bg=self.bg_color)
        batch_frame.pack(pady=20)

        tk.Label(
            batch_frame,
            text="Batch Size (comments per batch):",
            font=("Arial", 12, "bold"),
            bg=self.bg_color
        ).pack(anchor="w", pady=(0, 10))

        spinbox_frame = tk.Frame(batch_frame, bg=self.bg_color)
        spinbox_frame.pack(pady=10)

        spinbox = tk.Spinbox(
            spinbox_frame,
            from_=1,
            to=50,
            textvariable=self.batch_size,
            font=("Arial", 14),
            width=10
        )
        spinbox.pack()

        tk.Label(
            batch_frame,
            text="⚠️ Recommended for gemma2:9b model: 5\n⚠️ Recommended for gemma2:2b model: 3",
            font=("Arial", 10, "bold"),
            bg=self.bg_color,
            fg=self.success_color
        ).pack(pady=(5, 5))

        warning = tk.Label(
            batch_frame,
            text="Only increase if you have sufficient VRAM (8GB+)",
            font=("Arial", 9),
            bg=self.bg_color,
            fg="#ff9800",
            wraplength=400
        )
        warning.pack(pady=(0, 20))

        # Model selection section
        model_frame = tk.Frame(main_frame, bg=self.bg_color)
        model_frame.pack(pady=20)

        tk.Label(
            model_frame,
            text="Select AI Model:",
            font=("Arial", 12, "bold"),
            bg=self.bg_color
        ).pack(anchor="w", pady=(0, 10))

        # Radio buttons for model selection
        ttk.Radiobutton(
            model_frame,
            text="gemma2:9b (Higher accuracy, ~5GB VRAM)",
            variable=self.model_var,
            value="gemma2:9b"
        ).pack(anchor="w", pady=5)

        ttk.Radiobutton(
            model_frame,
            text="gemma2:2b (Faster processing, ~2GB VRAM)",
            variable=self.model_var,
            value="gemma2:2b"
        ).pack(anchor="w", pady=5)

        tk.Label(
            model_frame,
            text="💡 Choose based on your available GPU memory",
            font=("Arial", 9),
            bg=self.bg_color,
            fg="#666",
            wraplength=400
        ).pack(pady=(5, 0))

        # Summary
        summary_frame = tk.Frame(main_frame, bg="white", relief="solid", borderwidth=1)
        summary_frame.pack(pady=20, padx=20)

        tk.Label(
            summary_frame,
            text="Analysis Summary:",
            font=("Arial", 11, "bold"),
            bg="white"
        ).pack(anchor="w", padx=15, pady=(10, 5))

        tk.Label(
            summary_frame,
            text=f"File: {self.file_path.name}",
            font=("Arial", 10),
            bg="white",
            anchor="w"
        ).pack(anchor="w", padx=15, pady=2)

        tk.Label(
            summary_frame,
            text=f"Sheet: {self.sheet_name}",
            font=("Arial", 10),
            bg="white",
            anchor="w"
        ).pack(anchor="w", padx=15, pady=2)

        tk.Label(
            summary_frame,
            text=f"Column: {self.column_name}",
            font=("Arial", 10),
            bg="white",
            anchor="w"
        ).pack(anchor="w", padx=15, pady=(2, 10))

        # Buttons
        btn_frame = tk.Frame(main_frame, bg=self.bg_color)
        btn_frame.pack(pady=30)

        tk.Button(
            btn_frame,
            text="← Back",
            font=("Arial", 11),
            bg="#757575",
            fg="white",
            width=15,
            relief="flat",
            cursor="hand2",
            command=self.show_column_selection
        ).pack(side="left", padx=10)

        tk.Button(
            btn_frame,
            text="🚀 Start Analysis",
            font=("Arial", 11, "bold"),
            bg=self.success_color,
            fg="white",
            width=18,
            height=2,
            relief="flat",
            cursor="hand2",
            command=self.start_analysis
        ).pack(side="left", padx=10)

    def start_analysis(self):
        """Begin the analysis process"""
        self.show_progress_page()
        thread = threading.Thread(target=self.run_analysis, daemon=True)
        thread.start()

    def show_progress_page(self):
        """Display progress page"""
        self.clear_window()

        main_frame = tk.Frame(self.root, bg=self.bg_color)
        main_frame.place(relx=0.5, rely=0.5, anchor="center")

        title = tk.Label(
            main_frame,
            text="Analyzing Comments...",
            font=("Arial", 20, "bold"),
            bg=self.bg_color
        )
        title.pack(pady=(0, 30))

        self.progress_label = tk.Label(
            main_frame,
            text="Initializing analysis...",
            font=("Arial", 12),
            bg=self.bg_color
        )
        self.progress_label.pack(pady=10)

        self.progress_bar = ttk.Progressbar(
            main_frame,
            length=500,
            mode='determinate'
        )
        self.progress_bar.pack(pady=20)

        self.batch_label = tk.Label(
            main_frame,
            text="Batch 0 of 0",
            font=("Arial", 11),
            bg=self.bg_color,
            fg="#666"
        )
        self.batch_label.pack(pady=5)

        # Add timer labels
        self.elapsed_time_label = tk.Label(
            main_frame,
            text="Elapsed: 0s",
            font=("Arial", 10, "bold"),
            bg=self.bg_color,
            fg="#2196F3"
        )
        self.elapsed_time_label.pack(pady=(10, 2))

        self.estimated_time_label = tk.Label(
            main_frame,
            text="Estimated remaining: Calculating...",
            font=("Arial", 10),
            bg=self.bg_color,
            fg="#666"
        )
        self.estimated_time_label.pack(pady=(2, 10))

        self.status_label = tk.Label(
            main_frame,
            text="Please wait... This may take several minutes.",
            font=("Arial", 10),
            bg=self.bg_color,
            fg="#666"
        )
        self.status_label.pack(pady=20)

    def update_progress(self, current, total, message="Processing..."):
        """Update progress bar and labels"""
        self.root.after(0, lambda: self._update_progress_ui(current, total, message))

    def _update_progress_ui(self, current, total, message):
        """Update UI elements"""
        import time

        # Calculate elapsed time
        if self.start_time is None:
            self.start_time = time.time()

        elapsed = time.time() - self.start_time
        self.elapsed_time_label.config(text=f"Elapsed: {self.format_time(elapsed)}")

        # Calculate estimated remaining time
        if current > 0 and total > 0:
            progress_percent = current / total
            if progress_percent > 0:
                estimated_total = elapsed / progress_percent
                remaining = estimated_total - elapsed

                if remaining > 0:
                    self.estimated_time_label.config(
                        text=f"Estimated remaining: {self.format_time(remaining)}"
                    )
                else:
                    self.estimated_time_label.config(text="Estimated remaining: Almost done!")
            else:
                self.estimated_time_label.config(text="Estimated remaining: Calculating...")

        # Update progress bar and labels
        progress_percent_display = (current / total * 100) if total > 0 else 0
        self.progress_bar['value'] = progress_percent_display
        self.batch_label.config(text=f"Batch {current} of {total}")
        self.progress_label.config(text=message)
        self.root.update_idletasks()

    def format_time(self, seconds):
        """Format seconds into readable time string"""
        if seconds < 60:
            return f"{int(seconds)}s"
        elif seconds < 3600:
            minutes = int(seconds // 60)
            secs = int(seconds % 60)
            return f"{minutes}m {secs}s"
        else:
            hours = int(seconds // 3600)
            minutes = int((seconds % 3600) // 60)
            return f"{hours}h {minutes}m"

    def run_analysis(self):
        """Run the actual analysis - SAME LOGIC AS MAIN PROGRAM"""
        # Reset timer
        self.start_time = None
        self.estimated_total_time = None
        try:
            # Load data using the same method as main program
            df = pd.read_excel(self.file_path, sheet_name=self.sheet_name)

            # Normalize headers just like main program
            df.columns = df.columns.astype(str).str.strip()

            # Use ID column if found, otherwise use default
            id_col = self.id_column_name if self.id_column_name else "Response ID"
            text_col = self.column_name

            # Check columns exist
            if id_col not in df.columns:
                # Create Response ID from index like main program would
                df.insert(0, 'Response ID', df.index)
                id_col = 'Response ID'

            if text_col not in df.columns:
                raise ValueError(f"Column '{text_col}' not found in sheet")

            # Prepare comments EXACTLY like main program
            comments = []
            for idx, row in df.iterrows():
                comment_text = str(row[text_col]).strip()
                if comment_text and comment_text.lower() not in ['nan', 'none', '']:
                    comments.append({
                        'id': row[id_col],
                        'text': comment_text
                    })

            if not comments:
                raise ValueError("No valid comments found in the selected column")

            print(f"Found {len(comments)} comments to analyze")
            print(f"Using model: {self.model_var.get()}")

            # Analyze in batches using USER'S selected batch size
            all_results = []
            batch_size = self.batch_size.get()  # Use GUI selected batch size
            total_batches = (len(comments) - 1) // batch_size + 1

            for i in range(0, len(comments), batch_size):
                batch = comments[i:i + batch_size]
                batch_num = i // batch_size + 1

                self.update_progress(
                    batch_num - 1,
                    total_batches,
                    f"Analyzing batch {batch_num} of {total_batches}..."
                )

                # Call analysis function with user's batch
                batch_results = analyze_comments_batch(batch, self.model_var.get(), self.ollama_url)
                all_results.extend(batch_results)

                print(f"✓ Batch {batch_num}/{total_batches} ({len(batch_results)} analyzed)")

            # Create results DataFrame EXACTLY like main program
            results_data = []
            result_map = {r['id']: r for r in all_results}

            for idx, row in df.iterrows():
                comment_id = row[id_col]
                comment_text = str(row[text_col])

                if comment_id in result_map:
                    res = result_map[comment_id]
                    topics = '; '.join(res.get('temas', []))
                    sentiment = res.get('sentimiento', 'NEUTRAL')
                    score = res.get('puntuacion', 0)
                else:
                    topics = ''
                    sentiment = 'NO_RESPONSE'
                    score = 0

                results_data.append({
                    'response_id': comment_id,
                    'comment': comment_text,
                    'topics': topics,
                    'sentiment': sentiment,
                    'sentiment_score': score
                })

            results_df = pd.DataFrame(results_data)

            # Save results to TEMPORARY location first
            self.update_progress(total_batches, total_batches, "Generating visualizations...")

            import tempfile
            import os

            # Create temporary file for results
            temp_dir = tempfile.gettempdir()
            temp_xlsx = Path(temp_dir) / f"{self.file_path.stem}_topic_analysis_temp.xlsx"
            temp_csv = Path(temp_dir) / f"{self.file_path.stem}_topic_analysis_temp.csv"

            results_df.to_excel(temp_xlsx, index=False, sheet_name='Resultados Detallados')
            results_df.to_csv(temp_csv, index=False, encoding='utf-8-sig')

            # Create visualizations using same functions as main program
            print("Creating visualizations...")
            create_visualizations(temp_xlsx, results_df)
            create_topic_analysis(temp_xlsx, results_df)

            # Store the temp path - will be moved when user downloads
            self.results_path = temp_xlsx
            self.results_csv_path = temp_csv

            # Print summary like main program
            print("\n" + "=" * 60)
            print("RESUMEN DEL ANÁLISIS")
            print("=" * 60)

            responded = results_df[results_df['sentiment'] != 'NO_RESPONSE']
            print(f"Total de comentarios analizados: {len(responded)}")
            print(f"Puntuación promedio: {responded['sentiment_score'].mean():.2f} / 10")
            print(f"\nDistribución de sentimientos:")
            print(f"  Positivos: {(responded['sentiment'] == 'POSITIVO').sum()}")
            print(f"  Neutrales: {(responded['sentiment'] == 'NEUTRAL').sum()}")
            print(f"  Negativos: {(responded['sentiment'] == 'NEGATIVO').sum()}")

            print(f"\n✓ Resultados guardados temporalmente")

            # Unload Ollama model to free VRAM
            self.unload_ollama_model()

            # Show completion page
            self.root.after(0, self.show_completion_page)

        except Exception as e:
            import traceback
            error_details = traceback.format_exc()
            print(f"ERROR: {error_details}")
            self.root.after(0, lambda: self.show_error(f"{str(e)}\n\nDetails:\n{error_details}"))

    def unload_ollama_model(self):
        """Unload the Ollama model to free up VRAM"""
        try:
            print("\nUnloading Ollama model to free VRAM...")
            url = f"{self.ollama_url}/api/generate"

            # Send request with empty prompt and keep_alive=0 to unload
            payload = {
                "model": self.model_var.get(),
                "prompt": "",
                "keep_alive": 0
            }

            response = requests.post(url, json=payload, timeout=10)
            if response.status_code == 200:
                print("✓ Ollama model unloaded successfully")
            else:
                print(f"⚠ Could not unload model (status {response.status_code})")
        except Exception as e:
            print(f"⚠ Error unloading model: {e}")

    def show_error(self, error_message):
        """Display error message"""
        messagebox.showerror("Analysis Error", f"An error occurred:\n\n{error_message}")
        self.show_main_page()

    def show_completion_page(self):
        """Display completion page"""
        self.clear_window()

        main_frame = tk.Frame(self.root, bg=self.bg_color)
        main_frame.place(relx=0.5, rely=0.5, anchor="center")

        tk.Label(
            main_frame,
            text="✓",
            font=("Arial", 60, "bold"),
            bg=self.bg_color,
            fg=self.success_color
        ).pack(pady=(0, 20))

        title = tk.Label(
            main_frame,
            text="Analysis Complete!",
            font=("Arial", 22, "bold"),
            bg=self.bg_color,
            fg=self.success_color
        )
        title.pack(pady=(0, 10))

        # Show total time if available
        if self.start_time:
            import time
            total_time = time.time() - self.start_time
            time_label = tk.Label(
                main_frame,
                text=f"Total time: {self.format_time(total_time)}",
                font=("Arial", 11, "bold"),
                bg=self.bg_color,
                fg="#2196F3"
            )
            time_label.pack(pady=(0, 10))

        subtitle = tk.Label(
            main_frame,
            text="Your results are ready to download",
            font=("Arial", 12),
            bg=self.bg_color,
            fg="#666"
        )
        subtitle.pack(pady=(0, 40))

        download_btn = tk.Button(
            main_frame,
            text="📥 Download Results",
            font=("Arial", 13, "bold"),
            bg=self.primary_color,
            fg="white",
            width=25,
            height=2,
            relief="flat",
            cursor="hand2",
            command=self.download_results
        )
        download_btn.pack(pady=(0, 20))

        return_btn = tk.Button(
            main_frame,
            text="Return to Main Page",
            font=("Arial", 11),
            bg="#757575",
            fg="white",
            width=25,
            relief="flat",
            cursor="hand2",
            command=self.show_main_page
        )
        return_btn.pack()

    def download_results(self):
        """Allow user to save results file"""
        if not self.results_path or not self.results_path.exists():
            messagebox.showerror("Error", "Results file not found")
            return

        save_path = filedialog.asksaveasfilename(
            title="Save Analysis Results",
            defaultextension=".xlsx",
            filetypes=[("Excel files", "*.xlsx"), ("All files", "*.*")],
            initialfile=self.results_path.name
        )

        if save_path:
            try:
                import shutil
                shutil.copy2(self.results_path, save_path)
                self.show_download_complete()
            except Exception as e:
                messagebox.showerror("Error", f"Failed to save file:\n{str(e)}")

    def show_download_complete(self):
        """Show download complete message"""
        self.clear_window()

        main_frame = tk.Frame(self.root, bg=self.bg_color)
        main_frame.place(relx=0.5, rely=0.5, anchor="center")

        tk.Label(
            main_frame,
            text="✓",
            font=("Arial", 60, "bold"),
            bg=self.bg_color,
            fg=self.success_color
        ).pack(pady=(0, 20))

        title = tk.Label(
            main_frame,
            text="Download Complete",
            font=("Arial", 22, "bold"),
            bg=self.bg_color,
            fg=self.success_color
        )
        title.pack(pady=(0, 10))

        subtitle = tk.Label(
            main_frame,
            text="Your analysis results have been saved successfully",
            font=("Arial", 12),
            bg=self.bg_color,
            fg="#666"
        )
        subtitle.pack(pady=(0, 40))

        return_btn = tk.Button(
            main_frame,
            text="Return to Main Page",
            font=("Arial", 12, "bold"),
            bg=self.primary_color,
            fg="white",
            width=25,
            height=2,
            relief="flat",
            cursor="hand2",
            command=self.show_main_page
        )
        return_btn.pack()

    def exit_app(self):
        """Safely exit the application"""
        if messagebox.askokcancel("Exit", "Are you sure you want to exit?"):
            # Unload Ollama model before exiting
            self.unload_ollama_model()

            # Clean up any temp files
            try:
                if hasattr(self, 'results_path') and self.results_path and self.results_path.exists():
                    self.results_path.unlink()
                if hasattr(self, 'results_csv_path') and self.results_csv_path and self.results_csv_path.exists():
                    self.results_csv_path.unlink()
            except:
                pass

            self.root.quit()
            self.root.destroy()


# ==================== COMMAND LINE INTERFACE ====================

def parse_args():
    ap = argparse.ArgumentParser(description="Analyze Spanish comments for topics and sentiment")
    ap.add_argument("--input", default=None, help="Excel file path")
    ap.add_argument("--sheet", default=None, help="Sheet name")
    ap.add_argument("--id-col", default="Response ID")
    ap.add_argument("--text-col", default="Comments")
    ap.add_argument("--ollama-url", default="http://localhost:11434")
    ap.add_argument("--model", default="gemma2:2b")
    ap.add_argument("--batch-size", type=int, default=15)
    ap.add_argument("--gui", action="store_true", help="Launch GUI mode")
    return ap.parse_args()


def main_cli(args):
    """Command line interface - original program logic"""
    in_path = Path(args.input)

    if not in_path.exists():
        raise SystemExit(f"File not found: {in_path}")

    print(f"Loading data from {in_path}...")
    df, chosen_sheet, all_sheets = load_single_sheet_excel(in_path, args.sheet)

    if df.empty:
        raise SystemExit(f"Sheet '{chosen_sheet}' is empty. Available: {all_sheets}")

    df.columns = df.columns.astype(str).str.strip()

    id_col = args.id_col
    text_col = args.text_col

    if id_col not in df.columns or text_col not in df.columns:
        raise SystemExit(f"Required columns not found. Available: {list(df.columns)}")

    # Prepare comments
    comments = []
    for idx, row in df.iterrows():
        comment_text = str(row[text_col]).strip()
        if comment_text and comment_text.lower() not in ['nan', 'none', '']:
            comments.append({
                'id': row[id_col],
                'text': comment_text
            })

    print(f"Found {len(comments)} comments to analyze")
    print(f"Using model: {args.model}")
    print(f"Ollama URL: {args.ollama_url}")

    # Analyze in batches
    all_results = []
    batch_size = args.batch_size

    for i in range(0, len(comments), batch_size):
        batch = comments[i:i + batch_size]
        print(f"Analyzing batch {i // batch_size + 1}/{(len(comments) - 1) // batch_size + 1}...", end=" ")

        batch_results = analyze_comments_batch(batch, args.model, args.ollama_url)
        all_results.extend(batch_results)
        print(f"✓ ({len(batch_results)} analyzed)")

    # Create results DataFrame
    results_data = []
    result_map = {r['id']: r for r in all_results}

    for idx, row in df.iterrows():
        comment_id = row[id_col]
        comment_text = str(row[text_col])

        if comment_id in result_map:
            res = result_map[comment_id]
            topics = '; '.join(res.get('temas', []))
            sentiment = res.get('sentimiento', 'NEUTRAL')
            score = res.get('puntuacion', 0)
        else:
            topics = ''
            sentiment = 'NO_RESPONSE'
            score = 0

        results_data.append({
            'response_id': comment_id,
            'comment': comment_text,
            'topics': topics,
            'sentiment': sentiment,
            'sentiment_score': score
        })

    results_df = pd.DataFrame(results_data)

    # Save results
    out_xlsx = in_path.with_name(in_path.stem + "_topic_analysis.xlsx")
    out_csv = in_path.with_name(in_path.stem + "_topic_analysis.csv")

    results_df.to_excel(out_xlsx, index=False, sheet_name='Resultados Detallados')
    results_df.to_csv(out_csv, index=False, encoding='utf-8-sig')

    print("Creating visualizations...")
    create_visualizations(out_xlsx, results_df)
    create_topic_analysis(out_xlsx, results_df)

    # Print summary
    print("\n" + "=" * 60)
    print("RESUMEN DEL ANÁLISIS")
    print("=" * 60)

    responded = results_df[results_df['sentiment'] != 'NO_RESPONSE']
    print(f"Total de comentarios analizados: {len(responded)}")
    print(f"Puntuación promedio: {responded['sentiment_score'].mean():.2f} / 10")
    print(f"\nDistribución de sentimientos:")
    print(f"  Positivos: {(responded['sentiment'] == 'POSITIVO').sum()}")
    print(f"  Neutrales: {(responded['sentiment'] == 'NEUTRAL').sum()}")
    print(f"  Negativos: {(responded['sentiment'] == 'NEGATIVO').sum()}")

    print(f"\n✓ Resultados guardados en:")
    print(f"  - {out_xlsx}")
    print(f"  - {out_csv}")


def main():
    args = parse_args()

    # Launch GUI if --gui flag or no input file specified
    if args.gui or args.input is None:
        root = tk.Tk()
        app = TopicAnalyzerGUI(root)
        root.mainloop()
    else:
        # Run command line version
        main_cli(args)


if __name__ == "__main__":
    main()