# ==========================================================
# METHODOLOGY EXTRACTION MODULE (FINAL SV APPROACH)
# FILE: models/methodology_extraction.py
#
# Purpose:
# - Extract methodology steps for poster flowchart
# - Hybrid robust approach:
#   (A) Bullet/numbered detection (high priority)
#   (B) Regex keyword-based action sentence detection
#   (C) Optional spaCy action sentence filtering
#   (D) Optional Flan-T5-base refinement (reformatting)
#   (E) Fallback safe steps if everything fails
#
# Output:
# [
#   "Step 1 ...",
#   "Step 2 ...",
#   ...
# ]
#
# Notes (SV-style design):
# - We DO NOT try to generate a perfect flowchart directly.
# - We ensure system always returns usable steps.
# - We keep steps short, actionable, poster-friendly.
# ==========================================================

import re


# ==========================================================
# NORMALIZE TEXT
# ==========================================================
def normalize_text(text: str) -> str:
    if not text:
        return ""
    text = text.replace("\r", "\n")
    text = re.sub(r"\n{3,}", "\n\n", text)
    text = re.sub(r"[ \t]{2,}", " ", text)
    return text.strip()


# ==========================================================
# REMOVE IEEE / PDF NOISE
# ==========================================================
def remove_ieee_noise(text: str) -> str:
    if not text:
        return ""

    noise_patterns = [
        r"authorized licensed use.*",
        r"downloaded on.*",
        r"ieee xplore.*",
        r"copyright.*",
        r"doi:\s*\S+",
        r"\(cid:\d+\)",
        r"this article has been accepted.*",
        r"personal use is permitted.*",
        r"vol\.\s*\d+.*",
        r"pp\.\s*\d+.*",
    ]

    for pat in noise_patterns:
        text = re.sub(pat, "", text, flags=re.IGNORECASE)

    text = re.sub(r"\s{2,}", " ", text)
    return text.strip()


# ==========================================================
# CLEAN STEP TEXT
# ==========================================================
def clean_step_text(step: str) -> str:
    if not step:
        return ""

    step = step.strip()

    # remove bullets
    step = step.replace("•", "")
    step = step.replace("◦", "")
    step = step.replace("▪", "")
    step = step.replace("–", "-")
    step = step.replace("—", "-")

    # remove numbering prefix
    step = re.sub(r"^\(?[0-9]{1,2}\)?[\.\)]\s*", "", step)
    step = re.sub(r"^[A-Za-z]\)\s*", "", step)
    step = re.sub(r"^[-]\s*", "", step)

    # remove "Step X:" prefix
    step = re.sub(r"^step\s*\d+\s*[:\-]\s*", "", step, flags=re.IGNORECASE)

    # remove extra spaces
    step = re.sub(r"\s{2,}", " ", step).strip()

    # fix broken hyphen words
    step = step.replace("pre- processing", "preprocessing")
    step = step.replace("pre - processing", "preprocessing")

    # cut too long step
    if len(step) > 260:
        step = step[:260].rsplit(" ", 1)[0] + "..."

    return step.strip()


# ==========================================================
# VALIDATE STEP QUALITY
# ==========================================================
def is_valid_step(step: str) -> bool:
    if not step:
        return False

    step = step.strip()

    if len(step) < 12:
        return False

    # reject pure symbols/numbers
    if re.fullmatch(r"[0-9\.\-\s]+", step):
        return False

    low = step.lower()

    # reject typical noise
    bad_phrases = [
        "authorized licensed use",
        "downloaded on",
        "ieee xplore",
        "all rights reserved",
        "doi:",
        "references",
        "bibliography",
    ]

    if any(bp in low for bp in bad_phrases):
        return False

    # reject figure/table headers
    if re.match(r"^(table|fig|figure)\s+[0-9ivx]+", low):
        return False

    # reject if too many words (paragraph-like)
    if len(step.split()) > 40:
        return False

    return True


# ==========================================================
# EXTRACT BULLET / NUMBERED STEPS (HIGH PRIORITY)
# ==========================================================
def extract_numbered_or_bullet_steps(text: str, max_steps=10):
    """
    Detect common patterns:
    - 1. ...
    - 1) ...
    - (1) ...
    - Step 1: ...
    - • ...
    - - ...
    """

    if not text:
        return []

    text = normalize_text(text)
    text = remove_ieee_noise(text)

    lines = [l.strip() for l in text.split("\n") if l.strip()]

    bullet_patterns = [
        r"^\d+\.\s+",
        r"^\d+\)\s+",
        r"^\(\d+\)\s+",
        r"^[A-Za-z]\)\s+",
        r"^[-•▪◦]\s+",
        r"^step\s*\d+\s*[:\-]\s+",
    ]

    bullet_regex = re.compile("|".join(bullet_patterns), re.IGNORECASE)

    steps = []

    for ln in lines:
        if bullet_regex.search(ln):
            step = clean_step_text(ln)
            if is_valid_step(step):
                steps.append(step)

        if len(steps) >= max_steps:
            break

    return steps[:max_steps]


