import numpy as np, pandas as pd
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.linear_model import LogisticRegression
from sklearn.model_selection import train_test_split, cross_val_score, StratifiedKFold
from sklearn.metrics import classification_report, confusion_matrix
from sklearn.dummy import DummyClassifier

rng = np.random.default_rng(7)
# Heavy shared vocabulary: "access", "community", "cost", "program" appear in BOTH senses.
shared  = ["students","policy","education","study","data","analysis","results","school",
           "program","evaluation","access","community","cost","outcomes","district","state"]
equity  = ["equity","barriers","inclusion","marginalized","disparity","underserved"]
effic   = ["optimization","efficiency","scalable","throughput","productivity","allocation"]

def make_doc(rng, topical, p_topical):
    n = rng.integers(22, 34)
    words = []
    for _ in range(n):
        if rng.random() < p_topical:
            words.append(rng.choice(topical))
        else:
            words.append(rng.choice(shared))
    return " ".join(words)

n_eq, n_ef = 80, 420           # 16% minority: the marginal-cases setting
docs, labels = [], []
for _ in range(n_eq):
    docs.append(make_doc(rng, equity, 0.10)); labels.append("equity")
for _ in range(n_ef):
    docs.append(make_doc(rng, effic, 0.10)); labels.append("efficiency")

df = pd.DataFrame({"doc_id":[f"d{i:03d}" for i in range(len(docs))],"text":docs,"label":labels})
# label noise: 8% of documents were coded inconsistently upstream
flip = rng.choice(len(df), size=int(0.08*len(df)), replace=False)
df.loc[flip,"label"] = df.loc[flip,"label"].map({"equity":"efficiency","efficiency":"equity"})
df = df.sample(frac=1, random_state=1).reset_index(drop=True)

X_train, X_test, y_train, y_test = train_test_split(
    df["text"], df["label"], test_size=0.3, stratify=df["label"], random_state=42)
vec = TfidfVectorizer(min_df=2)
Xtr, Xte = vec.fit_transform(X_train), vec.transform(X_test)
clf = LogisticRegression(max_iter=1000, random_state=42).fit(Xtr, y_train)
pred = clf.predict(Xte)

print("shapes:", Xtr.shape, Xte.shape, "| vocab:", len(vec.get_feature_names_out()))
print("class balance (full):", df["label"].value_counts().to_dict())
print("\nCONFUSION rows=actual[efficiency,equity] cols=predicted")
print(confusion_matrix(y_test, pred, labels=["efficiency","equity"]))
print("\nREPORT"); print(classification_report(y_test, pred, digits=3, zero_division=0))
dum = DummyClassifier(strategy="most_frequent").fit(Xtr,y_train)
print("majority-baseline accuracy:", round((dum.predict(Xte)==y_test).mean(),3))
cv = cross_val_score(LogisticRegression(max_iter=1000,random_state=42),
                     vec.fit_transform(df["text"]), df["label"],
                     cv=StratifiedKFold(5,shuffle=True,random_state=42), scoring="f1_macro")
print("5-fold macro F1:", np.round(cv,3), "mean", round(cv.mean(),3), "sd", round(cv.std(),3))

# balanced version
clf2 = LogisticRegression(max_iter=1000, class_weight="balanced", random_state=42).fit(Xtr,y_train)
p2 = clf2.predict(Xte)
print("\nWITH class_weight='balanced'")
print(confusion_matrix(y_test,p2,labels=["efficiency","equity"]))
print(classification_report(y_test,p2,digits=3,zero_division=0))
feat = pd.DataFrame(Xtr[:3].toarray(), columns=vec.get_feature_names_out())
print("FEATURE PEEK\n", feat[[c for c in ["equity","access","optimization","students"] if c in feat.columns]].round(3))