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

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)

# Contributions are written here by scripts/export_contrib.py during CI retrain.
_REPO_CONTRIB_FILE = Path(__file__).parent / "superpitch" / "assets" / "data" / "contrib.jsonl"


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 load_contrib():
    """Load contributed decisions as training examples (Layer 3 — federated signal)."""
    if not _REPO_CONTRIB_FILE.exists():
        return []
    examples = load_examples(str(_REPO_CONTRIB_FILE))
    for ex in examples:
        ex["source"] = "contrib"
        ex.setdefault("source_file", "contrib.jsonl")
        ex.setdefault("text", pitch_to_text(ex.get("input", {})) if ex.get("input") else ex.get("text", ""))
        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)

    # Merge contributed decisions (Layer 3 — federated signal from other operators)
    contrib_examples = load_contrib()
    if contrib_examples:
        print(f"Contrib: {len(contrib_examples)} decisions from contributed data")
        train_examples.extend(contrib_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()
