
"""PKTron Chemistry — Molecule creation, properties, VSEPR."""
import numpy as np
from .elements import molar_mass, get_element

class Molecule:
    """Represents a simple chemical molecule."""

    def __init__(self, name, formula_dict, bonds=None):
        """
        Parameters
        ----------
        name         : str   common name e.g. "Water"
        formula_dict : dict  {symbol: count}  e.g. {"H":2, "O":1}
        bonds        : list[tuple]  optional bond list [(sym1,sym2,order),...]
        """
        self.name         = name
        self.formula_dict = formula_dict
        self.bonds        = bonds or []
        self.molar_mass   = molar_mass(formula_dict)

    def formula(self):
        """Return chemical formula string e.g. H2O"""
        parts = []
        for sym, cnt in sorted(self.formula_dict.items()):
            parts.append(sym + (str(cnt) if cnt > 1 else ""))
        return "".join(parts)

    def n_atoms(self):
        return sum(self.formula_dict.values())

    def n_electrons(self):
        total = 0
        for sym, cnt in self.formula_dict.items():
            total += get_element(sym)["Z"] * cnt
        return total

    def is_polar(self):
        """Crude polarity flag — True if more than one distinct element."""
        return len(self.formula_dict) > 1

    def __repr__(self):
        return f"Molecule({self.formula()}, M={self.molar_mass:.3f} g/mol)"


# Common molecules library
WATER   = Molecule("Water",   {"H":2,"O":1},  [("H","O",1),("O","H",1)])
CO2     = Molecule("Carbon dioxide",{"C":1,"O":2},[("C","O",2),("C","O",2)])
NH3     = Molecule("Ammonia", {"N":1,"H":3},  [("N","H",1)]*3)
CH4     = Molecule("Methane", {"C":1,"H":4},  [("C","H",1)]*4)
O2      = Molecule("Oxygen",  {"O":2},        [("O","O",2)])
N2      = Molecule("Nitrogen",{"N":2},        [("N","N",3)])
GLUCOSE = Molecule("Glucose", {"C":6,"H":12,"O":6})
ETHANOL = Molecule("Ethanol", {"C":2,"H":6,"O":1})
NaCl    = Molecule("Salt",    {"Na":1,"Cl":1},[("Na","Cl",1)])

def molecules_library():
    """Return dict of pre-built common molecules."""
    return {"H2O":WATER,"CO2":CO2,"NH3":NH3,"CH4":CH4,
            "O2":O2,"N2":N2,"glucose":GLUCOSE,
            "ethanol":ETHANOL,"NaCl":NaCl}

def van_der_waals_pressure(n, T, V, a, b, R=8.314):
    """(P + a*n^2/V^2)*(V - n*b) = n*R*T  -> P"""
    import numpy as np
    n,T,V,a,b = map(float,[n,T,V,a,b])
    return (n*R*T/(V-n*b)) - a*n**2/V**2
