from src.api import MaterialsProjectAdapter
from mendeleev import element

class UniversalMaterialsOptimizer:
    def __init__(self, host_formula: str, host_element_symbol: str, api_key: str = None):
        self.adapter = MaterialsProjectAdapter(api_key)
        self.props = self.adapter.get_material_properties(host_formula)
        self.base_voltage = self.props['band_gap']
        
        host = element(host_element_symbol)
        self.host_radius = host.ionic_radii[0].ionic_radius
        self.host_en = host.en_pauling
        self.host_symbol = host_element_symbol

    def analyze_dopant(self, dopant_symbol: str, concentration: float):
        dopant = element(dopant_symbol)
        
        # 1. Physics Calculations
        delta_en = dopant.en_pauling - self.host_en
        # Simple ionic radius lookup (using first available radius for simplicity)
        dop_rad = dopant.ionic_radii[0].ionic_radius if dopant.ionic_radii else self.host_radius
        radius_diff = dop_rad - self.host_radius
        
        voltage_gain = 1.5 * concentration * delta_en
        predicted_voltage = self.base_voltage + voltage_gain
        strain_energy = (radius_diff ** 2) * concentration * 100

        # 2. Verdict Logic
        status = "Stable"
        if strain_energy > 500: status = "Critical Strain - Phase Separation"
        elif predicted_voltage < 1.0: status = "Voltage Collapse"

        return {
            "dopant": dopant_symbol,
            "concentration": concentration,
            "predicted_voltage": round(predicted_voltage, 3),
            "lattice_strain": round(strain_energy, 2),
            "stability_status": status,
            "data_confidence": "Low (Estimated)" if self.props['is_estimated'] else "High (Real Data)"
        }

    def analyze_mixture(self, dopant_recipe: dict):
        """
        Analyzes a mixture of dopants (Co-doping).
        Input: {"Cl": 0.1, "Br": 0.1} -> implies 0.2 total doping
        """
        total_conc = sum(dopant_recipe.values())
        
        # 1. Edge Case: Empty Recipe
        if total_conc == 0:
            return {"error": "Recipe is empty. Please add at least one dopant."}

        # 2. Calculate Weighted Properties (Vegard's Law)
        weighted_radius = 0.0
        weighted_en = 0.0
        
        detailed_breakdown = []

        for symbol, amount in dopant_recipe.items():
            atom = element(symbol)
            
            # Get radius (default to host radius if data missing to prevent crash)
            r = atom.ionic_radii[0].ionic_radius if atom.ionic_radii else self.host_radius
            en = atom.en_pauling if atom.en_pauling else self.host_en
            
            # Contribution to the average
            fraction = amount / total_conc
            weighted_radius += r * fraction
            weighted_en += en * fraction
            
            detailed_breakdown.append(f"{symbol} ({amount*100:.1f}%)")

        # 3. Compare 'Virtual Dopant' vs Host
        radius_diff = weighted_radius - self.host_radius
        delta_en = weighted_en - self.host_en
        
        # 4. Physics Engine
        # Voltage is boosted by the total amount of dopant * average electronegativity gain
        voltage_gain = 1.5 * total_conc * delta_en
        predicted_voltage = self.base_voltage + voltage_gain
        
        # Strain is proportional to the mismatch squared * total concentration
        strain_energy = (radius_diff ** 2) * total_conc * 100

        # 5. Verdict
        status = "Stable"
        if strain_energy > 500: status = "Critical Strain - Phase Separation"
        elif predicted_voltage < 1.0: status = "Voltage Collapse"

        return {
            "recipe_description": ", ".join(detailed_breakdown),
            "total_doping_level": round(total_conc, 3),
            "predicted_voltage": round(predicted_voltage, 3),
            "lattice_strain": round(strain_energy, 2),
            "stability_status": status,
            "data_confidence": "Low (Estimated)" if self.props['is_estimated'] else "High (Real Data)"
        }