# ============================================================
# app.py — Tortured Phrase Detector (ZeroGPU-compatible)
# ============================================================

import spaces          # ✅ MUST be the very first import

import os
import re
import json
import torch
import gradio as gr
import pandas as pd
from datetime import datetime
from transformers import AutoModelForCausalLM, AutoTokenizer
from peft import PeftModel

# ════════════════════════════════════════════════════════════
#  CONFIG
# ════════════════════════════════════════════════════════════

BASE_MODEL_ID        = "meta-llama/Llama-3.2-1B-Instruct"
STANDARD_ADAPTER_ID  = "skygg/paperlytix_non_reasoning_medical_trained_adapter"
REASONING_ADAPTER_ID = "skygg/paperlytix_reasoning_medical_trained_adapter"
SONIC_DB_PATH        = "./search_tortured_correct.csv"

SONIC_RESULTS_DIR      = "./sonic_text_results"
DEEPER_RESULTS_DIR     = "./go_deeper_text_results"
HYBRIONIX_RESULTS_DIR  = "./hybrionix_text_results"
GO_BEYOND_RESULTS_DIR  = "./go_beyond_text_results"

MAX_NEW_TOKENS            = 512
MAX_NEW_TOKENS_REASONING  = 1024
MAX_INPUT_CHARS           = 3000

HF_TOKEN = os.environ.get("HF_TOKEN", None)


# ════════════════════════════════════════════════════════════
#  SYSTEM PROMPTS
# ════════════════════════════════════════════════════════════

STANDARD_SYSTEM_PROMPT = (
    "You are a scientific language correction assistant specializing in "
    "Clinical Medicine. Your task is to detect 'tortured phrases' in the "
    "given sentence — these are unnatural, paraphrased substitutions of "
    "standard medical terminology, often introduced by paraphrasing tools "
    "or paper mills to evade plagiarism detection. Identify each tortured "
    "phrase, replace it with the correct and widely accepted clinical term, "
    "and briefly explain why the original phrasing is non-standard. If no "
    "tortured phrases are found, explicitly state that the sentence uses "
    "correct clinical terminology."
)

REASONING_SYSTEM_PROMPT = (
    "You are a scientific language correction assistant specializing in "
    "Clinical Medicine. Your task is to detect 'tortured phrases' in the "
    "given sentence — these are unnatural, paraphrased substitutions of "
    "standard medical terminology, often introduced by paraphrasing tools "
    "or paper mills to evade plagiarism detection.\n\n"
    "Follow this strict chain-of-thought reasoning process before giving "
    "your final answer:\n\n"
    "Step 1 - Segment the sentence: Break the sentence into individual "
    "medical phrases or noun groups.\n"
    "Step 2 - Evaluate each phrase: For each phrase, ask — Is this a "
    "recognized, standard clinical/medical term used in peer-reviewed "
    "literature or clinical guidelines? If not, flag it as a potential "
    "tortured phrase.\n"
    "Step 3 - Identify the distortion pattern: For each flagged phrase, "
    "explain what standard term it appears to be paraphrasing and why the "
    "substitution is linguistically or clinically non-standard.\n"
    "Step 4 - Provide the correction: Replace each tortured phrase with "
    "the correct, widely accepted clinical terminology.\n"
    "Step 5 - Output the corrected sentence: Present the fully corrected "
    "sentence.\n\n"
    "If no tortured phrases are found after this reasoning, explicitly "
    "state that the sentence uses correct clinical terminology."
)


# ════════════════════════════════════════════════════════════
#  GLOBAL MODEL STATE
# ════════════════════════════════════════════════════════════

_tokenizer = None
_model     = None
_sonic_df  = None


