src/data/taxonomy.py

"""
Taxonomic Hierarchy Engine for Western Australia Flora.
Builds and maps Species -> Genus -> Family for hierarchical multi-task classification.
"""
import pandas as pd
from pathlib import Path

# Common WA Flora Genus -> Family mapping dictionary
GENUS_TO_FAMILY = {
    # Proteaceae
    "Banksia": "Proteaceae", "Grevillea": "Proteaceae", "Hakea": "Proteaceae",
    "Dryandra": "Proteaceae", "Adenanthos": "Proteaceae", "Isopogon": "Proteaceae",
    "Petrophile": "Proteaceae", "Synaphea": "Proteaceae", "Conospermum": "Proteaceae",
    "Lambertia": "Proteaceae", "Persoonia": "Proteaceae", "Xylomelum": "Proteaceae",
    "Franklandia": "Proteaceae", "Stirlingia": "Proteaceae",
    
    # Myrtaceae
    "Eucalyptus": "Myrtaceae", "Corymbia": "Myrtaceae", "Melaleuca": "Myrtaceae",
    "Verticordia": "Myrtaceae", "Calothamnus": "Myrtaceae", "Darwinia": "Myrtaceae",
    "Chamelaucium": "Myrtaceae", "Thryptomene": "Myrtaceae", "Beaufortia": "Myrtaceae",
    "Baeckea": "Myrtaceae", "Regelia": "Myrtaceae", "Calytrix": "Myrtaceae",
    "Eremaea": "Myrtaceae", "Kunzea": "Myrtaceae", "Micromyrtus": "Myrtaceae",
    "Hypocalymma": "Myrtaceae", "Scholtzia": "Myrtaceae", "Phymatocarpus": "Myrtaceae",
    "Leptospermum": "Myrtaceae", "Astartea": "Myrtaceae", "Taxandria": "Myrtaceae",

    # Fabaceae (Legumes)
    "Acacia": "Fabaceae", "Gastrolobium": "Fabaceae", "Daviesia": "Fabaceae",
    "Jacksonia": "Fabaceae", "Bossiaea": "Fabaceae", "Hovea": "Fabaceae",
    "Kennedia": "Fabaceae", "Hardenbergia": "Fabaceae", "Gompholobium": "Fabaceae",
    "Mirbelia": "Fabaceae", "Chorizema": "Fabaceae", "Pultenaea": "Fabaceae",
    "Isotropis": "Fabaceae", "Swainsona": "Fabaceae", "Senna": "Fabaceae",
    "Indigofera": "Fabaceae", "Templetonia": "Fabaceae", "Latrobea": "Fabaceae",

    # Orchidaceae (Orchids)
    "Caladenia": "Orchidaceae", "Thelymitra": "Orchidaceae", "Diuris": "Orchidaceae",
    "Pterostylis": "Orchidaceae", "Drakaea": "Orchidaceae", "Microtis": "Orchidaceae",
    "Prasophyllum": "Orchidaceae", "Cyanicula": "Orchidaceae", "Ericksonella": "Orchidaceae",
    "Leporella": "Orchidaceae", "Paracaleana": "Orchidaceae", "Pyrorchis": "Orchidaceae",
    "Cryptostylis": "Orchidaceae", "Elythranthera": "Orchidaceae", "Spiculaea": "Orchidaceae",

    # Haemodoraceae (Kangaroo Paws)
    "Anigozanthos": "Haemodoraceae", "Macropidia": "Haemodoraceae", "Conostylis": "Haemodoraceae",
    "Tribonanthes": "Haemodoraceae", "Haemodorum": "Haemodoraceae", "Phlebocarya": "Haemodoraceae",

    # Ericaceae / Epacridaceae (Heaths)
    "Leucopogon": "Ericaceae", "Astroloma": "Ericaceae", "Lysinema": "Ericaceae",
    "Conostephium": "Ericaceae", "Andersonia": "Ericaceae", "Sphenotoma": "Ericaceae",
    "Styphelia": "Ericaceae", "Cosmelia": "Ericaceae",

    # Goodeniaceae
    "Goodenia": "Goodeniaceae", "Dampiera": "Goodeniaceae", "Scaevola": "Goodeniaceae",
    "Lechenaultia": "Goodeniaceae", "Velleia": "Goodeniaceae", "Anthotium": "Goodeniaceae",

    # Stylidiaceae (Triggerplants)
    "Stylidium": "Stylidiaceae", "Levenhookia": "Stylidiaceae",

    # Droseraceae (Sundews / Carnivorous)
    "Drosera": "Droseraceae",

    # Asteraceae (Daisies)
    "Rhodanthe": "Asteraceae", "Waitzia": "Asteraceae", "Schoenia": "Asteraceae",
    "Podolepis": "Asteraceae", "Olearia": "Asteraceae", "Senecio": "Asteraceae",
    "Brachyscome": "Asteraceae", "Lawrencella": "Asteraceae", "Xerochrysum": "Asteraceae",

    # Xanthorrhoeaceae / Asphodelaceae
    "Xanthorrhoea": "Asphodelaceae",

    # Dasypogonaceae
    "Kingia": "Dasypogonaceae", "Dasypogon": "Dasypogonaceae", "Calectasia": "Dasypogonaceae",

    # Casuarinaceae (Sheoaks)
    "Allocasuarina": "Casuarinaceae", "Casuarina": "Casuarinaceae",

    # Pittosporaceae
    "Marianthus": "Pittosporaceae", "Billardiera": "Pittosporaceae", "Pittosporum": "Pittosporaceae",
    "Cheiranthera": "Pittosporaceae",

    # Restionaceae (Rushes / Sedges)
    "Desmocladus": "Restionaceae", "Alexgeorgea": "Restionaceae", "Loxocarya": "Restionaceae",
    "Chordifex": "Restionaceae", "Lyginia": "Restionaceae",

    # Chenopodiaceae / Amaranthaceae
    "Atriplex": "Amaranthaceae", "Maireana": "Amaranthaceae", "Tecticornia": "Amaranthaceae",
    "Ptilotus": "Amaranthaceae",

    # Malvaceae
    "Alyogyne": "Malvaceae", "Thomasia": "Malvaceae", "Guichenotia": "Malvaceae",
    "Lasiopetalum": "Malvaceae", "Seringia": "Malvaceae", "Androcalva": "Malvaceae",

    # Rutaceae (Boronias / Waxflowers)
    "Boronia": "Rutaceae", "Philotheca": "Rutaceae", "Diplolaena": "Rutaceae",
    "Correa": "Rutaceae", "Crowea": "Rutaceae", "Geleznowia": "Rutaceae",
    "Asterolasia": "Rutaceae", "Eriostemon": "Rutaceae",

    # Loranthaceae (Mistletoes / Christmas Tree)
    "Nuytsia": "Loranthaceae", "Amyema": "Loranthaceae", "Lysiana": "Loranthaceae",
}