# ==========================================================
# SENTENCE SPLITTER (ROBUST)
# ==========================================================
def split_into_sentences(text: str):
    if not text:
        return []

    text = normalize_text(text)

    # split by period + newline + semicolon
    sentences = re.split(r"(?<=[\.\!\?])\s+|\n+", text)

    cleaned = []
    for s in sentences:
        s = s.strip()
        if len(s) >= 10:
            cleaned.append(s)

    return cleaned


# ==========================================================
# RULE-BASED KEYWORD METHOD STEPS (SV SAFE RULES)
# ==========================================================
def extract_keyword_based_steps(text: str, max_steps=10):
    """
    Extract methodology-like sentences based on strong action keywords.
    This works when paper has no explicit bullets.
    """

    if not text:
        return []

    text = normalize_text(text)
    text = remove_ieee_noise(text)

    sentences = split_into_sentences(text)

    # Strong methodology action indicators
    keywords = [
        "we propose", "we present", "we develop", "we design", "we introduce",
        "we implement", "we apply", "we applied", "we use", "we utilized",
        "we employ", "we evaluate", "we test", "we trained", "we train",
        "we validate", "we compare", "we compute", "we extract",
        "dataset", "data set", "preprocess", "pre-processing",
        "feature extraction", "segmentation", "classification",
        "training", "testing", "evaluation", "hyperparameter",
        "cross validation", "k-fold", "architecture", "pipeline",
        "framework", "optimizer", "loss function", "learning rate",
        "confusion matrix", "accuracy", "precision", "recall", "f1-score",
        "rouge", "bertscore"
    ]

    steps = []

    for s in sentences:
        low = s.lower()

        if any(k in low for k in keywords):
            step = clean_step_text(s)
            if is_valid_step(step):
                steps.append(step)

        if len(steps) >= max_steps:
            break

    return steps[:max_steps]


# ==========================================================
# OPTIONAL SPACY EXTRACTION (ACTION VERB FILTERING)
# ==========================================================
def extract_steps_spacy(text: str, max_steps=10):
    """
    Use spaCy if installed to detect verb-driven action sentences.
    """

    if not text:
        return []

    try:
        import spacy
    except Exception:
        return []

    try:
        nlp = spacy.load("en_core_web_sm")
    except Exception:
        return []

    text = normalize_text(text)
    text = remove_ieee_noise(text)

    doc = nlp(text)

    action_verbs = {
        "use", "apply", "employ", "train", "test", "evaluate",
        "extract", "clean", "preprocess", "classify", "detect",
        "generate", "summarize", "measure", "compare", "validate",
        "propose", "design", "implement"
    }

    steps = []

    for sent in doc.sents:
        s = sent.text.strip()

        if len(s) < 15:
            continue

        has_action = False
        for token in sent:
            if token.pos_ == "VERB" and token.lemma_.lower() in action_verbs:
                has_action = True
                break

        if has_action:
            step = clean_step_text(s)
            if is_valid_step(step):
                steps.append(step)

        if len(steps) >= max_steps:
            break

    return steps[:max_steps]


# ==========================================================
# OPTIONAL FLAN-T5 REFINEMENT (REFORMAT, NOT GENERATE NEW IDEA)
# ==========================================================
def refine_steps_with_flan_t5(method_text: str, max_steps=8):
    """
    Flan-T5 is used ONLY for reformatting the extracted methodology
    into short steps for poster flowchart.
    """

    if not method_text:
        return []

    try:
        from transformers import pipeline
    except Exception:
        return []

    try:
        model = pipeline(
            "text2text-generation",
            model="google/flan-t5-base",
            tokenizer="google/flan-t5-base"
        )
    except Exception:
        return []

    method_text = normalize_text(method_text)
    method_text = remove_ieee_noise(method_text)

    # truncate for safety
    if len(method_text) > 3500:
        method_text = method_text[:3500]

    prompt = (
        f"Extract the research methodology as {max_steps} short steps for a poster flowchart. "
        f"Each step must be a short sentence. "
        f"Do not add new information.\n\n"
        f"Methodology:\n{method_text}"
    )

    try:
        out = model(
            prompt,
            max_new_tokens=220,
            do_sample=False
        )

        if not out or not isinstance(out, list):
            return []

        generated = out[0].get("generated_text", "").strip()
        if not generated:
            return []

        # split by numbering OR line breaks
        parts = re.split(r"\n+|\d+\.\s*", generated)

        steps = []
        for p in parts:
            p = p.strip()
            if not p:
                continue

            step = clean_step_text(p)
            if is_valid_step(step):
                steps.append(step)

            if len(steps) >= max_steps:
                break

        return steps[:max_steps]

    except Exception:
        return []


