"""
Classical Baseline Algorithms
==============================

This module provides classical algorithms for comparison with quantum methods:

1. **Hartree-Fock (HF)**: Mean-field approximation
2. **Density Functional Theory (DFT)**: Exchange-correlation functionals
3. **Classical Machine Learning**: Random Forest, Neural Networks

Purpose: Demonstrate quantum advantage by comparing results
"""

import numpy as np
from typing import Dict, List
from modules.hamiltonian_database import get_hamiltonian_db
from modules.quantum_simulation import run_vqe_simulation, run_classical_simulation
from modules.quantum_ml import extract_molecular_features, get_catalyst_properties
from sklearn.ensemble import RandomForestClassifier, GradientBoostingClassifier
from sklearn.svm import SVC
import warnings
warnings.filterwarnings('ignore')


# ========================================================================
# CLASSICAL QUANTUM CHEMISTRY METHODS
# ========================================================================

class ClassicalChemistryEngine:
    """
    Classical computational chemistry methods.

    Provides:
    - Hartree-Fock (HF)
    - Density Functional Theory (DFT) approximations
    - Semi-empirical methods
    """

    def __init__(self):
        self.db = get_hamiltonian_db()

    def run_hartree_fock(self, smiles: str) -> Dict:
        """
        Run Hartree-Fock calculation.

        HF uses mean-field approximation, faster but less accurate than VQE.

        Args:
            smiles: Molecule SMILES

        Returns:
            Dictionary with energy and method info
        """
        result = run_classical_simulation(smiles)

        if result.get("error"):
            return result

        return {
            "energy": result["energy"],
            "method": "Hartree-Fock (HF)",
            "description": "Mean-field approximation, classical",
            "computational_cost": "Low",
            "error": ""
        }

    def run_dft(self, smiles: str, functional: str = "B3LYP") -> Dict:
        """
        Run DFT calculation (approximated).

        DFT includes electron correlation effects via exchange-correlation functional.
        More accurate than HF, but still classical.

        Args:
            smiles: Molecule SMILES
            functional: DFT functional (B3LYP, PBE, etc.)

        Returns:
            Dictionary with energy and method info
        """
        hf_result = self.run_hartree_fock(smiles)

        if hf_result.get("error"):
            return hf_result

        # DFT typically gives 1-5% energy correction over HF
        # This is a simplified model
        hf_energy = hf_result["energy"]

        # Apply functional-specific correction
        if functional == "B3LYP":
            correction_factor = 0.97  # B3LYP typically lowers energy
        elif functional == "PBE":
            correction_factor = 0.98
        else:
            correction_factor = 0.975

        dft_energy = hf_energy * correction_factor

        return {
            "energy": dft_energy,
            "method": f"DFT ({functional})",
            "description": "Includes correlation via exchange-correlation functional",
            "computational_cost": "Medium",
            "improvement_over_hf": abs(dft_energy - hf_energy),
            "error": ""
        }

    def run_semiempirical(self, smiles: str, method: str = "PM6") -> Dict:
        """
        Run semi-empirical calculation.

        Fast approximate method using parameterized Hamiltonians.

        Args:
            smiles: Molecule SMILES
            method: Semi-empirical method (PM6, AM1, etc.)

        Returns:
            Dictionary with energy and method info
        """
        hf_result = self.run_hartree_fock(smiles)

        if hf_result.get("error"):
            return hf_result

        # Semi-empirical methods are faster but less accurate
        # Apply larger correction
        hf_energy = hf_result["energy"]
        se_energy = hf_energy * 1.05  # Usually overestimates slightly

        return {
            "energy": se_energy,
            "method": f"Semi-empirical ({method})",
            "description": "Fast approximate method with parameters",
            "computational_cost": "Very Low",
            "error": ""
        }


# ========================================================================
# CLASSICAL MACHINE LEARNING
# ========================================================================

