import os

import joblib
import numpy as np

import config
from core.deep_learning_model import load_lstm_mlp_bundle, predict_lstm_mlp
from core.model_features import (
    align_feature_frame,
    apply_label_encoders,
    build_prediction_dataframe,
    engineer_features,
    to_float_array,
)


MODEL_INFO = {
    'random_forest': {'name': 'random_forest', 'name_cn': '随机森林', 'description': '稳健的树模型集成'},
    'xgboost': {'name': 'xgboost', 'name_cn': 'XGBoost', 'description': '梯度提升树模型'},
    'lightgbm': {'name': 'lightgbm', 'name_cn': 'LightGBM', 'description': '轻量级梯度提升树'},
    'gradient_boosting': {'name': 'gradient_boosting', 'name_cn': 'GBDT', 'description': '梯度提升决策树'},
    'extra_trees': {'name': 'extra_trees', 'name_cn': '极端随机树', 'description': '高随机性的树模型'},
    'stacking': {'name': 'stacking', 'name_cn': 'Stacking集成', 'description': '多模型融合'},
    'lstm_mlp': {
        'name': 'lstm_mlp',
        'name_cn': '时序注意力融合网络',
        'description': 'Transformer时序编码 + 静态特征门控融合的深度学习模型',
    },
}


class PredictService:
    def __init__(self):
        self.models = {}
        self.classifiers = {}
        self.classification_metrics = {}
        self.scaler = None
        self.feature_names = None
        self.selected_features = None
        self.label_encoders = {}
        self.model_metrics = {}
        self.training_metadata = {}
        self.default_model = 'random_forest'

    def _ensure_models_loaded(self):
        if not self.models:
            self.load_models()

    def load_models(self):
        metadata_path = os.path.join(config.MODELS_DIR, 'training_metadata.pkl')
        if os.path.exists(metadata_path):
            self.training_metadata = joblib.load(metadata_path)

        model_files = {
            'random_forest': 'random_forest_model.pkl',
            'xgboost': 'xgboost_model.pkl',
            'lightgbm': 'lightgbm_model.pkl',
            'gradient_boosting': 'gradient_boosting_model.pkl',
            'extra_trees': 'extra_trees_model.pkl',
            'stacking': 'stacking_model.pkl',
            'lstm_mlp': 'lstm_mlp_model.pt',
        }
        allowed_models = self.training_metadata.get('available_models')
        if allowed_models:
            model_files = {k: v for k, v in model_files.items() if k in allowed_models}

        for name, filename in model_files.items():
            path = os.path.join(config.MODELS_DIR, filename)
            if os.path.exists(path):
                try:
                    if name == 'lstm_mlp':
                        bundle = load_lstm_mlp_bundle(path)
                        if bundle is not None:
                            self.models[name] = bundle
                    else:
                        self.models[name] = joblib.load(path)
                except Exception as exc:
                    print(f'Failed to load model {name}: {exc}')

        if os.path.exists(config.SCALER_PATH):
            self.scaler = joblib.load(config.SCALER_PATH)
        for filename, attr in [
            ('feature_names.pkl', 'feature_names'),
            ('selected_features.pkl', 'selected_features'),
            ('label_encoders.pkl', 'label_encoders'),
            ('model_metrics.pkl', 'model_metrics'),
        ]:
            path = os.path.join(config.MODELS_DIR, filename)
            if os.path.exists(path):
                try:
                    setattr(self, attr, joblib.load(path))
                except Exception as exc:
                    print(f'Failed to load artifact {filename}: {exc}')

        valid_metrics = {key: value for key, value in self.model_metrics.items() if key in self.models}
        if valid_metrics:
            self.default_model = max(valid_metrics.items(), key=lambda item: item[1]['r2'])[0]

        # 加载风险分类模型
        for name in ['random_forest', 'gradient_boosting', 'lightgbm', 'xgboost']:
            path = os.path.join(config.MODELS_DIR, f'risk_{name}_classifier.pkl')
            if os.path.exists(path):
                try:
                    self.classifiers[name] = joblib.load(path)
                except Exception:
                    pass
        cls_metrics_path = os.path.join(config.MODELS_DIR, 'classification_metrics.pkl')
        if os.path.exists(cls_metrics_path):
            try:
                self.classification_metrics = joblib.load(cls_metrics_path)
            except Exception:
                pass

    def get_available_models(self):
        self._ensure_models_loaded()
        models = []
        for name in self.models.keys():
            info = MODEL_INFO.get(name, {'name': name, 'name_cn': name, 'description': ''}).copy()
            info['is_available'] = True
            info['is_default'] = name == self.default_model
            info['metrics'] = self.model_metrics.get(name, {'r2': 0, 'rmse': 0, 'mae': 0})
            models.append(info)
        models.sort(key=lambda item: item['metrics']['r2'], reverse=True)
        return models

    def predict_single(self, data, model_type=None):
        self._ensure_models_loaded()
        model_type = model_type or self.default_model
        if model_type not in self.models:
            fallback = next(iter(self.models), None)
            if fallback is None:
                return self._get_default_prediction(data)
            model_type = fallback
        if self.scaler is None or self.feature_names is None:
            return self._get_default_prediction(data)

        features = self._prepare_features(data)
        try:
            if model_type == 'lstm_mlp':
                current_df = build_prediction_dataframe(data)
                predicted_hours = predict_lstm_mlp(self.models[model_type], current_df)
            else:
                predicted_hours = self.models[model_type].predict([features])[0]
                predicted_hours = self._inverse_transform_prediction(predicted_hours)
            predicted_hours = max(0.5, float(predicted_hours))
        except Exception:
            return self._get_default_prediction(data)

        risk_level, risk_label = self._get_risk_level(predicted_hours)
        confidence = max(0.5, self.model_metrics.get(model_type, {}).get('r2', 0.82))

        # 风险分类概率
        risk_probability = self._get_risk_probability(features, model_type)

        return {
            'predicted_hours': round(predicted_hours, 2),
            'risk_level': risk_level,
            'risk_label': risk_label,
            'risk_probability': risk_probability,
            'confidence': round(confidence, 2),
            'model_used': model_type,
            'model_name_cn': MODEL_INFO.get(model_type, {}).get('name_cn', model_type),
        }

    def predict_compare(self, data):
        self._ensure_models_loaded()
        results = []
        for name in self.models.keys():
            result = self.predict_single(data, name)
            result['model'] = name
            result['model_name_cn'] = MODEL_INFO.get(name, {}).get('name_cn', name)
            result['r2'] = self.model_metrics.get(name, {}).get('r2', 0)
            results.append(result)
        results.sort(key=lambda item: item.get('r2', 0), reverse=True)
        if results:
            results[0]['recommended'] = True
        return results

    def _prepare_features(self, data):
        X_df = build_prediction_dataframe(data)
        X_df = engineer_features(X_df)
        X_df = apply_label_encoders(X_df, self.label_encoders)
        X_df = align_feature_frame(X_df, self.feature_names)
        features = self.scaler.transform(to_float_array(X_df))[0]
        if self.selected_features:
            selected_indices = [self.feature_names.index(name) for name in self.selected_features if name in self.feature_names]
            if selected_indices:
                features = features[selected_indices]
        return features

    def _inverse_transform_prediction(self, prediction):
        if self.training_metadata.get('target_transform') == 'log1p':
            return float(np.expm1(prediction))
        return float(prediction)

    def _get_risk_level(self, hours):
        if hours < 4:
            return 'low', '低风险'
        if hours <= 8:
            return 'medium', '中风险'
        return 'high', '高风险'

    def _get_default_prediction(self, data):
        base_hours = 3.8
        base_hours += min(float(data.get('monthly_overtime_hours', 24)) / 20, 3.0)
        base_hours += min(float(data.get('commute_minutes', 40)) / 50, 2.0)
        base_hours += 1.6 if int(data.get('is_night_shift', 0)) == 1 else 0
        base_hours += 1.8 if int(data.get('chronic_disease_flag', 0)) == 1 else 0
        base_hours += 0.9 if int(data.get('near_holiday_flag', 0)) == 1 else 0
        base_hours += 0.8 if int(data.get('medical_certificate_flag', 0)) == 1 else 0
        base_hours += 0.5 * int(data.get('children_count', 0))
        if data.get('leave_type') in ['病假', '工伤假', '婚假', '丧假']:
            base_hours += 2.5
        if data.get('stress_level') == '高':
            base_hours += 0.9
        if data.get('performance_level') == 'A':
            base_hours -= 0.5
        risk_level, risk_label = self._get_risk_level(base_hours)
        return {
            'predicted_hours': round(max(0.5, base_hours), 2),
            'risk_level': risk_level,
            'risk_label': risk_label,
            'risk_probability': {'low': 0.0, 'medium': 1.0, 'high': 0.0},
            'confidence': 0.72,
            'model_used': 'default',
            'model_name_cn': '默认规则',
        }

    def _get_risk_probability(self, features, model_type):
        """获取分类器预测的风险概率"""
        classifier = self.classifiers.get(model_type)
        if classifier is None:
            classifier = self.classifiers.get('random_forest')
        if classifier is None:
            return {'low': 0.0, 'medium': 1.0, 'high': 0.0}
        try:
            proba = classifier.predict_proba([features])[0]
            classes = list(classifier.classes_)
            result = {'low': 0.0, 'medium': 0.0, 'high': 0.0}
            label_map = {0: 'low', 1: 'medium', 2: 'high'}
            for idx, cls in enumerate(classes):
                if cls in label_map:
                    result[label_map[cls]] = round(float(proba[idx]), 4)
            return result
        except Exception:
            return {'low': 0.0, 'medium': 1.0, 'high': 0.0}

    def predict_risk_classification(self, data, model_type=None):
        """使用分类模型直接预测风险等级"""
        self._ensure_models_loaded()
        model_type = model_type or self.default_model
        classifier = self.classifiers.get(model_type)
        if classifier is None:
            classifier = self.classifiers.get('random_forest')
        if classifier is None or self.scaler is None:
            return None

        features = self._prepare_features(data)
        try:
            pred_class = int(classifier.predict([features])[0])
            proba = classifier.predict_proba([features])[0]
            label_map = {0: 'low', 1: 'medium', 2: 'high'}
            risk_labels_map = {'low': '低风险', 'medium': '中风险', 'high': '高风险'}
            risk_level = label_map.get(pred_class, 'medium')

            classes = list(classifier.classes_)
            probabilities = {'low': 0.0, 'medium': 0.0, 'high': 0.0}
            for idx, cls in enumerate(classes):
                if cls in label_map:
                    probabilities[label_map[cls]] = round(float(proba[idx]), 4)

            return {
                'risk_level': risk_level,
                'risk_label': risk_labels_map[risk_level],
                'risk_probability': probabilities,
                'model_used': model_type,
                'classification_metrics': self.classification_metrics.get(model_type, {}),
            }
        except Exception:
            return None

    def get_model_info(self):
        self._ensure_models_loaded()
        return {
            'models': self.get_available_models(),
            'training_info': {
                'train_samples': self.training_metadata.get('train_samples', 0),
                'test_samples': self.training_metadata.get('test_samples', 0),
                'feature_count': self.training_metadata.get('feature_count_after_selection', 0),
                'training_date': self.training_metadata.get('training_date', ''),
                'sequence_window_size': self.training_metadata.get('sequence_window_size', 0),
                'deep_learning_available': self.training_metadata.get('deep_learning_available', False),
            },
        }


predict_service = PredictService()
