from __future__ import annotations

from typing import Optional

from chemrecon.chem.elements import Element
from chemrecon.chem import elements as chem


class SumFormula:
    """ Encapsulates a molecular formula.
    """
    formula: dict[Element, int]
    charge: Optional[int]

    def __init__(self, formula: dict[Element, int], charge: Optional[int] = 0):
        """ Assumes 0 (neutral charge) rather than None (unknown). If charge is unknown, specify explicitly. """
        self.formula = formula
        self.charge = charge

    def __getitem__(self, element: Element):
        return self.formula[element]

    def __iter__(self):
        yield from self.formula.__iter__()

    def __or__(self, other: SumFormula):
        return self.formula | other.formula

    def get(self, element: Element, default: int):
        return self.formula.get(element, default)

    def elements(self) -> set[Element]:
        return set(self.formula.keys())

    def is_zero(self) -> bool:
        """ Returns true if all elements are 0, and charge is 0 or unknown.
        """
        if any(i != 0 for _, i in self.formula.items()):
            return False
        if self.charge is not None and self.charge != 0:
            return False
        return True

    def has_negative(self) -> bool:
        """ Returns true if the formula has any negative atom counts (e.g. when it represents a difference).
        """
        return any(v < 0 for v in self.formula.values())

    # Arithmetic operations
    def __add__(self, other: SumFormula):
        """ Add the elements and charges. If at least one has undefined charge, the result has undefined charge.
        """
        sumformula = {
            e: self.get(e, 0) + other.get(e, 0)
            for e in self.elements() | other.elements()
        }
        if self.charge is not None and other.charge is not None:
            charge: Optional[int] = self.charge + other.charge
        else:
            charge: Optional[int] = None

        return SumFormula(sumformula, charge)

    def __sub__(self, other: SumFormula):
        """ Computes the difference between sum formulae
        """
        sumformula = {
            e: self.get(e, 0) - other.get(e, 0)
            for e in self.elements() | other.elements()
        }
        if self.charge is not None and other.charge is not None:
            charge: Optional[int] = self.charge + other.charge
        else:
            charge: Optional[int] = None

        return SumFormula(sumformula, charge)

    def __eq__(self, other: SumFormula):
        # TODO how to handle equality when one has charge 0 and other has unknown charge?
        return self.formula == other.formula and self.charge == other.charge

    # Misc
    # ------------------------------------------------------------------------------------------------------------------
    def __str__(self):
        """ Prints output in Hill order
        """
        # https://en.wikipedia.org/wiki/Chemical_formula#Condensed_formula
        elist: list[tuple[Element, int]] = list()
        if chem.C in self.elements():
            # C, H first order
            elist.append((chem.C, self[chem.C]))
            if chem.H in self.elements():
                elist.append((chem.H, self[chem.H]))
            elist.extend([
                (e, i)
                for e, i in sorted(self.formula.items(), key = lambda pair: pair[0].symbol)
                if e not in {chem.C, chem.H}
            ])
        else:
            # Alphabetical order
            elist = [
                (e, i)
                for e, i in sorted(self.formula.items(), key = lambda pair: pair[0].symbol)
            ]

        chargestr: str = ''
        if self.charge is not None:
            if self.charge > 0:
                chargestr = f' {self.charge}+'
            elif self.charge < 0:
                chargestr = f' {abs(self.charge)}-'

        # Print as string
        return ''.join(f'{e.symbol}{i}' for e, i in elist) + f'{chargestr}'

def molformula_from_str(s: str) -> SumFormula:
    """ Convert a molecular formula string to an object representation.
    """
    # TODO
    raise NotImplementedError(s)
