"""Learn conservative YNAB category suggestions from transaction CSV history.""" import csv from dataclasses import dataclass from decimal import Decimal, InvalidOperation, ROUND_HALF_UP from difflib import SequenceMatcher import json import math from pathlib import Path import re import unicodedata _MEMORY_VERSION = 1 _MIN_PAYEE_SIMILARITY = 0.70 _MIN_CATEGORY_SHARE = 0.80 _AMOUNT_SMOOTHING_CENTS = 100 @dataclass(frozen=True, order=True) class CategoryPattern: """An aggregated payee/category/amount pattern, without transaction dates.""" payee: str category: str direction: str amount_cents: int count: int = 1 @dataclass(frozen=True) class CategorySuggestion: """A category candidate and whether evidence is strong enough to apply it.""" category_name: str confidence: float category_id: str | None applied: bool def normalize_payee(payee: str) -> str: """Normalize formatting and discard long numeric or mixed reference codes.""" normalized = unicodedata.normalize("NFKD", payee).encode("ascii", "ignore").decode() tokens = re.findall(r"[a-z0-9]+", normalized.casefold()) stable_tokens = [ token for token in tokens if not ( (token.isdigit() and len(token) >= 3) or ( len(token) >= 5 and re.search(r"[a-z]", token) and re.search(r"[0-9]", token) ) ) ] return " ".join(stable_tokens) def _amount_cents(value: str) -> int | None: cleaned = value.strip().replace("$", "").replace(",", "") if not cleaned: return None try: amount = Decimal(cleaned) except InvalidOperation: return None if not amount.is_finite(): return None cents = (amount * 100).quantize(Decimal("1"), rounding=ROUND_HALF_UP) return int(cents) class CategoryHistory: """Memory of unique merchant/category/amount patterns and snapshot support counts.""" def __init__( self, patterns: dict[tuple[str, str, str, int], CategoryPattern] | None = None, ) -> None: self.patterns = patterns or {} @staticmethod def _key(pattern: CategoryPattern) -> tuple[str, str, str, int]: return ( pattern.payee, pattern.category.casefold(), pattern.direction, pattern.amount_cents, ) @staticmethod def memory_path_for_csv(csv_path: str | Path) -> Path: """Return the portable category-memory sidecar path beside the CSV.""" return Path(csv_path).with_name("category_memory.json") @classmethod def load(cls, path: str | Path) -> "CategoryHistory": """Load and validate a category-memory sidecar.""" memory_path = Path(path) payload = json.loads(memory_path.read_text(encoding="utf-8")) if not isinstance(payload, dict) or payload.get("version") != _MEMORY_VERSION: raise ValueError(f"Unsupported category memory version in {memory_path}") patterns: dict[tuple[str, str, str, int], CategoryPattern] = {} for item in payload.get("patterns", []): payee = str(item["payee"]).strip() category = str(item["category"]).strip() direction = str(item["direction"]).strip().lower() amount_cents = int(item["amount_cents"]) count = max(1, int(item.get("count", 1))) if payee and category and direction in ("inflow", "outflow") and amount_cents > 0: pattern = CategoryPattern(payee, category, direction, amount_cents, count) patterns[cls._key(pattern)] = pattern return cls(patterns) @classmethod def from_csv(cls, csv_path: str | Path) -> "CategoryHistory": """Read the repository's four-column transaction CSV as category examples.""" path = Path(csv_path) patterns: dict[tuple[str, str, str, int], CategoryPattern] = {} with path.open("r", encoding="utf-8-sig", newline="") as csv_file: reader = csv.DictReader(csv_file) required = {"payee", "category", "inflow", "outflow"} if not reader.fieldnames or not required.issubset(reader.fieldnames): raise ValueError( f"{path} must contain payee, category, inflow, and outflow columns" ) for row in reader: payee = normalize_payee(row.get("payee") or "") category = (row.get("category") or "").strip() inflow = _amount_cents(row.get("inflow") or "") outflow = _amount_cents(row.get("outflow") or "") if not payee or not category: continue if inflow is not None and inflow > 0 and (outflow is None or outflow == 0): pattern = CategoryPattern(payee, category, "inflow", inflow) key = cls._key(pattern) existing = patterns.get(key) patterns[key] = CategoryPattern( payee, category, "inflow", inflow, (existing.count if existing else 0) + 1, ) elif outflow is not None and outflow > 0 and (inflow is None or inflow == 0): pattern = CategoryPattern(payee, category, "outflow", outflow) key = cls._key(pattern) existing = patterns.get(key) patterns[key] = CategoryPattern( payee, category, "outflow", outflow, (existing.count if existing else 0) + 1, ) return cls(patterns) def merge(self, other: "CategoryHistory") -> int: """Union patterns, returning how many new patterns were added.""" changes = 0 for key, pattern in other.patterns.items(): existing = self.patterns.get(key) if existing is None or pattern.count > existing.count: self.patterns[key] = pattern changes += 1 return changes def save(self, path: str | Path) -> None: """Write a deterministic, portable JSON sidecar using an atomic replace.""" memory_path = Path(path) memory_path.parent.mkdir(parents=True, exist_ok=True) payload = { "version": _MEMORY_VERSION, "patterns": [ { "payee": pattern.payee, "category": pattern.category, "direction": pattern.direction, "amount_cents": pattern.amount_cents, "count": pattern.count, } for pattern in sorted( self.patterns.values(), key=lambda item: ( item.payee, item.category.casefold(), item.direction, item.amount_cents, ), ) ], } temporary_path = memory_path.with_name(memory_path.name + ".tmp") temporary_path.write_text( json.dumps(payload, indent=2, ensure_ascii=False) + "\n", encoding="utf-8", ) temporary_path.replace(memory_path) def suggest( self, payee: str, amount_milli: int, category_ids_by_name: dict[str, list[str]], ) -> CategorySuggestion | None: """Rank historical categories using payee similarity, direction, and amount.""" payee_key = normalize_payee(payee) if not payee_key or not self.patterns: return None direction = "inflow" if amount_milli > 0 else "outflow" amount_cents = round(abs(amount_milli) / 10) category_scores: dict[str, float] = {} category_names: dict[str, str] = {} category_support: dict[str, int] = {} category_exact: dict[str, bool] = {} for pattern in self.patterns.values(): if pattern.direction != direction: continue payee_similarity = SequenceMatcher( None, payee_key, pattern.payee, autojunk=False ).ratio() if payee_similarity < _MIN_PAYEE_SIMILARITY: continue ratio = (amount_cents + _AMOUNT_SMOOTHING_CENTS) / ( pattern.amount_cents + _AMOUNT_SMOOTHING_CENTS ) amount_similarity = 1 / (1 + abs(math.log(ratio))) weight = payee_similarity * (amount_similarity ** 2) * math.sqrt(pattern.count) category_key = pattern.category.casefold() category_scores[category_key] = category_scores.get(category_key, 0.0) + weight category_names.setdefault(category_key, pattern.category) category_support[category_key] = ( category_support.get(category_key, 0) + pattern.count ) if payee_key == pattern.payee: category_exact[category_key] = True else: category_exact.setdefault(category_key, False) total_score = sum(category_scores.values()) if total_score <= 0: return None best_category = max(category_scores, key=category_scores.get) confidence = category_scores[best_category] / total_score sufficiently_supported = ( category_exact.get(best_category, False) or category_support[best_category] >= 2 ) applied = ( confidence >= _MIN_CATEGORY_SHARE and sufficiently_supported ) category_name = category_names[best_category] matching_ids = category_ids_by_name.get(category_name.casefold(), []) category_id = matching_ids[0] if applied and len(matching_ids) == 1 else None return CategorySuggestion( category_name=category_name, confidence=confidence, category_id=category_id, applied=category_id is not None, )