#!/usr/bin/env python3
"""
Basin Dictionary Shell - FML Frequency-Based Basin Assignment
Standalone version for cubesquare package.

HEX_ID: 0BF1DD03

CLI interface for building and querying the FML basin dictionary.
Basin assignment uses frequency-derived wavelength mapping (NOT Fernandez iteration).

Hypothesis: Word frequency (Zipf's law) maps to wavelength:
  - High frequency words (the, is, a) -> short wavelength -> violet -> low basin
  - Low frequency words (fiber, healer) -> long wavelength -> red -> high basin

Domain-specific terms co-occur and have similar frequency profiles,
so they should cluster in nearby basins.

Commands:
    status   - Show dictionary status
    query    - Query basin for a term
    rank     - Get frequency rank for a term
    build    - Build/update vocabulary from terms
    validate - Validate basin assignments
    stats    - Show vocabulary statistics
"""

import json
import re
from typing import Dict, Any, List, Tuple, Optional
from collections import Counter
from dataclasses import dataclass


# =============================================================================
# CONSTANTS
# =============================================================================

PRIME_131 = 131

# Wavelength range (visible spectrum)
LAMBDA_MIN = 380.0  # nm (violet, high frequency)
LAMBDA_MAX = 750.0  # nm (red, low frequency)
LAMBDA_RANGE = LAMBDA_MAX - LAMBDA_MIN  # 370nm

# Python keywords to exclude
PYTHON_KEYWORDS = {
    'def', 'class', 'import', 'from', 'return', 'if', 'else', 'elif',
    'for', 'while', 'try', 'except', 'with', 'as', 'in', 'is', 'not',
    'and', 'or', 'True', 'False', 'None', 'self', 'lambda', 'yield',
    'pass', 'break', 'continue', 'raise', 'finally', 'global', 'assert'
}

# Common English stopwords
STOPWORDS = {
    'the', 'a', 'an', 'is', 'are', 'was', 'were', 'be', 'been', 'being',
    'have', 'has', 'had', 'do', 'does', 'did', 'will', 'would', 'could',
    'should', 'may', 'might', 'must', 'shall', 'can', 'need', 'dare',
    'to', 'of', 'in', 'for', 'on', 'with', 'at', 'by', 'from', 'up',
    'about', 'into', 'over', 'after', 'this', 'that', 'these', 'those',
    'it', 'its', 'they', 'them', 'their', 'we', 'us', 'our', 'you', 'your',
    'he', 'him', 'his', 'she', 'her', 'who', 'which', 'what', 'where',
    'when', 'why', 'how', 'all', 'each', 'every', 'both', 'few', 'more',
    'most', 'other', 'some', 'such', 'no', 'nor', 'not', 'only', 'own',
    'same', 'so', 'than', 'too', 'very', 'just', 'also', 'now', 'here'
}


# =============================================================================
# CORE FUNCTIONS
# =============================================================================

def extract_terms(text: str) -> List[str]:
    """Extract meaningful terms from text, excluding keywords and stopwords."""
    if not text:
        return []

    tokens = re.findall(r'[a-zA-Z_][a-zA-Z0-9_]*', text.lower())

    terms = []
    for t in tokens:
        if len(t) < 3:
            continue
        if t in PYTHON_KEYWORDS or t in STOPWORDS:
            continue
        parts = re.findall(r'[a-z]+', t)
        for p in parts:
            if len(p) >= 3 and p not in STOPWORDS:
                terms.append(p)

    return terms


def frequency_to_wavelength(rank: int, vocab_size: int) -> float:
    """
    Map Zipf rank to wavelength (380-750nm).

    High rank (rare word) -> long wavelength (red)
    Low rank (common word) -> short wavelength (violet)
    """
    normalized = (rank - 1) / max(vocab_size - 1, 1)
    wavelength = LAMBDA_MIN + (LAMBDA_RANGE * normalized)
    return wavelength


def wavelength_to_basin(wavelength: float) -> int:
    """Map wavelength to basin (0-130)."""
    normalized = (wavelength - LAMBDA_MIN) / LAMBDA_RANGE
    normalized = max(0, min(1, normalized))
    basin = int(normalized * (PRIME_131 - 1))
    return basin


def get_wavelength_range_for_basin(basin: int) -> Tuple[float, float]:
    """Get wavelength range for a basin."""
    min_wl = LAMBDA_MIN + (basin / PRIME_131) * LAMBDA_RANGE
    max_wl = LAMBDA_MIN + ((basin + 1) / PRIME_131) * LAMBDA_RANGE
    return (round(min_wl, 2), round(max_wl, 2))


