"""
CDES v1.5 Terpene Ratio Analysis Module

Provides compositional data analysis methods for terpene profiles:
- Dominance Index
- Balance Score (Shannon Entropy)
- CLR (Centered Log-Ratio) Coordinates
- Ratio Signatures
- Archetype Classification

Version: 1.5.0
CDES Standard: Cannabis Data Exchange Standard v1.5
"""

from dataclasses import dataclass, field
from typing import Optional, List, Dict, Any, Tuple
from enum import Enum
import math


# =============================================================================
# ENUMERATIONS
# =============================================================================

class TerpeneArchetype(str, Enum):
    """CDES v1.5: Terpene profile archetypes based on clustering analysis."""
    ARCH_MYR = "ARCH_MYR"  # Myrcene Dominant
    ARCH_LIM = "ARCH_LIM"  # Limonene Forward
    ARCH_BAL = "ARCH_BAL"  # Balanced Trio
    ARCH_PIN = "ARCH_PIN"  # Pinene Sharp
    ARCH_CAR = "ARCH_CAR"  # Caryophyllene Spicy
    ARCH_TER = "ARCH_TER"  # Terpinolene Rare
    ARCH_LIN = "ARCH_LIN"  # Linalool Floral
    ARCH_OCI = "ARCH_OCI"  # Ocimene Sweet
    ARCH_UNK = "ARCH_UNK"  # Unknown/Unclassified


class DominanceLevel(str, Enum):
    """CDES v1.5: Qualitative dominance level based on Dominance Index."""
    BALANCED = "balanced"      # DI < 0.25
    SLIGHT = "slight"          # DI 0.25-0.50
    MODERATE = "moderate"      # DI 0.50-0.75
    STRONG = "strong"          # DI > 0.75


# =============================================================================
# TERPENE CODES AND ORDERING
# =============================================================================

TERPENE_CODES = {
    "myrcene": "MYR",
    "limonene": "LIM",
    "caryophyllene": "CAR",
    "pinene": "PIN",
    "linalool": "LIN",
    "humulene": "HUM",
    "terpinolene": "TER",
    "ocimene": "OCI",
    "bisabolol": "BIS",
}

TERPENE_NAMES = {
    "myrcene": "β-Myrcene",
    "limonene": "δ-Limonene",
    "caryophyllene": "β-Caryophyllene",
    "pinene": "α-Pinene",
    "linalool": "Linalool",
    "humulene": "α-Humulene",
    "terpinolene": "Terpinolene",
    "ocimene": "Ocimene",
    "bisabolol": "α-Bisabolol",
}

FIXED_TERPENE_ORDER = [
    "myrcene", "limonene", "caryophyllene", "pinene", "linalool",
    "humulene", "terpinolene", "ocimene", "bisabolol"
]

# Archetype centroids (CLR coordinates) - trained from Terprint data
# These are preliminary values to be refined with actual clustering
ARCHETYPE_CENTROIDS = {
    TerpeneArchetype.ARCH_MYR: {"myrcene": 1.2, "limonene": -0.3, "caryophyllene": -0.4, "pinene": -0.5, "linalool": -0.2, "humulene": -0.6, "terpinolene": -0.8, "ocimene": -0.7, "bisabolol": -0.5},
    TerpeneArchetype.ARCH_LIM: {"myrcene": -0.2, "limonene": 1.1, "caryophyllene": -0.2, "pinene": 0.3, "linalool": -0.4, "humulene": -0.5, "terpinolene": 0.2, "ocimene": -0.3, "bisabolol": -0.4},
    TerpeneArchetype.ARCH_BAL: {"myrcene": 0.3, "limonene": 0.2, "caryophyllene": 0.2, "pinene": 0.1, "linalool": 0.1, "humulene": -0.1, "terpinolene": -0.2, "ocimene": -0.2, "bisabolol": -0.1},
    TerpeneArchetype.ARCH_PIN: {"myrcene": -0.4, "limonene": 0.2, "caryophyllene": -0.3, "pinene": 1.0, "linalool": -0.3, "humulene": -0.4, "terpinolene": 0.1, "ocimene": -0.2, "bisabolol": -0.3},
    TerpeneArchetype.ARCH_CAR: {"myrcene": -0.3, "limonene": -0.4, "caryophyllene": 1.0, "pinene": -0.2, "linalool": 0.2, "humulene": 0.6, "terpinolene": -0.5, "ocimene": -0.4, "bisabolol": -0.3},
    TerpeneArchetype.ARCH_TER: {"myrcene": -0.5, "limonene": -0.3, "caryophyllene": -0.4, "pinene": 0.2, "linalool": -0.3, "humulene": -0.5, "terpinolene": 1.3, "ocimene": 0.1, "bisabolol": -0.4},
    TerpeneArchetype.ARCH_LIN: {"myrcene": 0.2, "limonene": -0.2, "caryophyllene": 0.1, "pinene": -0.3, "linalool": 1.0, "humulene": -0.2, "terpinolene": -0.4, "ocimene": -0.3, "bisabolol": 0.3},
    TerpeneArchetype.ARCH_OCI: {"myrcene": -0.3, "limonene": -0.2, "caryophyllene": -0.3, "pinene": -0.2, "linalool": -0.2, "humulene": -0.4, "terpinolene": 0.2, "ocimene": 1.1, "bisabolol": -0.2},
}