def load_models():
    """
    Load base model + attach both adapters, then move to 'cuda'
    directly — this is the officially documented ZeroGPU pattern.
    The `spaces` package (imported first, above) patches torch so
    this call is safe even though no physical GPU is attached to
    this process yet; the real transfer happens automatically when
    a @spaces.GPU function is later invoked.

    NOTE: the adapters are explicitly loaded onto CPU (torch_device="cpu")
    because PEFT auto-infers "cuda" as the load device (since `spaces`
    reports torch.cuda.is_available() == True) and tries to load the
    safetensors weights *directly* onto a real CUDA device — which fails
    with "No CUDA GPUs are available" since no physical GPU is attached
    at this point in the process. Loading to CPU first and moving the
    whole composed model with `.to("cuda")` afterward is the supported
    deferred-placement pattern.
    """
    global _tokenizer, _model

    print("📂 Loading tokenizer ...")
    _tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL_ID, token=HF_TOKEN)
    if _tokenizer.pad_token is None:
        _tokenizer.pad_token    = _tokenizer.eos_token
        _tokenizer.pad_token_id = _tokenizer.eos_token_id

    print("📂 Loading base LLaMA model ...")
    base_model = AutoModelForCausalLM.from_pretrained(
        BASE_MODEL_ID,
        dtype=torch.float16,      # ✅ dtype instead of torch_dtype
        token=HF_TOKEN
    )

    print(f"📂 Attaching standard adapter: {STANDARD_ADAPTER_ID}")
    _model = PeftModel.from_pretrained(
        base_model,
        STANDARD_ADAPTER_ID,
        adapter_name="standard",
        token=HF_TOKEN,
        torch_device="cpu"
    )

    print(f"📂 Attaching reasoning adapter: {REASONING_ADAPTER_ID}")
    _model.load_adapter(
        REASONING_ADAPTER_ID,
        adapter_name="reasoning",
        token=HF_TOKEN,
        torch_device="cpu"
    )

    # ── Official ZeroGPU pattern: move to cuda right after load ─
    _model.to("cuda")
    _model.eval()
    print("✅ Model ready with both adapters, moved to CUDA.\n")


def load_sonic_db():
    global _sonic_df
    if _sonic_df is None:
        print("📂 Loading Sonic helper database ...")
        _sonic_df = pd.read_csv(SONIC_DB_PATH, encoding="utf-8")
        print(f"✅ Sonic database loaded — {len(_sonic_df)} pairs.\n")
    return _sonic_df


# ════════════════════════════════════════════════════════════
#  SHARED UTILITIES
# ════════════════════════════════════════════════════════════

def split_into_sentences(text):
    raw       = re.split(r'(?<=[.!?])\s+', text.strip())
    sentences = [s.strip() for s in raw if s.strip()]
    return sentences


def validate_input(text):
    if not text or not text.strip():
        return False, "Please enter some text to analyze."
    if len(text) > MAX_INPUT_CHARS:
        return False, (
            f"Input too long ({len(text)} characters). "
            f"Please limit to {MAX_INPUT_CHARS} characters per request."
        )
    return True, None


def build_pattern(phrase_lower):
    escaped = re.escape(phrase_lower)
    return re.compile(r'\b' + escaped + r"s?\b", re.IGNORECASE)


def get_next_result_index(folder_path):
    os.makedirs(folder_path, exist_ok=True)
    existing = [
        f for f in os.listdir(folder_path)
        if f.startswith("test_") and f.endswith("_results.json")
    ]
    if not existing:
        return 1
    indices = []
    for fname in existing:
        try:
            idx = int(fname.replace("test_", "").replace("_results.json", ""))
            indices.append(idx)
        except ValueError:
            continue
    return max(indices) + 1 if indices else 1


def save_results(folder_path, engine_name, input_text,
                 sentences, results_per_sentence, extra_data=None):
    os.makedirs(folder_path, exist_ok=True)
    idx      = get_next_result_index(folder_path)
    filepath = os.path.join(folder_path, f"test_{idx}_results.json")

    output = {
        "test_index"      : idx,
        "engine"          : engine_name,
        "timestamp"       : datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
        "input_text"      : input_text,
        "total_sentences" : len(sentences),
        "results"         : []
    }
    if extra_data:
        output.update(extra_data)

    for i, (sentence, result) in enumerate(
        zip(sentences, results_per_sentence), start=1
    ):
        output["results"].append({
            "sentence_index" : i,
            "original"       : sentence,
            "output"         : result
        })

    with open(filepath, "w", encoding="utf-8") as f:
        json.dump(output, f, indent=2, ensure_ascii=False)

    return filepath


# ════════════════════════════════════════════════════════════
#  ENGINE : SONIC  (pure CPU — CSV lookup, no GPU needed)
# ════════════════════════════════════════════════════════════

