generation.py

"""Grounded recommendation explanations for REmatch Part 4."""

from __future__ import annotations

import re
from typing import Any

import pandas as pd
import spaces
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

MODEL_ID = "Qwen/Qwen2.5-1.5B-Instruct"

SEVERE_UNSUPPORTED_CLAIMS = (
    "guaranteed return",
    "guaranteed profit",
    "risk-free investment",
    "cannot lose money",
    "will definitely appreciate",
    "guaranteed appreciation",
    "guaranteed rental income",
)

FORBIDDEN_ALIGNMENT_TERMS = (
    "yield",
    "rent",
    "rental",
    "forecast",
    "volatility",
    "liquidity",
    "financing",
    "property age",
    "construction year",
    "days on market",
    "recent listing",
    "market demand",
    "market conditions",
    "market environment",
    "appreciation",
    "profit",
    "profitability",
    "cash flow",
    "model score",
    "compatibility score",
    "risk",
)

ALIGNMENT_RELATIONS = (
    "align with",
    "fit",
    "suitable for",
    "relevant to",
    "suit",
    "match",
)

YIELD_PHRASES = {
    "high_yield": "high rental yield",
    "medium_yield": "medium rental yield",
    "low_yield": "low rental yield",
}

FORECAST_PHRASES = {
    "positive_forecast": "positive twelve-month forecast",
    "stable_forecast": "stable twelve-month forecast",
    "negative_forecast": "negative twelve-month forecast",
}

VOLATILITY_PHRASES = {
    "low_volatility": "low price volatility",
    "medium_volatility": "medium price volatility",
    "high_volatility": "high price volatility",
}

LIQUIDITY_PHRASES = {
    "high_liquidity": "high property liquidity",
    "medium_liquidity": "medium property liquidity",
    "balanced_liquidity": "balanced property liquidity",
    "low_liquidity": "low property liquidity",
}

VERIFICATION_ITEMS = (
    "property condition",
    "operating expenses",
    "financing terms",
    "rental assumptions",
)


def _is_missing(value: Any) -> bool:
    if value is None:
        return True

    try:
        return bool(pd.isna(value))
    except (TypeError, ValueError):
        return False


def _clean_text(value: Any, limit: int = 300) -> str:
    if _is_missing(value):
        return "Not available"

    return " ".join(str(value).split())[:limit]

def _naturalize_label(value: Any) -> str:
    """Convert machine labels into natural-language labels."""

    return _clean_text(value).replace("_", " ")

def _property_reference(
    property_record: dict[str, Any],
) -> str:
    for key in (
        "formattedAddress",
        "city_rentcast",
        "city",
        "propertyType",
        "id",
    ):
        value = property_record.get(key)

        if not _is_missing(value):
            return _clean_text(value)

    return "This property"


def _mapped_phrase(
    property_record: dict[str, Any],
    field: str,
    mapping: dict[str, str],
    fallback: str,
) -> str:
    value = _clean_text(property_record.get(field))
    return mapping.get(value, fallback)


