predict.py

"""
Predictive modeling for synthesis pool using scikit-learn.
"""
import pandas as pd
import numpy as np
from pathlib import Path
from datetime import datetime, timedelta
from typing import Dict, List, Optional
import joblib
from sklearn.ensemble import RandomForestClassifier
from sklearn.model_selection import train_test_split

from models import PoolHealth, SynthesisType


def sanitize_for_json(obj):
    """Convert numpy/pandas types to native Python types for JSON serialization."""
    if isinstance(obj, dict):
        return {k: sanitize_for_json(v) for k, v in obj.items()}
    elif isinstance(obj, list):
        return [sanitize_for_json(v) for v in obj]
    elif isinstance(obj, (np.integer, np.int64, np.int32)):
        return int(obj)
    elif isinstance(obj, (np.floating, np.float64, np.float32)):
        return float(obj)
    elif isinstance(obj, np.ndarray):
        return obj.tolist()
    elif pd.isna(obj):
        return None
    else:
        return obj


# Use same data directory as operations.py
DATA_DIR = Path(__file__).parent / "data"
MEMORY_MATTER_CSV = DATA_DIR / "memory_matter.csv"
SYNTHESIS_OUTPUTS_CSV = DATA_DIR / "synthesis_outputs.csv"
POOL_STATES_CSV = DATA_DIR / "pool_states.csv"
MODEL_POOL_HEALTH = DATA_DIR / "model_pool_health.pkl"
MODEL_SYNTHESIS_TYPE = DATA_DIR / "model_synthesis_type.pkl"


