#!/usr/bin/env python3
from __future__ import annotations

import argparse
import json
from pathlib import Path

import joblib
import pandas as pd
from sklearn.base import clone
from sklearn.calibration import CalibratedClassifierCV
from sklearn.metrics import accuracy_score
from sklearn.pipeline import Pipeline
from sklearn.preprocessing import FunctionTransformer, StandardScaler

from targeted_optimize_leagues import (
    REPO_ROOT,
    build_feature_frames,
    build_model_wrappers,
    load_df,
    optimize_threshold,
    select_columns,
)
from train_models import LEAGUES

SERVING_DIR = REPO_ROOT / 'artifacts' / 'v4_models_serving_current'


def build_pipeline(model_name: str):
    wrappers = {w.name: w for w in build_model_wrappers()}
    calibrated = model_name.endswith('_cal')
    base_name = model_name[:-4] if calibrated else model_name
    wrapper = wrappers[base_name]
    if calibrated:
        estimator = CalibratedClassifierCV(clone(wrapper.estimator), method='sigmoid', cv=3)
    else:
        estimator = clone(wrapper.estimator)
    scaler = StandardScaler() if wrapper.scale else FunctionTransformer(validate=False)
    model = estimator
    return scaler, model


def save_candidate(code: str, league: str, feature_mode: str, feature_set: str, train_window: str, model_name: str, threshold: float, out_dir: Path):
    df = load_df(code)
    frames = build_feature_frames(df)
    features = frames[feature_mode]
    split_idx = int(len(df) * 0.8)
    y_train_all = df['asian_result'].iloc[:split_idx].reset_index(drop=True)
    y_test = df['asian_result'].iloc[split_idx:].reset_index(drop=True)
    X_train_all = features.iloc[:split_idx].reset_index(drop=True)
    X_test_all = features.iloc[split_idx:].reset_index(drop=True)

    if train_window == 'all':
        X_train = X_train_all
        y_train = y_train_all
    else:
        window = int(train_window.replace('last_', ''))
        X_train = X_train_all.iloc[-window:].reset_index(drop=True)
        y_train = y_train_all.iloc[-window:].reset_index(drop=True)

    Xtr, Xte, cols = select_columns(X_train, y_train, X_test_all, feature_mode, feature_set)
    scaler, model = build_pipeline(model_name)
    Xtr_scaled = scaler.fit_transform(Xtr)
    Xte_scaled = scaler.transform(Xte)
    model.fit(Xtr_scaled, y_train)

    probs = model.predict_proba(Xte_scaled)[:, 1]
    preds = (probs >= 0.5).astype(int)
    test_acc = float(accuracy_score(y_test, preds))
    opt_threshold, opt_bet_acc, opt_bet_count, high_conf_acc, high_conf_count = optimize_threshold(y_test, probs, preds)

    out_dir.mkdir(parents=True, exist_ok=True)
    joblib.dump(scaler, out_dir / f'{code}_pro_v4_scaler.pkl')
    joblib.dump(model, out_dir / f'{code}_pro_v4_model.pkl')
    config = {
        'threshold': float(threshold),
        'selected_threshold': float(threshold),
        'opt_threshold': float(opt_threshold),
        'test_acc': test_acc,
        'bet_acc': float(opt_bet_acc),
        'bet_count': int(opt_bet_count),
        'high_conf_acc': float(high_conf_acc),
        'high_conf_count': int(high_conf_count),
        'features': cols,
        'model_name': model_name,
        'data_count': int(len(df)),
        'league': league,
        'feature_mode': feature_mode,
        'feature_set': feature_set,
        'train_window': train_window,
        'source': 'train_candidate_model.py',
    }
    joblib.dump(config, out_dir / f'{code}_pro_v4_config.pkl')
    return config


def update_manifest(code: str, out_dir: Path):
    manifest_path = out_dir / 'manifest.json'
    cfg = joblib.load(out_dir / f'{code}_pro_v4_config.pkl')
    rows = []
    if manifest_path.exists():
        rows = json.loads(manifest_path.read_text(encoding='utf-8'))
    rows = [r for r in rows if r.get('code') != code]
    rows.append({
        'code': code,
        'league': cfg.get('league', LEAGUES.get(code, code)),
        'model_name': cfg.get('model_name'),
        'bet_acc': float(cfg.get('bet_acc', 0.0)),
        'threshold': float(cfg.get('selected_threshold', cfg.get('threshold', 0.55))),
        'feature_mode': cfg.get('feature_mode', 'baseline_v4'),
        'feature_count': len(cfg.get('features', [])),
        'source': 'stable_candidate',
    })
    rows.sort(key=lambda x: x['code'])
    manifest_path.write_text(json.dumps(rows, ensure_ascii=False, indent=2), encoding='utf-8')


def main():
    parser = argparse.ArgumentParser(description='Train one selected candidate and optionally write into serving dir.')
    parser.add_argument('--code', required=True)
    parser.add_argument('--league', default=None)
    parser.add_argument('--feature-mode', required=True)
    parser.add_argument('--feature-set', required=True)
    parser.add_argument('--train-window', required=True)
    parser.add_argument('--model-name', required=True)
    parser.add_argument('--threshold', type=float, required=True)
    parser.add_argument('--output-dir', type=Path, default=SERVING_DIR)
    args = parser.parse_args()

    league = args.league or LEAGUES.get(args.code, args.code)
    cfg = save_candidate(
        code=args.code,
        league=league,
        feature_mode=args.feature_mode,
        feature_set=args.feature_set,
        train_window=args.train_window,
        model_name=args.model_name,
        threshold=args.threshold,
        out_dir=args.output_dir,
    )
    update_manifest(args.code, args.output_dir)
    print(json.dumps({'code': args.code, 'league': league, 'output_dir': str(args.output_dir), 'config': cfg}, ensure_ascii=False, indent=2))


if __name__ == '__main__':
    main()
