
from flask import Blueprint, render_template, request, redirect, url_for, flash, session, jsonify
from app import db
from app.models import User, PredictionRecord
import joblib
import pandas as pd
import numpy as np
import os

main = Blueprint('main', __name__)

# 加载模型和预处理器
MODEL_PATH = 'models/stroke_model.pkl'
SCALER_PATH = 'models/scaler.pkl'
ENCODERS_PATH = 'models/label_encoders.pkl'

model = None
scaler = None
label_encoders = None

def load_artifacts():
    global model, scaler, label_encoders
    if os.path.exists(MODEL_PATH):
        model = joblib.load(MODEL_PATH)
        scaler = joblib.load(SCALER_PATH)
        label_encoders = joblib.load(ENCODERS_PATH)

# 在首次请求前尝试加载模型
@main.before_app_request
def before_first_request():
    if model is None:
        try:
            load_artifacts()
        except Exception as e:
            print(f"Error loading model: {e}")

@main.route('/')
def index():
    return render_template('index.html')

@main.route('/login', methods=['GET', 'POST'])
def login():
    if request.method == 'POST':
        username = request.form.get('username')
        password = request.form.get('password')
        
        user = User.query.filter_by(username=username).first()
        if user and user.password == password:
            session['user_id'] = user.id
            session['username'] = user.username
            return redirect(url_for('main.dashboard'))
        else:
            flash('用户名或密码错误')
            
    return render_template('login.html')

@main.route('/register', methods=['GET', 'POST'])
def register():
    if request.method == 'POST':
        username = request.form.get('username')
        password = request.form.get('password')
        
        # 简单验证
        if User.query.filter_by(username=username).first():
            flash('用户名已存在')
        else:
            new_user = User(username=username, password=password)
            db.session.add(new_user)
            db.session.commit()
            flash('注册成功，请登录')
            return redirect(url_for('main.login'))
            
    return render_template('register.html')

@main.route('/logout')
def logout():
    session.pop('user_id', None)
    session.pop('username', None)
    return redirect(url_for('main.index'))

@main.route('/dashboard')
def dashboard():
    if 'user_id' not in session:
        return redirect(url_for('main.login'))
    return render_template('dashboard.html')

@main.route('/analysis')
def analysis():
    if 'user_id' not in session:
        return redirect(url_for('main.login'))
    
    # 读取数据
    df = pd.read_csv('healthcare-dataset-stroke-data.csv')
    
    # 1. 目标变量分布
    stroke_counts = df['stroke'].value_counts()
    target_dist = [
        {'value': int(stroke_counts.get(1, 0)), 'name': '中风'},
        {'value': int(stroke_counts.get(0, 0)), 'name': '未中风'}
    ]
    
    # 2. 性别分布
    gender_counts = df['gender'].value_counts()
    gender_dist = [{'value': int(v), 'name': k} for k, v in gender_counts.items()]
    
    # 3. 年龄分布 (分箱)
    age_bins = [0, 20, 40, 60, 80, 100]
    age_labels = ['0-20', '20-40', '40-60', '60-80', '80+']
    df['age_group'] = pd.cut(df['age'], bins=age_bins, labels=age_labels)
    age_dist = df['age_group'].value_counts().sort_index()
    age_data = {'categories': age_labels, 'values': age_dist.values.tolist()}
    
    # 4. 疾病关联 (高血压、心脏病与中风)
    hypertension_stroke = df[df['stroke'] == 1]['hypertension'].value_counts()
    heart_disease_stroke = df[df['stroke'] == 1]['heart_disease'].value_counts()
    
    disease_stats = {
        'hypertension': int(hypertension_stroke.get(1, 0)),
        'heart_disease': int(heart_disease_stroke.get(1, 0)),
        'total_stroke': int(stroke_counts.get(1, 0))
    }
    
    # 5. 工作类型分布
    work_counts = df['work_type'].value_counts()
    work_dist = {'categories': work_counts.index.tolist(), 'values': work_counts.values.tolist()}
    
    # 6. 吸烟状况分布
    smoking_counts = df['smoking_status'].value_counts()
    smoking_dist = {'categories': smoking_counts.index.tolist(), 'values': smoking_counts.values.tolist()}

    # 7. 散点图数据 (BMI vs Glucose) - 抽样以避免过多点
    sample_df = df.sample(n=min(1000, len(df)), random_state=42)
    # 过滤掉 N/A 的 BMI
    sample_df = sample_df[sample_df['bmi'] != 'N/A']
    scatter_data = []
    for _, row in sample_df.iterrows():
        try:
             bmi = float(row['bmi'])
             scatter_data.append([
                 row['avg_glucose_level'],
                 bmi,
                 int(row['stroke']) 
             ])
        except:
            continue

    data = {
        'target_dist': target_dist,
        'gender_dist': gender_dist,
        'age_data': age_data,
        'disease_stats': disease_stats,
        'work_dist': work_dist,
        'smoking_dist': smoking_dist,
        'scatter_data': scatter_data
    }
    
    return render_template('analysis.html', data=data)

@main.route('/predict', methods=['GET', 'POST'])
def predict():
    if 'user_id' not in session:
        return redirect(url_for('main.login'))
    
    result = None
    if request.method == 'POST':
        try:
            # 获取表单数据
            data = {
                'gender': request.form.get('gender'),
                'age': float(request.form.get('age')),
                'hypertension': int(request.form.get('hypertension')),
                'heart_disease': int(request.form.get('heart_disease')),
                'ever_married': request.form.get('ever_married'),
                'work_type': request.form.get('work_type'),
                'Residence_type': request.form.get('Residence_type'),
                'avg_glucose_level': float(request.form.get('avg_glucose_level')),
                'bmi': float(request.form.get('bmi')),
                'smoking_status': request.form.get('smoking_status')
            }
            
            # 预处理数据
            input_df = pd.DataFrame([data])
            
            # 编码类别特征
            for col, le in label_encoders.items():
                input_df[col] = le.transform(input_df[col])
                
            # 标准化
            input_scaled = scaler.transform(input_df)
            
            # 预测
            prediction = model.predict(input_scaled)[0]
            prob = model.predict_proba(input_scaled)[0][1]
            
            result = {
                'prediction': int(prediction),
                'probability': round(prob * 100, 2)
            }
            
            # 保存记录
            record = PredictionRecord(
                user_id=session['user_id'],
                gender=data['gender'],
                age=data['age'],
                hypertension=data['hypertension'],
                heart_disease=data['heart_disease'],
                ever_married=data['ever_married'],
                work_type=data['work_type'],
                residence_type=data['Residence_type'],
                avg_glucose_level=data['avg_glucose_level'],
                bmi=data['bmi'],
                smoking_status=data['smoking_status'],
                prediction_result=int(prediction),
                prediction_prob=prob
            )
            db.session.add(record)
            db.session.commit()
            
        except Exception as e:
            flash(f'预测出错: {str(e)}')
            
    return render_template('predict.html', result=result)