def build_evidence_packet(
    investor_profile: dict[str, Any],
    property_record: dict[str, Any],
) -> dict[str, Any]:
    """Select all approved evidence before language generation."""

    goal = _naturalize_label(
        investor_profile.get("primary_goal")
    )
    risk_profile = _naturalize_label(
        investor_profile.get("risk_profile")
    )

    yield_evidence = _mapped_phrase(
        property_record,
        "yield_band",
        YIELD_PHRASES,
        "available rental-yield signal",
    )
    forecast_evidence = _mapped_phrase(
        property_record,
        "forecast_band",
        FORECAST_PHRASES,
        "available twelve-month forecast signal",
    )
    volatility_evidence = _mapped_phrase(
        property_record,
        "volatility_band",
        VOLATILITY_PHRASES,
        "available price-volatility signal",
    )
    liquidity_evidence = _mapped_phrase(
        property_record,
        "liquidity_band",
        LIQUIDITY_PHRASES,
        "available property-liquidity signal",
    )

    if goal == "income":
        goal_evidence = yield_evidence
    elif goal == "growth":
        goal_evidence = forecast_evidence
    elif goal == "preservation":
        goal_evidence = (
            f"{volatility_evidence} and "
            f"{forecast_evidence}"
        )
    else:
        goal_evidence = (
            f"{yield_evidence} and "
            f"{forecast_evidence}"
        )

    volatility_band = _clean_text(
        property_record.get("volatility_band")
    )
    forecast_band = _clean_text(
        property_record.get("forecast_band")
    )
    liquidity_band = _clean_text(
        property_record.get("liquidity_band")
    )
    yield_band = _clean_text(
        property_record.get("yield_band")
    )

    if volatility_band == "high_volatility":
        limitation_sentence = (
            "High price volatility is an important "
            "consideration for this recommendation."
        )
    elif forecast_band == "negative_forecast":
        limitation_sentence = (
            "The negative twelve-month forecast is an "
            "important consideration for this recommendation."
        )
    elif liquidity_band == "low_liquidity":
        limitation_sentence = (
            "Low property liquidity is an important "
            "consideration for this recommendation."
        )
    elif goal == "income" and yield_band == "low_yield":
        limitation_sentence = (
            "Low rental yield is an important consideration "
            "for an income-focused investor."
        )
    else:
        limitation_sentence = (
            "Forecast and rental estimates remain uncertain "
            "and should be verified."
        )

    property_reference = _property_reference(
        property_record
    )

    article = (
        "an"
        if risk_profile[:1].lower() in "aeiou"
        else "a"
    )

    alignment_source_sentence = (
        f"This property may align with {article} "
        f"{risk_profile} investor focused on {goal}."
    )

    return {
        "property_reference": property_reference,
        "goal": goal,
        "risk_profile": risk_profile,
        "goal_evidence": goal_evidence,
        "volatility_evidence": volatility_evidence,
        "forecast_evidence": forecast_evidence,
        "liquidity_evidence": liquidity_evidence,
        "alignment_source_sentence": (
            alignment_source_sentence
        ),
        "limitation_sentence": limitation_sentence,
        "verification_items": list(
            VERIFICATION_ITEMS
        ),
    }

def build_prompt(
    evidence: dict[str, Any],
) -> str:
    """Ask Qwen for one controlled but natural alignment sentence."""

    return f"""
Generate one concise and natural English sentence describing only the
possible alignment between the selected property and the investor.
Required anchors:
- Refer to the selected listing as "This property".
- Include the exact phrase "{evidence['risk_profile']} investor".
- Express the goal exactly in one of these two forms:
  "focused on {evidence['goal']}"
  or "who focuses on {evidence['goal']}".
- Do not replace, expand, or qualify the goal with other words.
- The goal expression must end the sentence.
Requirements:
- Use a neutral alignment relation such as "align with", "fit for",
  "suitable for", "relevant to", "suit", or "match".
- Use cautious wording: "may", "might", or "could".
- State only possible alignment between the property and investor.
- Do not discuss yield, rent, forecast, volatility, liquidity, market
  conditions, appreciation, profitability, financing, property age,
  days on market, model scores, or any other property benefit.
- Do not add facts, financial numbers, reasons, guarantees, or buying advice.
- Return only one sentence.
- Do not return a heading, JSON, list, or markdown.
""".strip()

def build_retry_prompt(
    evidence: dict[str, Any],
) -> str:
    """Use a fully constrained prompt for the single retry."""

    return f"""
Return exactly this sentence and nothing else:
{evidence['alignment_source_sentence']}
""".strip()


def _financial_tokens(text: str) -> set[str]:
    return set(
        re.findall(
            r"\$[\d,]+(?:\.\d+)?|\b\d+(?:\.\d+)?%",
            text,
        )
    )


def _clean_generated_sentence(raw_text: str) -> str:
    """Extract exactly one clean sentence."""

    text = " ".join(raw_text.strip().split())
    text = re.sub(
        r"^(answer|response|explanation|rewrite)\s*:\s*",
        "",
        text,
        flags=re.IGNORECASE,
    )
    text = text.strip("`\"' ")

    if not text:
        raise ValueError(
            "The model returned an empty response."
        )

    sentences = [
        part.strip()
        for part in re.split(
            r"(?<=[.!?])\s+",
            text,
        )
        if part.strip()
    ]

    if len(sentences) != 1:
        raise ValueError(
            "The model must return exactly one sentence."
        )

    sentence = sentences[0]

    if sentence[-1] not in ".!?":
        sentence += "."

    return sentence


def _deterministic_evidence(
    evidence: dict[str, Any],
) -> str:
    return (
        f"Goal indicator: {evidence['goal_evidence']}. "
        f"Risk indicators: {evidence['volatility_evidence']} "
        f"and {evidence['forecast_evidence']}. "
        f"Liquidity indicator: "
        f"{evidence['liquidity_evidence']}."
    )


