import numpy as np
import pandas as pd
from sklearn.model_selection import train_test_split, cross_val_score
from sklearn.preprocessing import StandardScaler
from sklearn.ensemble import RandomForestClassifier
from sklearn.metrics import roc_auc_score, classification_report
import matplotlib.pyplot as plt
from sklearn.metrics import roc_curve, precision_recall_curve

class DiseasePredictor:
    """遺伝的変異データからの疾患リスク予測"""
    
    def __init__(self, model_type='random_forest'):
        self.model_type = model_type
        self.model = None
        self.scaler = StandardScaler()
        self.feature_importance = None
    
    def prepare_features(self, genotype_data, clinical_data=None):
        """
        特徴量の準備
        
        Args:
            genotype_data: SNPデータ (n_samples, n_snps)
            clinical_data: 臨床データ (n_samples, n_clinical_features)
        
        Returns:
            np.array: 統合特徴量行列
        """
        features = []
        
        # SNP特徴量（0,1,2エンコーディング）
        features.append(genotype_data)
        
        # 遺伝子間相互作用特徴量
        n_snps = genotype_data.shape[1]
        interactions = []
        for i in range(min(n_snps, 100)):  # 計算量を考慮して制限
            for j in range(i+1, min(n_snps, 100)):
                interactions.append(genotype_data[:, i] * genotype_data[:, j])
        
        if interactions:
            features.append(np.column_stack(interactions))
        
        # 臨床データがある場合は追加
        if clinical_data is not None:
            features.append(clinical_data)
        
        return np.hstack(features)
    
    def train(self, X, y, validation_split=0.2):
        """
        モデルの訓練
        
        Args:
            X: 特徴量行列
            y: ラベル（0: 健常, 1: 疾患）
            validation_split: 検証データの割合
        
        Returns:
            dict: 訓練結果の統計
        """
        # データ分割
        X_train, X_val, y_train, y_val = train_test_split(
            X, y, test_size=validation_split, stratify=y, random_state=42
        )
        
        # 特徴量の標準化
        X_train_scaled = self.scaler.fit_transform(X_train)
        X_val_scaled = self.scaler.transform(X_val)
        
        # モデルの訓練
        if self.model_type == 'random_forest':
            self.model = RandomForestClassifier(
                n_estimators=500,
                max_depth=10,
                min_samples_split=5,
                min_samples_leaf=2,
                class_weight='balanced',
                random_state=42,
                n_jobs=-1
            )
        
        self.model.fit(X_train_scaled, y_train)
        
        # 特徴量重要度の取得
        if hasattr(self.model, 'feature_importances_'):
            self.feature_importance = self.model.feature_importances_
        
        # 性能評価
        train_pred = self.model.predict_proba(X_train_scaled)[:, 1]
        val_pred = self.model.predict_proba(X_val_scaled)[:, 1]
        
        train_auc = roc_auc_score(y_train, train_pred)
        val_auc = roc_auc_score(y_val, val_pred)
        
        # 交差検証
        cv_scores = cross_val_score(
            self.model, X_train_scaled, y_train, cv=5, scoring='roc_auc'
        )
        
        return {
            'train_auc': train_auc,
            'val_auc': val_auc,
            'cv_mean': cv_scores.mean(),
            'cv_std': cv_scores.std(),
            'n_features': X.shape[1]
        }
    
    def plot_performance(self, X_test, y_test):
        """性能評価の可視化"""
        X_scaled = self.scaler.transform(X_test)
        y_pred_proba = self.model.predict_proba(X_scaled)[:, 1]
        
        fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(12, 5))
        
        # ROC曲線
        fpr, tpr, _ = roc_curve(y_test, y_pred_proba)
        auc = roc_auc_score(y_test, y_pred_proba)
        
        ax1.plot(fpr, tpr, label=f'ROC curve (AUC = {auc:.3f})')
        ax1.plot([0, 1], [0, 1], 'k--', label='Random')
        ax1.set_xlabel('False Positive Rate')
        ax1.set_ylabel('True Positive Rate')
        ax1.set_title('ROC Curve')
        ax1.legend()
        ax1.grid(True)
        
        # Precision-Recall曲線
        precision, recall, _ = precision_recall_curve(y_test, y_pred_proba)
        
        ax2.plot(recall, precision)
        ax2.set_xlabel('Recall')
        ax2.set_ylabel('Precision')
        ax2.set_title('Precision-Recall Curve')
        ax2.grid(True)
        
        plt.tight_layout()
        plt.show()