# ==========================================================
# FALLBACK SAFE STEPS (ENSURE NEVER EMPTY)
# ==========================================================
def fallback_generic_steps(max_steps=6):
    base = [
        "Collect and prepare the dataset from the paper.",
        "Perform preprocessing and cleaning of the data.",
        "Apply the proposed model or framework for analysis.",
        "Train the model using selected hyperparameters.",
        "Evaluate performance using standard evaluation metrics.",
        "Compare results against baseline or existing methods."
    ]
    return base[:max_steps]


# ==========================================================
# MAIN FUNCTION (FINAL PIPELINE)
# ==========================================================
def extract_methodology_steps(
    methodology_text: str,
    max_steps=8,
    use_spacy=True,
    use_flan_refine=True
):
    """
    Final SV hybrid pipeline:
    1) Extract bullet/numbered steps (most accurate)
    2) Extract keyword-based action sentences
    3) Extract spaCy action sentences (optional)
    4) Flan-T5 refinement (optional) to clean + restructure
    5) Merge + remove duplicates
    6) If still weak -> fallback generic steps
    """

    methodology_text = normalize_text(methodology_text)
    methodology_text = remove_ieee_noise(methodology_text)

    if not methodology_text:
        return fallback_generic_steps(max_steps=max_steps)

    steps = []

    # -------------------------------------------------------
    # 1) Bullet / numbered detection
    # -------------------------------------------------------
    bullet_steps = extract_numbered_or_bullet_steps(methodology_text, max_steps=max_steps)
    steps.extend(bullet_steps)

    # -------------------------------------------------------
    # 2) Keyword-based extraction
    # -------------------------------------------------------
    if len(steps) < max_steps:
        keyword_steps = extract_keyword_based_steps(methodology_text, max_steps=max_steps)
        steps.extend(keyword_steps)

    # -------------------------------------------------------
    # 3) spaCy extraction
    # -------------------------------------------------------
    if use_spacy and len(steps) < max_steps:
        spacy_steps = extract_steps_spacy(methodology_text, max_steps=max_steps)
        steps.extend(spacy_steps)

    # -------------------------------------------------------
    # 4) Remove duplicates early
    # -------------------------------------------------------
    deduped = []
    seen = set()

    for s in steps:
        s_clean = clean_step_text(s)
        key = re.sub(r"[^a-z0-9 ]", "", s_clean.lower()).strip()

        if not key or key in seen:
            continue

        if is_valid_step(s_clean):
            deduped.append(s_clean)
            seen.add(key)

        if len(deduped) >= max_steps:
            break

    steps = deduped[:max_steps]

    # -------------------------------------------------------
    # 5) Flan refinement (only if steps still weak)
    # -------------------------------------------------------
    if use_flan_refine:
        flan_steps = refine_steps_with_flan_t5(methodology_text, max_steps=max_steps)

        # accept flan only if looks strong
        if len(flan_steps) >= 4:
            steps = flan_steps

    # -------------------------------------------------------
    # 6) Final fallback if too weak
    # -------------------------------------------------------
    if len(steps) < 3:
        return fallback_generic_steps(max_steps=max_steps)

    return steps[:max_steps]


# ==========================================================
# QUICK TEST
# ==========================================================
if __name__ == "__main__":
    sample_method = """
    III. METHODOLOGY
    We collected datasets from Kaggle.
    1. Data preprocessing and cleaning.
    2. Training Flan-T5 base model.
    3. Evaluation using ROUGE-L and BERTScore.
    • Comparison against baseline model.
    """

    steps = extract_methodology_steps(sample_method, max_steps=8)

    print("Extracted Methodology Steps:")
    for i, s in enumerate(steps, 1):
        print(f"{i}. {s}")