# =============================================================================
# DATA CLASSES
# =============================================================================

@dataclass
class RankedTerpene:
    """CDES v1.5: A terpene with its rank position in the profile."""
    rank: int
    terpene_id: str
    name: str
    concentration: float
    unit: str = "percent"
    percent_of_total: float = 0.0

    def to_dict(self) -> Dict[str, Any]:
        return {
            "rank": self.rank,
            "terpeneId": self.terpene_id,
            "name": self.name,
            "concentration": self.concentration,
            "unit": self.unit,
            "percentOfTotal": self.percent_of_total,
        }


@dataclass
class RatioSignature:
    """CDES v1.5: Compact ratio representation for similarity matching."""
    encoded: str                    # "MYR:LIM:CAR|0.57:0.38:0.24"
    t1_t2_ratio: float              # T1/T2
    t1_t3_ratio: float              # T1/T3
    t2_t3_ratio: float              # T2/T3

    def to_dict(self) -> Dict[str, Any]:
        return {
            "encoded": self.encoded,
            "t1t2Ratio": self.t1_t2_ratio,
            "t1t3Ratio": self.t1_t3_ratio,
            "t2t3Ratio": self.t2_t3_ratio,
        }


@dataclass
class ArchetypeClassification:
    """CDES v1.5: Archetype classification results."""
    primary_archetype: TerpeneArchetype
    archetype_confidence: float
    archetype_distances: Dict[str, float] = field(default_factory=dict)

    def to_dict(self) -> Dict[str, Any]:
        return {
            "primaryArchetype": self.primary_archetype.value,
            "archetypeConfidence": self.archetype_confidence,
            "archetypeDistances": self.archetype_distances,
        }


@dataclass
class TerpeneRatioMetrics:
    """CDES v1.5: Computed ratio analysis metrics for terpene profiles."""
    
    # Core metrics
    dominance_index: float          # 0.0 (balanced) to 1.0 (dominated)
    dominance_level: DominanceLevel # Qualitative level
    balance_score: float            # Shannon entropy (raw)
    balance_score_normalized: float # 0-1 normalized entropy
    hhi: float                      # Herfindahl-Hirschman Index
    top_terpene_count: int          # Terpenes > 5% of total
    
    # Ranked terpenes
    ranked_terpenes: List[RankedTerpene] = field(default_factory=list)
    
    # Signature
    ratio_signature: Optional[RatioSignature] = None
    
    # CLR coordinates (for clustering/ML)
    clr_coordinates: Dict[str, float] = field(default_factory=dict)
    
    # Key pairwise log-ratios
    key_ratios: Dict[str, float] = field(default_factory=dict)
    
    # Archetype classification
    archetype_classification: Optional[ArchetypeClassification] = None

    def to_dict(self) -> Dict[str, Any]:
        return {
            "dominanceIndex": self.dominance_index,
            "dominanceLevel": self.dominance_level.value,
            "balanceScore": self.balance_score,
            "balanceScoreNormalized": self.balance_score_normalized,
            "hhi": self.hhi,
            "topTerpeneCount": self.top_terpene_count,
            "rankedTerpenes": [t.to_dict() for t in self.ranked_terpenes],
            "ratioSignature": self.ratio_signature.to_dict() if self.ratio_signature else None,
            "clrCoordinates": self.clr_coordinates,
            "keyRatios": self.key_ratios,
            "archetypeClassification": self.archetype_classification.to_dict() if self.archetype_classification else None,
        }


# =============================================================================
# COMPUTATION FUNCTIONS
# =============================================================================

