"""
train.py — Wedding Decor Price Pipeline
  • Pulls training data from Supabase DB (or a local CSV as fallback)
  • Generates CLIP embeddings with parallel download + batched GPU/CPU inference
  • Builds stacking ensemble: XGBoost (CPU) + LightGBM + ExtraTrees + MLP → LogisticRegression
  • Saves artifacts to LOCAL_ARTIFACTS dir; app.py is responsible for uploading to Supabase Storage

CPU optimisations (no accuracy loss):
  XGBoost device='cpu', tree_method='hist'
  CLIP freed from RAM before stacking trains
  All transformers persisted in a single transforms.joblib dict
"""

import os, io, json, re, warnings, argparse
from pathlib import Path
from concurrent.futures import ThreadPoolExecutor

import numpy as np
import pandas as pd
import requests
from PIL import Image
import joblib

import torch
import open_clip

from sklearn.preprocessing import LabelEncoder, StandardScaler, OneHotEncoder, MinMaxScaler
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.decomposition import PCA
from sklearn.model_selection import train_test_split
from sklearn.ensemble import StackingClassifier, ExtraTreesClassifier
from sklearn.neural_network import MLPClassifier
from sklearn.linear_model import LogisticRegression
from sklearn.metrics import accuracy_score, classification_report

from xgboost import XGBClassifier
from lightgbm import LGBMClassifier

warnings.filterwarnings('ignore', message='.*Falling back to prediction using DMatrix.*')
warnings.filterwarnings('ignore', message='.*X does not have valid feature names.*')

# ── DEFAULTS ──────────────────────────────────────────────────────────────────
DEFAULT_ARTIFACTS = os.getenv("ARTIFACTS_DIR", "/tmp/artifacts")
DEFAULT_CACHE     = os.getenv("EMBEDDINGS_CACHE", "/tmp/embeddings_cache.json")
MAX_ROWS          = int(os.getenv("MAX_ROWS", 10000))
MAX_WORKERS       = int(os.getenv("MAX_WORKERS", 8))
BATCH_SIZE        = int(os.getenv("BATCH_SIZE", 32))
RANDOM_STATE      = 42
PRICE_LOW, PRICE_HIGH = 1000, 500_000
FILTER_TERMS      = ['verified_icon', 'location']
PRICE_BUDGET_MAX  = 15_000
PRICE_MID_MAX     = 80_000

device = "cuda" if torch.cuda.is_available() else "cpu"

# ── HELPERS ───────────────────────────────────────────────────────────────────

def parse_price(text):
    if pd.isna(text) or text == '': return None
    text = str(text).lower().replace(',', '')
    if any(x in text for x in ['request', 'contact', 'onwards', 'call']): return None
    m = re.search(r'\d+', text)
    if m:
        val = float(m.group())
        return val if PRICE_LOW <= val <= PRICE_HIGH else None
    return None


def normalize_price_tier(tag):
    if pd.isna(tag) or str(tag).strip() == "":
        return None

    t = str(tag).strip().lower().replace("_", "-")
    if t in {"budget", "low", "economy"}:
        return "Budget"
    if t in {"mid", "mid-range", "midrange", "standard"}:
        return "Mid-Range"
    if t in {"premium", "high", "luxury"}:
        return "Premium"
    return None


def tier_from_price(price):
    if pd.isna(price):
        return None
    p = float(price)
    if p < PRICE_LOW or p > PRICE_HIGH:
        return None
    if p < PRICE_BUDGET_MAX:
        return "Budget"
    if p < PRICE_MID_MAX:
        return "Mid-Range"
    return "Premium"

def extract_domain(url):
    if not url or pd.isna(url): return 'unknown'
    m = re.search(r'https?://(?:www\.)?([^/]+)', str(url))
    return m.group(1) if m else 'unknown'

# ── STEP 1: LOAD DATA ─────────────────────────────────────────────────────────