def _deterministic_verification(
    evidence: dict[str, Any],
) -> str:
    return (
        "Verify "
        + ", ".join(evidence["verification_items"])
        + " before making an investment decision."
    )


def _render_explanation(
    generated_sentence: str,
    evidence: dict[str, Any],
) -> str:
    return (
        f"Why this property may fit:\n"
        f"Property: {evidence['property_reference']}\n"
        f"{generated_sentence}\n\n"
        f"Key consideration:\n"
        f"{evidence['limitation_sentence']}\n\n"
        f"Key property indicators:\n"
        f"{_deterministic_evidence(evidence)}\n\n"
        f"What to verify:\n"
        f"{_deterministic_verification(evidence)}\n\n"
        "This explanation is based only on the available "
        "data and is not financial advice."
    )


def _normalize_match_text(text: str) -> str:
    """Normalize text for deterministic anchor matching."""

    return " ".join(
        re.findall(
            r"[a-z0-9]+",
            text.lower(),
        )
    )

def _contains_normalized_phrase(
    text: str,
    phrase: str,
) -> bool:
    normalized_text = (
        f" {_normalize_match_text(text)} "
    )
    normalized_phrase = (
        f" {_normalize_match_text(phrase)} "
    )

    return normalized_phrase in normalized_text


def _normalized_phrase_count(
    text: str,
    phrase: str,
) -> int:
    normalized_text = (
        f" {_normalize_match_text(text)} "
    )
    normalized_phrase = (
        f" {_normalize_match_text(phrase)} "
    )

    return normalized_text.count(
        normalized_phrase
    )


def _uses_alignment_relation(
    sentence: str,
) -> bool:
    sentence_normalized = (
        f" {_normalize_match_text(sentence)} "
    )

    return any(
        (
            f" {_normalize_match_text(relation)} "
            in sentence_normalized
        )
        for relation in ALIGNMENT_RELATIONS
    )


def _uses_cautious_language(
    sentence: str,
) -> bool:
    return bool(
        re.search(
            r"\b(?:may|might|could)\b",
            sentence.lower(),
        )
    )


def _contains_forbidden_alignment_content(
    sentence: str,
) -> bool:
    sentence_normalized = (
        f" {_normalize_match_text(sentence)} "
    )

    return any(
        (
            f" {_normalize_match_text(term)} "
            in sentence_normalized
        )
        for term in FORBIDDEN_ALIGNMENT_TERMS
    )


def validate_explanation(
    explanation: str,
    generated_sentence: str,
    prompt: str,
    evidence: dict[str, Any],
) -> dict[str, Any]:
    """Validate controlled generation and deterministic grounding."""

    normalized = " ".join(explanation.split())
    lower = normalized.lower()
    
    sentence_normalized = (
        _normalize_match_text(
            generated_sentence
        )
    )

    sentence_word_count = len(
        generated_sentence.split()
    )
    word_count = len(normalized.split())

    unsupported_financial_tokens = sorted(
        _financial_tokens(generated_sentence)
        - _financial_tokens(prompt)
    )

    checks = {
        "non_empty": bool(
            generated_sentence.strip()
        ),
        "generated_sentence_length": (
            7 <= sentence_word_count <= 40
        ),
        "final_word_count_in_range": (
            45 <= word_count <= 180
        ),
        "required_sections_present": all(
            section in lower
            for section in (
                "why this property may fit:",
                "key consideration:",
                "key property indicators:",
                "what to verify:",
                "not financial advice",
            )
        ),
        
        "property_referenced": (
            evidence["property_reference"].lower()
            in lower
        ),
        
        "property_anchor_used": (
            sentence_normalized.startswith(
                "this property "
            )
        ),
        "risk_profile_referenced": (
            _contains_normalized_phrase(
                generated_sentence,
                f"{evidence['risk_profile']} investor",
            )
        ),
        "goal_referenced": any(
            sentence_normalized.endswith(
                _normalize_match_text(goal_phrase)
            )
            for goal_phrase in (
                f"focused on {evidence['goal']}",
                f"who focuses on {evidence['goal']}",
            )
        ),
        "goal_used_once": (
            _normalized_phrase_count(
                generated_sentence,
                evidence["goal"],
            )
            == 1
        ),
        "alignment_relation_used": (
            _uses_alignment_relation(
                generated_sentence
            )
        ),
        "cautious_language_used": (
            _uses_cautious_language(
                generated_sentence
            )
        ),
        "no_severe_unsupported_claims": not any(
            claim in generated_sentence.lower()
            for claim in SEVERE_UNSUPPORTED_CLAIMS
        ),
        "no_forbidden_alignment_content": not (
            _contains_forbidden_alignment_content(
                generated_sentence
            )
        ),
        "no_unsupported_financial_numbers": (
            not unsupported_financial_tokens
        ),
        "limitation_preserved": (
            evidence["limitation_sentence"].lower()
            in lower
        ),
        "goal_evidence_preserved": (
            evidence["goal_evidence"].lower()
            in lower
        ),
        "volatility_evidence_preserved": (
            evidence["volatility_evidence"].lower()
            in lower
        ),
        "forecast_evidence_preserved": (
            evidence["forecast_evidence"].lower()
            in lower
        ),
        "liquidity_evidence_preserved": (
            evidence["liquidity_evidence"].lower()
            in lower
        ),
        "verification_items_preserved": all(
            item.lower() in lower
            for item in evidence["verification_items"]
        ),
    }

    return {
        "passed": all(checks.values()),
        "checks": checks,
        "word_count": word_count,
        "unsupported_financial_tokens": (
            unsupported_financial_tokens
        ),
    }