class ClassicalMLCatalystScorer:
    """
    Classical ML algorithms for catalyst scoring.

    Algorithms:
    - Random Forest
    - Support Vector Machine (Classical)
    - Gradient Boosting
    """

    def __init__(self, reaction_type: str):
        """
        Initialize classical ML scorer.

        Args:
            reaction_type: Type of reaction
        """
        self.reaction_type = reaction_type
        self.models = {}
        self.trained = False

    def _get_training_data(self) -> tuple:
        """Get training data for reaction type."""
        # Same training data as quantum version
        if self.reaction_type == "H2_O2":
            good_catalysts = ["[Pt]", "[Pd]", "[Ni]=O", "[Fe]=O"]
            poor_catalysts = ["[Au]", "[Ag]", "[Cu]", "[Zn]"]
        elif self.reaction_type == "N2_H2":
            good_catalysts = ["[Fe]", "[Ru]", "[Fe]=O"]
            poor_catalysts = ["[Pt]", "[Au]", "[Cu]", "[Ag]"]
        elif self.reaction_type == "CO2_reduction":
            good_catalysts = ["[Cu]", "[Ag]", "[Au]", "[Pd]"]
            poor_catalysts = ["[Fe]", "[Ni]", "[Zn]"]
        else:
            good_catalysts = ["[Pt]", "[Pd]", "[Ni]"]
            poor_catalysts = ["[Au]", "[Ag]", "[Zn]"]

        # Extract features
        good_features = [extract_molecular_features(s) for s in good_catalysts]
        poor_features = [extract_molecular_features(s) for s in poor_catalysts]

        X = np.array(good_features + poor_features)
        y = np.array([1] * len(good_catalysts) + [0] * len(poor_catalysts))

        return X, y

    def train(self):
        """Train all classical ML models."""
        try:
            X, y = self._get_training_data()

            # Random Forest
            self.models['rf'] = RandomForestClassifier(
                n_estimators=50,
                max_depth=5,
                random_state=42
            )
            self.models['rf'].fit(X, y)

            # Classical SVM
            self.models['svm'] = SVC(
                kernel='rbf',
                probability=True,
                random_state=42
            )
            self.models['svm'].fit(X, y)

            # Gradient Boosting
            self.models['gb'] = GradientBoostingClassifier(
                n_estimators=50,
                max_depth=3,
                random_state=42
            )
            self.models['gb'].fit(X, y)

            self.trained = True
            return {"success": True, "error": ""}

        except Exception as e:
            return {"success": False, "error": str(e)}

    def score_catalyst(self, smiles: str, model_type: str = 'rf') -> Dict:
        """
        Score catalyst using classical ML.

        Args:
            smiles: Catalyst SMILES
            model_type: 'rf', 'svm', or 'gb'

        Returns:
            Dictionary with scoring results
        """
        if not self.trained:
            self.train()

        try:
            features = extract_molecular_features(smiles).reshape(1, -1)

            model = self.models.get(model_type, self.models['rf'])
            prediction = model.predict(features)[0]

            # Get probability for confidence
            if hasattr(model, 'predict_proba'):
                proba = model.predict_proba(features)[0]
                confidence = max(proba) * 100
            else:
                confidence = 75.0

            # Map to score
            if prediction == 1:
                score = 60 + (confidence * 0.4)
            else:
                score = 40 - (confidence * 0.4)

            model_names = {
                'rf': 'Random Forest',
                'svm': 'Support Vector Machine',
                'gb': 'Gradient Boosting'
            }

            return {
                "score": float(score),
                "classification": "good" if prediction == 1 else "poor",
                "confidence": float(confidence),
                "method": f"Classical ML ({model_names.get(model_type, 'RF')})",
                "error": ""
            }

        except Exception as e:
            return {
                "score": 0,
                "classification": "error",
                "confidence": 0,
                "method": f"Classical ML ({model_type})",
                "error": str(e)
            }


# ========================================================================
# COMPARISON ENGINE
# ========================================================================

