"""Retrain the local Superpitch model from bundled data and user feedback."""
import argparse
import os
import pickle
from collections import Counter, defaultdict

import numpy as np
from sklearn.neural_network import MLPRegressor

from model_a import (
    fit_model_a_bundle,
    predict_model_a_bundle,
)
from superpitch.app_paths import (
    FEEDBACK_FILE,
    LEARNED_MODEL_FILE,
    LATEST_MODEL_FILE,
    TRAIN_FILE,
    VAL_FILE,
    bootstrap_user_files,
    ensure_app_dirs,
)
from superpitch.state import count_feedback_entries, load_config, mark_train_complete
from prepare import (
    embedding_backend_available,
    example_training_weight,
    FEATURE_NAMES, CASE_FEATURE_NAMES,
    load_examples,
    embed_texts,
    preferred_embedding_backend,
    normalize_example_source, pitch_to_text,
)

HIDDEN = (128, 64)


def train_model_b(train_texts, val_texts, s_train):
    emb_train = embed_texts(train_texts, data_path=str(TRAIN_FILE), n_texts=len(train_texts)).astype(np.float32)
    emb_val = embed_texts(val_texts, data_path=str(VAL_FILE)).astype(np.float32)
    model_b = MLPRegressor(
        hidden_layer_sizes=HIDDEN,
        activation="relu",
        solver="adam",
        alpha=1e-4,
        batch_size=min(32, len(train_texts)),
        learning_rate_init=1e-3,
        max_iter=400,
        early_stopping=True,
        validation_fraction=0.15,
        random_state=42,
    )
    model_b.fit(emb_train, s_train)
    pred_b_train = np.clip(model_b.predict(emb_train), 1, 10)
    pred_b_val = np.clip(model_b.predict(emb_val), 1, 10)
    return model_b, pred_b_train, pred_b_val


def unpack_examples(examples):
    texts, labels, scores, sub_scores, inputs, sources = [], [], [], [], [], []
    weights = []
    for ex in examples:
        inp = ex.get("input") or {"one_liner": ex.get("text", "")[:200]}
        text = ex.get("text") or pitch_to_text(inp)
        texts.append(text)
        labels.append(1 if ex["decision"] == "approve" else 0)
        scores.append(float(ex.get("overall_score", ex.get("score", 5))))
        sub_scores.append([
            float(ex.get("team_score", -1)),
            float(ex.get("traction_score", -1)),
            float(ex.get("market_score", -1)),
            float(ex.get("risk_score", -1)),
        ])
        inputs.append(inp)
        sources.append(normalize_example_source(ex))
        weights.append(example_training_weight(ex))
    return texts, labels, scores, sub_scores, inputs, sources, np.array(weights, dtype=np.float64)


def load_feedback():
    """Load user feedback as additional training examples (Layer 2)."""
    if not os.path.exists(FEEDBACK_FILE):
        return []
    examples = load_examples(FEEDBACK_FILE)
    for ex in examples:
        ex["source"] = "feedback"
        ex.setdefault("source_file", "feedback.jsonl")
        ex.setdefault("text", pitch_to_text(ex.get("input", {})))
        ex.setdefault("team_score", -1)
        ex.setdefault("traction_score", -1)
        ex.setdefault("market_score", -1)
        ex.setdefault("risk_score", -1)
    return examples


