import gradio as gr
import spaces
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
from znum import Znum, Topsis, Promethee, Beast
import re
from helpers.utils import DEFAULT_QUERY, DEFAULT_QUERY2, DEFAULT_QUERY3

# Z-number mappings: value/confidence (1-5) to fuzzy trapezoidal numbers
A_MAP = {
    1: [2, 3, 3, 4],
    2: [4, 5, 5, 6],
    3: [6, 7, 7, 8],
    4: [8, 9, 9, 10],
    5: [10, 11, 11, 12],
}

B_MAP = {
    1: [0.2, 0.3, 0.3, 0.4],
    2: [0.3, 0.4, 0.4, 0.5],
    3: [0.4, 0.5, 0.5, 0.6],
    4: [0.5, 0.6, 0.6, 0.7],
    5: [0.6, 0.7, 0.7, 0.8],
}

SYSTEM_PROMPT = """\
Extract a Z-number decision matrix from the following user input.

## Z-Number Scales:
- Value (A-part):
    - benefit: 5 (excellent) → 4 (good) → 3 (moderate) → 2 (poor) → 1 (very poor)
    - neutral: 0
    - cost: -1 (very low cost) → -2 (low) → -3 (moderate) → -4 (high) → -5 (very high cost)
- Confidence (B-part): 5 (very confident) → 4 (confident) → 3 (somewhat confident) → 2 (uncertain) → 1 (very uncertain)

## Output Format:
Return ONLY a Markdown table in this exact format:

| | criterion_1 | criterion_2 | ... |
|---|---|---|---|
| type | benefit | cost | ... |
| alt_1 | 4:3 | -3:4 | ... |
| alt_2 | 3:4 | -2:5 | ... |
| ... | ... | ... | ... |
| weight | 3:2 | 4:3 | ... |

## Rules:
1. First row: empty cell, then criterion names (alphanumeric + underscores only)
2. Second row: "type", then either "benefit" or "cost" for each criterion
3. Middle rows: alternative names, then VALUE:CONFIDENCE pairs
4. Last row: "weight", then importance weights as VALUE:CONFIDENCE (always use positive values 1-5 for weights)
5. VALUE must be positive (1-5) for benefits, negative (-1 to -5) for costs
6. CONFIDENCE is always positive (1-5) regardless of criterion type
"""

# Global model and tokenizer (loaded once)
model = None
tokenizer = None


def load_model():
    """Load model and tokenizer (called once on first inference)."""
    global model, tokenizer
    if model is None:
        model_name = "nuriyev/Qwen3-4B-znum-decision-matrix"
        tokenizer = AutoTokenizer.from_pretrained(model_name)
        model = AutoModelForCausalLM.from_pretrained(
            model_name,
            device_map="auto",
            torch_dtype=torch.bfloat16,
        )
    return model, tokenizer


def parse_znum_pair(pair_str: str) -> Znum | None:
    """Convert 'N:M' string to Znum object using A_MAP and B_MAP."""
    try:
        parts = pair_str.strip().split(':')
        if len(parts) != 2:
            return None
        a_val = abs(int(parts[0]))
        b_val = int(parts[1])
        if a_val not in A_MAP or b_val not in B_MAP:
            return None
        return Znum(A_MAP[a_val], B_MAP[b_val])
    except (ValueError, KeyError):
        return None


def parse_markdown_table(text: str) -> dict:
    """Parse markdown table from model output into structured dict."""
    lines = [l.strip() for l in text.strip().split('\n') if l.strip() and '|' in l]
    lines = [l for l in lines if not re.match(r'^\|[-:\s|]+\|$', l)]

    if len(lines) < 4:
        return {}

    def split_row(row: str) -> list:
        cells = [c.strip() for c in row.split('|')]
        return [c for c in cells if c]

    headers = split_row(lines[0])
    criteria = headers  # All headers are criteria (empty first cell is filtered out)

    types_row = split_row(lines[1])
    types = types_row[1:] if len(types_row) > 1 else []

    weights_row = split_row(lines[-1])
    weights = weights_row[1:] if len(weights_row) > 1 else []

    alternatives = {}
    for line in lines[2:-1]:
        row = split_row(line)
        if row:
            alt_name = row[0]
            values = row[1:]
            alternatives[alt_name] = values

    return {
        'criteria': criteria,
        'types': types,
        'alternatives': alternatives,
        'weights': weights
    }