def load_data(csv_path: str) -> pd.DataFrame:
    """Load from CSV (which was built from Supabase by app.py or provided directly)."""
    print(f"[train] Loading data from {csv_path}")
    df = pd.read_csv(csv_path)

    # Backward-compatible column normalization.
    if 'Name' not in df.columns:
        title = df.get('title', df.get('name', pd.Series('unknown', index=df.index))).fillna('').astype(str)
        desc = df.get('description', pd.Series('', index=df.index)).fillna('').astype(str)
        df['Name'] = np.where(title.str.strip() != '', title, desc.str.slice(0, 180))

    if 'Description' not in df.columns:
        df['Description'] = df.get('description', pd.Series('', index=df.index)).fillna('').astype(str)

    if 'Seed Price' not in df.columns:
        df['Seed Price'] = df.get('price_text', df.get('seed_price', pd.Series('', index=df.index))).fillna('').astype(str)

    if 'Price INR' not in df.columns:
        df['Price INR'] = df.get('price_inr', pd.Series(np.nan, index=df.index))

    if 'Price Range Tag' not in df.columns:
        df['Price Range Tag'] = df.get('price_range_tag', pd.Series('', index=df.index)).fillna('').astype(str)

    if 'Image URL' not in df.columns:
        df['Image URL'] = df.get('image_url', pd.Series('', index=df.index)).fillna('').astype(str)

    if 'Storage URL' not in df.columns:
        df['Storage URL'] = pd.Series('', index=df.index)

    if 'Source URL' not in df.columns:
        df['Source URL'] = df.get('source_url', pd.Series('', index=df.index)).fillna('').astype(str)

    if 'Source Domain' not in df.columns:
        df['Source Domain'] = df.get('source_domain', pd.Series('unknown', index=df.index)).fillna('unknown').astype(str)

    if 'row_id' not in df.columns:
        source_id = df.get('id', df.index)
        df['row_id'] = source_id.astype(str)

    df['Name'] = df['Name'].fillna('unknown').astype(str)
    name_and_desc = (df['Name'].fillna('') + ' ' + df.get('Description', pd.Series('', index=df.index)).fillna('')).str.lower()
    df = df[~name_and_desc.str.contains('|'.join(FILTER_TERMS), na=False)].copy()

    df['price_inr'] = pd.to_numeric(df['Price INR'], errors='coerce')
    missing_price = df['price_inr'].isna()
    df.loc[missing_price, 'price_inr'] = df.loc[missing_price, 'Seed Price'].apply(parse_price)
    df['price_tier_tagged'] = df['Price Range Tag'].apply(normalize_price_tier)

    # Keep rows that have either direct tier labels or parseable numeric prices.
    df = df[df['price_tier_tagged'].notna() | df['price_inr'].notna()].copy()
    df = df.head(MAX_ROWS).reset_index(drop=True)

    if len(df) == 0:
        raise ValueError("No usable rows after data cleaning. Need tagged tiers or parseable prices.")

    pmin = float(df['price_inr'].dropna().min()) if df['price_inr'].notna().any() else 0.0
    pmax = float(df['price_inr'].dropna().max()) if df['price_inr'].notna().any() else 0.0
    tagged_cnt = int(df['price_tier_tagged'].notna().sum())
    print(f"[train] Clean rows: {len(df)} | tagged: {tagged_cnt} | ₹{pmin:,.0f}–₹{pmax:,.0f}")
    return df

# ── STEP 2: CLIP EMBEDDINGS ───────────────────────────────────────────────────

def load_clip():
    print(f"[train] Loading CLIP on {device}…")
    m, _, p = open_clip.create_model_and_transforms("ViT-B-32", pretrained="laion2b_s34b_b79k")
    return m.to(device).eval(), p

def _download(row):
    headers = {'User-Agent': 'Mozilla/5.0'}
    for url in [row.get('Image URL'), row.get('Storage URL'), row.get('Source URL')]:
        if not url or pd.isna(url): continue
        if not str(url).startswith(('http://', 'https://')): continue
        try:
            r = requests.get(url, timeout=10, headers=headers)
            r.raise_for_status()
            if 'text/html' in r.headers.get('Content-Type', ''): continue
            return row['row_id'], Image.open(io.BytesIO(r.content)).convert('RGB')
        except Exception: continue
    return row['row_id'], None

