"""
FedComp Index - Federal contractor posture classification.

Classifies federal contractors into four Posture Classes using two-axis
thresholds: base contract dollars ($5M) x base contract count (3).
No composite score. No weighted sum.

Class 1: high volume + high frequency (systematic winners)
Class 2: high volume + low frequency (concentrated risk)
Class 3: low volume + high frequency (growth pipeline)
Class 4: low volume + low frequency (entry level)

Methodology: https://fedcompindex.org/methodology/
Data: https://huggingface.co/datasets/npetro6/nevada-federal-contractors
"""

__version__ = "2026.4.0"

import math
from datetime import date, datetime
from enum import Enum
from collections import defaultdict


class PostureClass(Enum):
    """Federal contractor posture classification. Two-axis: volume x frequency."""
    CLASS_1 = "Class 1"  # high vol ($5M+) + high freq (3+ base contracts)
    CLASS_2 = "Class 2"  # high vol ($5M+) + low freq (<3 base contracts)
    CLASS_3 = "Class 3"  # low vol (<$5M) + high freq (3+ base contracts)
    CLASS_4 = "Class 4"  # low vol (<$5M) + low freq (<3 base contracts)


# Classification thresholds
VOL_THRESHOLD = 5_000_000    # $5M in base contract dollars over 5 years
FREQ_THRESHOLD = 3           # 3 distinct base contracts over 5 years


class FedCompResult:
    """Result of a FedComp Index classification."""

    def __init__(self, posture_class, base_dollars, base_contract_count,
                 velocity=None, proximity=None, neighborhood_density=None,
                 proximity_pressure=None):
        self.posture_class = posture_class
        self.base_dollars = base_dollars
        self.base_contract_count = base_contract_count
        self.velocity = velocity
        self.proximity = proximity or []
        self.neighborhood_density = neighborhood_density
        self.proximity_pressure = proximity_pressure

    def __repr__(self):
        return (
            f"FedCompResult(posture_class={self.posture_class.value}, "
            f"base_dollars={self.base_dollars}, "
            f"base_contract_count={self.base_contract_count})"
        )


def classify(base_dollars, base_contract_count):
    """
    Classify a contractor into a Posture Class using two-axis thresholds.

    Args:
        base_dollars: Total base contract dollars over trailing 5-year window.
        base_contract_count: Number of distinct base contracts over 5 years.

    Returns:
        PostureClass enum value.

    Example::

        from fedcomp_index_scoring import classify, PostureClass

        cls = classify(12_000_000, 5)
        assert cls == PostureClass.CLASS_1

        cls = classify(8_000_000, 1)
        assert cls == PostureClass.CLASS_2
    """
    high_vol = base_dollars >= VOL_THRESHOLD
    high_freq = base_contract_count >= FREQ_THRESHOLD
    if high_vol and high_freq:
        return PostureClass.CLASS_1
    elif high_vol and not high_freq:
        return PostureClass.CLASS_2
    elif not high_vol and high_freq:
        return PostureClass.CLASS_3
    return PostureClass.CLASS_4


def classify_contractor(base_dollars, base_contract_count):
    """
    Classify a contractor and return a full FedCompResult.

    Args:
        base_dollars: Total base contract dollars over trailing 5-year window.
        base_contract_count: Number of distinct base contracts over 5 years.

    Returns:
        FedCompResult with posture_class, base_dollars, base_contract_count.

    Example::

        from fedcomp_index_scoring import classify_contractor

        result = classify_contractor(12_000_000, 5)
        print(result.posture_class)         # PostureClass.CLASS_1
        print(result.base_dollars)          # 12000000
        print(result.base_contract_count)   # 5
    """
    return FedCompResult(
        posture_class=classify(base_dollars, base_contract_count),
        base_dollars=base_dollars,
        base_contract_count=base_contract_count,
    )


# ---- Velocity (cadence-based) ------------------------------------------------