def _prepare_for_clr(values: List[float], delta: float = 0.001) -> List[float]:
    """
    Multiplicative replacement of zeros before CLR transform.
    
    Args:
        values: List of terpene concentrations
        delta: Replacement value for zeros (default 0.001%)
        
    Returns:
        List with zeros replaced and non-zeros scaled
    """
    result = list(values)
    zeros = [i for i, v in enumerate(result) if v == 0 or v is None]
    
    if zeros:
        n_zeros = len(zeros)
        total_delta = delta * n_zeros
        non_zero_sum = sum(v for v in result if v and v > 0)
        
        if non_zero_sum > 0:
            scale = (1 - total_delta) / non_zero_sum
            for i, v in enumerate(result):
                if v and v > 0:
                    result[i] = v * scale
                else:
                    result[i] = delta
        else:
            # All zeros - distribute evenly
            result = [delta] * len(values)
    
    return result


def _compute_clr(values: List[float]) -> List[float]:
    """
    Compute centered log-ratio transformation.
    
    CLR_i = ln(x_i) - (1/n) * sum(ln(x_j))
    
    Args:
        values: Positive values (after zero replacement)
        
    Returns:
        CLR-transformed values
    """
    log_values = [math.log(v) for v in values]
    geo_mean_log = sum(log_values) / len(log_values)
    return [lv - geo_mean_log for lv in log_values]


def _aitchison_distance(clr1: Dict[str, float], clr2: Dict[str, float]) -> float:
    """
    Compute Aitchison distance between two CLR coordinate sets.
    
    d_A(x, y) = sqrt(sum((CLR_x_i - CLR_y_i)^2))
    """
    squared_diff_sum = 0.0
    for key in FIXED_TERPENE_ORDER:
        diff = clr1.get(key, 0.0) - clr2.get(key, 0.0)
        squared_diff_sum += diff * diff
    return math.sqrt(squared_diff_sum)


def compute_clr_coordinates(terpene_values: Dict[str, float]) -> Dict[str, float]:
    """
    Compute CLR coordinates for a terpene profile.
    
    Args:
        terpene_values: Dict mapping terpene name to concentration
        
    Returns:
        Dict mapping terpene name to CLR coordinate
    """
    # Extract values in fixed order
    values = [terpene_values.get(t, 0.0) for t in FIXED_TERPENE_ORDER]
    
    # Replace zeros
    prepared = _prepare_for_clr(values)
    
    # Compute CLR
    clr_values = _compute_clr(prepared)
    
    # Map back to dict
    return {t: clr_values[i] for i, t in enumerate(FIXED_TERPENE_ORDER)}