def compute_embeddings(df: pd.DataFrame, clip_model, preprocess, cache_path: str) -> dict:
    embed_map = {}
    if os.path.exists(cache_path):
        with open(cache_path) as f:
            raw = json.load(f)
        embed_map = {k: np.array(v, dtype=np.float32) for k, v in raw.items() if k in df['row_id'].values}
        print(f"[train] Cache hit: {len(embed_map)} embeddings.")

    to_do = df[~df['row_id'].isin(embed_map.keys())].to_dict('records')
    if not to_do:
        return embed_map

    print(f"[train] Downloading {len(to_do)} images…")
    with ThreadPoolExecutor(max_workers=MAX_WORKERS) as ex:
        results = list(ex.map(_download, to_do))

    valid = [(rid, img) for rid, img in results if img is not None]
    print(f"[train] Encoding {len(valid)} images…")

    for i in range(0, len(valid), BATCH_SIZE):
        batch = valid[i:i+BATCH_SIZE]
        rids  = [b[0] for b in batch]
        imgs  = [b[1] for b in batch]
        t     = torch.stack([preprocess(img) for img in imgs]).to(device)
        with torch.no_grad():
            embs = clip_model.encode_image(t)
            embs = embs / embs.norm(dim=-1, keepdim=True)
            embs = embs.cpu().float().numpy()
        for rid, emb in zip(rids, embs):
            embed_map[rid] = emb

    with open(cache_path, 'w') as f:
        json.dump({k: v.tolist() for k, v in embed_map.items()}, f)
    print(f"[train] Cache saved → {cache_path}")
    return embed_map

# ── STEP 3: FEATURE ENGINEERING ───────────────────────────────────────────────

def build_features(df_ready, embed_map, fit=True, saved=None):
    T = saved or {}

    # A) CLIP → StandardScaler → PCA(95%)
    X_clip = np.vstack([embed_map[str(i)] for i in df_ready['row_id']])
    if fit:
        cs = StandardScaler();  X_cs  = cs.fit_transform(X_clip)
        pc = PCA(n_components=0.95, random_state=RANDOM_STATE); X_pca = pc.fit_transform(X_cs)
        T['clip_scaler'] = cs;  T['pca'] = pc
    else:
        X_pca = T['pca'].transform(T['clip_scaler'].transform(X_clip))

    # B) TF-IDF (250 features, 1-2 gram)
    names = df_ready['Name'].fillna('unknown')
    if fit:
        tf = TfidfVectorizer(max_features=250, stop_words='english', ngram_range=(1,2))
        X_text = tf.fit_transform(names).toarray();  T['tfidf'] = tf
    else:
        X_text = T['tfidf'].transform(names).toarray()

    # C) Meta-features
    df_ready = df_ready.copy()
    df_ready['name_length']     = df_ready['Name'].astype(str).apply(len)
    df_ready['word_count']      = df_ready['Name'].astype(str).apply(lambda x: len(x.split()))
    df_ready['avg_word_length'] = df_ready['name_length'] / (df_ready['word_count'] + 1e-5)
    df_ready['caps_ratio']      = df_ready['Name'].astype(str).apply(lambda x: sum(1 for c in x if c.isupper()) / (len(x)+1e-5))
    mc = ['name_length','word_count','avg_word_length','caps_ratio']
    if fit:
        ms = MinMaxScaler();  X_meta = ms.fit_transform(df_ready[mc]);  T['meta_scaler'] = ms
    else:
        X_meta = T['meta_scaler'].transform(df_ready[mc])

    # D) Domain OHE
    df_ready['Domain_Source'] = df_ready.get('Source Domain', pd.Series('unknown', index=df_ready.index)).fillna('unknown')
    df_ready['Domain_URL']    = df_ready.get('Source URL', pd.Series('unknown', index=df_ready.index)).apply(extract_domain)
    if fit:
        ohe = OneHotEncoder(sparse_output=False, handle_unknown='ignore')
        X_dom = ohe.fit_transform(df_ready[['Domain_Source','Domain_URL']]);  T['ohe'] = ohe
    else:
        X_dom = T['ohe'].transform(df_ready[['Domain_Source','Domain_URL']])

    return np.hstack([X_pca, X_text, X_meta, X_dom]), T

# ── STEP 4: STACKING CLASSIFIER ───────────────────────────────────────────────