def format_table_html(matrix: dict) -> str:
    """Convert parsed matrix to a nicely formatted HTML table."""
    if not matrix or not matrix.get('criteria'):
        return "<p style='color:#666;'>No decision matrix generated yet.</p>"

    html = """
    <div style="overflow-x:auto;">
    <table style="border-collapse:collapse;width:100%;font-family:system-ui,sans-serif;font-size:13px;background:#fff;">
    <thead>
        <tr>
            <th style="border:1px solid #ddd;padding:10px 12px;text-align:left;background:#111;color:#fff;font-weight:500;">Alternative</th>
    """

    th_style = "border:1px solid #ddd;padding:10px 12px;text-align:center;background:#111;color:#fff;font-weight:500;"
    for crit in matrix['criteria']:
        html += f'<th style="{th_style}">{crit.replace("_", " ").title()}</th>'
    html += "</tr></thead><tbody>"

    td_style = "border:1px solid #ddd;padding:10px 12px;text-align:center;color:#333;"
    td_left = "border:1px solid #ddd;padding:10px 12px;text-align:left;color:#333;font-weight:500;"

    # Type row
    html += f'<tr style="background:#f9f9f9;"><td style="{td_left}">Type</td>'
    for t in matrix['types']:
        color = "#111" if t.lower() == "benefit" else "#666"
        html += f'<td style="{td_style}color:{color};">{t.capitalize()}</td>'
    html += "</tr>"

    # Alternative rows
    for i, (alt_name, values) in enumerate(matrix['alternatives'].items()):
        bg = "#fff" if i % 2 == 0 else "#f9f9f9"
        html += f'<tr style="background:{bg};"><td style="{td_left}">{alt_name.replace("_", " ").title()}</td>'
        for v in values:
            html += f'<td style="{td_style}font-family:monospace;">{v}</td>'
        html += "</tr>"

    # Weight row
    html += f'<tr style="background:#111;"><td style="border:1px solid #333;padding:10px 12px;text-align:left;color:#fff;font-weight:500;">Weight</td>'
    for w in matrix['weights']:
        html += f'<td style="border:1px solid #333;padding:10px 12px;text-align:center;color:#fff;font-family:monospace;">{w}</td>'
    html += "</tr>"

    html += "</tbody></table></div>"
    return html


def format_results_html(alt_names: list, solver, method: str) -> str:
    """Format MCDM results as HTML."""
    html = f"""
    <div style="font-family:system-ui,sans-serif;padding:20px;background:#111;border-radius:6px;color:#fff;">
        <div style="font-size:11px;font-weight:500;margin-bottom:16px;text-transform:uppercase;letter-spacing:1px;color:#888;">{method.upper()} Ranking</div>
    """

    for rank, idx in enumerate(solver.ordered_indices, 1):
        # Different circle colors for ranks
        if rank == 1:
            circle_bg, circle_color = "#fff", "#000"
        elif rank == 2:
            circle_bg, circle_color = "#666", "#fff"
        else:
            circle_bg, circle_color = "#444", "#fff"

        badge = '<span style="background:#fff;color:#000;padding:3px 10px;border-radius:3px;font-size:10px;font-weight:600;text-transform:uppercase;">BEST</span>' if rank == 1 else ""

        border = "border-bottom:1px solid #333;" if rank < len(solver.ordered_indices) else ""
        html += f"""
        <div style="display:flex;align-items:center;padding:12px 0;{border}">
            <span style="width:26px;height:26px;background:{circle_bg};color:{circle_color};border-radius:50%;display:inline-flex;align-items:center;justify-content:center;font-weight:600;font-size:12px;margin-right:14px;">{rank}</span>
            <span style="flex-grow:1;font-size:14px;">{alt_names[idx].replace('_', ' ').title()}</span>
            {badge}
        </div>
        """

    html += "</div>"
    return html