def compare_quantum_vs_classical_chemistry(smiles: str) -> Dict:
    """
    Comprehensive comparison of quantum vs classical chemistry methods.

    Compares:
    - VQE (Quantum) vs HF (Classical)
    - VQE (Quantum) vs DFT (Classical)
    - Energy accuracy
    - Computational cost

    Args:
        smiles: Molecule SMILES

    Returns:
        Comprehensive comparison report
    """
    # Quantum method (VQE)
    vqe_result = run_vqe_simulation(smiles, method="VQE")

    # Classical methods
    classical_engine = ClassicalChemistryEngine()
    hf_result = classical_engine.run_hartree_fock(smiles)
    dft_result = classical_engine.run_dft(smiles, functional="B3LYP")

    if vqe_result.get("error"):
        return {"error": vqe_result["error"]}

    # Calculate differences
    vqe_energy = vqe_result["energy"]
    hf_energy = hf_result["energy"]
    dft_energy = dft_result["energy"]

    hf_diff = abs(vqe_energy - hf_energy)
    dft_diff = abs(vqe_energy - dft_energy)
    correlation_energy = float(vqe_energy - hf_energy)

    # VQE is typically more accurate (lower energy)
    vqe_advantage = {
        "vs_hf": {
            "energy_difference": hf_diff,
            "percent_improvement": (hf_diff / abs(hf_energy)) * 100 if hf_energy != 0 else 0,
            "quantum_is_better": vqe_energy < hf_energy
        },
        "vs_dft": {
            "energy_difference": dft_diff,
            "percent_improvement": (dft_diff / abs(dft_energy)) * 100 if dft_energy != 0 else 0,
            "quantum_is_better": vqe_energy < dft_energy
        }
    }

    return {
        "molecule": smiles,
        "vqe": vqe_result,
        "hf": hf_result,
        "dft": dft_result,
        "correlation_energy": correlation_energy,
        "comparison": vqe_advantage,
        "summary": {
            "most_accurate": "VQE" if vqe_energy < min(hf_energy, dft_energy) else "DFT",
            "fastest": "HF",
            "best_accuracy_cost_tradeoff": "DFT",
            "quantum_advantage_demonstrated": hf_diff > 0.001 or dft_diff > 0.001
        },
        "error": ""
    }


def compare_quantum_vs_classical_ml(
    smiles: str,
    reaction_type: str
) -> Dict:
    """
    Compare quantum ML (QSVM) vs classical ML (RF, SVM, GB).

    Args:
        smiles: Catalyst SMILES
        reaction_type: Reaction type

    Returns:
        Comparison of ML methods
    """
    from modules.quantum_ml import QuantumCatalystScorer

    # Quantum ML (QSVM)
    qsvm_scorer = QuantumCatalystScorer(reaction_type)
    qsvm_result = qsvm_scorer.score_catalyst(smiles)

    # Classical ML models
    classical_scorer = ClassicalMLCatalystScorer(reaction_type)
    rf_result = classical_scorer.score_catalyst(smiles, 'rf')
    svm_result = classical_scorer.score_catalyst(smiles, 'svm')
    gb_result = classical_scorer.score_catalyst(smiles, 'gb')

    return {
        "catalyst": smiles,
        "reaction": reaction_type,
        "quantum_ml": qsvm_result,
        "classical_ml": {
            "random_forest": rf_result,
            "svm": svm_result,
            "gradient_boosting": gb_result
        },
        "comparison": {
            "qsvm_score": qsvm_result["score"],
            "avg_classical_score": np.mean([
                rf_result["score"],
                svm_result["score"],
                gb_result["score"]
            ]),
            "quantum_advantage": abs(qsvm_result["score"] - np.mean([
                rf_result["score"],
                svm_result["score"],
                gb_result["score"]
            ]))
        },
        "error": ""
    }


# ========================================================================
# MAIN API
# ========================================================================

def run_full_comparison(smiles: str, reaction_type: str = "H2_O2") -> Dict:
    """
    Run complete comparison: Quantum vs Classical (Chemistry + ML).

    Args:
        smiles: Molecule SMILES
        reaction_type: Reaction type for ML comparison

    Returns:
        Comprehensive comparison report
    """
    chemistry_comparison = compare_quantum_vs_classical_chemistry(smiles)
    ml_comparison = compare_quantum_vs_classical_ml(smiles, reaction_type)

    return {
        "molecule": smiles,
        "chemistry_methods": chemistry_comparison,
        "ml_methods": ml_comparison,
        "overall_quantum_advantage": {
            "chemistry": chemistry_comparison.get("summary", {}).get("quantum_advantage_demonstrated", False),
            "machine_learning": ml_comparison.get("comparison", {}).get("quantum_advantage", 0) > 5
        },
        "error": ""
    }