"""
formulation_agent.py
Zein-Protein Nanoparticle Anti-TB Drug Delivery AI Formulation Engine
Uses NumPy and Pandas to simulate, evaluate, and optimize nanoparticle formulations.
"""

import numpy as np
import pandas as pd

class ZeinFormulationEngine:
    def __init__(self):
        # Anti-TB Drugs typical for loading
        self.drugs = ['Rifampicin', 'Isoniazid', 'Pyrazinamide']
        
    def generate_synthetic_experimental_data(self, num_samples=200, seed=42):
        """Generates a historical dataset of formulations for training/evaluation."""
        np.random.seed(seed)
        
        data = {
            'zein_concentration_mg_ml': np.random.uniform(5.0, 30.0, num_samples),
            'surfactant_concentration_w_v': np.random.uniform(0.1, 2.0, num_samples),
            'drug_to_zein_ratio': np.random.uniform(0.05, 0.3, num_samples),
            'ph': np.random.uniform(4.0, 9.0, num_samples),
            'drug_type': np.random.choice(self.drugs, num_samples)
        }
        
        df = pd.DataFrame(data)
        
        # Simulating outputs based on biochemical relationships
        # Particle size (nm): lower is better, usually increases with zein concentration
        df['particle_size_nm'] = (
            120 + 5.5 * df['zein_concentration_mg_ml'] 
            - 15 * df['surfactant_concentration_w_v'] 
            + 40 * (df['ph'] - 7.0)**2 
            + np.random.normal(0, 15, num_samples)
        )
        # Ensure realistic bounds
        df['particle_size_nm'] = df['particle_size_nm'].clip(lower=50, upper=450)
        
        # Polydispersity Index (PDI): lower is better (< 0.3 is good)
        df['pdi'] = (
            0.15 + 0.005 * df['zein_concentration_mg_ml'] 
            + 0.02 * (df['surfactant_concentration_w_v'] - 1.0)**2 
            + np.random.normal(0, 0.04, num_samples)
        )
        df['pdi'] = df['pdi'].clip(lower=0.05, upper=0.6)
        
        # Encapsulation Efficiency (EE%): higher is better
        df['encapsulation_efficiency_pct'] = (
            50 + 1.2 * df['zein_concentration_mg_ml'] 
            - 80 * (df['drug_to_zein_ratio'] - 0.15)**2 
            - 2.5 * (df['ph'] - 6.5)**2 
            + np.random.normal(0, 5, num_samples)
        )
        df['encapsulation_efficiency_pct'] = df['encapsulation_efficiency_pct'].clip(lower=10, upper=98)
        
        return df

    def fitness_function(self, particle_size, pdi, ee, target_size=150, target_pdi=0.2):
        """Calculates a composite score for optimization (higher is better)."""
        # Penalize deviation from target size
        size_penalty = np.abs(particle_size - target_size) / target_size
        # Penalize high PDI
        pdi_penalty = max(0, pdi - target_pdi) / target_pdi
        # Reward high Encapsulation Efficiency
        ee_reward = ee / 100.0
        
        score = ee_reward - 0.5 * size_penalty - 0.5 * pdi_penalty
        return score

    def predict_performance(self, zein_conc, surf_conc, ratio, ph, drug_type):
        """Simplistic heuristic predictive engine imitating an AI response surface model."""
        # Baseline coefficients depending on drug characteristics
        drug_mod = 0 if drug_type == 'Rifampicin' else (5 if drug_type == 'Isoniazid' else -3)
        
        pred_size = 120 + 5.5 * zein_conc - 15 * surf_conc + 40 * (ph - 7.0)**2 + drug_mod
        pred_pdi = 0.15 + 0.005 * zein_conc + 0.02 * (surf_conc - 1.0)**2
        pred_ee = 50 + 1.2 * zein_conc - 80 * (ratio - 0.15)**2 - 2.5 * (ph - 6.5)**2 - drug_mod
        
        return float(np.clip(pred_size, 50, 450)), float(np.clip(pred_pdi, 0.05, 0.6)), float(np.clip(pred_ee, 10, 98))

    def optimize_formulation(self, drug_type, target_size=150, steps=30):
        """Grid-search optimization agent using numpy array grids to locate best parameters."""
        zein_range = np.linspace(5.0, 30.0, steps)
        surf_range = np.linspace(0.1, 2.0, steps)
        ratio_range = np.linspace(0.05, 0.3, steps)
        ph_range = np.linspace(4.0, 9.0, steps)
        
        best_score = -float('inf')
        best_params = {}
        
        # Monte Carlo random grid search optimization across the continuous landscape
        for _ in range(5000):
            z = np.random.choice(zein_range)
            s = np.random.choice(surf_range)
            r = np.random.choice(ratio_range)
            p = np.random.choice(ph_range)
            
            size, pdi, ee = self.predict_performance(z, s, r, p, drug_type)
            score = self.fitness_function(size, pdi, ee, target_size)
            
            if score > best_score:
                best_score = score
                best_params = {
                    'zein_concentration_mg_ml': z,
                    'surfactant_concentration_w_v': s,
                    'drug_to_zein_ratio': r,
                    'ph': p,
                    'predicted_size_nm': size,
                    'predicted_pdi': pdi,
                    'predicted_ee_pct': ee,
                    'optimization_score': score
                }
                
        return best_params

if __name__ == "__main__":
    print("="*70)
    print("Initializing Zein-Protein Nanoparticle Anti-TB AI Formulation Engine...")
    print("="*70)
    
    engine = ZeinFormulationEngine()
    
    # 1. Generate historical experimental log
    df_history = engine.generate_synthetic_experimental_data()
    print(f"\nSuccessfully generated simulated historical log with {len(df_history)} runs.")
    print("Sample data preview:")
    print(df_history.head(3).to_string(index=False))
    
    # 2. Run optimization routine for Rifampicin delivery systems
    print("\nOptimizing formulation parameters for [Rifampicin] delivery...")
    optimal_rifampicin = engine.optimize_formulation(drug_type='Rifampicin', target_size=140)
    
    print("\nTop Optimized Formulation Parameters Found:")
    for key, val in optimal_rifampicin.items():
        print(f" - {key.replace('_', ' ').title()}: {val:.4f}" if isinstance(val, float) else f" - {key}: {val}")
        
    print("\nThis is for informational purposes only. For medical advice or diagnosis, consult a professional. AI responses may include mistakes.")