@spaces.GPU
def process_decision(query: str, method: str, progress=gr.Progress()):
    """Main processing function with ZeroGPU support."""
    if not query.strip():
        return "<p>Please enter a decision query.</p>", "<p>No results yet.</p>", ""

    progress(0.1, desc="Loading model...")
    model, tokenizer = load_model()

    progress(0.2, desc="Preparing input...")
    messages = [
        {"role": "system", "content": SYSTEM_PROMPT},
        {"role": "user", "content": query},
    ]

    prompt = tokenizer.apply_chat_template(
        messages,
        tokenize=False,
        add_generation_prompt=True,
        enable_thinking=False
    )
    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)

    progress(0.3, desc="Generating decision matrix...")
    with torch.no_grad():
        output_ids = model.generate(
            **inputs,
            max_new_tokens=2048,
            temperature=0.7,
            do_sample=True,
            pad_token_id=tokenizer.eos_token_id
        )

    generated_ids = output_ids[0][inputs['input_ids'].shape[1]:]
    generated_text = tokenizer.decode(generated_ids, skip_special_tokens=True)

    progress(0.6, desc="Parsing decision matrix...")
    matrix = parse_markdown_table(generated_text)

    if not matrix or not matrix.get('criteria'):
        return (
            "<p style='color: red;'>Failed to generate a valid decision matrix. Please try again with a clearer query.</p>",
            "<p>No results available.</p>",
            generated_text
        )

    # Format table HTML
    table_html = format_table_html(matrix)

    progress(0.8, desc=f"Applying {method.upper()}...")

    # Convert to Znum objects
    znum_weights = [parse_znum_pair(w) for w in matrix['weights']]
    znum_alternatives = {}
    for alt_name, values in matrix['alternatives'].items():
        znum_alternatives[alt_name] = [parse_znum_pair(v) for v in values]

    # Check for parsing errors
    if None in znum_weights or any(None in vals for vals in znum_alternatives.values()):
        return (
            table_html,
            "<p style='color: orange;'>Warning: Some Z-numbers could not be parsed. Results may be incomplete.</p>",
            generated_text
        )

    # Build criteria types
    criteria_types = [
        Beast.CriteriaType.BENEFIT if t.lower() == 'benefit' else Beast.CriteriaType.COST
        for t in matrix['types']
    ]

    # Build decision table
    alt_names = list(znum_alternatives.keys())
    alt_rows = [znum_alternatives[name] for name in alt_names]
    table = [znum_weights] + alt_rows + [criteria_types]

    # Apply MCDM method
    if method == "TOPSIS":
        solver = Topsis(table)
    else:
        solver = Promethee(table)

    solver.solve()

    progress(1.0, desc="Done!")

    results_html = format_results_html(alt_names, solver, method)

    return table_html, results_html, generated_text


# Build Gradio interface
with gr.Blocks(
    title="Text2MCDM",
    theme=gr.themes.Default(
        primary_hue="neutral",
        neutral_hue="slate",
    ),
    css="""
        .gradio-container { max-width: 960px !important; }
        .header { text-align: center; padding: 24px 0; margin-bottom: 16px; }
        .header h1 { font-size: 1.75rem; font-weight: 600; color: #111; margin: 0 0 4px 0; }
        .header p { color: #666; font-size: 0.9rem; margin: 0; }
    """
) as demo:
    gr.HTML('''
        <div class="header">
            <h1>Text2MCDM</h1>
            <p>Transform decision narratives into structured Z-number analysis</p>
        </div>
    ''')

    query_input = gr.Textbox(
        label="Decision Narrative",
        placeholder="Describe your decision: What are your options? What factors matter? How confident are you about each?",
        lines=5,
        value=DEFAULT_QUERY
    )

    with gr.Row():
        method_dropdown = gr.Dropdown(
            choices=["TOPSIS", "PROMETHEE"],
            value="TOPSIS",
            label="Method",
            scale=1
        )
        submit_btn = gr.Button("Analyze", variant="primary", scale=2)

    gr.Markdown("---")

    with gr.Row():
        with gr.Column():
            gr.Markdown("**Decision Matrix**")
            table_output = gr.HTML(value="<p style='color:#888;'>Results will appear here.</p>")

        with gr.Column():
            gr.Markdown("**Ranking**")
            results_output = gr.HTML(value="<p style='color:#888;'>Results will appear here.</p>")

    with gr.Accordion("Raw Model Output", open=False):
        raw_output = gr.Textbox(label="Generated Text", lines=6, interactive=False)

    with gr.Accordion("How it works", open=False):
        gr.Markdown("""
        1. Describe your decision problem in natural language
        2. The LLM extracts alternatives, criteria, and ratings
        3. Z-numbers capture both **value** and **confidence** (format: `value:confidence`)
        4. MCDM algorithm (TOPSIS or PROMETHEE) ranks your options

        **Scale:** Values 1-5 for benefits, -1 to -5 for costs. Confidence always 1-5.
        """)

    gr.Examples(
        examples=[
            [DEFAULT_QUERY, "TOPSIS"],
            [DEFAULT_QUERY2, "TOPSIS"],
            [DEFAULT_QUERY3, "PROMETHEE"],
        ],
        inputs=[query_input, method_dropdown],
        label="Examples"
    )

    submit_btn.click(
        fn=process_decision,
        inputs=[query_input, method_dropdown],
        outputs=[table_output, results_output, raw_output]
    )

if __name__ == "__main__":
    demo.launch()