import sys, subprocess, getpass, os
from openai import OpenAI

# Upgrade SDK
subprocess.run([sys.executable, "-m", "pip", "install", "-q", "--upgrade", "openai>=1.40.0"])

# Get API key
OPENAI_API_KEY = getpass.getpass("Paste your OpenAI API key (starts with 'sk-'): ")
assert OPENAI_API_KEY.startswith("sk-"), "Key must start with 'sk-'"

# Initialize client
client = OpenAI(api_key=OPENAI_API_KEY)
print("✅ OpenAI client ready.")

resp = client.responses.create(
    model="gpt-4o",
    input=[{"role": "user", "content": "Say ok"}],
    temperature=0
)
print(resp.output_text)

rows = [
    # --- originals (4) ---
    ("Borrowers with FICO >= 700 qualify for best pricing.", "Pricing"),
    ("Applications with DTI above 45% are not eligible.", "Eligibility"),
    ("Accounts with two missed payments move to collections.", "Collections"),
    ("Flag transactions with mismatched addresses for review.", "Fraud"),

    # --- Pricing vs Eligibility confusers ---
    ("Prime tier applies to FICO 720+ with a 25 bps rate reduction.", "Pricing"),
    ("Minimum credit score 680 required to apply.", "Eligibility"),
    ("Applicants with FICO 740 are offered a points discount.", "Pricing"),
    ("DTI must be <= 36% for prime offers.", "Eligibility"),

    # --- Collections vs Fraud boundary ---
    ("If payment is 60 days late, send a delinquency notice.", "Collections"),
    ("Unrecognized device + inconsistent SSN → escalate for identity verification.", "Fraud"),

    # --- Mixed sentences (contain multiple signals) ---
    ("DTI must be under 40%; rate is 7.1% for approved borrowers.", "Eligibility"),
    ("Late fees apply after 15 days; repeated chargebacks may indicate account takeover.", "Fraud"),

    # --- Negations & hedges ---
    ("No change to APR tiers this quarter.", "Pricing"),
    ("This is not suspected fraud; route to normal collections.", "Collections"),

    # --- Colloquial / indirect phrasing ---
    ("Score in the high 600s? You’re probably in the better-price bucket.", "Pricing"),
    ("When debt eats half your paycheck, we decline the app.", "Eligibility"),

    # --- Lookalikes that should separate ---
    ("Two consecutive NSF events move the account to collections workflow.", "Collections"),
    ("Synthetic name patterns (e.g., repeated vowels) should be auto-flagged.", "Fraud"),

    # --- Edge constraints: LTV/DTI mixes ---
    ("Max LTV is 80% and DTI must be <= 43% for jumbo loans.", "Eligibility"),
    ("Preferred tier gets 0.25% rate cut with autopay enrollment.", "Pricing"),
]

df_ext = pd.DataFrame(rows, columns=["text","label"])
Path("data").mkdir(exist_ok=True, parents=True)
df_ext.to_csv("data/credit_notes_extended.csv", index=False, encoding="utf-8")
len(df_ext), df_ext.head(8)

from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.metrics.pairwise import cosine_similarity

snips = {
 "001": "DTI cap is typically 36% for prime borrowers.",
 "002": "Minimum credit score 680+ is required for standard eligibility.",
 "003": "Collections process begins after 60 days of nonpayment.",
 "004": "Pricing discounts of 25 bps apply for autopay enrollment."
}

ids, corpus = list(snips.keys()), list(snips.values())
vec = TfidfVectorizer().fit(corpus)
X = vec.transform(corpus)

def retrieve(query, k=1):
    q = vec.transform([query])
    sims = cosine_similarity(q, X).ravel()
    idx = sims.argsort()[::-1][:k]
    return [(ids[i], corpus[i], float(sims[i])) for i in idx]

def answer_with_grounding(question):
    hits = retrieve(question, k=1)
    doc_id, passage, score = hits[0]
    prompt = f"""Answer the question **only** using the passage below.
If the passage is insufficient, say "insufficient".
Return JSON: answer

Passage (doc_id={doc_id}):
{passage}

Question: {question}
"""
    r = client.responses.create(model="gpt-4o", input=[{"role":"user","content":prompt}], temperature=0)
    import json, re
    try:
        obj = json.loads(r.output_text)
    except Exception:
        m = re.search(r"\{[\s\S]*\}", r.output_text)
        obj = json.loads(m.group(0)) if m else {"answer": r.output_text.strip(), "doc_id": doc_id}
    obj["retrieval_score"] = round(score,3)
    return obj

print(answer_with_grounding("What is the typical DTI cap?"))
print(answer_with_grounding("When do collections begin?"))