core/evolution/evolution.py

from .population import Population
from .molecule import Molecule
from core.predictors.pure_component.property_predictor import PropertyPredictor
from core.config import EvolutionConfig
from crem.crem import mutate_mol
from rdkit import Chem
import pandas as pd
import numpy as np
import random
from typing import List, Tuple, Dict
from core.data_prep import df
from pathlib import Path
from joblib import Parallel, delayed
from .anomaly_detection import EnsembleUncertaintyFilter

class MolecularEvolution:
    """Main evolutionary algorithm coordinator with uncertainty filtering."""
    BASE_DIR = Path(__file__).resolve().parent.parent.parent
    REP_DB_PATH = BASE_DIR / "data" / "fragments" / "diesel_fragments.db"

    # Calibrated uncertainty thresholds — shared across all instances in the same process
    _cached_uncertainty_thresholds: dict = {}

    def __init__(self, config: EvolutionConfig):
        self.config = config
        self.predictor = PropertyPredictor(config)
        self.population = Population(config)

        # Initialize and calibrate uncertainty filters
        self._setup_uncertainty_filters()
    
    def _setup_uncertainty_filters(self):
        """Setup and calibrate uncertainty filters for CN and YSI."""
        from core.shared_features import featurize_df

        cache_key = (
            id(self.predictor.predictors.get('cn')),
            id(self.predictor.predictors.get('ysi')),
            self.config.minimize_ysi,
        )
        if cache_key in MolecularEvolution._cached_uncertainty_thresholds:
            self.uncertainty_filters = MolecularEvolution._cached_uncertainty_thresholds[cache_key]
            print("✓ Uncertainty filters loaded from cache\n")
            return

        print("\nCalibrating ensemble uncertainty filters...")
        
        val_df = df.sample(200, random_state=42)  # Sample DataFrame, not just SMILES
        val_smiles = val_df["SMILES"].tolist()    # Extract SMILES
        y_val_cn = val_df["cn"].values

        X_val = featurize_df(val_smiles, return_df=False)
        
        if X_val is None:
            print("⚠ Warning: Could not featurize validation set, uncertainty filtering disabled")
            self.uncertainty_filters = {}
            return
        
        self.uncertainty_filters = {}
        
        # Calibrate CN uncertainty filter
        if 'cn' in self.predictor.predictors:
            print("\n  Calibrating CN uncertainty filter...")
            
            
            # Get predictions with uncertainty
            cn_predictor = self.predictor.predictors['cn']
            X_cn = cn_predictor.selector.transform(X_val)
            
            all_tree_preds = np.array(
                Parallel(n_jobs=-1, prefer="threads")(
                    delayed(tree.predict)(X_cn)
                    for tree in cn_predictor.model.estimators_
                )
            )
            cn_variances = np.std(all_tree_preds, axis=0)
            
            # Calculate errors for analysis
            cn_preds = np.mean(all_tree_preds, axis=0)
            cn_errors = np.abs(cn_preds - y_val_cn)
            
            print(f"    Mean prediction error: {cn_errors.mean():.3f}")
            print(f"    Mean variance: {cn_variances.mean():.4f}")
            
            # Calibrate filter
            self.uncertainty_filters['cn'] = EnsembleUncertaintyFilter(percentile_threshold=90)
            self.uncertainty_filters['cn'].calibrate(cn_variances)
            
            # Show effectiveness
            filtered_mask = cn_variances > self.uncertainty_filters['cn'].variance_threshold
            if filtered_mask.sum() > 0:
                print(f"    Will filter: {filtered_mask.sum()}/{len(cn_variances)} ({filtered_mask.mean()*100:.1f}%)")
                print(f"    Filtered have error: {cn_errors[filtered_mask].mean():.3f}")
                print(f"    Kept have error: {cn_errors[~filtered_mask].mean():.3f}")
        
        # Calibrate YSI uncertainty filter (if used)
        if self.config.minimize_ysi and 'ysi' in self.predictor.predictors:
            print("\n  Calibrating YSI uncertainty filter...")
            y_val_ysi = val_df["YSI_Unified_Measured"].values

            
            # Get predictions with uncertainty
            ysi_predictor = self.predictor.predictors['ysi']
            X_ysi = ysi_predictor.selector.transform(X_val)
            
            all_tree_preds = np.array(
                Parallel(n_jobs=-1, prefer="threads")(
                    delayed(tree.predict)(X_ysi)
                    for tree in ysi_predictor.model.estimators_
                )
            )
            ysi_variances = np.std(all_tree_preds, axis=0)
            
            # Calculate errors for analysis
            ysi_preds = np.mean(all_tree_preds, axis=0)
            ysi_errors = np.abs(ysi_preds - y_val_ysi)
            
            print(f"    Mean prediction error: {ysi_errors.mean():.3f}")
            print(f"    Mean variance: {ysi_variances.mean():.4f}")
            
            # Calibrate filter
            self.uncertainty_filters['ysi'] = EnsembleUncertaintyFilter(percentile_threshold=90)
            self.uncertainty_filters['ysi'].calibrate(ysi_variances)
            
            # Show effectiveness
            filtered_mask = ysi_variances > self.uncertainty_filters['ysi'].variance_threshold
            if filtered_mask.sum() > 0:
                print(f"    Will filter: {filtered_mask.sum()}/{len(ysi_variances)} ({filtered_mask.mean()*100:.1f}%)")
                print(f"    Filtered have error: {ysi_errors[filtered_mask].mean():.3f}")
                print(f"    Kept have error: {ysi_errors[~filtered_mask].mean():.3f}")
        
        MolecularEvolution._cached_uncertainty_thresholds[cache_key] = self.uncertainty_filters
        print("\n✓ Uncertainty filtering ready\n")

    def _mutate_molecule(self, mol: Chem.Mol) -> List[str]:
        """Generate mutations for a molecule using CREM."""
        try:
            mutants = list(mutate_mol(
                mol,
                db_name=str(self.REP_DB_PATH),
                max_size=self.config.max_size,
                max_replacements=100,
                min_freq=self.config.min_freq,
                return_mol=False
            ))
            return [m for m in mutants if m and m not in self.population.seen_smiles]
        except Exception:
            return []
    
    def _predict_with_uncertainty(self, smiles_list: List[str]) -> Tuple[Dict, Dict]:
        """
        Predict all properties with uncertainty estimates.
        
        Returns:
            predictions: Dict of property predictions
            uncertainties: Dict of uncertainty estimates (std across trees)
        """
        from core.shared_features import featurize_df
        
        if not smiles_list:
            return {}, {}
        
        # Featurize once
        X_full = featurize_df(smiles_list, return_df=False)
        
        if X_full is None:
            return {}, {}
        
        predictions = {}
        uncertainties = {}
        
        # Predict each property with uncertainty
        for prop_name, predictor_obj in self.predictor.predictors.items():
            X_selected = predictor_obj.selector.transform(X_full)

            # Thread-parallel over trees (sklearn trees release the GIL in C)
            all_tree_preds = np.array(
                Parallel(n_jobs=-1, prefer="threads")(
                    delayed(tree.predict)(X_selected)
                    for tree in predictor_obj.model.estimators_
                )
            )

            mean_preds = np.mean(all_tree_preds, axis=0)
            if predictor_obj.uses_log_transform:
                mean_preds = np.power(10.0, mean_preds)
            predictions[prop_name] = mean_preds
            uncertainties[prop_name] = np.std(all_tree_preds, axis=0)
        
        # Add Tanimoto (no uncertainty)
        predictions["tanimoto"] = self.predictor.compute_tanimoto(smiles_list)
        
        return predictions, uncertainties
    
    def _create_molecules(self, smiles_list: List[str]) -> Tuple[List[Molecule], Dict]:
        """Create Molecule objects with UNCERTAINTY FILTERING."""
        if not smiles_list:
            return [], {}
        
        # Track filtering stats
        filter_stats = {
            'total': len(smiles_list),
            'cn_none': 0,
            'ysi_none': 0,
            'tanimoto_fail': 0,
            'cn_uncertainty_fail': 0,
            'ysi_uncertainty_fail': 0,
            'passed': 0
        }
        
        # Predict with uncertainty
        predictions, uncertainties = self._predict_with_uncertainty(smiles_list)
        
        molecules = []
        
        for i, smiles in enumerate(smiles_list):
            # Extract predictions
            props = {k: predictions[k][i] for k in predictions}

            uncerts = {k: v[i] for k, v in uncertainties.items() if k in uncertainties}
            
            # Validate CN exists
            if props.get('cn') is None or not np.isfinite(props['cn']):
                filter_stats['cn_none'] += 1
                continue
            
            # NEW: Filter by CN uncertainty
            if 'cn' in self.uncertainty_filters and 'cn' in uncerts:
                cn_variance = uncerts['cn']
                is_reliable = self.uncertainty_filters['cn'].is_reliable(
                    np.array([cn_variance])
                )[0]
                
                if not is_reliable:
                    filter_stats['cn_uncertainty_fail'] += 1
                    continue
            
            # YSI validation
            if self.config.minimize_ysi:
                if props.get('ysi') is None or not np.isfinite(props['ysi']):
                    filter_stats['ysi_none'] += 1
                    continue
                
                # NEW: Filter by YSI uncertainty
                if 'ysi' in self.uncertainty_filters and 'ysi' in uncerts:
                    ysi_variance = uncerts['ysi']
                    is_reliable = self.uncertainty_filters['ysi'].is_reliable(
                        np.array([ysi_variance])
                    )[0]
                    
                    if not is_reliable:
                        filter_stats['ysi_uncertainty_fail'] += 1
                        continue
            
            # Tanimoto filter
            tanimoto = props.get("tanimoto")
            if tanimoto is None or tanimoto < 0.7:
                filter_stats['tanimoto_fail'] += 1
                continue

            filter_stats['passed'] += 1
            
            mol = Molecule(
            smiles=smiles,
            cn=float(props['cn']),
            cn_error=abs(float(props['cn']) - self.config.target_cn),
            cn_score=0.0,  # placeholder, set below
            bp=float(props.get('bp')) if props.get('bp') is not None else None,
            ysi=float(props.get('ysi')) if props.get('ysi') is not None else None,
            density=float(props.get('density')) if props.get('density') is not None else None,
            lhv=float(props.get('lhv')) if props.get('lhv') is not None else None,
            dynamic_viscosity=float(props.get('dynamic_viscosity')) if props.get('dynamic_viscosity') is not None else None
        )
            mol.cn_score = mol.fitness(self.config)
            molecules.append(mol)
        
        return molecules, filter_stats
    
    def initialize_population(self, initial_smiles: List[str]) -> int:
        """Initialize the population from initial SMILES."""
        print("Predicting properties for initial population...")
        molecules, filter_stats = self._create_molecules(initial_smiles)
        
        # Log initial filtering
        if filter_stats['total'] > 0:
            print(f"\n  Initial filtering: {filter_stats['total']} → {filter_stats['passed']} passed")
            print(f"    Tanimoto: {filter_stats['tanimoto_fail']} | "
                  f"CN Uncertainty: {filter_stats['cn_uncertainty_fail']} | "
                  f"YSI Uncertainty: {filter_stats['ysi_uncertainty_fail']}")
        
        return self.population.add_molecules(molecules)
    
    def _log_generation_stats(self, generation: int):
        """Log statistics for the current generation."""
        mols = self.population.molecules
        if not mols:
            print(f"Gen {generation}: Population empty after filtering.")
            return

        
        if self.config.maximize_cn:
            best_cn = max(mols, key=lambda m: m.cn)
            avg_cn = np.mean([m.cn for m in mols])
            
            print_msg = (f"Gen {generation}/{self.config.generations} | "
                        f"Pop {len(mols)} | "
                        f"Best CN: {best_cn.cn:.3f} | "
                        f"Avg CN: {avg_cn:.3f}")
        else:
            best_cn = min(mols, key=lambda m: m.cn_error)
            avg_cn_err = np.mean([m.cn_error for m in mols])
            
            print_msg = (f"Gen {generation}/{self.config.generations} | "
                        f"Pop {len(mols)} | "
                        f"Best CN err: {best_cn.cn_error:.3f} | "
                        f"Avg CN err: {avg_cn_err:.3f}")
        
        if self.config.minimize_ysi:
            front = self.population.pareto_front()
            best_ysi = min(mols, key=lambda m: m.ysi)
            avg_ysi = np.mean([m.ysi for m in mols])
            
            print_msg += (f" | Best YSI: {best_ysi.ysi:.3f} | "
                         f"Avg YSI: {avg_ysi:.3f} | "
                         f"Pareto: {len(front)}")
        
        print(print_msg)
    
    def _generate_offspring(self, survivors: List[Molecule]) -> Tuple[List[Molecule], Dict]:
        """Generates offspring from survivors with filtering stats."""
        target_count = self.config.population_size - len(survivors)
        max_attempts = target_count * self.config.max_offspring_attempts
        
        all_children = []
        new_molecules = []
        cumulative_stats = {
            'total': 0,
            'cn_none': 0,
            'ysi_none': 0,
            'tanimoto_fail': 0,
            'cn_uncertainty_fail': 0,
            'ysi_uncertainty_fail': 0,
            'passed': 0
        }
        
        print(f"  → Generating offspring (target: {target_count})...")
        
        for attempt in range(max_attempts):
            if len(new_molecules) >= target_count:
                break
            
            # Generate mutations
            # AFTER — fitness-proportionate (roulette wheel) selection
            weights = np.array([m.fitness(self.config) for m in survivors])
            weights /= weights.sum()
            parent = survivors[np.random.choice(len(survivors), p=weights)]
            mol = Chem.MolFromSmiles(parent.smiles)
            if mol is None:
                continue
            
            children = self._mutate_molecule(mol)
            all_children.extend(children[:self.config.mutations_per_parent])
            
            # Process in batches
            if len(all_children) >= self.config.batch_size:
                batch_mols, batch_stats = self._create_molecules(all_children)
                new_molecules.extend(batch_mols)
                
                # Accumulate stats
                for key in cumulative_stats:
                    cumulative_stats[key] += batch_stats[key]
                
                all_children = []
        
        # Process remaining children
        if all_children:
            batch_mols, batch_stats = self._create_molecules(all_children)
            new_molecules.extend(batch_mols)
            
            for key in cumulative_stats:
                cumulative_stats[key] += batch_stats[key]
        
        print(f"  ✓ Generated {len(new_molecules)} valid offspring")
        print(f"    Filtering: {cumulative_stats['total']} → {cumulative_stats['passed']} | "
              f"Tanimoto: {cumulative_stats['tanimoto_fail']} | "
              f"CN Unc: {cumulative_stats['cn_uncertainty_fail']} | "
              f"YSI Unc: {cumulative_stats['ysi_uncertainty_fail']} "
              f"(property constraints applied at end)")
        
        return new_molecules, cumulative_stats
    
    def _run_evolution_loop(self):
        """Run the main evolution loop."""
        for gen in range(1, self.config.generations + 1):
            self._log_generation_stats(gen)
            
            survivors = self.population.get_survivors()
            offspring, _ = self._generate_offspring(survivors)
            
            # Create new population
            new_pop = Population(self.config)
            new_pop.add_molecules(survivors + offspring)
            self.population = new_pop
    
    def _apply_property_filters(self, df: pd.DataFrame) -> pd.DataFrame:
        """Apply bp, density, lhv, dynamic_viscosity constraints to a DataFrame."""
        mask = pd.Series(True, index=df.index)
        for prop, (lo, hi) in self.config.filters.items():
            if prop not in df.columns:
                continue
            col = df[prop]
            if lo is not None:
                mask &= col >= lo
            if hi is not None:
                mask &= col <= hi
        filtered = df[mask]
        removed = len(df) - len(filtered)
        if removed > 0:
            print(f"  Property filters removed {removed}/{len(df)} molecules")
        return filtered

    def _sort_df(self, df: pd.DataFrame) -> pd.DataFrame:
        if self.config.maximize_cn:
            if self.config.minimize_ysi and "ysi" in df.columns:
                candidates = df[df["cn"] > 30] if len(df[df["cn"] > 50]) < 10 else df[df["cn"] > 50]
                return candidates.sort_values(["cn", "ysi"], ascending=[False, True])
            else:
                candidates = df[df["cn"] > 30] if len(df[df["cn"] > 50]) < 10 else df[df["cn"] > 50]
                return candidates.sort_values("cn", ascending=False)
        else:
            if self.config.minimize_ysi and "ysi" in df.columns:
                return df[df["cn_error"] < 15].sort_values(["cn_error", "ysi"], ascending=True)
            else:
                return df[df["cn_error"] < 5].sort_values("cn_error", ascending=True)

    def _generate_results(self) -> Tuple[pd.DataFrame, pd.DataFrame, pd.DataFrame]:
        """Generate final results DataFrames (filtered, unfiltered, pareto)."""
        raw_df = self.population.to_dataframe()

        # Unfiltered: only CN/YSI objectives applied, no property constraints
        unfiltered_df = self._sort_df(raw_df).copy()
        unfiltered_df["rank"] = range(1, len(unfiltered_df) + 1)

        # Filtered: property constraints (bp, density, lhv, dynamic_viscosity) applied first
        final_df = self._sort_df(self._apply_property_filters(raw_df)).copy()
        final_df["rank"] = range(1, len(final_df) + 1)

        if self.config.minimize_ysi:
            pareto_mols = self.population.pareto_front()

            pareto_df = pd.DataFrame([m.to_dict() for m in pareto_mols])
            if not pareto_df.empty:
                pareto_df = self._sort_df(pareto_df).copy()
                pareto_df.insert(0, 'rank', range(1, len(pareto_df) + 1))
        else:
            pareto_df = pd.DataFrame()

        return final_df, pareto_df, unfiltered_df
    
    def evolve(self) -> Tuple[pd.DataFrame, pd.DataFrame]:
        """Run the evolutionary algorithm."""
        # Initialize
        df_bins = pd.qcut(df["cn"], q=30)
        initial_smiles = (
            df.groupby(df_bins, observed=False)
            .apply(lambda x: x.sample(20, random_state=42))
            .reset_index(drop=True)["SMILES"]
            .tolist()
        )
        init_count = self.initialize_population(initial_smiles)

        if init_count == 0:
            print("No valid initial molecules")
            return pd.DataFrame(), pd.DataFrame(), pd.DataFrame()

        print(f"✓ Initial population size: {init_count}\n")

        # Evolution
        self._run_evolution_loop()

        # Results
        return self._generate_results()