class PoolPredictor:
    """Predictive modeling for pool health and synthesis types."""

    def __init__(self):
        self.pool_health_model = None
        self.synthesis_type_model = None
        self.pool_health_features = None
        self.synthesis_type_features = None

    def _aggregate_weekly_stats(self, df: pd.DataFrame) -> pd.DataFrame:
        """Aggregate memory_matter by week for pool health prediction."""
        df = df.copy()
        df['timestamp_actual'] = pd.to_datetime(df['timestamp_actual'])
        df['week'] = df['timestamp_actual'].dt.to_period('W')

        weekly_stats = df.groupby('week').agg({
            'toxicity': ['mean', 'max', 'std'],
            'nutrient_density': ['mean', 'min'],
            'id': 'count'
        }).reset_index()

        # Flatten column names
        weekly_stats.columns = ['week', 'toxicity_mean', 'toxicity_max', 'toxicity_std',
                                'nutrient_mean', 'nutrient_min', 'matter_count']

        # Fill NaN std with 0 (happens when only 1 entry in week)
        weekly_stats['toxicity_std'] = weekly_stats['toxicity_std'].fillna(0)

        # Convert week period to timestamp for merging
        weekly_stats['week'] = weekly_stats['week'].dt.to_timestamp()

        return weekly_stats

    def _extract_synthesis_features(self, synthesis_df: pd.DataFrame, matter_df: pd.DataFrame) -> pd.DataFrame:
        """Extract features for synthesis type prediction from source memories."""
        features = []

        for _, synthesis in synthesis_df.iterrows():
            # Get source memory IDs (handle both pipe and comma separators)
            separator = '|' if '|' in synthesis['derived_from'] else ','
            source_ids = [id.strip() for id in synthesis['derived_from'].split(separator)]

            # Get source memories
            source_memories = matter_df[matter_df['id'].isin(source_ids)]

            if len(source_memories) == 0:
                continue

            # Aggregate toxicity/nutrient stats
            feature_dict = {
                'toxicity_mean': source_memories['toxicity'].mean(),
                'toxicity_max': source_memories['toxicity'].max(),
                'toxicity_std': source_memories['toxicity'].std() if len(source_memories) > 1 else 0,
                'nutrient_mean': source_memories['nutrient_density'].mean(),
                'nutrient_max': source_memories['nutrient_density'].max(),
                'nutrient_std': source_memories['nutrient_density'].std() if len(source_memories) > 1 else 0,
                'matter_count': len(source_memories)
            }

            # Layer presence booleans
            all_layers = set()
            for layers_str in source_memories['layers'].dropna():
                all_layers.update(layers_str.split('|'))

            for layer in ['work', 'self', 'relational', 'environmental', 'financial']:
                feature_dict[f'layer_{layer}'] = 1 if layer in all_layers else 0

            # Target
            feature_dict['type'] = synthesis['type']

            features.append(feature_dict)

        return pd.DataFrame(features)

    def train_pool_health(self) -> Dict:
        """Train model to predict pool health from weekly memory stats."""
        # Load data
        matter_df = pd.read_csv(MEMORY_MATTER_CSV)
        pool_states_df = pd.read_csv(POOL_STATES_CSV)

        if len(matter_df) < 5:
            return {
                'error': f'Insufficient data: need at least 5 memory entries, found {len(matter_df)}',
                'success': False
            }

        if len(pool_states_df) < 5:
            return {
                'error': f'Insufficient pool states: need at least 5 snapshots, found {len(pool_states_df)}',
                'success': False
            }

        # Aggregate weekly stats
        weekly_stats = self._aggregate_weekly_stats(matter_df)

        # Merge with pool states
        pool_states_df['snapshot_date'] = pd.to_datetime(pool_states_df['snapshot_date'])

        # Match weekly stats to nearest pool state
        training_data = []
        for _, week_row in weekly_stats.iterrows():
            # Find closest pool state snapshot
            time_diffs = abs(pool_states_df['snapshot_date'] - week_row['week'])
            closest_idx = time_diffs.idxmin()
            closest_state = pool_states_df.loc[closest_idx]

            # Only use if within reasonable time window (7 days)
            if time_diffs[closest_idx].days <= 7:
                training_data.append({
                    **week_row.to_dict(),
                    'health': closest_state['health']
                })

        if len(training_data) < 5:
            return {
                'error': f'Insufficient matched data: need at least 5 week-state pairs, found {len(training_data)}',
                'success': False
            }

        train_df = pd.DataFrame(training_data)

        # Prepare features and target
        feature_cols = ['toxicity_mean', 'toxicity_max', 'toxicity_std',
                       'nutrient_mean', 'nutrient_min', 'matter_count']
        X = train_df[feature_cols].values
        y = train_df['health'].values

        # Train model (use all data if too small for split)
        if len(train_df) < 10:
            # Too small for train/test split - use all data
            X_train, y_train = X, y
            accuracy = None
            print(f"[Training] Using all {len(train_df)} samples (too small for train/test split)")
        else:
            # Split for validation
            X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)
            print(f"[Training] Split: {len(X_train)} train, {len(X_test)} test")

            # Train on training set for validation
            temp_model = RandomForestClassifier(n_estimators=50, random_state=42)
            temp_model.fit(X_train, y_train)
            accuracy = temp_model.score(X_test, y_test)
            print(f"[Training] Validation accuracy: {accuracy:.2%}")

            # Now retrain on ALL data for final model
            X_train, y_train = X, y
            print(f"[Training] Retraining on all {len(train_df)} samples for final model")

        # Train final model on all available data
        self.pool_health_model = RandomForestClassifier(n_estimators=50, random_state=42)
        self.pool_health_model.fit(X_train, y_train)
        self.pool_health_features = feature_cols

        # Show class distribution
        unique, counts = np.unique(y, return_counts=True)
        print(f"[Training] Target distribution: {dict(zip(unique, counts))}")

        # Save model
        joblib.dump({
            'model': self.pool_health_model,
            'features': self.pool_health_features
        }, MODEL_POOL_HEALTH)

        result = {
            'success': True,
            'samples': len(train_df),
            'accuracy': accuracy,
            'features': feature_cols,
            'model_path': str(MODEL_POOL_HEALTH)
        }
        return sanitize_for_json(result)

    def train_synthesis_type(self) -> Dict:
        """Train model to predict synthesis type from memory combinations."""
        # Load data
        matter_df = pd.read_csv(MEMORY_MATTER_CSV)
        synthesis_df = pd.read_csv(SYNTHESIS_OUTPUTS_CSV)

        if len(synthesis_df) < 5:
            return {
                'error': f'Insufficient data: need at least 5 synthesis outputs, found {len(synthesis_df)}',
                'success': False
            }

        # Extract features
        features_df = self._extract_synthesis_features(synthesis_df, matter_df)

        if len(features_df) < 5:
            return {
                'error': f'Insufficient valid data: need at least 5 synthesis-memory pairs, found {len(features_df)}',
                'success': False
            }

        # Prepare features and target
        feature_cols = ['toxicity_mean', 'toxicity_max', 'toxicity_std',
                       'nutrient_mean', 'nutrient_max', 'nutrient_std',
                       'matter_count', 'layer_work', 'layer_self',
                       'layer_relational', 'layer_environmental', 'layer_financial']
        X = features_df[feature_cols].values
        y = features_df['type'].values

        # Train model (use all data if too small for split)
        if len(features_df) < 10:
            # Too small for train/test split - use all data
            X_train, y_train = X, y
            accuracy = None
            print(f"[Training] Using all {len(features_df)} samples (too small for train/test split)")
        else:
            # Split for validation
            X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)
            print(f"[Training] Split: {len(X_train)} train, {len(X_test)} test")

            # Train on training set for validation
            temp_model = RandomForestClassifier(n_estimators=50, random_state=42)
            temp_model.fit(X_train, y_train)
            accuracy = temp_model.score(X_test, y_test)
            print(f"[Training] Validation accuracy: {accuracy:.2%}")

            # Now retrain on ALL data for final model
            X_train, y_train = X, y
            print(f"[Training] Retraining on all {len(features_df)} samples for final model")

        # Train final model on all available data
        self.synthesis_type_model = RandomForestClassifier(n_estimators=50, random_state=42)
        self.synthesis_type_model.fit(X_train, y_train)
        self.synthesis_type_features = feature_cols

        # Show class distribution
        unique, counts = np.unique(y, return_counts=True)
        print(f"[Training] Target distribution: {dict(zip(unique, counts))}")

        # Save model
        joblib.dump({
            'model': self.synthesis_type_model,
            'features': self.synthesis_type_features
        }, MODEL_SYNTHESIS_TYPE)

        result = {
            'success': True,
            'samples': len(features_df),
            'accuracy': accuracy,
            'features': feature_cols,
            'model_path': str(MODEL_SYNTHESIS_TYPE)
        }
        return sanitize_for_json(result)

    def predict_health(self, days_back: int = 7) -> Dict:
        """Predict next week's pool health from recent memory stats."""
        # Load model
        if not MODEL_POOL_HEALTH.exists():
            return {
                'error': 'Pool health model not trained. Run: pool.py train --target pool_health',
                'success': False
            }

        model_data = joblib.load(MODEL_POOL_HEALTH)
        self.pool_health_model = model_data['model']
        self.pool_health_features = model_data['features']

        # Load recent memories
        matter_df = pd.read_csv(MEMORY_MATTER_CSV)
        matter_df['timestamp_actual'] = pd.to_datetime(matter_df['timestamp_actual'])

        cutoff_date = datetime.now() - timedelta(days=days_back)
        recent = matter_df[matter_df['timestamp_actual'] >= cutoff_date]

        if len(recent) == 0:
            return {
                'error': f'No memories found in past {days_back} days',
                'success': False
            }

        # Aggregate stats for recent period
        features = {
            'toxicity_mean': recent['toxicity'].mean(),
            'toxicity_max': recent['toxicity'].max(),
            'toxicity_std': recent['toxicity'].std() if len(recent) > 1 else 0,
            'nutrient_mean': recent['nutrient_density'].mean(),
            'nutrient_min': recent['nutrient_density'].min(),
            'matter_count': len(recent)
        }

        # Predict
        X = np.array([[features[col] for col in self.pool_health_features]])
        prediction = self.pool_health_model.predict(X)[0]
        probabilities = self.pool_health_model.predict_proba(X)[0]

        # Map probabilities to classes (convert numpy types to Python types)
        class_probs = {
            str(cls): float(prob) for cls, prob in zip(self.pool_health_model.classes_, probabilities)
        }

        result = {
            'success': True,
            'prediction': prediction,
            'confidence': max(probabilities),
            'probabilities': class_probs,
            'features': features,
            'days_analyzed': days_back,
            'memory_count': len(recent)
        }
        return sanitize_for_json(result)

    def predict_synthesis(self, matter_ids: List[str]) -> Dict:
        """Predict synthesis type from given memory matter IDs."""
        # Load model
        if not MODEL_SYNTHESIS_TYPE.exists():
            return {
                'error': 'Synthesis type model not trained. Run: pool.py train --target synthesis_type',
                'success': False
            }

        model_data = joblib.load(MODEL_SYNTHESIS_TYPE)
        self.synthesis_type_model = model_data['model']
        self.synthesis_type_features = model_data['features']

        # Load memories
        matter_df = pd.read_csv(MEMORY_MATTER_CSV)
        source_memories = matter_df[matter_df['id'].isin(matter_ids)]

        if len(source_memories) == 0:
            return {
                'error': f'No valid memory IDs found: {matter_ids}',
                'success': False
            }

        # Extract features
        features = {
            'toxicity_mean': source_memories['toxicity'].mean(),
            'toxicity_max': source_memories['toxicity'].max(),
            'toxicity_std': source_memories['toxicity'].std() if len(source_memories) > 1 else 0,
            'nutrient_mean': source_memories['nutrient_density'].mean(),
            'nutrient_max': source_memories['nutrient_density'].max(),
            'nutrient_std': source_memories['nutrient_density'].std() if len(source_memories) > 1 else 0,
            'matter_count': len(source_memories)
        }

        # Layer presence
        all_layers = set()
        for layers_str in source_memories['layers'].dropna():
            all_layers.update(layers_str.split('|'))

        for layer in ['work', 'self', 'relational', 'environmental', 'financial']:
            features[f'layer_{layer}'] = 1 if layer in all_layers else 0

        # Predict
        X = np.array([[features[col] for col in self.synthesis_type_features]])
        prediction = self.synthesis_type_model.predict(X)[0]
        probabilities = self.synthesis_type_model.predict_proba(X)[0]

        # Map probabilities to classes (convert numpy types to Python types)
        class_probs = {
            str(cls): float(prob) for cls, prob in zip(self.synthesis_type_model.classes_, probabilities)
        }

        result = {
            'success': True,
            'prediction': prediction,
            'confidence': max(probabilities),
            'probabilities': class_probs,
            'features': features,
            'matter_ids': matter_ids,
            'found_ids': source_memories['id'].tolist()
        }
        return sanitize_for_json(result)