def build_stacking():
    return StackingClassifier(
        estimators=[
            ('xgb', XGBClassifier(n_estimators=600, learning_rate=0.03, max_depth=6,
                                   subsample=0.8, colsample_bytree=0.7,
                                   tree_method='hist', device='cpu',   # CPU for HF free tier
                                   random_state=RANDOM_STATE, eval_metric='mlogloss')),
            ('lgbm', LGBMClassifier(n_estimators=600, learning_rate=0.03, num_leaves=31,
                                    subsample=0.8, colsample_bytree=0.7,
                                    random_state=RANDOM_STATE, verbose=-1)),
            ('et',   ExtraTreesClassifier(n_estimators=500, max_depth=15,
                                          max_features='sqrt', n_jobs=-1, random_state=RANDOM_STATE)),
            ('mlp',  MLPClassifier(hidden_layer_sizes=(128,64), activation='relu', solver='adam',
                                   alpha=0.01, max_iter=300, early_stopping=True, random_state=RANDOM_STATE)),
        ],
        final_estimator=LogisticRegression(max_iter=1000),
        cv=5, n_jobs=1, passthrough=False,
    )

# ── STEP 5: MAIN ──────────────────────────────────────────────────────────────

def train(csv_path: str, artifacts_dir: str = DEFAULT_ARTIFACTS,
          cache_path: str = DEFAULT_CACHE) -> dict:

    out = Path(artifacts_dir)
    out.mkdir(parents=True, exist_ok=True)

    # 1. Load data
    df = load_data(csv_path)

    # 2. CLIP embeddings
    clip_model, preprocess = load_clip()
    embed_map = compute_embeddings(df, clip_model, preprocess, cache_path)

    df_ready = df[df['row_id'].isin(embed_map.keys())].copy()
    print(f"[train] Rows with embeddings: {len(df_ready)}")

    # Prepare target labels before feature building so X and y always align.
    df_ready['price_tier'] = df_ready['price_tier_tagged']
    fill_mask = df_ready['price_tier'].isna()
    df_ready.loc[fill_mask, 'price_tier'] = df_ready.loc[fill_mask, 'price_inr'].apply(tier_from_price)
    df_ready = df_ready.dropna(subset=['price_tier']).copy()

    # Free CLIP memory before ensemble training
    del clip_model
    if device == 'cuda': torch.cuda.empty_cache()

    # 3. Features
    X, transforms = build_features(df_ready, embed_map, fit=True)
    print(f"[train] Feature matrix: {X.shape}")

    # 4. Target

    if len(df_ready) < 30:
        raise ValueError(f"Only {len(df_ready)} rows with embeddings and target labels. Need at least 30.")

    class_counts = df_ready['price_tier'].value_counts()
    if class_counts.shape[0] < 2:
        raise ValueError("Need at least 2 price tiers in training data.")

    le = LabelEncoder()
    y  = le.fit_transform(df_ready['price_tier'])

    use_stratify = class_counts.min() >= 2
    X_train, X_test, y_train, y_test = train_test_split(
        X, y, test_size=0.2, random_state=RANDOM_STATE, stratify=y if use_stratify else None
    )

    # 5. Train
    print("[train] Fitting stacking ensemble (5-fold CV) — 5–15 min on CPU…")
    clf = build_stacking()
    clf.fit(X_train, y_train)

    # 6. Evaluate
    y_pred = clf.predict(X_test)
    acc = accuracy_score(y_test, y_pred)
    report_dict = classification_report(y_test, y_pred, target_names=le.classes_, output_dict=True, zero_division=0)
    print(f"[train] ✅ Accuracy: {acc:.2%}")
    print(classification_report(y_test, y_pred, target_names=le.classes_, zero_division=0))

    # 7. Save — only 3 files needed by app.py
    joblib.dump(clf,        out / "model.joblib",        compress=3)
    joblib.dump(transforms, out / "transforms.joblib",   compress=3)
    joblib.dump(le,         out / "label_encoder.joblib", compress=3)
    print(f"[train] Saved artifacts to {out}")

    return {
        "accuracy": float(acc),
        "precision": float(report_dict["macro avg"]["precision"]),
        "recall": float(report_dict["macro avg"]["recall"]),
        "f1": float(report_dict["macro avg"]["f1-score"]),
        "support": int(report_dict["macro avg"]["support"]),
        "training_image_count": int(len(df_ready)),
    }

# ── CLI ───────────────────────────────────────────────────────────────────────

if __name__ == "__main__":
    ap = argparse.ArgumentParser()
    ap.add_argument("--csv",      default=os.getenv("CSV_PATH", "decor_multisite.csv"))
    ap.add_argument("--out",      default=DEFAULT_ARTIFACTS)
    ap.add_argument("--cache",    default=DEFAULT_CACHE)
    args = ap.parse_args()
    train(args.csv, args.out, args.cache)