def sonic_process_sentence(sentence, df):
    found = []
    for _, row in df.iterrows():
        tortured_lower = str(row["tortured_lower"]).strip()
        correct        = str(row["correct"]).strip()
        tortured_orig  = str(row["tortured"]).strip()

        pattern = build_pattern(tortured_lower)
        match   = pattern.search(sentence)
        if match:
            found.append((tortured_orig, match.group(0), correct))

    if not found:
        return "No tortured phrases detected."

    corrected = sentence
    for tortured_orig, matched_text, correct in found:
        pattern = build_pattern(tortured_orig.lower())
        replacement = correct + "s" if matched_text.lower() == \
            (tortured_orig.lower() + "s") else correct
        if matched_text[0].isupper():
            replacement = replacement[0].upper() + replacement[1:]
        corrected = pattern.sub(replacement, corrected, count=1)

    lines = [f"Corrected: {corrected}", "", "Tortured phrases detected:"]
    for idx, (tortured_orig, matched_text, correct) in enumerate(found, 1):
        lines.append(f'{idx}. [Tortured Phrase] "{matched_text}" → "{correct}"')
    return "\n".join(lines)


def run_sonic(sentences):
    df = load_sonic_db()
    return [sonic_process_sentence(s, df) for s in sentences]


# ════════════════════════════════════════════════════════════
#  GPU-DECORATED GENERATION FUNCTION
# ════════════════════════════════════════════════════════════

@spaces.GPU(duration=90)
def generate_batch(sentences, adapter_name, system_prompt, max_new_tokens):
    """
    ZeroGPU attaches a real GPU only for the duration of this call.
    Model is already on 'cuda' (set at load time) — no device
    transfer needed here.
    """
    _model.set_adapter(adapter_name)

    results = []
    for sentence in sentences:
        messages = [
            {"role": "system", "content": system_prompt},
            {"role": "user",   "content": sentence}
        ]

        input_text = _tokenizer.apply_chat_template(
            messages,
            tokenize=False,
            add_generation_prompt=True
        )

        inputs = _tokenizer(input_text, return_tensors="pt").to("cuda")

        with torch.no_grad():
            output_ids = _model.generate(
                **inputs,
                max_new_tokens = max_new_tokens,
                do_sample      = False,
                pad_token_id   = _tokenizer.eos_token_id
            )

        generated_ids = output_ids[0][inputs["input_ids"].shape[-1]:]
        response = _tokenizer.decode(generated_ids, skip_special_tokens=True)
        results.append(response.strip())

    return results


# ════════════════════════════════════════════════════════════
#  MAIN DISPATCH
# ════════════════════════════════════════════════════════════

def run_detection(text, engine_choice, progress=gr.Progress()):

    is_valid, error = validate_input(text)
    if not is_valid:
        return f"⚠️ {error}"

    sentences = split_into_sentences(text)
    progress(0.2, desc="Analyzing your text...")

    if engine_choice.startswith("⚡"):
        results = run_sonic(sentences)
        save_results(SONIC_RESULTS_DIR, "Sonic", text, sentences, results)

    elif engine_choice.startswith("🧠"):
        progress(0.4, desc="Running AI detection (Go Deeper)...")
        results = generate_batch(
            sentences, "standard", STANDARD_SYSTEM_PROMPT, MAX_NEW_TOKENS
        )
        save_results(DEEPER_RESULTS_DIR, "Go Deeper", text, sentences, results)

    elif engine_choice.startswith("🔮"):
        progress(0.4, desc="Running AI reasoning (Go Beyond)...")
        results = generate_batch(
            sentences, "reasoning", REASONING_SYSTEM_PROMPT,
            MAX_NEW_TOKENS_REASONING
        )
        save_results(GO_BEYOND_RESULTS_DIR, "Go Beyond", text, sentences, results)

    elif engine_choice.startswith("🔀"):
        progress(0.3, desc="Running Go Deeper (AI)...")
        deeper_results = generate_batch(
            sentences, "standard", STANDARD_SYSTEM_PROMPT, MAX_NEW_TOKENS
        )
        progress(0.7, desc="Running Sonic cross-check...")
        sonic_results = run_sonic(sentences)

        results = []
        for d, s_r in zip(deeper_results, sonic_results):
            agree = "No tortured phrases" in d and "No tortured phrases" in s_r
            if agree:
                results.append(f"{d}\n\n✅ Confirmed by both Sonic and Go Deeper.")
            else:
                results.append(f"{d}\n\n---\n[Sonic cross-check]\n{s_r}")

        save_results(
            HYBRIONIX_RESULTS_DIR, "Hybrionix", text, sentences, results,
            extra_data={
                "go_deeper_results": deeper_results,
                "sonic_results": sonic_results
            }
        )

    else:
        return "⚠️ Please select a valid engine."

    progress(0.95, desc="Formatting results...")

    output_lines = []
    for i, (sentence, result) in enumerate(zip(sentences, results), start=1):
        output_lines.append(f"### Sentence {i}")
        output_lines.append(f"**Original:** {sentence}")
        output_lines.append("")
        output_lines.append(result)
        output_lines.append("\n---\n")

    return "\n".join(output_lines)