def train(skip_model_b: bool = False):
    bootstrap_user_files()
    ensure_app_dirs()
    print("Loading data...")
    train_examples = load_examples(TRAIN_FILE)
    val_examples = load_examples(VAL_FILE)

    # Merge user feedback into training data (Layer 2 — your preferences)
    fb_examples = load_feedback()
    if fb_examples:
        print(f"Feedback: {len(fb_examples)} labeled pitches from you")
        train_examples.extend(fb_examples)

    X_train_raw, y_train, scores_train, sub_train, inputs_train, train_sources, train_weights = unpack_examples(train_examples)
    X_val_raw,   y_val,   scores_val,   _sub_val,   inputs_val,   val_sources,   _ = unpack_examples(val_examples)

    approve_train = sum(y_train)
    approve_val = sum(y_val)
    print(f"Train: {len(X_train_raw)} examples  ({approve_train} approve / {len(y_train)-approve_train} reject)")
    print(f"Val:   {len(X_val_raw)} examples  ({approve_val} approve / {len(y_val)-approve_val} reject)")
    train_mix = Counter(train_sources)
    val_mix = Counter(val_sources)
    print(f"Train sources: {dict(train_mix)}")
    print(f"Val sources:   {dict(val_mix)}")
    by_source_weights = defaultdict(list)
    for source, weight in zip(train_sources, train_weights):
        by_source_weights[source].append(float(weight))
    print("Train weights:", ", ".join(
        f"{source}={np.mean(values):.2f}" for source, values in sorted(by_source_weights.items())
    ))

    s_train = np.array(scores_train, dtype=np.float64)
    s_val   = np.array(scores_val,   dtype=np.float64)
    sub_train_np = np.array(sub_train, dtype=np.float64)

    print(f"\n{'='*50}")
    print(f"  Superpitch model: {len(FEATURE_NAMES)} features + Ridge")
    print(f"{'='*50}")
    bundle = fit_model_a_bundle(
        X_train_raw,
        inputs_train,
        train_examples,
        sub_train_np,
        s_train,
        train_weights,
    )
    case_index = bundle.get("case_feature_index")
    pred_a_train, _, _, _, _, _ = predict_model_a_bundle(
        bundle,
        X_train_raw,
        inputs_train,
        query_matrix=case_index["matrix"] if case_index is not None else None,
        self_indices=np.arange(len(X_train_raw)),
    )
    pred_a_val, _, _, _, _, _ = predict_model_a_bundle(bundle, X_val_raw, inputs_val)
    mae_a_train = np.mean(np.abs(s_train - pred_a_train))
    mae_a_val   = np.mean(np.abs(s_val - pred_a_val))

    # Feature importance (show original + case-memory features)
    model_a = bundle["model"]
    primary_feature_names = FEATURE_NAMES + CASE_FEATURE_NAMES
    coefs = model_a.coef_[:len(primary_feature_names)]
    importance = sorted(zip(primary_feature_names, coefs), key=lambda x: abs(x[1]), reverse=True)
    print(f"  Train MAE: {mae_a_train:.2f}  |  Val MAE: {mae_a_val:.2f}")
    print(f"\n  Top features:")
    for name, coef in importance[:5]:
        print(f"    {coef:+.3f}  {name}")

    model_b = None
    pred_b_train = None
    pred_b_val = None

    print(f"\n{'='*50}")
    print(f"  Optional Model B: embeddings + MLP")
    print(f"{'='*50}")
    if skip_model_b:
        print("  Skipped (--skip-model-b).")
    elif not embedding_backend_available():
        print("  Skipped (no embedding backend available).")
    else:
        backend = preferred_embedding_backend()
        print(f"  Embedding backend: {backend}")
        model_b, pred_b_train, pred_b_val = train_model_b(X_train_raw, X_val_raw, s_train)
        mae_b_train = np.mean(np.abs(s_train - pred_b_train))
        mae_b_val = np.mean(np.abs(s_val - pred_b_val))
        print(f"  Train MAE: {mae_b_train:.2f}  |  Val MAE: {mae_b_val:.2f}")

    os.makedirs(LATEST_MODEL_FILE.parent, exist_ok=True)
    artifacts = {
        **bundle,
        "train_source_mix": dict(train_mix),
    }
    with open(LATEST_MODEL_FILE, "wb") as f:
        pickle.dump(artifacts, f)

    if model_b is not None:
        with open(LEARNED_MODEL_FILE, "wb") as f:
            pickle.dump(
                {
                    "model": model_b,
                    "embedding_backend": preferred_embedding_backend(),
                },
                f,
            )

    mark_train_complete(load_config(), feedback_count=count_feedback_entries())

    print("\nDone.")
    print(f"  Saved model -> {LATEST_MODEL_FILE}")
    if model_b is not None:
        print(f"  Saved Model B weights -> {LEARNED_MODEL_FILE}")
    print(f"  Training set -> {len(X_train_raw)} pitches")
    print(f"  Validation set -> {len(X_val_raw)} pitches")


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument(
        "--skip-model-b",
        action="store_true",
        help="Train only the shipped Model A and skip the optional embedding model.",
    )
    args = parser.parse_args()
    train(skip_model_b=args.skip_model_b)


if __name__ == "__main__":
    main()
