"""Recover unsaved classifiers from frozen Dev parameters; never select on Test."""
from pathlib import Path
import os, sys, json, hashlib, time
os.environ.setdefault('OPENBLAS_NUM_THREADS', '2')
os.environ.setdefault('OMP_NUM_THREADS', '2')
ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(ROOT))
import numpy as np
import pandas as pd
import joblib
from gensim.models import Word2Vec, FastText
from sklearn.linear_model import LogisticRegression
from sklearn.svm import LinearSVC
from sklearn.naive_bayes import ComplementNB
from compact_fasttext import CompactFastText


def matrix(texts, kv, fasttext=False):
    values = []
    for text in texts:
        vectors = [kv[w] for w in text.split() if fasttext or w in kv]
        values.append(np.mean(vectors, axis=0).astype(np.float32) if vectors else np.zeros(200, np.float32))
    return np.vstack(values)


def main():
    out = ROOT / 'models/baseline'
    train_path = ROOT / 'data/processed/train.csv'
    train = pd.read_csv(train_path, keep_default_na=False)
    selection = pd.read_csv(ROOT / 'reports/original/baseline_dev_selection.csv')
    vectorizer = joblib.load(out / 'tfidf_vectorizer.joblib')
    w2v = Word2Vec.load(str(out / 'word2vec.model')).wv
    ft_path = out / 'fasttext_recovered.npz'
    if not ft_path.exists():
        print('Original FastText is missing its external n-gram weights. Training a NEW recovery model on train only.', flush=True)
        ft = FastText(sentences=[x.split() for x in train.seg_text], vector_size=200,
                      window=5, min_count=2, sg=1, negative=10, epochs=20, seed=42,
                      workers=1, min_n=2, max_n=5)
        print('Saving exact compact inference weights...', flush=True)
        CompactFastText.export(ft, ft_path)
        del ft
    ft = CompactFastText(ft_path)
    features = {'TF-IDF': vectorizer.transform(train.nostop_text),
                'Word2Vec': matrix(train.seg_text, w2v),
                'FastText': matrix(train.seg_text, ft, True)}
    metadata = {'data_sha256': hashlib.sha256(train_path.read_bytes()).hexdigest(),
                'seed': 42, 'selection': 'original frozen Dev parameters; no new search',
                'fasttext_note': 'New FastText run: original vectors_ngrams.npy was not supplied. Historical metrics are kept separately.',
                'models': {}}
    for _, row in selection.iterrows():
        key = row['key']; params = json.loads(row.best_params)
        if key == 'TF-IDF+LogisticRegression':
            fitted = joblib.load(out / 'best_baseline_model.joblib'); source = 'original_checkpoint'
        else:
            if row.model == 'LogisticRegression':
                fitted = LogisticRegression(max_iter=2000, class_weight='balanced', random_state=42, **params)
            elif row.model == 'LinearSVM':
                fitted = LinearSVC(class_weight='balanced', random_state=42, **params)
            else:
                fitted = ComplementNB(**params)
            fitted.fit(features[row.feature], train.label_bin)
            source = 'retrained_fasttext_and_classifier' if row.feature == 'FastText' else 'refit_classifier_frozen_parameters'
        joblib.dump(fitted, out / (key + '.joblib'), compress=3)
        metadata['models'][key] = {'source': source, 'params': params}
        print('Saved', key, source, flush=True)
    (out / 'recovery_manifest.json').write_text(json.dumps(metadata, ensure_ascii=False, indent=2), encoding='utf-8')


if __name__ == '__main__':
    main()