def get_layer_for_basin(basin: int) -> Tuple[int, str]:
    """Get layer for a basin number."""
    layer_ranges = {
        1: (0, 14, 'GEOMETRIC'),
        2: (15, 29, 'SEMANTIC'),
        3: (30, 43, 'SPATIAL'),
        4: (44, 58, 'FUNCTIONAL'),
        5: (59, 72, 'PROBABILISTIC'),
        6: (73, 87, 'EMOTIONAL'),
        7: (88, 101, 'FREQUENCY'),
        8: (102, 116, 'ELECTROMAGNETIC'),
        9: (117, 130, 'MEMBRANE'),
    }

    for layer, (start, end, name) in layer_ranges.items():
        if start <= basin <= end:
            return (layer, name)

    return (9, 'MEMBRANE')


# =============================================================================
# BASIN DICTIONARY
# =============================================================================

class BasinDictionary:
    """
    Basin Dictionary - Maps terms to basins via frequency.

    The dictionary maintains:
    - term_counts: {term: count} - how often each term appears
    - term_ranks: {term: rank} - Zipf rank (1 = most frequent)

    Basin assignment is O(1) after vocabulary is built:
        term -> rank -> wavelength -> basin
    """

    def __init__(self):
        """Initialize basin dictionary."""
        self.term_counts: Dict[str, int] = {}
        self.term_ranks: Dict[str, int] = {}
        self.vocab_size: int = 0
        self.basin_terms: Dict[int, List[str]] = {i: [] for i in range(PRIME_131)}

        # Build default vocabulary
        self._build_default_vocabulary()

    def _build_default_vocabulary(self):
        """Build default vocabulary from common terms."""
        # Common terms ranked by estimated frequency
        common_terms = [
            # High frequency (low rank = violet = low basin)
            ('file', 100), ('data', 95), ('type', 90), ('value', 85), ('name', 80),
            ('path', 75), ('key', 70), ('node', 65), ('list', 60), ('dict', 55),
            ('class', 50), ('function', 45), ('method', 40), ('object', 35), ('string', 30),
            # Medium frequency (green = middle basins)
            ('token', 25), ('basin', 24), ('layer', 23), ('agent', 22), ('state', 21),
            ('logic', 20), ('chain', 19), ('parse', 18), ('encode', 17), ('decode', 16),
            ('validate', 15), ('process', 14), ('execute', 13), ('query', 12), ('index', 11),
            # Lower frequency (yellow/orange = higher basins)
            ('healer', 10), ('validator', 9), ('expander', 8), ('crystallizer', 7), ('absorber', 6),
            ('semantic', 5), ('geometric', 4), ('fernandez', 3), ('wavelength', 2), ('frequency', 1),
            # Low frequency (red = high basins) - condition words
            ('stuck', 5), ('corrupted', 4), ('broken', 3), ('isolated', 2), ('orphaned', 1),
            ('missing', 3), ('degraded', 2), ('unstable', 1), ('failing', 2), ('error', 4),
            ('problem', 5), ('issue', 4),
        ]

        # Build counts
        for term, count in common_terms:
            self.term_counts[term] = count

        # Compute ranks (1 = most frequent)
        sorted_terms = sorted(self.term_counts.items(), key=lambda x: -x[1])
        self.term_ranks = {term: rank + 1 for rank, (term, _) in enumerate(sorted_terms)}
        self.vocab_size = len(self.term_ranks)

        # Build basin-to-terms mapping
        self._rebuild_basin_mapping()

    def _rebuild_basin_mapping(self):
        """Rebuild the basin -> terms mapping."""
        self.basin_terms = {i: [] for i in range(PRIME_131)}

        for term, rank in self.term_ranks.items():
            wavelength = frequency_to_wavelength(rank, self.vocab_size)
            basin = wavelength_to_basin(wavelength)
            self.basin_terms[basin].append({
                'term': term,
                'rank': rank,
                'count': self.term_counts.get(term, 0),
                'wavelength': round(wavelength, 2),
            })

        # Sort terms within each basin by frequency
        for basin in self.basin_terms:
            self.basin_terms[basin].sort(key=lambda x: -x['count'])

    def add_terms(self, terms: List[str]):
        """Add terms to the dictionary and update ranks."""
        for term in terms:
            term_lower = term.lower()
            self.term_counts[term_lower] = self.term_counts.get(term_lower, 0) + 1

        # Recompute ranks
        sorted_terms = sorted(self.term_counts.items(), key=lambda x: -x[1])
        self.term_ranks = {term: rank + 1 for rank, (term, _) in enumerate(sorted_terms)}
        self.vocab_size = len(self.term_ranks)

        # Rebuild basin mapping
        self._rebuild_basin_mapping()

    def get_rank(self, term: str) -> int:
        """Get frequency rank for a term (1 = most common)."""
        term_lower = term.lower()
        if term_lower in self.term_ranks:
            return self.term_ranks[term_lower]
        # Unknown terms get median rank
        return self.vocab_size // 2

    def get_wavelength(self, term: str) -> float:
        """Get wavelength for a term."""
        rank = self.get_rank(term)
        return frequency_to_wavelength(rank, self.vocab_size)

    def get_basin(self, term: str) -> int:
        """Get basin for a term."""
        wavelength = self.get_wavelength(term)
        return wavelength_to_basin(wavelength)

    def get_term_info(self, term: str) -> Dict[str, Any]:
        """Get complete information for a term."""
        term_lower = term.lower()
        rank = self.get_rank(term)
        wavelength = self.get_wavelength(term)
        basin = wavelength_to_basin(wavelength)
        layer, layer_name = get_layer_for_basin(basin)

        return {
            'term': term_lower,
            'in_vocabulary': term_lower in self.term_ranks,
            'rank': rank,
            'count': self.term_counts.get(term_lower, 0),
            'wavelength': round(wavelength, 2),
            'basin': basin,
            'layer': layer,
            'layer_name': layer_name,
        }

    def get_basin_info(self, basin: int) -> Dict[str, Any]:
        """Get information about a basin."""
        if basin < 0 or basin >= PRIME_131:
            return {'error': f'Basin must be 0-130, got {basin}'}

        terms = self.basin_terms.get(basin, [])
        wavelength_range = get_wavelength_range_for_basin(basin)
        layer, layer_name = get_layer_for_basin(basin)

        return {
            'basin': basin,
            'layer': layer,
            'layer_name': layer_name,
            'wavelength_range': wavelength_range,
            'term_count': len(terms),
            'top_terms': [t['term'] for t in terms[:5]],
            'terms': terms[:10],  # First 10 terms
        }