def compute_velocity(award_dates):
    """
    Compute cadence-based velocity from a list of award dates.

    Requires 2+ dates. Returns None for fewer than 2.

    Velocity labels: accelerating, on pace, slowing, declining, inactive.

    Args:
        award_dates: List of date objects or "YYYY-MM-DD" strings.

    Returns:
        dict with keys: cadence (str or None), direction (str), dv, df, magnitude.
        Returns None if fewer than 2 dates.

    Example::

        from fedcomp_index_scoring import compute_velocity

        v = compute_velocity(["2024-01-15", "2024-06-01", "2025-02-10"])
        print(v["cadence"])   # "on pace" / "accelerating" / etc.
    """
    if not award_dates or len(award_dates) < 2:
        return None

    parsed = []
    for d in award_dates:
        if isinstance(d, str):
            d = datetime.strptime(d[:10], "%Y-%m-%d").date()
        parsed.append(d)
    parsed.sort()

    today = date.today()
    avg_gap = (parsed[-1] - parsed[0]).days / (len(parsed) - 1)
    if avg_gap == 0:
        avg_gap = 1
    days_since = (today - parsed[-1]).days
    gap_ratio = days_since / avg_gap

    mid_idx = len(parsed) // 2
    if mid_idx > 0 and mid_idx < len(parsed) - 1:
        early_gap = (parsed[mid_idx] - parsed[0]).days / mid_idx
        late_gap = (parsed[-1] - parsed[mid_idx]).days / (len(parsed) - mid_idx - 1)
        pace_ratio = late_gap / early_gap if early_gap > 0 else 1.0
    else:
        pace_ratio = 1.0

    if gap_ratio > 5:
        cadence = "inactive"
    elif pace_ratio < 0.7 and gap_ratio < 2:
        cadence = "accelerating"
    elif gap_ratio < 2 and pace_ratio < 1.5:
        cadence = "on pace"
    elif gap_ratio < 4 or pace_ratio < 2.5:
        cadence = "slowing"
    else:
        cadence = "declining"

    return {"cadence": cadence}


# ---- Proximity Map -----------------------------------------------------------

def build_proximity_map(target_codes, target_volume, candidates, top_n=6):
    """
    Find competitive neighbors using inverse-frequency-squared weighting on
    shared NAICS/PSC codes with scale filtering.

    Args:
        target_codes: dict mapping code -> count_of_contractors_sharing_that_code.
            Example: {"541511": 12, "R425": 3}
        target_volume: float, target contractor's base_dollars_5yr.
        candidates: list of dicts, each with keys:
            - codes: set of NAICS/PSC codes
            - volume: float base_dollars_5yr
            - id: any identifier
        top_n: number of neighbors to return (default 6).

    Returns:
        list of (id, score) tuples, sorted by score descending, top_n entries.

    The weighting is inverse-frequency-squared: rare shared codes produce
    disproportionately high proximity. Scale filter applies ratio^2 penalty
    for volume mismatches.

    Example::

        from fedcomp_index_scoring import build_proximity_map

        neighbors = build_proximity_map(
            target_codes={"541511": 12, "R425": 3},
            target_volume=8_000_000,
            candidates=[
                {"id": "ABC", "codes": {"541511", "R425"}, "volume": 6_000_000},
                {"id": "DEF", "codes": {"541511"}, "volume": 50_000_000},
            ],
        )
    """
    if not target_codes or not candidates:
        return []

    results = []
    for cand in candidates:
        shared = set(target_codes.keys()) & cand.get("codes", set())
        if not shared:
            continue
        overlap = sum(1.0 / (target_codes[c] ** 2) for c in shared)
        if overlap <= 0:
            continue

        cand_vol = cand.get("volume", 0)
        if target_volume > 0 and cand_vol > 0:
            ratio = min(target_volume, cand_vol) / max(target_volume, cand_vol)
            score = overlap * (ratio ** 2)
        else:
            score = overlap * 0.01

        results.append((cand["id"], score))

    results.sort(key=lambda x: x[1], reverse=True)
    return results[:top_n]


def neighborhood_density(neighbor_sets, target_neighbors):
    """
    Compute neighborhood density: fraction of a contractor's neighbors that
    are also neighbors of each other.

    Args:
        neighbor_sets: dict mapping contractor_id -> set of neighbor_ids.
        target_neighbors: set of neighbor_ids for the target contractor.

    Returns:
        float between 0.0 and 1.0. Higher = tighter competitive cluster.

    Example::

        from fedcomp_index_scoring import neighborhood_density

        sets = {
            "A": {"B", "C"},
            "B": {"A", "C"},
            "C": {"A"},
        }
        d = neighborhood_density(sets, {"A", "B", "C"})
    """
    if len(target_neighbors) < 2:
        return 0.0

    cross = 0
    possible = 0
    for n in target_neighbors:
        nn = neighbor_sets.get(n, set())
        cross += len(target_neighbors & nn - {n})
        possible += len(target_neighbors) - 1

    return round(cross / possible, 3) if possible > 0 else 0.0


# ---- Backward compatibility --------------------------------------------------
# score_contractor is kept as an alias that maps old-style calls to the new
# classification system. It returns a FedCompResult.

def score_contractor(total_awards_usd, last_award_date=None, base_contract_count=1):
    """
    Backward-compatible entry point. Classifies a contractor into a Posture Class.

    In v1.0, this computed a composite 0-100 score. In v1.1, classification is
    two-axis (volume x frequency) with no composite score.

    Args:
        total_awards_usd: Total base contract dollars over trailing 5-year window.
        last_award_date: Ignored in v1.1 (kept for API compatibility).
        base_contract_count: Number of distinct base contracts (default 1).

    Returns:
        FedCompResult with posture_class, base_dollars, base_contract_count.
    """
    return classify_contractor(total_awards_usd, base_contract_count)
