from __future__ import annotations

import json
import logging
import re
from collections import Counter
from pathlib import Path
from typing import Any

logger = logging.getLogger("pino.odor_descriptors")

# Target vocabulary size to align with PIMT output dimension.
VOCAB_SIZE = 138

# Chemical / functional-group terms that are structure, not smell.
CHEMICAL_TOKENS = {
    "methyl", "ethyl", "propyl", "isopropyl", "butyl", "isobutyl", "amyl", "isoamyl",
    "phenyl", "benzyl", "phenethyl", "cinnamyl", "salicyl", "anisyl", "veratryl",
    "aldehyde", "aldehydic", "ketone", "ester", "acetate", "formate", "propionate",
    "butyrate", "isobutyrate", "valerate", "isovalerate", "benzoate", "salicylate",
    "cinnamate", "anthranilate", "alcohol", "carbinol", "carbinyl", "phenol", "oxide",
    "lactone", "ether", "anhydride", "acid", "acetic", "formic", "benzoic", "cinnamic",
    "hydroxy", "methoxy", "dimethoxy", "trimethyl", "tert", "para", "meta", "ortho",
    "alpha", "beta", "gamma", "delta", "omega", "cis", "trans", "neo", "iso", "sec",
    "di", "tri", "tetra", "mono", "bis", "terpenic", "terpene", "sesquiterpene",
    "hydrocarbon", "aromatic", "aliphatic", "cyclic", "acyclic", "saturated", "unsaturated",
    "molecule", "molecular", "compound", "compounds", "constituents", "composition",
    "solution", "solvent", "concentration", "dilution", "percent", "percentage", "parts",
    "contains", "containing", "component", "components", "mixture", "mixtures", "preparation",
    "synthetic", "nature", "identical", "chemical", "chemically", "organic", "volatile",
    "vapor", "steam", "distillation", "extracted", "extraction", "isolated", "isolation",
    "derivatives", "derivative", "homologue", "homolog", "analogue", "analog",
}

# Prose / grammar terms that carry no scent meaning.
STOPWORDS = {
    "the", "and", "with", "without", "slightly", "very", "more", "less", "somewhat",
    "moderately", "highly", "extremely", "quite", "rather", "freshly", "freshness",
    "dry", "dried", "powdery", "note", "notes", "odor", "odour", "smell", "smells",
    "scent", "fragrance", "aroma", "character", "like", "type", "hint", "touch",
    "suggestion", "reminiscent", "of", "in", "on", "at", "by", "to", "from", "a", "an",
    "is", "are", "was", "were", "be", "been", "being", "have", "has", "had", "do",
    "does", "did", "will", "would", "could", "should", "may", "might", "can", "shall",
    "it", "its", "this", "that", "these", "those", "such", "also", "but", "or", "as",
    "so", "than", "too", "much", "many", "most", "some", "any", "all", "no", "not",
    "only", "just", "even", "still", "yet", "however", "therefore", "thus", "hence",
    "moreover", "furthermore", "nevertheless", "otherwise", "meanwhile", "whereas",
    "although", "though", "while", "because", "since", "until", "unless", "whether",
    "either", "neither", "both", "each", "every", "one", "two", "three", "first",
    "second", "third", "last", "next", "other", "another", "same", "different",
    "similar", "various", "several", "certain", "particular", "specific", "general",
    "overall", "main", "major", "minor", "key", "important", "pronounced", "distinct",
    "definite", "clear", "vague", "faint", "subtle", "delicate", "refined", "crude",
    "raw", "processed", "manufactured", "genuine", "authentic", "real", "true", "false",
    "apparent", "seeming", "usually", "often", "sometimes", "frequently", "rarely",
    "occasionally", "generally", "commonly", "mostly", "partly", "almost", "nearly",
    "approximately", "about", "around", "extremely", "intensely", "strongly", "weakly",
    "highly", "slightly", "very", "too", "so", "such", "rather", "quite", "fairly",
    "pretty", "really", "truly", "actually", "certainly", "probably", "possibly",
    "perhaps", "maybe", "apparently", "seemingly", "evidently", "obviously", "clearly",
    "undoubtedly", "definitely", "absolutely", "relatively", "comparatively", "exceptionally",
    "unusually", "remarkably", "particularly", "especially", "notably", "significantly",
    "considerably", "substantially", "markedly", "noticeably", "distinctly", "sharply",
    "strong", "weak", "powerful", "intense", "mild", "soft", "hard", "sharp", "round",
    "rich", "poor", "full", "light", "heavy", "deep", "high", "low", "fine", "coarse",
    "pure", "impure", "clean", "dirty", "clear", "opaque", "transparent", "cloudy",
    "warm", "cool", "cold", "hot", "tenacious", "fugitive", "volatile", "evanescent",
    "lasting", "lingering", "fading", "receding", "emerging", "developing", "mellowing",
    "ageing", "young", "old", "new", "novel", "normal", "abnormal", "natural", "artificial",
    "which", "following", "follows", "where", "when", "what", "who", "how", "why", "there",
    "here", "then", "now", "again", "once", "twice", "upon", "into", "onto", "up", "down",
    "out", "over", "under", "off", "through", "between", "among", "within", "without",
    "above", "below", "before", "after", "behind", "beyond", "across", "against", "along",
    "around", "beside", "besides", "beyond", "despite", "during", "except", "inside",
    "outside", "toward", "towards", "underneath", "until", "via", "worth", "per",
}