# =============================================================================
# BASIN DICTIONARY SHELL
# =============================================================================

class BasinDictShell:
    """
    Basin Dictionary Shell - FML Frequency-Based Basin Assignment.

    Commands:
        status   - Show dictionary status
        query    - Query basin for a term
        rank     - Get frequency rank for a term
        build    - Build/update vocabulary from terms
        validate - Validate basin assignments
        stats    - Show vocabulary statistics
    """

    def __init__(self):
        """Initialize basin dictionary shell."""
        self.hex_id = '0BF1DD03'
        self.dictionary = BasinDictionary()

    def status(self, as_json: bool = False) -> Dict[str, Any]:
        """Get dictionary status."""
        # Count non-empty basins
        non_empty = sum(1 for b in self.dictionary.basin_terms.values() if b)

        result = {
            'hex_id': self.hex_id,
            'system': 'Basin Dictionary Shell',
            'version': '1.0.0',
            'method': 'FML frequency-based wavelength mapping',
            'description': 'Basin = wavelength band, NOT Fernandez iteration',
            'wavelength_range': f'{LAMBDA_MIN}-{LAMBDA_MAX}nm',
            'basin_count': PRIME_131,
            'vocab_size': self.dictionary.vocab_size,
            'non_empty_basins': non_empty,
            'mapping': 'term -> rank (Zipf) -> wavelength -> basin',
        }

        if as_json:
            return json.dumps(result, indent=2)
        return result

    def query(self, term: str, as_json: bool = False) -> Dict[str, Any]:
        """Query basin for a term."""
        result = self.dictionary.get_term_info(term)

        # Add basin details
        basin_info = self.dictionary.get_basin_info(result['basin'])
        result['basin_label'] = '/'.join(basin_info['top_terms'][:3]) if basin_info['top_terms'] else 'EMPTY'
        result['other_terms_in_basin'] = basin_info['top_terms']

        if as_json:
            return json.dumps(result, indent=2)
        return result

    def rank(self, term: str, as_json: bool = False) -> Dict[str, Any]:
        """Get frequency rank for a term."""
        rank = self.dictionary.get_rank(term)
        count = self.dictionary.term_counts.get(term.lower(), 0)

        result = {
            'term': term.lower(),
            'rank': rank,
            'count': count,
            'vocab_size': self.dictionary.vocab_size,
            'frequency_type': 'HIGH' if rank < self.dictionary.vocab_size // 3 else 'MEDIUM' if rank < 2 * self.dictionary.vocab_size // 3 else 'LOW',
            'in_vocabulary': term.lower() in self.dictionary.term_ranks,
        }

        if as_json:
            return json.dumps(result, indent=2)
        return result

    def build(self, text: str, as_json: bool = False) -> Dict[str, Any]:
        """Build/update vocabulary from text."""
        terms = extract_terms(text)
        if not terms:
            return {'error': 'No meaningful terms found in text'}

        initial_size = self.dictionary.vocab_size
        self.dictionary.add_terms(terms)
        final_size = self.dictionary.vocab_size

        result = {
            'terms_extracted': len(terms),
            'unique_terms': len(set(terms)),
            'initial_vocab_size': initial_size,
            'final_vocab_size': final_size,
            'new_terms_added': final_size - initial_size,
            'sample_terms': list(set(terms))[:10],
        }

        if as_json:
            return json.dumps(result, indent=2)
        return result

    def validate(self, as_json: bool = False) -> Dict[str, Any]:
        """Validate basin assignments."""
        # Count terms per basin
        term_counts = [len(terms) for terms in self.dictionary.basin_terms.values()]

        # Basic statistics
        min_count = min(term_counts)
        max_count = max(term_counts)
        mean_count = sum(term_counts) / len(term_counts)
        non_empty = sum(1 for c in term_counts if c > 0)

        # Variance
        variance = sum((c - mean_count) ** 2 for c in term_counts) / len(term_counts)
        std_dev = variance ** 0.5

        # Sample basins for coherence check
        sample_basins = [0, 32, 65, 97, 130]
        coherence_samples = {}
        for basin in sample_basins:
            info = self.dictionary.get_basin_info(basin)
            coherence_samples[basin] = {
                'layer': info['layer_name'],
                'wavelength_range': info['wavelength_range'],
                'top_terms': info['top_terms'],
            }

        result = {
            'validation': 'PASSED' if non_empty > 0 else 'FAILED',
            'statistics': {
                'min_terms_per_basin': min_count,
                'max_terms_per_basin': max_count,
                'mean_terms_per_basin': round(mean_count, 2),
                'std_dev': round(std_dev, 2),
                'non_empty_basins': non_empty,
                'total_basins': PRIME_131,
            },
            'coherence_samples': coherence_samples,
        }

        if as_json:
            return json.dumps(result, indent=2)
        return result

    def stats(self, as_json: bool = False) -> Dict[str, Any]:
        """Show vocabulary statistics."""
        # Top 10 terms by frequency
        sorted_terms = sorted(self.dictionary.term_counts.items(), key=lambda x: -x[1])
        top_terms = [{'term': t, 'count': c, 'rank': self.dictionary.term_ranks[t]}
                     for t, c in sorted_terms[:10]]

        # Terms per layer
        layer_counts = {}
        for basin, terms in self.dictionary.basin_terms.items():
            layer, layer_name = get_layer_for_basin(basin)
            if layer_name not in layer_counts:
                layer_counts[layer_name] = 0
            layer_counts[layer_name] += len(terms)

        result = {
            'vocab_size': self.dictionary.vocab_size,
            'total_occurrences': sum(self.dictionary.term_counts.values()),
            'top_terms': top_terms,
            'terms_per_layer': layer_counts,
        }

        if as_json:
            return json.dumps(result, indent=2)
        return result

    def basin_info(self, basin: int, as_json: bool = False) -> Dict[str, Any]:
        """Get information about a specific basin."""
        result = self.dictionary.get_basin_info(basin)

        if as_json:
            return json.dumps(result, indent=2)
        return result

    def run(self, command: str, *args, as_json: bool = False) -> Any:
        """Run a basin dictionary shell command."""
        commands = {
            'status': lambda: self.status(as_json),
            'query': lambda: self.query(args[0], as_json) if args else {'error': 'Term required'},
            'rank': lambda: self.rank(args[0], as_json) if args else {'error': 'Term required'},
            'build': lambda: self.build(args[0], as_json) if args else {'error': 'Text required'},
            'validate': lambda: self.validate(as_json),
            'stats': lambda: self.stats(as_json),
            'basin': lambda: self.basin_info(int(args[0]), as_json) if args else {'error': 'Basin number required'},
        }

        if command in commands:
            return commands[command]()
        else:
            return {'error': f'Unknown command: {command}', 'available': list(commands.keys())}
