"""Reproduce training from the source Excel. Run stages in a machine with enough time/RAM.
Produces: vertical_model.joblib, subvertical_model.joblib, parent_map.json
Vertical trains on real verticals only (excludes 'others'); sub-vertical is downsampled per
class for speed. Feature vectorizers are shared across both stages."""
import re, json, gc, joblib
import numpy as np, pandas as pd
from scipy.sparse import hstack
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.linear_model import SGDClassifier
from sklearn.svm import LinearSVC
from sklearn.model_selection import train_test_split

SRC = "Raw_data_vertical_subvertical.xlsx"
clean = lambda s: " ".join(str(s).lower().split())

df = pd.read_excel(SRC, usecols=["category_name","vertical","sub_vertical"], dtype=str).fillna("")
for c in df.columns: df[c] = df[c].str.strip()
df["text"] = df["category_name"].map(clean)

# shared vectorizers, fit on real-vertical text
real = df[df.vertical.str.lower()!="others"]
wv = TfidfVectorizer(analyzer="word", ngram_range=(1,2), min_df=5, max_features=100000,
                     sublinear_tf=True, dtype=np.float32)
cv = TfidfVectorizer(analyzer="char_wb", ngram_range=(2,4), min_df=5, max_features=100000,
                     sublinear_tf=True, dtype=np.float32)
wv.fit(real.text); cv.fit(real.text)
vec = lambda s: hstack([wv.transform(s), cv.transform(s)]).tocsr()

# ---- vertical (21-class, probabilistic) ----
vclf = SGDClassifier(loss="log_loss", alpha=2e-6, max_iter=30, tol=1e-4, random_state=42)
vclf.fit(vec(real.text), real.vertical.to_numpy())
joblib.dump({"word_vec":wv,"char_vec":cv,"clf":vclf,"classes":list(vclf.classes_)},
            "vertical_model.joblib", compress=3)

# ---- sub-vertical (LinearSVC, downsampled per class) ----
sub = df[(df.vertical.str.lower()!="others") & (df.sub_vertical.str.lower()!="others")]
parent = sub.groupby("sub_vertical")["vertical"].agg(lambda s: s.value_counts().index[0]).to_dict()
json.dump(parent, open("parent_map.json","w"))
CAP=2500; rng=np.random.default_rng(42); idx=[]
for c, g in sub.groupby("sub_vertical"):
    ix=g.index.to_numpy()
    idx.append(rng.choice(ix,CAP,replace=False) if len(ix)>CAP else ix)
idx=np.concatenate(idx)
s=sub.loc[idx]
sclf=LinearSVC(C=1.0)
sclf.fit(vec(s.text), s.sub_vertical.to_numpy())
joblib.dump({"clf":sclf,"classes":list(sclf.classes_)}, "subvertical_model.joblib", compress=3)
print("done: vertical_model.joblib, subvertical_model.joblib, parent_map.json")