def compute_ratio_metrics(terpene_values: Dict[str, float]) -> TerpeneRatioMetrics:
    """
    Compute all CDES v1.5 ratio metrics for a terpene profile.
    
    Args:
        terpene_values: Dict mapping terpene name to concentration (%)
        
    Returns:
        TerpeneRatioMetrics with all computed values
    """
    # Get values in order
    values = [(t, terpene_values.get(t, 0.0)) for t in FIXED_TERPENE_ORDER]
    
    # Sort by concentration descending
    sorted_values = sorted(values, key=lambda x: x[1], reverse=True)
    
    # Total terpenes
    total = sum(v for _, v in values)
    if total == 0:
        total = 1.0  # Avoid division by zero
    
    # Create ranked terpenes
    ranked = []
    for i, (terpene, conc) in enumerate(sorted_values):
        ranked.append(RankedTerpene(
            rank=i + 1,
            terpene_id=f"terpene:{terpene}",
            name=TERPENE_NAMES.get(terpene, terpene),
            concentration=conc,
            unit="percent",
            percent_of_total=(conc / total) * 100 if total > 0 else 0.0,
        ))
    
    # T1, T2, T3 values
    t1 = sorted_values[0][1] if len(sorted_values) > 0 else 0.0
    t2 = sorted_values[1][1] if len(sorted_values) > 1 else 0.0
    t3 = sorted_values[2][1] if len(sorted_values) > 2 else 0.0
    t4 = sorted_values[3][1] if len(sorted_values) > 3 else 0.0
    
    # Dominance Index: (T1 - T2) / (T1 + T2)
    if t1 + t2 > 0:
        dominance_index = (t1 - t2) / (t1 + t2)
    else:
        dominance_index = 0.0
    
    # Dominance Level
    if dominance_index < 0.25:
        dominance_level = DominanceLevel.BALANCED
    elif dominance_index < 0.50:
        dominance_level = DominanceLevel.SLIGHT
    elif dominance_index < 0.75:
        dominance_level = DominanceLevel.MODERATE
    else:
        dominance_level = DominanceLevel.STRONG
    
    # Balance Score (Shannon Entropy)
    balance_score = 0.0
    for _, conc in values:
        if conc > 0 and total > 0:
            p = conc / total
            balance_score -= p * math.log(p)
    
    # Normalized balance score (0-1)
    n = len([v for _, v in values if v > 0])
    if n > 1:
        max_entropy = math.log(n)
        balance_score_normalized = balance_score / max_entropy if max_entropy > 0 else 0.0
    else:
        balance_score_normalized = 0.0
    
    # HHI (Herfindahl-Hirschman Index)
    hhi = sum((conc / total) ** 2 for _, conc in values if total > 0)
    
    # Top terpene count (> 5% of total)
    top_count = sum(1 for _, conc in values if (conc / total) * 100 > 5)
    
    # Ratio Signature
    t1_name = TERPENE_CODES.get(sorted_values[0][0], "UNK")
    t2_name = TERPENE_CODES.get(sorted_values[1][0], "UNK") if len(sorted_values) > 1 else "UNK"
    t3_name = TERPENE_CODES.get(sorted_values[2][0], "UNK") if len(sorted_values) > 2 else "UNK"
    
    t1_t2_ratio = t1 / t2 if t2 > 0 else float('inf')
    t1_t3_ratio = t1 / t3 if t3 > 0 else float('inf')
    t2_t3_ratio = t2 / t3 if t3 > 0 else float('inf')
    
    # Ratios as percent of T1
    r2 = t2 / t1 if t1 > 0 else 0.0
    r3 = t3 / t1 if t1 > 0 else 0.0
    r4 = t4 / t1 if t1 > 0 else 0.0
    
    encoded = f"{t1_name}:{t2_name}:{t3_name}|{r2:.2f}:{r3:.2f}:{r4:.2f}"
    
    ratio_signature = RatioSignature(
        encoded=encoded,
        t1_t2_ratio=t1_t2_ratio if t1_t2_ratio != float('inf') else 99.99,
        t1_t3_ratio=t1_t3_ratio if t1_t3_ratio != float('inf') else 99.99,
        t2_t3_ratio=t2_t3_ratio if t2_t3_ratio != float('inf') else 99.99,
    )
    
    # CLR Coordinates
    clr_coords = compute_clr_coordinates(terpene_values)
    
    # Key Ratios (pairwise log-ratios)
    key_ratios = {}
    myrcene = terpene_values.get("myrcene", 0.001)
    limonene = terpene_values.get("limonene", 0.001)
    pinene = terpene_values.get("pinene", 0.001)
    caryophyllene = terpene_values.get("caryophyllene", 0.001)
    humulene_val = terpene_values.get("humulene", 0.001)
    terpinolene = terpene_values.get("terpinolene", 0.001)
    
    # Replace zeros for log calculation
    myrcene = max(myrcene, 0.001)
    limonene = max(limonene, 0.001)
    pinene = max(pinene, 0.001)
    caryophyllene = max(caryophyllene, 0.001)
    humulene_val = max(humulene_val, 0.001)
    terpinolene = max(terpinolene, 0.001)
    
    key_ratios["myrceneLimoneneLr"] = round(math.log(myrcene / limonene), 3)
    key_ratios["myrcenePineneLr"] = round(math.log(myrcene / pinene), 3)
    key_ratios["caryophylleneHumuleneLr"] = round(math.log(caryophyllene / humulene_val), 3)
    key_ratios["limoneneTerpinoleneLr"] = round(math.log(limonene / terpinolene), 3)
    
    # Classify archetype
    archetype = classify_archetype(clr_coords)
    
    return TerpeneRatioMetrics(
        dominance_index=round(dominance_index, 4),
        dominance_level=dominance_level,
        balance_score=round(balance_score, 4),
        balance_score_normalized=round(balance_score_normalized, 4),
        hhi=round(hhi, 4),
        top_terpene_count=top_count,
        ranked_terpenes=ranked,
        ratio_signature=ratio_signature,
        clr_coordinates={k: round(v, 4) for k, v in clr_coords.items()},
        key_ratios=key_ratios,
        archetype_classification=archetype,
    )