# ============================================================
# QWEN MODEL LOADING FOR HUGGING FACE ZEROGPU
# ============================================================

# In a ZeroGPU Space, CUDA is emulated during startup and becomes a real GPU
# inside functions marked with @spaces.GPU. Keep this placement at module level.
DEVICE = "cuda"
MODEL_DTYPE = torch.bfloat16

_TOKENIZER = AutoTokenizer.from_pretrained(
    MODEL_ID
)

_MODEL = AutoModelForCausalLM.from_pretrained(
    MODEL_ID,
    torch_dtype=MODEL_DTYPE,
    low_cpu_mem_usage=True,
)

_MODEL = _MODEL.to(DEVICE)
_MODEL.eval()

if _TOKENIZER.pad_token_id is None:
    _TOKENIZER.pad_token_id = (
        _TOKENIZER.eos_token_id
    )


@spaces.GPU(duration=60)
def _run_generation(prompt: str) -> str:
    messages = [
        {
            "role": "system",
            "content": (
                "Generate one short, cautious alignment sentence. "
                "Refer to the selected listing as 'This property'. "
                "Preserve the provided risk phrase and goal phrase "
                "exactly, and never add property benefits or "
                "financial facts."
            ),
        },
        {
            "role": "user",
            "content": prompt,
        },
    ]

    rendered = _TOKENIZER.apply_chat_template(
        messages,
        tokenize=False,
        add_generation_prompt=True,
    )
    inputs = _TOKENIZER(
        rendered,
        return_tensors="pt",
    ).to(DEVICE)

    with torch.inference_mode():
        output = _MODEL.generate(
            **inputs,
            max_new_tokens=60,
            do_sample=False,
            repetition_penalty=1.05,
            pad_token_id=_TOKENIZER.eos_token_id,
        )

    generated_tokens = output[
        0,
        inputs["input_ids"].shape[1]:,
    ]

    return _TOKENIZER.decode(
        generated_tokens,
        skip_special_tokens=True,
    ).strip()

def _fallback_explanation(
    evidence: dict[str, Any],
) -> str:
    return _render_explanation(
        evidence["alignment_source_sentence"],
        evidence,
    )

def generate_property_explanation(
    investor_profile: dict[str, Any],
    property_record: dict[str, Any],
) -> str:
    """Generate one grounded explanation with one controlled retry."""

    evidence = build_evidence_packet(
        investor_profile,
        property_record,
    )

    prompts = (
        build_prompt(evidence),
        build_retry_prompt(evidence),
    )

    for current_prompt in prompts:
        try:
            raw_output = _run_generation(
                current_prompt
            )

            generated_sentence = (
                _clean_generated_sentence(
                    raw_output
                )
            )

            explanation = _render_explanation(
                generated_sentence,
                evidence,
            )

            qa = validate_explanation(
                explanation=explanation,
                generated_sentence=generated_sentence,
                prompt=current_prompt,
                evidence=evidence,
            )

            if qa["passed"]:
                return explanation

        except Exception:
            continue

    return _fallback_explanation(
        evidence
    )