class TaxonomyHierarchy:
    """Manages species -> genus -> family indices and mappings."""
    def __init__(self, species_list: list[str]):
        self.species_list = sorted(list(set(s for s in species_list if isinstance(s, str) and s.strip())))
        
        # 1. Species Index
        self.species_to_idx = {s: i for i, s in enumerate(self.species_list)}
        self.idx_to_species = {i: s for s, i in self.species_to_idx.items()}
        self.n_species = len(self.species_list)

        # 2. Genus Index
        self.species_to_genus = {}
        genera_set = set()
        for s in self.species_list:
            genus = s.split()[0] if s else "Unknown"
            self.species_to_genus[s] = genus
            genera_set.add(genus)
        
        self.genus_list = sorted(list(genera_set))
        self.genus_to_idx = {g: i for i, g in enumerate(self.genus_list)}
        self.idx_to_genus = {i: g for g, i in self.genus_to_idx.items()}
        self.n_genera = len(self.genus_list)

        # 3. Family Index
        # FIX: avoid synthetic per-genus families (previous f"{g}_Family" exploded to ~250 singletons).
        # Unknown genera now collapsed to a single "Unknown_Family" bucket to keep loss meaningful.
        self.genus_to_family = {}
        families_set = set()
        unknown_genera: list[str] = []
        for g in self.genus_list:
            fam = GENUS_TO_FAMILY.get(g)
            if fam is None:
                fam = "Unknown_Family"
                unknown_genera.append(g)
            self.genus_to_family[g] = fam
            families_set.add(fam)
        if unknown_genera:
            print(
                f"[taxonomy] {len(unknown_genera)}/{len(self.genus_list)} genera fallback to Unknown_Family "
                f"(e.g. {unknown_genera[:5]})",
                flush=True,
            )

        self.family_list = sorted(list(families_set))
        self.family_to_idx = {f: i for i, f in enumerate(self.family_list)}
        self.idx_to_family = {i: f for f, i in self.family_to_idx.items()}
        self.n_families = len(self.family_list)

        # Precompute lookup tables
        self.species_idx_to_genus_idx = [
            self.genus_to_idx[self.species_to_genus[self.idx_to_species[i]]]
            for i in range(self.n_species)
        ]
        self.species_idx_to_family_idx = [
            self.family_to_idx[self.genus_to_family[self.species_to_genus[self.idx_to_species[i]]]]
            for i in range(self.n_species)
        ]

    def get_taxonomy_indices(self, species_name: str) -> tuple[int, int, int]:
        """Returns (species_idx, genus_idx, family_idx)."""
        s_idx = self.species_to_idx.get(species_name, -1)
        if s_idx == -1:
            return -1, -1, -1
        return s_idx, self.species_idx_to_genus_idx[s_idx], self.species_idx_to_family_idx[s_idx]