def _tokenize(text: str) -> list[str]:
    """Extract normalized word tokens from a raw description."""
    text = text.lower()
    # Keep hyphens as word separators, keep letters and numbers.
    tokens = re.findall(r"[a-z0-9]+", text.replace("-", " "))
    return [
        t for t in tokens
        if t and t not in STOPWORDS and t not in CHEMICAL_TOKENS and len(t) > 2
    ]


def _extract_poucher_descriptors(records: list[dict[str, Any]]) -> Counter[str]:
    """Collect descriptors from Poucher structured profiles."""
    counts = Counter[str]()
    for rec in records:
        for key in ("top", "middle", "base"):
            for item in rec.get(key, []):
                # Each item is a raw material like "bergamot oil" or "linalool".
                # Extract simple descriptors (skip "oil", "absolute", etc.).
                for token in _tokenize(item):
                    if token in {"oil", "absolute", "resinoid", "concrete", "extract", "otto", "essence"}:
                        continue
                    counts[token] += 2  # Poucher descriptors are structured; weight higher.
    return counts


def _extract_arctander_descriptors(records: list[dict[str, Any]]) -> Counter[str]:
    """Collect descriptors from Arctander free-text monographs."""
    counts = Counter[str]()
    for rec in records:
        text = rec.get("odor_description") or ""
        if not text:
            continue
        for token in _tokenize(text):
            counts[token] += 1
    return counts


def build_descriptor_vocab(
    poucher_path: str | Path | None,
    arctander_path: str | Path | None,
    vocab_size: int = VOCAB_SIZE,
    extra_terms: list[str] | None = None,
) -> list[str]:
    """Build a canonical descriptor vocabulary from available literature sources."""
    counts = Counter[str]()

    if poucher_path and Path(poucher_path).exists():
        records = [
            json.loads(line) for line in Path(poucher_path).read_text().strip().splitlines()
        ]
        counts.update(_extract_poucher_descriptors(records))

    if arctander_path and Path(arctander_path).exists():
        records = [
            json.loads(line) for line in Path(arctander_path).read_text().strip().splitlines()
        ]
        counts.update(_extract_arctander_descriptors(records))

    if extra_terms:
        for term in extra_terms:
            counts[term.lower()] += 3

    # Drop numeric-only tokens and overly rare terms.
    filtered = {
        term: c
        for term, c in counts.items()
        if not term.isdigit() and c >= 2 and len(term) > 2
    }

    # Take the most frequent descriptors up to vocab_size.
    most_common = Counter(filtered).most_common(vocab_size)
    vocab = [term for term, _ in most_common]

    # Pad with generic descriptors if needed, so the model dimension stays fixed.
    pad_descriptors = [
        "citrus", "lemon", "bergamot", "orange", "lime", "grapefruit", "mandarin",
        "floral", "rose", "jasmine", "lily", "ylang", "neroli", "orange blossom",
        "tubereuse", "gardenia", "violet", "iris", "mimosa", "acacia", "hawthorn",
        "woody", "sandalwood", "cedar", "vetiver", "patchouli", "oakmoss", "pine",
        "musk", "amber", "vanilla", "benzoin", "tolu", "labdanum", "styrax",
        "spicy", "cinnamon", "clove", "nutmeg", "pepper", "ginger", "cardamom",
        "herbal", "lavender", "sage", "rosemary", "thyme", "basil", "mint",
        "green", "leafy", "grassy", "galbanum", "tomato leaf", "fig",
        "fruity", "apple", "pear", "peach", "plum", "berry", "tropical", "coconut",
        "aldehydic", "powdery", "mossy", "earthy", "rooty", "leathery", "tobacco",
        "honey", "sweet", "balsamic", "resinous", "incense", "smoky", "burnt",
        "marine", "aquatic", "ozonic", "iodine", "seaweed",
        "animalic", "civet", "castoreum", "musk tonkin", "ambergris",
        "metallic", "camphoraceous", "eucalyptus", "menthol", "medicinal",
        "fatty", "waxy", "lactonic", "coconut", "creamy", "milky",
        "dry", "pungent", "sour", "sulfurous", "rubbery", "plastic",
        "winey", "brandy", "rum", "boozy", "fermented", "yeasty",
        "chocolate", "coffee", "cocoa", "vanilla", "caramel", "toffee",
        "nutty", "almond", "hazelnut", "peanut", "walnut",
        "coconut", "tropical", "pineapple", "banana", "mango", "papaya",
    ]
    for d in pad_descriptors:
        if d not in vocab:
            vocab.append(d)
        if len(vocab) >= vocab_size:
            break

    return vocab[:vocab_size]


def save_vocab(vocab: list[str], path: str | Path) -> None:
    """Persist the vocabulary as a JSON list."""
    Path(path).write_text(json.dumps(vocab, indent=2, ensure_ascii=False))


def load_vocab(path: str | Path) -> list[str]:
    """Load a previously saved vocabulary."""
    return json.loads(Path(path).read_text())


if __name__ == "__main__":
    import argparse

    parser = argparse.ArgumentParser(description="Build PINO descriptor vocabulary")
    parser.add_argument("--poucher", default="data/literature_profiles_poucher.jsonl")
    parser.add_argument("--arctander", default="data/arctander_monographs.jsonl")
    parser.add_argument("--output", default="data/odor_descriptor_vocab.json")
    parser.add_argument("--vocab-size", type=int, default=VOCAB_SIZE)
    args = parser.parse_args()

    vocab = build_descriptor_vocab(args.poucher, args.arctander, args.vocab_size)
    save_vocab(vocab, args.output)
    print(f"Built descriptor vocabulary: {len(vocab)} terms -> {args.output}")
    print("Top 30:", vocab[:30])