def classify_archetype(clr_coordinates: Dict[str, float]) -> ArchetypeClassification:
    """
    Classify a terpene profile into an archetype based on CLR coordinates.
    
    Args:
        clr_coordinates: CLR-transformed coordinates
        
    Returns:
        ArchetypeClassification with primary archetype and confidence
    """
    distances = {}
    
    for archetype, centroid in ARCHETYPE_CENTROIDS.items():
        dist = _aitchison_distance(clr_coordinates, centroid)
        distances[archetype.value] = round(dist, 4)
    
    # Find minimum distance
    sorted_distances = sorted(distances.items(), key=lambda x: x[1])
    
    if len(sorted_distances) >= 2:
        primary_archetype = TerpeneArchetype(sorted_distances[0][0])
        primary_dist = sorted_distances[0][1]
        second_dist = sorted_distances[1][1]
        
        # Confidence: how much closer to primary than second
        # confidence = 1 - (primary_dist / second_dist) if second_dist > 0
        if second_dist > 0 and primary_dist < second_dist:
            confidence = 1 - (primary_dist / second_dist)
        else:
            confidence = 0.5
        
        # Clamp confidence to [0, 1]
        confidence = max(0.0, min(1.0, confidence))
    else:
        primary_archetype = TerpeneArchetype.ARCH_UNK
        confidence = 0.0
    
    return ArchetypeClassification(
        primary_archetype=primary_archetype,
        archetype_confidence=round(confidence, 4),
        archetype_distances=distances,
    )


# =============================================================================
# UTILITY FUNCTIONS
# =============================================================================

def get_archetype_description(archetype: TerpeneArchetype) -> Dict[str, Any]:
    """
    Get human-readable description of an archetype.
    
    Returns:
        Dict with name, description, typical_effects, signature
    """
    descriptions = {
        TerpeneArchetype.ARCH_MYR: {
            "name": "Myrcene Dominant",
            "description": "High myrcene content (>40%), classic indica-leaning profile",
            "typical_effects": ["sedating", "body-focused", "relaxing"],
            "signature_pattern": "MYR > 40%, others < 15% each",
        },
        TerpeneArchetype.ARCH_LIM: {
            "name": "Limonene Forward",
            "description": "Limonene-led profile with citrus notes",
            "typical_effects": ["uplifting", "energizing", "mood-enhancing"],
            "signature_pattern": "LIM > 25%, MYR secondary",
        },
        TerpeneArchetype.ARCH_BAL: {
            "name": "Balanced Trio",
            "description": "Top 3 terpenes within 80-100% of each other",
            "typical_effects": ["complex", "nuanced", "full-spectrum"],
            "signature_pattern": "Top 3 close in concentration",
        },
        TerpeneArchetype.ARCH_PIN: {
            "name": "Pinene Sharp",
            "description": "Pinene-dominant with cognitive clarity effects",
            "typical_effects": ["alert", "focused", "clear-headed"],
            "signature_pattern": "PIN > 15%, low MYR",
        },
        TerpeneArchetype.ARCH_CAR: {
            "name": "Caryophyllene Spicy",
            "description": "Caryophyllene-led with humulene support",
            "typical_effects": ["peppery", "therapeutic", "anti-inflammatory"],
            "signature_pattern": "CAR > 20%, HUM present",
        },
        TerpeneArchetype.ARCH_TER: {
            "name": "Terpinolene Rare",
            "description": "Unusual terpinolene-dominant profile",
            "typical_effects": ["unique", "energizing", "creative"],
            "signature_pattern": "TER dominant (rare)",
        },
        TerpeneArchetype.ARCH_LIN: {
            "name": "Linalool Floral",
            "description": "Linalool-forward with floral/lavender notes",
            "typical_effects": ["calming", "floral", "anxiety-reducing"],
            "signature_pattern": "LIN > 15%",
        },
        TerpeneArchetype.ARCH_OCI: {
            "name": "Ocimene Sweet",
            "description": "Ocimene-led sweet and herbaceous profile",
            "typical_effects": ["sweet", "herbaceous", "uplifting"],
            "signature_pattern": "OCI > 10%",
        },
        TerpeneArchetype.ARCH_UNK: {
            "name": "Unknown",
            "description": "Profile does not match known archetypes",
            "typical_effects": ["variable"],
            "signature_pattern": "N/A",
        },
    }
    return descriptions.get(archetype, descriptions[TerpeneArchetype.ARCH_UNK])