# ════════════════════════════════════════════════════════════
#  GRADIO INTERFACE
# ════════════════════════════════════════════════════════════

ENGINE_CHOICES = [
    "⚡ Sonic (fastest)",
    "🧠 Go Deeper (AI, direct)",
    "🔀 Hybrionix (AI + Sonic)",
    "🔮 Go Beyond (AI, reasoning)",
]

EXAMPLE_INPUTS = [
    ["The patient was diagnosed with bosom malignant growth and referred "
     "for assessment of cerebrum dead tissue.", "🔮 Go Beyond (AI, reasoning)"],
    ["The patient was diagnosed with myocardial infarction and started on "
     "dual antiplatelet therapy.", "⚡ Sonic (fastest)"],
    ["Cardiovascular breakdown was noted alongside glucose narrow "
     "mindedness in the follow-up report.", "🔀 Hybrionix (AI + Sonic)"],
]

DESCRIPTION = """
# 🔬 Tortured Phrase Detector — Clinical Medicine

**What is this?**
A free tool that scans clinical/medical research text for **"tortured phrases"** —
unnatural, paraphrased substitutions of standard medical terminology often
introduced by paraphrasing tools or paper mills to evade plagiarism detection.

**How to use it:**
1. Paste a sentence or paragraph from a clinical research paper below
2. Choose a detection engine (see guide below)
3. Click **Detect & Correct**

**Choosing an engine:**
| Engine | Speed | Best for |
|---|---|---|
| ⚡ Sonic | Instant | Quick checks, no AI needed |
| 🧠 Go Deeper | Fast (GPU) | AI-powered detection with brief explanations |
| 🔀 Hybrionix | Fast (GPU) | Cross-checked results from both Sonic + AI |
| 🔮 Go Beyond | Fast (GPU) | Full step-by-step reasoning — best for review/audit |

⚡ *Powered by ZeroGPU — a GPU is allocated on-demand for each AI request.*
"""

with gr.Blocks(title="Tortured Phrase Detector") as demo:
    gr.Markdown(DESCRIPTION)

    with gr.Row():
        with gr.Column():
            text_input = gr.Textbox(
                label="Paste your sentence(s) here",
                placeholder="e.g. The patient was diagnosed with bosom "
                            "malignant growth...",
                lines=6
            )
            engine_choice = gr.Radio(
                choices=ENGINE_CHOICES,
                value="⚡ Sonic (fastest)",
                label="Choose your detection engine"
            )
            submit_btn = gr.Button("🔍 Detect & Correct", variant="primary")

        with gr.Column():
            output_box = gr.Markdown(label="Results")

    gr.Examples(
        examples=EXAMPLE_INPUTS,
        inputs=[text_input, engine_choice],
        label="Try an example ⬇️"
    )

    submit_btn.click(
        fn=run_detection,
        inputs=[text_input, engine_choice],
        outputs=output_box
    )

    gr.Markdown(
        "---\n"
        "⚠️ *This is a research-assistance tool, not a substitute for "
        "professional editorial or clinical review.*"
    )


# ════════════════════════════════════════════════════════════
#  STARTUP
# ════════════════════════════════════════════════════════════

for folder in [SONIC_RESULTS_DIR, DEEPER_RESULTS_DIR,
               HYBRIONIX_RESULTS_DIR, GO_BEYOND_RESULTS_DIR]:
    os.makedirs(folder, exist_ok=True)

print("🚀 Loading models at startup ...")
load_models()
load_sonic_db()
print("✅ App ready to launch!")

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