"""Optional NEW training of unsaved configurations; never impersonates old runs.
Example: python scripts/train_missing_deep.py --model Embedding+CNN
All missing: python scripts/train_missing_deep.py --all-missing
Use a GPU environment for the full sweep. Test is not read during training.
"""
from pathlib import Path
import sys,os,json,pickle,argparse,gc
ROOT=Path(__file__).resolve().parents[1];sys.path.insert(0,str(ROOT))
os.environ.setdefault('TF_CPP_MIN_LOG_LEVEL','2')
import numpy as np
import pandas as pd
import tensorflow as tf
from model_registry import registry,slug

def build(spec,vocab,embedding):
    L=tf.keras.layers;inputs=L.Input(shape=(52,))
    kwargs={'weights':[embedding]} if embedding is not None else {}
    x=L.Embedding(vocab,300,trainable=True,**kwargs)(inputs)
    architecture=spec['architecture'];family=spec['family']
    if family=='CNN First':x=L.Conv1D(128,3,activation='relu')(x)
    recurrent=[s for s in architecture.split('+') if s in ('BiLSTM','BiGRU')]
    for i,name in enumerate(recurrent):
        cls=L.LSTM if name=='BiLSTM' else L.GRU
        returns=(family!='No CNN' or i<len(recurrent)-1)
        x=L.Bidirectional(cls(128,return_sequences=returns))(x)
    if family=='CNN Last':x=L.Conv1D(128,3,activation='relu')(x)
    if family!='No CNN':x=L.GlobalMaxPooling1D()(x)
    x=L.Dropout(.3)(x);outputs=L.Dense(1,activation='sigmoid')(x)
    model=tf.keras.Model(inputs,outputs)
    model.compile(optimizer='adam',loss='binary_crossentropy',metrics=['accuracy'])
    return model

def main():
    parser=argparse.ArgumentParser();parser.add_argument('--model');parser.add_argument('--all-missing',action='store_true')
    parser.add_argument('--epochs',type=int,default=20);args=parser.parse_args()
    specs=[m for m in registry() if m['kind']=='keras' and not m['available'] and (args.all_missing or m['key']==args.model or m['id']==args.model)]
    if not specs:raise SystemExit('Select an unsaved model using --model KEY or --all-missing. Existing checkpoints are never overwritten.')
    train=pd.read_csv(ROOT/'data/processed/train.csv',keep_default_na=False)
    dev=pd.read_csv(ROOT/'data/processed/dev.csv',keep_default_na=False)
    # New runs use the leakage-filtered development subset, unlike old DL runs.
    dev=dev.loc[~dev.clean.isin(set(train.clean))]
    with (ROOT/'models/deep/tokenizer.pkl').open('rb') as f:tokenizer=pickle.load(f)
    def sequences(df):return tf.keras.preprocessing.sequence.pad_sequences(tokenizer.texts_to_sequences(df.seg_text),maxlen=52,padding='post',truncating='post')
    x=sequences(train);xd=sequences(dev);y=train.label_bin.to_numpy();yd=dev.label_bin.to_numpy()
    vocab=min(tokenizer.num_words,len(tokenizer.word_index)+1)
    out=ROOT/'models/extras';out.mkdir(exist_ok=True)
    for spec in specs:
        tf.keras.backend.clear_session();gc.collect();tf.keras.utils.set_random_seed(42)
        embedding=None
        if spec['feature']!='Embedding':embedding=np.load(ROOT/('models/deep/'+('word2vec' if spec['feature']=='Word2Vec' else 'fasttext')+'_matrix.npy'))
        model=build(spec,vocab,embedding)
        print('Training new checkpoint:',spec['key'],flush=True)
        history=model.fit(x,y,validation_data=(xd,yd),epochs=args.epochs,batch_size=32,
                          callbacks=[tf.keras.callbacks.EarlyStopping(monitor='val_loss',patience=5,restore_best_weights=True)],verbose=2)
        model.save(ROOT/spec['path'])
        metadata=dict(key=spec['key'],origin='new_training_run',seed=42,selection='dev_tune',max_len=52,
                      epochs=len(history.history['loss']),history=history.history,note='Not the missing historical checkpoint. Test was not used for training/selection.')
        (out/(spec['key']+'.json')).write_text(json.dumps(metadata,ensure_ascii=False,indent=2),encoding='utf-8')
        print('Saved. Restart web to see this model; regenerate evaluation report to include its Test metrics.',flush=True)

if __name__=='__main__':main()